Merge branch 'master' into inpaint_textual_inversion

1 files changed, 208 insertions, 137 deletions

diff --git a/modules/processing.py b/modules/processing.py
index 57d3a523..c03e77e7 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -2,6 +2,7 @@ import json
 import math

 import os

 import sys

+import warnings

 

 import torch

 import numpy as np

@@ -12,15 +13,21 @@ from skimage import exposure
 from typing import Any, Dict, List, Optional

 

 import modules.sd_hijack

-from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste

+from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, script_callbacks

 from modules.sd_hijack import model_hijack

 from modules.shared import opts, cmd_opts, state

 import modules.shared as shared

 import modules.face_restoration

 import modules.images as images

 import modules.styles

+import modules.sd_models as sd_models

+import modules.sd_vae as sd_vae

 import logging

+from ldm.data.util import AddMiDaS

+from ldm.models.diffusion.ddpm import LatentDepth2ImageDiffusion

 

+from einops import repeat, rearrange

+from blendmodes.blend import blendLayers, BlendType

 

 # some of those options should not be changed at all because they would break the model, so I removed them from options.

 opt_C = 4

@@ -33,17 +40,19 @@ def setup_color_correction(image):
     return correction_target

 

 

-def apply_color_correction(correction, image):

+def apply_color_correction(correction, original_image):

     logging.info("Applying color correction.")

     image = Image.fromarray(cv2.cvtColor(exposure.match_histograms(

         cv2.cvtColor(

-            np.asarray(image),

+            np.asarray(original_image),

             cv2.COLOR_RGB2LAB

         ),

         correction,

         channel_axis=2

     ), cv2.COLOR_LAB2RGB).astype("uint8"))

-

+    

+    image = blendLayers(image, original_image, BlendType.LUMINOSITY)

+    

     return image

 

 

@@ -66,19 +75,15 @@ def apply_overlay(image, paste_loc, index, overlays):
 

     return image

 

-def get_correct_sampler(p):

-    if isinstance(p, modules.processing.StableDiffusionProcessingTxt2Img):

-        return sd_samplers.samplers

-    elif isinstance(p, modules.processing.StableDiffusionProcessingImg2Img):

-        return sd_samplers.samplers_for_img2img

-    elif isinstance(p, modules.api.processing.StableDiffusionProcessingAPI):

-        return sd_samplers.samplers

 

 class StableDiffusionProcessing():

     """

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

     """

-    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_index: int = 0, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, width: int = 512, height: int = 512, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, ddim_discretize: str = None, s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0, override_settings: Dict[str, Any] = None):

+    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, width: int = 512, height: int = 512, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, ddim_discretize: str = None, s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None):

+        if sampler_index is not None:

+            print("sampler_index argument for StableDiffusionProcessing does not do anything; use sampler_name", file=sys.stderr)

+

         self.sd_model = sd_model

         self.outpath_samples: str = outpath_samples

         self.outpath_grids: str = outpath_grids

@@ -91,7 +96,7 @@ class StableDiffusionProcessing():
         self.subseed_strength: float = subseed_strength

         self.seed_resize_from_h: int = seed_resize_from_h

         self.seed_resize_from_w: int = seed_resize_from_w

-        self.sampler_index: int = sampler_index

+        self.sampler_name: str = sampler_name

         self.batch_size: int = batch_size

         self.n_iter: int = n_iter

         self.steps: int = steps

@@ -116,6 +121,8 @@ class StableDiffusionProcessing():
         self.s_tmax = s_tmax or float('inf')  # not representable as a standard ui option

         self.s_noise = s_noise or opts.s_noise

         self.override_settings = {k: v for k, v in (override_settings or {}).items() if k not in shared.restricted_opts}

+        self.override_settings_restore_afterwards = override_settings_restore_afterwards

+        self.is_using_inpainting_conditioning = False

 

         if not seed_enable_extras:

             self.subseed = -1

@@ -126,41 +133,54 @@ class StableDiffusionProcessing():
         self.scripts = None

         self.script_args = None

         self.all_prompts = None

+        self.all_negative_prompts = None

         self.all_seeds = None

         self.all_subseeds = None

+        self.iteration = 0

 

     def txt2img_image_conditioning(self, x, width=None, height=None):

         if self.sampler.conditioning_key not in {'hybrid', 'concat'}:

             # Dummy zero conditioning if we're not using inpainting model.

             # Still takes up a bit of memory, but no encoder call.

             # Pretty sure we can just make this a 1x1 image since its not going to be used besides its batch size.

-            return torch.zeros(

-                x.shape[0], 5, 1, 1, 

-                dtype=x.dtype, 

-                device=x.device

-            )

+            return x.new_zeros(x.shape[0], 5, 1, 1)

+

+        self.is_using_inpainting_conditioning = True

 

         height = height or self.height

         width = width or self.width

 

         # The "masked-image" in this case will just be all zeros since the entire image is masked.

         image_conditioning = torch.zeros(x.shape[0], 3, height, width, device=x.device)

-        image_conditioning = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image_conditioning)) 

+        image_conditioning = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image_conditioning))

 

         # Add the fake full 1s mask to the first dimension.

         image_conditioning = torch.nn.functional.pad(image_conditioning, (0, 0, 0, 0, 1, 0), value=1.0)

-        image_conditioning = image_conditioning.to(x.dtype)            

+        image_conditioning = image_conditioning.to(x.dtype)

 

         return image_conditioning

 

-    def img2img_image_conditioning(self, source_image, latent_image, image_mask = None):

-        if self.sampler.conditioning_key not in {'hybrid', 'concat'}:

-            # Dummy zero conditioning if we're not using inpainting model.

-            return torch.zeros(

-                latent_image.shape[0], 5, 1, 1,

-                dtype=latent_image.dtype,

-                device=latent_image.device

-            )

+    def depth2img_image_conditioning(self, source_image):

+        # Use the AddMiDaS helper to Format our source image to suit the MiDaS model

+        transformer = AddMiDaS(model_type="dpt_hybrid")

+        transformed = transformer({"jpg": rearrange(source_image[0], "c h w -> h w c")})

+        midas_in = torch.from_numpy(transformed["midas_in"][None, ...]).to(device=shared.device)

+        midas_in = repeat(midas_in, "1 ... -> n ...", n=self.batch_size)

+

+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(source_image))

+        conditioning = torch.nn.functional.interpolate(

+            self.sd_model.depth_model(midas_in),

+            size=conditioning_image.shape[2:],

+            mode="bicubic",

+            align_corners=False,

+        )

+

+        (depth_min, depth_max) = torch.aminmax(conditioning)

+        conditioning = 2. * (conditioning - depth_min) / (depth_max - depth_min) - 1.

+        return conditioning

+

+    def inpainting_image_conditioning(self, source_image, latent_image, image_mask = None):

+        self.is_using_inpainting_conditioning = True

 

         # Handle the different mask inputs

         if image_mask is not None:

@@ -174,17 +194,17 @@ class StableDiffusionProcessing():
                 # Inpainting model uses a discretized mask as input, so we round to either 1.0 or 0.0

                 conditioning_mask = torch.round(conditioning_mask)

         else:

-            conditioning_mask = torch.ones(1, 1, *source_image.shape[-2:])

+            conditioning_mask = source_image.new_ones(1, 1, *source_image.shape[-2:])

 

         # Create another latent image, this time with a masked version of the original input.

         # Smoothly interpolate between the masked and unmasked latent conditioning image using a parameter.

-        conditioning_mask = conditioning_mask.to(source_image.device)

+        conditioning_mask = conditioning_mask.to(source_image.device).to(source_image.dtype)

         conditioning_image = torch.lerp(

             source_image,

             source_image * (1.0 - conditioning_mask),

             getattr(self, "inpainting_mask_weight", shared.opts.inpainting_mask_weight)

         )

-        

+

         # Encode the new masked image using first stage of network.

         conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(conditioning_image))

 

@@ -196,15 +216,31 @@ class StableDiffusionProcessing():
 

         return image_conditioning

 

+    def img2img_image_conditioning(self, source_image, latent_image, image_mask=None):

+        # HACK: Using introspection as the Depth2Image model doesn't appear to uniquely

+        # identify itself with a field common to all models. The conditioning_key is also hybrid.

+        if isinstance(self.sd_model, LatentDepth2ImageDiffusion):

+            return self.depth2img_image_conditioning(source_image)

+

+        if self.sampler.conditioning_key in {'hybrid', 'concat'}:

+            return self.inpainting_image_conditioning(source_image, latent_image, image_mask=image_mask)

+

+        # Dummy zero conditioning if we're not using inpainting or depth model.

+        return latent_image.new_zeros(latent_image.shape[0], 5, 1, 1)

+

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

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

         self.images = images_list

         self.prompt = p.prompt

         self.negative_prompt = p.negative_prompt

@@ -212,10 +248,10 @@ class Processed:
         self.subseed = subseed

         self.subseed_strength = p.subseed_strength

         self.info = info

+        self.comments = comments

         self.width = p.width

         self.height = p.height

-        self.sampler_index = p.sampler_index

-        self.sampler = sd_samplers.samplers[p.sampler_index].name

+        self.sampler_name = p.sampler_name

         self.cfg_scale = p.cfg_scale

         self.steps = p.steps

         self.batch_size = p.batch_size

@@ -242,17 +278,20 @@ class Processed:
         self.negative_prompt = self.negative_prompt if type(self.negative_prompt) != list else self.negative_prompt[0]

         self.seed = int(self.seed if type(self.seed) != list else self.seed[0]) if self.seed is not None else -1

         self.subseed = int(self.subseed if type(self.subseed) != list else self.subseed[0]) if self.subseed is not None else -1

+        self.is_using_inpainting_conditioning = p.is_using_inpainting_conditioning

 

-        self.all_prompts = all_prompts or [self.prompt]

-        self.all_seeds = all_seeds or [self.seed]

-        self.all_subseeds = all_subseeds or [self.subseed]

+        self.all_prompts = all_prompts or p.all_prompts or [self.prompt]

+        self.all_negative_prompts = all_negative_prompts or p.all_negative_prompts or [self.negative_prompt]

+        self.all_seeds = all_seeds or p.all_seeds or [self.seed]

+        self.all_subseeds = all_subseeds or p.all_subseeds or [self.subseed]

         self.infotexts = infotexts or [info]

 

     def js(self):

         obj = {

-            "prompt": self.prompt,

+            "prompt": self.all_prompts[0],

             "all_prompts": self.all_prompts,

-            "negative_prompt": self.negative_prompt,

+            "negative_prompt": self.all_negative_prompts[0],

+            "all_negative_prompts": self.all_negative_prompts,

             "seed": self.seed,

             "all_seeds": self.all_seeds,

             "subseed": self.subseed,

@@ -260,8 +299,7 @@ class Processed:
             "subseed_strength": self.subseed_strength,

             "width": self.width,

             "height": self.height,

-            "sampler_index": self.sampler_index,

-            "sampler": self.sampler,

+            "sampler_name": self.sampler_name,

             "cfg_scale": self.cfg_scale,

             "steps": self.steps,

             "batch_size": self.batch_size,

@@ -277,11 +315,12 @@ class Processed:
             "styles": self.styles,

             "job_timestamp": self.job_timestamp,

             "clip_skip": self.clip_skip,

+            "is_using_inpainting_conditioning": self.is_using_inpainting_conditioning,

         }

 

         return json.dumps(obj)

 

-    def infotext(self,  p: StableDiffusionProcessing, index):

+    def infotext(self, p: StableDiffusionProcessing, index):

         return create_infotext(p, self.all_prompts, self.all_seeds, self.all_subseeds, comments=[], position_in_batch=index % self.batch_size, iteration=index // self.batch_size)

 

 

@@ -301,13 +340,14 @@ def slerp(val, low, high):
 

 

 def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, seed_resize_from_h=0, seed_resize_from_w=0, p=None):

+    eta_noise_seed_delta = opts.eta_noise_seed_delta or 0

     xs = []

 

     # if we have multiple seeds, this means we are working with batch size>1; this then

     # enables the generation of additional tensors with noise that the sampler will use during its processing.

     # Using those pre-generated tensors instead of simple torch.randn allows a batch with seeds [100, 101] to

     # produce the same images as with two batches [100], [101].

-    if p is not None and p.sampler is not None and (len(seeds) > 1 and opts.enable_batch_seeds or opts.eta_noise_seed_delta > 0):

+    if p is not None and p.sampler is not None and (len(seeds) > 1 and opts.enable_batch_seeds or eta_noise_seed_delta > 0):

         sampler_noises = [[] for _ in range(p.sampler.number_of_needed_noises(p))]

     else:

         sampler_noises = None

@@ -347,8 +387,8 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
         if sampler_noises is not None:

             cnt = p.sampler.number_of_needed_noises(p)

 

-            if opts.eta_noise_seed_delta > 0:

-                torch.manual_seed(seed + opts.eta_noise_seed_delta)

+            if eta_noise_seed_delta > 0:

+                torch.manual_seed(seed + eta_noise_seed_delta)

 

             for j in range(cnt):

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

@@ -388,7 +428,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
 

     generation_params = {

         "Steps": p.steps,

-        "Sampler": get_correct_sampler(p)[p.sampler_index].name,

+        "Sampler": p.sampler_name,

         "CFG scale": p.cfg_scale,

         "Seed": all_seeds[index],

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

@@ -396,6 +436,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
         "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(':', '')),

         "Hypernet": (None if shared.loaded_hypernetwork is None else shared.loaded_hypernetwork.name),

+        "Hypernet hash": (None if shared.loaded_hypernetwork is None else sd_models.model_hash(shared.loaded_hypernetwork.filename)),

         "Hypernet strength": (None if shared.loaded_hypernetwork is None or shared.opts.sd_hypernetwork_strength >= 1 else shared.opts.sd_hypernetwork_strength),

         "Batch size": (None if p.batch_size < 2 else p.batch_size),

         "Batch pos": (None if p.batch_size < 2 else position_in_batch),

@@ -403,6 +444,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
         "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),

+        "Conditional mask weight": getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None,

         "Eta": (None if p.sampler is None or p.sampler.eta == p.sampler.default_eta else p.sampler.eta),

         "Clip skip": None if clip_skip <= 1 else clip_skip,

         "ENSD": None if opts.eta_noise_seed_delta == 0 else opts.eta_noise_seed_delta,

@@ -412,7 +454,7 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
 

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

 

-    negative_prompt_text = "\nNegative prompt: " + p.negative_prompt if p.negative_prompt else ""

+    negative_prompt_text = "\nNegative prompt: " + p.all_negative_prompts[index] if p.all_negative_prompts[index] else ""

 

     return f"{all_prompts[index]}{negative_prompt_text}\n{generation_params_text}".strip()

 

@@ -422,13 +464,21 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 

     try:

         for k, v in p.override_settings.items():

-            opts.data[k] = v  # we don't call onchange for simplicity which makes changing model, hypernet impossible

+            setattr(opts, k, v)

+            if k == 'sd_hypernetwork': shared.reload_hypernetworks()  # make onchange call for changing hypernet

+            if k == 'sd_model_checkpoint': sd_models.reload_model_weights()  # make onchange call for changing SD model

+            if k == 'sd_vae': sd_vae.reload_vae_weights()  # make onchange call for changing VAE

 

         res = process_images_inner(p)

 

     finally:

-        for k, v in stored_opts.items():

-            opts.data[k] = v

+        # restore opts to original state

+        if p.override_settings_restore_afterwards:

+            for k, v in stored_opts.items():

+                setattr(opts, k, v)

+                if k == 'sd_hypernetwork': shared.reload_hypernetworks()

+                if k == 'sd_model_checkpoint': sd_models.reload_model_weights()

+                if k == 'sd_vae': sd_vae.reload_vae_weights()

 

     return res

 

@@ -441,10 +491,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     else:

         assert p.prompt is not None

 

-    with open(os.path.join(shared.script_path, "params.txt"), "w", encoding="utf8") as file:

-        processed = Processed(p, [], p.seed, "")

-        file.write(processed.infotext(p, 0))

-

     devices.torch_gc()

 

     seed = get_fixed_seed(p.seed)

@@ -455,12 +501,15 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 

     comments = {}

 

-    shared.prompt_styles.apply_styles(p)

-

     if type(p.prompt) == list:

-        p.all_prompts = p.prompt

+        p.all_prompts = [shared.prompt_styles.apply_styles_to_prompt(x, p.styles) for x in p.prompt]

+    else:

+        p.all_prompts = p.batch_size * p.n_iter * [shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles)]

+

+    if type(p.negative_prompt) == list:

+        p.all_negative_prompts = [shared.prompt_styles.apply_negative_styles_to_prompt(x, p.styles) for x in p.negative_prompt]

     else:

-        p.all_prompts = p.batch_size * p.n_iter * [p.prompt]

+        p.all_negative_prompts = p.batch_size * p.n_iter * [shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles)]

 

     if type(seed) == list:

         p.all_seeds = seed

@@ -475,6 +524,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     def infotext(iteration=0, position_in_batch=0):

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

 

+    with open(os.path.join(shared.script_path, "params.txt"), "w", encoding="utf8") as file:

+        processed = Processed(p, [], p.seed, "")

+        file.write(processed.infotext(p, 0))

+

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

         model_hijack.embedding_db.load_textual_inversion_embeddings()

 

@@ -492,21 +545,27 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             state.job_count = p.n_iter

 

         for n in range(p.n_iter):

+            p.iteration = n

+

             if state.skipped:

                 state.skipped = False

-            

+

             if state.interrupted:

                 break

 

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

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

             seeds = p.all_seeds[n * p.batch_size:(n + 1) * p.batch_size]

             subseeds = p.all_subseeds[n * p.batch_size:(n + 1) * p.batch_size]

 

             if len(prompts) == 0:

                 break

 

+            if p.scripts is not None:

+                p.scripts.process_batch(p, batch_number=n, prompts=prompts, seeds=seeds, subseeds=subseeds)

+

             with devices.autocast():

-                uc = prompt_parser.get_learned_conditioning(shared.sd_model, len(prompts) * [p.negative_prompt], p.steps)

+                uc = prompt_parser.get_learned_conditioning(shared.sd_model, negative_prompts, p.steps)

                 c = prompt_parser.get_multicond_learned_conditioning(shared.sd_model, prompts, p.steps)

 

             if len(model_hijack.comments) > 0:

@@ -517,10 +576,10 @@ 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)

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

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

 

             del samples_ddim

@@ -530,9 +589,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 

             devices.torch_gc()

 

-            if opts.filter_nsfw:

-                import modules.safety as safety

-                x_samples_ddim = modules.safety.censor_batch(x_samples_ddim)

+            if p.scripts is not None:

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

 

             for i, x_sample in enumerate(x_samples_ddim):

                 x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)

@@ -566,7 +624,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     image.info["parameters"] = text

                 output_images.append(image)

 

-            del x_samples_ddim 

+            del x_samples_ddim

 

             devices.torch_gc()

 

@@ -592,28 +650,29 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 

     devices.torch_gc()

 

-    res = Processed(p, output_images, p.all_seeds[0], infotext() + "".join(["\n\n" + x for x in comments]), subseed=p.all_subseeds[0], all_prompts=p.all_prompts, all_seeds=p.all_seeds, all_subseeds=p.all_subseeds, index_of_first_image=index_of_first_image, infotexts=infotexts)

+    res = Processed(p, output_images, p.all_seeds[0], infotext(), comments="".join(["\n\n" + x for x in comments]), subseed=p.all_subseeds[0], index_of_first_image=index_of_first_image, infotexts=infotexts)

 

     if p.scripts is not None:

         p.scripts.postprocess(p, res)

 

-    p.sd_model = None

-    p.sampler = None

-

     return res

 

 

 class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):

     sampler = None

 

-    def __init__(self, enable_hr: bool=False, denoising_strength: float=0.75, firstphase_width: int=0, firstphase_height: int=0, **kwargs):

+    def __init__(self, enable_hr: bool = False, denoising_strength: float = 0.75, firstphase_width: int = 0, firstphase_height: int = 0, hr_scale: float = 2.0, hr_upscaler: str = None, **kwargs):

         super().__init__(**kwargs)

         self.enable_hr = enable_hr

         self.denoising_strength = denoising_strength

-        self.firstphase_width = firstphase_width

-        self.firstphase_height = firstphase_height

-        self.truncate_x = 0

-        self.truncate_y = 0

+        self.hr_scale = hr_scale

+        self.hr_upscaler = hr_upscaler

+

+        if firstphase_width != 0 or firstphase_height != 0:

+            print("firstphase_width/firstphase_height no longer supported; use hr_scale", file=sys.stderr)

+            self.hr_scale = self.width / firstphase_width

+            self.width = firstphase_width

+            self.height = firstphase_height

 

     def init(self, all_prompts, all_seeds, all_subseeds):

         if self.enable_hr:

@@ -622,48 +681,50 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             else:

                 state.job_count = state.job_count * 2

 

-            self.extra_generation_params["First pass size"] = f"{self.firstphase_width}x{self.firstphase_height}"

+            self.extra_generation_params["Hires upscale"] = self.hr_scale

+            if self.hr_upscaler is not None:

+                self.extra_generation_params["Hires upscaler"] = self.hr_upscaler

 

-            if self.firstphase_width == 0 or self.firstphase_height == 0:

-                desired_pixel_count = 512 * 512

-                actual_pixel_count = self.width * self.height

-                scale = math.sqrt(desired_pixel_count / actual_pixel_count)

-                self.firstphase_width = math.ceil(scale * self.width / 64) * 64

-                self.firstphase_height = math.ceil(scale * self.height / 64) * 64

-                firstphase_width_truncated = int(scale * self.width)

-                firstphase_height_truncated = int(scale * self.height)

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

+        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)

 

-            else:

+        latent_scale_mode = shared.latent_upscale_modes.get(self.hr_upscaler, None) if self.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "nearest")

+        if self.enable_hr and latent_scale_mode is None:

+            assert len([x for x in shared.sd_upscalers if x.name == self.hr_upscaler]) > 0, f"could not find upscaler named {self.hr_upscaler}"

+

+        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(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))

 

-                width_ratio = self.width / self.firstphase_width

-                height_ratio = self.height / self.firstphase_height

+        if not self.enable_hr:

+            return samples

 

-                if width_ratio > height_ratio:

-                    firstphase_width_truncated = self.firstphase_width

-                    firstphase_height_truncated = self.firstphase_width * self.height / self.width

-                else:

-                    firstphase_width_truncated = self.firstphase_height * self.width / self.height

-                    firstphase_height_truncated = self.firstphase_height

+        target_width = int(self.width * self.hr_scale)

+        target_height = int(self.height * self.hr_scale)

 

-            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 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"""

 

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

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

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

+                return

 

-        if not self.enable_hr:

-            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(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))

-            return samples

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

+                image = sd_samplers.sample_to_image(image, index, approximation=0)

 

-        x = create_random_tensors([opt_C, self.firstphase_height // opt_f, self.firstphase_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(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x, self.firstphase_width, self.firstphase_height))

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

 

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

+        if latent_scale_mode is not None:

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

+                save_intermediate(samples, i)

 

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

+            samples = torch.nn.functional.interpolate(samples, size=(target_height // opt_f, target_width // opt_f), mode=latent_scale_mode["mode"], antialias=latent_scale_mode["antialias"])

 

+            # Avoid making the inpainting conditioning unless necessary as

+            # this does need some extra compute to decode / encode the image again.

+            if getattr(self, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) < 1.0:

+                image_conditioning = self.img2img_image_conditioning(decode_first_stage(self.sd_model, samples), samples)

+            else:

+                image_conditioning = self.txt2img_image_conditioning(samples)

         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)

@@ -673,7 +734,10 @@ 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)

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

+

+                save_intermediate(image, i)

+

+                image = images.resize_image(0, image, target_width, target_height, upscaler_name=self.hr_upscaler)

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

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

                 batch_images.append(image)

@@ -684,13 +748,13 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

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

 

-        shared.state.nextjob()

+            image_conditioning = self.img2img_image_conditioning(decoded_samples, samples)

 

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

+        shared.state.nextjob()

 

-        noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

+        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)

 

-        image_conditioning = self.txt2img_image_conditioning(x)

+        noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, p=self)

 

         # GC now before running the next img2img to prevent running out of memory

         x = None

@@ -704,7 +768,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):

     sampler = None

 

-    def __init__(self, init_images: list=None, resize_mode: int=0, denoising_strength: float=0.75, mask: Any=None, mask_blur: int=4, inpainting_fill: int=0, inpaint_full_res: bool=True, inpaint_full_res_padding: int=0, inpainting_mask_invert: int=0, **kwargs):

+    def __init__(self, init_images: list = None, resize_mode: int = 0, denoising_strength: float = 0.75, mask: Any = None, mask_blur: int = 4, inpainting_fill: int = 0, inpaint_full_res: bool = True, inpaint_full_res_padding: int = 0, inpainting_mask_invert: int = 0, initial_noise_multiplier: float = None, **kwargs):

         super().__init__(**kwargs)

 

         self.init_images = init_images

@@ -712,7 +776,6 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         self.denoising_strength: float = denoising_strength

         self.init_latent = None

         self.image_mask = mask

-        #self.image_unblurred_mask = None

         self.latent_mask = None

         self.mask_for_overlay = None

         self.mask_blur = mask_blur

@@ -720,66 +783,68 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         self.inpaint_full_res = inpaint_full_res

         self.inpaint_full_res_padding = inpaint_full_res_padding

         self.inpainting_mask_invert = inpainting_mask_invert

+        self.initial_noise_multiplier = opts.initial_noise_multiplier if initial_noise_multiplier is None else initial_noise_multiplier

         self.mask = None

         self.nmask = None

         self.image_conditioning = None

 

     def init(self, all_prompts, all_seeds, all_subseeds):

-        self.sampler = sd_samplers.create_sampler_with_index(sd_samplers.samplers_for_img2img, self.sampler_index, self.sd_model)

+        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)

         crop_region = None

 

-        if self.image_mask is not None:

-            self.image_mask = self.image_mask.convert('L')

+        image_mask = self.image_mask

 

-            if self.inpainting_mask_invert:

-                self.image_mask = ImageOps.invert(self.image_mask)

+        if image_mask is not None:

+            image_mask = image_mask.convert('L')

 

-            #self.image_unblurred_mask = self.image_mask

+            if self.inpainting_mask_invert:

+                image_mask = ImageOps.invert(image_mask)

 

             if self.mask_blur > 0:

-                self.image_mask = self.image_mask.filter(ImageFilter.GaussianBlur(self.mask_blur))

+                image_mask = image_mask.filter(ImageFilter.GaussianBlur(self.mask_blur))

 

             if self.inpaint_full_res:

-                self.mask_for_overlay = self.image_mask

-                mask = self.image_mask.convert('L')

+                self.mask_for_overlay = image_mask

+                mask = image_mask.convert('L')

                 crop_region = masking.get_crop_region(np.array(mask), self.inpaint_full_res_padding)

                 crop_region = masking.expand_crop_region(crop_region, self.width, self.height, mask.width, mask.height)

                 x1, y1, x2, y2 = crop_region

 

                 mask = mask.crop(crop_region)

-                self.image_mask = images.resize_image(2, mask, self.width, self.height)

+                image_mask = images.resize_image(2, mask, self.width, self.height)

                 self.paste_to = (x1, y1, x2-x1, y2-y1)

             else:

-                self.image_mask = images.resize_image(self.resize_mode, self.image_mask, self.width, self.height)

-                np_mask = np.array(self.image_mask)

+                image_mask = images.resize_image(self.resize_mode, image_mask, self.width, self.height)

+                np_mask = np.array(image_mask)

                 np_mask = np.clip((np_mask.astype(np.float32)) * 2, 0, 255).astype(np.uint8)

                 self.mask_for_overlay = Image.fromarray(np_mask)

 

             self.overlay_images = []

 

-        latent_mask = self.latent_mask if self.latent_mask is not None else self.image_mask

+        latent_mask = self.latent_mask if self.latent_mask is not None else image_mask

 

         add_color_corrections = opts.img2img_color_correction and self.color_corrections is None

         if add_color_corrections:

             self.color_corrections = []

         imgs = []

         for img in self.init_images:

-            image = img.convert("RGB")

+            image = images.flatten(img, opts.img2img_background_color)

 

-            if crop_region is None:

+            if crop_region is None and self.resize_mode != 3:

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

 

-            if self.image_mask is not None:

+            if image_mask is not None:

                 image_masked = Image.new('RGBa', (image.width, image.height))

                 image_masked.paste(image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(self.mask_for_overlay.convert('L')))

 

                 self.overlay_images.append(image_masked.convert('RGBA'))

 

+            # crop_region is not None if we are doing inpaint full res

             if crop_region is not None:

                 image = image.crop(crop_region)

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

 

-            if self.image_mask is not None:

+            if image_mask is not None:

                 if self.inpainting_fill != 1:

                     image = masking.fill(image, latent_mask)

 

@@ -811,7 +876,10 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
 

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

 

-        if self.image_mask is not None:

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

+

+        if image_mask is not None:

             init_mask = latent_mask

             latmask = init_mask.convert('RGB').resize((self.init_latent.shape[3], self.init_latent.shape[2]))

             latmask = np.moveaxis(np.array(latmask, dtype=np.float32), 2, 0) / 255

@@ -828,12 +896,15 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             elif self.inpainting_fill == 3:

                 self.init_latent = self.init_latent * self.mask

 

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

+        self.image_conditioning = self.img2img_image_conditioning(image, self.init_latent, 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)

 

+        if self.initial_noise_multiplier != 1.0:

+            self.extra_generation_params["Noise multiplier"] = self.initial_noise_multiplier

+            x *= self.initial_noise_multiplier

+

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

 

         if self.mask is not None:

@@ -842,4 +913,4 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         del x

         devices.torch_gc()

 

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