mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-26 02:42:36 +08:00
feat: Support MiniMax-H3 (CORE-375) (#15224)
This commit is contained in:
@@ -359,6 +359,35 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
# PixArt diffusers
|
||||
return None
|
||||
|
||||
if '{}video_patch_proj.weight'.format(key_prefix) in state_dict_keys and '{}audio_patch_proj.weight'.format(key_prefix) in state_dict_keys: # MiniMax H3
|
||||
dit_config = {}
|
||||
dit_config["image_model"] = "minimax_h3"
|
||||
dit_config["num_layers"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
|
||||
dit_config["token_refiner_num_layers"] = count_blocks(state_dict_keys, '{}token_refiner.blocks.'.format(key_prefix) + '{}.')
|
||||
dit_config["hidden_size"] = state_dict['{}video_patch_proj.weight'.format(key_prefix)].shape[0]
|
||||
dit_config["latents_dim"] = state_dict['{}final_layer.video_out.weight'.format(key_prefix)].shape[0] // 4 # patch 1x2x2
|
||||
dit_config["audio_latents_dim"] = state_dict['{}final_layer.audio_out.weight'.format(key_prefix)].shape[0]
|
||||
dit_config["attention_head_dim"] = state_dict['{}blocks.0.attn.q_norm.weight'.format(key_prefix)].shape[0]
|
||||
qkv = state_dict['{}blocks.0.attn.qkv_proj.weight'.format(key_prefix)]
|
||||
dit_config["num_attention_heads"] = qkv.shape[0] // (3 * dit_config["attention_head_dim"])
|
||||
dit_config["ffn_hidden_size"] = state_dict['{}blocks.0.mlp.fc1.weight'.format(key_prefix)].shape[0] // 2
|
||||
dit_config["text_dim"] = state_dict['{}condition_proj.weight'.format(key_prefix)].shape[1]
|
||||
table_key = '{}adaln_t_table'.format(key_prefix)
|
||||
if table_key in state_dict_keys:
|
||||
# adaln shipped over a precomputed curve basis: the adaln linears span a small shared basis of the time-embedding curve (no time embedder)
|
||||
table = state_dict[table_key].shape # [grid, k]
|
||||
dit_config["adaln_curve_grid"] = table[0]
|
||||
dit_config["time_embed_dim"] = table[1]
|
||||
else:
|
||||
te = state_dict['{}time_embedder.proj_in.weight'.format(key_prefix)]
|
||||
dit_config["timestep_input_dim"] = te.shape[1]
|
||||
dit_config["time_embed_hidden_size"] = te.shape[0]
|
||||
dit_config["time_embed_dim"] = state_dict['{}time_embedder.proj_out.weight'.format(key_prefix)].shape[0]
|
||||
dit_config["rope_inv_freq_len"] = state_dict['{}rope.inv_freq'.format(key_prefix)].shape[0]
|
||||
if metadata is not None and "config" in metadata:
|
||||
dit_config.update(json.loads(metadata["config"]).get("transformer", {}))
|
||||
return dit_config
|
||||
|
||||
if '{}adaln_single.emb.timestep_embedder.linear_1.bias'.format(key_prefix) in state_dict_keys: #Lightricks ltxv
|
||||
dit_config = {}
|
||||
dit_config["image_model"] = "ltxav" if f'{key_prefix}audio_adaln_single.linear.weight' in state_dict_keys else "ltxv"
|
||||
|
||||
Reference in New Issue
Block a user