Generalize SD torch load/save to implement safetensor merging compat

1 files changed, 8 insertions, 7 deletions

diff --git a/modules/extras.py b/modules/extras.py
index 71b93a06..820427de 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -249,7 +249,7 @@ def run_pnginfo(image):
     return '', geninfo, info

 

 

-def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_name, interp_method, multiplier, save_as_half, custom_name):

+def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_name, interp_method, multiplier, save_as_half, save_as_safetensors, custom_name):

     def weighted_sum(theta0, theta1, alpha):

         return ((1 - alpha) * theta0) + (alpha * theta1)

 

@@ -264,16 +264,16 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
     teritary_model_info = sd_models.checkpoints_list.get(teritary_model_name, None)

 

     print(f"Loading {primary_model_info.filename}...")

-    primary_model = torch.load(primary_model_info.filename, map_location='cpu')

+    primary_model = sd_models.torch_load(primary_model_info.filename, primary_model_info, map_override='cpu')

     theta_0 = sd_models.get_state_dict_from_checkpoint(primary_model)

 

     print(f"Loading {secondary_model_info.filename}...")

-    secondary_model = torch.load(secondary_model_info.filename, map_location='cpu')

+    secondary_model = sd_models.torch_load(secondary_model_info.filename, primary_model_info, map_override='cpu')

     theta_1 = sd_models.get_state_dict_from_checkpoint(secondary_model)

 

     if teritary_model_info is not None:

         print(f"Loading {teritary_model_info.filename}...")

-        teritary_model = torch.load(teritary_model_info.filename, map_location='cpu')

+        teritary_model = sd_models.torch_load(teritary_model_info.filename, teritary_model_info, map_override='cpu')

         theta_2 = sd_models.get_state_dict_from_checkpoint(teritary_model)

     else:

         teritary_model = None

@@ -314,12 +314,13 @@ def run_modelmerger(primary_model_name, secondary_model_name, teritary_model_nam
 

     ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path

 

-    filename = primary_model_info.model_name + '_' + str(round(1-multiplier, 2)) + '-' + secondary_model_info.model_name + '_' + str(round(multiplier, 2)) + '-' + interp_method.replace(" ", "_") + '-merged.ckpt'

-    filename = filename if custom_name == '' else (custom_name + '.ckpt')

+    output_exttype = '.safetensors' if save_as_safetensors else '.ckpt'

+    filename = primary_model_info.model_name + '_' + str(round(1-multiplier, 2)) + '-' + secondary_model_info.model_name + '_' + str(round(multiplier, 2)) + '-' + interp_method.replace(" ", "_") + '-merged' + output_exttype

+    filename = filename if custom_name == '' else (custom_name + output_exttype)

     output_modelname = os.path.join(ckpt_dir, filename)

 

     print(f"Saving to {output_modelname}...")

-    torch.save(primary_model, output_modelname)

+    sd_models.torch_save(primary_model, output_modelname)

 

     sd_models.list_models()