Ref2VA/video_vae/vae_cnn.py

# SPDX-License-Identifier: Apache-2.0
# 3D causal CNN encoder for the MiniMax H3 visual VAE (inference-only bundle).
import os
import torch.nn as nn
import torch.nn.functional as F

from .attention import maybe_checkpoint
from .conv import SpatialParallelConv3d
from .norm import get_spatial_norm_3d
from .parallel import get_parallel_state, exchange_strides
from .norm import get_group_norm_3d










# ============================================================================
# 3D CNN Components
# ============================================================================


def norm_silu(x, norm, cond=None):
    if cond is None:
        return F.silu(norm(x))
    else:
        return F.silu(norm(x, cond))


class Downsample3D(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        time_stride=1,
        space_stride=2,
        padding_mode="zeros",
        padding_mode_t=None,
        causal=True,
    ):
        super().__init__()
        self.time_stride = time_stride
        self.space_stride = space_stride

        assert time_stride in [1, 2]
        assert space_stride in [1, 2, 3]

        self.conv = SpatialParallelConv3d(
            in_channels,
            out_channels,
            kernel_size=3,
            padding=(1, 0, 0),
            stride=(time_stride, space_stride, space_stride),
            padding_mode=padding_mode,
            padding_mode_t=padding_mode_t,
            causal=causal,
        )
        self.causal = self.conv.causal
        self.pad_mode = self.conv.pad_mode

    def forward(self, x):
        if self.space_stride == 2:
            if getattr(self.conv, "spatial_parallel", False):
                state = get_parallel_state()
                x = exchange_strides(
                    x,
                    self.pad_mode,
                    state["sp_rank"],
                    state["sp_size"],
                    state["sp_process_group"],
                    self.conv.chunk_dim,
                )
            else:
                pad = (0, 1, 0, 1, 0, 0)
                x = F.pad(x, pad, mode=self.pad_mode)
        return self.conv(x)


class ResnetBlock3D(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels=None,
        zq_ch=None,
        padding_mode="zeros",
        padding_mode_t=None,
        causal=True,
        use_t_isolated_gn=False,
    ):
        super().__init__()
        self.in_channels = in_channels
        out_channels = in_channels if out_channels is None else out_channels
        self.out_channels = out_channels

        self.use_fused_norm = (
            os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true"
        )

        if zq_ch is None:
            self.norm1 = get_group_norm_3d(in_channels, use_t_isolated_gn=use_t_isolated_gn)
            self.norm2 = get_group_norm_3d(out_channels, use_t_isolated_gn=use_t_isolated_gn)
        else:
            self.norm1 = get_spatial_norm_3d(
                in_channels,
                zq_ch,
                padding_mode=padding_mode,
                padding_mode_t=padding_mode_t,
                causal=causal,
                use_t_isolated_gn=use_t_isolated_gn,
            )
            self.norm2 = get_spatial_norm_3d(
                out_channels,
                zq_ch,
                padding_mode=padding_mode,
                padding_mode_t=padding_mode_t,
                causal=causal,
                use_t_isolated_gn=use_t_isolated_gn,
            )

        self.conv1 = SpatialParallelConv3d(
            in_channels,
            out_channels,
            kernel_size=3,
            padding=1,
            padding_mode=padding_mode,
            padding_mode_t=padding_mode_t,
            causal=causal,
        )

        self.conv2 = SpatialParallelConv3d(
            out_channels,
            out_channels,
            kernel_size=3,
            padding=1,
            padding_mode=padding_mode,
            padding_mode_t=padding_mode_t,
            causal=causal,
        )

        if self.in_channels != self.out_channels:
            self.nin_shortcut = SpatialParallelConv3d(
                in_channels,
                out_channels,
                kernel_size=1,
                padding_mode=padding_mode,
                padding_mode_t=padding_mode_t,
                causal=causal,
            )

    def forward(self, x, zq=None):
        h = x

        if self.use_fused_norm:
            h = self.norm1(h, zq)
        else:
            h = norm_silu(h, self.norm1, zq)

        h = self.conv1(h)

        if self.use_fused_norm:
            h = self.norm2(h, zq)
        else:
            h = norm_silu(h, self.norm2, zq)

        h = self.conv2(h)

        if self.in_channels != self.out_channels:
            x = self.nin_shortcut(x)

        return x + h


class EncoderFCN3D(nn.Module):
    def __init__(
        self,
        ch,
        ch_mult,
        space_down,
        time_down,
        num_res_blocks,
        in_channels,
        z_channels,
        double_z=False,
        zq_ch=None,
        padding_mode="zeros",
        padding_mode_t=None,
        causal=True,
        use_t_isolated_gn=False,
    ):
        super().__init__()
        self.ch = ch
        self.num_levels = len(ch_mult)

        if isinstance(num_res_blocks, int):
            self.num_res_blocks = [num_res_blocks] * self.num_levels
        else:
            self.num_res_blocks = num_res_blocks

        self.space_down_factors = space_down
        self.time_down_factors = time_down
        self.in_channels = in_channels

        self.use_fused_norm = (
            os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true"
        )

        block_mid = [ch * ch_mult[i] for i in range(self.num_levels)]
        block_in = [block_mid[0]] + block_mid[:-1]
        block_out = block_mid

        conv_kwargs = dict(
            padding_mode=padding_mode,
            padding_mode_t=padding_mode_t,
            causal=causal,
        )

        self.conv_in = SpatialParallelConv3d(
            in_channels, block_in[0], kernel_size=3, padding=1, **conv_kwargs
        )

        self.down = nn.ModuleList()
        for i_level in range(self.num_levels):
            down = nn.Module()

            down.block = nn.ModuleList()
            for i in range(self.num_res_blocks[i_level]):
                down.block.append(
                    ResnetBlock3D(
                        in_channels=block_in[i_level] if i == 0 else block_mid[i_level],
                        out_channels=block_mid[i_level],
                        zq_ch=zq_ch,
                        use_t_isolated_gn=use_t_isolated_gn,
                        **conv_kwargs,
                    )
                )

            if space_down[i_level] * time_down[i_level] > 1:
                down.downsample = Downsample3D(
                    block_mid[i_level],
                    block_out[i_level],
                    time_stride=time_down[i_level],
                    space_stride=space_down[i_level],
                    **conv_kwargs,
                )
            else:
                if block_out[i_level] != block_mid[i_level]:
                    down.downsample = SpatialParallelConv3d(
                        block_mid[i_level],
                        block_out[i_level],
                        kernel_size=1,
                        **conv_kwargs,
                    )

            self.down.append(down)

        if zq_ch is None:
            self.norm_out = get_group_norm_3d(
                block_out[-1], use_t_isolated_gn=use_t_isolated_gn
            )
        else:
            self.norm_out = get_spatial_norm_3d(
                block_out[-1],
                zq_ch,
                use_t_isolated_gn=use_t_isolated_gn,
                **conv_kwargs,
            )

        self.conv_out = SpatialParallelConv3d(
            block_out[-1],
            2 * z_channels if double_z else z_channels,
            kernel_size=3,
            padding=1,
            **conv_kwargs,
        )

        self.gradient_checkpointing = False

    def _set_gradient_checkpointing(self, module, value=False):
        if hasattr(module, "gradient_checkpointing"):
            module.gradient_checkpointing = value

    def forward(self, x, zq=None):
        h = self.conv_in(x)
        for i_level in range(self.num_levels):
            for i_block in range(self.num_res_blocks[i_level]):
                h = maybe_checkpoint(self, self.down[i_level].block[i_block], h, zq)
            if hasattr(self.down[i_level], "downsample"):
                h = self.down[i_level].downsample(h)

        if self.use_fused_norm:
            h = self.norm_out(h, zq)
        else:
            h = norm_silu(h, self.norm_out, zq)

        h = self.conv_out(h)
        return h