Merge branch 'dev' into dev

1 files changed, 68 insertions, 25 deletions

diff --git a/modules/processing.py b/modules/processing.py
index 6567b3cf..8f34c8b4 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -14,7 +14,7 @@ from skimage import exposure
 from typing import Any, Dict, List

 

 import modules.sd_hijack

-from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet

+from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet, errors

 from modules.sd_hijack import model_hijack

 from modules.shared import opts, cmd_opts, state

 import modules.shared as shared

@@ -492,7 +492,7 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
         noise_shape = shape if seed_resize_from_h <= 0 or seed_resize_from_w <= 0 else (shape[0], seed_resize_from_h//8, seed_resize_from_w//8)

 

         subnoise = None

-        if subseeds is not None:

+        if subseeds is not None and subseed_strength != 0:

             subseed = 0 if i >= len(subseeds) else subseeds[i]

 

             subnoise = devices.randn(subseed, noise_shape)

@@ -524,7 +524,7 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
             cnt = p.sampler.number_of_needed_noises(p)

 

             if eta_noise_seed_delta > 0:

-                torch.manual_seed(seed + eta_noise_seed_delta)

+                devices.manual_seed(seed + eta_noise_seed_delta)

 

             for j in range(cnt):

                 sampler_noises[j].append(devices.randn_without_seed(tuple(noise_shape)))

@@ -538,6 +538,40 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
     return x

 

 

+def decode_latent_batch(model, batch, target_device=None, check_for_nans=False):

+    samples = []

+

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

+        sample = decode_first_stage(model, batch[i:i + 1])[0]

+

+        if check_for_nans:

+            try:

+                devices.test_for_nans(sample, "vae")

+            except devices.NansException as e:

+                if devices.dtype_vae == torch.float32 or not shared.opts.auto_vae_precision:

+                    raise e

+

+                errors.print_error_explanation(

+                    "A tensor with all NaNs was produced in VAE.\n"

+                    "Web UI will now convert VAE into 32-bit float and retry.\n"

+                    "To disable this behavior, disable the 'Automaticlly revert VAE to 32-bit floats' setting.\n"

+                    "To always start with 32-bit VAE, use --no-half-vae commandline flag."

+                )

+

+                devices.dtype_vae = torch.float32

+                model.first_stage_model.to(devices.dtype_vae)

+                batch = batch.to(devices.dtype_vae)

+

+                sample = decode_first_stage(model, batch[i:i + 1])[0]

+

+        if target_device is not None:

+            sample = sample.to(target_device)

+

+        samples.append(sample)

+

+    return samples

+

+

 def decode_first_stage(model, x):

     x = model.decode_first_stage(x.to(devices.dtype_vae))

 

@@ -566,8 +600,12 @@ def program_version():
     return res

 

 

-def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iteration=0, position_in_batch=0, use_main_prompt=False):

-    index = position_in_batch + iteration * p.batch_size

+def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iteration=0, position_in_batch=0, use_main_prompt=False, index=None, all_negative_prompts=None):

+    if index is None:

+        index = position_in_batch + iteration * p.batch_size

+

+    if all_negative_prompts is None:

+        all_negative_prompts = p.all_negative_prompts

 

     clip_skip = getattr(p, 'clip_skip', opts.CLIP_stop_at_last_layers)

     enable_hr = getattr(p, 'enable_hr', False)

@@ -583,12 +621,12 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Sampler": p.sampler_name,

         "CFG scale": p.cfg_scale,

         "Image CFG scale": getattr(p, 'image_cfg_scale', None),

-        "Seed": all_seeds[index],

+        "Seed": p.all_seeds[0] if use_main_prompt else all_seeds[index],

         "Face restoration": (opts.face_restoration_model if p.restore_faces else None),

         "Size": f"{p.width}x{p.height}",

         "Model hash": getattr(p, 'sd_model_hash', None if not opts.add_model_hash_to_info or not shared.sd_model.sd_model_hash else shared.sd_model.sd_model_hash),

-        "Model": (None if not opts.add_model_name_to_info or not shared.sd_model.sd_checkpoint_info.model_name else shared.sd_model.sd_checkpoint_info.model_name.replace(',', '').replace(':', '')),

-        "Variation seed": (None if p.subseed_strength == 0 else all_subseeds[index]),

+        "Model": (None if not opts.add_model_name_to_info else shared.sd_model.sd_checkpoint_info.name_for_extra),

+        "Variation seed": (None if p.subseed_strength == 0 else (p.all_subseeds[0] if use_main_prompt else all_subseeds[index])),

         "Variation seed strength": (None if p.subseed_strength == 0 else p.subseed_strength),

         "Seed resize from": (None if p.seed_resize_from_w <= 0 or p.seed_resize_from_h <= 0 else f"{p.seed_resize_from_w}x{p.seed_resize_from_h}"),

         "Denoising strength": getattr(p, 'denoising_strength', None),

@@ -598,7 +636,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
         "Token merging ratio": None if token_merging_ratio == 0 else token_merging_ratio,

         "Token merging ratio hr": None if not enable_hr or token_merging_ratio_hr == 0 else token_merging_ratio_hr,

         "Init image hash": getattr(p, 'init_img_hash', None),

-        "RNG": opts.randn_source if opts.randn_source != "GPU" else None,

+        "RNG": opts.randn_source if opts.randn_source != "GPU" and opts.randn_source != "NV" else None,

         "NGMS": None if p.s_min_uncond == 0 else p.s_min_uncond,

         **p.extra_generation_params,

         "Version": program_version() if opts.add_version_to_infotext else None,

@@ -608,7 +646,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
     generation_params_text = ", ".join([k if k == v else f'{k}: {generation_parameters_copypaste.quote(v)}' for k, v in generation_params.items() if v is not None])

 

     prompt_text = p.prompt if use_main_prompt else all_prompts[index]

-    negative_prompt_text = f"\nNegative prompt: {p.all_negative_prompts[index]}" if p.all_negative_prompts[index] else ""

+    negative_prompt_text = f"\nNegative prompt: {all_negative_prompts[index]}" if all_negative_prompts[index] else ""

 

     return f"{prompt_text}{negative_prompt_text}\n{generation_params_text}".strip()

 

@@ -682,9 +720,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     else:

         p.all_subseeds = [int(subseed) + x for x in range(len(p.all_prompts))]

 

-    def infotext(iteration=0, position_in_batch=0, use_main_prompt=False):

-        return create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, comments, iteration, position_in_batch, use_main_prompt)

-

     if os.path.exists(cmd_opts.embeddings_dir) and not p.do_not_reload_embeddings:

         model_hijack.embedding_db.load_textual_inversion_embeddings()

 

@@ -758,10 +793,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast():

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

 

-            x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]

-            for x in x_samples_ddim:

-                devices.test_for_nans(x, "vae")

-

+            x_samples_ddim = decode_latent_batch(p.sd_model, samples_ddim, target_device=devices.cpu, check_for_nans=True)

             x_samples_ddim = torch.stack(x_samples_ddim).float()

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

 

@@ -775,6 +807,16 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if p.scripts is not None:

                 p.scripts.postprocess_batch(p, x_samples_ddim, batch_number=n)

 

+                p.prompts = p.all_prompts[n * p.batch_size:(n + 1) * p.batch_size]

+                p.negative_prompts = p.all_negative_prompts[n * p.batch_size:(n + 1) * p.batch_size]

+

+                batch_params = scripts.PostprocessBatchListArgs(list(x_samples_ddim))

+                p.scripts.postprocess_batch_list(p, batch_params, batch_number=n)

+                x_samples_ddim = batch_params.images

+

+            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)

+

             for i, x_sample in enumerate(x_samples_ddim):

                 p.batch_index = i

 

@@ -783,7 +825,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 

                 if p.restore_faces:

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

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

+                        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()

 

@@ -800,15 +842,15 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 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:

                         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(n, i), p=p, suffix="-before-color-correction")

+                        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:

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

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

 

-                text = infotext(n, i)

+                text = infotext(i)

                 infotexts.append(text)

                 if opts.enable_pnginfo:

                     image.info["parameters"] = text

@@ -819,10 +861,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     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')

 

                     if opts.save_mask:

-                        images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-mask")

+                        images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-mask")

 

                     if opts.save_mask_composite:

-                        images.save_image(image_mask_composite, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-mask-composite")

+                        images.save_image(image_mask_composite, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p, suffix="-mask-composite")

 

                     if opts.return_mask:

                         output_images.append(image_mask)

@@ -863,7 +905,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         p,

         images_list=output_images,

         seed=p.all_seeds[0],

-        info=infotext(),

+        info=infotexts[0],

         comments="".join(f"{comment}\n" for comment in comments),

         subseed=p.all_subseeds[0],

         index_of_first_image=index_of_first_image,

@@ -1029,7 +1071,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
                 image = sd_samplers.sample_to_image(image, index, approximation=0)

 

             info = create_infotext(self, self.all_prompts, self.all_seeds, self.all_subseeds, [], iteration=self.iteration, position_in_batch=index)

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

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

 

         if latent_scale_mode is not None:

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

@@ -1306,6 +1348,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         image = image.to(shared.device, dtype=devices.dtype_vae)

 

         self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image))

+        devices.torch_gc()

 

         if self.resize_mode == 3:

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