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

More generic unet prefix detection code.

Browse Source
This commit is contained in:
comfyanonymous 2024-07-23 14:13:32 -04:00
parent 14764aa2e2
commit 334ba48cea
1 changed files with 15 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
comfy/model_detection.py
View File

@ -261,13 +261,22 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal
return model_config
def unet_prefix_from_state_dict(state_dict):
if "model.model.postprocess_conv.weight" in state_dict: #audio models
unet_key_prefix = "model.model."
elif "model.double_layers.0.attn.w1q.weight" in state_dict: #aura flow
unet_key_prefix = "model."
candidates = ["model.diffusion_model.", #ldm/sgm models
"model.model.", #audio models
]
counts = {k: 0 for k in candidates}
for k in state_dict:
for c in candidates:
if k.startswith(c):
counts[c] += 1
break
top = max(counts, key=counts.get)
if counts[top] > 5:
return top
else:
unet_key_prefix = "model.diffusion_model."
return unet_key_prefix
return "model." #aura flow and others
def convert_config(unet_config):
new_config = unet_config.copy()