feat: Support Pixal3d and TRELLIS2 (CORE-278) (CORE-199) (CORE-236) (CORE-312) (#14718)

This commit is contained in:
Jukka Seppänen
2026-08-22 03:32:25 +03:00
committed by GitHub
parent 28864ca8fe
commit 0e65cb9071
28 changed files with 12397 additions and 83 deletions

View File

@@ -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"