mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-27 11:22:35 +08:00
[feat]Add JoyImageEdit native model support (#14428)
This commit is contained in:
@@ -1058,6 +1058,25 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
dit_config["image_model"] = "SAM31"
|
||||
return dit_config
|
||||
|
||||
if (
|
||||
'{}double_blocks.0.attn.img_attn_qkv.weight'.format(key_prefix) in state_dict_keys
|
||||
and '{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix) in state_dict_keys
|
||||
and '{}condition_embedder.time_embedder.linear_1.weight'.format(key_prefix) in state_dict_keys
|
||||
and '{}img_in.weight'.format(key_prefix) in state_dict_keys
|
||||
and len(state_dict['{}img_in.weight'.format(key_prefix)].shape) == 5
|
||||
):
|
||||
img_in = state_dict['{}img_in.weight'.format(key_prefix)]
|
||||
head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0]
|
||||
return {
|
||||
"image_model": "joyimage",
|
||||
"in_channels": img_in.shape[1],
|
||||
"hidden_size": img_in.shape[0],
|
||||
"patch_size": list(img_in.shape[2:]),
|
||||
"num_layers": count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.'),
|
||||
"num_attention_heads": img_in.shape[0] // head_dim,
|
||||
"text_dim": 4096,
|
||||
}
|
||||
|
||||
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
|
||||
return None
|
||||
|
||||
|
||||
Reference in New Issue
Block a user