Add option to use noisy latent for hires fix

modules/generation_parameters_copypaste.py

@@ -289,6 +289,11 @@ def parse_generation_parameters(x: str):
     if "Hires negative prompt" not in res:
         res["Hires negative prompt"] = ""
 
+    if "Hires noisy latent" not in res:
+        res["Hires noisy latent"] = False
+    elif res["Hires noisy latent"] is True:
+        res["Denoising strength"] = 1.0
+
     restore_old_hires_fix_params(res)
 
     # Missing RNG means the default was set, which is GPU RNG

modules/processing.py

@@ -943,10 +943,10 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
     cached_hr_uc = [None, None]
     cached_hr_c = [None, None]
 
-    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, hr_checkpoint_name: str = None, hr_sampler_name: str = None, hr_prompt: str = '', hr_negative_prompt: str = '', **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, hr_checkpoint_name: str = None, hr_sampler_name: str = None, hr_use_noisy: bool = False, hr_prompt: str = '', hr_negative_prompt: str = '', **kwargs):
         super().__init__(**kwargs)
         self.enable_hr = enable_hr
-        self.denoising_strength = denoising_strength
+        self.denoising_strength = denoising_strength if not hr_use_noisy else 1.0
         self.hr_scale = hr_scale
         self.hr_upscaler = hr_upscaler
         self.hr_second_pass_steps = hr_second_pass_steps
@@ -957,6 +957,7 @@ def __init__(self, enable_hr: bool = False, denoising_strength: float = 0.75, fi
         self.hr_checkpoint_name = hr_checkpoint_name
         self.hr_checkpoint_info = None
         self.hr_sampler_name = hr_sampler_name
+        self.hr_use_noisy = hr_use_noisy
         self.hr_prompt = hr_prompt
         self.hr_negative_prompt = hr_negative_prompt
         self.all_hr_prompts = None
@@ -1152,7 +1153,10 @@ def save_intermediate(image, index):
 
         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, p=self)
+        if self.hr_use_noisy:
+            noise = torch.zeros(samples.shape[1:], device=shared.device)
+        else:
+            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
         devices.torch_gc()

modules/sd_samplers_kdiffusion.py

@@ -340,6 +340,17 @@ def initialize(self, p):
         return extra_params_kwargs
 
     def get_sigmas(self, p, steps):
+        if getattr(p, 'enable_hr', False) and getattr(p, 'hr_use_noisy', False) and getattr(p, 'hr_second_pass_steps', 0) > 0:
+            p.extra_generation_params["Hires noisy latent"] = True
+            if p.is_hr_pass:
+                trim_steps = p.steps
+            else:
+                hr_second_pass_steps, hr_second_pass_t_enc = sd_samplers_common.setup_img2img_steps(p, p.hr_second_pass_steps)
+                trim_steps = hr_second_pass_t_enc + 1
+            steps += trim_steps
+        else:
+            trim_steps = 0
+
         discard_next_to_last_sigma = self.config is not None and self.config.options.get('discard_next_to_last_sigma', False)
         if opts.always_discard_next_to_last_sigma and not discard_next_to_last_sigma:
             discard_next_to_last_sigma = True
@@ -384,6 +395,12 @@ def get_sigmas(self, p, steps):
         if discard_next_to_last_sigma:
             sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
 
+        if trim_steps > 0:
+            if p.is_hr_pass:
+                sigmas = sigmas[trim_steps:]
+            else:
+                sigmas = sigmas[:-trim_steps]
+
         return sigmas
 
     def create_noise_sampler(self, x, sigmas, p):

modules/txt2img.py

@@ -9,7 +9,7 @@
 import gradio as gr
 
 
-def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, 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_index: int, hr_prompt: str, hr_negative_prompt, override_settings_texts, request: gr.Request, *args):
+def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, 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_index: int, hr_use_noisy: bool, hr_prompt: str, hr_negative_prompt, override_settings_texts, request: gr.Request, *args):
     override_settings = create_override_settings_dict(override_settings_texts)
 
     p = processing.StableDiffusionProcessingTxt2Img(
@@ -43,6 +43,7 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step
         hr_resize_y=hr_resize_y,
         hr_checkpoint_name=None if hr_checkpoint_name == 'Use same checkpoint' else hr_checkpoint_name,
         hr_sampler_name=sd_samplers.samplers_for_img2img[hr_sampler_index - 1].name if hr_sampler_index != 0 else None,
+        hr_use_noisy=hr_use_noisy,
         hr_prompt=hr_prompt,
         hr_negative_prompt=hr_negative_prompt,
         override_settings=override_settings,

modules/ui.py

@@ -482,6 +482,8 @@ def create_ui():
 
                                 hr_sampler_index = gr.Dropdown(label='Hires sampling method', elem_id="hr_sampler", choices=["Use same sampler"] + [x.name for x in samplers_for_img2img], value="Use same sampler", type="index")
 
+                                hr_use_noisy = gr.Checkbox(label='Hires noisy latent', value=False, elem_id="hr_use_noisy_latent")
+
                             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():
@@ -559,6 +561,7 @@ def create_ui():
                     hr_resize_y,
                     hr_checkpoint_name,
                     hr_sampler_index,
+                    hr_use_noisy,
                     hr_prompt,
                     hr_negative_prompt,
                     override_settings,
@@ -611,6 +614,13 @@ def create_ui():
                 show_progress = False,
             )
 
+            hr_use_noisy.change(
+                fn=lambda x: gr_show(not x),
+                inputs=[hr_use_noisy],
+                outputs=[denoising_strength],
+                show_progress = False,
+            )
+
             txt2img_paste_fields = [
                 (txt2img_prompt, "Prompt"),
                 (txt2img_negative_prompt, "Negative prompt"),
@@ -637,6 +647,7 @@ def create_ui():
                 (hr_resize_y, "Hires resize-2"),
                 (hr_checkpoint_name, "Hires checkpoint"),
                 (hr_sampler_index, "Hires sampler"),
+                (hr_use_noisy, "Hires noisy latent"),
                 (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" else gr.update()),
                 (hr_prompt, "Hires prompt"),
                 (hr_negative_prompt, "Hires negative prompt"),