From c4fb8e304bb008714923aa36810518623cc1a766 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Thu, 13 Oct 2022 12:06:12 -0400
Subject: [PATCH] fix noisy images at high step counts

At step counts greater than ~75, the ksamplers start producing noisy
images when using the Karras noise schedule. This PR reverts to using
the model's own noise schedule, which eliminates the problem at the
cost of slowing convergence at lower step counts.

This PR also introduces a new CLI `--save_intermediates <n>' argument,
which will save every nth intermediate image into a subdirectory
named `intermediates/<image_prefix>'.

Addresses issue #1083.
---
 ldm/invoke/args.py               |  7 +++++++
 ldm/invoke/readline.py           |  1 +
 ldm/models/diffusion/ksampler.py |  3 ++-
 ldm/models/diffusion/sampler.py  |  2 +-
 scripts/invoke.py                | 13 +++++++++++++
 5 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/ldm/invoke/args.py b/ldm/invoke/args.py
index 09073679c03..5dbe885d2d2 100644
--- a/ldm/invoke/args.py
+++ b/ldm/invoke/args.py
@@ -636,6 +636,13 @@ def _create_dream_cmd_parser(self):
             dest='hires_fix',
             help='Create hires image using img2img to prevent duplicated objects'
         )
+        render_group.add_argument(
+            '--save_intermediates',
+            type=int,
+            default=0,
+            dest='save_intermediates',
+            help='Save every nth intermediate image into an "intermediates" directory within the output directory'
+        )
         img2img_group.add_argument(
             '-I',
             '--init_img',
diff --git a/ldm/invoke/readline.py b/ldm/invoke/readline.py
index 73664ef82ce..7975ec56434 100644
--- a/ldm/invoke/readline.py
+++ b/ldm/invoke/readline.py
@@ -31,6 +31,7 @@
     '--perlin',
     '--grid','-g',
     '--individual','-i',
+    '--save_intermediates',
     '--init_img','-I',
     '--init_mask','-M',
     '--init_color',
diff --git a/ldm/models/diffusion/ksampler.py b/ldm/models/diffusion/ksampler.py
index 0bc6ccd2968..ac0615b30cc 100644
--- a/ldm/models/diffusion/ksampler.py
+++ b/ldm/models/diffusion/ksampler.py
@@ -98,7 +98,8 @@ def make_schedule(
             rho=7.,
             device=self.device,
         )
-        self.sigmas = self.karras_sigmas
+        self.sigmas = self.model_sigmas
+        #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
diff --git a/ldm/models/diffusion/sampler.py b/ldm/models/diffusion/sampler.py
index 88cdc019740..ff705513f87 100644
--- a/ldm/models/diffusion/sampler.py
+++ b/ldm/models/diffusion/sampler.py
@@ -140,7 +140,7 @@ def sample(
         conditioning=None,
         callback=None,
         normals_sequence=None,
-        img_callback=None,
+        img_callback=None,   # TODO: this is very confusing because it is called "step_callback" elsewhere. Change.
         quantize_x0=False,
         eta=0.0,
         mask=None,
diff --git a/scripts/invoke.py b/scripts/invoke.py
index 100ab2413b5..3752680d4b8 100644
--- a/scripts/invoke.py
+++ b/scripts/invoke.py
@@ -289,6 +289,7 @@ def main_loop(gen, opt, infile):
             grid_images      = dict()  # seed -> Image, only used if `opt.grid`
             prior_variations = opt.with_variations or []
             prefix = file_writer.unique_prefix()
+            step_callback = make_step_callback(gen, opt, prefix) if opt.save_intermediates > 0 else None
 
             def image_writer(image, seed, upscaled=False, first_seed=None, use_prefix=None):
                 # note the seed is the seed of the current image
@@ -350,6 +351,7 @@ def image_writer(image, seed, upscaled=False, first_seed=None, use_prefix=None):
                 opt.last_operation='generate'
                 gen.prompt2image(
                     image_callback=image_writer,
+                    step_callback=step_callback,
                     catch_interrupts=catch_ctrl_c,
                     **vars(opt)
                 )
@@ -547,6 +549,17 @@ def split_variations(variations_string) -> list:
     else:
         return parts
 
+def make_step_callback(gen, opt, prefix):
+    destination = os.path.join(opt.outdir,'intermediates',prefix)
+    os.makedirs(destination,exist_ok=True)
+    print(f'>> Intermediate images will be written into {destination}')
+    def callback(img, step):
+        if step % opt.save_intermediates == 0 or step == opt.steps-1:
+            filename = os.path.join(destination,f'{step:04}.png')
+            image = gen.sample_to_image(img)
+            image.save(filename,'PNG')
+    return callback
+    
 def retrieve_dream_command(opt,file_path,completer):
     '''
     Given a full or partial path to a previously-generated image file,