Allow saving "before-highres-fix. (#4150)

* Save image/s before doing highres fix.

3 files changed, 18 insertions, 5 deletions

diff --git a/modules/processing.py b/modules/processing.py
index b541ee2b..2dcf4879 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -521,7 +521,11 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 shared.state.job = f"Batch {n+1} out of {p.n_iter}"

 

             with devices.autocast():

-                samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength)

+                # Only Txt2Img needs an extra argument, n, when saving intermediate images pre highres fix. 

+                if isinstance(p, StableDiffusionProcessingTxt2Img):

+                    samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, n=n)

+                else:

+                    samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength)

 

             samples_ddim = samples_ddim.to(devices.dtype_vae)

             x_samples_ddim = decode_first_stage(p.sd_model, samples_ddim)

@@ -649,7 +653,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             self.truncate_x = int(self.firstphase_width - firstphase_width_truncated) // opt_f

             self.truncate_y = int(self.firstphase_height - firstphase_height_truncated) // opt_f

 

-    def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength):

+    def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, n=0):

         self.sampler = sd_samplers.create_sampler_with_index(sd_samplers.samplers, self.sampler_index, self.sd_model)

 

         if not self.enable_hr:

@@ -685,6 +689,15 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

             samples = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(decoded_samples))

 

+        # Save a copy of the image/s before doing highres fix, if applicable.

+        if opts.save and not self.do_not_save_samples and opts.save_images_before_highres_fix:

+            for i in range(self.batch_size):

+                # This batch's ith image.

+                img = sd_samplers.sample_to_image(samples, i)

+                # Index that accounts for both batch size and batch count.

+                ind = i + self.batch_size*n

+                images.save_image(img, self.outpath_samples, "", self.all_seeds[ind], self.all_prompts[ind], opts.samples_format, suffix=f"-before-highres-fix")

+

         shared.state.nextjob()

 

         self.sampler = sd_samplers.create_sampler_with_index(sd_samplers.samplers, self.sampler_index, self.sd_model)

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 44d4c189..d7fa89a0 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -93,9 +93,8 @@ def single_sample_to_image(sample):
     return Image.fromarray(x_sample)

 

 

-def sample_to_image(samples):

-    return single_sample_to_image(samples[0])

-

+def sample_to_image(samples, index=0):

+    return single_sample_to_image(samples[index])

 

 def samples_to_image_grid(samples):

     return images.image_grid([single_sample_to_image(sample) for sample in samples])

diff --git a/modules/shared.py b/modules/shared.py
index e65f6080..ce991424 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -255,6 +255,7 @@ options_templates.update(options_section(('saving-images', "Saving images/grids"
     "enable_pnginfo": OptionInfo(True, "Save text information about generation parameters as chunks to png files"),

     "save_txt": OptionInfo(False, "Create a text file next to every image with generation parameters."),

     "save_images_before_face_restoration": OptionInfo(False, "Save a copy of image before doing face restoration."),

+    "save_images_before_highres_fix": OptionInfo(False, "Save a copy of image before applying highres fix."),

     "jpeg_quality": OptionInfo(80, "Quality for saved jpeg images", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),

     "export_for_4chan": OptionInfo(True, "If PNG image is larger than 4MB or any dimension is larger than 4000, downscale and save copy as JPG"),