full LightVAE: lite encoder (remote-code class) + lite decoder; index -> AutoModel

modular_model_index.json

@@ -46,11 +46,11 @@
   ],
   "vae": [
     "diffusers",
-    "AutoencoderKLMiniMaxH3",
+    "AutoModel",
     {
       "type_hint": [
         "diffusers",
-        "AutoencoderKLMiniMaxH3"
+        "AutoModel"
       ],
       "pretrained_model_name_or_path": "diffusers-modular/MiniMax-H3-LightVAE",
       "subfolder": "vae",

vae/config.json

@@ -1,5 +1,5 @@
 {
-  "_class_name": "AutoencoderKLMiniMaxH3",
+  "_class_name": "AutoencoderKLMiniMaxH3LightVAE",
   "_diffusers_version": "0.40.0.dev0",
   "block_out_channels": [
     128,
@@ -93,5 +93,8 @@
     1,
     1
   ],
-  "token_drop": 3
-}
+  "token_drop": 3,
+  "auto_map": {
+    "AutoModel": "modeling_minimax_h3_lightvae.AutoencoderKLMiniMaxH3LightVAE"
+  }
+}

vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f1c69d47b24b0d5f8bbfba28bb04e93b4c16a3db1c4928e8b94dd7c4bfa3bf1d
-size 3969719952
+oid sha256:bccba90f8dfd77cd446e965d70596ab0d52d60a620d5867ddcbdbb7d4f857f6e
+size 3654687872

vae/modeling_minimax_h3_lightvae.py

@@ -0,0 +1,171 @@
+"""Remote-code VAE for the LightVAE build of MiniMax-H3: stock-shaped H3 video VAE whose encoder is replaced by
+lightx2v's block-pruned Lite encoder and whose ViT decoder is pruned 36 -> 12 blocks.
+
+`AutoencoderKLMiniMaxH3LightVAE` subclasses the merged `AutoencoderKLMiniMaxH3`, so it is a drop-in `vae` component
+(the latent interface is identical). The decoder depth is set by `decoder_num_layers=12` in the config; the encoder is
+rebuilt here because lightx2v's Lite encoder is structurally pruned (block 5's first resnet is a bare 1x1 channel
+projection), which a scalar `layers_per_block` cannot express.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from diffusers import AutoencoderKLMiniMaxH3
+from diffusers.models.autoencoders.autoencoder_kl_minimax_h3 import (
+    MiniMaxH3VideoCausalConv3d,
+    MiniMaxH3VideoDownBlock3d,
+    MiniMaxH3VideoEncoder3d,
+    MiniMaxH3VideoGroupNorm,
+    MiniMaxH3VideoResnetBlock3d,
+)
+
+
+class _ShortcutResnet(nn.Module):
+    """A resnet pruned to its 1x1 channel-projection shortcut alone; forward is `conv_shortcut(x)`."""
+
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.conv_shortcut = MiniMaxH3VideoCausalConv3d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.conv_shortcut(hidden_states)
+
+
+class _CustomDownBlock(nn.Module):
+    """A down block with an explicit resnet list and no downsampler."""
+
+    def __init__(self, resnets: list[nn.Module]) -> None:
+        super().__init__()
+        self.resnets = nn.ModuleList(resnets)
+        self.downsamplers = None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states)
+        return hidden_states
+
+
+class _LiteEncoder3d(MiniMaxH3VideoEncoder3d):
+    """lightx2v Lite H3 encoder. Blocks 0-4 are standard single-resnet down blocks (0-3 keep their downsampler);
+    block 5 is a projection-only resnet (512 -> 1024) followed by one full resnet (1024 -> 1024), no downsampler."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        block_out_channels,
+        spatial_downsample_factors,
+        temporal_downsample_factors,
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-6,
+        spatial_padding_mode: str = "reflect",
+    ) -> None:
+        nn.Module.__init__(self)
+        self.conv_in = MiniMaxH3VideoCausalConv3d(
+            in_channels, block_out_channels[0], kernel_size=3,
+            spatial_padding=1, temporal_padding=2, spatial_padding_mode=spatial_padding_mode,
+        )
+        block_in_channels = (block_out_channels[0],) + tuple(block_out_channels[:-1])
+        down_blocks = []
+        for i in range(len(block_out_channels) - 1):
+            down_blocks.append(
+                MiniMaxH3VideoDownBlock3d(
+                    in_channels=block_in_channels[i],
+                    out_channels=block_out_channels[i],
+                    num_layers=1,
+                    temporal_downsample_factor=temporal_downsample_factors[i],
+                    spatial_downsample_factor=spatial_downsample_factors[i],
+                    norm_num_groups=norm_num_groups,
+                    norm_eps=norm_eps,
+                    spatial_padding_mode=spatial_padding_mode,
+                )
+            )
+        down_blocks.append(
+            _CustomDownBlock(
+                [
+                    _ShortcutResnet(block_in_channels[-1], block_out_channels[-1]),
+                    MiniMaxH3VideoResnetBlock3d(
+                        in_channels=block_out_channels[-1],
+                        out_channels=block_out_channels[-1],
+                        norm_num_groups=norm_num_groups,
+                        norm_eps=norm_eps,
+                        spatial_padding_mode=spatial_padding_mode,
+                    ),
+                ]
+            )
+        )
+        self.down_blocks = nn.ModuleList(down_blocks)
+        self.norm_out = MiniMaxH3VideoGroupNorm(norm_num_groups, block_out_channels[-1], eps=norm_eps, affine=True)
+        self.conv_out = MiniMaxH3VideoCausalConv3d(
+            block_out_channels[-1], out_channels, kernel_size=3,
+            spatial_padding=1, temporal_padding=2, spatial_padding_mode=spatial_padding_mode,
+        )
+
+
+class AutoencoderKLMiniMaxH3LightVAE(AutoencoderKLMiniMaxH3):
+    """Stock H3 VAE with the Lite encoder swapped in. Decoder depth comes from `decoder_num_layers` in the config.
+
+    The parameter list is declared explicitly (not ``**kwargs``) so diffusers' ``from_config`` introspection forwards
+    every config key to ``super().__init__`` (which is ``@register_to_config``); the Lite encoder is then built in place.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        latent_channels: int = 24,
+        block_out_channels=(128, 256, 256, 512, 512, 1024),
+        layers_per_block: int = 2,
+        spatial_downsample_factors=(2, 2, 2, 2, 1, 1),
+        temporal_downsample_factors=(1, 2, 2, 1, 1, 1),
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-6,
+        spatial_padding_mode: str = "reflect",
+        decoder_num_layers: int = 36,
+        decoder_num_attention_heads: int = 32,
+        decoder_attention_head_dim: int = 64,
+        decoder_num_register_tokens: int = 4,
+        decoder_ffn_mult: int = 4,
+        decoder_rope_theta: float = 100.0,
+        decoder_rope_dim_ratio: float = 0.75,
+        decoder_norm_eps: float = 1e-5,
+        clip_length: int = 17,
+        token_drop: int = 3,
+        latents_mean=(0.0,) * 24,
+        latents_std=(1.0,) * 24,
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            latent_channels=latent_channels,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            spatial_downsample_factors=spatial_downsample_factors,
+            temporal_downsample_factors=temporal_downsample_factors,
+            norm_num_groups=norm_num_groups,
+            norm_eps=norm_eps,
+            spatial_padding_mode=spatial_padding_mode,
+            decoder_num_layers=decoder_num_layers,
+            decoder_num_attention_heads=decoder_num_attention_heads,
+            decoder_attention_head_dim=decoder_attention_head_dim,
+            decoder_num_register_tokens=decoder_num_register_tokens,
+            decoder_ffn_mult=decoder_ffn_mult,
+            decoder_rope_theta=decoder_rope_theta,
+            decoder_rope_dim_ratio=decoder_rope_dim_ratio,
+            decoder_norm_eps=decoder_norm_eps,
+            clip_length=clip_length,
+            token_drop=token_drop,
+            latents_mean=latents_mean,
+            latents_std=latents_std,
+        )
+        self.encoder = _LiteEncoder3d(
+            in_channels=in_channels,
+            out_channels=2 * latent_channels,
+            block_out_channels=block_out_channels,
+            spatial_downsample_factors=spatial_downsample_factors,
+            temporal_downsample_factors=temporal_downsample_factors,
+            norm_num_groups=norm_num_groups,
+            norm_eps=norm_eps,
+            spatial_padding_mode=spatial_padding_mode,
+        )