Merge branch 'AUTOMATIC1111:master' into master

1 files changed, 24 insertions, 6 deletions

diff --git a/modules/processing.py b/modules/processing.py
index 512c484f..2168208c 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -191,9 +191,13 @@ class StableDiffusionProcessing():
     def init(self, all_prompts, all_seeds, all_subseeds):

         pass

 

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

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

         raise NotImplementedError()

 

+    def close(self):

+        self.sd_model = None

+        self.sampler = None

+

 

 class Processed:

     def __init__(self, p: StableDiffusionProcessing, images_list, seed=-1, info="", subseed=None, all_prompts=None, all_seeds=None, all_subseeds=None, index_of_first_image=0, infotexts=None):

@@ -509,7 +513,7 @@ 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)

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

 

             samples_ddim = samples_ddim.to(devices.dtype_vae)

             x_samples_ddim = decode_first_stage(p.sd_model, samples_ddim)

@@ -637,7 +641,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, prompts):

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

 

         if not self.enable_hr:

@@ -650,6 +654,16 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

         samples = samples[:, :, self.truncate_y//2:samples.shape[2]-self.truncate_y//2, self.truncate_x//2:samples.shape[3]-self.truncate_x//2]

 

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

+        def save_intermediate(image, index):

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

+                return

+

+            if not isinstance(image, Image.Image):

+                image = sd_samplers.sample_to_image(image, index)

+

+            images.save_image(image, self.outpath_samples, "", seeds[index], prompts[index], opts.samples_format, suffix="-before-highres-fix")

+

         if opts.use_scale_latent_for_hires_fix:

             samples = torch.nn.functional.interpolate(samples, size=(self.height // opt_f, self.width // opt_f), mode="bilinear")

             

@@ -660,6 +674,8 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             else:

                 image_conditioning = self.txt2img_image_conditioning(samples)

 

+            for i in range(samples.shape[0]):

+                save_intermediate(samples, i)

         else:

             decoded_samples = decode_first_stage(self.sd_model, samples)

             lowres_samples = torch.clamp((decoded_samples + 1.0) / 2.0, min=0.0, max=1.0)

@@ -669,6 +685,9 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
                 x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)

                 x_sample = x_sample.astype(np.uint8)

                 image = Image.fromarray(x_sample)

+

+                save_intermediate(image, i)

+

                 image = images.resize_image(0, image, self.width, self.height)

                 image = np.array(image).astype(np.float32) / 255.0

                 image = np.moveaxis(image, 2, 0)

@@ -826,8 +845,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
 

         self.image_conditioning = self.img2img_image_conditioning(image, self.init_latent, self.image_mask)

 

-

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

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

         x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

 

         samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning, image_conditioning=self.image_conditioning)

@@ -838,4 +856,4 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         del x

         devices.torch_gc()

 

-        return samples
\ No newline at end of file
+        return samples