update sd training scripts

This commit is contained in:
Artiprocher
2026-04-24 14:30:09 +08:00
parent 5cdab9ed01
commit 3799bdc23a
23 changed files with 323 additions and 612 deletions

View File

@@ -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}