Merge pull request #15293 from light-and-ray/extras_upscaler_limit_target_resolution

Extras upscaler: option limit target resolution

modules/postprocessing.py

@@ -131,13 +131,14 @@ def run_postprocessing_webui(id_task, *args, **kwargs):
     return run_postprocessing(*args, **kwargs)
 
 
-def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, upscale_first: bool, save_output: bool = True):
+def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, upscale_first: bool, save_output: bool = True, limit_target_resolution = 0):
     """old handler for API"""
 
     args = scripts.scripts_postproc.create_args_for_run({
         "Upscale": {
             "upscale_mode": resize_mode,
             "upscale_by": upscaling_resize,
+            "limit_target_resolution": limit_target_resolution,
             "upscale_to_width": upscaling_resize_w,
             "upscale_to_height": upscaling_resize_h,
             "upscale_crop": upscaling_crop,

modules/shared_options.py

@@ -102,6 +102,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."),
+    "show_limit_target_resolution_in_extras_upscale": OptionInfo(False, 'Show "Limit target resolution" slider in "Upscale" extras script. Useful for batches where can be big images.').needs_reload_ui(),
 }))
 
 options_templates.update(options_section(('face-restoration', "Face restoration", "postprocessing"), {

scripts/postprocessing_upscale.py

@@ -12,6 +12,19 @@ from modules.ui import switch_values_symbol
 upscale_cache = {}
 
 
+def limitSizeByOneDemention(size: tuple, limit: int):
+    w, h = size
+    if h > w:
+        if h > limit:
+            w = limit / h * w
+            h = limit
+    else:
+        if w > limit:
+            h = limit / w * h
+            w = limit
+    return (int(w), int(h))
+
+
 class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
     name = "Upscale"
     order = 1000
@@ -31,6 +44,11 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
                 with gr.Tabs(elem_id="extras_resize_mode"):
                     with gr.TabItem('Scale by', elem_id="extras_scale_by_tab") as tab_scale_by:
                         upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id="extras_upscaling_resize")
+                        if shared.opts.show_limit_target_resolution_in_extras_upscale:
+                            limit_target_resolution = gr.Slider(minimum=0, maximum=10000, step=8, label="Limit target resolution", value=8000, elem_id="extras_upscale_limit_target_resolution",
+                                    tooltip="0 = no limit. Limit maximal target resolution by the biggest demension. Useful for batches where can be big images.")
+                        else:
+                            limit_target_resolution = gr.Number(0, visible=False)
 
                     with gr.TabItem('Scale to', elem_id="extras_scale_to_tab") as tab_scale_to:
                         with FormRow():
@@ -61,6 +79,7 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
             "upscale_enabled": upscale_enabled,
             "upscale_mode": selected_tab,
             "upscale_by": upscaling_resize,
+            "limit_target_resolution": limit_target_resolution,
             "upscale_to_width": upscaling_resize_w,
             "upscale_to_height": upscaling_resize_h,
             "upscale_crop": upscaling_crop,
@@ -69,12 +88,18 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
             "upscaler_2_visibility": extras_upscaler_2_visibility,
         }
 
-    def upscale(self, image, info, upscaler, upscale_mode, upscale_by,  upscale_to_width, upscale_to_height, upscale_crop):
+    def upscale(self, image, info, upscaler, upscale_mode, upscale_by, limit_target_resolution, upscale_to_width, upscale_to_height, upscale_crop):
         if upscale_mode == 1:
             upscale_by = max(upscale_to_width/image.width, upscale_to_height/image.height)
             info["Postprocess upscale to"] = f"{upscale_to_width}x{upscale_to_height}"
         else:
             info["Postprocess upscale by"] = upscale_by
+            if limit_target_resolution != 0 and max(*image.size)*upscale_by > limit_target_resolution:
+                upscale_mode = 1
+                upscale_crop = False
+                upscale_to_width, upscale_to_height = limitSizeByOneDemention((image.width*upscale_by, image.height*upscale_by), limit_target_resolution)
+                upscale_by = max(upscale_to_width/image.width, upscale_to_height/image.height)
+                info["Limit target resolution"] = limit_target_resolution
 
         cache_key = (hash(np.array(image.getdata()).tobytes()), upscaler.name, upscale_mode, upscale_by,  upscale_to_width, upscale_to_height, upscale_crop)
         cached_image = upscale_cache.pop(cache_key, None)
@@ -96,18 +121,21 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
 
         return image
 
-    def process_firstpass(self, pp: scripts_postprocessing.PostprocessedImage, upscale_enabled=True, upscale_mode=1, upscale_by=2.0, upscale_to_width=None, upscale_to_height=None, upscale_crop=False, upscaler_1_name=None, upscaler_2_name=None, upscaler_2_visibility=0.0):
-        if upscale_mode == 1:
-            pp.shared.target_width = upscale_to_width
-            pp.shared.target_height = upscale_to_height
+    def process_firstpass(self, pp: scripts_postprocessing.PostprocessedImage, **args):
+        if args['upscale_mode'] == 1:
+            pp.shared.target_width = args['upscale_to_width']
+            pp.shared.target_height = args['upscale_to_height']
         else:
-            pp.shared.target_width = int(pp.image.width * upscale_by)
-            pp.shared.target_height = int(pp.image.height * upscale_by)
+            pp.shared.target_width = int(pp.image.width * args['upscale_by'])
+            pp.shared.target_height = int(pp.image.height * args['upscale_by'])
+            if args['limit_target_resolution'] != 0:
+                pp.shared.target_width, pp.shared.target_height = limitSizeByOneDemention((pp.shared.target_width, pp.shared.target_height), args['limit_target_resolution'])
 
-    def process(self, pp: scripts_postprocessing.PostprocessedImage, upscale_enabled=True, upscale_mode=1, upscale_by=2.0, upscale_to_width=None, upscale_to_height=None, upscale_crop=False, upscaler_1_name=None, upscaler_2_name=None, upscaler_2_visibility=0.0):
-        if not upscale_enabled:
+    def process(self, pp: scripts_postprocessing.PostprocessedImage, **args):
+        if not args['upscale_enabled']:
             return
 
+        upscaler_1_name = args['upscaler_1_name']
         if upscaler_1_name == "None":
             upscaler_1_name = None
 
@@ -117,18 +145,21 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
         if not upscaler1:
             return
 
+        upscaler_2_name = args['upscaler_2_name']
         if upscaler_2_name == "None":
             upscaler_2_name = None
 
         upscaler2 = next(iter([x for x in shared.sd_upscalers if x.name == upscaler_2_name and x.name != "None"]), None)
         assert upscaler2 or (upscaler_2_name is None), f'could not find upscaler named {upscaler_2_name}'
 
-        upscaled_image = self.upscale(pp.image, pp.info, upscaler1, upscale_mode, upscale_by, upscale_to_width, upscale_to_height, upscale_crop)
+        upscaled_image = self.upscale(pp.image, pp.info, upscaler1, args['upscale_mode'], args['upscale_by'], args['limit_target_resolution'], args['upscale_to_width'],
+                    args['upscale_to_height'], args['upscale_crop'])
         pp.info["Postprocess upscaler"] = upscaler1.name
 
-        if upscaler2 and upscaler_2_visibility > 0:
-            second_upscale = self.upscale(pp.image, pp.info, upscaler2, upscale_mode, upscale_by, upscale_to_width, upscale_to_height, upscale_crop)
-            upscaled_image = Image.blend(upscaled_image, second_upscale, upscaler_2_visibility)
+        if upscaler2 and args['upscaler_2_visibility'] > 0:
+            second_upscale = self.upscale(pp.image, pp.info, upscaler2, args['upscale_mode'], args['upscale_by'], args['upscale_to_width'],
+                    args['upscale_to_height'], args['upscale_crop'])
+            upscaled_image = Image.blend(upscaled_image, second_upscale, args['upscaler_2_visibility'])
 
             pp.info["Postprocess upscaler 2"] = upscaler2.name
 
@@ -163,5 +194,5 @@ class ScriptPostprocessingUpscaleSimple(ScriptPostprocessingUpscale):
         upscaler1 = next(iter([x for x in shared.sd_upscalers if x.name == upscaler_name]), None)
         assert upscaler1, f'could not find upscaler named {upscaler_name}'
 
-        pp.image = self.upscale(pp.image, pp.info, upscaler1, 0, upscale_by, 0, 0, False)
+        pp.image = self.upscale(pp.image, pp.info, upscaler1, 0, upscale_by, 0, 0, 0, False)
         pp.info["Postprocess upscaler"] = upscaler1.name