From 06cd20610765aeb563700f377f1698a6e981b17d Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Tue, 20 Sep 2022 19:32:26 +0300
Subject: [PATCH] Enable neural network upscalers for highres. fix

---
 modules/processing.py | 22 +++++++++++++++++++++-
 modules/shared.py     |  1 +
 modules/ui.py         |  2 +-
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/modules/processing.py b/modules/processing.py
index 2bc19f6b..c9ba6eb3 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -450,7 +450,27 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             samples = torch.nn.functional.interpolate(samples, size=(self.height // opt_f, self.width // opt_f), mode="bilinear")
         else:
             decoded_samples = self.sd_model.decode_first_stage(samples)
-            decoded_samples = torch.nn.functional.interpolate(decoded_samples, size=(self.height, self.width), mode="bilinear")
+
+            if opts.upscaler_for_hires_fix is None or opts.upscaler_for_hires_fix == "None":
+                decoded_samples = torch.nn.functional.interpolate(decoded_samples, size=(self.height, self.width), mode="bilinear")
+            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)
+                    upscaler = [x for x in shared.sd_upscalers if x.name == opts.upscaler_for_hires_fix][0]
+                    image = upscaler.upscale(image, self.width, self.height)
+                    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)
+                decoded_samples = 2. * decoded_samples - 1.
+
             samples = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(decoded_samples))
 
         shared.state.nextjob()
diff --git a/modules/shared.py b/modules/shared.py
index 0b81d7f8..ae4efbee 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -144,6 +144,7 @@ class Options:
         "ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for upscaling. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),
         "random_artist_categories": OptionInfo([], "Allowed categories for random artists selection when using the Roll button", gr.CheckboxGroup, {"choices": artist_db.categories()}),
         "upscale_at_full_resolution_padding": OptionInfo(16, "Inpainting at full resolution: padding, in pixels, for the masked region.", gr.Slider, {"minimum": 0, "maximum": 128, "step": 4}),
+        "upscaler_for_hires_fix": OptionInfo(None, "Upscaler for highres. fix", gr.Radio, lambda: {"choices": [x.name for x in sd_upscalers]}),
         "show_progressbar": OptionInfo(True, "Show progressbar"),
         "show_progress_every_n_steps": OptionInfo(0, "Show show image creation progress every N sampling steps. Set 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}),
         "multiple_tqdm": OptionInfo(True, "Add a second progress bar to the console that shows progress for an entire job. Broken in PyCharm console."),
diff --git a/modules/ui.py b/modules/ui.py
index ec6f247e..752bc97b 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -396,7 +396,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     enable_hr = gr.Checkbox(label='Highres. fix', value=False)
 
                 with gr.Row(visible=False) as hr_options:
-                    scale_latent = gr.Checkbox(label='Scale latent', value=True)
+                    scale_latent = gr.Checkbox(label='Scale latent', value=False)
                     denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.7)
 
                 with gr.Row():