From 91643f651d2794349876b12abbf2449cdc4f30b6 Mon Sep 17 00:00:00 2001 From: William Moorehouse Date: Sun, 25 Sep 2022 19:22:12 -0400 Subject: [PATCH 1/2] 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 382ffa7df..2c5b1fd6c 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"

{message}

" 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 "

Model saved to " + output_modelname + "

" diff --git a/modules/ui.py b/modules/ui.py index f7ca55883..6b3ba2f77 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="

A merger of the two checkpoints will be generated in your /models directory.

") + + 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 9ea5f5a32..c70a11c7c 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( From dc11748dea36e7618a7cdad55526fa9d6faaa6cf Mon Sep 17 00:00:00 2001 From: William Moorehouse Date: Mon, 26 Sep 2022 10:50:21 -0400 Subject: [PATCH 2/2] Added smoothstep interpolation to checkpoint merging --- modules/extras.py | 19 +++++++++++++++++-- modules/ui.py | 6 ++++-- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/modules/extras.py b/modules/extras.py index 2c5b1fd6c..a9788e7de 100644 --- a/modules/extras.py +++ b/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: diff --git a/modules/ui.py b/modules/ui.py index 6b3ba2f77..6525676ca 100644 --- a/modules/ui.py +++ b/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,