vae/modeling_minimax_h3_lightvae.py

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