Merge branch 'master' into saving

1 files changed, 27 insertions, 19 deletions

diff --git a/modules/extras.py b/modules/extras.py
index b8ebc619..c2543fcf 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -6,13 +6,14 @@ from PIL import Image
 import torch

 import tqdm

 

-from modules import processing, shared, images, devices

+from modules import processing, shared, images, devices, sd_models

 from modules.shared import opts

 import modules.gfpgan_model

 from modules.ui import plaintext_to_html

 import modules.codeformer_model

 import piexif

 import piexif.helper

+import gradio as gr

 

 

 cached_images = {}

@@ -140,7 +141,7 @@ def run_pnginfo(image):
     return '', geninfo, info

 

 

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

+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):

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

@@ -150,23 +151,20 @@ def run_modelmerger(primary_model_name, secondary_model_name, interp_method, int
         alpha = alpha * alpha * (3 - (2 * alpha))

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

 

-    if os.path.exists(primary_model_name):

-        primary_model_filename = primary_model_name

-        primary_model_name = os.path.splitext(os.path.basename(primary_model_name))[0]

-    else:

-        primary_model_filename = 'models/' + primary_model_name + '.ckpt'

+    # 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)

 

-    if os.path.exists(secondary_model_name):

-        secondary_model_filename = secondary_model_name

-        secondary_model_name = os.path.splitext(os.path.basename(secondary_model_name))[0]

-    else:

-        secondary_model_filename = 'models/' + secondary_model_name + '.ckpt'

+    primary_model_info = sd_models.checkpoints_list[primary_model_name]

+    secondary_model_info = sd_models.checkpoints_list[secondary_model_name]

 

-    print(f"Loading {primary_model_filename}...")

-    primary_model = torch.load(primary_model_filename, map_location='cpu')

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

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

 

-    print(f"Loading {secondary_model_filename}...")

-    secondary_model = torch.load(secondary_model_filename, map_location='cpu')

+    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']

@@ -174,21 +172,31 @@ def run_modelmerger(primary_model_name, secondary_model_name, interp_method, int
     theta_funcs = {

         "Weighted Sum": weighted_sum,

         "Sigmoid": sigmoid,

+        "Inverse Sigmoid": inv_sigmoid,

     }

     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

+            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]

+            if save_as_half:

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

+

+    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 = filename if custom_name == '' else (custom_name + '.ckpt')

+    output_modelname = os.path.join(shared.cmd_opts.ckpt_dir, filename)

 

-    output_modelname = 'models/' + primary_model_name + '_' + str(round(interp_amount,2)) + '-' + secondary_model_name + '_' + str(round((float(1.0) - interp_amount),2)) + '-' + interp_method.replace(" ", "_") + '-merged.ckpt'

     print(f"Saving to {output_modelname}...")

     torch.save(primary_model, output_modelname)

 

+    sd_models.list_models()

+

     print(f"Checkpoint saved.")

-    return "Checkpoint saved to " + output_modelname
\ No newline at end of file
+    return ["Checkpoint saved to " + output_modelname] + [gr.Dropdown.update(choices=sd_models.checkpoint_tiles()) for _ in range(3)]