mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-04-24 06:46:13 +00:00
update sd training scripts
This commit is contained in:
@@ -196,19 +196,14 @@ class SDUnit_InputImageEmbedder(PipelineUnit):
|
||||
|
||||
def process(self, pipe: StableDiffusionPipeline, input_image, noise):
|
||||
if input_image is None:
|
||||
return {"latents": noise * pipe.scheduler.init_noise_sigma, "input_latents": None}
|
||||
return {"latents": noise}
|
||||
pipe.load_models_to_device(self.onload_model_names)
|
||||
input_tensor = pipe.preprocess_image(input_image)
|
||||
input_latents = pipe.vae.encode(input_tensor).sample() * pipe.vae.scaling_factor
|
||||
latents = pipe.scheduler.add_noise(input_latents, noise, pipe.scheduler.timesteps[0])
|
||||
if pipe.scheduler.training:
|
||||
pipe.load_models_to_device(self.onload_model_names)
|
||||
input_tensor = pipe.preprocess_image(input_image)
|
||||
input_latents = pipe.vae.encode(input_tensor).sample()
|
||||
latents = noise * pipe.scheduler.init_noise_sigma
|
||||
return {"latents": latents, "input_latents": input_latents}
|
||||
else:
|
||||
# Inference mode: VAE encode input image, add noise for initial latent
|
||||
pipe.load_models_to_device(self.onload_model_names)
|
||||
input_tensor = pipe.preprocess_image(input_image)
|
||||
input_latents = pipe.vae.encode(input_tensor).sample()
|
||||
latents = pipe.scheduler.add_noise(input_latents, noise, pipe.scheduler.timesteps[0])
|
||||
return {"latents": latents}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user