add safetensors support for model merging #4869

1 files changed, 15 insertions, 11 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 77236480..a1ea5611 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -160,6 +160,20 @@ def get_state_dict_from_checkpoint(pl_sd):
     return pl_sd

 

 

+def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):

+    _, extension = os.path.splitext(checkpoint_file)

+    if extension.lower() == ".safetensors":

+        pl_sd = safetensors.torch.load_file(checkpoint_file, device=map_location or shared.weight_load_location)

+    else:

+        pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)

+

+    if print_global_state and "global_step" in pl_sd:

+        print(f"Global Step: {pl_sd['global_step']}")

+

+    sd = get_state_dict_from_checkpoint(pl_sd)

+    return sd

+

+

 def load_model_weights(model, checkpoint_info, vae_file="auto"):

     checkpoint_file = checkpoint_info.filename

     sd_model_hash = checkpoint_info.hash

@@ -174,17 +188,7 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
         # load from file

         print(f"Loading weights [{sd_model_hash}] from {checkpoint_file}")

 

-        _, extension = os.path.splitext(checkpoint_file)

-        if extension.lower() == ".safetensors":

-            pl_sd = safetensors.torch.load_file(checkpoint_file, device=shared.weight_load_location)

-        else:

-            pl_sd = torch.load(checkpoint_file, map_location=shared.weight_load_location)

-

-        if "global_step" in pl_sd:

-            print(f"Global Step: {pl_sd['global_step']}")

-

-        sd = get_state_dict_from_checkpoint(pl_sd)

-        del pl_sd

+        sd = read_state_dict(checkpoint_file)

         model.load_state_dict(sd, strict=False)

         del sd