/diffusion/model/nets/sana_multi_scale_video_v2v.py

# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0


import os
from typing import List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.attention.flex_attention import create_block_mask

from diffusion.model.builder import MODELS
from diffusion.model.nets.sana_blocks import (
    CaptionEmbedder,
    ClipVisionProjection,
    PatchEmbedMS3D,
    T2IFinalLayer,
)
from diffusion.model.nets.sana_multi_scale_video import (
    SanaMSVideo,
    SanaVideoMSBlock,
)
from diffusion.model.registry import ATTENTION_BLOCKS, FFN_BLOCKS
from diffusion.model.utils import auto_grad_checkpoint
from diffusion.utils.dist_utils import get_rank
from diffusion.utils.import_utils import is_xformers_available

_xformers_available = False if os.environ.get("DISABLE_XFORMERS", "0") == "1" else is_xformers_available()
if _xformers_available:
    import xformers.ops


def get_softmax_layer_indices(depth: int, softmax_ratio: float = 0.25) -> List[int]:
    """
    Calculate which layer indices should use softmax attention.

    By default, 25% of layers use softmax attention, evenly distributed.
    For a 20-layer model: layers [4, 9, 14, 19] would use softmax.

    Args:
        depth: Total number of layers
        softmax_ratio: Ratio of layers to use softmax attention (default 0.25)

    Returns:
        List of layer indices that should use softmax attention
    """
    if softmax_ratio == 0:
        return []
    num_softmax_layers = max(1, int(depth * softmax_ratio))
    step = depth / num_softmax_layers
    indices = [int((i + 1) * step) - 1 for i in range(num_softmax_layers)]
    return indices


class SanaV2VVideoMSBlock(SanaVideoMSBlock):
    """Sana video block with V2V-only registry fallbacks.

    This keeps custom V2V attention/FFN names local to ``SanaMSVideoV2V``
    instead of changing the base Sana-Video block used by other releases.
    """

    def __init__(
        self,
        hidden_size,
        num_heads,
        mlp_ratio=4.0,
        drop_path=0.0,
        qk_norm=False,
        attn_type="flash",
        ffn_type="mlp",
        mlp_acts=("silu", "silu", None),
        linear_head_dim=32,
        cross_norm=False,
        cross_attn_image_embeds=False,
        t_kernel_size=3,
        additional_flash_attn=False,
        flash_attn_window_count=None,
        **block_kwargs,
    ):
        super().__init__(
            hidden_size,
            num_heads,
            mlp_ratio=mlp_ratio,
            drop_path=drop_path,
            qk_norm=qk_norm,
            attn_type=attn_type,
            ffn_type=ffn_type,
            mlp_acts=mlp_acts,
            linear_head_dim=linear_head_dim,
            cross_norm=cross_norm,
            cross_attn_image_embeds=cross_attn_image_embeds,
            t_kernel_size=t_kernel_size,
            additional_flash_attn=additional_flash_attn,
            flash_attn_window_count=flash_attn_window_count,
            **block_kwargs,
        )

        if self.attn is None:
            attn_cls = ATTENTION_BLOCKS.get(attn_type) if attn_type else None
            if attn_cls is not None:
                self.attn = attn_cls(
                    in_dim=hidden_size,
                    out_dim=hidden_size,
                    heads=hidden_size // linear_head_dim,
                    eps=1e-8,
                    qk_norm=qk_norm,
                )

        if self.mlp is None:
            ffn_cls = FFN_BLOCKS.get(ffn_type) if ffn_type else None
            if ffn_cls is not None:
                self.mlp = ffn_cls(
                    in_features=hidden_size,
                    hidden_features=int(hidden_size * mlp_ratio),
                    use_bias=(True, True, False),
                    norm=(None, None, None),
                    act=mlp_acts,
                    t_kernel_size=t_kernel_size,
                )


#############################################################################
#                                 Core Sana Model                                #
#################################################################################
@MODELS.register_module()
class SanaMSVideoV2V(SanaMSVideo):
    """
    Diffusion model with a Transformer backbone.
    """

    def __init__(
        self,
        input_size=32,
        patch_size=(1, 2, 2),
        in_channels=4,
        hidden_size=1152,
        depth=28,
        num_heads=16,
        mlp_ratio=4.0,
        class_dropout_prob=0.1,
        learn_sigma=True,
        pred_sigma=True,
        drop_path: float = 0.0,
        caption_channels=2304,
        pe_interpolation=1.0,
        config=None,
        model_max_length=300,
        qk_norm=False,
        y_norm=False,
        norm_eps=1e-5,
        attn_type="flash",
        ffn_type="mlp",
        use_pe=True,
        y_norm_scale_factor=1.0,
        patch_embed_kernel=None,
        mlp_acts=("silu", "silu", None),
        linear_head_dim=32,
        cross_norm=False,
        cross_attn_type="flash",
        cross_attn_image_embeds=False,
        image_embed_channels=1152,
        pos_embed_type="wan_rope",
        rope_fhw_dim=None,
        t_kernel_size=3,
        flash_attn_layer_idx=None,
        flash_attn_layer_type=None,
        flash_attn_window_count=None,
        addition_layers_num=0,
        pack_latents=False,
        additional_inchannels=0,
        softmax_ratio: float = 0.0,
        softmax_layer_indices: Optional[List[int]] = None,
        softmax_attn_type="GDNSoftmaxAttention",
        **kwargs,
    ):
        super().__init__(
            input_size=input_size,
            patch_size=patch_size,
            in_channels=in_channels,
            hidden_size=hidden_size,
            depth=depth,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            class_dropout_prob=class_dropout_prob,
            learn_sigma=learn_sigma,
            pred_sigma=pred_sigma,
            drop_path=drop_path,
            caption_channels=caption_channels,
            pe_interpolation=pe_interpolation,
            config=config,
            model_max_length=model_max_length,
            qk_norm=qk_norm,
            y_norm=y_norm,
            norm_eps=norm_eps,
            attn_type=attn_type,
            ffn_type=ffn_type,
            use_pe=use_pe,
            y_norm_scale_factor=y_norm_scale_factor,
            patch_embed_kernel=patch_embed_kernel,
            mlp_acts=mlp_acts,
            linear_head_dim=linear_head_dim,
            cross_norm=cross_norm,
            cross_attn_type=cross_attn_type,
            pos_embed_type=pos_embed_type,
            rope_fhw_dim=rope_fhw_dim,
            t_kernel_size=t_kernel_size,
            flash_attn_layer_idx=flash_attn_layer_idx,
            flash_attn_layer_type=flash_attn_layer_type,
            flash_attn_window_count=flash_attn_window_count,
            addition_layers_num=addition_layers_num,
            pack_latents=pack_latents,
            additional_inchannels=additional_inchannels,
            **kwargs,
        )
        self.patch_size = patch_size
        self.h = self.w = 0
        approx_gelu = lambda: nn.GELU(approximate="tanh")
        self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
        self.pos_embed_ms = None
        self.pack_latents = pack_latents
        self.addition_layers_num = addition_layers_num

        kernel_size = patch_embed_kernel or patch_size
        x_embedder_in_channels = (
            in_channels + additional_inchannels
            if additional_inchannels is not None and additional_inchannels > 0
            else in_channels
        )
        if self.pack_latents:
            x_embedder_in_channels = x_embedder_in_channels * 2 * 2
            self.out_channels = in_channels * 2 * 2
        elif self.addition_layers_num > 0:
            self.out_channels = in_channels

        self.x_embedder = PatchEmbedMS3D(
            patch_size, x_embedder_in_channels, hidden_size, kernel_size=kernel_size, bias=True
        )

        self.y_embedder = CaptionEmbedder(
            in_channels=caption_channels,
            hidden_size=hidden_size,
            uncond_prob=class_dropout_prob,
            act_layer=approx_gelu,
            token_num=model_max_length,
        )
        if cross_attn_image_embeds:
            self.image_embedder = ClipVisionProjection(image_embed_channels, hidden_size)
        else:
            self.image_embedder = None

        # Calculate which layers use softmax attention
        if softmax_layer_indices is not None:
            self.softmax_layer_indices = softmax_layer_indices
        else:
            self.softmax_layer_indices = get_softmax_layer_indices(depth, softmax_ratio)

        if attn_type in ["flash", "FlexLinearAttention", "flex"]:
            attention_head_dim = hidden_size // num_heads
        else:
            attention_head_dim = linear_head_dim
        if self.use_pe:
            self.rope = self.get_rope(pos_embed_type, attention_head_dim, patch_size, rope_fhw_dim)
        else:
            self.rope = None

        drop_path = [x.item() for x in torch.linspace(0, drop_path, depth)]  # stochastic depth decay rule

        # insert flash attention layers
        if flash_attn_layer_idx is not None and flash_attn_layer_type is not None:
            assert int(flash_attn_layer_idx[-1]) < depth
            additional_flash_attn = [
                flash_attn_layer_type if i in flash_attn_layer_idx else False for i in range(depth)
            ]
        else:
            additional_flash_attn = [False] * depth

        # visualize qkv
        self.save_qkv = False
        self.qkv_store_buffer = {}

        # diagonal mask
        self.diagonal_mask = None
        attn_type_list = [attn_type] * depth
        if attn_type in ["flex", "FlexLinearAttention"]:
            attn_type_list[0] = "flash"
            attn_type_list[1] = "flash"

        for i in self.softmax_layer_indices:
            attn_type_list[i] = softmax_attn_type

        self.use_flex_attention = len([_attn_type for _attn_type in attn_type_list if "flex" in _attn_type.lower()]) > 0

        self.blocks = nn.ModuleList(
            [
                SanaV2VVideoMSBlock(
                    hidden_size,
                    num_heads,
                    mlp_ratio=mlp_ratio,
                    drop_path=drop_path[i],
                    qk_norm=qk_norm,
                    attn_type=attn_type_list[i],
                    ffn_type=ffn_type,
                    mlp_acts=mlp_acts,
                    linear_head_dim=linear_head_dim,
                    cross_norm=cross_norm,
                    cross_attn_image_embeds=cross_attn_image_embeds,
                    t_kernel_size=t_kernel_size,
                    additional_flash_attn=additional_flash_attn[i],
                    flash_attn_window_count=flash_attn_window_count,
                )
                for i in range(depth)
            ]
        )

        self.final_layer = T2IFinalLayer(hidden_size, patch_size, self.out_channels)

        if get_rank() == 0:
            if ffn_type == "GLUMBConvTemp":
                self.logger(f"{ffn_type} Temporal kernal: {t_kernel_size}")
            if flash_attn_layer_idx is not None:
                self.logger(f"additional flash attn layer idx: {flash_attn_layer_idx}, type: {flash_attn_layer_type}")
                if flash_attn_layer_type == "window_flash":
                    self.logger(f"flash attn window count: {flash_attn_window_count}")

        self.initialize()
        self.save_block_output = False
        self.block_output_buffer = {}

    def create_flexattention_chunkcausal_mask(self, x, THW, chunk_index=None):
        """
        Args:
            x: input tensor, shape (B, N, C)
            THW: tuple (f, h, w)
            chunk_indices: list or tensor, containing the start frame index of each chunk. If None, view each frame as a separate chunk.
        Returns:
            block_mask: BlockMask object
        """
        B, N, C = x.shape
        f, h, w = THW
        BLOCK_SIZE = 128
        chunk_id_map = torch.zeros(f, h, w, dtype=torch.long, device=x.device)
        if chunk_index is None:
            chunk_indices = range(f)
        else:
            chunk_indices = chunk_index

        for i, start_idx in enumerate(chunk_indices):
            chunk_id_map[start_idx:] = i
        chunk_id_map = chunk_id_map.view(f * h * w)

        pad_len = (BLOCK_SIZE - (N % BLOCK_SIZE)) % BLOCK_SIZE
        if pad_len > 0:
            # chunk_indices always larger than chunk_id_map, since the start index cannot be smaller than the frame index
            padding = torch.full((pad_len,), chunk_indices[-1] + 1, device=chunk_id_map.device)
            chunk_id_map = torch.cat([chunk_id_map, padding])

        def chunk_causal_mask_mod(b, h, q, kv):
            return chunk_id_map[q] >= chunk_id_map[kv]

        block_mask = create_block_mask(
            chunk_causal_mask_mod, B=None, H=None, Q_LEN=N, KV_LEN=N, device=x.device, BLOCK_SIZE=BLOCK_SIZE
        )

        return block_mask

    def forward(self, x, timestep, y, mask=None, **kwargs):
        """
        Forward pass of Sana.
        x: (N, C, T, H, W) tensor of spatial inputs (images or latent representations of images)
        t: (N,) tensor of diffusion timesteps or (N, 1, F) tensor of diffusion timesteps
        y: (N, 1, 120, C) tensor of class labels
        """
        bs = x.shape[0]
        x = x.to(self.dtype)
        if self.timestep_norm_scale_factor != 1.0:
            timestep = (timestep.float() / self.timestep_norm_scale_factor).to(torch.float32)
        else:
            timestep = timestep.long().to(torch.float32)
        y = y.to(self.dtype)
        self.f, self.h, self.w = (
            x.shape[-3] // self.patch_size[0],
            x.shape[-2] // self.patch_size[1],
            x.shape[-1] // self.patch_size[2],
        )

        data_info = kwargs.get("data_info", {})
        if data_info.get("image_vae_embeds", None) is not None:
            x = torch.cat([x, data_info["image_vae_embeds"].to(self.dtype)], dim=1)
        if data_info.get("image_embeds", None) is not None:
            image_embeds = data_info["image_embeds"].to(self.dtype)
            image_embeds = self.image_embedder(image_embeds)
            kwargs["image_embeds"] = image_embeds
        if self.save_qkv:
            self.qkv_store_buffer[int(timestep[0].item())] = {}
        if self.save_block_output:
            self.inference_timestep = int(timestep[0].item())

        if self.pack_latents:
            x = self._pack_latents(x, bs, self.in_channels, self.h, self.w, self.f)
            self.h = self.h // 2
            self.w = self.w // 2
        if self.x_embedder.patch_size != self.x_embedder.kernel_size and self.x_embedder.kernel_size == (1, 2, 2):
            x = F.pad(x, (0, 1, 0, 1, 0, 0))

        x = self.x_embedder(x)
        image_pos_embed = None
        if self.use_pe:
            x, image_pos_embed = self._apply_positional_embedding(x, bs)

        t = self.t_embedder(timestep.flatten())  # (N, D)
        t0 = self.t_block(t)
        t = t.unflatten(dim=0, sizes=timestep.shape)
        t0 = t0.unflatten(dim=0, sizes=timestep.shape)
        y = self.y_embedder(y, self.training, mask=mask)  # (N, D)
        if self.y_norm:
            y = self.attention_y_norm(y)

        if mask is not None:
            mask = mask.to(torch.int16)
            mask = mask.repeat(y.shape[0] // mask.shape[0], 1) if mask.shape[0] != y.shape[0] else mask
            mask = mask.squeeze(1).squeeze(1)
            if _xformers_available:
                y = y.squeeze(1).masked_select(mask.unsqueeze(-1) != 0).view(1, -1, x.shape[-1])
                y_lens = mask.sum(dim=1).tolist()
            else:
                y_lens = mask
        elif _xformers_available:
            y_lens = [y.shape[2]] * y.shape[0]
            y = y.squeeze(1).view(1, -1, x.shape[-1])
        else:
            raise ValueError(f"Attention type is not available due to _xformers_available={_xformers_available}.")

        if self.use_flex_attention:
            block_mask = self.create_flexattention_chunkcausal_mask(
                x, (self.f, self.h, self.w), kwargs.get("chunk_index", None)
            )
        else:
            block_mask = None

        for i, block in enumerate(self.blocks):
            if self.save_qkv:
                block.attn.qkv_store_buffer = {}

            x = auto_grad_checkpoint(
                block,
                x,
                y,
                t0,
                y_lens,
                (self.f, self.h, self.w),
                image_pos_embed,
                block_mask=block_mask,
                **kwargs,
                use_reentrant=False,
            )  # (N, T, D) #support grad checkpoint

            if self.save_qkv:
                self.qkv_store_buffer[int(timestep[0].item())][f"block_{i}"] = block.attn.qkv_store_buffer
                block.attn.qkv_store_buffer = None

        if self.addition_layers_num > 0:
            x = self.upsample_layer(x)
            x = self._unpack_latents_additional_layers(x, self.h * 2, self.w * 2, self.f)
            if self.pos_embed_type == "wan_rope":
                image_pos_embed = self.rope((self.f, self.h * 2, self.w * 2), x.device)
            else:
                raise ValueError(f"Unknown pos_embed_type: {self.pos_embed_type}")
            for i, block in enumerate(self.addition_layers):
                x = auto_grad_checkpoint(
                    block,
                    x,
                    y,
                    t0,
                    y_lens,
                    (self.f, self.h * 2, self.w * 2),
                    image_pos_embed,
                    block_mask=block_mask if i > 1 else None,
                    **kwargs,
                    use_reentrant=False,
                )

        x = self.final_layer(x, t)  # (N, T, patch_size ** 2 * out_channels)
        x = self.unpatchify(x)  # (N, out_channels, H, W)
        if self.pack_latents:
            x = self._unpack_latents(x, self.h * 2, self.w * 2, self.f)

        if self.save_block_output:
            block_output = self.get_block_output()
            self.block_output_buffer[self.inference_timestep] = block_output
        return x


#################################################################################
#                             Sana Multi-scale Configs                          #
#################################################################################


@MODELS.register_module()
def SanaMSVideoV2V_2000M_P1_D20(**kwargs):
    # 20 layers, 2B
    return SanaMSVideoV2V(depth=20, hidden_size=2240, patch_size=(1, 1, 1), num_heads=20, **kwargs)