From 57aaa068bb20ea15925d08c7736845a5213c1cb2 Mon Sep 17 00:00:00 2001
From: Logan <loganbooker@gmail.com>
Date: Thu, 8 Aug 2024 15:20:48 +1000
Subject: [PATCH] Add fast prescale option to upscaler settings (off by
 default)

* By default, upscaling will loop up to 3 times to upscale the incoming image to the target dimensions. This is necessary, as upscaling models work in fixed increments (x4 is common).
* For very small images, such as those generated by ADetailer for inpainting, this can result in additional upscaling steps, which can be expensive. Usually the incoming image is only off by a small amount, so it can be preferable to do a minor upscale via Lanczos before the main upscaling step.
* We introduce an optional value to the upscaler settings to allow this minor upscale, should the incoming image fall below a certain threshold compared to the fixed scaling value of the upscaler model. By default, this setting is set to 1.0, effectively disabling it.
---
 modules/shared_options.py |  1 +
 modules/upscaler.py       | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/modules/shared_options.py b/modules/shared_options.py
index 9f4520274..c7a9017ab 100644
--- a/modules/shared_options.py
+++ b/modules/shared_options.py
@@ -105,6 +105,7 @@ options_templates.update(options_section(('upscaling', "Upscaling", "postprocess
     "DAT_tile_overlap": OptionInfo(8, "Tile overlap for DAT upscalers.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}).info("Low values = visible seam"),
     "upscaler_for_img2img": OptionInfo(None, "Upscaler for img2img", gr.Dropdown, lambda: {"choices": [x.name for x in shared.sd_upscalers]}),
     "set_scale_by_when_changing_upscaler": OptionInfo(False, "Automatically set the Scale by factor based on the name of the selected Upscaler."),
+    "upscaler_fast_prescale_threshold": OptionInfo(1.0, "Maximum threshold for fast pre-scaling of very small images.", gr.Slider, {"minimum": 1, "maximum": 8, "step": 0.01}).info("The maximum scale difference for performing fast Lanczos pre-scaling before the first upscaling step. This can prevent expensive additional steps if the source image is very small. 1 = disabled; 1.25-3.0 recommended"),
 }))
 
 options_templates.update(options_section(('face-restoration', "Face restoration", "postprocessing"), {
diff --git a/modules/upscaler.py b/modules/upscaler.py
index 507881fed..df0cd25c8 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -6,6 +6,7 @@ from PIL import Image
 
 import modules.shared
 from modules import modelloader, shared
+import math
 
 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
 NEAREST = (Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.NEAREST)
@@ -56,6 +57,40 @@ class Upscaler:
         dest_w = int((img.width * scale) // 8 * 8)
         dest_h = int((img.height * scale) // 8 * 8)
 
+        # Attempt a cheap resize of the source image, if it falls below the fixed scaling size of the upscaling model.
+        # We resize the image by the smallest amount necessary for the fixed scaling to meet the target dimensions.
+        
+        prescale_threshold = modules.shared.opts.upscaler_fast_prescale_threshold
+        if prescale_threshold > 1 and self.name and self.name not in ["Nearest", "Lanczos"]:
+            
+            # Get the matching upscaler
+            upscaler_data = next((x for x in self.scalers if x.data_path == selected_model), None)
+
+            if upscaler_data is not None:
+                upscaler_scale = upscaler_data.scale
+                if scale > upscaler_scale:
+
+                    # Calculate the minimum intermediate dimensions.
+                    min_intermediate_w = math.ceil(dest_w / upscaler_scale)
+                    min_intermediate_h = math.ceil(dest_h / upscaler_scale)
+
+                    # Preserve aspect ratio and make sure any adjustments don't drop us below the
+                    # minimum scaling needed.
+                    aspect_ratio = img.width / img.height
+
+                    intermediate_w = max(min_intermediate_w, int(math.ceil(min_intermediate_h * aspect_ratio)))
+                    intermediate_h = max(min_intermediate_h, int(math.ceil(min_intermediate_w / aspect_ratio)))
+
+                    if intermediate_w / aspect_ratio > intermediate_h:
+                        intermediate_w = int(math.ceil(intermediate_h * aspect_ratio))
+                    else:
+                        intermediate_h = int(math.ceil(intermediate_w / aspect_ratio))
+
+                    scale_diff = max(intermediate_w / img.width, intermediate_h / img.height)
+
+                    if scale_diff <= prescale_threshold:
+                        img = img.resize((intermediate_w, intermediate_h), resample=LANCZOS)
+
         for i in range(3):
             if img.width >= dest_w and img.height >= dest_h and (i > 0 or scale != 1):
                 break