aboutsummaryrefslogtreecommitdiff
path: root/modules/extras.py
diff options
context:
space:
mode:
Diffstat (limited to 'modules/extras.py')
-rw-r--r--modules/extras.py9
1 files changed, 8 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