Added smoothstep interpolation to checkpoint merging

modules/extras.py

@@ -139,16 +139,31 @@ def run_pnginfo(image):
     return '', geninfo, info
 
 
-def run_modelmerger(modelname_0, modelname_1, alpha):
+def run_modelmerger(modelname_0, modelname_1, interp_method, interp_amount):
+    # Linear interpolation (https://en.wikipedia.org/wiki/Linear_interpolation)
+    def weighted_sum(theta0, theta1, 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)
+
     model_0 = torch.load('models/' + modelname_0 + '.ckpt')
     model_1 = torch.load('models/' + modelname_1 + '.ckpt')
     
     theta_0 = model_0['state_dict']
     theta_1 = model_1['state_dict']
+    theta_func = weighted_sum
+    
+    if interp_method == "Weighted Sum":
+        theta_func = weighted_sum
+    if interp_method == "Sigmoid":
+        theta_func = sigmoid
     
     for key in theta_0.keys():
         if 'model' in key and key in theta_1:
-            theta_0[key] = (1 - alpha) * theta_0[key] + alpha * theta_1[key]
+            theta_0[key] = theta_func(theta_0[key], theta_1[key], interp_amount)
     
     for key in theta_1.keys():
         if 'model' in key and key not in theta_0:

modules/ui.py

@@ -860,7 +860,8 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo, run_modelmerger):
                 
                 modelname_0 = gr.Textbox(elem_id="modelmerger_modelname_0", label="Model Name (to)")
                 modelname_1 = gr.Textbox(elem_id="modelmerger_modelname_1", label="Model Name (from)")
-                alpha = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Alpha', value=0.3)
+                interp_method = gr.Radio(choices=["Weighted Sum", "Sigmoid"], value="Weighted Sum", label="Interpolation Method")
+                interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Interpolation Amount', value=0.3)
                 submit = gr.Button(elem_id="modelmerger_merge", label="Merge", variant='primary')
             
             with gr.Column(variant='panel'):
@@ -871,7 +872,8 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo, run_modelmerger):
                 inputs=[
                     modelname_0,
                     modelname_1,
-                    alpha
+                    interp_method,
+                    interp_amount
                 ],
                 outputs=[
                     submit_result,