wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-20 07:22:34 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

Add detection for TwinFlow-Z-Image checkpoints

Browse Source
This commit is contained in:
azazeal04 2026-04-04 15:57:19 +02:00 committed by GitHub
parent 21ed4a2242
commit 84c983a779
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/model_detection.py
View File

@ -44,6 +44,12 @@ def calculate_transformer_depth(prefix, state_dict_keys, state_dict):
def detect_unet_config(state_dict, key_prefix, metadata=None): def detect_unet_config(state_dict, key_prefix, metadata=None):
state_dict_keys = list(state_dict.keys()) state_dict_keys = list(state_dict.keys())
# TwinFlow-Z-Image: detect dual timestep embedder checkpoints first.
if any(k.startswith('{}t_embedder_2.'.format(key_prefix)) for k in state_dict_keys):
return {
"image_model": "twinflow_z_image",
"architecture": "TwinFlow_Z_Image",
}
if '{}joint_blocks.0.context_block.attn.qkv.weight'.format(key_prefix) in state_dict_keys: #mmdit model if '{}joint_blocks.0.context_block.attn.qkv.weight'.format(key_prefix) in state_dict_keys: #mmdit model
unet_config = {} unet_config = {}
unet_config["in_channels"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[1] unet_config["in_channels"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[1]