Add "Scale to" option to Extras

3 files changed, 50 insertions, 19 deletions

diff --git a/javascript/ui.js b/javascript/ui.js
index b1053201..4100944e 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -101,7 +101,8 @@ function create_tab_index_args(tabId, args){
 }
 
 function get_extras_tab_index(){
-    return create_tab_index_args('mode_extras', arguments)
+    const [,,...args] = [...arguments]
+    return [get_tab_index('mode_extras'), get_tab_index('extras_resize_mode'), ...args]
 }
 
 function create_submit_args(args){
diff --git a/modules/extras.py b/modules/extras.py
index 41e8612c..83ca7049 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,3 +1,4 @@
+import math

 import os

 

 import numpy as np

@@ -19,7 +20,7 @@ import gradio as gr
 cached_images = {}

 

 

-def run_extras(extras_mode, image, image_folder, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):

+def run_extras(extras_mode, resize_mode, image, image_folder, 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):

     devices.torch_gc()

 

     imageArr = []

@@ -67,8 +68,23 @@ def run_extras(extras_mode, image, image_folder, gfpgan_visibility, codeformer_v
             info += f"CodeFormer w: {round(codeformer_weight, 2)}, CodeFormer visibility:{round(codeformer_visibility, 2)}\n"

             image = res

 

+        if resize_mode == 1:

+            upscaling_resize = max(upscaling_resize_w/image.width, upscaling_resize_h/image.height)

+            crop_info = " (crop)" if upscaling_crop else ""

+            info += f"Resize to: {upscaling_resize_w:g}x{upscaling_resize_h:g}{crop_info}\n"

+

+            def crop_upscaled_center(image, resize_w, resize_h):

+                left = int(math.ceil((image.width - resize_w) / 2))

+                right = image.width - int(math.floor((image.width - resize_w) / 2))

+                top = int(math.ceil((image.height - resize_h) / 2))

+                bottom = image.height - int(math.floor((image.height - resize_h) / 2))

+

+                image = image.crop((left, top, right, bottom))

+                return image

+

+

         if upscaling_resize != 1.0:

-            def upscale(image, scaler_index, resize):

+            def upscale(image, scaler_index, resize, mode, resize_w, resize_h, crop):

                 small = image.crop((image.width // 2, image.height // 2, image.width // 2 + 10, image.height // 2 + 10))

                 pixels = tuple(np.array(small).flatten().tolist())

                 key = (resize, scaler_index, image.width, image.height, gfpgan_visibility, codeformer_visibility, codeformer_weight) + pixels

@@ -77,15 +93,17 @@ def run_extras(extras_mode, image, image_folder, gfpgan_visibility, codeformer_v
                 if c is None:

                     upscaler = shared.sd_upscalers[scaler_index]

                     c = upscaler.scaler.upscale(image, resize, upscaler.data_path)

+                    if mode == 1 and crop:

+                        c = crop_upscaled_center(c, resize_w, resize_h)

                     cached_images[key] = c

 

                 return c

 

             info += f"Upscale: {round(upscaling_resize, 3)}, model:{shared.sd_upscalers[extras_upscaler_1].name}\n"

-            res = upscale(image, extras_upscaler_1, upscaling_resize)

+            res = upscale(image, extras_upscaler_1, upscaling_resize, resize_mode, upscaling_resize_w, upscaling_resize_h, upscaling_crop)

 

             if extras_upscaler_2 != 0 and extras_upscaler_2_visibility > 0:

-                res2 = upscale(image, extras_upscaler_2, upscaling_resize)

+                res2 = upscale(image, extras_upscaler_2, upscaling_resize, resize_mode, upscaling_resize_w, upscaling_resize_h, upscaling_crop)

                 info += f"Upscale: {round(upscaling_resize, 3)}, visibility: {round(extras_upscaler_2_visibility, 3)}, model:{shared.sd_upscalers[extras_upscaler_2].name}\n"

                 res = Image.blend(res, res2, extras_upscaler_2_visibility)

 

@@ -190,7 +208,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, interp_method, int
             theta_0[key] = theta_func(theta_0[key], theta_1[key], (float(1.0) - interp_amount))  # Need to reverse the interp_amount to match the desired mix ration in the merged checkpoint

             if save_as_half:

                 theta_0[key] = theta_0[key].half()

-    

+

     for key in theta_1.keys():

         if 'model' in key and key not in theta_0:

             theta_0[key] = theta_1[key]

diff --git a/modules/ui.py b/modules/ui.py
index 2231a8ed..4bb2892b 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -101,7 +101,7 @@ def send_gradio_gallery_to_image(x):
 

 

 def save_files(js_data, images, do_make_zip, index):

-    import csv    

+    import csv

     filenames = []

     fullfns = []

 

@@ -551,7 +551,7 @@ def create_ui(wrap_gradio_gpu_call):
 

                     with gr.Row():

                         do_make_zip = gr.Checkbox(label="Make Zip when Save?", value=False)

-                    

+

                     with gr.Row():

                         download_files = gr.File(None, file_count="multiple", interactive=False, show_label=False, visible=False)

 

@@ -739,7 +739,7 @@ def create_ui(wrap_gradio_gpu_call):
 

                     with gr.Row():

                         do_make_zip = gr.Checkbox(label="Make Zip when Save?", value=False)

-                    

+

                     with gr.Row():

                         download_files = gr.File(None, file_count="multiple", interactive=False, show_label=False, visible=False)

 

@@ -903,7 +903,15 @@ def create_ui(wrap_gradio_gpu_call):
                     with gr.TabItem('Batch Process'):

                         image_batch = gr.File(label="Batch Process", file_count="multiple", interactive=True, type="file")

 

-                upscaling_resize = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Resize", value=2)

+                with gr.Tabs(elem_id="extras_resize_mode"):

+                    with gr.TabItem('Scale by'):

+                        upscaling_resize = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Resize", value=2)

+                    with gr.TabItem('Scale to'):

+                        with gr.Group():

+                            with gr.Row():

+                                upscaling_resize_w = gr.Number(label="Width", value=512)

+                                upscaling_resize_h = gr.Number(label="Height", value=512)

+                            upscaling_crop = gr.Checkbox(label='Crop to fit', value=True)

 

                 with gr.Group():

                     extras_upscaler_1 = gr.Radio(label='Upscaler 1', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index")

@@ -935,12 +943,16 @@ def create_ui(wrap_gradio_gpu_call):
             _js="get_extras_tab_index",

             inputs=[

                 dummy_component,

+                dummy_component,

                 extras_image,

                 image_batch,

                 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,

@@ -951,14 +963,14 @@ def create_ui(wrap_gradio_gpu_call):
                 html_info,

             ]

         )

-     

+

         extras_send_to_img2img.click(

             fn=lambda x: image_from_url_text(x),

             _js="extract_image_from_gallery_img2img",

             inputs=[result_images],

             outputs=[init_img],

         )

-        

+

         extras_send_to_inpaint.click(

             fn=lambda x: image_from_url_text(x),

             _js="extract_image_from_gallery_img2img",

@@ -1286,7 +1298,7 @@ Requested path was: {f}
             outputs=[],

             _js='function(){restart_reload()}'

         )

-        

+

         if column is not None:

             column.__exit__()

 

@@ -1318,12 +1330,12 @@ Requested path was: {f}
                 component_dict[k] = component

 

         settings_interface.gradio_ref = demo

-        

+

         with gr.Tabs() as tabs:

             for interface, label, ifid in interfaces:

                 with gr.TabItem(label, id=ifid):

                     interface.render()

-        

+

         if os.path.exists(os.path.join(script_path, "notification.mp3")):

             audio_notification = gr.Audio(interactive=False, value=os.path.join(script_path, "notification.mp3"), elem_id="audio_notification", visible=False)

 

@@ -1456,10 +1468,10 @@ Requested path was: {f}
 

             if getattr(obj,'custom_script_source',None) is not None:

               key = 'customscript/' + obj.custom_script_source + '/' + key

-            

+

             if getattr(obj, 'do_not_save_to_config', False):

                 return

-            

+

             saved_value = ui_settings.get(key, None)

             if saved_value is None:

                 ui_settings[key] = getattr(obj, field)

@@ -1483,10 +1495,10 @@ Requested path was: {f}
 

         if type(x) == gr.Textbox:

             apply_field(x, 'value')

-        

+

         if type(x) == gr.Number:

             apply_field(x, 'value')

-        

+

     visit(txt2img_interface, loadsave, "txt2img")

     visit(img2img_interface, loadsave, "img2img")

     visit(extras_interface, loadsave, "extras")