Merge remote-tracking branch 'origin/master'

3 years ago · 45cc0ce3c4
parent 706d5944a0 616b7218f7
commit 45cc0ce3c4
1 changed files with 5 additions and 1 deletions
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@ -122,7 +122,11 @@ def load_model_weights(model, checkpoint_file, sd_model_hash):
    pl_sd = torch.load(checkpoint_file, map_location="cpu")
    if "global_step" in pl_sd:
        print(f"Global Step: {pl_sd['global_step']}")
-    sd = pl_sd["state_dict"]
+    
+    if "state_dict" in pl_sd:
+        sd = pl_sd["state_dict"]
+    else:
+        sd = pl_sd

    model.load_state_dict(sd, strict=False)