Merge branch 'master' of github.com:AUTOMATIC1111/stable-diffusion-webui into gamepad

1 files changed, 228 insertions, 101 deletions

diff --git a/modules/processing.py b/modules/processing.py
index 4a406084..5072fc40 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -13,10 +13,11 @@ from skimage import exposure
 from typing import Any, Dict, List, Optional

 

 import modules.sd_hijack

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

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

 from modules.sd_hijack import model_hijack

 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

@@ -50,9 +51,9 @@ def apply_color_correction(correction, original_image):
         correction,

         channel_axis=2

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

-    

+

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

-    

+

     return image

 

 

@@ -76,15 +77,32 @@ def apply_overlay(image, paste_loc, index, overlays):
     return image

 

 

-class StableDiffusionProcessing():

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

+    if sd_model.model.conditioning_key not in {'hybrid', 'concat'}:

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

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

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

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

+

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

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

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

+

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

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

+    image_conditioning = image_conditioning.to(x.dtype)

+

+    return image_conditioning

+

+

+class StableDiffusionProcessing:

     """

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

     """

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

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

         if sampler_index is not None:

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

 

-        self.sd_model = sd_model

         self.outpath_samples: str = outpath_samples

         self.outpath_grids: str = outpath_grids

         self.prompt: str = prompt

@@ -123,6 +141,7 @@ class StableDiffusionProcessing():
         self.override_settings = {k: v for k, v in (override_settings or {}).items() if k not in shared.restricted_opts}

         self.override_settings_restore_afterwards = override_settings_restore_afterwards

         self.is_using_inpainting_conditioning = False

+        self.disable_extra_networks = False

 

         if not seed_enable_extras:

             self.subseed = -1

@@ -131,33 +150,21 @@ class StableDiffusionProcessing():
             self.seed_resize_from_w = 0

 

         self.scripts = None

-        self.script_args = None

+        self.script_args = script_args

         self.all_prompts = None

         self.all_negative_prompts = None

         self.all_seeds = None

         self.all_subseeds = None

+        self.iteration = 0

 

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

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

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

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

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

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

+    @property

+    def sd_model(self):

+        return shared.sd_model

 

-        self.is_using_inpainting_conditioning = True

-

-        height = height or self.height

-        width = width or self.width

-

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

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

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

-

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

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

+        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

@@ -166,7 +173,8 @@ class StableDiffusionProcessing():
         midas_in = torch.from_numpy(transformed["midas_in"][None, ...]).to(device=shared.device)

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

 

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

+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(source_image.to(devices.dtype_vae) if devices.unet_needs_upcast else source_image))

+        conditioning_image = conditioning_image.float() if devices.unet_needs_upcast else conditioning_image

         conditioning = torch.nn.functional.interpolate(

             self.sd_model.depth_model(midas_in),

             size=conditioning_image.shape[2:],

@@ -178,7 +186,12 @@ class StableDiffusionProcessing():
         conditioning = 2. * (conditioning - depth_min) / (depth_max - depth_min) - 1.

         return conditioning

 

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

+    def edit_image_conditioning(self, source_image):

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

+

+        return conditioning_image

+

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

         self.is_using_inpainting_conditioning = True

 

         # Handle the different mask inputs

@@ -197,7 +210,7 @@ class StableDiffusionProcessing():
 

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

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

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

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

         conditioning_image = torch.lerp(

             source_image,

             source_image * (1.0 - conditioning_mask),

@@ -205,7 +218,7 @@ class StableDiffusionProcessing():
         )

 

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

+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(conditioning_image.to(devices.dtype_vae) if devices.unet_needs_upcast else 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:])

@@ -219,10 +232,13 @@ class StableDiffusionProcessing():
         # 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)

+            return self.depth2img_image_conditioning(source_image.float() if devices.unet_needs_upcast else 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)

+            return self.inpainting_image_conditioning(source_image.float() if devices.unet_needs_upcast else 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)

@@ -234,12 +250,11 @@ class StableDiffusionProcessing():
         raise NotImplementedError()

 

     def close(self):

-        self.sd_model = None

         self.sampler = None

 

 

 class Processed:

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

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

         self.images = images_list

         self.prompt = p.prompt

         self.negative_prompt = p.negative_prompt

@@ -247,6 +262,7 @@ class Processed:
         self.subseed = subseed

         self.subseed_strength = p.subseed_strength

         self.info = info

+        self.comments = comments

         self.width = p.width

         self.height = p.height

         self.sampler_name = p.sampler_name

@@ -338,13 +354,14 @@ def slerp(val, low, high):
 

 

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

+    eta_noise_seed_delta = opts.eta_noise_seed_delta or 0

     xs = []

 

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

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

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

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

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

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

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

     else:

         sampler_noises = None

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

             cnt = p.sampler.number_of_needed_noises(p)

 

-            if opts.eta_noise_seed_delta > 0:

-                torch.manual_seed(seed + opts.eta_noise_seed_delta)

+            if eta_noise_seed_delta > 0:

+                torch.manual_seed(seed + eta_noise_seed_delta)

 

             for j in range(cnt):

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

@@ -401,7 +418,7 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
 

 def decode_first_stage(model, x):

     with devices.autocast(disable=x.dtype == devices.dtype_vae):

-        x = model.decode_first_stage(x)

+        x = model.decode_first_stage(x.to(devices.dtype_vae) if devices.unet_needs_upcast else x)

 

     return x

 

@@ -418,7 +435,7 @@ def fix_seed(p):
     p.subseed = get_fixed_seed(p.subseed)

 

 

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

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

     index = position_in_batch + iteration * p.batch_size

 

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

@@ -432,9 +449,6 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration
         "Size": f"{p.width}x{p.height}",

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

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

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

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

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

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

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

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

@@ -462,9 +476,12 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
     try:

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

             setattr(opts, k, v)

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

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

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

+

+            if k == 'sd_model_checkpoint':

+                sd_models.reload_model_weights()

+

+            if k == 'sd_vae':

+                sd_vae.reload_vae_weights()

 

         res = process_images_inner(p)

 

@@ -473,9 +490,11 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
         if p.override_settings_restore_afterwards:

             for k, v in stored_opts.items():

                 setattr(opts, k, v)

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

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

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

+                if k == 'sd_model_checkpoint':

+                    sd_models.reload_model_weights()

+

+                if k == 'sd_vae':

+                    sd_vae.reload_vae_weights()

 

     return res

 

@@ -521,27 +540,61 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     def infotext(iteration=0, position_in_batch=0):

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

 

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

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

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

-

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

         model_hijack.embedding_db.load_textual_inversion_embeddings()

 

+    _, extra_network_data = extra_networks.parse_prompts(p.all_prompts[0:1])

+

     if p.scripts is not None:

         p.scripts.process(p)

 

     infotexts = []

     output_images = []

 

+    cached_uc = [None, None]

+    cached_c = [None, None]

+

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

+        """

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

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

+

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

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

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

+        computed result is stored.

+        """

+

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

+            return cache[1]

+

+        with devices.autocast():

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

+

+        cache[0] = (required_prompts, steps)

+        return cache[1]

+

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

         with devices.autocast():

             p.init(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()

+

+            if not p.disable_extra_networks:

+                extra_networks.activate(p, extra_network_data)

+

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

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

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

+

         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

 

@@ -556,12 +609,13 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if len(prompts) == 0:

                 break

 

+            prompts, _ = extra_networks.parse_prompts(prompts)

+

             if p.scripts is not None:

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

 

-            with devices.autocast():

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

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

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

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

 

             if len(model_hijack.comments) > 0:

                 for comment in model_hijack.comments:

@@ -570,10 +624,13 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if p.n_iter > 1:

                 shared.state.job = f"Batch {n+1} out of {p.n_iter}"

 

-            with devices.autocast():

+            with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast():

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

 

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

+            for x in x_samples_ddim:

+                devices.test_for_nans(x, "vae")

+

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

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

 

@@ -602,6 +659,11 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 

                 image = Image.fromarray(x_sample)

 

+                if p.scripts is not None:

+                    pp = scripts.PostprocessImageArgs(image)

+                    p.scripts.postprocess_image(p, pp)

+                    image = pp.image

+

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

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

                         image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images)

@@ -643,9 +705,12 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             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(), short_filename=not opts.grid_extended_filename, p=p, grid=True)

 

+    if not p.disable_extra_networks:

+        extra_networks.deactivate(p, extra_network_data)

+

     devices.torch_gc()

 

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

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

 

     if p.scripts is not None:

         p.scripts.postprocess(p, res)

@@ -653,79 +718,138 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
     return res

 

 

+def old_hires_fix_first_pass_dimensions(width, height):

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

+

+    desired_pixel_count = 512 * 512

+    actual_pixel_count = width * height

+    scale = math.sqrt(desired_pixel_count / actual_pixel_count)

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

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

+

+    return width, height

+

+

 class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):

     sampler = None

 

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

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

         super().__init__(**kwargs)

         self.enable_hr = enable_hr

         self.denoising_strength = denoising_strength

-        self.firstphase_width = firstphase_width

-        self.firstphase_height = firstphase_height

+        self.hr_scale = hr_scale

+        self.hr_upscaler = hr_upscaler

+        self.hr_second_pass_steps = hr_second_pass_steps

+        self.hr_resize_x = hr_resize_x

+        self.hr_resize_y = hr_resize_y

+        self.hr_upscale_to_x = hr_resize_x

+        self.hr_upscale_to_y = hr_resize_y

+

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

+            self.hr_upscale_to_x = self.width

+            self.hr_upscale_to_y = self.height

+            self.width = firstphase_width

+            self.height = firstphase_height

+

         self.truncate_x = 0

         self.truncate_y = 0

+        self.applied_old_hires_behavior_to = None

 

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

         if self.enable_hr:

-            if state.job_count == -1:

-                state.job_count = self.n_iter * 2

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

+                self.hr_resize_x = self.width

+                self.hr_resize_y = self.height

+                self.hr_upscale_to_x = self.width

+                self.hr_upscale_to_y = self.height

+

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

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

+

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

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

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

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

             else:

-                state.job_count = state.job_count * 2

-

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

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

+

+                if self.hr_resize_y == 0:

+                    self.hr_upscale_to_x = self.hr_resize_x

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

+                elif self.hr_resize_x == 0:

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

+                    self.hr_upscale_to_y = self.hr_resize_y

+                else:

+                    target_w = self.hr_resize_x

+                    target_h = self.hr_resize_y

+                    src_ratio = self.width / self.height

+                    dst_ratio = self.hr_resize_x / self.hr_resize_y

+

+                    if src_ratio < dst_ratio:

+                        self.hr_upscale_to_x = self.hr_resize_x

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

+                    else:

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

+                        self.hr_upscale_to_y = self.hr_resize_y

+

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

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

+

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

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

+                self.enable_hr = False

+                self.denoising_strength = None

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

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

+                return

 

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

-                desired_pixel_count = 512 * 512

-                actual_pixel_count = self.width * self.height

-                scale = math.sqrt(desired_pixel_count / actual_pixel_count)

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

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

-                firstphase_width_truncated = int(scale * self.width)

-                firstphase_height_truncated = int(scale * self.height)

+            if not state.processing_has_refined_job_count:

+                if state.job_count == -1:

+                    state.job_count = self.n_iter

 

-            else:

-

-                width_ratio = self.width / self.firstphase_width

-                height_ratio = self.height / self.firstphase_height

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

+                state.job_count = state.job_count * 2

+                state.processing_has_refined_job_count = True

 

-                if width_ratio > height_ratio:

-                    firstphase_width_truncated = self.firstphase_width

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

-                else:

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

-                    firstphase_height_truncated = self.firstphase_height

+            if self.hr_second_pass_steps:

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

 

-            self.truncate_x = int(self.firstphase_width - firstphase_width_truncated) // opt_f

-            self.truncate_y = int(self.firstphase_height - firstphase_height_truncated) // opt_f

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

 

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

+        if self.enable_hr and latent_scale_mode is None:

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

+

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

+        samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))

+

         if not self.enable_hr:

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

-            samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))

             return samples

 

-        x = create_random_tensors([opt_C, self.firstphase_height // opt_f, self.firstphase_width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

-        samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x, self.firstphase_width, self.firstphase_height))

-

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

+        target_width = self.hr_upscale_to_x

+        target_height = self.hr_upscale_to_y

 

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

         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 opts.save or self.do_not_save_samples or not opts.save_images_before_highres_fix:

                 return

 

             if not isinstance(image, Image.Image):

-                image = sd_samplers.sample_to_image(image, index)

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

 

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

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

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

 

-        if opts.use_scale_latent_for_hires_fix:

+        if latent_scale_mode is not None:

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

                 save_intermediate(samples, i)

 

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

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

 

             # Avoid making the inpainting conditioning unless necessary as

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

@@ -745,7 +869,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

                 save_intermediate(image, i)

 

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

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

@@ -760,15 +884,18 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

         shared.state.nextjob()

 

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

+        img2img_sampler_name = self.sampler_name if self.sampler_name != 'PLMS' else 'DDIM'  # PLMS does not support img2img so we just silently switch ot DDIM

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

+

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

 

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

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

 

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

         x = None

         devices.torch_gc()

 

-        samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.steps, image_conditioning=image_conditioning)

+        samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.hr_second_pass_steps or self.steps, image_conditioning=image_conditioning)

 

         return samples

 

@@ -880,7 +1007,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
 

         image = torch.from_numpy(batch_images)

         image = 2. * image - 1.

-        image = image.to(shared.device)

+        image = image.to(device=shared.device, dtype=devices.dtype_vae if devices.unet_needs_upcast else None)

 

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