From 5157cbeda17de4231d46b2330d602d6b3ccace79 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Fri, 7 Oct 2022 14:43:59 -0400 Subject: [PATCH] restore ability of ksamplers to process -v variation options - supersedes #977 --- ldm/dream/generator/img2img.py | 1 + ldm/models/diffusion/ksampler.py | 14 ++++++++++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/ldm/dream/generator/img2img.py b/ldm/dream/generator/img2img.py index 09750b3748..27425a2b6b 100644 --- a/ldm/dream/generator/img2img.py +++ b/ldm/dream/generator/img2img.py @@ -49,6 +49,7 @@ class Img2Img(Generator): img_callback = step_callback, unconditional_guidance_scale=cfg_scale, unconditional_conditioning=uc, + init_latent = self.init_latent, # changes how noising is performed in ksampler ) return self.sample_to_image(samples) diff --git a/ldm/models/diffusion/ksampler.py b/ldm/models/diffusion/ksampler.py index 6692907d2d..3d84ba0c9a 100644 --- a/ldm/models/diffusion/ksampler.py +++ b/ldm/models/diffusion/ksampler.py @@ -97,6 +97,7 @@ class KSampler(Sampler): rho=7., device=self.device, ) + self.sigmas = self.karras_sigmas # ALERT: We are completely overriding the sample() method in the base class, which # means that inpainting will not work. To get this to work we need to be able to @@ -170,11 +171,16 @@ class KSampler(Sampler): img_callback(k_callback_values['x'],k_callback_values['i']) # sigmas are set up in make_schedule - we take the last steps items - total_steps = len(self.karras_sigmas) - sigmas = self.karras_sigmas[-S-1:] - + total_steps = len(self.sigmas) + sigmas = self.sigmas[-S-1:] + + # x_T is variation noise. When an init image is provided (in x0) we need to add + # more randomness to the starting image. if x_T is not None: - x = x_T + torch.randn([batch_size, *shape], device=self.device) * sigmas[0] + if x0 is not None: + x = x_T + torch.randn_like(x0, device=self.device) * sigmas[0] + else: + x = x_T * sigmas[0] else: x = torch.randn([batch_size, *shape], device=self.device) * sigmas[0]