"""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, )