Merge branch 'master' into tensorboard

1 files changed, 436 insertions, 145 deletions

diff --git a/modules/processing.py b/modules/processing.py
index bcb0c32c..ae04cab7 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

+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,34 +40,68 @@ 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

+

+

+def apply_overlay(image, paste_loc, index, overlays):

+    if overlays is None or index >= len(overlays):

+        return image

+

+    overlay = overlays[index]

+

+    if paste_loc is not None:

+        x, y, w, h = paste_loc

+        base_image = Image.new('RGBA', (overlay.width, overlay.height))

+        image = images.resize_image(1, image, w, h)

+        base_image.paste(image, (x, y))

+        image = base_image

+

+    image = image.convert('RGBA')

+    image.alpha_composite(overlay)

+    image = image.convert('RGB')

+

     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

+def txt2img_image_conditioning(sd_model, x, width, height):

+    if sd_model.model.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 x.new_zeros(x.shape[0], 5, 1, 1, dtype=x.dtype, device=x.device)

+

+    # 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 = sd_model.get_first_stage_encoding(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)

+

+    return image_conditioning

+

 

 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 = "uniform", s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0):

+    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, script_args: list = 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

@@ -73,7 +114,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

@@ -90,13 +131,16 @@ class StableDiffusionProcessing():
         self.do_not_reload_embeddings = do_not_reload_embeddings

         self.paste_to = None

         self.color_corrections = None

-        self.denoising_strength: float = 0

+        self.denoising_strength: float = denoising_strength

         self.sampler_noise_scheduler_override = None

-        self.ddim_discretize = opts.ddim_discretize

+        self.ddim_discretize = ddim_discretize or opts.ddim_discretize

         self.s_churn = s_churn or opts.s_churn

         self.s_tmin = s_tmin or opts.s_tmin

         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

@@ -104,16 +148,100 @@ class StableDiffusionProcessing():
             self.seed_resize_from_h = 0

             self.seed_resize_from_w = 0

 

+        self.scripts = None

+        self.script_args = script_args

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

+        self.is_using_inpainting_conditioning = self.sd_model.model.conditioning_key in {'hybrid', 'concat'}

+

+        return txt2img_image_conditioning(self.sd_model, x, width or self.width, height or self.height)

+

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

+            if torch.is_tensor(image_mask):

+                conditioning_mask = image_mask

+            else:

+                conditioning_mask = np.array(image_mask.convert("L"))

+                conditioning_mask = conditioning_mask.astype(np.float32) / 255.0

+                conditioning_mask = torch.from_numpy(conditioning_mask[None, None])

+

+                # 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 = 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).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))

+

+        # Create the concatenated conditioning tensor to be fed to `c_concat`

+        conditioning_mask = torch.nn.functional.interpolate(conditioning_mask, size=latent_image.shape[-2:])

+        conditioning_mask = conditioning_mask.expand(conditioning_image.shape[0], -1, -1, -1)

+        image_conditioning = torch.cat([conditioning_mask, conditioning_image], dim=1)

+        image_conditioning = image_conditioning.to(shared.device).type(self.sd_model.dtype)

+

+        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

@@ -121,10 +249,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

@@ -151,17 +279,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,

@@ -169,8 +300,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,

@@ -186,11 +316,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)

 

 

@@ -210,13 +341,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

@@ -256,8 +388,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)))

@@ -290,27 +422,30 @@ def fix_seed(p):
     p.subseed = get_fixed_seed(p.subseed)

 

 

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

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

     index = position_in_batch + iteration * p.batch_size

 

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

 

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

         "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(':', '')),

-        "Hypernet": (None if shared.loaded_hypernetwork is None else shared.loaded_hypernetwork.filename.split('\\')[-1].split('.')[0]),

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

         "Variation seed": (None if p.subseed_strength == 0 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),

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

@@ -318,14 +453,44 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
 

     generation_params.update(p.extra_generation_params)

 

-    generation_params_text = ", ".join([k if k == v else f'{k}: {v}' for k, v in generation_params.items() if v is not None])

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

 

 

 def process_images(p: StableDiffusionProcessing) -> Processed:

+    stored_opts = {k: opts.data[k] for k in p.override_settings.keys()}

+

+    try:

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

+            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

+                p.sd_model = shared.sd_model

+

+            if k == 'sd_vae':

+                sd_vae.reload_vae_weights()  # make onchange call for changing VAE

+

+        res = process_images_inner(p)

+

+    finally:

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

+

+

+def process_images_inner(p: StableDiffusionProcessing) -> Processed:

     """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch"""

 

     if type(p.prompt) == list:

@@ -333,10 +498,6 @@ def process_images(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)

@@ -347,58 +508,94 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 

     comments = {}

 

-    shared.prompt_styles.apply_styles(p)

-

     if type(p.prompt) == list:

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

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

-        all_seeds = seed

+        p.all_seeds = seed

     else:

-        all_seeds = [int(seed) + (x if p.subseed_strength == 0 else 0) for x in range(len(all_prompts))]

+        p.all_seeds = [int(seed) + (x if p.subseed_strength == 0 else 0) for x in range(len(p.all_prompts))]

 

     if type(subseed) == list:

-        all_subseeds = subseed

+        p.all_subseeds = subseed

     else:

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

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

 

     def infotext(iteration=0, position_in_batch=0):

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

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

 

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

         model_hijack.embedding_db.load_textual_inversion_embeddings()

 

+    if p.scripts is not None:

+        p.scripts.process(p)

+

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

+

     infotexts = []

     output_images = []

 

+    cached_uc = [None, None]

+    cached_c = [None, None]

+

+    def get_conds_with_caching(function, required_prompts, steps, cache):

+        """

+        Returns the result of calling function(shared.sd_model, required_prompts, steps)

+        using a cache to store the result if the same arguments have been used before.

+

+        cache is an array containing two elements. The first element is a tuple

+        representing the previously used arguments, or None if no arguments

+        have been used before. The second element is where the previously

+        computed result is stored.

+        """

+

+        if cache[0] is not None and (required_prompts, steps) == cache[0]:

+            return cache[1]

+

+        with devices.autocast():

+            cache[1] = function(shared.sd_model, required_prompts, steps)

+

+        cache[0] = (required_prompts, steps)

+        return cache[1]

+

     with torch.no_grad(), p.sd_model.ema_scope():

         with devices.autocast():

-            p.init(all_prompts, all_seeds, all_subseeds)

+            p.init(p.all_prompts, p.all_seeds, p.all_subseeds)

 

         if state.job_count == -1:

             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 = all_prompts[n * p.batch_size:(n + 1) * p.batch_size]

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

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

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

+            if len(prompts) == 0:

                 break

 

-            #uc = p.sd_model.get_learned_conditioning(len(prompts) * [p.negative_prompt])

-            #c = p.sd_model.get_learned_conditioning(prompts)

-            with devices.autocast():

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

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

+            if p.scripts is not None:

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

+

+            uc = get_conds_with_caching(prompt_parser.get_learned_conditioning, negative_prompts, p.steps, cached_uc)

+            c = get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, prompts, p.steps, cached_c)

 

             if len(model_hijack.comments) > 0:

                 for comment in model_hijack.comments:

@@ -408,10 +605,10 @@ def process_images(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

@@ -421,9 +618,8 @@ def process_images(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)

@@ -442,22 +638,11 @@ def process_images(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:

-                        images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-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, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-color-correction")

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

 

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

-                    overlay = p.overlay_images[i]

-

-                    if p.paste_to is not None:

-                        x, y, w, h = p.paste_to

-                        base_image = Image.new('RGBA', (overlay.width, overlay.height))

-                        image = images.resize_image(1, image, w, h)

-                        base_image.paste(image, (x, y))

-                        image = base_image

-

-                    image = image.convert('RGBA')

-                    image.alpha_composite(overlay)

-                    image = image.convert('RGB')

+                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, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p)

@@ -468,7 +653,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                     image.info["parameters"] = text

                 output_images.append(image)

 

-            del x_samples_ddim 

+            del x_samples_ddim

 

             devices.torch_gc()

 

@@ -490,73 +675,157 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                 index_of_first_image = 1

 

             if opts.grid_save:

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

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

 

     devices.torch_gc()

-    return Processed(p, output_images, all_seeds[0], infotext() + "".join(["\n\n" + x for x in comments]), subseed=all_subseeds[0], all_prompts=all_prompts, all_seeds=all_seeds, all_subseeds=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)

+

+    return res

+

+

+def old_hires_fix_first_pass_dimensions(width, height):

+    """old algorithm for auto-calculating first pass size"""

+

+    desired_pixel_count = 512 * 512

+    actual_pixel_count = width * height

+    scale = math.sqrt(desired_pixel_count / actual_pixel_count)

+    width = math.ceil(scale * width / 64) * 64

+    height = math.ceil(scale * height / 64) * 64

+

+    return width, height

 

 

 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, hr_second_pass_steps: int = 0, hr_resize_x: int = 0, hr_resize_y: int = 0, **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.hr_scale = hr_scale

+        self.hr_upscaler = hr_upscaler

+        self.hr_second_pass_steps = hr_second_pass_steps

+        self.hr_resize_x = hr_resize_x

+        self.hr_resize_y = hr_resize_y

+        self.hr_upscale_to_x = hr_resize_x

+        self.hr_upscale_to_y = hr_resize_y

+

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

+            self.hr_upscale_to_x = self.width

+            self.hr_upscale_to_y = self.height

+            self.width = firstphase_width

+            self.height = firstphase_height

+

         self.truncate_x = 0

         self.truncate_y = 0

+        self.applied_old_hires_behavior_to = None

 

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

         if self.enable_hr:

-            if state.job_count == -1:

-                state.job_count = self.n_iter * 2

+            if opts.use_old_hires_fix_width_height and self.applied_old_hires_behavior_to != (self.width, self.height):

+                self.hr_resize_x = self.width

+                self.hr_resize_y = self.height

+                self.hr_upscale_to_x = self.width

+                self.hr_upscale_to_y = self.height

+

+                self.width, self.height = old_hires_fix_first_pass_dimensions(self.width, self.height)

+                self.applied_old_hires_behavior_to = (self.width, self.height)

+

+            if self.hr_resize_x == 0 and self.hr_resize_y == 0:

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

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

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

             else:

+                self.extra_generation_params["Hires resize"] = f"{self.hr_resize_x}x{self.hr_resize_y}"

+

+                if self.hr_resize_y == 0:

+                    self.hr_upscale_to_x = self.hr_resize_x

+                    self.hr_upscale_to_y = self.hr_resize_x * self.height // self.width

+                elif self.hr_resize_x == 0:

+                    self.hr_upscale_to_x = self.hr_resize_y * self.width // self.height

+                    self.hr_upscale_to_y = self.hr_resize_y

+                else:

+                    target_w = self.hr_resize_x

+                    target_h = self.hr_resize_y

+                    src_ratio = self.width / self.height

+                    dst_ratio = self.hr_resize_x / self.hr_resize_y

+

+                    if src_ratio < dst_ratio:

+                        self.hr_upscale_to_x = self.hr_resize_x

+                        self.hr_upscale_to_y = self.hr_resize_x * self.height // self.width

+                    else:

+                        self.hr_upscale_to_x = self.hr_resize_y * self.width // self.height

+                        self.hr_upscale_to_y = self.hr_resize_y

+

+                    self.truncate_x = (self.hr_upscale_to_x - target_w) // opt_f

+                    self.truncate_y = (self.hr_upscale_to_y - target_h) // opt_f

+

+            # special case: the user has chosen to do nothing

+            if self.hr_upscale_to_x == self.width and self.hr_upscale_to_y == self.height:

+                self.enable_hr = False

+                self.denoising_strength = None

+                self.extra_generation_params.pop("Hires upscale", None)

+                self.extra_generation_params.pop("Hires resize", None)

+                return

+

+            if not state.processing_has_refined_job_count:

+                if state.job_count == -1:

+                    state.job_count = self.n_iter

+

+                shared.total_tqdm.updateTotal((self.steps + (self.hr_second_pass_steps or self.steps)) * state.job_count)

                 state.job_count = state.job_count * 2

+                state.processing_has_refined_job_count = True

 

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

+            if self.hr_second_pass_steps:

+                self.extra_generation_params["Hires steps"] = self.hr_second_pass_steps

 

-            else:

-

-                width_ratio = self.width / self.firstphase_width

-                height_ratio = self.height / self.firstphase_height

-

-                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

+            if self.hr_upscaler is not None:

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

 

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

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

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

 

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

 

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