298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-25 15:55:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support loading some checkpoint files with nested dicts.

Browse Source
This commit is contained in:
comfyanonymous 2024-12-11 08:04:54 -05:00
parent 44db978531
commit 7a7efe8424
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/utils.py
View File

@ -45,6 +45,12 @@ def load_torch_file(ckpt, safe_load=False, device=None):
logging.debug(f"Global Step: {pl_sd['global_step']}") logging.debug(f"Global Step: {pl_sd['global_step']}")
if "state_dict" in pl_sd: if "state_dict" in pl_sd:
sd = pl_sd["state_dict"] sd = pl_sd["state_dict"]
else:
if len(pl_sd) == 1:
key = list(pl_sd.keys())[0]
sd = pl_sd[key]
if not isinstance(sd, dict):
sd = pl_sd
else: else:
sd = pl_sd sd = pl_sd
return sd return sd