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 09073679c0..5dbe885d2d 100644
--- a/ldm/invoke/args.py
+++ b/ldm/invoke/args.py
@@ -636,6 +636,13 @@ class Args(object):
             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 73664ef82c..7975ec5643 100644
--- a/ldm/invoke/readline.py
+++ b/ldm/invoke/readline.py
@@ -31,6 +31,7 @@ COMMANDS = (
     '--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 0bc6ccd296..ac0615b30c 100644
--- a/ldm/models/diffusion/ksampler.py
+++ b/ldm/models/diffusion/ksampler.py
@@ -98,7 +98,8 @@ class KSampler(Sampler):
             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 88cdc01974..ff705513f8 100644
--- a/ldm/models/diffusion/sampler.py
+++ b/ldm/models/diffusion/sampler.py
@@ -140,7 +140,7 @@ class Sampler(object):
         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 100ab2413b..3752680d4b 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 main_loop(gen, opt, infile):
                 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,