Support loading some checkpoint files with nested dicts.

comfyanonymous · Dec 11, 2024 · 7a7efe8 · 7a7efe8
1 parent 44db978
commit 7a7efe8
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/comfy/utils.py b/comfy/utils.py
@@ -46,7 +46,13 @@ def load_torch_file(ckpt, safe_load=False, device=None):
         if "state_dict" in pl_sd:
             sd = pl_sd["state_dict"]
         else:
-            sd = pl_sd
+            if len(pl_sd) == 1:
+                key = list(pl_sd.keys())[0]
+                sd = pl_sd[key]
+                if not isinstance(sd, dict):
+                    sd = pl_sd
+            else:
+                sd = pl_sd
     return sd
 
 def save_torch_file(sd, ckpt, metadata=None):