mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-29 04:17:14 +08:00
feat: Support Pixal3d and TRELLIS2 (CORE-278) (CORE-199) (CORE-236) (CORE-312) (#14718)
This commit is contained in:
@@ -120,6 +120,27 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
unet_config['block_repeat'] = [[1, 1, 1, 1], [2, 2, 2, 2]]
|
||||
return unet_config
|
||||
|
||||
shape_key = '{}img2shape.t_embedder.mlp.0.weight'.format(key_prefix)
|
||||
tex_key = '{}shape2txt.t_embedder.mlp.0.weight'.format(key_prefix)
|
||||
if shape_key in state_dict_keys or tex_key in state_dict_keys: # trellis2 / pixal3d
|
||||
has_shape = shape_key in state_dict_keys
|
||||
has_tex = tex_key in state_dict_keys
|
||||
unet_config = {
|
||||
"image_model": "trellis2",
|
||||
"resolution": 32 if (metadata or {}).get("is_512") else 64,
|
||||
"init_txt_model": has_tex,
|
||||
"txt_only": has_tex and not has_shape,
|
||||
}
|
||||
# Per-submodel projection head (Pixal3D adds `proj_linear`; Trellis2 doesn't).
|
||||
for sub, name, proj_in_channels in (("img2shape", "shape", 2048),
|
||||
("shape2txt", "texture", 2048),
|
||||
("structure_model", "structure", 1024)):
|
||||
key = '{}{}.blocks.0.cross_attn.proj_linear.weight'.format(key_prefix, sub)
|
||||
if key in state_dict_keys:
|
||||
unet_config["image_attn_mode_{}".format(name)] = "proj"
|
||||
unet_config["proj_in_channels_{}".format(name)] = proj_in_channels
|
||||
return unet_config
|
||||
|
||||
if '{}transformer.rotary_pos_emb.inv_freq'.format(key_prefix) in state_dict_keys: #stable audio dit
|
||||
unet_config = {}
|
||||
unet_config["audio_model"] = "dit1.0"
|
||||
|
||||
Reference in New Issue
Block a user