diff options
| -rw-r--r-- | modules/extras.py | 9 | ||||
| -rw-r--r-- | modules/ui.py | 4 |
2 files changed, 12 insertions, 1 deletions
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,
|