add hiresfix_mode

camenduru · May 13, 2023 · 3196d74 · 3196d74
1 parent 07ad1a3
commit 3196d74
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 3 deletions.
diff --git a/api/models/diffusion.py b/api/models/diffusion.py
@@ -22,6 +22,7 @@ class ImageGenerationOptions:
 
     image: PIL.Image.Image = field(default_factory=PIL.Image.Image)
     hiresfix: bool = False
+    hiresfix_mode: str = "bilinear"
     hiresfix_scale: float = 1.5
 
     def dict(self):

diff --git a/modules/components/image_generation_options.py b/modules/components/image_generation_options.py
@@ -91,11 +91,17 @@ def common_options_ui():
 def hires_options_ui():
     with gr.Row():
         with gr.Accordion("Hires.fix", open=False):
-            enable_upscale = gr.Checkbox(label="Hires.fix")
+            with gr.Row():
+                enable_upscale = gr.Checkbox(label="Hires.fix")
+                upscaler_mode = gr.Dropdown(
+                    choices=["nearest", "bilinear", "bicubic"],
+                    value="bilinear",
+                    label="Latent upscaler mode",
+                )
             scale_slider = gr.Slider(
                 value=1.5, minimum=1, maximum=2, step=0.1, label="Scale"
             )
-    return enable_upscale, scale_slider
+    return enable_upscale, upscaler_mode, scale_slider
 
 
 def img2img_options_ui():

diff --git a/modules/diffusion/pipelines/diffusers.py b/modules/diffusion/pipelines/diffusers.py
@@ -405,7 +405,7 @@ def __call__(
             opts.height = int(opts.height * opts.hiresfix_scale)
             opts.width = int(opts.width * opts.hiresfix_scale)
             opts.image = torch.nn.functional.interpolate(
-                opts.image, (opts.height // 8, opts.width // 8), mode="bilinear"
+                opts.image, (opts.height // 8, opts.width // 8), mode=opts.hiresfix_mode
             )
             opts.image = self.create_output(opts.image, "pil", True).images[0]
 

diff --git a/modules/tabs/generate.py b/modules/tabs/generate.py
@@ -23,6 +23,7 @@ def wrapper(self, data):
             width,
             height,
             hiresfix,
+            hiresfix_mode,
             hiresfix_scale,
             init_image,
             strength,
@@ -44,6 +45,7 @@ def wrapper(self, data):
             seed=seed,
             image=init_image,
             hiresfix=hiresfix,
+            hiresfix_mode=hiresfix_mode,
             hiresfix_scale=hiresfix_scale,
         )
         yield from fn(self, opts, plugin_values)