diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index a07657159d36..becba32a5b59 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -50,6 +50,7 @@ convert_ltx_transformer_checkpoint_to_diffusers, convert_ltx_vae_checkpoint_to_diffusers, convert_lumina2_to_diffusers, + convert_minimax_h3_transformer_checkpoint_to_diffusers, convert_mochi_transformer_checkpoint_to_diffusers, convert_sana_transformer_to_diffusers, convert_sd3_transformer_checkpoint_to_diffusers, @@ -220,6 +221,10 @@ "checkpoint_mapping_fn": lambda checkpoint, **kwargs: checkpoint, "default_subfolder": "transformer", }, + "MiniMaxH3Transformer3DModel": { + "checkpoint_mapping_fn": convert_minimax_h3_transformer_checkpoint_to_diffusers, + "default_subfolder": "transformer", + }, } diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index ac707947d491..38a324e7cc8a 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -127,6 +127,7 @@ ], "z-image-turbo-controlnet": "control_all_x_embedder.2-1.weight", "z-image-turbo-controlnet-2.x": "control_layers.14.adaLN_modulation.0.weight", + "minimax-h3": "token_refiner.blocks.0.attn.qkv_proj.weight", "sana": [ "blocks.0.cross_attn.q_linear.weight", "blocks.0.cross_attn.q_linear.bias", @@ -242,6 +243,7 @@ "z-image-turbo-controlnet-2.0": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0"}, "z-image-turbo-controlnet-2.1": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.1"}, "ltx2-dev": {"pretrained_model_name_or_path": "Lightricks/LTX-2"}, + "minimax-h3": {"pretrained_model_name_or_path": "MiniMaxAI/MiniMax-H3"}, } # Use to configure model sample size when original config is provided @@ -824,6 +826,9 @@ def infer_diffusers_model_type(checkpoint): elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["ltx2"]): model_type = "ltx2-dev" + elif CHECKPOINT_KEY_NAMES["minimax-h3"] in checkpoint: + model_type = "minimax-h3" + else: model_type = "v1" @@ -4222,3 +4227,88 @@ def convert_ernie_image_transformer_checkpoint_to_diffusers(checkpoint, **kwargs checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) return checkpoint + + +def convert_minimax_h3_transformer_checkpoint_to_diffusers(checkpoint, config, qkv_layout="stacked", **kwargs): + if qkv_layout not in ("stacked", "interleaved"): + raise ValueError( + f'`qkv_layout` must be "stacked" or "interleaved", got {qkv_layout!r}. "stacked" is the `[q; k; v]` row ' + "order of every published single-file MiniMax-H3 checkpoint (Comfy-Org/MiniMax-H3 and the GGUFs derived " + 'from it). "interleaved" is the per-head `[q k v]` row order of the MiniMaxAI/MiniMax-H3 shards, for a ' + "file merged from those shards by hand. The two cannot be told apart from the checkpoint itself." + ) + if "adaln_t_table" in checkpoint: + raise ValueError( + "This is a pruned MiniMax-H3 checkpoint: it replaces `time_embedder` with `adaln_t_table` of shape " + f"{tuple(checkpoint['adaln_t_table'].shape)}, which `MiniMaxH3Transformer3DModel` does not support. " + "Use the unpruned checkpoint from https://huggingface.co/MiniMaxAI/MiniMax-H3." + ) + + MINIMAX_H3_KEYS_RENAME_DICT = { + "token_refiner.blocks.": "token_refiner.refiner_blocks.", + "time_embedder.proj_in.": "time_embedder.linear_1.", + "time_embedder.proj_out.": "time_embedder.linear_2.", + "video_patch_proj.": "proj_in.", + "audio_patch_proj.": "audio_proj_in.", + "condition_proj.": "context_embedder.", + "final_layer.norm.": "norm_out.norm.", + "final_layer.adaln_proj.linear.": "norm_out.linear.", + "final_layer.video_out.": "proj_out.", + "final_layer.audio_out.": "audio_proj_out.", + ".attn.q_norm.": ".attn.norm_q.", + ".attn.k_norm.": ".attn.norm_k.", + ".attn.out_proj.": ".attn.to_out.0.", + ".mlp.fc1.": ".ff.net.0.proj.", + ".mlp.fc2.": ".ff.net.2.", + } + + def convert_minimax_h3_fused_attention(key: str, state_dict: dict[str, object]) -> None: + # Published single files store the reference model's post-load `[q; k; v]` stack; the MiniMaxAI/MiniMax-H3 + # shards interleave the rows per head, `[head0: q k v, head1: q k v, ...]`. Keys and shapes are identical, so + # the caller has to say which one it is. + fused_qkv_weight = state_dict.pop(key) + if qkv_layout == "interleaved": + fused_qkv_weight = fused_qkv_weight.unflatten( + 0, (config["num_attention_heads"], 3, config["attention_head_dim"]) + ) + to_q_weight, to_k_weight, to_v_weight = [weight.flatten(0, 1) for weight in fused_qkv_weight.unbind(dim=1)] + else: + to_q_weight, to_k_weight, to_v_weight = torch.chunk(fused_qkv_weight, 3, dim=0) + state_dict[key.replace(".attn.qkv_proj.weight", ".attn.to_q.weight")] = to_q_weight + state_dict[key.replace(".attn.qkv_proj.weight", ".attn.to_k.weight")] = to_k_weight + state_dict[key.replace(".attn.qkv_proj.weight", ".attn.to_v.weight")] = to_v_weight + + def convert_minimax_h3_gated_ff(key: str, state_dict: dict[str, object]) -> None: + # The checkpoint fuses `[gate; value]`, `SwiGLU` reads `[value; gate]`. + gate, value = torch.chunk(state_dict[key], 2, dim=0) + state_dict[key] = torch.cat([value, gate], dim=0) + + TRANSFORMER_SPECIAL_KEYS_REMAP = { + ".attn.qkv_proj.weight": convert_minimax_h3_fused_attention, + ".ff.net.0.proj.weight": convert_minimax_h3_gated_ff, + } + + def update_state_dict(state_dict: dict[str, object], old_key: str, new_key: str) -> None: + state_dict[new_key] = state_dict.pop(old_key) + + converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} + + # `MiniMaxH3RotaryPosEmbed` recomputes this buffer from the config. + converted_state_dict.pop("rope.inv_freq", None) + + for key in list(converted_state_dict.keys()): + new_key = key[:] + if new_key.startswith("blocks."): + new_key = new_key.replace("blocks.", "transformer_blocks.", 1) + for replace_key, rename_key in MINIMAX_H3_KEYS_RENAME_DICT.items(): + new_key = new_key.replace(replace_key, rename_key) + + update_state_dict(converted_state_dict, key, new_key) + + for key in list(converted_state_dict.keys()): + for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): + if special_key not in key: + continue + handler_fn_inplace(key, converted_state_dict) + + return converted_state_dict diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index f49cdaca2eb6..95ee1f2ac511 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -19,7 +19,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import BaseOutput, apply_lora_scale, logging from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward @@ -373,7 +373,9 @@ def forward( return hidden_states -class MiniMaxH3Transformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftAdapterMixin, CacheMixin): +class MiniMaxH3Transformer3DModel( + ModelMixin, ConfigMixin, AttentionMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin +): r""" A Transformer model for joint video + audio generation, introduced in MiniMax-H3. diff --git a/tests/models/transformers/test_models_transformer_minimax_h3.py b/tests/models/transformers/test_models_transformer_minimax_h3.py index 00baa37c84a0..4165f9b5392d 100644 --- a/tests/models/transformers/test_models_transformer_minimax_h3.py +++ b/tests/models/transformers/test_models_transformer_minimax_h3.py @@ -27,6 +27,7 @@ LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, + SingleFileTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -189,3 +190,19 @@ class TestMiniMaxH3TransformerContextParallel(MiniMaxH3TransformerTesterConfig, class TestMiniMaxH3TransformerLoRA(MiniMaxH3TransformerTesterConfig, LoraTesterMixin): """LoRA tests for the MiniMax-H3 transformer.""" + + +class TestMiniMaxH3TransformerSingleFile(MiniMaxH3TransformerTesterConfig, SingleFileTesterMixin): + @property + def ckpt_path(self): + return ( + "https://huggingface.co/Comfy-Org/MiniMax-H3/blob/main/diffusion_models/minimax_h3_fl2va_bf16.safetensors" + ) + + @property + def pretrained_model_name_or_path(self): + return "MiniMaxAI/MiniMax-H3" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"}