Upload 7 files

sec/processing.py

@@ -0,0 +1,1860 @@
+from __future__ import annotations
+import json
+import logging
+import math
+import os
+import sys
+import hashlib
+from dataclasses import dataclass, field
+
+import torch
+import numpy as np
+from PIL import Image, ImageOps
+import random
+import cv2
+from skimage import exposure
+from typing import Any
+
+import modules.sd_hijack
+from modules import devices, prompt_parser, masking, sd_samplers, lowvram, infotext_utils, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet, errors, rng, profiling
+from modules.rng import slerp # noqa: F401
+from modules.sd_hijack import model_hijack
+from modules.sd_samplers_common import images_tensor_to_samples, decode_first_stage, approximation_indexes
+from modules.shared import opts, cmd_opts, state
+import modules.shared as shared
+import modules.paths as paths
+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
+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
+opt_f = 8
+
+
+def setup_color_correction(image):
+    logging.info("Calibrating color correction.")
+    correction_target = cv2.cvtColor(np.asarray(image.copy()), cv2.COLOR_RGB2LAB)
+    return correction_target
+
+
+def apply_color_correction(correction, original_image):
+    logging.info("Applying color correction.")
+    image = Image.fromarray(cv2.cvtColor(exposure.match_histograms(
+        cv2.cvtColor(
+            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.convert('RGB')
+
+
+def uncrop(image, dest_size, paste_loc):
+    x, y, w, h = paste_loc
+    base_image = Image.new('RGBA', dest_size)
+    image = images.resize_image(1, image, w, h)
+    base_image.paste(image, (x, y))
+    image = base_image
+
+    return image
+
+
+def apply_overlay(image, paste_loc, overlay):
+    if overlay is None:
+        return image, image.copy()
+
+    if paste_loc is not None:
+        image = uncrop(image, (overlay.width, overlay.height), paste_loc)
+
+    original_denoised_image = image.copy()
+
+    image = image.convert('RGBA')
+    image.alpha_composite(overlay)
+    image = image.convert('RGB')
+
+    return image, original_denoised_image
+
+def create_binary_mask(image, round=True):
+    if image.mode == 'RGBA' and image.getextrema()[-1] != (255, 255):
+        if round:
+            image = image.split()[-1].convert("L").point(lambda x: 255 if x > 128 else 0)
+        else:
+            image = image.split()[-1].convert("L")
+    else:
+        image = image.convert('L')
+    return image
+
+def txt2img_image_conditioning(sd_model, x, width, height):
+    if sd_model.model.conditioning_key in {'hybrid', 'concat'}: # Inpainting models
+
+        # The "masked-image" in this case will just be all 0.5 since the entire image is masked.
+        image_conditioning = torch.ones(x.shape[0], 3, height, width, device=x.device) * 0.5
+        image_conditioning = images_tensor_to_samples(image_conditioning, approximation_indexes.get(opts.sd_vae_encode_method))
+
+        # 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
+
+    elif sd_model.model.conditioning_key == "crossattn-adm": # UnCLIP models
+
+        return x.new_zeros(x.shape[0], 2*sd_model.noise_augmentor.time_embed.dim, dtype=x.dtype, device=x.device)
+
+    else:
+        if sd_model.is_sdxl_inpaint:
+            # The "masked-image" in this case will just be all 0.5 since the entire image is masked.
+            image_conditioning = torch.ones(x.shape[0], 3, height, width, device=x.device) * 0.5
+            image_conditioning = images_tensor_to_samples(image_conditioning,
+                                                            approximation_indexes.get(opts.sd_vae_encode_method))
+
+            # 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
+
+        # Dummy zero conditioning if we're not using inpainting or unclip models.
+        # 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)
+
+
+@dataclass(repr=False)
+class StableDiffusionProcessing:
+    sd_model: object = None
+    outpath_samples: str = None
+    outpath_grids: str = None
+    prompt: str = ""
+    prompt_for_display: str = None
+    negative_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
+    scheduler: 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 = None
+    tiling: bool = None
+    do_not_save_samples: bool = False
+    do_not_save_grid: bool = False
+    extra_generation_params: dict[str, Any] = None
+    overlay_images: list = None
+    eta: float = None
+    do_not_reload_embeddings: bool = False
+    denoising_strength: float = None
+    ddim_discretize: str = None
+    s_min_uncond: float = None
+    s_churn: float = None
+    s_tmax: float = None
+    s_tmin: float = None
+    s_noise: float = None
+    override_settings: dict[str, Any] = None
+    override_settings_restore_afterwards: bool = True
+    sampler_index: int = None
+    refiner_checkpoint: str = None
+    refiner_switch_at: float = None
+    token_merging_ratio = 0
+    token_merging_ratio_hr = 0
+    disable_extra_networks: bool = False
+    firstpass_image: Image = None
+
+    scripts_value: scripts.ScriptRunner = field(default=None, init=False)
+    script_args_value: list = field(default=None, init=False)
+    scripts_setup_complete: bool = field(default=False, init=False)
+
+    cached_uc = [None, None]
+    cached_c = [None, None]
+
+    comments: dict = None
+    sampler: sd_samplers_common.Sampler | None = field(default=None, init=False)
+    is_using_inpainting_conditioning: bool = field(default=False, init=False)
+    paste_to: tuple | None = field(default=None, init=False)
+
+    is_hr_pass: bool = field(default=False, init=False)
+
+    c: tuple = field(default=None, init=False)
+    uc: tuple = field(default=None, init=False)
+
+    rng: rng.ImageRNG | None = field(default=None, init=False)
+    step_multiplier: int = field(default=1, init=False)
+    color_corrections: list = field(default=None, init=False)
+
+    all_prompts: list = field(default=None, init=False)
+    all_negative_prompts: list = field(default=None, init=False)
+    all_seeds: list = field(default=None, init=False)
+    all_subseeds: list = field(default=None, init=False)
+    iteration: int = field(default=0, init=False)
+    main_prompt: str = field(default=None, init=False)
+    main_negative_prompt: str = field(default=None, init=False)
+
+    prompts: list = field(default=None, init=False)
+    negative_prompts: list = field(default=None, init=False)
+    seeds: list = field(default=None, init=False)
+    subseeds: list = field(default=None, init=False)
+    extra_network_data: dict = field(default=None, init=False)
+
+    user: str = field(default=None, init=False)
+
+    sd_model_name: str = field(default=None, init=False)
+    sd_model_hash: str = field(default=None, init=False)
+    sd_vae_name: str = field(default=None, init=False)
+    sd_vae_hash: str = field(default=None, init=False)
+
+    is_api: bool = field(default=False, init=False)
+
+    def __post_init__(self):
+        if self.sampler_index is not None:
+            print("sampler_index argument for StableDiffusionProcessing does not do anything; use sampler_name", file=sys.stderr)
+
+        self.comments = {}
+
+        if self.styles is None:
+            self.styles = []
+
+        self.sampler_noise_scheduler_override = None
+
+        self.extra_generation_params = self.extra_generation_params or {}
+        self.override_settings = self.override_settings or {}
+        self.script_args = self.script_args or {}
+
+        self.refiner_checkpoint_info = None
+
+        if not self.seed_enable_extras:
+            self.subseed = -1
+            self.subseed_strength = 0
+            self.seed_resize_from_h = 0
+            self.seed_resize_from_w = 0
+
+        self.cached_uc = StableDiffusionProcessing.cached_uc
+        self.cached_c = StableDiffusionProcessing.cached_c
+
+    def fill_fields_from_opts(self):
+        self.s_min_uncond = self.s_min_uncond if self.s_min_uncond is not None else opts.s_min_uncond
+        self.s_churn = self.s_churn if self.s_churn is not None else opts.s_churn
+        self.s_tmin = self.s_tmin if self.s_tmin is not None else opts.s_tmin
+        self.s_tmax = (self.s_tmax if self.s_tmax is not None else opts.s_tmax) or float('inf')
+        self.s_noise = self.s_noise if self.s_noise is not None else opts.s_noise
+
+    @property
+    def sd_model(self):
+        return shared.sd_model
+
+    @sd_model.setter
+    def sd_model(self, value):
+        pass
+
+    @property
+    def scripts(self):
+        return self.scripts_value
+
+    @scripts.setter
+    def scripts(self, value):
+        self.scripts_value = value
+
+        if self.scripts_value and self.script_args_value and not self.scripts_setup_complete:
+            self.setup_scripts()
+
+    @property
+    def script_args(self):
+        return self.script_args_value
+
+    @script_args.setter
+    def script_args(self, value):
+        self.script_args_value = value
+
+        if self.scripts_value and self.script_args_value and not self.scripts_setup_complete:
+            self.setup_scripts()
+
+    def setup_scripts(self):
+        self.scripts_setup_complete = True
+
+        self.scripts.setup_scrips(self, is_ui=not self.is_api)
+
+    def comment(self, text):
+        self.comments[text] = 1
+
+    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 = images_tensor_to_samples(source_image*0.5+0.5, approximation_indexes.get(opts.sd_vae_encode_method))
+        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 edit_image_conditioning(self, source_image):
+        conditioning_image = shared.sd_model.encode_first_stage(source_image).mode()
+
+        return conditioning_image
+
+    def unclip_image_conditioning(self, source_image):
+        c_adm = self.sd_model.embedder(source_image)
+        if self.sd_model.noise_augmentor is not None:
+            noise_level = 0 # TODO: Allow other noise levels?
+            c_adm, noise_level_emb = self.sd_model.noise_augmentor(c_adm, noise_level=repeat(torch.tensor([noise_level]).to(c_adm.device), '1 -> b', b=c_adm.shape[0]))
+            c_adm = torch.cat((c_adm, noise_level_emb), 1)
+        return c_adm
+
+    def inpainting_image_conditioning(self, source_image, latent_image, image_mask=None, round_image_mask=True):
+        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])
+
+                if round_image_mask:
+                    # Caller is requesting 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(device=source_image.device, dtype=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, round_image_mask=True):
+        source_image = devices.cond_cast_float(source_image)
+
+        # 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.sd_model.cond_stage_key == "edit":
+            return self.edit_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, round_image_mask=round_image_mask)
+
+        if self.sampler.conditioning_key == "crossattn-adm":
+            return self.unclip_image_conditioning(source_image)
+
+        if self.sampler.model_wrap.inner_model.is_sdxl_inpaint:
+            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, prompts):
+        raise NotImplementedError()
+
+    def close(self):
+        self.sampler = None
+        self.c = None
+        self.uc = None
+        if not opts.persistent_cond_cache:
+            StableDiffusionProcessing.cached_c = [None, None]
+            StableDiffusionProcessing.cached_uc = [None, None]
+
+    def get_token_merging_ratio(self, for_hr=False):
+        if for_hr:
+            return self.token_merging_ratio_hr or opts.token_merging_ratio_hr or self.token_merging_ratio or opts.token_merging_ratio
+
+        return self.token_merging_ratio or opts.token_merging_ratio
+
+    def setup_prompts(self):
+        if isinstance(self.prompt,list):
+            self.all_prompts = self.prompt
+        elif isinstance(self.negative_prompt, list):
+            self.all_prompts = [self.prompt] * len(self.negative_prompt)
+        else:
+            self.all_prompts = self.batch_size * self.n_iter * [self.prompt]
+
+        if isinstance(self.negative_prompt, list):
+            self.all_negative_prompts = self.negative_prompt
+        else:
+            self.all_negative_prompts = [self.negative_prompt] * len(self.all_prompts)
+
+        if len(self.all_prompts) != len(self.all_negative_prompts):
+            raise RuntimeError(f"Received a different number of prompts ({len(self.all_prompts)}) and negative prompts ({len(self.all_negative_prompts)})")
+
+        self.all_prompts = [shared.prompt_styles.apply_styles_to_prompt(x, self.styles) for x in self.all_prompts]
+        self.all_negative_prompts = [shared.prompt_styles.apply_negative_styles_to_prompt(x, self.styles) for x in self.all_negative_prompts]
+
+        self.main_prompt = self.all_prompts[0]
+        self.main_negative_prompt = self.all_negative_prompts[0]
+
+    def cached_params(self, required_prompts, steps, extra_network_data, hires_steps=None, use_old_scheduling=False):
+        """Returns parameters that invalidate the cond cache if changed"""
+
+        return (
+            required_prompts,
+            steps,
+            hires_steps,
+            use_old_scheduling,
+            opts.CLIP_stop_at_last_layers,
+            shared.sd_model.sd_checkpoint_info,
+            extra_network_data,
+            opts.sdxl_crop_left,
+            opts.sdxl_crop_top,
+            self.width,
+            self.height,
+            opts.fp8_storage,
+            opts.cache_fp16_weight,
+            opts.emphasis,
+        )
+
+    def get_conds_with_caching(self, function, required_prompts, steps, caches, extra_network_data, hires_steps=None):
+        """
+        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.
+
+        caches is a list with items described above.
+        """
+
+        if shared.opts.use_old_scheduling:
+            old_schedules = prompt_parser.get_learned_conditioning_prompt_schedules(required_prompts, steps, hires_steps, False)
+            new_schedules = prompt_parser.get_learned_conditioning_prompt_schedules(required_prompts, steps, hires_steps, True)
+            if old_schedules != new_schedules:
+                self.extra_generation_params["Old prompt editing timelines"] = True
+
+        cached_params = self.cached_params(required_prompts, steps, extra_network_data, hires_steps, shared.opts.use_old_scheduling)
+
+        for cache in caches:
+            if cache[0] is not None and cached_params == cache[0]:
+                return cache[1]
+
+        cache = caches[0]
+
+        with devices.autocast():
+            cache[1] = function(shared.sd_model, required_prompts, steps, hires_steps, shared.opts.use_old_scheduling)
+
+        cache[0] = cached_params
+        return cache[1]
+
+    def setup_conds(self):
+        prompts = prompt_parser.SdConditioning(self.prompts, width=self.width, height=self.height)
+        negative_prompts = prompt_parser.SdConditioning(self.negative_prompts, width=self.width, height=self.height, is_negative_prompt=True)
+
+        sampler_config = sd_samplers.find_sampler_config(self.sampler_name)
+        total_steps = sampler_config.total_steps(self.steps) if sampler_config else self.steps
+        self.step_multiplier = total_steps // self.steps
+        self.firstpass_steps = total_steps
+
+        self.uc = self.get_conds_with_caching(prompt_parser.get_learned_conditioning, negative_prompts, total_steps, [self.cached_uc], self.extra_network_data)
+        self.c = self.get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, prompts, total_steps, [self.cached_c], self.extra_network_data)
+
+    def get_conds(self):
+        return self.c, self.uc
+
+    def parse_extra_network_prompts(self):
+        self.prompts, self.extra_network_data = extra_networks.parse_prompts(self.prompts)
+
+    def save_samples(self) -> bool:
+        """Returns whether generated images need to be written to disk"""
+        return opts.samples_save and not self.do_not_save_samples and (opts.save_incomplete_images or not state.interrupted and not state.skipped)
+
+
+class Processed:
+    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
+        self.seed = seed
+        self.subseed = subseed
+        self.subseed_strength = p.subseed_strength
+        self.info = info
+        self.comments = "".join(f"{comment}\n" for comment in p.comments)
+        self.width = p.width
+        self.height = p.height
+        self.sampler_name = p.sampler_name
+        self.cfg_scale = p.cfg_scale
+        self.image_cfg_scale = getattr(p, 'image_cfg_scale', None)
+        self.steps = p.steps
+        self.batch_size = p.batch_size
+        self.restore_faces = p.restore_faces
+        self.face_restoration_model = opts.face_restoration_model if p.restore_faces else None
+        self.sd_model_name = p.sd_model_name
+        self.sd_model_hash = p.sd_model_hash
+        self.sd_vae_name = p.sd_vae_name
+        self.sd_vae_hash = p.sd_vae_hash
+        self.seed_resize_from_w = p.seed_resize_from_w
+        self.seed_resize_from_h = p.seed_resize_from_h
+        self.denoising_strength = getattr(p, 'denoising_strength', None)
+        self.extra_generation_params = p.extra_generation_params
+        self.index_of_first_image = index_of_first_image
+        self.styles = p.styles
+        self.job_timestamp = state.job_timestamp
+        self.clip_skip = opts.CLIP_stop_at_last_layers
+        self.token_merging_ratio = p.token_merging_ratio
+        self.token_merging_ratio_hr = p.token_merging_ratio_hr
+
+        self.eta = p.eta
+        self.ddim_discretize = p.ddim_discretize
+        self.s_churn = p.s_churn
+        self.s_tmin = p.s_tmin
+        self.s_tmax = p.s_tmax
+        self.s_noise = p.s_noise
+        self.s_min_uncond = p.s_min_uncond
+        self.sampler_noise_scheduler_override = p.sampler_noise_scheduler_override
+        self.prompt = self.prompt if not isinstance(self.prompt, list) else self.prompt[0]
+        self.negative_prompt = self.negative_prompt if not isinstance(self.negative_prompt, list) else self.negative_prompt[0]
+        self.seed = int(self.seed if not isinstance(self.seed, list) else self.seed[0]) if self.seed is not None else -1
+        self.subseed = int(self.subseed if not isinstance(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 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] * len(images_list)
+        self.version = program_version()
+
+    def js(self):
+        obj = {
+            "prompt": self.all_prompts[0],
+            "all_prompts": self.all_prompts,
+            "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,
+            "all_subseeds": self.all_subseeds,
+            "subseed_strength": self.subseed_strength,
+            "width": self.width,
+            "height": self.height,
+            "sampler_name": self.sampler_name,
+            "cfg_scale": self.cfg_scale,
+            "steps": self.steps,
+            "batch_size": self.batch_size,
+            "restore_faces": self.restore_faces,
+            "face_restoration_model": self.face_restoration_model,
+            "sd_model_name": self.sd_model_name,
+            "sd_model_hash": self.sd_model_hash,
+            "sd_vae_name": self.sd_vae_name,
+            "sd_vae_hash": self.sd_vae_hash,
+            "seed_resize_from_w": self.seed_resize_from_w,
+            "seed_resize_from_h": self.seed_resize_from_h,
+            "denoising_strength": self.denoising_strength,
+            "extra_generation_params": self.extra_generation_params,
+            "index_of_first_image": self.index_of_first_image,
+            "infotexts": self.infotexts,
+            "styles": self.styles,
+            "job_timestamp": self.job_timestamp,
+            "clip_skip": self.clip_skip,
+            "is_using_inpainting_conditioning": self.is_using_inpainting_conditioning,
+            "version": self.version,
+        }
+
+        return json.dumps(obj, default=lambda o: None)
+
+    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)
+
+    def get_token_merging_ratio(self, for_hr=False):
+        return self.token_merging_ratio_hr if for_hr else self.token_merging_ratio
+
+
+def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, seed_resize_from_h=0, seed_resize_from_w=0, p=None):
+    g = rng.ImageRNG(shape, seeds, subseeds=subseeds, subseed_strength=subseed_strength, seed_resize_from_h=seed_resize_from_h, seed_resize_from_w=seed_resize_from_w)
+    return g.next()
+
+
+class DecodedSamples(list):
+    already_decoded = True
+
+
+def decode_latent_batch(model, batch, target_device=None, check_for_nans=False):
+    samples = DecodedSamples()
+
+    if check_for_nans:
+        devices.test_for_nans(batch, "unet")
+
+    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 shared.opts.auto_vae_precision_bfloat16:
+                    autofix_dtype = torch.bfloat16
+                    autofix_dtype_text = "bfloat16"
+                    autofix_dtype_setting = "Automatically convert VAE to bfloat16"
+                    autofix_dtype_comment = ""
+                elif shared.opts.auto_vae_precision:
+                    autofix_dtype = torch.float32
+                    autofix_dtype_text = "32-bit float"
+                    autofix_dtype_setting = "Automatically revert VAE to 32-bit floats"
+                    autofix_dtype_comment = "\nTo always start with 32-bit VAE, use --no-half-vae commandline flag."
+                else:
+                    raise e
+
+                if devices.dtype_vae == autofix_dtype:
+                    raise e
+
+                errors.print_error_explanation(
+                    "A tensor with all NaNs was produced in VAE.\n"
+                    f"Web UI will now convert VAE into {autofix_dtype_text} and retry.\n"
+                    f"To disable this behavior, disable the '{autofix_dtype_setting}' setting.{autofix_dtype_comment}"
+                )
+
+                devices.dtype_vae = autofix_dtype
+                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 get_fixed_seed(seed):
+    if seed == '' or seed is None:
+        seed = -1
+    elif isinstance(seed, str):
+        try:
+            seed = int(seed)
+        except Exception:
+            seed = -1
+
+    if seed == -1:
+        return int(random.randrange(4294967294))
+
+    return seed
+
+
+def fix_seed(p):
+    p.seed = get_fixed_seed(p.seed)
+    p.subseed = get_fixed_seed(p.subseed)
+
+
+def program_version():
+    import launch
+
+    res = launch.git_tag()
+    if res == "<none>":
+        res = None
+
+    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=None, all_negative_prompts=None):
+    """
+    this function is used to generate the infotext that is stored in the generated images, it's contains the parameters that are required to generate the imagee
+    Args:
+        p: StableDiffusionProcessing
+        all_prompts: list[str]
+        all_seeds: list[int]
+        all_subseeds: list[int]
+        comments: list[str]
+        iteration: int
+        position_in_batch: int
+        use_main_prompt: bool
+        index: int
+        all_negative_prompts: list[str]
+
+    Returns: str
+
+    Extra generation params
+    p.extra_generation_params dictionary allows for additional parameters to be added to the infotext
+    this can be use by the base webui or extensions.
+    To add a new entry, add a new key value pair, the dictionary key will be used as the key of the parameter in the infotext
+    the value generation_params can be defined as:
+        - str | None
+        - List[str|None]
+        - callable func(**kwargs) -> str | None
+
+    When defined as a string, it will be used as without extra processing; this is this most common use case.
+
+    Defining as a list allows for parameter that changes across images in the job, for example, the 'Seed' parameter.
+    The list should have the same length as the total number of images in the entire job.
+
+    Defining as a callable function allows parameter cannot be generated earlier or when extra logic is required.
+    For example 'Hires prompt', due to reasons the hr_prompt might be changed by process in the pipeline or extensions
+    and may vary across different images, defining as a static string or list would not work.
+
+    The function takes locals() as **kwargs, as such will have access to variables like 'p' and 'index'.
+    the base signature of the function should be:
+        func(**kwargs) -> str | None
+    optionally it can have additional arguments that will be used in the function:
+        func(p, index, **kwargs) -> str | None
+    note: for better future compatibility even though this function will have access to all variables in the locals(),
+        it is recommended to only use the arguments present in the function signature of create_infotext.
+    For actual implementation examples, see StableDiffusionProcessingTxt2Img.init > get_hr_prompt.
+    """
+
+    if use_main_prompt:
+        index = 0
+    elif 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)
+    token_merging_ratio = p.get_token_merging_ratio()
+    token_merging_ratio_hr = p.get_token_merging_ratio(for_hr=True)
+
+    prompt_text = p.main_prompt if use_main_prompt else all_prompts[index]
+    negative_prompt = p.main_negative_prompt if use_main_prompt else all_negative_prompts[index]
+
+    uses_ensd = opts.eta_noise_seed_delta != 0
+    if uses_ensd:
+        uses_ensd = sd_samplers_common.is_sampler_using_eta_noise_seed_delta(p)
+
+    generation_params = {
+        "Steps": p.steps,
+        "Sampler": p.sampler_name,
+        "Schedule type": p.scheduler,
+        "CFG scale": p.cfg_scale,
+        "Image CFG scale": getattr(p, 'image_cfg_scale', None),
+        "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": p.sd_model_hash if opts.add_model_hash_to_info else None,
+        "Model": p.sd_model_name if opts.add_model_name_to_info else None,
+        "FP8 weight": opts.fp8_storage if devices.fp8 else None,
+        "Cache FP16 weight for LoRA": opts.cache_fp16_weight if devices.fp8 else None,
+        "VAE hash": p.sd_vae_hash if opts.add_vae_hash_to_info else None,
+        "VAE": p.sd_vae_name if opts.add_vae_name_to_info else None,
+        "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": p.extra_generation_params.get("Denoising strength"),
+        "Conditional mask weight": getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None,
+        "Clip skip": None if clip_skip <= 1 else clip_skip,
+        "ENSD": opts.eta_noise_seed_delta if uses_ensd else None,
+        "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,
+        "Tiling": "True" if p.tiling else None,
+        "Progressive Growing": "True" if p.enable_progressive_growing else None,
+        "Min Scale": p.progressive_growing_min_scale if p.enable_progressive_growing else None,
+        "Max Scale": p.progressive_growing_max_scale if p.enable_progressive_growing else None,
+        "Progressive Growing Steps": p.progressive_growing_steps if p.enable_progressive_growing else None,
+        "Refinement": "True" if p.progressive_growing_refinement and p.enable_progressive_growing else None,
+        **p.extra_generation_params,
+        "Version": program_version() if opts.add_version_to_infotext else None,
+        "User": p.user if opts.add_user_name_to_info else None,
+    }
+
+    for key, value in generation_params.items():
+        try:
+            if isinstance(value, list):
+                generation_params[key] = value[index]
+            elif callable(value):
+                generation_params[key] = value(**locals())
+        except Exception:
+            errors.report(f'Error creating infotext for key "{key}"', exc_info=True)
+            generation_params[key] = None
+
+    generation_params_text = ", ".join([k if k == v else f'{k}: {infotext_utils.quote(v)}' for k, v in generation_params.items() if v is not None])
+
+    negative_prompt_text = f"\nNegative prompt: {negative_prompt}" if negative_prompt else ""
+
+    return f"{prompt_text}{negative_prompt_text}\n{generation_params_text}".strip()
+
+
+def process_images(p: StableDiffusionProcessing) -> Processed:
+    if p.scripts is not None:
+        p.scripts.before_process(p)
+
+    stored_opts = {k: opts.data[k] if k in opts.data else opts.get_default(k) for k in p.override_settings.keys() if k in opts.data}
+
+    try:
+        # if no checkpoint override or the override checkpoint can't be found, remove override entry and load opts checkpoint
+        # and if after running refiner, the refiner model is not unloaded - webui swaps back to main model here, if model over is present it will be reloaded afterwards
+        if sd_models.checkpoint_aliases.get(p.override_settings.get('sd_model_checkpoint')) is None:
+            p.override_settings.pop('sd_model_checkpoint', None)
+            sd_models.reload_model_weights()
+
+        for k, v in p.override_settings.items():
+            opts.set(k, v, is_api=True, run_callbacks=False)
+
+            if k == 'sd_model_checkpoint':
+                sd_models.reload_model_weights()
+
+            if k == 'sd_vae':
+                sd_vae.reload_vae_weights()
+
+        sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio())
+
+        # backwards compatibility, fix sampler and scheduler if invalid
+        sd_samplers.fix_p_invalid_sampler_and_scheduler(p)
+
+        with profiling.Profiler():
+            res = process_images_inner(p)
+
+    finally:
+        sd_models.apply_token_merging(p.sd_model, 0)
+
+        # 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_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 isinstance(p.prompt, list):
+        assert(len(p.prompt) > 0)
+    else:
+        assert p.prompt is not None
+
+    devices.torch_gc()
+
+    seed = get_fixed_seed(p.seed)
+    subseed = get_fixed_seed(p.subseed)
+
+    if p.restore_faces is None:
+        p.restore_faces = opts.face_restoration
+
+    if p.tiling is None:
+        p.tiling = opts.tiling
+
+    if p.refiner_checkpoint not in (None, "", "None", "none"):
+        p.refiner_checkpoint_info = sd_models.get_closet_checkpoint_match(p.refiner_checkpoint)
+        if p.refiner_checkpoint_info is None:
+            raise Exception(f'Could not find checkpoint with name {p.refiner_checkpoint}')
+
+    if hasattr(shared.sd_model, 'fix_dimensions'):
+        p.width, p.height = shared.sd_model.fix_dimensions(p.width, p.height)
+
+    p.sd_model_name = shared.sd_model.sd_checkpoint_info.name_for_extra
+    p.sd_model_hash = shared.sd_model.sd_model_hash
+    p.sd_vae_name = sd_vae.get_loaded_vae_name()
+    p.sd_vae_hash = sd_vae.get_loaded_vae_hash()
+
+    modules.sd_hijack.model_hijack.apply_circular(p.tiling)
+    modules.sd_hijack.model_hijack.clear_comments()
+
+    p.fill_fields_from_opts()
+    p.setup_prompts()
+
+    if isinstance(seed, list):
+        p.all_seeds = seed
+    else:
+        p.all_seeds = [int(seed) + (x if p.subseed_strength == 0 else 0) for x in range(len(p.all_prompts))]
+
+    if isinstance(subseed, list):
+        p.all_subseeds = subseed
+    else:
+        p.all_subseeds = [int(subseed) + x for x in range(len(p.all_prompts))]
+
+    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)
+
+    infotexts = []
+    output_images = []
+    with torch.no_grad(), p.sd_model.ema_scope():
+        with devices.autocast():
+            p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
+
+            # for OSX, loading the model during sampling changes the generated picture, so it is loaded here
+            if shared.opts.live_previews_enable and opts.show_progress_type == "Approx NN":
+                sd_vae_approx.model()
+
+            sd_unet.apply_unet()
+
+        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 or state.stopping_generation:
+                break
+
+            sd_models.reload_model_weights()  # model can be changed for example by refiner
+
+            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]
+            p.seeds = p.all_seeds[n * p.batch_size:(n + 1) * p.batch_size]
+            p.subseeds = p.all_subseeds[n * p.batch_size:(n + 1) * p.batch_size]
+
+            latent_channels = getattr(shared.sd_model, 'latent_channels', opt_C)
+            p.rng = rng.ImageRNG((latent_channels, p.height // opt_f, p.width // opt_f), p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w)
+
+            if p.scripts is not None:
+                p.scripts.before_process_batch(p, batch_number=n, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds)
+
+            if len(p.prompts) == 0:
+                break
+
+            p.parse_extra_network_prompts()
+
+            if not p.disable_extra_networks:
+                with devices.autocast():
+                    extra_networks.activate(p, p.extra_network_data)
+
+            if p.scripts is not None:
+                p.scripts.process_batch(p, batch_number=n, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds)
+
+            p.setup_conds()
+
+            p.extra_generation_params.update(model_hijack.extra_generation_params)
+
+            # params.txt should be saved after scripts.process_batch, since the
+            # infotext could be modified by that callback
+            # Example: a wildcard processed by process_batch sets an extra model
+            # strength, which is saved as "Model Strength: 1.0" in the infotext
+            if n == 0 and not cmd_opts.no_prompt_history:
+                with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:
+                    processed = Processed(p, [])
+                    file.write(processed.infotext(p, 0))
+
+            for comment in model_hijack.comments:
+                p.comment(comment)
+
+            if p.n_iter > 1:
+                shared.state.job = f"Batch {n+1} out of {p.n_iter}"
+
+            sd_models.apply_alpha_schedule_override(p.sd_model, p)
+
+            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)
+
+            if p.scripts is not None:
+                ps = scripts.PostSampleArgs(samples_ddim)
+                p.scripts.post_sample(p, ps)
+                samples_ddim = ps.samples
+
+            if getattr(samples_ddim, 'already_decoded', False):
+                x_samples_ddim = samples_ddim
+            else:
+                devices.test_for_nans(samples_ddim, "unet")
+
+                if opts.sd_vae_decode_method != 'Full':
+                    p.extra_generation_params['VAE Decoder'] = opts.sd_vae_decode_method
+                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)
+
+            del samples_ddim
+
+            if lowvram.is_enabled(shared.sd_model):
+                lowvram.send_everything_to_cpu()
+
+            devices.torch_gc()
+
+            state.nextjob()
+
+            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)
+
+            save_samples = p.save_samples()
+
+            for i, x_sample in enumerate(x_samples_ddim):
+                p.batch_index = i
+
+                x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
+                x_sample = x_sample.astype(np.uint8)
+
+                if p.restore_faces:
+                    if 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(i), p=p, suffix="-before-face-restoration")
+
+                    devices.torch_gc()
+
+                    x_sample = modules.face_restoration.restore_faces(x_sample)
+                    devices.torch_gc()
+
+                image = Image.fromarray(x_sample)
+
+                if p.scripts is not None:
+                    pp = scripts.PostprocessImageArgs(image)
+                    p.scripts.postprocess_image(p, pp)
+                    image = pp.image
+
+                mask_for_overlay = getattr(p, "mask_for_overlay", None)
+
+                if not shared.opts.overlay_inpaint:
+                    overlay_image = None
+                elif getattr(p, "overlay_images", None) is not None and i < len(p.overlay_images):
+                    overlay_image = p.overlay_images[i]
+                else:
+                    overlay_image = None
+
+                if p.scripts is not None:
+                    ppmo = scripts.PostProcessMaskOverlayArgs(i, mask_for_overlay, overlay_image)
+                    p.scripts.postprocess_maskoverlay(p, ppmo)
+                    mask_for_overlay, overlay_image = ppmo.mask_for_overlay, ppmo.overlay_image
+
+                if p.color_corrections is not None and i < len(p.color_corrections):
+                    if save_samples and opts.save_images_before_color_correction:
+                        image_without_cc, _ = apply_overlay(image, p.paste_to, overlay_image)
+                        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)
+
+                # If the intention is to show the output from the model
+                # that is being composited over the original image,
+                # we need to keep the original image around
+                # and use it in the composite step.
+                image, original_denoised_image = apply_overlay(image, p.paste_to, overlay_image)
+
+                if p.scripts is not None:
+                    pp = scripts.PostprocessImageArgs(image)
+                    p.scripts.postprocess_image_after_composite(p, pp)
+                    image = pp.image
+
+                if save_samples:
+                    images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], opts.samples_format, info=infotext(i), p=p)
+
+                text = infotext(i)
+                infotexts.append(text)
+                if opts.enable_pnginfo:
+                    image.info["parameters"] = text
+                output_images.append(image)
+
+                if mask_for_overlay is not None:
+                    if opts.return_mask or opts.save_mask:
+                        image_mask = mask_for_overlay.convert('RGB')
+                        if save_samples and opts.save_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.return_mask:
+                            output_images.append(image_mask)
+
+                    if opts.return_mask_composite or opts.save_mask_composite:
+                        image_mask_composite = Image.composite(original_denoised_image.convert('RGBA').convert('RGBa'), Image.new('RGBa', image.size), images.resize_image(2, mask_for_overlay, image.width, image.height).convert('L')).convert('RGBA')
+                        if save_samples and opts.save_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_composite:
+                            output_images.append(image_mask_composite)
+
+            del x_samples_ddim
+
+            devices.torch_gc()
+
+        if not infotexts:
+            infotexts.append(Processed(p, []).infotext(p, 0))
+
+        p.color_corrections = None
+
+        index_of_first_image = 0
+        unwanted_grid_because_of_img_count = len(output_images) < 2 and opts.grid_only_if_multiple
+        if (opts.return_grid or opts.grid_save) and not p.do_not_save_grid and not unwanted_grid_because_of_img_count:
+            grid = images.image_grid(output_images, p.batch_size)
+
+            if opts.return_grid:
+                text = infotext(use_main_prompt=True)
+                infotexts.insert(0, text)
+                if opts.enable_pnginfo:
+                    grid.info["parameters"] = text
+                output_images.insert(0, grid)
+                index_of_first_image = 1
+            if opts.grid_save:
+                images.save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], opts.grid_format, info=infotext(use_main_prompt=True), short_filename=not opts.grid_extended_filename, p=p, grid=True)
+
+    if not p.disable_extra_networks and p.extra_network_data:
+        extra_networks.deactivate(p, p.extra_network_data)
+
+    devices.torch_gc()
+
+    res = Processed(
+        p,
+        images_list=output_images,
+        seed=p.all_seeds[0],
+        info=infotexts[0],
+        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
+
+
+@dataclass(repr=False)
+class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
+    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
+    hr_checkpoint_name: str = None
+    hr_sampler_name: str = None
+    hr_scheduler: str = None
+    hr_prompt: str = ''
+    hr_negative_prompt: str = ''
+    force_task_id: str = None
+
+    cached_hr_uc = [None, None]
+    cached_hr_c = [None, None]
+
+    hr_checkpoint_info: dict = field(default=None, init=False)
+    hr_upscale_to_x: int = field(default=0, init=False)
+    hr_upscale_to_y: int = field(default=0, init=False)
+    truncate_x: int = field(default=0, init=False)
+    truncate_y: int = field(default=0, init=False)
+    applied_old_hires_behavior_to: tuple = field(default=None, init=False)
+    latent_scale_mode: dict = field(default=None, init=False)
+    hr_c: tuple | None = field(default=None, init=False)
+    hr_uc: tuple | None = field(default=None, init=False)
+    all_hr_prompts: list = field(default=None, init=False)
+    all_hr_negative_prompts: list = field(default=None, init=False)
+    hr_prompts: list = field(default=None, init=False)
+    hr_negative_prompts: list = field(default=None, init=False)
+    hr_extra_network_data: list = field(default=None, init=False)
+    enable_progressive_growing: bool = field(default=False, init=False)
+    progressive_growing_min_scale: float = field(default=0.25, init=False)
+    progressive_growing_max_scale: float = field(default=1.0, init=False)
+    progressive_growing_steps: int = field(default=4, init=False)
+    progressive_growing_refinement: bool = field(default=True, init=False)
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        self.enable_progressive_growing = getattr(self, 'enable_progressive_growing', False)
+        self.progressive_growing_min_scale = getattr(self, 'progressive_growing_min_scale', 0.25)
+        self.progressive_growing_max_scale = getattr(self, 'progressive_growing_max_scale', 1.0)
+        self.progressive_growing_steps = getattr(self, 'progressive_growing_steps', 4)
+        self.progressive_growing_refinement = getattr(self, 'progressive_growing_refinement', True)
+        
+    def __post_init__(self):
+        super().__post_init__()
+
+        if self.firstphase_width != 0 or self.firstphase_height != 0:
+            self.hr_upscale_to_x = self.width
+            self.hr_upscale_to_y = self.height
+            self.width = self.firstphase_width
+            self.height = self.firstphase_height
+
+        self.cached_hr_uc = StableDiffusionProcessingTxt2Img.cached_hr_uc
+        self.cached_hr_c = StableDiffusionProcessingTxt2Img.cached_hr_c
+
+    def calculate_target_resolution(self):
+        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
+
+    def init(self, all_prompts, all_seeds, all_subseeds):
+        if self.enable_hr:
+            self.extra_generation_params["Denoising strength"] = self.denoising_strength
+
+            if self.hr_checkpoint_name and self.hr_checkpoint_name != 'Use same checkpoint':
+                self.hr_checkpoint_info = sd_models.get_closet_checkpoint_match(self.hr_checkpoint_name)
+
+                if self.hr_checkpoint_info is None:
+                    raise Exception(f'Could not find checkpoint with name {self.hr_checkpoint_name}')
+
+                self.extra_generation_params["Hires checkpoint"] = self.hr_checkpoint_info.short_title
+
+            if self.hr_sampler_name is not None and self.hr_sampler_name != self.sampler_name:
+                self.extra_generation_params["Hires sampler"] = self.hr_sampler_name
+
+            def get_hr_prompt(p, index, prompt_text, **kwargs):
+                hr_prompt = p.all_hr_prompts[index]
+                return hr_prompt if hr_prompt != prompt_text else None
+
+            def get_hr_negative_prompt(p, index, negative_prompt, **kwargs):
+                hr_negative_prompt = p.all_hr_negative_prompts[index]
+                return hr_negative_prompt if hr_negative_prompt != negative_prompt else None
+
+            self.extra_generation_params["Hires prompt"] = get_hr_prompt
+            self.extra_generation_params["Hires negative prompt"] = get_hr_negative_prompt
+
+            self.extra_generation_params["Hires schedule type"] = None  # to be set in sd_samplers_kdiffusion.py
+
+            if self.hr_scheduler is None:
+                self.hr_scheduler = self.scheduler
+
+            self.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 self.latent_scale_mode is None:
+                if not any(x.name == self.hr_upscaler for x in shared.sd_upscalers):
+                    raise Exception(f"could not find upscaler named {self.hr_upscaler}")
+
+            self.calculate_target_resolution()
+
+            if not state.processing_has_refined_job_count:
+                if state.job_count == -1:
+                    state.job_count = self.n_iter
+                if getattr(self, 'txt2img_upscale', False):
+                    total_steps = (self.hr_second_pass_steps or self.steps) * state.job_count
+                else:
+                    total_steps = (self.steps + (self.hr_second_pass_steps or self.steps)) * state.job_count
+                shared.total_tqdm.updateTotal(total_steps)
+                state.job_count = state.job_count * 2
+                state.processing_has_refined_job_count = True
+
+            if self.hr_second_pass_steps:
+                self.extra_generation_params["Hires steps"] = self.hr_second_pass_steps
+
+            if self.hr_upscaler is not None:
+                self.extra_generation_params["Hires upscaler"] = self.hr_upscaler
+
+    def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
+        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)
+
+        if self.enable_progressive_growing:
+            return self.sample_progressive(conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts)
+
+        if self.firstpass_image is not None and self.enable_hr:
+            # here we don't need to generate image, we just take self.firstpass_image and prepare it for hires fix
+
+            if self.latent_scale_mode is None:
+                image = np.array(self.firstpass_image).astype(np.float32) / 255.0 * 2.0 - 1.0
+                image = np.moveaxis(image, 2, 0)
+
+                samples = None
+                decoded_samples = torch.asarray(np.expand_dims(image, 0))
+
+            else:
+                image = np.array(self.firstpass_image).astype(np.float32) / 255.0
+                image = np.moveaxis(image, 2, 0)
+                image = torch.from_numpy(np.expand_dims(image, axis=0))
+                image = image.to(shared.device, dtype=devices.dtype_vae)
+
+                if opts.sd_vae_encode_method != 'Full':
+                    self.extra_generation_params['VAE Encoder'] = opts.sd_vae_encode_method
+
+                samples = images_tensor_to_samples(image, approximation_indexes.get(opts.sd_vae_encode_method), self.sd_model)
+                decoded_samples = None
+                devices.torch_gc()
+
+        else:
+            # here we generate an image normally
+
+            x = self.rng.next()
+            if self.scripts is not None:
+                self.scripts.process_before_every_sampling(
+                    p=self,
+                    x=x,
+                    noise=x,
+                    c=conditioning,
+                    uc=unconditional_conditioning
+                )
+
+            samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))
+            del x
+
+            if not self.enable_hr:
+                return samples
+
+            devices.torch_gc()
+
+            if self.latent_scale_mode is None:
+                decoded_samples = torch.stack(decode_latent_batch(self.sd_model, samples, target_device=devices.cpu, check_for_nans=True)).to(dtype=torch.float32)
+            else:
+                decoded_samples = None
+
+        with sd_models.SkipWritingToConfig():
+            sd_models.reload_model_weights(info=self.hr_checkpoint_info)
+
+        return self.sample_hr_pass(samples, decoded_samples, seeds, subseeds, subseed_strength, prompts)
+
+    def sample_progressive(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
+        is_sdxl = getattr(self.sd_model, 'is_sdxl', False)
+
+        if is_sdxl:
+            min_scale = max(0.5, self.progressive_growing_min_scale)
+        else:
+            min_scale = self.progressive_growing_min_scale
+
+        resolution_steps = np.linspace(min_scale, self.progressive_growing_max_scale, self.progressive_growing_steps)
+
+        initial_width = max(512 if is_sdxl else 64, int(self.width * resolution_steps[0]))
+        initial_height = max(512 if is_sdxl else 64, int(self.height * resolution_steps[0]))
+
+        x = create_random_tensors((opt_C, initial_height // opt_f, initial_width // opt_f), 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)
+        samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))
+
+        for i in range(1, len(resolution_steps)):
+            target_width = int(self.width * resolution_steps[i])
+            target_height = int(self.height * resolution_steps[i])
+
+            if is_sdxl:
+                target_width = max(512, min(1536, target_width))
+                target_height = max(512, min(1536, target_height))
+
+            samples = torch.nn.functional.interpolate(samples, size=(target_height // opt_f, target_width // opt_f), mode='bicubic', align_corners=False)
+
+            if self.progressive_growing_refinement:
+                steps_for_refinement = self.steps // len(resolution_steps)
+                noise = create_random_tensors(samples.shape[1:], 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)
+                decoded_samples = decode_latent_batch(self.sd_model, samples, target_device=devices.cpu, check_for_nans=True)
+                decoded_samples = torch.stack(decoded_samples).float()
+                decoded_samples = torch.clamp((decoded_samples + 1.0) / 2.0, min=0.0, max=1.0)
+                self.image_conditioning = self.img2img_image_conditioning(decoded_samples * 2 - 1, samples)
+
+                samples = self.sampler.sample_img2img(
+                    self,
+                    samples,
+                    noise,
+                    conditioning,
+                    unconditional_conditioning,
+                    steps=steps_for_refinement,
+                    image_conditioning=self.image_conditioning
+                )
+
+        return samples
+    
+    def sample_hr_pass(self, samples, decoded_samples, seeds, subseeds, subseed_strength, prompts):
+        if shared.state.interrupted:
+            return samples
+
+        self.is_hr_pass = True
+        target_width = self.hr_upscale_to_x
+        target_height = self.hr_upscale_to_y
+
+        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"""
+
+            if not self.save_samples() or not opts.save_images_before_highres_fix:
+                return
+
+            if not isinstance(image, Image.Image):
+                image = sd_samplers.sample_to_image(image, index, 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, p=self, suffix="-before-highres-fix")
+
+        img2img_sampler_name = self.hr_sampler_name or self.sampler_name
+
+        self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model)
+
+        if self.latent_scale_mode is not None:
+            for i in range(samples.shape[0]):
+                save_intermediate(samples, i)
+
+            samples = torch.nn.functional.interpolate(samples, size=(target_height // opt_f, target_width // opt_f), mode=self.latent_scale_mode["mode"], antialias=self.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:
+            lowres_samples = torch.clamp((decoded_samples + 1.0) / 2.0, min=0.0, max=1.0)
+
+            batch_images = []
+            for i, x_sample in enumerate(lowres_samples):
+                x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
+                x_sample = x_sample.astype(np.uint8)
+                image = Image.fromarray(x_sample)
+
+                save_intermediate(image, i)
+
+                image = images.resize_image(0, image, 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)
+
+            decoded_samples = torch.from_numpy(np.array(batch_images))
+            decoded_samples = decoded_samples.to(shared.device, dtype=devices.dtype_vae)
+
+            if opts.sd_vae_encode_method != 'Full':
+                self.extra_generation_params['VAE Encoder'] = opts.sd_vae_encode_method
+            samples = images_tensor_to_samples(decoded_samples, approximation_indexes.get(opts.sd_vae_encode_method))
+
+            image_conditioning = self.img2img_image_conditioning(decoded_samples, samples)
+
+        shared.state.nextjob()
+
+        samples = samples[:, :, self.truncate_y//2:samples.shape[2]-(self.truncate_y+1)//2, self.truncate_x//2:samples.shape[3]-(self.truncate_x+1)//2]
+
+        self.rng = rng.ImageRNG(samples.shape[1:], self.seeds, subseeds=self.subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w)
+        noise = self.rng.next()
+
+        # GC now before running the next img2img to prevent running out of memory
+        devices.torch_gc()
+
+        if not self.disable_extra_networks:
+            with devices.autocast():
+                extra_networks.activate(self, self.hr_extra_network_data)
+
+        with devices.autocast():
+            self.calculate_hr_conds()
+
+        sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio(for_hr=True))
+
+        if self.scripts is not None:
+            self.scripts.before_hr(self)
+            self.scripts.process_before_every_sampling(
+                p=self,
+                x=samples,
+                noise=noise,
+                c=self.hr_c,
+                uc=self.hr_uc,
+            )
+
+        samples = self.sampler.sample_img2img(self, samples, noise, self.hr_c, self.hr_uc, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)
+
+        sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio())
+
+        self.sampler = None
+        devices.torch_gc()
+
+        decoded_samples = decode_latent_batch(self.sd_model, samples, target_device=devices.cpu, check_for_nans=True)
+
+        self.is_hr_pass = False
+        return decoded_samples
+
+    def close(self):
+        super().close()
+        self.hr_c = None
+        self.hr_uc = None
+        if not opts.persistent_cond_cache:
+            StableDiffusionProcessingTxt2Img.cached_hr_uc = [None, None]
+            StableDiffusionProcessingTxt2Img.cached_hr_c = [None, None]
+
+    def setup_prompts(self):
+        super().setup_prompts()
+
+        if not self.enable_hr:
+            return
+
+        if self.hr_prompt == '':
+            self.hr_prompt = self.prompt
+
+        if self.hr_negative_prompt == '':
+            self.hr_negative_prompt = self.negative_prompt
+
+        if isinstance(self.hr_prompt, list):
+            self.all_hr_prompts = self.hr_prompt
+        else:
+            self.all_hr_prompts = self.batch_size * self.n_iter * [self.hr_prompt]
+
+        if isinstance(self.hr_negative_prompt, list):
+            self.all_hr_negative_prompts = self.hr_negative_prompt
+        else:
+            self.all_hr_negative_prompts = self.batch_size * self.n_iter * [self.hr_negative_prompt]
+
+        self.all_hr_prompts = [shared.prompt_styles.apply_styles_to_prompt(x, self.styles) for x in self.all_hr_prompts]
+        self.all_hr_negative_prompts = [shared.prompt_styles.apply_negative_styles_to_prompt(x, self.styles) for x in self.all_hr_negative_prompts]
+
+    def calculate_hr_conds(self):
+        if self.hr_c is not None:
+            return
+
+        hr_prompts = prompt_parser.SdConditioning(self.hr_prompts, width=self.hr_upscale_to_x, height=self.hr_upscale_to_y)
+        hr_negative_prompts = prompt_parser.SdConditioning(self.hr_negative_prompts, width=self.hr_upscale_to_x, height=self.hr_upscale_to_y, is_negative_prompt=True)
+
+        sampler_config = sd_samplers.find_sampler_config(self.hr_sampler_name or self.sampler_name)
+        steps = self.hr_second_pass_steps or self.steps
+        total_steps = sampler_config.total_steps(steps) if sampler_config else steps
+
+        self.hr_uc = self.get_conds_with_caching(prompt_parser.get_learned_conditioning, hr_negative_prompts, self.firstpass_steps, [self.cached_hr_uc, self.cached_uc], self.hr_extra_network_data, total_steps)
+        self.hr_c = self.get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, hr_prompts, self.firstpass_steps, [self.cached_hr_c, self.cached_c], self.hr_extra_network_data, total_steps)
+
+    def setup_conds(self):
+        if self.is_hr_pass:
+            # if we are in hr pass right now, the call is being made from the refiner, and we don't need to setup firstpass cons or switch model
+            self.hr_c = None
+            self.calculate_hr_conds()
+            return
+
+        super().setup_conds()
+
+        self.hr_uc = None
+        self.hr_c = None
+
+        if self.enable_hr and self.hr_checkpoint_info is None:
+            if shared.opts.hires_fix_use_firstpass_conds:
+                self.calculate_hr_conds()
+
+            elif lowvram.is_enabled(shared.sd_model) and shared.sd_model.sd_checkpoint_info == sd_models.select_checkpoint():  # if in lowvram mode, we need to calculate conds right away, before the cond NN is unloaded
+                with devices.autocast():
+                    extra_networks.activate(self, self.hr_extra_network_data)
+
+                self.calculate_hr_conds()
+
+                with devices.autocast():
+                    extra_networks.activate(self, self.extra_network_data)
+
+    def get_conds(self):
+        if self.is_hr_pass:
+            return self.hr_c, self.hr_uc
+
+        return super().get_conds()
+
+    def parse_extra_network_prompts(self):
+        res = super().parse_extra_network_prompts()
+
+        if self.enable_hr:
+            self.hr_prompts = self.all_hr_prompts[self.iteration * self.batch_size:(self.iteration + 1) * self.batch_size]
+            self.hr_negative_prompts = self.all_hr_negative_prompts[self.iteration * self.batch_size:(self.iteration + 1) * self.batch_size]
+
+            self.hr_prompts, self.hr_extra_network_data = extra_networks.parse_prompts(self.hr_prompts)
+
+        return res
+
+
+@dataclass(repr=False)
+class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
+    init_images: list = None
+    resize_mode: int = 0
+    denoising_strength: float = 0.75
+    image_cfg_scale: float = None
+    mask: Any = None
+    mask_blur_x: int = 4
+    mask_blur_y: int = 4
+    mask_blur: int = None
+    mask_round: bool = True
+    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
+    latent_mask: Image = None
+    force_task_id: str = None
+
+    image_mask: Any = field(default=None, init=False)
+
+    nmask: torch.Tensor = field(default=None, init=False)
+    image_conditioning: torch.Tensor = field(default=None, init=False)
+    init_img_hash: str = field(default=None, init=False)
+    mask_for_overlay: Image = field(default=None, init=False)
+    init_latent: torch.Tensor = field(default=None, init=False)
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        self.image_mask = self.mask
+        self.mask = None
+        self.initial_noise_multiplier = opts.initial_noise_multiplier if self.initial_noise_multiplier is None else self.initial_noise_multiplier
+
+    @property
+    def mask_blur(self):
+        if self.mask_blur_x == self.mask_blur_y:
+            return self.mask_blur_x
+        return None
+
+    @mask_blur.setter
+    def mask_blur(self, value):
+        if isinstance(value, int):
+            self.mask_blur_x = value
+            self.mask_blur_y = value
+
+    def init(self, all_prompts, all_seeds, all_subseeds):
+        self.extra_generation_params["Denoising strength"] = self.denoising_strength
+
+        self.image_cfg_scale: float = self.image_cfg_scale if shared.sd_model.cond_stage_key == "edit" else None
+
+        self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)
+        crop_region = None
+
+        image_mask = self.image_mask
+
+        if image_mask is not None:
+            # image_mask is passed in as RGBA by Gradio to support alpha masks,
+            # but we still want to support binary masks.
+            image_mask = create_binary_mask(image_mask, round=self.mask_round)
+
+            if self.inpainting_mask_invert:
+                image_mask = ImageOps.invert(image_mask)
+                self.extra_generation_params["Mask mode"] = "Inpaint not masked"
+
+            if self.mask_blur_x > 0:
+                np_mask = np.array(image_mask)
+                kernel_size = 2 * int(2.5 * self.mask_blur_x + 0.5) + 1
+                np_mask = cv2.GaussianBlur(np_mask, (kernel_size, 1), self.mask_blur_x)
+                image_mask = Image.fromarray(np_mask)
+
+            if self.mask_blur_y > 0:
+                np_mask = np.array(image_mask)
+                kernel_size = 2 * int(2.5 * self.mask_blur_y + 0.5) + 1
+                np_mask = cv2.GaussianBlur(np_mask, (1, kernel_size), self.mask_blur_y)
+                image_mask = Image.fromarray(np_mask)
+
+            if self.mask_blur_x > 0 or self.mask_blur_y > 0:
+                self.extra_generation_params["Mask blur"] = self.mask_blur
+
+            if self.inpaint_full_res:
+                self.mask_for_overlay = image_mask
+                mask = image_mask.convert('L')
+                crop_region = masking.get_crop_region_v2(mask, self.inpaint_full_res_padding)
+                if crop_region:
+                    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)
+                    image_mask = images.resize_image(2, mask, self.width, self.height)
+                    self.paste_to = (x1, y1, x2-x1, y2-y1)
+                    self.extra_generation_params["Inpaint area"] = "Only masked"
+                    self.extra_generation_params["Masked area padding"] = self.inpaint_full_res_padding
+                else:
+                    crop_region = None
+                    image_mask = None
+                    self.mask_for_overlay = None
+                    self.inpaint_full_res = False
+                    massage = 'Unable to perform "Inpaint Only mask" because mask is blank, switch to img2img mode.'
+                    model_hijack.comments.append(massage)
+                    logging.info(massage)
+            else:
+                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 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:
+
+            # Save init image
+            if opts.save_init_img:
+                self.init_img_hash = hashlib.md5(img.tobytes()).hexdigest()
+                images.save_image(img, path=opts.outdir_init_images, basename=None, forced_filename=self.init_img_hash, save_to_dirs=False, existing_info=img.info)
+
+            image = images.flatten(img, opts.img2img_background_color)
+
+            if crop_region is None and self.resize_mode != 3:
+                image = images.resize_image(self.resize_mode, image, self.width, self.height)
+
+            if image_mask is not None:
+                if self.mask_for_overlay.size != (image.width, image.height):
+                    self.mask_for_overlay = images.resize_image(self.resize_mode, self.mask_for_overlay, image.width, image.height)
+                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 image_mask is not None:
+                if self.inpainting_fill != 1:
+                    image = masking.fill(image, latent_mask)
+
+                    if self.inpainting_fill == 0:
+                        self.extra_generation_params["Masked content"] = 'fill'
+
+            if add_color_corrections:
+                self.color_corrections.append(setup_color_correction(image))
+
+            image = np.array(image).astype(np.float32) / 255.0
+            image = np.moveaxis(image, 2, 0)
+
+            imgs.append(image)
+
+        if len(imgs) == 1:
+            batch_images = np.expand_dims(imgs[0], axis=0).repeat(self.batch_size, axis=0)
+            if self.overlay_images is not None:
+                self.overlay_images = self.overlay_images * self.batch_size
+
+            if self.color_corrections is not None and len(self.color_corrections) == 1:
+                self.color_corrections = self.color_corrections * self.batch_size
+
+        elif len(imgs) <= self.batch_size:
+            self.batch_size = len(imgs)
+            batch_images = np.array(imgs)
+        else:
+            raise RuntimeError(f"bad number of images passed: {len(imgs)}; expecting {self.batch_size} or less")
+
+        image = torch.from_numpy(batch_images)
+        image = image.to(shared.device, dtype=devices.dtype_vae)
+
+        if opts.sd_vae_encode_method != 'Full':
+            self.extra_generation_params['VAE Encoder'] = opts.sd_vae_encode_method
+
+        self.init_latent = images_tensor_to_samples(image, approximation_indexes.get(opts.sd_vae_encode_method), self.sd_model)
+        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")
+
+        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
+            latmask = latmask[0]
+            if self.mask_round:
+                latmask = np.around(latmask)
+            latmask = np.tile(latmask[None], (self.init_latent.shape[1], 1, 1))
+
+            self.mask = torch.asarray(1.0 - latmask).to(shared.device).type(devices.dtype)
+            self.nmask = torch.asarray(latmask).to(shared.device).type(devices.dtype)
+
+            # this needs to be fixed to be done in sample() using actual seeds for batches
+            if self.inpainting_fill == 2:
+                self.init_latent = self.init_latent * self.mask + create_random_tensors(self.init_latent.shape[1:], all_seeds[0:self.init_latent.shape[0]]) * self.nmask
+                self.extra_generation_params["Masked content"] = 'latent noise'
+
+            elif self.inpainting_fill == 3:
+                self.init_latent = self.init_latent * self.mask
+                self.extra_generation_params["Masked content"] = 'latent nothing'
+
+        self.image_conditioning = self.img2img_image_conditioning(image * 2 - 1, self.init_latent, image_mask, self.mask_round)
+
+    def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
+        x = self.rng.next()
+
+        if self.initial_noise_multiplier != 1.0:
+            self.extra_generation_params["Noise multiplier"] = self.initial_noise_multiplier
+            x *= self.initial_noise_multiplier
+
+        if self.scripts is not None:
+            self.scripts.process_before_every_sampling(
+                p=self,
+                x=self.init_latent,
+                noise=x,
+                c=conditioning,
+                uc=unconditional_conditioning
+            )
+        samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning, image_conditioning=self.image_conditioning)
+
+        if self.mask is not None:
+            blended_samples = samples * self.nmask + self.init_latent * self.mask
+
+            if self.scripts is not None:
+                mba = scripts.MaskBlendArgs(samples, self.nmask, self.init_latent, self.mask, blended_samples)
+                self.scripts.on_mask_blend(self, mba)
+                blended_samples = mba.blended_latent
+
+            samples = blended_samples
+
+        del x
+        devices.torch_gc()
+
+        return samples
+
+    def get_token_merging_ratio(self, for_hr=False):
+        return self.token_merging_ratio or ("token_merging_ratio" in self.override_settings and opts.token_merging_ratio) or opts.token_merging_ratio_img2img or opts.token_merging_ratio

sec/sampling.py

@@ -0,0 +1,726 @@
+import math
+
+from scipy import integrate
+import torch
+from torch import nn
+from torchdiffeq import odeint
+import torchsde
+from tqdm.auto import trange, tqdm
+
+from . import utils
+
+
+def append_zero(x):
+    return torch.cat([x, x.new_zeros([1])])
+
+
+def get_sigmas_karras(n, sigma_min, sigma_max, rho=7., device='cpu'):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    ramp = torch.linspace(0, 1, n)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+    return append_zero(sigmas).to(device)
+
+
+def get_sigmas_exponential(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs an exponential noise schedule."""
+    sigmas = torch.linspace(math.log(sigma_max), math.log(sigma_min), n, device=device).exp()
+    return append_zero(sigmas)
+
+
+def get_sigmas_polyexponential(n, sigma_min, sigma_max, rho=1., device='cpu'):
+    """Constructs an polynomial in log sigma noise schedule."""
+    ramp = torch.linspace(1, 0, n, device=device) ** rho
+    sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min))
+    return append_zero(sigmas)
+
+
+def get_sigmas_vp(n, beta_d=19.9, beta_min=0.1, eps_s=1e-3, device='cpu'):
+    """Constructs a continuous VP noise schedule."""
+    t = torch.linspace(1, eps_s, n, device=device)
+    sigmas = torch.sqrt(torch.exp(beta_d * t ** 2 / 2 + beta_min * t) - 1)
+    return append_zero(sigmas)
+
+
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / utils.append_dims(sigma, x.ndim)
+
+
+def get_ancestral_step(sigma_from, sigma_to, eta=1.):
+    """Calculates the noise level (sigma_down) to step down to and the amount
+    of noise to add (sigma_up) when doing an ancestral sampling step."""
+    if not eta:
+        return sigma_to, 0.
+    sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
+    sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
+    return sigma_down, sigma_up
+
+
+def default_noise_sampler(x):
+    return lambda sigma, sigma_next: torch.randn_like(x)
+
+
+class BatchedBrownianTree:
+    """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
+
+    def __init__(self, x, t0, t1, seed=None, **kwargs):
+        t0, t1, self.sign = self.sort(t0, t1)
+        w0 = kwargs.get('w0', torch.zeros_like(x))
+        if seed is None:
+            seed = torch.randint(0, 2 ** 63 - 1, []).item()
+        self.batched = True
+        try:
+            assert len(seed) == x.shape[0]
+            w0 = w0[0]
+        except TypeError:
+            seed = [seed]
+            self.batched = False
+        self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
+
+    @staticmethod
+    def sort(a, b):
+        return (a, b, 1) if a < b else (b, a, -1)
+
+    def __call__(self, t0, t1):
+        t0, t1, sign = self.sort(t0, t1)
+        w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
+        return w if self.batched else w[0]
+
+
+class BrownianTreeNoiseSampler:
+    """A noise sampler backed by a torchsde.BrownianTree.
+
+    Args:
+        x (Tensor): The tensor whose shape, device and dtype to use to generate
+            random samples.
+        sigma_min (float): The low end of the valid interval.
+        sigma_max (float): The high end of the valid interval.
+        seed (int or List[int]): The random seed. If a list of seeds is
+            supplied instead of a single integer, then the noise sampler will
+            use one BrownianTree per batch item, each with its own seed.
+        transform (callable): A function that maps sigma to the sampler's
+            internal timestep.
+    """
+
+    def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x):
+        self.transform = transform
+        t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
+        self.tree = BatchedBrownianTree(x, t0, t1, seed)
+
+    def __call__(self, sigma, sigma_next):
+        t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
+        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
+
+
+@torch.no_grad()
+def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = to_d(x, sigma_hat, denoised)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+        dt = sigmas[i + 1] - sigma_hat
+        # Euler method
+        x = x + d * dt
+    return x
+
+
+@torch.no_grad()
+def sample_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+    """Ancestral sampling with Euler method steps."""
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        d = to_d(x, sigmas[i], denoised)
+        # Euler method
+        dt = sigma_down - sigmas[i]
+        x = x + d * dt
+        if sigmas[i + 1] > 0:
+            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+    return x
+
+
+@torch.no_grad()
+def sample_heun(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = to_d(x, sigma_hat, denoised)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+        dt = sigmas[i + 1] - sigma_hat
+        if sigmas[i + 1] == 0:
+            # Euler method
+            x = x + d * dt
+        else:
+            # Heun's method
+            x_2 = x + d * dt
+            denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
+            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
+            d_prime = (d + d_2) / 2
+            x = x + d_prime * dt
+    return x
+
+
+@torch.no_grad()
+def sample_dpm_2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = to_d(x, sigma_hat, denoised)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            # Euler method
+            dt = sigmas[i + 1] - sigma_hat
+            x = x + d * dt
+        else:
+            # DPM-Solver-2
+            sigma_mid = sigma_hat.log().lerp(sigmas[i + 1].log(), 0.5).exp()
+            dt_1 = sigma_mid - sigma_hat
+            dt_2 = sigmas[i + 1] - sigma_hat
+            x_2 = x + d * dt_1
+            denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
+            d_2 = to_d(x_2, sigma_mid, denoised_2)
+            x = x + d_2 * dt_2
+    return x
+
+
+@torch.no_grad()
+def sample_dpm_2_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+    """Ancestral sampling with DPM-Solver second-order steps."""
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        d = to_d(x, sigmas[i], denoised)
+        if sigma_down == 0:
+            # Euler method
+            dt = sigma_down - sigmas[i]
+            x = x + d * dt
+        else:
+            # DPM-Solver-2
+            sigma_mid = sigmas[i].log().lerp(sigma_down.log(), 0.5).exp()
+            dt_1 = sigma_mid - sigmas[i]
+            dt_2 = sigma_down - sigmas[i]
+            x_2 = x + d * dt_1
+            denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
+            d_2 = to_d(x_2, sigma_mid, denoised_2)
+            x = x + d_2 * dt_2
+            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+    return x
+
+
+def linear_multistep_coeff(order, t, i, j):
+    if order - 1 > i:
+        raise ValueError(f'Order {order} too high for step {i}')
+    def fn(tau):
+        prod = 1.
+        for k in range(order):
+            if j == k:
+                continue
+            prod *= (tau - t[i - k]) / (t[i - j] - t[i - k])
+        return prod
+    return integrate.quad(fn, t[i], t[i + 1], epsrel=1e-4)[0]
+
+
+@torch.no_grad()
+def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4):
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    sigmas_cpu = sigmas.detach().cpu().numpy()
+    ds = []
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        d = to_d(x, sigmas[i], denoised)
+        ds.append(d)
+        if len(ds) > order:
+            ds.pop(0)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        cur_order = min(i + 1, order)
+        coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)]
+        x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
+    return x
+
+
+@torch.no_grad()
+def log_likelihood(model, x, sigma_min, sigma_max, extra_args=None, atol=1e-4, rtol=1e-4):
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    v = torch.randint_like(x, 2) * 2 - 1
+    fevals = 0
+    def ode_fn(sigma, x):
+        nonlocal fevals
+        with torch.enable_grad():
+            x = x[0].detach().requires_grad_()
+            denoised = model(x, sigma * s_in, **extra_args)
+            d = to_d(x, sigma, denoised)
+            fevals += 1
+            grad = torch.autograd.grad((d * v).sum(), x)[0]
+            d_ll = (v * grad).flatten(1).sum(1)
+        return d.detach(), d_ll
+    x_min = x, x.new_zeros([x.shape[0]])
+    t = x.new_tensor([sigma_min, sigma_max])
+    sol = odeint(ode_fn, x_min, t, atol=atol, rtol=rtol, method='dopri5')
+    latent, delta_ll = sol[0][-1], sol[1][-1]
+    ll_prior = torch.distributions.Normal(0, sigma_max).log_prob(latent).flatten(1).sum(1)
+    return ll_prior + delta_ll, {'fevals': fevals}
+
+
+class PIDStepSizeController:
+    """A PID controller for ODE adaptive step size control."""
+    def __init__(self, h, pcoeff, icoeff, dcoeff, order=1, accept_safety=0.81, eps=1e-8):
+        self.h = h
+        self.b1 = (pcoeff + icoeff + dcoeff) / order
+        self.b2 = -(pcoeff + 2 * dcoeff) / order
+        self.b3 = dcoeff / order
+        self.accept_safety = accept_safety
+        self.eps = eps
+        self.errs = []
+
+    def limiter(self, x):
+        return 1 + math.atan(x - 1)
+
+    def propose_step(self, error):
+        inv_error = 1 / (float(error) + self.eps)
+        if not self.errs:
+            self.errs = [inv_error, inv_error, inv_error]
+        self.errs[0] = inv_error
+        factor = self.errs[0] ** self.b1 * self.errs[1] ** self.b2 * self.errs[2] ** self.b3
+        factor = self.limiter(factor)
+        accept = factor >= self.accept_safety
+        if accept:
+            self.errs[2] = self.errs[1]
+            self.errs[1] = self.errs[0]
+        self.h *= factor
+        return accept
+
+
+class DPMSolver(nn.Module):
+    """DPM-Solver. See https://arxiv.org/abs/2206.00927."""
+
+    def __init__(self, model, extra_args=None, eps_callback=None, info_callback=None):
+        super().__init__()
+        self.model = model
+        self.extra_args = {} if extra_args is None else extra_args
+        self.eps_callback = eps_callback
+        self.info_callback = info_callback
+
+    def t(self, sigma):
+        return -sigma.log()
+
+    def sigma(self, t):
+        return t.neg().exp()
+
+    def eps(self, eps_cache, key, x, t, *args, **kwargs):
+        if key in eps_cache:
+            return eps_cache[key], eps_cache
+        sigma = self.sigma(t) * x.new_ones([x.shape[0]])
+        eps = (x - self.model(x, sigma, *args, **self.extra_args, **kwargs)) / self.sigma(t)
+        if self.eps_callback is not None:
+            self.eps_callback()
+        return eps, {key: eps, **eps_cache}
+
+    def dpm_solver_1_step(self, x, t, t_next, eps_cache=None):
+        eps_cache = {} if eps_cache is None else eps_cache
+        h = t_next - t
+        eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
+        x_1 = x - self.sigma(t_next) * h.expm1() * eps
+        return x_1, eps_cache
+
+    def dpm_solver_2_step(self, x, t, t_next, r1=1 / 2, eps_cache=None):
+        eps_cache = {} if eps_cache is None else eps_cache
+        h = t_next - t
+        eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
+        s1 = t + r1 * h
+        u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
+        eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
+        x_2 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / (2 * r1) * h.expm1() * (eps_r1 - eps)
+        return x_2, eps_cache
+
+    def dpm_solver_3_step(self, x, t, t_next, r1=1 / 3, r2=2 / 3, eps_cache=None):
+        eps_cache = {} if eps_cache is None else eps_cache
+        h = t_next - t
+        eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
+        s1 = t + r1 * h
+        s2 = t + r2 * h
+        u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
+        eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
+        u2 = x - self.sigma(s2) * (r2 * h).expm1() * eps - self.sigma(s2) * (r2 / r1) * ((r2 * h).expm1() / (r2 * h) - 1) * (eps_r1 - eps)
+        eps_r2, eps_cache = self.eps(eps_cache, 'eps_r2', u2, s2)
+        x_3 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / r2 * (h.expm1() / h - 1) * (eps_r2 - eps)
+        return x_3, eps_cache
+
+    def dpm_solver_fast(self, x, t_start, t_end, nfe, eta=0., s_noise=1., noise_sampler=None):
+        noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+        if not t_end > t_start and eta:
+            raise ValueError('eta must be 0 for reverse sampling')
+
+        m = math.floor(nfe / 3) + 1
+        ts = torch.linspace(t_start, t_end, m + 1, device=x.device)
+
+        if nfe % 3 == 0:
+            orders = [3] * (m - 2) + [2, 1]
+        else:
+            orders = [3] * (m - 1) + [nfe % 3]
+
+        for i in range(len(orders)):
+            eps_cache = {}
+            t, t_next = ts[i], ts[i + 1]
+            if eta:
+                sd, su = get_ancestral_step(self.sigma(t), self.sigma(t_next), eta)
+                t_next_ = torch.minimum(t_end, self.t(sd))
+                su = (self.sigma(t_next) ** 2 - self.sigma(t_next_) ** 2) ** 0.5
+            else:
+                t_next_, su = t_next, 0.
+
+            eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
+            denoised = x - self.sigma(t) * eps
+            if self.info_callback is not None:
+                self.info_callback({'x': x, 'i': i, 't': ts[i], 't_up': t, 'denoised': denoised})
+
+            if orders[i] == 1:
+                x, eps_cache = self.dpm_solver_1_step(x, t, t_next_, eps_cache=eps_cache)
+            elif orders[i] == 2:
+                x, eps_cache = self.dpm_solver_2_step(x, t, t_next_, eps_cache=eps_cache)
+            else:
+                x, eps_cache = self.dpm_solver_3_step(x, t, t_next_, eps_cache=eps_cache)
+
+            x = x + su * s_noise * noise_sampler(self.sigma(t), self.sigma(t_next))
+
+        return x
+
+    def dpm_solver_adaptive(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None):
+        noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+        if order not in {2, 3}:
+            raise ValueError('order should be 2 or 3')
+        forward = t_end > t_start
+        if not forward and eta:
+            raise ValueError('eta must be 0 for reverse sampling')
+        h_init = abs(h_init) * (1 if forward else -1)
+        atol = torch.tensor(atol)
+        rtol = torch.tensor(rtol)
+        s = t_start
+        x_prev = x
+        accept = True
+        pid = PIDStepSizeController(h_init, pcoeff, icoeff, dcoeff, 1.5 if eta else order, accept_safety)
+        info = {'steps': 0, 'nfe': 0, 'n_accept': 0, 'n_reject': 0}
+
+        while s < t_end - 1e-5 if forward else s > t_end + 1e-5:
+            eps_cache = {}
+            t = torch.minimum(t_end, s + pid.h) if forward else torch.maximum(t_end, s + pid.h)
+            if eta:
+                sd, su = get_ancestral_step(self.sigma(s), self.sigma(t), eta)
+                t_ = torch.minimum(t_end, self.t(sd))
+                su = (self.sigma(t) ** 2 - self.sigma(t_) ** 2) ** 0.5
+            else:
+                t_, su = t, 0.
+
+            eps, eps_cache = self.eps(eps_cache, 'eps', x, s)
+            denoised = x - self.sigma(s) * eps
+
+            if order == 2:
+                x_low, eps_cache = self.dpm_solver_1_step(x, s, t_, eps_cache=eps_cache)
+                x_high, eps_cache = self.dpm_solver_2_step(x, s, t_, eps_cache=eps_cache)
+            else:
+                x_low, eps_cache = self.dpm_solver_2_step(x, s, t_, r1=1 / 3, eps_cache=eps_cache)
+                x_high, eps_cache = self.dpm_solver_3_step(x, s, t_, eps_cache=eps_cache)
+            delta = torch.maximum(atol, rtol * torch.maximum(x_low.abs(), x_prev.abs()))
+            error = torch.linalg.norm((x_low - x_high) / delta) / x.numel() ** 0.5
+            accept = pid.propose_step(error)
+            if accept:
+                x_prev = x_low
+                x = x_high + su * s_noise * noise_sampler(self.sigma(s), self.sigma(t))
+                s = t
+                info['n_accept'] += 1
+            else:
+                info['n_reject'] += 1
+            info['nfe'] += order
+            info['steps'] += 1
+
+            if self.info_callback is not None:
+                self.info_callback({'x': x, 'i': info['steps'] - 1, 't': s, 't_up': s, 'denoised': denoised, 'error': error, 'h': pid.h, **info})
+
+        return x, info
+
+
+@torch.no_grad()
+def sample_dpm_fast(model, x, sigma_min, sigma_max, n, extra_args=None, callback=None, disable=None, eta=0., s_noise=1., noise_sampler=None):
+    """DPM-Solver-Fast (fixed step size). See https://arxiv.org/abs/2206.00927."""
+    if sigma_min <= 0 or sigma_max <= 0:
+        raise ValueError('sigma_min and sigma_max must not be 0')
+    with tqdm(total=n, disable=disable) as pbar:
+        dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
+        if callback is not None:
+            dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
+        return dpm_solver.dpm_solver_fast(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), n, eta, s_noise, noise_sampler)
+
+
+@torch.no_grad()
+def sample_dpm_adaptive(model, x, sigma_min, sigma_max, extra_args=None, callback=None, disable=None, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None, return_info=False):
+    """DPM-Solver-12 and 23 (adaptive step size). See https://arxiv.org/abs/2206.00927."""
+    if sigma_min <= 0 or sigma_max <= 0:
+        raise ValueError('sigma_min and sigma_max must not be 0')
+    with tqdm(disable=disable) as pbar:
+        dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
+        if callback is not None:
+            dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
+        x, info = dpm_solver.dpm_solver_adaptive(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_safety, eta, s_noise, noise_sampler)
+    if return_info:
+        return x, info
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+    """Ancestral sampling with DPM-Solver++(2S) second-order steps."""
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigma_down == 0:
+            # Euler method
+            d = to_d(x, sigmas[i], denoised)
+            dt = sigma_down - sigmas[i]
+            x = x + d * dt
+        else:
+            # DPM-Solver++(2S)
+            t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
+            r = 1 / 2
+            h = t_next - t
+            s = t + r * h
+            x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * denoised
+            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
+            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_2
+        # Noise addition
+        if sigmas[i + 1] > 0:
+            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
+    """DPM-Solver++ (stochastic)."""
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            # Euler method
+            d = to_d(x, sigmas[i], denoised)
+            dt = sigmas[i + 1] - sigmas[i]
+            x = x + d * dt
+        else:
+            # DPM-Solver++
+            t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
+            h = t_next - t
+            s = t + h * r
+            fac = 1 / (2 * r)
+
+            # Step 1
+            sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(s), eta)
+            s_ = t_fn(sd)
+            x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - (t - s_).expm1() * denoised
+            x_2 = x_2 + noise_sampler(sigma_fn(t), sigma_fn(s)) * s_noise * su
+            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
+
+            # Step 2
+            sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(t_next), eta)
+            t_next_ = t_fn(sd)
+            denoised_d = (1 - fac) * denoised + fac * denoised_2
+            x = (sigma_fn(t_next_) / sigma_fn(t)) * x - (t - t_next_).expm1() * denoised_d
+            x = x + noise_sampler(sigma_fn(t), sigma_fn(t_next)) * s_noise * su
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None):
+    """DPM-Solver++(2M)."""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+    old_denoised = None
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
+        h = t_next - t
+        if old_denoised is None or sigmas[i + 1] == 0:
+            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
+        else:
+            h_last = t - t_fn(sigmas[i - 1])
+            r = h_last / h
+            denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
+            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d
+        old_denoised = denoised
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
+    """DPM-Solver++(2M) SDE."""
+
+    if solver_type not in {'heun', 'midpoint'}:
+        raise ValueError('solver_type must be \'heun\' or \'midpoint\'')
+
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    old_denoised = None
+    h_last = None
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            # Denoising step
+            x = denoised
+        else:
+            # DPM-Solver++(2M) SDE
+            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = s - t
+            eta_h = eta * h
+
+            x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised
+
+            if old_denoised is not None:
+                r = h_last / h
+                if solver_type == 'heun':
+                    x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised)
+                elif solver_type == 'midpoint':
+                    x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised)
+
+            if eta:
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
+
+        old_denoised = denoised
+        h_last = h
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+    """DPM-Solver++(3M) SDE."""
+
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    denoised_1, denoised_2 = None, None
+    h_1, h_2 = None, None
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            # Denoising step
+            x = denoised
+        else:
+            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = s - t
+            h_eta = h * (eta + 1)
+
+            x = torch.exp(-h_eta) * x + (-h_eta).expm1().neg() * denoised
+
+            if h_2 is not None:
+                r0 = h_1 / h
+                r1 = h_2 / h
+                d1_0 = (denoised - denoised_1) / r0
+                d1_1 = (denoised_1 - denoised_2) / r1
+                d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
+                d2 = (d1_0 - d1_1) / (r0 + r1)
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                phi_3 = phi_2 / h_eta - 0.5
+                x = x + phi_2 * d1 - phi_3 * d2
+            elif h_1 is not None:
+                r = h_1 / h
+                d = (denoised - denoised_1) / r
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                x = x + phi_2 * d
+
+            if eta:
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
+
+        denoised_1, denoised_2 = denoised, denoised_1
+        h_1, h_2 = h, h_1
+    return x
+
+@torch.no_grad()
+def sampler_dpmu(model, x, sigmas, extra_args=None, callback=None, disable=None):
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+    last_x = None
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = x if i == 0 else model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
+        h = t_next - t
+        if sigmas[i + 1] == 0:
+            return torch.lerp(denoised, last_x, 0.5) * dpmu_factor
+        else:
+            h_last = t - t_fn(sigmas[i - 1])
+            r = h_last / h
+            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * (1 + 1 / (2 * r)) * denoised / 2
+        if sigmas[i + 2] == 0:
+            last_x = x
+        torch.clamp(x, -1.0, 1.0)
+    return x

sec/sd_schedulers.py

@@ -0,0 +1,146 @@
+import dataclasses
+import torch
+import k_diffusion
+import numpy as np
+from scipy import stats
+import modules.simple_karras_exponential_scheduler as simple_kes
+from modules import shared
+
+
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / sigma
+
+
+k_diffusion.sampling.to_d = to_d
+
+
+@dataclasses.dataclass
+class Scheduler:
+    name: str
+    label: str
+    function: any
+
+    default_rho: float = -1
+    need_inner_model: bool = False
+    aliases: list = None
+
+
+def uniform(n, sigma_min, sigma_max, inner_model, device):
+    return inner_model.get_sigmas(n).to(device)
+
+
+def sgm_uniform(n, sigma_min, sigma_max, inner_model, device):
+    start = inner_model.sigma_to_t(torch.tensor(sigma_max))
+    end = inner_model.sigma_to_t(torch.tensor(sigma_min))
+    sigs = [
+        inner_model.t_to_sigma(ts)
+        for ts in torch.linspace(start, end, n + 1)[:-1]
+    ]
+    sigs += [0.0]
+    return torch.FloatTensor(sigs).to(device)
+
+
+def get_align_your_steps_sigmas(n, sigma_min, sigma_max, device):
+    # https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html
+    def loglinear_interp(t_steps, num_steps):
+        """
+        Performs log-linear interpolation of a given array of decreasing numbers.
+        """
+        xs = np.linspace(0, 1, len(t_steps))
+        ys = np.log(t_steps[::-1])
+
+        new_xs = np.linspace(0, 1, num_steps)
+        new_ys = np.interp(new_xs, xs, ys)
+
+        interped_ys = np.exp(new_ys)[::-1].copy()
+        return interped_ys
+
+    if shared.sd_model.is_sdxl:
+        sigmas = [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.029]
+    else:
+        # Default to SD 1.5 sigmas.
+        sigmas = [14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.029]
+
+    if n != len(sigmas):
+        sigmas = np.append(loglinear_interp(sigmas, n), [0.0])
+    else:
+        sigmas.append(0.0)
+
+    return torch.FloatTensor(sigmas).to(device)
+
+
+def kl_optimal(n, sigma_min, sigma_max, device):
+    alpha_min = torch.arctan(torch.tensor(sigma_min, device=device))
+    alpha_max = torch.arctan(torch.tensor(sigma_max, device=device))
+    step_indices = torch.arange(n + 1, device=device)
+    sigmas = torch.tan(step_indices / n * alpha_min + (1.0 - step_indices / n) * alpha_max)
+    return sigmas
+
+
+def simple_scheduler(n, sigma_min, sigma_max, inner_model, device):
+    sigs = []
+    ss = len(inner_model.sigmas) / n
+    for x in range(n):
+        sigs += [float(inner_model.sigmas[-(1 + int(x * ss))])]
+    sigs += [0.0]
+    return torch.FloatTensor(sigs).to(device)
+
+
+def normal_scheduler(n, sigma_min, sigma_max, inner_model, device, sgm=False, floor=False):
+    start = inner_model.sigma_to_t(torch.tensor(sigma_max))
+    end = inner_model.sigma_to_t(torch.tensor(sigma_min))
+
+    if sgm:
+        timesteps = torch.linspace(start, end, n + 1)[:-1]
+    else:
+        timesteps = torch.linspace(start, end, n)
+
+    sigs = []
+    for x in range(len(timesteps)):
+        ts = timesteps[x]
+        sigs.append(inner_model.t_to_sigma(ts))
+    sigs += [0.0]
+    return torch.FloatTensor(sigs).to(device)
+
+
+def ddim_scheduler(n, sigma_min, sigma_max, inner_model, device):
+    sigs = []
+    ss = max(len(inner_model.sigmas) // n, 1)
+    x = 1
+    while x < len(inner_model.sigmas):
+        sigs += [float(inner_model.sigmas[x])]
+        x += ss
+    sigs = sigs[::-1]
+    sigs += [0.0]
+    return torch.FloatTensor(sigs).to(device)
+
+
+def beta_scheduler(n, sigma_min, sigma_max, inner_model, device):
+    # From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024) """
+    alpha = shared.opts.beta_dist_alpha
+    beta = shared.opts.beta_dist_beta
+    timesteps = 1 - np.linspace(0, 1, n)
+    timesteps = [stats.beta.ppf(x, alpha, beta) for x in timesteps]
+    sigmas = [sigma_min + (x * (sigma_max-sigma_min)) for x in timesteps]
+    sigmas += [0.0]
+    return torch.FloatTensor(sigmas).to(device)
+
+
+schedulers = [
+    Scheduler('automatic', 'Automatic', None),
+    Scheduler('uniform', 'Uniform', uniform, need_inner_model=True),
+    Scheduler('karras', 'Karras', k_diffusion.sampling.get_sigmas_karras, default_rho=7.0),
+    Scheduler('exponential', 'Exponential', k_diffusion.sampling.get_sigmas_exponential),
+    Scheduler('polyexponential', 'Polyexponential', k_diffusion.sampling.get_sigmas_polyexponential, default_rho=1.0),
+    Scheduler('sgm_uniform', 'SGM Uniform', sgm_uniform, need_inner_model=True, aliases=["SGMUniform"]),
+    Scheduler('kl_optimal', 'KL Optimal', kl_optimal),
+    Scheduler('align_your_steps', 'Align Your Steps', get_align_your_steps_sigmas),
+    Scheduler('simple', 'Simple', simple_scheduler, need_inner_model=True),
+    Scheduler('normal', 'Normal', normal_scheduler, need_inner_model=True),
+    Scheduler('ddim', 'DDIM', ddim_scheduler, need_inner_model=True),
+    Scheduler('beta', 'Beta', beta_scheduler, need_inner_model=True),
+    Scheduler('karras_exponential', 'Karras Exponential', simple_kes.simple_karras_exponential_scheduler),
+]
+
+schedulers_map = {**{x.name: x for x in schedulers}, **{x.label: x for x in schedulers}}

sec/simple_karras_exponential_scheduler.py

@@ -0,0 +1,390 @@
+#simple_karras_exponential_scheduler.py
+import torch
+import logging
+from k_diffusion.sampling import get_sigmas_karras, get_sigmas_exponential
+import os
+import yaml
+import random
+from watchdog.observers import Observer
+from watchdog.events import FileSystemEventHandler
+from datetime import datetime
+import warnings
+import os
+import logging
+from datetime import datetime
+def get_random_or_default(scheduler_config, key_prefix, default_value, global_randomize):
+    """Helper function to either randomize a value based on conditions or return the default."""
+    
+    # Determine if we should randomize based on global and individual flags
+    randomize_flag = global_randomize or scheduler_config.get(f'{key_prefix}_rand', False)
+
+    if randomize_flag:
+        # Use specified min/max values for randomization if they exist, else use default range
+        rand_min = scheduler_config.get(f'{key_prefix}_rand_min', default_value * 0.8)
+        rand_max = scheduler_config.get(f'{key_prefix}_rand_max', default_value * 1.2)
+        value = random.uniform(rand_min, rand_max)
+        custom_logger.info(f"Randomized {key_prefix}: {value}")
+    else:
+        # Use default value if no randomization is applied
+        value = default_value
+        custom_logger.info(f"Using default {key_prefix}: {value}")
+
+    return value
+    
+
+class CustomLogger:
+    def __init__(self, log_name, print_to_console=False, debug_enabled=False):
+        self.print_to_console = print_to_console #prints to console
+        self.debug_enabled = debug_enabled #logs debug messages
+
+        # Create folders for generation info and error logs
+        gen_log_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'simple_kes_generation')
+        error_log_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'simple_kes_error')
+
+        os.makedirs(gen_log_dir, exist_ok=True)
+        os.makedirs(error_log_dir, exist_ok=True)
+
+        # Get current time in HH-MM-SS format
+        current_time = datetime.now().strftime('%H-%M-%S')
+
+        # Create file paths for the log files
+        gen_log_file_path = os.path.join(gen_log_dir, f'{current_time}.log')
+        error_log_file_path = os.path.join(error_log_dir, f'{current_time}.log')
+
+        # Set up generation logger
+        #self.gen_logger = logging.getLogger(f'{log_name}_generation')
+        self.gen_logger = logging.getLogger('simple_kes_generation')
+        self.gen_logger.setLevel(logging.DEBUG)
+        self._setup_file_handler(self.gen_logger, gen_log_file_path)
+
+        # Set up error logger
+        self.error_logger = logging.getLogger(f'{log_name}_error')
+        self.error_logger.setLevel(logging.ERROR)
+        self._setup_file_handler(self.error_logger, error_log_file_path)
+        
+        # Prevent log propagation to root logger (important to avoid accidental console logging)
+        self.gen_logger.propagate = False
+        self.error_logger.propagate = False
+        
+        
+        # Optionally print to console
+        if self.print_to_console:
+            self._setup_console_handler(self.gen_logger)
+            self._setup_console_handler(self.error_logger)
+
+    def _setup_file_handler(self, logger, file_path):
+        """Set up file handler for logging to a file."""
+        file_handler = logging.FileHandler(file_path, mode='a')
+        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+
+    def _setup_console_handler(self, logger):
+        """Optionally set up a console handler for logging to the console."""
+        console_handler = logging.StreamHandler()
+        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+        console_handler.setFormatter(formatter)
+        logger.addHandler(console_handler)
+
+    def log_debug(self, message):
+        """Log a debug message."""
+        if self.debug_enabled:
+            self.gen_logger.debug(message)
+
+    def log_info(self, message):
+        """Log an info message."""
+        self.gen_logger.info(message)
+    info=log_info #alias created
+
+    def log_error(self, message):
+        """Log an error message."""
+        self.error_logger.error(message)
+
+    def enable_console_logging(self):
+        """Enable console logging dynamically."""
+        if not any(isinstance(handler, logging.StreamHandler) for handler in self.gen_logger.handlers):
+            self._setup_console_handler(self.gen_logger)
+
+        if not any(isinstance(handler, logging.StreamHandler) for handler in self.error_logger.handlers):
+            self._setup_console_handler(self.error_logger)
+
+# Usage example
+custom_logger = CustomLogger('simple_kes', print_to_console=False, debug_enabled=True)
+
+# Logging examples
+#custom_logger.log_debug("Debug message: Using default sigma_min: 0.01")
+#custom_logger.info("Info message: Step completed successfully.")
+#custom_logger.log_error("Error message: Something went wrong!")
+      
+    
+class ConfigManagerYaml:
+    def __init__(self, config_path):
+        self.config_path = config_path
+        self.config_data = self.load_config()  # Initialize config_data here
+
+    def load_config(self):
+        try:
+            with open(self.config_path, 'r') as f:
+                user_config = yaml.safe_load(f)
+                return user_config
+        except FileNotFoundError:
+            print(f"Config file not found: {self.config_path}. Using empty config.")
+            return {}
+        except yaml.YAMLError as e:
+            print(f"Error loading config file: {e}")
+            return {}
+            
+
+#ConfigWatcher monitors changes to the config file and reloads during program use (so you can continue work without resetting the program)        
+class ConfigWatcher(FileSystemEventHandler):
+    def __init__(self, config_manager, config_path):
+        self.config_manager = config_manager
+        self.config_path = config_path
+
+    def on_modified(self, event):
+        if event.src_path == self.config_path:
+            logging.info(f"Config file {self.config_path} modified. Reloading config.")
+            self.config_manager.config_data = self.config_manager.load_config()                
+ 
+
+ 
+def start_config_watcher(config_manager, config_path):
+    event_handler = ConfigWatcher(config_manager, config_path)
+    observer = Observer()
+    observer.schedule(event_handler, os.path.dirname(config_path), recursive=False)
+    observer.start()
+    return observer
+
+
+"""
+    Scheduler function that blends sigma sequences using Karras and Exponential methods with adaptive parameters.
+
+    Parameters are dynamically updated if the config file changes during execution.
+"""
+# If user config is provided, update default config with user values
+config_path = "modules/simple_kes_scheduler.yaml"
+config_manager = ConfigManagerYaml(config_path)    
+
+
+# Start watching for config changes
+observer = start_config_watcher(config_manager, config_path)
+
+
+def simple_karras_exponential_scheduler(
+    n, device, sigma_min=0.01, sigma_max=50, start_blend=0.1, end_blend=0.5, 
+    sharpness=0.95, early_stopping_threshold=0.01, update_interval=10, initial_step_size=0.9, 
+    final_step_size=0.2, initial_noise_scale=1.25, final_noise_scale=0.8, smooth_blend_factor=11, step_size_factor=0.8, noise_scale_factor=0.9, randomize=False, user_config=None
+):    
+    """
+    Scheduler function that blends sigma sequences using Karras and Exponential methods with adaptive parameters.
+
+    Parameters:
+        n (int): Number of steps.
+        sigma_min (float): Minimum sigma value.
+        sigma_max (float): Maximum sigma value.
+        device (torch.device): The device on which to perform computations (e.g., 'cuda' or 'cpu').
+        start_blend (float): Initial blend factor for dynamic blending.
+        end_bend (float): Final blend factor for dynamic blending.
+        sharpen_factor (float): Sharpening factor to be applied adaptively.
+        early_stopping_threshold (float): Threshold to trigger early stopping.
+        update_interval (int): Interval to update blend factors.
+        initial_step_size (float): Initial step size for adaptive step size calculation.
+        final_step_size (float): Final step size for adaptive step size calculation.
+        initial_noise_scale (float): Initial noise scale factor.
+        final_noise_scale (float): Final noise scale factor.
+        step_size_factor: Adjust to compensate for oversmoothing
+        noise_scale_factor: Adjust to provide more variation
+        
+    Returns:
+        torch.Tensor: A tensor of blended sigma values.
+    """
+    config_path = os.path.join(os.path.dirname(__file__), 'simple_kes_scheduler.yaml')
+    config = config_manager.load_config()
+    scheduler_config = config.get('scheduler', {})
+    if not scheduler_config:
+        warnings.warn("Scheduler configuration is missing from the config file. Using default values.")
+    
+    # Global randomization flag
+    global_randomize = scheduler_config.get('randomize', False)
+    
+    #debug_log("Entered simple_karras_exponential_scheduler function")    
+    default_config = {
+        "debug": False,
+        "device": "cuda" if torch.cuda.is_available() else "cpu",        
+        "sigma_min": 0.01,
+        "sigma_max": 50, #if sigma_max is too low the resulting picture may be undesirable. 
+        "start_blend": 0.1,
+        "end_blend": 0.5,
+        "sharpness": 0.95,
+        "early_stopping_threshold": 0.01,
+        "update_interval": 10,
+        "initial_step_size": 0.9,
+        "final_step_size": 0.2,
+        "initial_noise_scale": 1.25,
+        "final_noise_scale": 0.8,
+        "smooth_blend_factor": 11,
+        "step_size_factor": 0.8, #suggested value to avoid oversmoothing
+        "noise_scale_factor": 0.9, #suggested value to add more variation 
+        "randomize": False,
+        "sigma_min_rand": False,
+        "sigma_min_rand_min": 0.001,
+        "sigma_min_rand_max": 0.05,
+        "sigma_max_rand": False,
+        "sigma_max_rand_min": 0.05,
+        "sigma_max_rand_max": 0.20,
+        "start_blend_rand": False,
+        "start_blend_rand_min": 0.05,
+        "start_blend_rand_max": 0.2, 
+        "end_blend_rand": False,
+        "end_blend_rand_min": 0.4,
+        "end_blend_rand_max": 0.6,
+        "sharpness_rand": False,
+        "sharpness_rand_min": 0.85,
+        "sharpness_rand_max": 1.0,
+        "early_stopping_rand": False,
+        "early_stopping_rand_min": 0.001,
+        "early_stopping_rand_max": 0.02,
+        "update_interval_rand": False,
+        "update_interval_rand_min": 5,
+        "update_interval_rand_max": 10,
+        "initial_step_rand": False,
+        "initial_step_rand_min": 0.7,
+        "initial_step_rand_max": 1.0,
+        "final_step_rand": False,
+        "final_step_rand_min": 0.1,
+        "final_step_rand_max": 0.3,
+        "initial_noise_rand": False,
+        "initial_noise_rand_min": 1.0,
+        "initial_noise_rand_max": 1.5,
+        "final_noise_rand": False,
+        "final_noise_rand_min": 0.6,
+        "final_noise_rand_max": 1.0,
+        "smooth_blend_factor_rand": False,
+        "smooth_blend_factor_rand_min": 6,    
+        "smooth_blend_factor_rand_max": 11,   
+        "step_size_factor_rand": False,
+        "step_size_factor_rand_min": 0.65,
+        "step_size_factor_rand_max": 0.85,   
+        "noise_scale_factor_rand": False,
+        "noise_scale_factor_rand_min": 0.75,
+        "noise_scale_factor_rand_max": 0.95,    
+    }    
+    custom_logger.info(f"Default Config create {default_config}")
+    config = config_manager.load_config().get('scheduler', {})
+    if not config:
+        warnings.warn("Scheduler configuration is missing from the config file.")
+    
+    # Log loaded YAML configuration
+    custom_logger.info(f"Configuration loaded from YAML: {config}")
+    
+    for key, value in config.items():
+        if key in default_config:
+            default_config[key] = value  # Override default with YAML value
+            custom_logger.info(f"Overriding default config: {key} = {value}")
+        else:
+            custom_logger.info(f"Ignoring unknown config option: {key}")
+        
+    custom_logger.info(f"Final configuration after merging with YAML: {default_config}")
+
+    global_randomize = default_config.get('randomize', False)
+    custom_logger.info(f"Global randomization flag set to: {global_randomize}")
+
+    custom_logger.info(f"Config loaded from yaml {config}")
+   
+    # Now using default_config, updated with valid YAML values
+    custom_logger.info(f"Final Config after overriding: {default_config}")
+    
+    # Example: Reading the randomization flags from the config
+    randomize = config.get('scheduler', {}).get('randomize', False)
+    
+    # Use the get_random_or_default function for each parameter
+    #if randomize = false, then it checks for each variable for randomize, if true, then that particular option is randomized, with the others using default or config defined values.
+    sigma_min = get_random_or_default(config, 'sigma_min', sigma_min, global_randomize)
+    sigma_max = get_random_or_default(config, 'sigma_max', sigma_max, global_randomize)
+    start_blend = get_random_or_default(config, 'start_blend', start_blend, global_randomize)
+    end_blend = get_random_or_default(config, 'end_blend', end_blend, global_randomize)
+    sharpness = get_random_or_default(config, 'sharpness', sharpness, global_randomize)
+    early_stopping_threshold = get_random_or_default(config, 'early_stopping', early_stopping_threshold, global_randomize)
+    update_interval = get_random_or_default(config, 'update_interval', update_interval, global_randomize)
+    initial_step_size = get_random_or_default(config, 'initial_step', initial_step_size, global_randomize)
+    final_step_size = get_random_or_default(config, 'final_step', final_step_size, global_randomize)
+    initial_noise_scale = get_random_or_default(config, 'initial_noise', initial_noise_scale, global_randomize)
+    final_noise_scale = get_random_or_default(config, 'final_noise', final_noise_scale, global_randomize)
+    smooth_blend_factor = get_random_or_default(config, 'smooth_blend_factor', smooth_blend_factor, global_randomize)
+    step_size_factor = get_random_or_default(config, 'step_size_factor', step_size_factor, global_randomize)
+    noise_scale_factor = get_random_or_default(config, 'noise_scale_factor', noise_scale_factor, global_randomize)
+    
+       
+    # Expand sigma_max slightly to account for smoother transitions
+    sigma_max = sigma_max * 1.1
+    custom_logger.info(f"Using device: {device}")
+    # Generate sigma sequences using Karras and Exponential methods
+    sigmas_karras = get_sigmas_karras(n=n, sigma_min=sigma_min, sigma_max=sigma_max, device=device)
+    sigmas_exponential = get_sigmas_exponential(n=n, sigma_min=sigma_min, sigma_max=sigma_max, device=device)
+    config = config_manager.config_data.get('scheduler', {})            
+    # Match lengths of sigma sequences
+    target_length = min(len(sigmas_karras), len(sigmas_exponential))  
+    sigmas_karras = sigmas_karras[:target_length]
+    sigmas_exponential = sigmas_exponential[:target_length]
+              
+    custom_logger.info(f"Generated sigma sequences. Karras: {sigmas_karras}, Exponential: {sigmas_exponential}")
+    if sigmas_karras is None:
+        raise ValueError("Sigmas Karras:{sigmas_karras} Failed to generate or assign sigmas correctly.")
+    if sigmas_exponential is None:    
+        raise ValueError("Sigmas Exponential: {sigmas_exponential} Failed to generate or assign sigmas correctly.")
+        #sigmas_karras = torch.zeros(n).to(device)
+        #sigmas_exponential = torch.zeros(n).to(device)   
+    try:
+        pass
+    except Exception as e:
+        error_log(f"Error generating sigmas: {e}")
+    finally:
+        # Stop the observer when done
+        observer.stop()
+        observer.join()
+  
+    # Define progress and initialize blend factor
+    progress = torch.linspace(0, 1, len(sigmas_karras)).to(device)
+    custom_logger.info(f"Progress created {progress}")
+    custom_logger.info(f"Progress Using device: {device}")
+    
+    sigs = torch.zeros_like(sigmas_karras).to(device)
+    custom_logger.info(f"Sigs created {sigs}")
+    custom_logger.info(f"Sigs Using device: {device}")
+
+    # Iterate through each step, dynamically adjust blend factor, step size, and noise scaling
+    for i in range(len(sigmas_karras)):
+        # Adaptive step size and blend factor calculations
+        step_size = initial_step_size * (1 - progress[i]) + final_step_size * progress[i] * step_size_factor  # 0.8 default value Adjusted to avoid over-smoothing
+        custom_logger.info(f"Step_size created {step_size}"   )
+        dynamic_blend_factor = start_blend * (1 - progress[i]) + end_blend * progress[i]
+        custom_logger.info(f"Dynamic_blend_factor created {dynamic_blend_factor}"  )
+        noise_scale = initial_noise_scale * (1 - progress[i]) + final_noise_scale * progress[i] * noise_scale_factor  # 0.9 default value Adjusted to keep more variation
+        custom_logger.info(f"noise_scale created {noise_scale}"   )
+        
+        # Calculate smooth blending between the two sigma sequences
+        smooth_blend = torch.sigmoid((dynamic_blend_factor - 0.5) * smooth_blend_factor) # Increase scaling factor to smooth transitions more
+        custom_logger.info(f"smooth_blend created {smooth_blend}"   )
+        
+        # Compute blended sigma values
+        blended_sigma = sigmas_karras[i] * (1 - smooth_blend) + sigmas_exponential[i] * smooth_blend
+        custom_logger.info(f"blended_sigma created {blended_sigma}"   )
+        
+        # Apply step size and noise scaling
+        sigs[i] = blended_sigma * step_size * noise_scale
+
+    # Optional: Adaptive sharpening based on sigma values
+    sharpen_mask = torch.where(sigs < sigma_min * 1.5, sharpness, 1.0).to(device)
+    custom_logger.info(f"sharpen_mask created {sharpen_mask} with device {device}"   )
+    sigs = sigs * sharpen_mask
+    
+    # Implement early stop criteria based on sigma convergence
+    change = torch.abs(sigs[1:] - sigs[:-1])
+    if torch.all(change < early_stopping_threshold):
+        custom_logger.info("Early stopping criteria met."   )
+        return sigs[:len(change) + 1].to(device)
+    
+    if torch.isnan(sigs).any() or torch.isinf(sigs).any():
+        raise ValueError("Invalid sigma values detected (NaN or Inf).")
+
+    return sigs.to(device)

sec/simple_kes_scheduler.yaml

@@ -0,0 +1,146 @@
+scheduler:
+
+  #Optionally print to a log file for debugging. If false, debug is turned off, and no log file will be created.
+  #config options: true or false
+  debug: false
+  
+  # The minimum value for the noise level (sigma) during image generation.
+  # Decreasing this value makes the image clearer but less detailed.
+  # Increasing it makes the image noisier but potentially more artistic or abstract.
+  sigma_min: 0.01  # Default: 0.01, Suggested range: 0.01 - 0.1
+
+  # The maximum value for the noise level (sigma) during image generation.
+  # Increasing this value can create more variation in the image details.
+  # Lower values keep the image more stable and less noisy.
+  sigma_max: 50  # Default: 50, Suggested range:10 - 60
+
+  # The device used for running the scheduler. If you have a GPU, set this to "cuda".
+  # Otherwise, use "cpu", but note that it will be significantly slower.
+  #device: "cuda"  # Options: "cuda" (GPU) or "cpu" (processor)
+
+  # Initial blend factor between Karras and Exponential noise methods.
+  # A higher initial blend makes the image sharper at the start.
+  # A lower initial blend makes the image smoother early on.
+  start_blend: 0.1  # Default: 0.1, Suggested range: 0.05 - 0.2
+
+  # Final blend factor between Karras and Exponential noise methods.
+  # Higher values blend more noise at the end, possibly adding more detail.
+  # Lower values blend less noise for smoother, simpler images at the end.
+  end_blend: 0.5  # Default: 0.5, Suggested range: 0.4 - 0.6
+
+  # Sharpening factor applied to images during generation.
+  # Higher values increase sharpness but can add unwanted artifacts.
+  # Lower values reduce sharpness but may make the image look blurry.
+  sharpness: 0.95  # Default: 0.95, Suggested range: 0.8 - 1.0
+
+  # Early stopping threshold for stopping the image generation when changes between steps are minimal.
+  # Lower values stop early, saving time, but might produce incomplete images.
+  # Higher values take longer but may give more detailed results.
+  early_stopping_threshold: 0.01  # Default: 0.01, Suggested range: 0.005 - 0.02
+
+  # The number of steps between updates of the blend factor.
+  # Smaller values update the blend more frequently for smoother transitions.
+  # Larger values update the blend less frequently for faster processing.
+  update_interval: 10  # Default: 10, Suggested range: 5 - 15
+
+  # Initial step size, which controls how quickly the image evolves early on.
+  # Higher values make big changes at the start, possibly generating faster but less refined images.
+  # Lower values make smaller changes, giving more control over details.
+  initial_step_size: 0.9  # Default, 0.9, Suggested range: 0.5 - 1.0
+
+  # Final step size, which controls how much the image changes towards the end.
+  # Higher values keep details more flexible until the end, which may add complexity.
+  # Lower values lock the details earlier, making the image simpler.
+  final_step_size: 0.2  # Default: 0.2, Suggested range: 0.1 - 0.3
+
+  # Initial noise scaling applied to the image generation process.
+  # Higher values add more noise early on, making the initial image more random.
+  # Lower values reduce noise early on, leading to a smoother initial image.
+  initial_noise_scale: 1.25  # Default, 1.25, Suggested range: 1.0 - 1.5
+
+  # Final noise scaling applied at the end of the image generation.
+  # Higher values add noise towards the end, possibly adding fine detail.
+  # Lower values reduce noise towards the end, making the final image smoother.
+  final_noise_scale: 0.8  # Default, 0.8, Suggested range: 0.6 - 1.0
+  
+  
+  smooth_blend_factor: 11 #Default: 11, try 6 for more variation 
+  step_size_factor: 0.75 #suggested value (0.8) to avoid oversmoothing
+  noise_scale_factor: 0.95 #suggested value (0.9) to add more variation 
+  
+  
+  # Enables global randomization. 
+  # If true, all parameters are randomized within specified min/max ranges.
+  # If false, individual parameters with _rand flags set to true will still be randomized.
+  randomize: true 
+  
+  #Sigma values typically start very small. Lowering this could allow more gradual noise reduction. Too large would overwhelm the process.   
+  sigma_min_rand: false
+  sigma_min_rand_min: 0.001
+  sigma_min_rand_max: 0.05
+  
+  #Sigma max controls the upper limit of the noise. A lower minimum could allow faster convergence, while a higher max gives more flexibility for noisier images.
+  sigma_max_rand: false
+  sigma_max_rand_min: 10
+  sigma_max_rand_max: 60
+  
+  #Start blend controls how strongly Karras and Exponential are blended at the start. A slightly lower value introduces more variety in the blending at the beginning.
+  start_blend_rand: false
+  start_blend_rand_min: 0.05
+  start_blend_rand_max: 0.2
+  
+  # End blend affects how much the blending changes towards the end. Increasing the upper limit would allow more variation.
+  end_blend_rand: false
+  end_blend_rand_min: 0.4
+  end_blend_rand_max: 0.6
+  
+  # Sharpness controls detail retention. You wouldn’t want to lower it too much, as it might lose detail.
+  sharpness_rand: false
+  sharpness_rand_min: 0.85
+  sharpness_rand_max: 1.0
+  
+  #A smaller early stopping threshold could lead to earlier stopping if the changes between sigma steps become too small, while the upper value would prevent early stopping until larger changes occur.
+  early_stopping_rand: false
+  early_stopping_rand_min: 0.001
+  early_stopping_rand_max: 0.02
+  
+  #Update intervals affect how frequently blending factors are updated. More frequent updates allow more flexibility in blending.
+  update_interval_rand: false
+  update_interval_rand_min: 5 
+  update_interval_rand_max: 10
+  
+  # The initial step size defines how large the steps are at the start. A slightly smaller value introduces more gradual transitions.
+  initial_step_rand: false
+  initial_step_rand_min: 0.7
+  initial_step_rand_max: 1.0
+  
+  # The final step size defines how small the steps become towards the end. A slightly larger range gives more control over the final convergence.
+  final_step_rand: false
+  final_step_rand_min: 0.1
+  final_step_rand_max: 0.3
+  
+  #Initial noise scale defines how much noise to introduce initially. Larger values make the process start with more randomness, while smaller values keep it controlled.
+  initial_noise_rand: false
+  initial_noise_rand_min: 1.0
+  initial_noise_rand_max: 1.5
+  
+  # Final noise scale affects how much noise is reduced at the end. A lower minimum allows more noise to persist, while a higher maximum ensures full convergence.
+  final_noise_rand: false
+  final_noise_rand_min: 0.6  
+  final_noise_rand_max: 1.0
+  
+  #The smooth blend factor controls how aggressively the blending is smoothed. Lower values allow more abrupt blending changes, while higher values give smoother transitions.
+  smooth_blend_factor_rand: false
+  smooth_blend_factor_rand_min: 6    
+  smooth_blend_factor_rand_max: 11   
+  
+  #Step size factor adjusts the step size dynamically to avoid oversmoothing. A lower minimum increases variety, while a higher max provides smoother results.
+  step_size_factor_rand: false
+  step_size_factor_rand_min: 0.65   
+  step_size_factor_rand_max: 0.85   
+  
+  # Noise scale factor controls how noise is scaled throughout the steps. A slightly lower minimum adds more variety, while keeping the maximum value near the suggested ensures more uniform results.
+  noise_scale_factor_rand: false
+  noise_scale_factor_rand_min: 0.75  
+  noise_scale_factor_rand_max: 0.95 
+  

sec/txt2img.py

@@ -0,0 +1,136 @@
+import json
+from contextlib import closing
+
+import modules.scripts
+from modules import processing, infotext_utils
+from modules.infotext_utils import create_override_settings_dict, parse_generation_parameters
+from modules.shared import opts
+import modules.shared as shared
+from modules.ui import plaintext_to_html
+from PIL import Image
+import gradio as gr
+
+
+def txt2img_create_processing(id_task: str, request: gr.Request, prompt: str, negative_prompt: str, prompt_styles,
+                              n_iter: int, batch_size: int, cfg_scale: float, height: int, width: int, enable_hr: bool,
+                              denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int,
+                              hr_resize_x: int, hr_resize_y: int, hr_checkpoint_name: str, hr_sampler_name: str, hr_scheduler: str,
+                              hr_prompt: str, hr_negative_prompt, override_settings_texts, enable_progressive_growing: bool,
+                              progressive_growing_min_scale: float, progressive_growing_max_scale: float, progressive_growing_steps: int,
+                              progressive_growing_refinement: bool, *args, force_enable_hr=False):
+    override_settings = create_override_settings_dict(override_settings_texts)
+
+    if force_enable_hr:
+        enable_hr = True
+
+
+    print(f"enable_progressive_growing: {enable_progressive_growing}")
+    print(f"progressive_growing_min_scale: {progressive_growing_min_scale}")
+
+    p = processing.StableDiffusionProcessingTxt2Img(
+        sd_model=shared.sd_model,
+        outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,
+        outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids,
+        prompt=prompt,
+        styles=prompt_styles,
+        negative_prompt=negative_prompt,
+        batch_size=batch_size,
+        n_iter=n_iter,
+        cfg_scale=cfg_scale,
+        width=width,
+        height=height,
+        enable_hr=enable_hr,
+        denoising_strength=denoising_strength,
+        hr_scale=hr_scale,
+        hr_upscaler=hr_upscaler,
+        hr_second_pass_steps=hr_second_pass_steps,
+        hr_resize_x=hr_resize_x,
+        hr_resize_y=hr_resize_y,
+        hr_checkpoint_name=None if hr_checkpoint_name == 'Use same checkpoint' else hr_checkpoint_name,
+        hr_sampler_name=None if hr_sampler_name == 'Use same sampler' else hr_sampler_name,
+        hr_scheduler=None if hr_scheduler == 'Use same scheduler' else hr_scheduler,
+        hr_prompt=hr_prompt,
+        hr_negative_prompt=hr_negative_prompt,
+        override_settings=override_settings,
+    )
+
+    p.id_task = id_task
+    p.enable_progressive_growing = enable_progressive_growing
+    p.progressive_growing_min_scale = progressive_growing_min_scale
+    p.progressive_growing_max_scale = progressive_growing_max_scale
+    p.progressive_growing_steps = progressive_growing_steps
+    p.progressive_growing_refinement = progressive_growing_refinement
+    p.scripts = modules.scripts.scripts_txt2img
+    p.script_args = args
+
+    p.user = request.username
+
+    if shared.opts.enable_console_prompts:
+        print(f"\ntxt2img: {prompt}", file=shared.progress_print_out)
+
+    return p
+
+
+def txt2img_upscale(id_task: str, request: gr.Request, gallery, gallery_index, generation_info, *args):
+    assert len(gallery) > 0, 'No image to upscale'
+    assert 0 <= gallery_index < len(gallery), f'Bad image index: {gallery_index}'
+
+    p = txt2img_create_processing(id_task, request, *args, force_enable_hr=True)
+    p.batch_size = 1
+    p.n_iter = 1
+    # txt2img_upscale attribute that signifies this is called by txt2img_upscale
+    p.txt2img_upscale = True
+
+    geninfo = json.loads(generation_info)
+
+    image_info = gallery[gallery_index] if 0 <= gallery_index < len(gallery) else gallery[0]
+    p.firstpass_image = infotext_utils.image_from_url_text(image_info)
+
+    parameters = parse_generation_parameters(geninfo.get('infotexts')[gallery_index], [])
+    p.seed = parameters.get('Seed', -1)
+    p.subseed = parameters.get('Variation seed', -1)
+
+    p.override_settings['save_images_before_highres_fix'] = False
+
+    with closing(p):
+        processed = modules.scripts.scripts_txt2img.run(p, *p.script_args)
+
+        if processed is None:
+            processed = processing.process_images(p)
+
+    shared.total_tqdm.clear()
+
+    new_gallery = []
+    for i, image in enumerate(gallery):
+        if i == gallery_index:
+            geninfo["infotexts"][gallery_index: gallery_index+1] = processed.infotexts
+            new_gallery.extend(processed.images)
+        else:
+            fake_image = Image.new(mode="RGB", size=(1, 1))
+            fake_image.already_saved_as = image["name"].rsplit('?', 1)[0]
+            new_gallery.append(fake_image)
+
+    geninfo["infotexts"][gallery_index] = processed.info
+
+    return new_gallery, json.dumps(geninfo), plaintext_to_html(processed.info), plaintext_to_html(processed.comments, classname="comments")
+
+
+def txt2img(id_task: str, request: gr.Request, *args):
+    p = txt2img_create_processing(id_task, request, *args)
+
+    with closing(p):
+        processed = modules.scripts.scripts_txt2img.run(p, *p.script_args)
+
+        if processed is None:
+            processed = processing.process_images(p)
+
+    shared.total_tqdm.clear()
+
+    generation_info_js = processed.js()
+    if opts.samples_log_stdout:
+        print(generation_info_js)
+
+    if opts.do_not_show_images:
+        processed.images = []
+
+    return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments, classname="comments")

sec/ui.py

@@ -0,0 +1,1249 @@
+import datetime
+import mimetypes
+import os
+import sys
+from functools import reduce
+import warnings
+from contextlib import ExitStack
+
+import gradio as gr
+import gradio.utils
+import numpy as np
+from PIL import Image, PngImagePlugin  # noqa: F401
+from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call, wrap_gradio_call_no_job # noqa: F401
+
+from modules import gradio_extensons, sd_schedulers  # noqa: F401
+from modules import sd_hijack, sd_models, script_callbacks, ui_extensions, deepbooru, extra_networks, ui_common, ui_postprocessing, progress, ui_loadsave, shared_items, ui_settings, timer, sysinfo, ui_checkpoint_merger, scripts, sd_samplers, processing, ui_extra_networks, ui_toprow, launch_utils
+from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML, InputAccordion, ResizeHandleRow
+from modules.paths import script_path
+from modules.ui_common import create_refresh_button
+from modules.ui_gradio_extensions import reload_javascript
+
+from modules.shared import opts, cmd_opts
+
+import modules.infotext_utils as parameters_copypaste
+import modules.hypernetworks.ui as hypernetworks_ui
+import modules.textual_inversion.ui as textual_inversion_ui
+import modules.textual_inversion.textual_inversion as textual_inversion
+import modules.shared as shared
+from modules import prompt_parser
+from modules.sd_hijack import model_hijack
+from modules.infotext_utils import image_from_url_text, PasteField
+
+create_setting_component = ui_settings.create_setting_component
+
+warnings.filterwarnings("default" if opts.show_warnings else "ignore", category=UserWarning)
+warnings.filterwarnings("default" if opts.show_gradio_deprecation_warnings else "ignore", category=gr.deprecation.GradioDeprecationWarning)
+
+# this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the browser will not show any UI
+mimetypes.init()
+mimetypes.add_type('application/javascript', '.js')
+mimetypes.add_type('application/javascript', '.mjs')
+
+# Likewise, add explicit content-type header for certain missing image types
+mimetypes.add_type('image/webp', '.webp')
+mimetypes.add_type('image/avif', '.avif')
+
+if not cmd_opts.share and not cmd_opts.listen:
+    # fix gradio phoning home
+    gradio.utils.version_check = lambda: None
+    gradio.utils.get_local_ip_address = lambda: '127.0.0.1'
+
+if cmd_opts.ngrok is not None:
+    import modules.ngrok as ngrok
+    print('ngrok authtoken detected, trying to connect...')
+    ngrok.connect(
+        cmd_opts.ngrok,
+        cmd_opts.port if cmd_opts.port is not None else 7860,
+        cmd_opts.ngrok_options
+        )
+
+
+def gr_show(visible=True):
+    return {"visible": visible, "__type__": "update"}
+
+
+sample_img2img = "assets/stable-samples/img2img/sketch-mountains-input.jpg"
+sample_img2img = sample_img2img if os.path.exists(sample_img2img) else None
+
+# Using constants for these since the variation selector isn't visible.
+# Important that they exactly match script.js for tooltip to work.
+random_symbol = '\U0001f3b2\ufe0f'  # 🎲️
+reuse_symbol = '\u267b\ufe0f'  # ♻️
+paste_symbol = '\u2199\ufe0f'  # ↙
+refresh_symbol = '\U0001f504'  # 🔄
+save_style_symbol = '\U0001f4be'  # 💾
+apply_style_symbol = '\U0001f4cb'  # 📋
+clear_prompt_symbol = '\U0001f5d1\ufe0f'  # 🗑️
+extra_networks_symbol = '\U0001F3B4'  # 🎴
+switch_values_symbol = '\U000021C5' # ⇅
+restore_progress_symbol = '\U0001F300' # 🌀
+detect_image_size_symbol = '\U0001F4D0'  # 📐
+
+
+plaintext_to_html = ui_common.plaintext_to_html
+
+
+def send_gradio_gallery_to_image(x):
+    if len(x) == 0:
+        return None
+    return image_from_url_text(x[0])
+
+
+def calc_resolution_hires(enable, width, height, hr_scale, hr_resize_x, hr_resize_y):
+    if not enable:
+        return ""
+
+    p = processing.StableDiffusionProcessingTxt2Img(width=width, height=height, enable_hr=True, hr_scale=hr_scale, hr_resize_x=hr_resize_x, hr_resize_y=hr_resize_y)
+    p.calculate_target_resolution()
+
+    return f"from <span class='resolution'>{p.width}x{p.height}</span> to <span class='resolution'>{p.hr_resize_x or p.hr_upscale_to_x}x{p.hr_resize_y or p.hr_upscale_to_y}</span>"
+
+
+def resize_from_to_html(width, height, scale_by):
+    target_width = int(width * scale_by)
+    target_height = int(height * scale_by)
+
+    if not target_width or not target_height:
+        return "no image selected"
+
+    return f"resize: from <span class='resolution'>{width}x{height}</span> to <span class='resolution'>{target_width}x{target_height}</span>"
+
+
+def process_interrogate(interrogation_function, mode, ii_input_dir, ii_output_dir, *ii_singles):
+    if mode in {0, 1, 3, 4}:
+        return [interrogation_function(ii_singles[mode]), None]
+    elif mode == 2:
+        return [interrogation_function(ii_singles[mode]["image"]), None]
+    elif mode == 5:
+        assert not shared.cmd_opts.hide_ui_dir_config, "Launched with --hide-ui-dir-config, batch img2img disabled"
+        images = shared.listfiles(ii_input_dir)
+        print(f"Will process {len(images)} images.")
+        if ii_output_dir != "":
+            os.makedirs(ii_output_dir, exist_ok=True)
+        else:
+            ii_output_dir = ii_input_dir
+
+        for image in images:
+            img = Image.open(image)
+            filename = os.path.basename(image)
+            left, _ = os.path.splitext(filename)
+            print(interrogation_function(img), file=open(os.path.join(ii_output_dir, f"{left}.txt"), 'a', encoding='utf-8'))
+
+        return [gr.update(), None]
+
+
+def interrogate(image):
+    prompt = shared.interrogator.interrogate(image.convert("RGB"))
+    return gr.update() if prompt is None else prompt
+
+
+def interrogate_deepbooru(image):
+    prompt = deepbooru.model.tag(image)
+    return gr.update() if prompt is None else prompt
+
+
+def connect_clear_prompt(button):
+    """Given clear button, prompt, and token_counter objects, setup clear prompt button click event"""
+    button.click(
+        _js="clear_prompt",
+        fn=None,
+        inputs=[],
+        outputs=[],
+    )
+
+
+def update_token_counter(text, steps, styles, *, is_positive=True):
+    params = script_callbacks.BeforeTokenCounterParams(text, steps, styles, is_positive=is_positive)
+    script_callbacks.before_token_counter_callback(params)
+    text = params.prompt
+    steps = params.steps
+    styles = params.styles
+    is_positive = params.is_positive
+
+    if shared.opts.include_styles_into_token_counters:
+        apply_styles = shared.prompt_styles.apply_styles_to_prompt if is_positive else shared.prompt_styles.apply_negative_styles_to_prompt
+        text = apply_styles(text, styles)
+
+    try:
+        text, _ = extra_networks.parse_prompt(text)
+
+        if is_positive:
+            _, prompt_flat_list, _ = prompt_parser.get_multicond_prompt_list([text])
+        else:
+            prompt_flat_list = [text]
+
+        prompt_schedules = prompt_parser.get_learned_conditioning_prompt_schedules(prompt_flat_list, steps)
+
+    except Exception:
+        # a parsing error can happen here during typing, and we don't want to bother the user with
+        # messages related to it in console
+        prompt_schedules = [[[steps, text]]]
+
+    flat_prompts = reduce(lambda list1, list2: list1+list2, prompt_schedules)
+    prompts = [prompt_text for step, prompt_text in flat_prompts]
+    token_count, max_length = max([model_hijack.get_prompt_lengths(prompt) for prompt in prompts], key=lambda args: args[0])
+    return f"<span class='gr-box gr-text-input'>{token_count}/{max_length}</span>"
+
+
+def update_negative_prompt_token_counter(*args):
+    return update_token_counter(*args, is_positive=False)
+
+
+def setup_progressbar(*args, **kwargs):
+    pass
+
+
+def apply_setting(key, value):
+    if value is None:
+        return gr.update()
+
+    if shared.cmd_opts.freeze_settings:
+        return gr.update()
+
+    # dont allow model to be swapped when model hash exists in prompt
+    if key == "sd_model_checkpoint" and opts.disable_weights_auto_swap:
+        return gr.update()
+
+    if key == "sd_model_checkpoint":
+        ckpt_info = sd_models.get_closet_checkpoint_match(value)
+
+        if ckpt_info is not None:
+            value = ckpt_info.title
+        else:
+            return gr.update()
+
+    comp_args = opts.data_labels[key].component_args
+    if comp_args and isinstance(comp_args, dict) and comp_args.get('visible') is False:
+        return
+
+    valtype = type(opts.data_labels[key].default)
+    oldval = opts.data.get(key, None)
+    opts.data[key] = valtype(value) if valtype != type(None) else value
+    if oldval != value and opts.data_labels[key].onchange is not None:
+        opts.data_labels[key].onchange()
+
+    opts.save(shared.config_filename)
+    return getattr(opts, key)
+
+
+def create_output_panel(tabname, outdir, toprow=None):
+    return ui_common.create_output_panel(tabname, outdir, toprow)
+
+
+def ordered_ui_categories():
+    user_order = {x.strip(): i * 2 + 1 for i, x in enumerate(shared.opts.ui_reorder_list)}
+
+    for _, category in sorted(enumerate(shared_items.ui_reorder_categories()), key=lambda x: user_order.get(x[1], x[0] * 2 + 0)):
+        yield category
+
+
+def create_override_settings_dropdown(tabname, row):
+    dropdown = gr.Dropdown([], label="Override settings", visible=False, elem_id=f"{tabname}_override_settings", multiselect=True)
+
+    dropdown.change(
+        fn=lambda x: gr.Dropdown.update(visible=bool(x)),
+        inputs=[dropdown],
+        outputs=[dropdown],
+    )
+
+    return dropdown
+
+
+def create_ui():
+    import modules.img2img
+    import modules.txt2img
+
+    reload_javascript()
+
+    parameters_copypaste.reset()
+
+    settings = ui_settings.UiSettings()
+    settings.register_settings()
+
+    scripts.scripts_current = scripts.scripts_txt2img
+    scripts.scripts_txt2img.initialize_scripts(is_img2img=False)
+
+    with gr.Blocks(analytics_enabled=False) as txt2img_interface:
+        toprow = ui_toprow.Toprow(is_img2img=False, is_compact=shared.opts.compact_prompt_box)
+
+        dummy_component = gr.Label(visible=False)
+
+        extra_tabs = gr.Tabs(elem_id="txt2img_extra_tabs", elem_classes=["extra-networks"])
+        extra_tabs.__enter__()
+
+        with gr.Tab("Generation", id="txt2img_generation") as txt2img_generation_tab, ResizeHandleRow(equal_height=False):
+            with ExitStack() as stack:
+                if shared.opts.txt2img_settings_accordion:
+                    stack.enter_context(gr.Accordion("Open for Settings", open=False))
+                stack.enter_context(gr.Column(variant='compact', elem_id="txt2img_settings"))
+
+                scripts.scripts_txt2img.prepare_ui()
+
+                for category in ordered_ui_categories():
+                    if category == "prompt":
+                        toprow.create_inline_toprow_prompts()
+
+                    elif category == "dimensions":
+                        with FormRow():
+                            with gr.Column(elem_id="txt2img_column_size", scale=4):
+                                width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="txt2img_width")
+                                height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="txt2img_height")
+
+                            with gr.Column(elem_id="txt2img_dimensions_row", scale=1, elem_classes="dimensions-tools"):
+                                res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="txt2img_res_switch_btn", tooltip="Switch width/height")
+
+                            if opts.dimensions_and_batch_together:
+                                with gr.Column(elem_id="txt2img_column_batch"):
+                                    batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="txt2img_batch_count")
+                                    batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="txt2img_batch_size")
+
+                    elif category == "cfg":
+                        with gr.Row():
+                            cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0, elem_id="txt2img_cfg_scale")
+
+                    elif category == "checkboxes":
+                        with FormRow(elem_classes="checkboxes-row", variant="compact"):
+                            pass
+
+                    elif category == "accordions":
+                        with gr.Row(elem_id="txt2img_accordions", elem_classes="accordions"):
+                            with InputAccordion(False, label="Hires. fix", elem_id="txt2img_hr") as enable_hr:
+                                with enable_hr.extra():
+                                    hr_final_resolution = FormHTML(value="", elem_id="txtimg_hr_finalres", label="Upscaled resolution", interactive=False, min_width=0)
+
+                                with FormRow(elem_id="txt2img_hires_fix_row1", variant="compact"):
+                                    hr_upscaler = gr.Dropdown(label="Upscaler", elem_id="txt2img_hr_upscaler", choices=[*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]], value=shared.latent_upscale_default_mode)
+                                    hr_second_pass_steps = gr.Slider(minimum=0, maximum=150, step=1, label='Hires steps', value=0, elem_id="txt2img_hires_steps")
+                                    denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.7, elem_id="txt2img_denoising_strength")
+
+                                with FormRow(elem_id="txt2img_hires_fix_row2", variant="compact"):
+                                    hr_scale = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Upscale by", value=2.0, elem_id="txt2img_hr_scale")
+                                    hr_resize_x = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize width to", value=0, elem_id="txt2img_hr_resize_x")
+                                    hr_resize_y = gr.Slider(minimum=0, maximum=2048, step=8, label="Resize height to", value=0, elem_id="txt2img_hr_resize_y")
+
+                                with FormRow(elem_id="txt2img_hires_fix_row3", variant="compact", visible=opts.hires_fix_show_sampler) as hr_sampler_container:
+
+                                    hr_checkpoint_name = gr.Dropdown(label='Checkpoint', elem_id="hr_checkpoint", choices=["Use same checkpoint"] + modules.sd_models.checkpoint_tiles(use_short=True), value="Use same checkpoint")
+                                    create_refresh_button(hr_checkpoint_name, modules.sd_models.list_models, lambda: {"choices": ["Use same checkpoint"] + modules.sd_models.checkpoint_tiles(use_short=True)}, "hr_checkpoint_refresh")
+
+                                    hr_sampler_name = gr.Dropdown(label='Hires sampling method', elem_id="hr_sampler", choices=["Use same sampler"] + sd_samplers.visible_sampler_names(), value="Use same sampler")
+                                    hr_scheduler = gr.Dropdown(label='Hires schedule type', elem_id="hr_scheduler", choices=["Use same scheduler"] + [x.label for x in sd_schedulers.schedulers], value="Use same scheduler")
+
+                                with FormRow(elem_id="txt2img_hires_fix_row4", variant="compact", visible=opts.hires_fix_show_prompts) as hr_prompts_container:
+                                    with gr.Column(scale=80):
+                                        with gr.Row():
+                                            hr_prompt = gr.Textbox(label="Hires prompt", elem_id="hires_prompt", show_label=False, lines=3, placeholder="Prompt for hires fix pass.\nLeave empty to use the same prompt as in first pass.", elem_classes=["prompt"])
+                                    with gr.Column(scale=80):
+                                        with gr.Row():
+                                            hr_negative_prompt = gr.Textbox(label="Hires negative prompt", elem_id="hires_neg_prompt", show_label=False, lines=3, placeholder="Negative prompt for hires fix pass.\nLeave empty to use the same negative prompt as in first pass.", elem_classes=["prompt"])
+
+                            with InputAccordion(False, label="Progressive Growing", elem_id="txt2img_progressive_growing") as enable_progressive_growing:
+                                with FormRow(elem_id="txt2img_progressive_growing_row1", variant="compact"):
+                                    progressive_growing_min_scale = gr.Slider(minimum=0.1, maximum=1.0, step=0.05, label="Min Scale", value=0.25, elem_id="txt2img_progressive_growing_min_scale")
+                                    progressive_growing_max_scale = gr.Slider(minimum=0.1, maximum=1.0, step=0.05, label="Max Scale", value=1.0, elem_id="txt2img_progressive_growing_max_scale")
+
+                                with FormRow(elem_id="txt2img_progressive_growing_row2", variant="compact"):
+                                    progressive_growing_steps = gr.Slider(minimum=2, maximum=10, step=1, label="Steps", value=4, elem_id="txt2img_progressive_growing_steps")
+                                    progressive_growing_refinement = gr.Checkbox(label="Enable Refinement", value=True, elem_id="txt2img_progressive_growing_refinement")
+
+                            scripts.scripts_txt2img.setup_ui_for_section(category)
+
+                    elif category == "batch":
+                        if not opts.dimensions_and_batch_together:
+                            with FormRow(elem_id="txt2img_column_batch"):
+                                batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="txt2img_batch_count")
+                                batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="txt2img_batch_size")
+
+                    elif category == "override_settings":
+                        with FormRow(elem_id="txt2img_override_settings_row") as row:
+                            override_settings = create_override_settings_dropdown('txt2img', row)
+
+                    elif category == "scripts":
+                        with FormGroup(elem_id="txt2img_script_container"):
+                            custom_inputs = scripts.scripts_txt2img.setup_ui()
+
+                    if category not in {"accordions"}:
+                        scripts.scripts_txt2img.setup_ui_for_section(category)
+
+            hr_resolution_preview_inputs = [enable_hr, width, height, hr_scale, hr_resize_x, hr_resize_y]
+
+            for component in hr_resolution_preview_inputs:
+                event = component.release if isinstance(component, gr.Slider) else component.change
+
+                event(
+                    fn=calc_resolution_hires,
+                    inputs=hr_resolution_preview_inputs,
+                    outputs=[hr_final_resolution],
+                    show_progress=False,
+                )
+                event(
+                    None,
+                    _js="onCalcResolutionHires",
+                    inputs=hr_resolution_preview_inputs,
+                    outputs=[],
+                    show_progress=False,
+                )
+
+            output_panel = create_output_panel("txt2img", opts.outdir_txt2img_samples, toprow)
+
+            txt2img_inputs = [
+                dummy_component,
+                toprow.prompt,
+                toprow.negative_prompt,
+                toprow.ui_styles.dropdown,
+                batch_count,
+                batch_size,
+                cfg_scale,
+                height,
+                width,
+                enable_hr,
+                denoising_strength,
+                hr_scale,
+                hr_upscaler,
+                hr_second_pass_steps,
+                hr_resize_x,
+                hr_resize_y,
+                hr_checkpoint_name,
+                hr_sampler_name,
+                hr_scheduler,
+                hr_prompt,
+                hr_negative_prompt,
+                override_settings,
+                enable_progressive_growing,
+                progressive_growing_min_scale,
+                progressive_growing_max_scale,
+                progressive_growing_steps,
+                progressive_growing_refinement,
+            ] + custom_inputs
+
+            txt2img_outputs = [
+                output_panel.gallery,
+                output_panel.generation_info,
+                output_panel.infotext,
+                output_panel.html_log,
+            ]
+
+            txt2img_args = dict(
+                fn=wrap_gradio_gpu_call(modules.txt2img.txt2img, extra_outputs=[None, '', '']),
+                _js="submit",
+                inputs=txt2img_inputs,
+                outputs=txt2img_outputs,
+                show_progress=False,
+            )
+
+            toprow.prompt.submit(**txt2img_args)
+            toprow.submit.click(**txt2img_args)
+
+            output_panel.button_upscale.click(
+                fn=wrap_gradio_gpu_call(modules.txt2img.txt2img_upscale, extra_outputs=[None, '', '']),
+                _js="submit_txt2img_upscale",
+                inputs=txt2img_inputs[0:1] + [output_panel.gallery, dummy_component, output_panel.generation_info] + txt2img_inputs[1:],
+                outputs=txt2img_outputs,
+                show_progress=False,
+            )
+
+            res_switch_btn.click(fn=None, _js="function(){switchWidthHeight('txt2img')}", inputs=None, outputs=None, show_progress=False)
+
+            toprow.restore_progress_button.click(
+                fn=progress.restore_progress,
+                _js="restoreProgressTxt2img",
+                inputs=[dummy_component],
+                outputs=[
+                    output_panel.gallery,
+                    output_panel.generation_info,
+                    output_panel.infotext,
+                    output_panel.html_log,
+                ],
+                show_progress=False,
+            )
+
+            txt2img_paste_fields = [
+                PasteField(toprow.prompt, "Prompt", api="prompt"),
+                PasteField(toprow.negative_prompt, "Negative prompt", api="negative_prompt"),
+                PasteField(cfg_scale, "CFG scale", api="cfg_scale"),
+                PasteField(width, "Size-1", api="width"),
+                PasteField(height, "Size-2", api="height"),
+                PasteField(batch_size, "Batch size", api="batch_size"),
+                PasteField(toprow.ui_styles.dropdown, lambda d: d["Styles array"] if isinstance(d.get("Styles array"), list) else gr.update(), api="styles"),
+                PasteField(denoising_strength, "Denoising strength", api="denoising_strength"),
+                PasteField(enable_hr, lambda d: "Denoising strength" in d and ("Hires upscale" in d or "Hires upscaler" in d or "Hires resize-1" in d), api="enable_hr"),
+                PasteField(hr_scale, "Hires upscale", api="hr_scale"),
+                PasteField(hr_upscaler, "Hires upscaler", api="hr_upscaler"),
+                PasteField(hr_second_pass_steps, "Hires steps", api="hr_second_pass_steps"),
+                PasteField(hr_resize_x, "Hires resize-1", api="hr_resize_x"),
+                PasteField(hr_resize_y, "Hires resize-2", api="hr_resize_y"),
+                PasteField(hr_checkpoint_name, "Hires checkpoint", api="hr_checkpoint_name"),
+                PasteField(hr_sampler_name, sd_samplers.get_hr_sampler_from_infotext, api="hr_sampler_name"),
+                PasteField(hr_scheduler, sd_samplers.get_hr_scheduler_from_infotext, api="hr_scheduler"),
+                PasteField(hr_sampler_container, lambda d: gr.update(visible=True) if d.get("Hires sampler", "Use same sampler") != "Use same sampler" or d.get("Hires checkpoint", "Use same checkpoint") != "Use same checkpoint" or d.get("Hires schedule type", "Use same scheduler") != "Use same scheduler" else gr.update()),
+                PasteField(hr_prompt, "Hires prompt", api="hr_prompt"),
+                PasteField(hr_negative_prompt, "Hires negative prompt", api="hr_negative_prompt"),
+                PasteField(hr_prompts_container, lambda d: gr.update(visible=True) if d.get("Hires prompt", "") != "" or d.get("Hires negative prompt", "") != "" else gr.update()),
+                *scripts.scripts_txt2img.infotext_fields
+            ]
+            parameters_copypaste.add_paste_fields("txt2img", None, txt2img_paste_fields, override_settings)
+            parameters_copypaste.register_paste_params_button(parameters_copypaste.ParamBinding(
+                paste_button=toprow.paste, tabname="txt2img", source_text_component=toprow.prompt, source_image_component=None,
+            ))
+
+            steps = scripts.scripts_txt2img.script('Sampler').steps
+
+            txt2img_preview_params = [
+                toprow.prompt,
+                toprow.negative_prompt,
+                steps,
+                scripts.scripts_txt2img.script('Sampler').sampler_name,
+                cfg_scale,
+                scripts.scripts_txt2img.script('Seed').seed,
+                width,
+                height,
+            ]
+
+            toprow.ui_styles.dropdown.change(fn=wrap_queued_call(update_token_counter), inputs=[toprow.prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.token_counter])
+            toprow.ui_styles.dropdown.change(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.negative_token_counter])
+            toprow.token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[toprow.prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.token_counter])
+            toprow.negative_token_button.click(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.negative_token_counter])
+
+        extra_networks_ui = ui_extra_networks.create_ui(txt2img_interface, [txt2img_generation_tab], 'txt2img')
+        ui_extra_networks.setup_ui(extra_networks_ui, output_panel.gallery)
+
+        extra_tabs.__exit__()
+
+    scripts.scripts_current = scripts.scripts_img2img
+    scripts.scripts_img2img.initialize_scripts(is_img2img=True)
+
+    with gr.Blocks(analytics_enabled=False) as img2img_interface:
+        toprow = ui_toprow.Toprow(is_img2img=True, is_compact=shared.opts.compact_prompt_box)
+
+        extra_tabs = gr.Tabs(elem_id="img2img_extra_tabs", elem_classes=["extra-networks"])
+        extra_tabs.__enter__()
+
+        with gr.Tab("Generation", id="img2img_generation") as img2img_generation_tab, ResizeHandleRow(equal_height=False):
+            with ExitStack() as stack:
+                if shared.opts.img2img_settings_accordion:
+                    stack.enter_context(gr.Accordion("Open for Settings", open=False))
+                stack.enter_context(gr.Column(variant='compact', elem_id="img2img_settings"))
+
+                copy_image_buttons = []
+                copy_image_destinations = {}
+
+                def add_copy_image_controls(tab_name, elem):
+                    with gr.Row(variant="compact", elem_id=f"img2img_copy_to_{tab_name}"):
+                        gr.HTML("Copy image to: ", elem_id=f"img2img_label_copy_to_{tab_name}")
+
+                        for title, name in zip(['img2img', 'sketch', 'inpaint', 'inpaint sketch'], ['img2img', 'sketch', 'inpaint', 'inpaint_sketch']):
+                            if name == tab_name:
+                                gr.Button(title, interactive=False)
+                                copy_image_destinations[name] = elem
+                                continue
+
+                            button = gr.Button(title)
+                            copy_image_buttons.append((button, name, elem))
+
+                scripts.scripts_img2img.prepare_ui()
+
+                for category in ordered_ui_categories():
+                    if category == "prompt":
+                        toprow.create_inline_toprow_prompts()
+
+                    if category == "image":
+                        with gr.Tabs(elem_id="mode_img2img"):
+                            img2img_selected_tab = gr.Number(value=0, visible=False)
+
+                            with gr.TabItem('img2img', id='img2img', elem_id="img2img_img2img_tab") as tab_img2img:
+                                init_img = gr.Image(label="Image for img2img", elem_id="img2img_image", show_label=False, source="upload", interactive=True, type="pil", tool="editor", image_mode="RGBA", height=opts.img2img_editor_height)
+                                add_copy_image_controls('img2img', init_img)
+
+                            with gr.TabItem('Sketch', id='img2img_sketch', elem_id="img2img_img2img_sketch_tab") as tab_sketch:
+                                sketch = gr.Image(label="Image for img2img", elem_id="img2img_sketch", show_label=False, source="upload", interactive=True, type="pil", tool="color-sketch", image_mode="RGB", height=opts.img2img_editor_height, brush_color=opts.img2img_sketch_default_brush_color)
+                                add_copy_image_controls('sketch', sketch)
+
+                            with gr.TabItem('Inpaint', id='inpaint', elem_id="img2img_inpaint_tab") as tab_inpaint:
+                                init_img_with_mask = gr.Image(label="Image for inpainting with mask", show_label=False, elem_id="img2maskimg", source="upload", interactive=True, type="pil", tool="sketch", image_mode="RGBA", height=opts.img2img_editor_height, brush_color=opts.img2img_inpaint_mask_brush_color)
+                                add_copy_image_controls('inpaint', init_img_with_mask)
+
+                            with gr.TabItem('Inpaint sketch', id='inpaint_sketch', elem_id="img2img_inpaint_sketch_tab") as tab_inpaint_color:
+                                inpaint_color_sketch = gr.Image(label="Color sketch inpainting", show_label=False, elem_id="inpaint_sketch", source="upload", interactive=True, type="pil", tool="color-sketch", image_mode="RGB", height=opts.img2img_editor_height, brush_color=opts.img2img_inpaint_sketch_default_brush_color)
+                                inpaint_color_sketch_orig = gr.State(None)
+                                add_copy_image_controls('inpaint_sketch', inpaint_color_sketch)
+
+                                def update_orig(image, state):
+                                    if image is not None:
+                                        same_size = state is not None and state.size == image.size
+                                        has_exact_match = np.any(np.all(np.array(image) == np.array(state), axis=-1))
+                                        edited = same_size and has_exact_match
+                                        return image if not edited or state is None else state
+
+                                inpaint_color_sketch.change(update_orig, [inpaint_color_sketch, inpaint_color_sketch_orig], inpaint_color_sketch_orig)
+
+                            with gr.TabItem('Inpaint upload', id='inpaint_upload', elem_id="img2img_inpaint_upload_tab") as tab_inpaint_upload:
+                                init_img_inpaint = gr.Image(label="Image for img2img", show_label=False, source="upload", interactive=True, type="pil", elem_id="img_inpaint_base")
+                                init_mask_inpaint = gr.Image(label="Mask", source="upload", interactive=True, type="pil", image_mode="RGBA", elem_id="img_inpaint_mask")
+
+                            with gr.TabItem('Batch', id='batch', elem_id="img2img_batch_tab") as tab_batch:
+                                with gr.Tabs(elem_id="img2img_batch_source"):
+                                    img2img_batch_source_type = gr.Textbox(visible=False, value="upload")
+                                    with gr.TabItem('Upload', id='batch_upload', elem_id="img2img_batch_upload_tab") as tab_batch_upload:
+                                        img2img_batch_upload = gr.Files(label="Files", interactive=True, elem_id="img2img_batch_upload")
+                                    with gr.TabItem('From directory', id='batch_from_dir', elem_id="img2img_batch_from_dir_tab") as tab_batch_from_dir:
+                                        hidden = '<br>Disabled when launched with --hide-ui-dir-config.' if shared.cmd_opts.hide_ui_dir_config else ''
+                                        gr.HTML(
+                                            "<p style='padding-bottom: 1em;' class=\"text-gray-500\">Process images in a directory on the same machine where the server is running." +
+                                            "<br>Use an empty output directory to save pictures normally instead of writing to the output directory." +
+                                            f"<br>Add inpaint batch mask directory to enable inpaint batch processing."
+                                            f"{hidden}</p>"
+                                        )
+                                        img2img_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs, elem_id="img2img_batch_input_dir")
+                                        img2img_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs, elem_id="img2img_batch_output_dir")
+                                        img2img_batch_inpaint_mask_dir = gr.Textbox(label="Inpaint batch mask directory (required for inpaint batch processing only)", **shared.hide_dirs, elem_id="img2img_batch_inpaint_mask_dir")
+                                tab_batch_upload.select(fn=lambda: "upload", inputs=[], outputs=[img2img_batch_source_type])
+                                tab_batch_from_dir.select(fn=lambda: "from dir", inputs=[], outputs=[img2img_batch_source_type])
+                                with gr.Accordion("PNG info", open=False):
+                                    img2img_batch_use_png_info = gr.Checkbox(label="Append png info to prompts", elem_id="img2img_batch_use_png_info")
+                                    img2img_batch_png_info_dir = gr.Textbox(label="PNG info directory", **shared.hide_dirs, placeholder="Leave empty to use input directory", elem_id="img2img_batch_png_info_dir")
+                                    img2img_batch_png_info_props = gr.CheckboxGroup(["Prompt", "Negative prompt", "Seed", "CFG scale", "Sampler", "Steps", "Model hash"], label="Parameters to take from png info", info="Prompts from png info will be appended to prompts set in ui.")
+
+                            img2img_tabs = [tab_img2img, tab_sketch, tab_inpaint, tab_inpaint_color, tab_inpaint_upload, tab_batch]
+
+                            for i, tab in enumerate(img2img_tabs):
+                                tab.select(fn=lambda tabnum=i: tabnum, inputs=[], outputs=[img2img_selected_tab])
+
+                        def copy_image(img):
+                            if isinstance(img, dict) and 'image' in img:
+                                return img['image']
+
+                            return img
+
+                        for button, name, elem in copy_image_buttons:
+                            button.click(
+                                fn=copy_image,
+                                inputs=[elem],
+                                outputs=[copy_image_destinations[name]],
+                            )
+                            button.click(
+                                fn=lambda: None,
+                                _js=f"switch_to_{name.replace(' ', '_')}",
+                                inputs=[],
+                                outputs=[],
+                            )
+
+                        with FormRow():
+                            resize_mode = gr.Radio(label="Resize mode", elem_id="resize_mode", choices=["Just resize", "Crop and resize", "Resize and fill", "Just resize (latent upscale)"], type="index", value="Just resize")
+
+                    elif category == "dimensions":
+                        with FormRow():
+                            with gr.Column(elem_id="img2img_column_size", scale=4):
+                                selected_scale_tab = gr.Number(value=0, visible=False)
+
+                                with gr.Tabs(elem_id="img2img_tabs_resize"):
+                                    with gr.Tab(label="Resize to", id="to", elem_id="img2img_tab_resize_to") as tab_scale_to:
+                                        with FormRow():
+                                            with gr.Column(elem_id="img2img_column_size", scale=4):
+                                                width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="img2img_width")
+                                                height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="img2img_height")
+                                            with gr.Column(elem_id="img2img_dimensions_row", scale=1, elem_classes="dimensions-tools"):
+                                                res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="img2img_res_switch_btn", tooltip="Switch width/height")
+                                                detect_image_size_btn = ToolButton(value=detect_image_size_symbol, elem_id="img2img_detect_image_size_btn", tooltip="Auto detect size from img2img")
+
+                                    with gr.Tab(label="Resize by", id="by", elem_id="img2img_tab_resize_by") as tab_scale_by:
+                                        scale_by = gr.Slider(minimum=0.05, maximum=4.0, step=0.05, label="Scale", value=1.0, elem_id="img2img_scale")
+
+                                        with FormRow():
+                                            scale_by_html = FormHTML(resize_from_to_html(0, 0, 0.0), elem_id="img2img_scale_resolution_preview")
+                                            gr.Slider(label="Unused", elem_id="img2img_unused_scale_by_slider")
+                                            button_update_resize_to = gr.Button(visible=False, elem_id="img2img_update_resize_to")
+
+                                    on_change_args = dict(
+                                        fn=resize_from_to_html,
+                                        _js="currentImg2imgSourceResolution",
+                                        inputs=[dummy_component, dummy_component, scale_by],
+                                        outputs=scale_by_html,
+                                        show_progress=False,
+                                    )
+
+                                    scale_by.release(**on_change_args)
+                                    button_update_resize_to.click(**on_change_args)
+
+                            tab_scale_to.select(fn=lambda: 0, inputs=[], outputs=[selected_scale_tab])
+                            tab_scale_by.select(fn=lambda: 1, inputs=[], outputs=[selected_scale_tab])
+
+                            if opts.dimensions_and_batch_together:
+                                with gr.Column(elem_id="img2img_column_batch"):
+                                    batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="img2img_batch_count")
+                                    batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="img2img_batch_size")
+
+                    elif category == "denoising":
+                        denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.75, elem_id="img2img_denoising_strength")
+
+                    elif category == "cfg":
+                        with gr.Row():
+                            cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0, elem_id="img2img_cfg_scale")
+                            image_cfg_scale = gr.Slider(minimum=0, maximum=3.0, step=0.05, label='Image CFG Scale', value=1.5, elem_id="img2img_image_cfg_scale", visible=False)
+
+                    elif category == "checkboxes":
+                        with FormRow(elem_classes="checkboxes-row", variant="compact"):
+                            pass
+
+                    elif category == "accordions":
+                        with gr.Row(elem_id="img2img_accordions", elem_classes="accordions"):
+                            scripts.scripts_img2img.setup_ui_for_section(category)
+
+                    elif category == "batch":
+                        if not opts.dimensions_and_batch_together:
+                            with FormRow(elem_id="img2img_column_batch"):
+                                batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="img2img_batch_count")
+                                batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="img2img_batch_size")
+
+                    elif category == "override_settings":
+                        with FormRow(elem_id="img2img_override_settings_row") as row:
+                            override_settings = create_override_settings_dropdown('img2img', row)
+
+                    elif category == "scripts":
+                        with FormGroup(elem_id="img2img_script_container"):
+                            custom_inputs = scripts.scripts_img2img.setup_ui()
+
+                    elif category == "inpaint":
+                        with FormGroup(elem_id="inpaint_controls", visible=False) as inpaint_controls:
+                            with FormRow():
+                                mask_blur = gr.Slider(label='Mask blur', minimum=0, maximum=64, step=1, value=4, elem_id="img2img_mask_blur")
+                                mask_alpha = gr.Slider(label="Mask transparency", visible=False, elem_id="img2img_mask_alpha")
+
+                            with FormRow():
+                                inpainting_mask_invert = gr.Radio(label='Mask mode', choices=['Inpaint masked', 'Inpaint not masked'], value='Inpaint masked', type="index", elem_id="img2img_mask_mode")
+
+                            with FormRow():
+                                inpainting_fill = gr.Radio(label='Masked content', choices=['fill', 'original', 'latent noise', 'latent nothing'], value='original', type="index", elem_id="img2img_inpainting_fill")
+
+                            with FormRow():
+                                with gr.Column():
+                                    inpaint_full_res = gr.Radio(label="Inpaint area", choices=["Whole picture", "Only masked"], type="index", value="Whole picture", elem_id="img2img_inpaint_full_res")
+
+                                with gr.Column(scale=4):
+                                    inpaint_full_res_padding = gr.Slider(label='Only masked padding, pixels', minimum=0, maximum=256, step=4, value=32, elem_id="img2img_inpaint_full_res_padding")
+
+                    if category not in {"accordions"}:
+                        scripts.scripts_img2img.setup_ui_for_section(category)
+
+            # the code below is meant to update the resolution label after the image in the image selection UI has changed.
+            # as it is now the event keeps firing continuously for inpaint edits, which ruins the page with constant requests.
+            # I assume this must be a gradio bug and for now we'll just do it for non-inpaint inputs.
+            for component in [init_img, sketch]:
+                component.change(fn=lambda: None, _js="updateImg2imgResizeToTextAfterChangingImage", inputs=[], outputs=[], show_progress=False)
+
+            def select_img2img_tab(tab):
+                return gr.update(visible=tab in [2, 3, 4]), gr.update(visible=tab == 3),
+
+            for i, elem in enumerate(img2img_tabs):
+                elem.select(
+                    fn=lambda tab=i: select_img2img_tab(tab),
+                    inputs=[],
+                    outputs=[inpaint_controls, mask_alpha],
+                )
+
+            output_panel = create_output_panel("img2img", opts.outdir_img2img_samples, toprow)
+
+            img2img_args = dict(
+                fn=wrap_gradio_gpu_call(modules.img2img.img2img, extra_outputs=[None, '', '']),
+                _js="submit_img2img",
+                inputs=[
+                    dummy_component,
+                    dummy_component,
+                    toprow.prompt,
+                    toprow.negative_prompt,
+                    toprow.ui_styles.dropdown,
+                    init_img,
+                    sketch,
+                    init_img_with_mask,
+                    inpaint_color_sketch,
+                    inpaint_color_sketch_orig,
+                    init_img_inpaint,
+                    init_mask_inpaint,
+                    mask_blur,
+                    mask_alpha,
+                    inpainting_fill,
+                    batch_count,
+                    batch_size,
+                    cfg_scale,
+                    image_cfg_scale,
+                    denoising_strength,
+                    selected_scale_tab,
+                    height,
+                    width,
+                    scale_by,
+                    resize_mode,
+                    inpaint_full_res,
+                    inpaint_full_res_padding,
+                    inpainting_mask_invert,
+                    img2img_batch_input_dir,
+                    img2img_batch_output_dir,
+                    img2img_batch_inpaint_mask_dir,
+                    override_settings,
+                    img2img_batch_use_png_info,
+                    img2img_batch_png_info_props,
+                    img2img_batch_png_info_dir,
+                    img2img_batch_source_type,
+                    img2img_batch_upload,
+                ] + custom_inputs,
+                outputs=[
+                    output_panel.gallery,
+                    output_panel.generation_info,
+                    output_panel.infotext,
+                    output_panel.html_log,
+                ],
+                show_progress=False,
+            )
+
+            interrogate_args = dict(
+                _js="get_img2img_tab_index",
+                inputs=[
+                    dummy_component,
+                    img2img_batch_input_dir,
+                    img2img_batch_output_dir,
+                    init_img,
+                    sketch,
+                    init_img_with_mask,
+                    inpaint_color_sketch,
+                    init_img_inpaint,
+                ],
+                outputs=[toprow.prompt, dummy_component],
+            )
+
+            toprow.prompt.submit(**img2img_args)
+            toprow.submit.click(**img2img_args)
+
+            res_switch_btn.click(fn=None, _js="function(){switchWidthHeight('img2img')}", inputs=None, outputs=None, show_progress=False)
+
+            detect_image_size_btn.click(
+                fn=lambda w, h, _: (w or gr.update(), h or gr.update()),
+                _js="currentImg2imgSourceResolution",
+                inputs=[dummy_component, dummy_component, dummy_component],
+                outputs=[width, height],
+                show_progress=False,
+            )
+
+            toprow.restore_progress_button.click(
+                fn=progress.restore_progress,
+                _js="restoreProgressImg2img",
+                inputs=[dummy_component],
+                outputs=[
+                    output_panel.gallery,
+                    output_panel.generation_info,
+                    output_panel.infotext,
+                    output_panel.html_log,
+                ],
+                show_progress=False,
+            )
+
+            toprow.button_interrogate.click(
+                fn=lambda *args: process_interrogate(interrogate, *args),
+                **interrogate_args,
+            )
+
+            toprow.button_deepbooru.click(
+                fn=lambda *args: process_interrogate(interrogate_deepbooru, *args),
+                **interrogate_args,
+            )
+
+            steps = scripts.scripts_img2img.script('Sampler').steps
+
+            toprow.ui_styles.dropdown.change(fn=wrap_queued_call(update_token_counter), inputs=[toprow.prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.token_counter])
+            toprow.ui_styles.dropdown.change(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.negative_token_counter])
+            toprow.token_button.click(fn=update_token_counter, inputs=[toprow.prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.token_counter])
+            toprow.negative_token_button.click(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.negative_token_counter])
+
+            img2img_paste_fields = [
+                (toprow.prompt, "Prompt"),
+                (toprow.negative_prompt, "Negative prompt"),
+                (cfg_scale, "CFG scale"),
+                (image_cfg_scale, "Image CFG scale"),
+                (width, "Size-1"),
+                (height, "Size-2"),
+                (batch_size, "Batch size"),
+                (toprow.ui_styles.dropdown, lambda d: d["Styles array"] if isinstance(d.get("Styles array"), list) else gr.update()),
+                (denoising_strength, "Denoising strength"),
+                (mask_blur, "Mask blur"),
+                (inpainting_mask_invert, 'Mask mode'),
+                (inpainting_fill, 'Masked content'),
+                (inpaint_full_res, 'Inpaint area'),
+                (inpaint_full_res_padding, 'Masked area padding'),
+                *scripts.scripts_img2img.infotext_fields
+            ]
+            parameters_copypaste.add_paste_fields("img2img", init_img, img2img_paste_fields, override_settings)
+            parameters_copypaste.add_paste_fields("inpaint", init_img_with_mask, img2img_paste_fields, override_settings)
+            parameters_copypaste.register_paste_params_button(parameters_copypaste.ParamBinding(
+                paste_button=toprow.paste, tabname="img2img", source_text_component=toprow.prompt, source_image_component=None,
+            ))
+
+        extra_networks_ui_img2img = ui_extra_networks.create_ui(img2img_interface, [img2img_generation_tab], 'img2img')
+        ui_extra_networks.setup_ui(extra_networks_ui_img2img, output_panel.gallery)
+
+        extra_tabs.__exit__()
+
+    scripts.scripts_current = None
+
+    with gr.Blocks(analytics_enabled=False) as extras_interface:
+        ui_postprocessing.create_ui()
+
+    with gr.Blocks(analytics_enabled=False) as pnginfo_interface:
+        with ResizeHandleRow(equal_height=False):
+            with gr.Column(variant='panel'):
+                image = gr.Image(elem_id="pnginfo_image", label="Source", source="upload", interactive=True, type="pil")
+
+            with gr.Column(variant='panel'):
+                html = gr.HTML()
+                generation_info = gr.Textbox(visible=False, elem_id="pnginfo_generation_info")
+                html2 = gr.HTML()
+                with gr.Row():
+                    buttons = parameters_copypaste.create_buttons(["txt2img", "img2img", "inpaint", "extras"])
+
+                for tabname, button in buttons.items():
+                    parameters_copypaste.register_paste_params_button(parameters_copypaste.ParamBinding(
+                        paste_button=button, tabname=tabname, source_text_component=generation_info, source_image_component=image,
+                    ))
+
+        image.change(
+            fn=wrap_gradio_call_no_job(modules.extras.run_pnginfo),
+            inputs=[image],
+            outputs=[html, generation_info, html2],
+        )
+
+    modelmerger_ui = ui_checkpoint_merger.UiCheckpointMerger()
+
+    with gr.Blocks(analytics_enabled=False) as train_interface:
+        with gr.Row(equal_height=False):
+            gr.HTML(value="<p style='margin-bottom: 0.7em'>See <b><a href=\"https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Textual-Inversion\">wiki</a></b> for detailed explanation.</p>")
+
+        with ResizeHandleRow(variant="compact", equal_height=False):
+            with gr.Tabs(elem_id="train_tabs"):
+
+                with gr.Tab(label="Create embedding", id="create_embedding"):
+                    new_embedding_name = gr.Textbox(label="Name", elem_id="train_new_embedding_name")
+                    initialization_text = gr.Textbox(label="Initialization text", value="*", elem_id="train_initialization_text")
+                    nvpt = gr.Slider(label="Number of vectors per token", minimum=1, maximum=75, step=1, value=1, elem_id="train_nvpt")
+                    overwrite_old_embedding = gr.Checkbox(value=False, label="Overwrite Old Embedding", elem_id="train_overwrite_old_embedding")
+
+                    with gr.Row():
+                        with gr.Column(scale=3):
+                            gr.HTML(value="")
+
+                        with gr.Column():
+                            create_embedding = gr.Button(value="Create embedding", variant='primary', elem_id="train_create_embedding")
+
+                with gr.Tab(label="Create hypernetwork", id="create_hypernetwork"):
+                    new_hypernetwork_name = gr.Textbox(label="Name", elem_id="train_new_hypernetwork_name")
+                    new_hypernetwork_sizes = gr.CheckboxGroup(label="Modules", value=["768", "320", "640", "1280"], choices=["768", "1024", "320", "640", "1280"], elem_id="train_new_hypernetwork_sizes")
+                    new_hypernetwork_layer_structure = gr.Textbox("1, 2, 1", label="Enter hypernetwork layer structure", placeholder="1st and last digit must be 1. ex:'1, 2, 1'", elem_id="train_new_hypernetwork_layer_structure")
+                    new_hypernetwork_activation_func = gr.Dropdown(value="linear", label="Select activation function of hypernetwork. Recommended : Swish / Linear(none)", choices=hypernetworks_ui.keys, elem_id="train_new_hypernetwork_activation_func")
+                    new_hypernetwork_initialization_option = gr.Dropdown(value = "Normal", label="Select Layer weights initialization. Recommended: Kaiming for relu-like, Xavier for sigmoid-like, Normal otherwise", choices=["Normal", "KaimingUniform", "KaimingNormal", "XavierUniform", "XavierNormal"], elem_id="train_new_hypernetwork_initialization_option")
+                    new_hypernetwork_add_layer_norm = gr.Checkbox(label="Add layer normalization", elem_id="train_new_hypernetwork_add_layer_norm")
+                    new_hypernetwork_use_dropout = gr.Checkbox(label="Use dropout", elem_id="train_new_hypernetwork_use_dropout")
+                    new_hypernetwork_dropout_structure = gr.Textbox("0, 0, 0", label="Enter hypernetwork Dropout structure (or empty). Recommended : 0~0.35 incrementing sequence: 0, 0.05, 0.15", placeholder="1st and last digit must be 0 and values should be between 0 and 1. ex:'0, 0.01, 0'")
+                    overwrite_old_hypernetwork = gr.Checkbox(value=False, label="Overwrite Old Hypernetwork", elem_id="train_overwrite_old_hypernetwork")
+
+                    with gr.Row():
+                        with gr.Column(scale=3):
+                            gr.HTML(value="")
+
+                        with gr.Column():
+                            create_hypernetwork = gr.Button(value="Create hypernetwork", variant='primary', elem_id="train_create_hypernetwork")
+
+                def get_textual_inversion_template_names():
+                    return sorted(textual_inversion.textual_inversion_templates)
+
+                with gr.Tab(label="Train", id="train"):
+                    gr.HTML(value="<p style='margin-bottom: 0.7em'>Train an embedding or Hypernetwork; you must specify a directory with a set of 1:1 ratio images <a href=\"https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Textual-Inversion\" style=\"font-weight:bold;\">[wiki]</a></p>")
+                    with FormRow():
+                        train_embedding_name = gr.Dropdown(label='Embedding', elem_id="train_embedding", choices=sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys()))
+                        create_refresh_button(train_embedding_name, sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings, lambda: {"choices": sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys())}, "refresh_train_embedding_name")
+
+                        train_hypernetwork_name = gr.Dropdown(label='Hypernetwork', elem_id="train_hypernetwork", choices=sorted(shared.hypernetworks))
+                        create_refresh_button(train_hypernetwork_name, shared.reload_hypernetworks, lambda: {"choices": sorted(shared.hypernetworks)}, "refresh_train_hypernetwork_name")
+
+                    with FormRow():
+                        embedding_learn_rate = gr.Textbox(label='Embedding Learning rate', placeholder="Embedding Learning rate", value="0.005", elem_id="train_embedding_learn_rate")
+                        hypernetwork_learn_rate = gr.Textbox(label='Hypernetwork Learning rate', placeholder="Hypernetwork Learning rate", value="0.00001", elem_id="train_hypernetwork_learn_rate")
+
+                    with FormRow():
+                        clip_grad_mode = gr.Dropdown(value="disabled", label="Gradient Clipping", choices=["disabled", "value", "norm"])
+                        clip_grad_value = gr.Textbox(placeholder="Gradient clip value", value="0.1", show_label=False)
+
+                    with FormRow():
+                        batch_size = gr.Number(label='Batch size', value=1, precision=0, elem_id="train_batch_size")
+                        gradient_step = gr.Number(label='Gradient accumulation steps', value=1, precision=0, elem_id="train_gradient_step")
+
+                    dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images", elem_id="train_dataset_directory")
+                    log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value="textual_inversion", elem_id="train_log_directory")
+
+                    with FormRow():
+                        template_file = gr.Dropdown(label='Prompt template', value="style_filewords.txt", elem_id="train_template_file", choices=get_textual_inversion_template_names())
+                        create_refresh_button(template_file, textual_inversion.list_textual_inversion_templates, lambda: {"choices": get_textual_inversion_template_names()}, "refrsh_train_template_file")
+
+                    training_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="train_training_width")
+                    training_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="train_training_height")
+                    varsize = gr.Checkbox(label="Do not resize images", value=False, elem_id="train_varsize")
+                    steps = gr.Number(label='Max steps', value=100000, precision=0, elem_id="train_steps")
+
+                    with FormRow():
+                        create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable', value=500, precision=0, elem_id="train_create_image_every")
+                        save_embedding_every = gr.Number(label='Save a copy of embedding to log directory every N steps, 0 to disable', value=500, precision=0, elem_id="train_save_embedding_every")
+
+                    use_weight = gr.Checkbox(label="Use PNG alpha channel as loss weight", value=False, elem_id="use_weight")
+
+                    save_image_with_stored_embedding = gr.Checkbox(label='Save images with embedding in PNG chunks', value=True, elem_id="train_save_image_with_stored_embedding")
+                    preview_from_txt2img = gr.Checkbox(label='Read parameters (prompt, etc...) from txt2img tab when making previews', value=False, elem_id="train_preview_from_txt2img")
+
+                    shuffle_tags = gr.Checkbox(label="Shuffle tags by ',' when creating prompts.", value=False, elem_id="train_shuffle_tags")
+                    tag_drop_out = gr.Slider(minimum=0, maximum=1, step=0.1, label="Drop out tags when creating prompts.", value=0, elem_id="train_tag_drop_out")
+
+                    latent_sampling_method = gr.Radio(label='Choose latent sampling method', value="once", choices=['once', 'deterministic', 'random'], elem_id="train_latent_sampling_method")
+
+                    with gr.Row():
+                        train_embedding = gr.Button(value="Train Embedding", variant='primary', elem_id="train_train_embedding")
+                        interrupt_training = gr.Button(value="Interrupt", elem_id="train_interrupt_training")
+                        train_hypernetwork = gr.Button(value="Train Hypernetwork", variant='primary', elem_id="train_train_hypernetwork")
+
+                params = script_callbacks.UiTrainTabParams(txt2img_preview_params)
+
+                script_callbacks.ui_train_tabs_callback(params)
+
+            with gr.Column(elem_id='ti_gallery_container'):
+                ti_output = gr.Text(elem_id="ti_output", value="", show_label=False)
+                gr.Gallery(label='Output', show_label=False, elem_id='ti_gallery', columns=4)
+                gr.HTML(elem_id="ti_progress", value="")
+                ti_outcome = gr.HTML(elem_id="ti_error", value="")
+
+        create_embedding.click(
+            fn=textual_inversion_ui.create_embedding,
+            inputs=[
+                new_embedding_name,
+                initialization_text,
+                nvpt,
+                overwrite_old_embedding,
+            ],
+            outputs=[
+                train_embedding_name,
+                ti_output,
+                ti_outcome,
+            ]
+        )
+
+        create_hypernetwork.click(
+            fn=hypernetworks_ui.create_hypernetwork,
+            inputs=[
+                new_hypernetwork_name,
+                new_hypernetwork_sizes,
+                overwrite_old_hypernetwork,
+                new_hypernetwork_layer_structure,
+                new_hypernetwork_activation_func,
+                new_hypernetwork_initialization_option,
+                new_hypernetwork_add_layer_norm,
+                new_hypernetwork_use_dropout,
+                new_hypernetwork_dropout_structure
+            ],
+            outputs=[
+                train_hypernetwork_name,
+                ti_output,
+                ti_outcome,
+            ]
+        )
+
+        train_embedding.click(
+            fn=wrap_gradio_gpu_call(textual_inversion_ui.train_embedding, extra_outputs=[gr.update()]),
+            _js="start_training_textual_inversion",
+            inputs=[
+                dummy_component,
+                train_embedding_name,
+                embedding_learn_rate,
+                batch_size,
+                gradient_step,
+                dataset_directory,
+                log_directory,
+                training_width,
+                training_height,
+                varsize,
+                steps,
+                clip_grad_mode,
+                clip_grad_value,
+                shuffle_tags,
+                tag_drop_out,
+                latent_sampling_method,
+                use_weight,
+                create_image_every,
+                save_embedding_every,
+                template_file,
+                save_image_with_stored_embedding,
+                preview_from_txt2img,
+                *txt2img_preview_params,
+            ],
+            outputs=[
+                ti_output,
+                ti_outcome,
+            ]
+        )
+
+        train_hypernetwork.click(
+            fn=wrap_gradio_gpu_call(hypernetworks_ui.train_hypernetwork, extra_outputs=[gr.update()]),
+            _js="start_training_textual_inversion",
+            inputs=[
+                dummy_component,
+                train_hypernetwork_name,
+                hypernetwork_learn_rate,
+                batch_size,
+                gradient_step,
+                dataset_directory,
+                log_directory,
+                training_width,
+                training_height,
+                varsize,
+                steps,
+                clip_grad_mode,
+                clip_grad_value,
+                shuffle_tags,
+                tag_drop_out,
+                latent_sampling_method,
+                use_weight,
+                create_image_every,
+                save_embedding_every,
+                template_file,
+                preview_from_txt2img,
+                *txt2img_preview_params,
+            ],
+            outputs=[
+                ti_output,
+                ti_outcome,
+            ]
+        )
+
+        interrupt_training.click(
+            fn=lambda: shared.state.interrupt(),
+            inputs=[],
+            outputs=[],
+        )
+
+    loadsave = ui_loadsave.UiLoadsave(cmd_opts.ui_config_file)
+    ui_settings_from_file = loadsave.ui_settings.copy()
+
+    settings.create_ui(loadsave, dummy_component)
+
+    interfaces = [
+        (txt2img_interface, "txt2img", "txt2img"),
+        (img2img_interface, "img2img", "img2img"),
+        (extras_interface, "Extras", "extras"),
+        (pnginfo_interface, "PNG Info", "pnginfo"),
+        (modelmerger_ui.blocks, "Checkpoint Merger", "modelmerger"),
+        (train_interface, "Train", "train"),
+    ]
+
+    interfaces += script_callbacks.ui_tabs_callback()
+    interfaces += [(settings.interface, "Settings", "settings")]
+
+    extensions_interface = ui_extensions.create_ui()
+    interfaces += [(extensions_interface, "Extensions", "extensions")]
+
+    shared.tab_names = []
+    for _interface, label, _ifid in interfaces:
+        shared.tab_names.append(label)
+
+    with gr.Blocks(theme=shared.gradio_theme, analytics_enabled=False, title="Stable Diffusion") as demo:
+        settings.add_quicksettings()
+
+        parameters_copypaste.connect_paste_params_buttons()
+
+        with gr.Tabs(elem_id="tabs") as tabs:
+            tab_order = {k: i for i, k in enumerate(opts.ui_tab_order)}
+            sorted_interfaces = sorted(interfaces, key=lambda x: tab_order.get(x[1], 9999))
+
+            for interface, label, ifid in sorted_interfaces:
+                if label in shared.opts.hidden_tabs:
+                    continue
+                with gr.TabItem(label, id=ifid, elem_id=f"tab_{ifid}"):
+                    interface.render()
+
+                if ifid not in ["extensions", "settings"]:
+                    loadsave.add_block(interface, ifid)
+
+            loadsave.add_component(f"webui/Tabs@{tabs.elem_id}", tabs)
+
+            loadsave.setup_ui()
+
+        if os.path.exists(os.path.join(script_path, "notification.mp3")) and shared.opts.notification_audio:
+            gr.Audio(interactive=False, value=os.path.join(script_path, "notification.mp3"), elem_id="audio_notification", visible=False)
+
+        footer = shared.html("footer.html")
+        footer = footer.format(versions=versions_html(), api_docs="/docs" if shared.cmd_opts.api else "https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/API")
+        gr.HTML(footer, elem_id="footer")
+
+        settings.add_functionality(demo)
+
+        update_image_cfg_scale_visibility = lambda: gr.update(visible=shared.sd_model and shared.sd_model.cond_stage_key == "edit")
+        settings.text_settings.change(fn=update_image_cfg_scale_visibility, inputs=[], outputs=[image_cfg_scale])
+        demo.load(fn=update_image_cfg_scale_visibility, inputs=[], outputs=[image_cfg_scale])
+
+        modelmerger_ui.setup_ui(dummy_component=dummy_component, sd_model_checkpoint_component=settings.component_dict['sd_model_checkpoint'])
+
+    if ui_settings_from_file != loadsave.ui_settings:
+        loadsave.dump_defaults()
+    demo.ui_loadsave = loadsave
+
+    return demo
+
+
+def versions_html():
+    import torch
+    import launch
+
+    python_version = ".".join([str(x) for x in sys.version_info[0:3]])
+    commit = launch.commit_hash()
+    tag = launch.git_tag()
+
+    if shared.xformers_available:
+        import xformers
+        xformers_version = xformers.__version__
+    else:
+        xformers_version = "N/A"
+
+    return f"""
+version: <a href="https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/{commit}">{tag}</a>
+&#x2000;•&#x2000;
+python: <span title="{sys.version}">{python_version}</span>
+&#x2000;•&#x2000;
+torch: {getattr(torch, '__long_version__',torch.__version__)}
+&#x2000;•&#x2000;
+xformers: {xformers_version}
+&#x2000;•&#x2000;
+gradio: {gr.__version__}
+&#x2000;•&#x2000;
+checkpoint: <a id="sd_checkpoint_hash">N/A</a>
+"""
+
+
+def setup_ui_api(app):
+    from pydantic import BaseModel, Field
+
+    class QuicksettingsHint(BaseModel):
+        name: str = Field(title="Name of the quicksettings field")
+        label: str = Field(title="Label of the quicksettings field")
+
+    def quicksettings_hint():
+        return [QuicksettingsHint(name=k, label=v.label) for k, v in opts.data_labels.items()]
+
+    app.add_api_route("/internal/quicksettings-hint", quicksettings_hint, methods=["GET"], response_model=list[QuicksettingsHint])
+
+    app.add_api_route("/internal/ping", lambda: {}, methods=["GET"])
+
+    app.add_api_route("/internal/profile-startup", lambda: timer.startup_record, methods=["GET"])
+
+    def download_sysinfo(attachment=False):
+        from fastapi.responses import PlainTextResponse
+
+        text = sysinfo.get()
+        filename = f"sysinfo-{datetime.datetime.utcnow().strftime('%Y-%m-%d-%H-%M')}.json"
+
+        return PlainTextResponse(text, headers={'Content-Disposition': f'{"attachment" if attachment else "inline"}; filename="{filename}"'})
+
+    app.add_api_route("/internal/sysinfo", download_sysinfo, methods=["GET"])
+    app.add_api_route("/internal/sysinfo-download", lambda: download_sysinfo(attachment=True), methods=["GET"])
+
+    import fastapi.staticfiles
+    app.mount("/webui-assets", fastapi.staticfiles.StaticFiles(directory=launch_utils.repo_dir('stable-diffusion-webui-assets')), name="webui-assets")