feat: Support MiniMax-H3 (CORE-375) (#15224)

This commit is contained in:
Jukka Seppänen
2026-08-03 05:28:29 +03:00
committed by GitHub
parent b53e247c94
commit 57500fc5bc
16 changed files with 2599 additions and 6 deletions

View File

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