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,