Config, modular index and the pruned transformer class

modular_model_index.json

@@ -0,0 +1,131 @@
+{
+  "_class_name": "MiniMaxH3ModularPipeline",
+  "_diffusers_version": "0.40.0.dev0",
+  "_blocks_class_name": "MiniMaxH3Blocks",
+  "text_encoder": [
+    "transformers",
+    "Qwen3VLForConditionalGeneration",
+    {
+      "type_hint": [
+        "transformers",
+        "Qwen3VLForConditionalGeneration"
+      ],
+      "pretrained_model_name_or_path": "MiniMaxAI/MiniMax-H3",
+      "subfolder": "text_encoder",
+      "variant": null,
+      "revision": null
+    }
+  ],
+  "tokenizer": [
+    "transformers",
+    "Qwen2TokenizerFast",
+    {
+      "type_hint": [
+        "transformers",
+        "Qwen2TokenizerFast"
+      ],
+      "pretrained_model_name_or_path": "MiniMaxAI/MiniMax-H3",
+      "subfolder": "tokenizer",
+      "variant": null,
+      "revision": null
+    }
+  ],
+  "processor": [
+    "transformers",
+    "Qwen3VLProcessor",
+    {
+      "type_hint": [
+        "transformers",
+        "Qwen3VLProcessor"
+      ],
+      "pretrained_model_name_or_path": "MiniMaxAI/MiniMax-H3",
+      "subfolder": "processor",
+      "variant": null,
+      "revision": null
+    }
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLMiniMaxH3",
+    {
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLMiniMaxH3"
+      ],
+      "pretrained_model_name_or_path": "MiniMaxAI/MiniMax-H3",
+      "subfolder": "vae",
+      "variant": null,
+      "revision": null
+    }
+  ],
+  "audio_vae": [
+    "diffusers",
+    "AutoencoderKLMiniMaxH3Audio",
+    {
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLMiniMaxH3Audio"
+      ],
+      "pretrained_model_name_or_path": "MiniMaxAI/MiniMax-H3",
+      "subfolder": "audio_vae",
+      "variant": null,
+      "revision": null
+    }
+  ],
+  "transformer": [
+    "diffusers",
+    "AutoModel",
+    {
+      "type_hint": [
+        "diffusers",
+        "AutoModel"
+      ],
+      "pretrained_model_name_or_path": "multimodalart/MiniMax-H3-Pruned",
+      "subfolder": "transformer",
+      "variant": null,
+      "revision": null
+    }
+  ],
+  "transformer_ref": [
+    "diffusers",
+    "AutoModel",
+    {
+      "type_hint": [
+        "diffusers",
+        "AutoModel"
+      ],
+      "pretrained_model_name_or_path": "multimodalart/MiniMax-H3-Pruned",
+      "subfolder": "transformer_ref",
+      "variant": null,
+      "revision": null
+    }
+  ],
+  "scheduler": [
+    "diffusers",
+    "MiniMaxH3Scheduler",
+    {
+      "type_hint": [
+        "diffusers",
+        "MiniMaxH3Scheduler"
+      ],
+      "pretrained_model_name_or_path": "MiniMaxAI/MiniMax-H3",
+      "subfolder": "scheduler",
+      "variant": null,
+      "revision": null
+    }
+  ],
+  "audio_scheduler": [
+    "diffusers",
+    "MiniMaxH3Scheduler",
+    {
+      "type_hint": [
+        "diffusers",
+        "MiniMaxH3Scheduler"
+      ],
+      "pretrained_model_name_or_path": "MiniMaxAI/MiniMax-H3",
+      "subfolder": "audio_scheduler",
+      "variant": null,
+      "revision": null
+    }
+  ]
+}

transformer/config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "MiniMaxH3PrunedTransformer3DModel",
+  "_diffusers_version": "0.40.0.dev0",
+  "adaln_rank": 8,
+  "attention_head_dim": 128,
+  "audio_in_channels": 32,
+  "auto_map": {
+    "AutoModel": "modeling_minimax_h3_pruned.MiniMaxH3PrunedTransformer3DModel"
+  },
+  "ffn_dim": 14336,
+  "final_norm_eps": 1e-05,
+  "freq_dim": 256,
+  "hidden_size": 5376,
+  "in_channels": 24,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 56,
+  "num_layers": 50,
+  "num_refiner_layers": 2,
+  "patch_size": [
+    1,
+    2,
+    2
+  ],
+  "qk_norm_eps": 1e-05,
+  "rope_freq_dim": 16,
+  "rope_theta": 10000.0,
+  "text_dim": 5120,
+  "time_embed_dim": 2688,
+  "time_embed_hidden_dim": 5376,
+  "time_table_size": 1025
+}

transformer/modeling_minimax_h3_pruned.py

@@ -0,0 +1,243 @@
+# Copyright 2025 The MiniMax Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""AdaLN-pruned MiniMax-H3 transformer.
+
+Everything outside the timestep path is inherited from `MiniMaxH3Transformer3DModel`: the attention, the blocks, the
+token refiner, the output heads and `forward` itself are the released implementation, unmodified. Only what feeds the
+AdaLN projections changes.
+"""
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import register_to_config
+from diffusers.models.modeling_utils import get_parameter_dtype
+from diffusers.models.transformers.transformer_minimax_h3 import (
+    MINIMAX_H3_MODALITY_NUM,
+    MiniMaxH3RotaryPosEmbed,
+    MiniMaxH3TokenRefiner,
+    MiniMaxH3Transformer3DModel,
+    MiniMaxH3TransformerBlock,
+)
+
+
+class MiniMaxH3PrunedTimeEmbedder(nn.Module):
+    r"""The released timestep MLP, replaced by an interpolated table of AdaLN coordinates.
+
+    Every AdaLN projection in the released model consumes `silu(time_embedder(time_proj(t)))`, which depends on the
+    scalar timestep alone: over `t` in `[0, 1]` it traces a one-dimensional curve in `R^{time_embed_dim}`. A rank-8
+    affine subspace reproduces that curve to about 1.5e-5 relative RMS, so only the curve's coordinates in that
+    subspace are stored - sampled on a uniform grid of `table_size` timesteps and linearly interpolated in between.
+    The subspace offset is folded into the AdaLN biases and its basis into the AdaLN weights, which is why the
+    projections take an `adaln_rank`-wide input here instead of `time_embed_dim`.
+
+    The module stands in for `time_proj` and `time_embedder` together: it consumes the raw timestep, so the released
+    `forward` needs no change once `time_proj` is an identity.
+    """
+
+    def __init__(self, table_size: int = 1025, adaln_rank: int = 8) -> None:
+        super().__init__()
+        self.register_buffer("table", torch.zeros(table_size, adaln_rank), persistent=True)
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        table = self.table
+        steps = table.shape[0] - 1
+        position = timestep.to(table.dtype).flatten().clamp(0.0, 1.0) * steps
+        lower = position.floor().clamp(max=steps - 1).long()
+        weight = (position - lower).unsqueeze(-1)
+        return torch.lerp(table.index_select(0, lower), table.index_select(0, lower + 1), weight)
+
+
+class MiniMaxH3PrunedTimeProj(nn.Module):
+    r"""Identity stand-in for `Timesteps`: the pruned time embedder indexes the raw timestep."""
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        return timestep
+
+
+class MiniMaxH3PrunedAdaLayerNormModulation(nn.Module):
+    r"""`MiniMaxH3AdaLayerNormModulation` over the pruned timestep coordinates.
+
+    Two differences from the released module. It applies no `silu` - the table already holds the coordinates of the
+    activated curve. And the folded bias is a float32 buffer applied outside the projection rather than the
+    projection's own bias: it carries almost the entire modulation (the coordinate term contributes a few tenths of
+    it), so storing it in bfloat16 would put a full output-scale rounding step into every evaluation. Kept in
+    float32 it costs 0.4 MB per block and leaves the pruned AdaLN function closer to an exact float64 evaluation
+    than the released bfloat16 checkpoint's own arithmetic is.
+
+    `linear` stays a bias-free `nn.Linear` so PEFT wraps it exactly as it wraps the released projection.
+    """
+
+    def __init__(self, adaln_rank: int, hidden_size: int) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        out_features = 6 * hidden_size * MINIMAX_H3_MODALITY_NUM
+        self.linear = nn.Linear(adaln_rank, out_features, bias=False)
+        self.register_buffer("folded_bias", torch.zeros(out_features), persistent=True)
+
+    def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        dtype = get_parameter_dtype(self.linear)
+        temb = self.linear(temb.to(dtype))
+        temb = (temb.float() + self.folded_bias).to(dtype)
+        temb = temb.view(-1, 6 * self.hidden_size)
+        return temb.chunk(6, dim=-1)
+
+
+class MiniMaxH3PrunedAdaLayerNormOut(nn.Module):
+    r"""`MiniMaxH3AdaLayerNormOut` over the pruned timestep coordinates; see the modulation module above."""
+
+    def __init__(self, hidden_size: int, adaln_rank: int, eps: float) -> None:
+        super().__init__()
+        self.norm = nn.RMSNorm(hidden_size, eps=eps)
+        self.linear = nn.Linear(adaln_rank, 2 * hidden_size, bias=False)
+        self.register_buffer("folded_bias", torch.zeros(2 * hidden_size), persistent=True)
+
+    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, timestep_indices: torch.Tensor) -> torch.Tensor:
+        dtype = get_parameter_dtype(self.linear)
+        temb = self.linear(temb.to(dtype))
+        shift, scale = (temb.float() + self.folded_bias).to(dtype).chunk(2, dim=-1)
+        hidden_states = self.norm(hidden_states)
+        return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select(
+            0, timestep_indices
+        )
+
+
+class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
+    r"""MiniMax-H3's DiT with the AdaLN input projections reduced to their reachable rank.
+
+    The released checkpoint spends 13.03B of its 33.14B parameters on the 50 per-block `adaln_proj.linear` matrices
+    plus `norm_out.linear`, all of which read the same 2688-wide timestep embedding. Because that embedding is a
+    function of the scalar timestep, its reachable set is a curve an 8-dimensional affine subspace covers to ~1.5e-5
+    relative RMS - far below one bfloat16 rounding step of the weights themselves. Folding the subspace into the
+    projections leaves an 8-wide input and removes 26 GB per partition.
+
+    Only what builds the timestep path differs from [`MiniMaxH3Transformer3DModel`]: `time_proj` becomes an identity,
+    `time_embedder` becomes [`MiniMaxH3PrunedTimeEmbedder`], and the AdaLN projections take `adaln_rank` inputs.
+    `forward` is inherited unchanged. The module names are the released ones, so LoRAs trained against a pruned
+    checkpoint - what the common trainers use by default - load natively.
+
+    Args:
+        adaln_rank (`int`, defaults to `8`):
+            The width of the timestep coordinates every AdaLN projection consumes.
+        time_table_size (`int`, defaults to `1025`):
+            The number of uniformly spaced timesteps the coordinate table holds; values in between are interpolated
+            linearly.
+
+        Every other argument is [`MiniMaxH3Transformer3DModel`]'s and carries the same meaning. `freq_dim` and
+        `time_embed_hidden_dim` are kept in the config, unused, so a pruned config still records the shape of the
+        released timestep MLP it was folded from.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["MiniMaxH3TransformerBlock", "MiniMaxH3TokenRefinerBlock", "MiniMaxH3PrunedAdaLayerNormOut"]
+    _repeated_blocks = ["MiniMaxH3TransformerBlock", "MiniMaxH3TokenRefinerBlock"]
+    _skip_layerwise_casting_patterns = ["norm"]
+    # The released checkpoint's mixed-precision split - patch projections, output heads and the timestep path in
+    # float32, the block stack in bfloat16 - plus the folded AdaLN biases, for the reason given on the modulation
+    # module. Entries are matched against the dot-separated segments of each parameter name.
+    _keep_in_fp32_modules = [
+        "proj_in",
+        "audio_proj_in",
+        "time_embedder",
+        "proj_out",
+        "audio_proj_out",
+        "rope",
+        "folded_bias",
+    ]
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 56,
+        attention_head_dim: int = 128,
+        hidden_size: int = 5376,
+        num_layers: int = 50,
+        num_refiner_layers: int = 2,
+        ffn_dim: int = 14336,
+        in_channels: int = 24,
+        audio_in_channels: int = 32,
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        text_dim: int = 5120,
+        freq_dim: int = 256,
+        time_embed_hidden_dim: int = 5376,
+        time_embed_dim: int = 2688,
+        rope_freq_dim: int = 16,
+        rope_theta: float = 10000.0,
+        norm_eps: float = 1e-5,
+        qk_norm_eps: float = 1e-5,
+        final_norm_eps: float = 1e-5,
+        adaln_rank: int = 8,
+        time_table_size: int = 1025,
+    ) -> None:
+        # `MiniMaxH3Transformer3DModel.__init__` is itself wrapped by `register_to_config`, so calling it would
+        # register the released config over this one - and would allocate the 26 GB of AdaLN projections this class
+        # exists to avoid. The module tree is built here instead; everything but the timestep path is verbatim.
+        nn.Module.__init__(self)
+
+        video_patch_dim = in_channels * patch_size[0] * patch_size[1] * patch_size[2]
+
+        # 1. Per-modality input projections
+        self.proj_in = nn.Linear(video_patch_dim, hidden_size, bias=True)
+        self.audio_proj_in = nn.Linear(audio_in_channels, hidden_size, bias=True)
+        self.context_embedder = nn.Linear(text_dim, hidden_size, bias=True)
+
+        # 2. Timestep coordinates, shared by every AdaLN projection
+        self.time_proj = MiniMaxH3PrunedTimeProj()
+        self.time_embedder = MiniMaxH3PrunedTimeEmbedder(table_size=time_table_size, adaln_rank=adaln_rank)
+
+        # 3. Rotary embedding over the packed (t, h, w) grid
+        self.rope = MiniMaxH3RotaryPosEmbed(rope_freq_dim=rope_freq_dim, rope_theta=rope_theta)
+
+        # 4. Text stream refiner
+        self.token_refiner = MiniMaxH3TokenRefiner(
+            hidden_size=hidden_size,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            ffn_dim=ffn_dim,
+            num_layers=num_refiner_layers,
+            norm_eps=norm_eps,
+            qk_norm_eps=qk_norm_eps,
+            final_norm_eps=final_norm_eps,
+        )
+
+        # 5. The block stack, with each block's AdaLN projection narrowed to the timestep coordinates. The block is
+        #    built with `time_embed_dim=adaln_rank` so its own projection is already the right shape, then swapped
+        #    for the pruned module, which drops the `silu` and moves the bias to float32.
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MiniMaxH3TransformerBlock(
+                    hidden_size=hidden_size,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    ffn_dim=ffn_dim,
+                    time_embed_dim=adaln_rank,
+                    norm_eps=norm_eps,
+                    qk_norm_eps=qk_norm_eps,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        for block in self.transformer_blocks:
+            block.adaln_proj = MiniMaxH3PrunedAdaLayerNormModulation(adaln_rank=adaln_rank, hidden_size=hidden_size)
+
+        # 6. Shared output norm and the two per-modality output heads
+        self.norm_out = MiniMaxH3PrunedAdaLayerNormOut(
+            hidden_size=hidden_size, adaln_rank=adaln_rank, eps=final_norm_eps
+        )
+        self.proj_out = nn.Linear(hidden_size, video_patch_dim, bias=True)
+        self.audio_proj_out = nn.Linear(hidden_size, audio_in_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+
+MiniMaxH3PrunedTransformer3DModel.register_for_auto_class("AutoModel")

transformer_ref/config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "MiniMaxH3PrunedTransformer3DModel",
+  "_diffusers_version": "0.40.0.dev0",
+  "adaln_rank": 8,
+  "attention_head_dim": 128,
+  "audio_in_channels": 32,
+  "auto_map": {
+    "AutoModel": "modeling_minimax_h3_pruned.MiniMaxH3PrunedTransformer3DModel"
+  },
+  "ffn_dim": 14336,
+  "final_norm_eps": 1e-05,
+  "freq_dim": 256,
+  "hidden_size": 5376,
+  "in_channels": 24,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 56,
+  "num_layers": 50,
+  "num_refiner_layers": 2,
+  "patch_size": [
+    1,
+    2,
+    2
+  ],
+  "qk_norm_eps": 1e-05,
+  "rope_freq_dim": 16,
+  "rope_theta": 10000.0,
+  "text_dim": 5120,
+  "time_embed_dim": 2688,
+  "time_embed_hidden_dim": 5376,
+  "time_table_size": 1025
+}

transformer_ref/modeling_minimax_h3_pruned.py

@@ -0,0 +1,243 @@
+# Copyright 2025 The MiniMax Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""AdaLN-pruned MiniMax-H3 transformer.
+
+Everything outside the timestep path is inherited from `MiniMaxH3Transformer3DModel`: the attention, the blocks, the
+token refiner, the output heads and `forward` itself are the released implementation, unmodified. Only what feeds the
+AdaLN projections changes.
+"""
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import register_to_config
+from diffusers.models.modeling_utils import get_parameter_dtype
+from diffusers.models.transformers.transformer_minimax_h3 import (
+    MINIMAX_H3_MODALITY_NUM,
+    MiniMaxH3RotaryPosEmbed,
+    MiniMaxH3TokenRefiner,
+    MiniMaxH3Transformer3DModel,
+    MiniMaxH3TransformerBlock,
+)
+
+
+class MiniMaxH3PrunedTimeEmbedder(nn.Module):
+    r"""The released timestep MLP, replaced by an interpolated table of AdaLN coordinates.
+
+    Every AdaLN projection in the released model consumes `silu(time_embedder(time_proj(t)))`, which depends on the
+    scalar timestep alone: over `t` in `[0, 1]` it traces a one-dimensional curve in `R^{time_embed_dim}`. A rank-8
+    affine subspace reproduces that curve to about 1.5e-5 relative RMS, so only the curve's coordinates in that
+    subspace are stored - sampled on a uniform grid of `table_size` timesteps and linearly interpolated in between.
+    The subspace offset is folded into the AdaLN biases and its basis into the AdaLN weights, which is why the
+    projections take an `adaln_rank`-wide input here instead of `time_embed_dim`.
+
+    The module stands in for `time_proj` and `time_embedder` together: it consumes the raw timestep, so the released
+    `forward` needs no change once `time_proj` is an identity.
+    """
+
+    def __init__(self, table_size: int = 1025, adaln_rank: int = 8) -> None:
+        super().__init__()
+        self.register_buffer("table", torch.zeros(table_size, adaln_rank), persistent=True)
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        table = self.table
+        steps = table.shape[0] - 1
+        position = timestep.to(table.dtype).flatten().clamp(0.0, 1.0) * steps
+        lower = position.floor().clamp(max=steps - 1).long()
+        weight = (position - lower).unsqueeze(-1)
+        return torch.lerp(table.index_select(0, lower), table.index_select(0, lower + 1), weight)
+
+
+class MiniMaxH3PrunedTimeProj(nn.Module):
+    r"""Identity stand-in for `Timesteps`: the pruned time embedder indexes the raw timestep."""
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        return timestep
+
+
+class MiniMaxH3PrunedAdaLayerNormModulation(nn.Module):
+    r"""`MiniMaxH3AdaLayerNormModulation` over the pruned timestep coordinates.
+
+    Two differences from the released module. It applies no `silu` - the table already holds the coordinates of the
+    activated curve. And the folded bias is a float32 buffer applied outside the projection rather than the
+    projection's own bias: it carries almost the entire modulation (the coordinate term contributes a few tenths of
+    it), so storing it in bfloat16 would put a full output-scale rounding step into every evaluation. Kept in
+    float32 it costs 0.4 MB per block and leaves the pruned AdaLN function closer to an exact float64 evaluation
+    than the released bfloat16 checkpoint's own arithmetic is.
+
+    `linear` stays a bias-free `nn.Linear` so PEFT wraps it exactly as it wraps the released projection.
+    """
+
+    def __init__(self, adaln_rank: int, hidden_size: int) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        out_features = 6 * hidden_size * MINIMAX_H3_MODALITY_NUM
+        self.linear = nn.Linear(adaln_rank, out_features, bias=False)
+        self.register_buffer("folded_bias", torch.zeros(out_features), persistent=True)
+
+    def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        dtype = get_parameter_dtype(self.linear)
+        temb = self.linear(temb.to(dtype))
+        temb = (temb.float() + self.folded_bias).to(dtype)
+        temb = temb.view(-1, 6 * self.hidden_size)
+        return temb.chunk(6, dim=-1)
+
+
+class MiniMaxH3PrunedAdaLayerNormOut(nn.Module):
+    r"""`MiniMaxH3AdaLayerNormOut` over the pruned timestep coordinates; see the modulation module above."""
+
+    def __init__(self, hidden_size: int, adaln_rank: int, eps: float) -> None:
+        super().__init__()
+        self.norm = nn.RMSNorm(hidden_size, eps=eps)
+        self.linear = nn.Linear(adaln_rank, 2 * hidden_size, bias=False)
+        self.register_buffer("folded_bias", torch.zeros(2 * hidden_size), persistent=True)
+
+    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, timestep_indices: torch.Tensor) -> torch.Tensor:
+        dtype = get_parameter_dtype(self.linear)
+        temb = self.linear(temb.to(dtype))
+        shift, scale = (temb.float() + self.folded_bias).to(dtype).chunk(2, dim=-1)
+        hidden_states = self.norm(hidden_states)
+        return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select(
+            0, timestep_indices
+        )
+
+
+class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
+    r"""MiniMax-H3's DiT with the AdaLN input projections reduced to their reachable rank.
+
+    The released checkpoint spends 13.03B of its 33.14B parameters on the 50 per-block `adaln_proj.linear` matrices
+    plus `norm_out.linear`, all of which read the same 2688-wide timestep embedding. Because that embedding is a
+    function of the scalar timestep, its reachable set is a curve an 8-dimensional affine subspace covers to ~1.5e-5
+    relative RMS - far below one bfloat16 rounding step of the weights themselves. Folding the subspace into the
+    projections leaves an 8-wide input and removes 26 GB per partition.
+
+    Only what builds the timestep path differs from [`MiniMaxH3Transformer3DModel`]: `time_proj` becomes an identity,
+    `time_embedder` becomes [`MiniMaxH3PrunedTimeEmbedder`], and the AdaLN projections take `adaln_rank` inputs.
+    `forward` is inherited unchanged. The module names are the released ones, so LoRAs trained against a pruned
+    checkpoint - what the common trainers use by default - load natively.
+
+    Args:
+        adaln_rank (`int`, defaults to `8`):
+            The width of the timestep coordinates every AdaLN projection consumes.
+        time_table_size (`int`, defaults to `1025`):
+            The number of uniformly spaced timesteps the coordinate table holds; values in between are interpolated
+            linearly.
+
+        Every other argument is [`MiniMaxH3Transformer3DModel`]'s and carries the same meaning. `freq_dim` and
+        `time_embed_hidden_dim` are kept in the config, unused, so a pruned config still records the shape of the
+        released timestep MLP it was folded from.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["MiniMaxH3TransformerBlock", "MiniMaxH3TokenRefinerBlock", "MiniMaxH3PrunedAdaLayerNormOut"]
+    _repeated_blocks = ["MiniMaxH3TransformerBlock", "MiniMaxH3TokenRefinerBlock"]
+    _skip_layerwise_casting_patterns = ["norm"]
+    # The released checkpoint's mixed-precision split - patch projections, output heads and the timestep path in
+    # float32, the block stack in bfloat16 - plus the folded AdaLN biases, for the reason given on the modulation
+    # module. Entries are matched against the dot-separated segments of each parameter name.
+    _keep_in_fp32_modules = [
+        "proj_in",
+        "audio_proj_in",
+        "time_embedder",
+        "proj_out",
+        "audio_proj_out",
+        "rope",
+        "folded_bias",
+    ]
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 56,
+        attention_head_dim: int = 128,
+        hidden_size: int = 5376,
+        num_layers: int = 50,
+        num_refiner_layers: int = 2,
+        ffn_dim: int = 14336,
+        in_channels: int = 24,
+        audio_in_channels: int = 32,
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        text_dim: int = 5120,
+        freq_dim: int = 256,
+        time_embed_hidden_dim: int = 5376,
+        time_embed_dim: int = 2688,
+        rope_freq_dim: int = 16,
+        rope_theta: float = 10000.0,
+        norm_eps: float = 1e-5,
+        qk_norm_eps: float = 1e-5,
+        final_norm_eps: float = 1e-5,
+        adaln_rank: int = 8,
+        time_table_size: int = 1025,
+    ) -> None:
+        # `MiniMaxH3Transformer3DModel.__init__` is itself wrapped by `register_to_config`, so calling it would
+        # register the released config over this one - and would allocate the 26 GB of AdaLN projections this class
+        # exists to avoid. The module tree is built here instead; everything but the timestep path is verbatim.
+        nn.Module.__init__(self)
+
+        video_patch_dim = in_channels * patch_size[0] * patch_size[1] * patch_size[2]
+
+        # 1. Per-modality input projections
+        self.proj_in = nn.Linear(video_patch_dim, hidden_size, bias=True)
+        self.audio_proj_in = nn.Linear(audio_in_channels, hidden_size, bias=True)
+        self.context_embedder = nn.Linear(text_dim, hidden_size, bias=True)
+
+        # 2. Timestep coordinates, shared by every AdaLN projection
+        self.time_proj = MiniMaxH3PrunedTimeProj()
+        self.time_embedder = MiniMaxH3PrunedTimeEmbedder(table_size=time_table_size, adaln_rank=adaln_rank)
+
+        # 3. Rotary embedding over the packed (t, h, w) grid
+        self.rope = MiniMaxH3RotaryPosEmbed(rope_freq_dim=rope_freq_dim, rope_theta=rope_theta)
+
+        # 4. Text stream refiner
+        self.token_refiner = MiniMaxH3TokenRefiner(
+            hidden_size=hidden_size,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            ffn_dim=ffn_dim,
+            num_layers=num_refiner_layers,
+            norm_eps=norm_eps,
+            qk_norm_eps=qk_norm_eps,
+            final_norm_eps=final_norm_eps,
+        )
+
+        # 5. The block stack, with each block's AdaLN projection narrowed to the timestep coordinates. The block is
+        #    built with `time_embed_dim=adaln_rank` so its own projection is already the right shape, then swapped
+        #    for the pruned module, which drops the `silu` and moves the bias to float32.
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MiniMaxH3TransformerBlock(
+                    hidden_size=hidden_size,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    ffn_dim=ffn_dim,
+                    time_embed_dim=adaln_rank,
+                    norm_eps=norm_eps,
+                    qk_norm_eps=qk_norm_eps,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        for block in self.transformer_blocks:
+            block.adaln_proj = MiniMaxH3PrunedAdaLayerNormModulation(adaln_rank=adaln_rank, hidden_size=hidden_size)
+
+        # 6. Shared output norm and the two per-modality output heads
+        self.norm_out = MiniMaxH3PrunedAdaLayerNormOut(
+            hidden_size=hidden_size, adaln_rank=adaln_rank, eps=final_norm_eps
+        )
+        self.proj_out = nn.Linear(hidden_size, video_patch_dim, bias=True)
+        self.audio_proj_out = nn.Linear(hidden_size, audio_in_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+
+MiniMaxH3PrunedTransformer3DModel.register_for_auto_class("AutoModel")