From 112416d04171e4bee673f0adc9bd3aeba87ec71a Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 22 Jan 2023 10:17:12 +0300
Subject: [PATCH] add option to discard weights in checkpoint merger UI

---
 modules/extras.py | 9 ++++++++-
 modules/ui.py     | 4 ++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/modules/extras.py b/modules/extras.py
index 1218f88f..385430dc 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 import math
 import os
+import re
 import sys
 import traceback
 import shutil
@@ -285,7 +286,7 @@ def to_half(tensor, enable):
     return tensor
 
 
-def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae):
+def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights):
     shared.state.begin()
     shared.state.job = 'model-merge'
 
@@ -430,6 +431,12 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_
         for key in theta_0.keys():
             theta_0[key] = to_half(theta_0[key], save_as_half)
 
+    if discard_weights:
+        regex = re.compile(discard_weights)
+        for key in list(theta_0):
+            if re.search(regex, key):
+                theta_0.pop(key, None)
+
     ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path
 
     filename = filename_generator() if custom_name == '' else custom_name
diff --git a/modules/ui.py b/modules/ui.py
index af6dfb21..eb4b7e6b 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1248,6 +1248,9 @@ def create_ui():
                             bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None", label="Bake in VAE", elem_id="modelmerger_bake_in_vae")
                             create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["None"] + list(sd_vae.vae_dict)}, "modelmerger_refresh_bake_in_vae")
 
+                with FormRow():
+                    discard_weights = gr.Textbox(value="", label="Discard weights with matching name", elem_id="modelmerger_discard_weights")
+
                 with gr.Row():
                     modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')
 
@@ -1838,6 +1841,7 @@ def create_ui():
                 checkpoint_format,
                 config_source,
                 bake_in_vae,
+                discard_weights,
             ],
             outputs=[
                 primary_model_name,