Skip to content

Commit 7a7efe8

Browse files
Support loading some checkpoint files with nested dicts.
1 parent 44db978 commit 7a7efe8

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

comfy/utils.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,13 @@ def load_torch_file(ckpt, safe_load=False, device=None):
4646
if "state_dict" in pl_sd:
4747
sd = pl_sd["state_dict"]
4848
else:
49-
sd = pl_sd
49+
if len(pl_sd) == 1:
50+
key = list(pl_sd.keys())[0]
51+
sd = pl_sd[key]
52+
if not isinstance(sd, dict):
53+
sd = pl_sd
54+
else:
55+
sd = pl_sd
5056
return sd
5157

5258
def save_torch_file(sd, ckpt, metadata=None):

0 commit comments

Comments
 (0)