File size: 7,325 Bytes
fc62d8e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """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,
)
|