Merge branch 'master' into master

1 files changed, 48 insertions, 28 deletions

diff --git a/modules/extras.py b/modules/extras.py
index 6a0d5cb0..f2f5a7b0 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 = []

@@ -29,7 +30,7 @@ def run_extras(extras_mode, image, image_folder, gfpgan_visibility, codeformer_v
     if extras_mode == 1:

         #convert file to pillow image

         for img in image_folder:

-            image = Image.fromarray(np.array(Image.open(img)))

+            image = Image.open(img)

             imageArr.append(image)

             imageNameArr.append(os.path.splitext(img.orig_name)[0])

     else:

@@ -67,8 +68,13 @@ 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"

+

         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 +83,19 @@ 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:

+                        cropped = Image.new("RGB", (resize_w, resize_h))

+                        cropped.paste(c, box=(resize_w // 2 - c.width // 2, resize_h // 2 - c.height // 2))

+                        c = cropped

                     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)

 

@@ -98,8 +108,14 @@ def run_extras(extras_mode, image, image_folder, gfpgan_visibility, codeformer_v
                           no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=existing_pnginfo,

                           forced_filename=image_name if opts.use_original_name_batch else None)

 

+        if opts.enable_pnginfo:

+            image.info = existing_pnginfo

+            image.info["extras"] = info

+

         outputs.append(image)

 

+    devices.torch_gc()

+

     return outputs, plaintext_to_html(info), ''

 

 

@@ -143,48 +159,52 @@ def run_pnginfo(image):
     return '', geninfo, info

 

 

-def run_modelmerger(primary_model_name, secondary_model_name, interp_method, interp_amount, save_as_half, custom_name):

-    # Linear interpolation (https://en.wikipedia.org/wiki/Linear_interpolation)

-    def weighted_sum(theta0, theta1, alpha):

+def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_name, interp_method, multiplier, save_as_half, custom_name):

+    def weighted_sum(theta0, theta1, theta2, alpha):

         return ((1 - alpha) * theta0) + (alpha * theta1)

 

-    # Smoothstep (https://en.wikipedia.org/wiki/Smoothstep)

-    def sigmoid(theta0, theta1, alpha):

-        alpha = alpha * alpha * (3 - (2 * alpha))

-        return theta0 + ((theta1 - theta0) * alpha)

-

-    # Inverse Smoothstep (https://en.wikipedia.org/wiki/Smoothstep)

-    def inv_sigmoid(theta0, theta1, alpha):

-        import math

-        alpha = 0.5 - math.sin(math.asin(1.0 - 2.0 * alpha) / 3.0)

-        return theta0 + ((theta1 - theta0) * alpha)

+    def add_difference(theta0, theta1, theta2, alpha):

+        return theta0 + (theta1 - theta2) * alpha

 

     primary_model_info = sd_models.checkpoints_list[primary_model_name]

     secondary_model_info = sd_models.checkpoints_list[secondary_model_name]

+    teritary_model_info = sd_models.checkpoints_list.get(teritary_model_name, None)

 

     print(f"Loading {primary_model_info.filename}...")

     primary_model = torch.load(primary_model_info.filename, map_location='cpu')

+    theta_0 = sd_models.get_state_dict_from_checkpoint(primary_model)

 

     print(f"Loading {secondary_model_info.filename}...")

     secondary_model = torch.load(secondary_model_info.filename, map_location='cpu')

-   

-    theta_0 = primary_model['state_dict']

-    theta_1 = secondary_model['state_dict']

+    theta_1 = sd_models.get_state_dict_from_checkpoint(secondary_model)

+

+    if teritary_model_info is not None:

+        print(f"Loading {teritary_model_info.filename}...")

+        teritary_model = torch.load(teritary_model_info.filename, map_location='cpu')

+        theta_2 = sd_models.get_state_dict_from_checkpoint(teritary_model)

+    else:

+        theta_2 = None

 

     theta_funcs = {

-        "Weighted Sum": weighted_sum,

-        "Sigmoid": sigmoid,

-        "Inverse Sigmoid": inv_sigmoid,

+        "Weighted sum": weighted_sum,

+        "Add difference": add_difference,

     }

     theta_func = theta_funcs[interp_method]

 

     print(f"Merging...")

+

     for key in tqdm.tqdm(theta_0.keys()):

         if 'model' in key and key in theta_1:

-            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

+            t2 = (theta_2 or {}).get(key)

+            if t2 is None:

+                t2 = torch.zeros_like(theta_0[key])

+

+            theta_0[key] = theta_func(theta_0[key], theta_1[key], t2, multiplier)

+

             if save_as_half:

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

-    

+

+    # I believe this part should be discarded, but I'll leave it for now until I am sure

     for key in theta_1.keys():

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

             theta_0[key] = theta_1[key]

@@ -193,7 +213,7 @@ def run_modelmerger(primary_model_name, secondary_model_name, interp_method, int
 

     ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path

 

-    filename = primary_model_info.model_name + '_' + str(round(interp_amount, 2)) + '-' + secondary_model_info.model_name + '_' + str(round((float(1.0) - interp_amount), 2)) + '-' + interp_method.replace(" ", "_") + '-merged.ckpt'

+    filename = primary_model_info.model_name + '_' + str(round(1-multiplier, 2)) + '-' + secondary_model_info.model_name + '_' + str(round(multiplier, 2)) + '-' + interp_method.replace(" ", "_") + '-merged.ckpt'

     filename = filename if custom_name == '' else (custom_name + '.ckpt')

     output_modelname = os.path.join(ckpt_dir, filename)

 

@@ -203,4 +223,4 @@ def run_modelmerger(primary_model_name, secondary_model_name, interp_method, int
     sd_models.list_models()

 

     print(f"Checkpoint saved.")

-    return ["Checkpoint saved to " + output_modelname] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(3)]

+    return ["Checkpoint saved to " + output_modelname] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(4)]