From 91643f651d2794349876b12abbf2449cdc4f30b6 Mon Sep 17 00:00:00 2001
From: William Moorehouse <moorehousew@gmail.com>
Date: Sun, 25 Sep 2022 19:22:12 -0400
Subject: Add support for checkpoint merging

---
 modules/extras.py | 24 ++++++++++++++++++++++++
 modules/ui.py     | 28 +++++++++++++++++++++++++++-
 webui.py          |  3 ++-
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index 382ffa7d..2c5b1fd6 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -3,6 +3,8 @@ import os
 import numpy as np
 from PIL import Image
 
+import torch
+
 from modules import processing, shared, images, devices
 from modules.shared import opts
 import modules.gfpgan_model
@@ -135,3 +137,25 @@ def run_pnginfo(image):
         info = f"<div><p>{message}<p></div>"
 
     return '', geninfo, info
+
+
+def run_modelmerger(modelname_0, modelname_1, 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']
+    
+    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]
+    
+    for key in theta_1.keys():
+        if 'model' in key and key not in theta_0:
+            theta_0[key] = theta_1[key]
+    
+    output_modelname = 'models/' + modelname_0 + '-' + modelname_1 + '-merged.ckpt';
+    
+    torch.save(model_0, output_modelname)
+    
+    return "<p>Model saved to " + output_modelname + "</p>"
diff --git a/modules/ui.py b/modules/ui.py
index f7ca5588..6b3ba2f7 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -393,7 +393,7 @@ def setup_progressbar(progressbar, preview, id_part):
     )
 
 
-def create_ui(txt2img, img2img, run_extras, run_pnginfo):
+def create_ui(txt2img, img2img, run_extras, run_pnginfo, run_modelmerger):
     with gr.Blocks(analytics_enabled=False) as txt2img_interface:
         txt2img_prompt, roll, txt2img_prompt_style, txt2img_negative_prompt, txt2img_prompt_style2, submit, _, txt2img_prompt_style_apply, txt2img_save_style, paste = create_toprow(is_img2img=False)
         dummy_component = gr.Label(visible=False)
@@ -853,6 +853,31 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
             outputs=[html, generation_info, html2],
         )
 
+    with gr.Blocks() as modelmerger_interface:
+        with gr.Row().style(equal_height=False):
+            with gr.Column(variant='panel'):
+                gr.HTML(value="<p>A merger of the two checkpoints will be generated in your <b>/models</b> directory.</p>")
+                
+                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)
+                submit = gr.Button(elem_id="modelmerger_merge", label="Merge", variant='primary')
+            
+            with gr.Column(variant='panel'):
+                submit_result = gr.HTML(elem_id="modelmerger_result")
+
+            submit.click(
+                fn=run_modelmerger,
+                inputs=[
+                    modelname_0,
+                    modelname_1,
+                    alpha
+                ],
+                outputs=[
+                    submit_result,
+                ]
+            )
+
     def create_setting_component(key):
         def fun():
             return opts.data[key] if key in opts.data else opts.data_labels[key].default
@@ -950,6 +975,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
         (img2img_interface, "img2img", "img2img"),
         (extras_interface, "Extras", "extras"),
         (pnginfo_interface, "PNG Info", "pnginfo"),
+        (modelmerger_interface, "Checkpoint Merger", "modelmerger"),
         (settings_interface, "Settings", "settings"),
     ]
 
diff --git a/webui.py b/webui.py
index 9ea5f5a3..c70a11c7 100644
--- a/webui.py
+++ b/webui.py
@@ -85,7 +85,8 @@ def webui():
         txt2img=wrap_gradio_gpu_call(modules.txt2img.txt2img),
         img2img=wrap_gradio_gpu_call(modules.img2img.img2img),
         run_extras=wrap_gradio_gpu_call(modules.extras.run_extras),
-        run_pnginfo=modules.extras.run_pnginfo
+        run_pnginfo=modules.extras.run_pnginfo,
+        run_modelmerger=modules.extras.run_modelmerger
     )
 
     demo.launch(
-- 
cgit v1.2.1