fix noisy images at high step counts

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',

ldm/invoke/readline.py

@@ -31,6 +31,7 @@
     '--perlin',
     '--grid','-g',
     '--individual','-i',
+    '--save_intermediates',
     '--init_img','-I',
     '--init_mask','-M',
     '--init_color',

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

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,

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,