From 7a3eae457239119814ee6f43ec6dab5b5f0b7f6a Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Tue, 4 Oct 2022 17:48:16 -0400
Subject: [PATCH 1/2] revert to original k* noise schedule

---
 ldm/dream/generator/txt2img.py   |  2 +-
 ldm/models/diffusion/ksampler.py | 26 +++++++++++++++-----------
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/ldm/dream/generator/txt2img.py b/ldm/dream/generator/txt2img.py
index 790f098f84..bffe73745b 100644
--- a/ldm/dream/generator/txt2img.py
+++ b/ldm/dream/generator/txt2img.py
@@ -31,7 +31,7 @@ class Txt2Img(Generator):
             if self.free_gpu_mem and self.model.model.device != self.model.device:
                 self.model.model.to(self.model.device)
                                 
-            sampler.make_schedule(ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=True)
+            sampler.make_schedule(ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False)
 
             samples, _ = sampler.sample(
                 batch_size                   = 1,
diff --git a/ldm/models/diffusion/ksampler.py b/ldm/models/diffusion/ksampler.py
index d1064f06b8..ea17e6fe6b 100644
--- a/ldm/models/diffusion/ksampler.py
+++ b/ldm/models/diffusion/ksampler.py
@@ -45,6 +45,7 @@ class KSampler(Sampler):
             ddim_eta=0.0,
             verbose=False,
     ):
+        ddim_num_steps += 1
         outer_model = self.model
         self.model  = outer_model.inner_model
         super().make_schedule(
@@ -53,17 +54,19 @@ class KSampler(Sampler):
             ddim_eta=0.0,
             verbose=False,
         )
-        self.model = outer_model
+        self.model          = outer_model
         self.ddim_num_steps = ddim_num_steps
-        sigmas = K.sampling.get_sigmas_karras(
-            n=ddim_num_steps,
-            sigma_min=self.model.sigmas[0].item(),
-            sigma_max=self.model.sigmas[-1].item(),
-            rho=7.,
-            device=self.device,
-            # Birch-san recommends this, but it doesn't match the call signature in his branch of k-diffusion
-            # concat_zero=False
-        )
+        # not working quite right
+        # sigmas = K.sampling.get_sigmas_karras(
+        #     n=ddim_num_steps,
+        #     sigma_min=self.model.sigmas[0].item(),
+        #     sigma_max=self.model.sigmas[-1].item(),
+        #     rho=7.,
+        #     device=self.device,
+        #     # Birch-san recommends this, but it doesn't match the call signature in his branch of k-diffusion
+        #     # concat_zero=False
+        # )
+        sigmas = self.model.get_sigmas(ddim_num_steps)
         self.sigmas = sigmas
         
     # ALERT: We are completely overriding the sample() method in the base class, which
@@ -99,6 +102,7 @@ class KSampler(Sampler):
         # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
         **kwargs,
     ):
+        S += 1
         def route_callback(k_callback_values):
             if img_callback is not None:
                 img_callback(k_callback_values['x'],k_callback_values['i'])
@@ -119,7 +123,7 @@ class KSampler(Sampler):
             'uncond': unconditional_conditioning,
             'cond_scale': unconditional_guidance_scale,
         }
-        print(f'>> Sampling with k__{self.schedule}')
+        print(f'>> Sampling with k_{self.schedule}')
         return (
             K.sampling.__dict__[f'sample_{self.schedule}'](
                 model_wrap_cfg, x, sigmas, extra_args=extra_args,

From 483097f31c68648f0fc47da351ed786a37789dff Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Tue, 4 Oct 2022 18:02:12 -0400
Subject: [PATCH 2/2] fix off-by-one error

---
 ldm/models/diffusion/ksampler.py | 15 ++-------------
 1 file changed, 2 insertions(+), 13 deletions(-)

diff --git a/ldm/models/diffusion/ksampler.py b/ldm/models/diffusion/ksampler.py
index ea17e6fe6b..ffdb041eaf 100644
--- a/ldm/models/diffusion/ksampler.py
+++ b/ldm/models/diffusion/ksampler.py
@@ -45,7 +45,6 @@ class KSampler(Sampler):
             ddim_eta=0.0,
             verbose=False,
     ):
-        ddim_num_steps += 1
         outer_model = self.model
         self.model  = outer_model.inner_model
         super().make_schedule(
@@ -56,16 +55,6 @@ class KSampler(Sampler):
         )
         self.model          = outer_model
         self.ddim_num_steps = ddim_num_steps
-        # not working quite right
-        # sigmas = K.sampling.get_sigmas_karras(
-        #     n=ddim_num_steps,
-        #     sigma_min=self.model.sigmas[0].item(),
-        #     sigma_max=self.model.sigmas[-1].item(),
-        #     rho=7.,
-        #     device=self.device,
-        #     # Birch-san recommends this, but it doesn't match the call signature in his branch of k-diffusion
-        #     # concat_zero=False
-        # )
         sigmas = self.model.get_sigmas(ddim_num_steps)
         self.sigmas = sigmas
         
@@ -102,14 +91,14 @@ class KSampler(Sampler):
         # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
         **kwargs,
     ):
-        S += 1
         def route_callback(k_callback_values):
             if img_callback is not None:
                 img_callback(k_callback_values['x'],k_callback_values['i'])
 
         # sigmas = self.model.get_sigmas(S)
         # sigmas are now set up in make_schedule - we take the last steps items
-        sigmas = self.sigmas[-S:]
+        sigmas = self.sigmas[-S-1:]
+
         if x_T is not None:
             x = x_T * sigmas[0]
         else: