Ref2VA/video_vae/vae_vit.py

# SPDX-License-Identifier: Apache-2.0
# ViT3D decoder for the MiniMax H3 visual VAE (inference-only bundle).
import torch
import torch.nn as nn
import torch.distributed as dist
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
from diffusers.utils import logging

from .attention import maybe_checkpoint
from .base_module import TransformerBlock, RotaryEmbeddingND
from .flash import make_block_causal_mask_mod
from .func import create_token_ids
from .parallel import get_subseq, gather_subseq, get_parallel_state

logger = logging.get_logger(__name__)


def _linear_with_module_dtype(linear, tensor, out_dtype=None):
    weight = getattr(linear, "weight", None)
    target_dtype = getattr(weight, "dtype", tensor.dtype)
    output = linear(tensor.to(target_dtype))
    if out_dtype is not None and output.dtype != out_dtype:
        output = output.to(out_dtype)
    return output


def _make_seq_len_mask_mod(seq_len, base_mask_mod=None):
    if base_mask_mod is None:
        def mask_mod(b, h, q_idx, kv_idx, seqlen_info, aux_tensors):
            return (q_idx < seq_len) & (kv_idx < seq_len)

        mask_mod.block_sparse_cache_key = ("seq_len", seq_len)
        return mask_mod

    def mask_mod(b, h, q_idx, kv_idx, seqlen_info, aux_tensors):
        return (
            (q_idx < seq_len)
            & (kv_idx < seq_len)
            & base_mask_mod(b, h, q_idx, kv_idx, seqlen_info, aux_tensors)
        )

    base_cache_key = getattr(base_mask_mod, "block_sparse_cache_key", None)
    if base_cache_key is not None:
        mask_mod.block_sparse_cache_key = ("seq_len", seq_len, base_cache_key)
    if hasattr(base_mask_mod, "use_fast_sampling"):
        mask_mod.use_fast_sampling = base_mask_mod.use_fast_sampling
    return mask_mod






def _pack_tensors_3d(tensors, patch_size, patch_size_t):
    batch_size, num_channels_tensors, temporal, height, width = tensors.shape

    tensors = tensors.view(
        batch_size,
        num_channels_tensors,
        temporal // patch_size_t,
        patch_size_t,
        height // patch_size,
        patch_size,
        width // patch_size,
        patch_size,
    )
    tensors = tensors.permute(0, 2, 4, 6, 1, 3, 5, 7)
    tensors = tensors.reshape(
        batch_size,
        (temporal // patch_size_t) * (height // patch_size) * (width // patch_size),
        num_channels_tensors * patch_size_t * patch_size * patch_size,
    )
    return tensors


def _unpack_tensors_3d(tensors, patch_size, patch_size_t, temporal, height, width):
    batch_size, num_patches, channels = tensors.shape
    num_channels_tensors = channels // (patch_size_t * patch_size * patch_size)

    tensors = tensors.view(
        batch_size,
        temporal // patch_size_t,
        height // patch_size,
        width // patch_size,
        num_channels_tensors,
        patch_size_t,
        patch_size,
        patch_size,
    )
    tensors = tensors.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()
    tensors = tensors.reshape(batch_size, num_channels_tensors, temporal, height, width)
    return tensors


class ViTBase(ModelMixin, ConfigMixin):
    """Base class for ViT Encoder and Decoder with common functionality."""

    _supports_gradient_checkpointing = True
    _no_split_modules = ["TransformerBlock"]
    gradient_checkpointing_mode = "full"

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

    def set_spatial_parallel(self, enabled):
        self.spatial_parallel = enabled
        if hasattr(self, "transformer_blocks"):
            for block in self.transformer_blocks:
                block.attn.spatial_parallel = enabled

    def _init_weights(self):
        def basic_init(m):
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

        self.apply(basic_init)

    def init_mask_config(self, dim, is_3d=False):
        self._mask_dim = dim
        self._mask_is_3d = is_3d
        self.register_buffer("mask_token", torch.zeros(1, 1, dim))

    def set_mask_config(self, mask_config):
        self.mask_prob = mask_config.get("mask_prob", 0.0)
        self.mask_enabled = self.mask_prob > 0
        self.mask_style = mask_config.get("mask_style", "replace")
        if self.mask_enabled and self.mask_style == "drop" and self.mask_prob < 1.0:
            logger.warning("mask_style='drop' with mask_prob < 1.0")
        if self._mask_is_3d:
            self.temporal_scale_range = mask_config.get("temporal_scale_range", (0.3, 0.5))
            self.spatial_scale_range = mask_config.get("spatial_scale_range", (0.1, 0.25))
            self.min_mask_ratio = mask_config.get("min_mask_ratio", 0.75)
            self.max_mask_ratio = mask_config.get("max_mask_ratio", 0.95)
        else:
            self.spatial_scale_range = mask_config.get("spatial_scale_range", (0.15, 0.15))
            self.min_mask_ratio = mask_config.get("min_mask_ratio", 0.5)
            self.max_mask_ratio = mask_config.get("max_mask_ratio", 0.75)
        self.aspect_ratio_range = mask_config.get("aspect_ratio_range", (0.75, 1.5))
        self.max_retries = mask_config.get("max_retries", 100)
        if self.mask_enabled and self.mask_style == "drop" and getattr(self, "t_causal", False):
            logger.warning("mask_style='drop' with t_causal may cause issues")
        if self.mask_enabled and "mask_token" in self._buffers:
            del self._buffers["mask_token"]
            self.mask_token = nn.Parameter(torch.randn(1, 1, self._mask_dim) * 0.02)

    def init_suffix_tokens(self, dim, num_register_tokens, has_cls_token=True):
        self.num_register_tokens = num_register_tokens
        if num_register_tokens > 0:
            self.register_tokens = nn.Parameter(torch.randn(1, num_register_tokens, dim) * 0.02)
        else:
            self.register_tokens = None
        if has_cls_token:
            self.cls_token = nn.Parameter(torch.randn(1, 1, dim) * 0.02)

    def apply_mask_preprocess(self, hidden_states, img_ids, patch_dims, num_suffix):
        if self.training and self.mask_enabled:
            raise NotImplementedError(
                "mask modeling is not supported in this inference-only bundle"
            )
        return hidden_states, img_ids

    def forward_transformer_blocks(self, hidden_states, rotary_pos_emb, pack_info=None):
        if pack_info is None:
            pack_info = {}
        for block in self.transformer_blocks:
            hidden_states = maybe_checkpoint(
                self, block, hidden_states, rotary_pos_emb, pack_info
            )
        return hidden_states

    def _pad_for_sp(self, hidden_states, img_ids, pack_info=None):
        if pack_info is None:
            pack_info = {}
        if not self.spatial_parallel:
            return hidden_states, img_ids, pack_info, 0

        seq_len = hidden_states.shape[1]
        sp_size = get_parallel_state().get("sp_size", 1)
        pad_len = (-seq_len) % sp_size
        if pad_len == 0:
            return hidden_states, img_ids, pack_info, 0

        hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, pad_len))
        img_ids = torch.nn.functional.pad(img_ids, (0, 0, 0, pad_len))

        pack_info = dict(pack_info)
        base_mask_mod = pack_info.get("mask_mod")
        pack_info["mask_mod"] = _make_seq_len_mask_mod(seq_len, base_mask_mod)
        pack_info.pop("block_sparse", None)
        return hidden_states, img_ids, pack_info, pad_len

    @staticmethod
    def _unpad_for_sp(hidden_states, pad_len):
        if pad_len == 0:
            return hidden_states
        return hidden_states[:, :-pad_len, :]

    def apply_mask_postprocess(self, hidden_states, num_patches):
        if self.training and self.mask_enabled and self.mask_style == "drop":
            raise NotImplementedError(
                "mask modeling is not supported in this inference-only bundle"
            )
        return hidden_states








class ViT3DDecoder(ViTBase):
    """Vision Transformer Video Decoder using TransformerBlock."""

    @register_to_config
    def __init__(
        self,
        patch_size: int = 16,
        patch_size_t: int = 4,
        t_causal: bool = False,
        in_channels: int = 16,
        out_channels: int = 3,
        num_layers: int = 24,
        heads: int = 16,
        dim_head: int = 64,
        norm_type: str = "layer_norm",
        norm_affine: bool = True,
        qk_norm_type: str = None,
        qk_norm_affine: bool = False,
        ffn_activation_fn: str = "gelu",
        ffn_use_gated: bool = False,
        rope_theta: float = 100.0,
        rope_dim_ratio: float = 1.0,
        bias: bool = True,
        eps: float = 1e-5,
        num_register_tokens: int = 4,
        mask_config: dict = {},
        **kwargs,
    ):
        super().__init__()

        dim = heads * dim_head
        rope_apply_dim = int(dim_head * rope_dim_ratio)

        self.pos_embed = RotaryEmbeddingND(rope_apply_dim, rope_theta, n_dim=3, use_angle=True)

        self.x_embedder = nn.Linear(in_channels, dim)

        self.init_suffix_tokens(dim, num_register_tokens, has_cls_token=False)

        self.t_causal = t_causal

        self.transformer_blocks = nn.ModuleList(
            [
                TransformerBlock(
                    heads=heads,
                    dim_head=dim_head,
                    norm_type=norm_type,
                    norm_affine=norm_affine,
                    qk_norm_type=qk_norm_type,
                    qk_norm_affine=qk_norm_affine,
                    ffn_activation_fn=ffn_activation_fn,
                    ffn_use_gated=ffn_use_gated,
                    bias=bias,
                    eps=eps,
                    **kwargs,
                )
                for _ in range(num_layers)
            ]
        )

        self.spatial_parallel = False
        for block in self.transformer_blocks:
            block.attn.spatial_parallel = False

        self.norm_out = nn.LayerNorm(dim, elementwise_affine=norm_affine, eps=eps)
        patch_dim = out_channels * patch_size_t * patch_size * patch_size
        self.proj_out = nn.Linear(dim, patch_dim)

        self.init_mask_config(dim, is_3d=True)
        self.set_mask_config(mask_config)

        self._init_weights()
        self.gradient_checkpointing = False

        if len(kwargs) > 0 and (not dist.is_initialized() or dist.get_rank() == 0):
            logger.warning(f"Unused kwargs: {kwargs}")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        self.loss_info = {}

        B, C, latent_T, latent_H, latent_W = x.shape
        patch_size = self.config.patch_size
        patch_size_t = self.config.patch_size_t
        num_suffix = 1 + self.num_register_tokens

        hidden_states = _pack_tensors_3d(x, 1, 1)
        latent_size = (latent_T, latent_H, latent_W)

        with torch.autocast("cuda", enabled=False):
            hidden_states = _linear_with_module_dtype(self.x_embedder, hidden_states, hidden_states.dtype)

        num_patches = hidden_states.shape[1]

        tokens = [hidden_states]

        if self.register_tokens is not None:
            register_tokens = self.register_tokens.expand(B, -1, -1)
            tokens.append(register_tokens)

        cls_token = torch.zeros_like(hidden_states[:, 0:1, :])
        tokens.append(cls_token)
        hidden_states = torch.cat(tokens, dim=1)

        patch_dims = [latent_T, latent_H, latent_W]
        img_ids = create_token_ids(latent_size, x.device, x.dtype).expand(B, -1, -1)
        suffix_ids = torch.zeros((B, num_suffix, 3), device=x.device, dtype=img_ids.dtype)
        img_ids = torch.cat([img_ids, suffix_ids], dim=1)

        hidden_states, img_ids = self.apply_mask_preprocess(hidden_states, img_ids, patch_dims, num_suffix)

        pack_info = {}
        if self.t_causal:
            spatial_size = latent_H * latent_W
            mask_mod = make_block_causal_mask_mod(
                num_tokens=num_patches,
                block_size=spatial_size,
                suffix=True,
            )
            pack_info["mask_mod"] = mask_mod

        hidden_states, img_ids, pack_info, sp_pad_len = self._pad_for_sp(hidden_states, img_ids, pack_info)

        rotary_pos_emb = self.pos_embed(img_ids)

        if self.spatial_parallel:
            hidden_states = get_subseq(hidden_states)

        for block in self.transformer_blocks:
            hidden_states = maybe_checkpoint(
                self, block, hidden_states, rotary_pos_emb, pack_info
            )

        if self.spatial_parallel:
            hidden_states = gather_subseq(hidden_states)
        hidden_states = self._unpad_for_sp(hidden_states, sp_pad_len)

        hidden_states = self.norm_out(hidden_states)

        hidden_states = self.apply_mask_postprocess(hidden_states, num_patches)

        with torch.autocast("cuda", enabled=False):
            output = _linear_with_module_dtype(self.proj_out, hidden_states, hidden_states.dtype)

        output = output[:, :num_patches, :]

        video_t = latent_size[0] * patch_size_t
        video_h = latent_size[1] * patch_size
        video_w = latent_size[2] * patch_size
        output = _unpack_tensors_3d(output, patch_size, patch_size_t, video_t, video_h, video_w)

        return output