Option to not save incomplete images

1 files changed, 11 insertions, 8 deletions

diff --git a/modules/processing.py b/modules/processing.py
index 43cb763f..bf4f938b 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -103,6 +103,10 @@ def txt2img_image_conditioning(sd_model, x, width, height):
         return x.new_zeros(x.shape[0], 5, 1, 1, dtype=x.dtype, device=x.device)

 

 

+def save_images_if_interrupted():

+    return not (opts.dont_save_interrupted_images and (state.interrupted or state.skipped))

+

+

 class StableDiffusionProcessing:

     """

     The first set of paramaters: sd_models -> do_not_reload_embeddings represent the minimum required to create a StableDiffusionProcessing

@@ -821,6 +825,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             def infotext(index=0, use_main_prompt=False):

                 return create_infotext(p, p.prompts, p.seeds, p.subseeds, use_main_prompt=use_main_prompt, index=index, all_negative_prompts=p.negative_prompts)

 

+            save_images_if_interrupt = save_images_if_interrupted()

+

             for i, x_sample in enumerate(x_samples_ddim):

                 p.batch_index = i

 

@@ -828,7 +834,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 x_sample = x_sample.astype(np.uint8)

 

                 if p.restore_faces:

-                    if opts.save and not p.do_not_save_samples and opts.save_images_before_face_restoration:

+                    if opts.save and not p.do_not_save_samples and opts.save_images_before_face_restoration and save_images_if_interrupt:

                         images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-before-face-restoration")

 

                     devices.torch_gc()

@@ -842,16 +848,15 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     pp = scripts.PostprocessImageArgs(image)

                     p.scripts.postprocess_image(p, pp)

                     image = pp.image

-

                 if p.color_corrections is not None and i < len(p.color_corrections):

-                    if opts.save and not p.do_not_save_samples and opts.save_images_before_color_correction:

+                    if opts.save and not p.do_not_save_samples and opts.save_images_before_color_correction and save_images_if_interrupt:

                         image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images)

                         images.save_image(image_without_cc, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-before-color-correction")

                     image = apply_color_correction(p.color_corrections[i], image)

 

                 image = apply_overlay(image, p.paste_to, i, p.overlay_images)

 

-                if opts.samples_save and not p.do_not_save_samples:

+                if opts.samples_save and not p.do_not_save_samples and save_images_if_interrupt:

                     images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p)

 

                 text = infotext(i)

@@ -859,8 +864,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 if opts.enable_pnginfo:

                     image.info["parameters"] = text

                 output_images.append(image)

-

-                if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([opts.save_mask, opts.save_mask_composite, opts.return_mask, opts.return_mask_composite]):

+                if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([opts.save_mask, opts.save_mask_composite, opts.return_mask, opts.return_mask_composite]) and save_images_if_interrupt:

                     image_mask = p.mask_for_overlay.convert('RGB')

                     image_mask_composite = Image.composite(image.convert('RGBA').convert('RGBa'), Image.new('RGBa', image.size), images.resize_image(2, p.mask_for_overlay, image.width, image.height).convert('L')).convert('RGBA')

 

@@ -896,7 +900,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     grid.info["parameters"] = text

                 output_images.insert(0, grid)

                 index_of_first_image = 1

-

             if opts.grid_save:

                 images.save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], opts.grid_format, info=infotext(use_main_prompt=True), short_filename=not opts.grid_extended_filename, p=p, grid=True)

 

@@ -1091,7 +1094,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         def save_intermediate(image, index):

             """saves image before applying hires fix, if enabled in options; takes as an argument either an image or batch with latent space images"""

 

-            if not opts.save or self.do_not_save_samples or not opts.save_images_before_highres_fix:

+            if not opts.save or self.do_not_save_samples or not opts.save_images_before_highres_fix or not save_images_if_interrupted():

                 return

 

             if not isinstance(image, Image.Image):