cosmos_predict1/diffusion/training/networks/general_dit.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

"""
A general implementation of adaln-modulated VIT-like~(DiT) transformer for video processing.
It allows us easy to switch building blocks used and their order. Its instantiation includes
* transformer on fully flattened tokens
* factored spatial and temporal attention
* factored non-overlap spatial and temporal attention
* mixing of above attention types

Limitations:

* In favor of simplicity and cleanness, many ops are not fused and we can do better
* such as combining mutiple adaln MLPs into one inside one transformer block.
* we use reshape heavily, which may be not efficient when its occurs unnecessary CUDA memory copy

Purpose:
* A prototype for testing different attention types and their combinations
* Idealy, we want to know where we should allocate our resources / FLOPS / memory via extensive empirical studies
"""

from collections.abc import Container
from typing import List, Optional, Tuple

import torch
from einops import rearrange
from megatron.core import parallel_state
from torch import nn
from torch.distributed import ProcessGroup, get_process_group_ranks
from torchvision import transforms

from cosmos_predict1.diffusion.module.attention import get_normalization
from cosmos_predict1.diffusion.training.conditioner import DataType
from cosmos_predict1.diffusion.training.module.blocks import (
    DITBuildingBlock,
    FinalLayer,
    GeneralDITTransformerBlock,
    PatchEmbed,
    SDXLTimestepEmbedding,
    SDXLTimesteps,
)
from cosmos_predict1.diffusion.training.module.position_embedding import (
    LearnableEmb3D,
    LearnableEmb3D_FPS_Aware,
    LearnablePosEmbAxis,
    SinCosPosEmb,
    SinCosPosEmb_FPS_Aware,
    SinCosPosEmbAxis,
    VideoRopePosition3DEmb,
    VideoRopePositionEmb,
)
from cosmos_predict1.diffusion.training.tensor_parallel import gather_along_first_dim, scatter_along_first_dim
from cosmos_predict1.utils import log


class GeneralDIT(nn.Module):
    """
    A general implementation of adaln-modulated VIT-like~(DiT) transformer for video processing.
        Attributes:
        max_img_h (int): Maximum height of the input images.
        max_img_w (int): Maximum width of the input images.
        max_frames (int): Maximum number of frames in the video sequence.
        in_channels (int): Number of input channels (e.g., RGB channels for color images).
        out_channels (int): Number of output channels.
        patch_spatial (tuple of int): Spatial resolution of patches for input processing.
        patch_temporal (int): Temporal resolution of patches for input processing.
        concat_padding_mask (bool): If True, includes a mask channel in the input to handle padding.
        block_config (str): Configuration of the transformer block, e.g., 'FA-CA-MLP', means
            full attention, cross attention, and MLP in sequence in one transformer block.
        model_channels (int): Base number of channels used throughout the model.
        num_blocks (int): Number of residual blocks per resolution in the transformer.
        num_heads (int): Number of heads in the multi-head self-attention layers.
        spatial_attn_win_size (int): Window size for the spatial attention mechanism.
        temporal_attn_win_size (int): Window size for the temporal attention mechanism.
        mlp_ratio (float): Expansion ratio for the MLP (multi-layer perceptron) blocks in the transformer.
        use_memory_save (bool): If True, utilizes checkpointing to reduce memory usage during training. (Deprecated)
        use_checkpoint (bool): If True, utilizes checkpointing to reduce memory usage during training for all blocks.
        crossattn_emb_channels (int): Number of embedding channels used in the cross-attention layers.
        use_cross_attn_mask (bool): If True, applies a mask during cross-attention operations to manage sequence alignment.
        pos_emb_cls (str): Type of positional embeddings used ('sincos' for sinusoidal or other types).
        pos_emb_learnable (bool): Specifies if positional embeddings are learnable.
        pos_emb_interpolation (str): Method used for interpolating positional embeddings, e.g., 'crop' for cropping adjustments.
        block_x_format (str, optional): The format of the input tensor for the transformer block. Defaults to "BTHWD". Only support 'BTHWD' and 'THWBD'.
        legacy_patch_emb (bool): If True, applies 3D convolutional layers for video inputs, otherwise, use Linear! This is for backward compatibility.
        rope_h_extrapolation_ratio (float): Ratio of the height extrapolation for the rope positional embedding.
        rope_w_extrapolation_ratio (float): Ratio of the width extrapolation for the rope positional embedding.
        rope_t_extrapolation_ratio (float): Ratio of the temporal extrapolation for the rope positional embedding.
    Note:
        block_config support block type:
        * spatial_sa, ssa: spatial self attention
        * temporal_sa, tsa: temporal self attention
        * cross_attn, ca: cross attention
        * full_attn: full attention on all flatten tokens
        * mlp, ff: feed forward block
        * use '-' to separate different building blocks, e.g., 'FA-CA-MLP' means full attention, cross attention, and MLP in sequence in one transformer block.

    Example:
        >>> # full attention, cross attention, and MLP
        >>> option1_block_config = 'FA-CA-MLP'
        >>> model_1 = GeneralDIT(
                max_img_h=64, max_img_w=64, max_frames=32, in_channels=16, out_channels=16,
                patch_spatial=2, patch_temporal=1, model_channels=768, num_blocks=10,
                num_heads=16, mlp_ratio=4.0,
                spatial_attn_win_size=1, temporal_attn_win_size=1,
                block_config=option1_block_config
            )
        >>> option2_block_config = 'SSA-CA-MLP-TSA-CA-MLP'
        >>> model_2 = GeneralDIT(
                max_img_h=64, max_img_w=64, max_frames=32, in_channels=16, out_channels=16,
                patch_spatial=2, patch_temporal=1, model_channels=768, num_blocks=10,
                num_heads=16, mlp_ratio=4.0,
                spatial_attn_win_size=1, temporal_attn_win_size=1,
                block_config=option2_block_config
            )
        >>> # option3 model
        >>> model_3 = GeneralDIT(
                max_img_h=64, max_img_w=64, max_frames=32, in_channels=16, out_channels=16,
                patch_spatial=2, patch_temporal=1, model_channels=768, num_blocks=10,
                num_heads=16, mlp_ratio=4.0,
                spatial_attn_win_size=1, temporal_attn_win_size=2,
                block_config=option2_block_config
            )
        >>> # Process input tensor through the model
        >>> output = model(input_tensor)
    """

    def __init__(
        self,
        max_img_h: int,
        max_img_w: int,
        max_frames: int,
        in_channels: int,
        out_channels: int,
        patch_spatial: tuple,
        patch_temporal: int,
        concat_padding_mask: bool = True,
        # attention settings
        block_config: str = "FA-CA-MLP",
        model_channels: int = 768,
        num_blocks: int = 10,
        num_heads: int = 16,
        window_block_indexes: list = [],  # index for window attention block
        window_sizes: list = [],  # window size for window attention block in the order of T, H, W
        spatial_attn_win_size: int = 1,
        temporal_attn_win_size: int = 1,
        mlp_ratio: float = 4.0,
        use_memory_save: bool = False,
        use_checkpoint: bool = False,
        block_x_format: str = "BTHWD",
        # cross attention settings
        crossattn_emb_channels: int = 1024,
        use_cross_attn_mask: bool = False,
        # positional embedding settings
        pos_emb_cls: str = "sincos",
        pos_emb_learnable: bool = False,
        pos_emb_interpolation: str = "crop",
        min_fps: int = 1,  # 1 for getty video
        max_fps: int = 30,  # 120 for getty video but let's use 30
        additional_timestamp_channels: dict = None,  # Follow SDXL, in format of {condition_name : dimension}
        affline_emb_norm: bool = False,  # whether or not to normalize the affine embedding
        use_adaln_lora: bool = False,
        adaln_lora_dim: int = 256,
        layer_mask: list = None,  # whether or not a layer is used. For controlnet encoder
        legacy_patch_emb: bool = True,
        rope_h_extrapolation_ratio: float = 1.0,
        rope_w_extrapolation_ratio: float = 1.0,
        rope_t_extrapolation_ratio: float = 1.0,
        extra_per_block_abs_pos_emb: bool = True,
        extra_per_block_abs_pos_emb_type: str = "learnable",
        extra_h_extrapolation_ratio: float = 1.0,
        extra_w_extrapolation_ratio: float = 1.0,
        extra_t_extrapolation_ratio: float = 1.0,
    ) -> None:
        super().__init__()
        self.max_img_h = max_img_h
        self.max_img_w = max_img_w
        self.max_frames = max_frames
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.patch_spatial = patch_spatial
        self.patch_temporal = patch_temporal
        self.num_heads = num_heads
        self.num_blocks = num_blocks
        self.model_channels = model_channels
        self.use_cross_attn_mask = use_cross_attn_mask
        self.concat_padding_mask = concat_padding_mask
        # positional embedding settings
        self.pos_emb_cls = pos_emb_cls
        self.pos_emb_learnable = pos_emb_learnable
        self.pos_emb_interpolation = pos_emb_interpolation
        self.min_fps = min_fps
        self.max_fps = max_fps
        self.additional_timestamp_channels = additional_timestamp_channels
        self.affline_emb_norm = affline_emb_norm
        self.legacy_patch_emb = legacy_patch_emb
        self.rope_h_extrapolation_ratio = rope_h_extrapolation_ratio
        self.rope_w_extrapolation_ratio = rope_w_extrapolation_ratio
        self.rope_t_extrapolation_ratio = rope_t_extrapolation_ratio
        self.extra_per_block_abs_pos_emb = extra_per_block_abs_pos_emb
        self.extra_per_block_abs_pos_emb_type = extra_per_block_abs_pos_emb_type.lower()
        self.extra_h_extrapolation_ratio = extra_h_extrapolation_ratio
        self.extra_w_extrapolation_ratio = extra_w_extrapolation_ratio
        self.extra_t_extrapolation_ratio = extra_t_extrapolation_ratio

        self.build_patch_embed()
        self.build_pos_embed()
        self.cp_group = None
        self.sequence_parallel = getattr(parallel_state, "sequence_parallel", False)
        self.block_x_format = block_x_format
        self.use_adaln_lora = use_adaln_lora
        self.adaln_lora_dim = adaln_lora_dim
        self.t_embedder = nn.Sequential(
            SDXLTimesteps(model_channels),
            SDXLTimestepEmbedding(model_channels, model_channels, use_adaln_lora=use_adaln_lora),
        )

        self.blocks = nn.ModuleDict()
        self.block_config = block_config
        self.use_memory_save = use_memory_save
        self.use_checkpoint = use_checkpoint

        assert (
            len(window_block_indexes) == 0 or block_config == "FA-CA-MLP"
        ), "Block config must be FA-CA-MLP if using a combination of window attention and global attention"

        layer_mask = [False] * num_blocks if layer_mask is None else layer_mask
        assert (
            len(layer_mask) == num_blocks
        ), f"Layer mask length {len(layer_mask)} does not match num_blocks {num_blocks}"
        for idx in range(num_blocks):
            if layer_mask[idx]:
                continue
            self.blocks[f"block{idx}"] = GeneralDITTransformerBlock(
                x_dim=model_channels,
                context_dim=crossattn_emb_channels,
                num_heads=num_heads,
                block_config=block_config,
                window_sizes=(
                    window_sizes if idx in window_block_indexes else []
                ),  # There will be bug if using "WA-CA-MLP"
                mlp_ratio=mlp_ratio,
                spatial_attn_win_size=spatial_attn_win_size,
                temporal_attn_win_size=temporal_attn_win_size,
                x_format=self.block_x_format,
                use_adaln_lora=use_adaln_lora,
                adaln_lora_dim=adaln_lora_dim,
                use_checkpoint=use_checkpoint,
            )

        self.build_decode_head()
        self.build_additional_timestamp_embedder()
        if self.affline_emb_norm:
            log.critical("Building affine embedding normalization layer")
            self.affline_norm = get_normalization("R", model_channels)
        else:
            self.affline_norm = nn.Identity()
        self.init_weights()

        if self.use_memory_save:
            log.critical("Using checkpointing to save memory! only verified in 14B base model training!")
            for block in self.blocks.values():
                block.set_memory_save()

    def init_weights(self):
        # Initialize transformer layers:
        def _basic_init(module):
            if isinstance(module, nn.Linear):
                torch.nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

        self.apply(_basic_init)

        # Initialize timestep embedding
        nn.init.normal_(self.t_embedder[1].linear_1.weight, std=0.02)
        if self.t_embedder[1].linear_1.bias is not None:
            nn.init.constant_(self.t_embedder[1].linear_1.bias, 0)
        nn.init.normal_(self.t_embedder[1].linear_2.weight, std=0.02)
        if self.t_embedder[1].linear_2.bias is not None:
            nn.init.constant_(self.t_embedder[1].linear_2.bias, 0)

        # Zero-out adaLN modulation layers in DiT blocks:
        for transformer_block in self.blocks.values():
            for block in transformer_block.blocks:
                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
                if block.adaLN_modulation[-1].bias is not None:
                    nn.init.constant_(block.adaLN_modulation[-1].bias, 0)

        # Tensor parallel
        if parallel_state.is_initialized() and parallel_state.get_tensor_model_parallel_world_size() > 1:
            self.initialize_tensor_parallel_weights()

    def initialize_tensor_parallel_weights(self):
        """
        Initialize weights for tensor parallel layers.

        This function performs the following steps:
        1. Retrieves the tensor parallel rank.
        2. Saves the current random state.
        3. Sets a new random seed based on the tensor parallel rank.
        4. Initializes weights for attention and MLP layers in each block.
        5. Restores the original random state.

        The use of different random seeds for each rank ensures
        unique initializations across parallel processes.
        """
        tp_rank = parallel_state.get_tensor_model_parallel_rank()

        # Save the current random state
        rng_state = torch.get_rng_state()

        # Set a new random seed based on the tensor parallel rank
        torch.manual_seed(tp_rank)

        for block in self.blocks.values():
            for layer in block.blocks:
                if layer.block_type in ["full_attn", "fa", "cross_attn", "ca"]:
                    # Initialize weights for attention layers
                    torch.nn.init.xavier_uniform_(layer.block.attn.to_q[0].weight)
                    torch.nn.init.xavier_uniform_(layer.block.attn.to_k[0].weight)
                    torch.nn.init.xavier_uniform_(layer.block.attn.to_v[0].weight)
                    torch.nn.init.xavier_uniform_(layer.block.attn.to_out[0].weight)
                elif layer.block_type in ["mlp", "ff"]:
                    # Initialize weights for MLP layers
                    torch.nn.init.xavier_uniform_(layer.block.layer1.weight)
                    torch.nn.init.xavier_uniform_(layer.block.layer2.weight)
                else:
                    raise ValueError(f"Unknown block type {layer.block_type}")

        # Restore the original random state
        torch.set_rng_state(rng_state)

    def build_decode_head(self):
        self.final_layer = FinalLayer(
            hidden_size=self.model_channels,
            spatial_patch_size=self.patch_spatial,
            temporal_patch_size=self.patch_temporal,
            out_channels=self.out_channels,
            use_adaln_lora=self.use_adaln_lora,
            adaln_lora_dim=self.adaln_lora_dim,
        )

    def build_patch_embed(self):
        (
            concat_padding_mask,
            in_channels,
            patch_spatial,
            patch_temporal,
            model_channels,
        ) = (
            self.concat_padding_mask,
            self.in_channels,
            self.patch_spatial,
            self.patch_temporal,
            self.model_channels,
        )
        in_channels = in_channels + 1 if concat_padding_mask else in_channels
        self.x_embedder = PatchEmbed(
            spatial_patch_size=patch_spatial,
            temporal_patch_size=patch_temporal,
            in_channels=in_channels,
            out_channels=model_channels,
            bias=False,
            keep_spatio=True,
            legacy_patch_emb=self.legacy_patch_emb,
        )
        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d)
        if self.legacy_patch_emb:
            w = self.x_embedder.proj.weight.data
            nn.init.xavier_uniform_(w.view([w.shape[0], -1]))

    def build_additional_timestamp_embedder(self):
        if self.additional_timestamp_channels:
            self.additional_timestamp_embedder = nn.ModuleDict()
            for cond_name, cond_emb_channels in self.additional_timestamp_channels.items():
                log.critical(
                    f"Building additional timestamp embedder for {cond_name} with {cond_emb_channels} channels"
                )
                self.additional_timestamp_embedder[cond_name] = nn.Sequential(
                    SDXLTimesteps(cond_emb_channels),
                    SDXLTimestepEmbedding(cond_emb_channels, cond_emb_channels),
                )

    def prepare_additional_timestamp_embedder(self, **kwargs):
        condition_concat = []

        for cond_name, embedder in self.additional_timestamp_embedder.items():
            condition_concat.append(embedder(kwargs[cond_name])[0])
        embedding = torch.cat(condition_concat, dim=1)
        if embedding.shape[1] < self.model_channels:
            embedding = nn.functional.pad(embedding, (0, self.model_channels - embedding.shape[1]))
        return embedding

    def build_pos_embed(self):
        if self.pos_emb_cls == "sincos":
            cls_type = SinCosPosEmb
        elif self.pos_emb_cls == "learnable":
            cls_type = LearnableEmb3D
        elif self.pos_emb_cls == "sincos_fps_aware":
            cls_type = SinCosPosEmb_FPS_Aware
        elif self.pos_emb_cls == "learnable_fps_aware":
            cls_type = LearnableEmb3D_FPS_Aware
        elif self.pos_emb_cls == "rope":
            cls_type = VideoRopePositionEmb
        elif self.pos_emb_cls == "rope3d":
            cls_type = VideoRopePosition3DEmb
        else:
            raise ValueError(f"Unknown pos_emb_cls {self.pos_emb_cls}")

        log.critical(f"Building positional embedding with {self.pos_emb_cls} class, impl {cls_type}")
        kwargs = dict(
            model_channels=self.model_channels,
            len_h=self.max_img_h // self.patch_spatial,
            len_w=self.max_img_w // self.patch_spatial,
            len_t=self.max_frames // self.patch_temporal,
            max_fps=self.max_fps,
            min_fps=self.min_fps,
            is_learnable=self.pos_emb_learnable,
            interpolation=self.pos_emb_interpolation,
            head_dim=self.model_channels // self.num_heads,
            h_extrapolation_ratio=self.rope_h_extrapolation_ratio,
            w_extrapolation_ratio=self.rope_w_extrapolation_ratio,
            t_extrapolation_ratio=self.rope_t_extrapolation_ratio,
        )
        self.pos_embedder = cls_type(
            **kwargs,
        )

        assert self.extra_per_block_abs_pos_emb is True, "extra_per_block_abs_pos_emb must be True"

        if self.extra_per_block_abs_pos_emb:
            assert self.extra_per_block_abs_pos_emb_type in [
                "learnable",
            ], f"Unknown extra_per_block_abs_pos_emb_type {self.extra_per_block_abs_pos_emb_type}"
            kwargs["h_extrapolation_ratio"] = self.extra_h_extrapolation_ratio
            kwargs["w_extrapolation_ratio"] = self.extra_w_extrapolation_ratio
            kwargs["t_extrapolation_ratio"] = self.extra_t_extrapolation_ratio
            self.extra_pos_embedder = LearnablePosEmbAxis(**kwargs)

    def prepare_embedded_sequence(
        self,
        x_B_C_T_H_W: torch.Tensor,
        fps: Optional[torch.Tensor] = None,
        padding_mask: Optional[torch.Tensor] = None,
        latent_condition: Optional[torch.Tensor] = None,
        latent_condition_sigma: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Prepares an embedded sequence tensor by applying positional embeddings and handling padding masks.

        Args:
            x_B_C_T_H_W (torch.Tensor): video
            fps (Optional[torch.Tensor]): Frames per second tensor to be used for positional embedding when required.
                                    If None, a default value (`self.base_fps`) will be used.
            padding_mask (Optional[torch.Tensor]): current it is not used

        Returns:
            Tuple[torch.Tensor, Optional[torch.Tensor]]:
                - A tensor of shape (B, T, H, W, D) with the embedded sequence.
                - An optional positional embedding tensor, returned only if the positional embedding class
                (`self.pos_emb_cls`) includes 'rope'. Otherwise, None.

        Notes:
            - If `self.concat_padding_mask` is True, a padding mask channel is concatenated to the input tensor.
            - The method of applying positional embeddings depends on the value of `self.pos_emb_cls`.
            - If 'rope' is in `self.pos_emb_cls` (case insensitive), the positional embeddings are generated using
                the `self.pos_embedder` with the shape [T, H, W].
            - If "fps_aware" is in `self.pos_emb_cls`, the positional embeddings are generated using the `self.pos_embedder`
                with the fps tensor.
            - Otherwise, the positional embeddings are generated without considering fps.
        """
        if self.concat_padding_mask:
            padding_mask = transforms.functional.resize(
                padding_mask, list(x_B_C_T_H_W.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
            )
            x_B_C_T_H_W = torch.cat(
                [x_B_C_T_H_W, padding_mask.unsqueeze(1).repeat(1, 1, x_B_C_T_H_W.shape[2], 1, 1)], dim=1
            )
        x_B_T_H_W_D = self.x_embedder(x_B_C_T_H_W)

        if self.extra_per_block_abs_pos_emb:
            extra_pos_emb = self.extra_pos_embedder(x_B_T_H_W_D, fps=fps)
        else:
            extra_pos_emb = None

        if "rope" in self.pos_emb_cls.lower():
            return x_B_T_H_W_D, self.pos_embedder(x_B_T_H_W_D, fps=fps), extra_pos_emb

        if "fps_aware" in self.pos_emb_cls:
            x_B_T_H_W_D = x_B_T_H_W_D + self.pos_embedder(x_B_T_H_W_D, fps=fps)  # [B, T, H, W, D]
        else:
            x_B_T_H_W_D = x_B_T_H_W_D + self.pos_embedder(x_B_T_H_W_D)  # [B, T, H, W, D]
        return x_B_T_H_W_D, None, extra_pos_emb

    def decoder_head(
        self,
        x_B_T_H_W_D: torch.Tensor,
        emb_B_D: torch.Tensor,
        crossattn_emb: torch.Tensor,
        origin_shape: Tuple[int, int, int, int, int],  # [B, C, T, H, W]
        crossattn_mask: Optional[torch.Tensor] = None,
        adaln_lora_B_3D: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        del crossattn_emb, crossattn_mask
        B, C, T_before_patchify, H_before_patchify, W_before_patchify = origin_shape
        x_BT_HW_D = rearrange(x_B_T_H_W_D, "B T H W D -> (B T) (H W) D")
        x_BT_HW_D = self.final_layer(x_BT_HW_D, emb_B_D, adaln_lora_B_3D=adaln_lora_B_3D)
        # This is to ensure x_BT_HW_D has the correct shape because
        # when we merge T, H, W into one dimension, x_BT_HW_D has shape (B * T * H * W, 1*1, D).
        x_BT_HW_D = x_BT_HW_D.view(
            B * T_before_patchify // self.patch_temporal,
            H_before_patchify // self.patch_spatial * W_before_patchify // self.patch_spatial,
            -1,
        )
        x_B_D_T_H_W = rearrange(
            x_BT_HW_D,
            "(B T) (H W) (p1 p2 t C) -> B C (T t) (H p1) (W p2)",
            p1=self.patch_spatial,
            p2=self.patch_spatial,
            H=H_before_patchify // self.patch_spatial,
            W=W_before_patchify // self.patch_spatial,
            t=self.patch_temporal,
            B=B,
        )
        return x_B_D_T_H_W

    def forward_before_blocks(
        self,
        x: torch.Tensor,
        timesteps: torch.Tensor,
        crossattn_emb: torch.Tensor,
        crossattn_mask: Optional[torch.Tensor] = None,
        fps: Optional[torch.Tensor] = None,
        image_size: Optional[torch.Tensor] = None,
        padding_mask: Optional[torch.Tensor] = None,
        scalar_feature: Optional[torch.Tensor] = None,
        data_type: Optional[DataType] = DataType.VIDEO,
        latent_condition: Optional[torch.Tensor] = None,
        latent_condition_sigma: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> torch.Tensor:
        """
        Args:
            x: (B, C, T, H, W) tensor of spatial-temp inputs
            timesteps: (B, ) tensor of timesteps
            crossattn_emb: (B, N, D) tensor of cross-attention embeddings
            crossattn_mask: (B, N) tensor of cross-attention masks
        """
        del kwargs
        assert isinstance(
            data_type, DataType
        ), f"Expected DataType, got {type(data_type)}. We need discuss this flag later."
        original_shape = x.shape
        x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = self.prepare_embedded_sequence(
            x,
            fps=fps,
            padding_mask=padding_mask,
            latent_condition=latent_condition,
            latent_condition_sigma=latent_condition_sigma,
        )
        # logging affline scale information
        affline_scale_log_info = {}

        timesteps_B_D, adaln_lora_B_3D = self.t_embedder(timesteps.flatten())
        affline_emb_B_D = timesteps_B_D
        affline_scale_log_info["timesteps_B_D"] = timesteps_B_D.detach()

        if scalar_feature is not None:
            raise NotImplementedError("Scalar feature is not implemented yet.")
            timesteps_B_D = timesteps_B_D + scalar_feature.mean(dim=1)

        if self.additional_timestamp_channels:
            additional_cond_B_D = self.prepare_additional_timestamp_embedder(
                bs=x.shape[0],
                fps=fps,
                h=image_size[:, 0],
                w=image_size[:, 1],
                org_h=image_size[:, 2],
                org_w=image_size[:, 3],
            )

            affline_emb_B_D += additional_cond_B_D
            affline_scale_log_info["additional_cond_B_D"] = additional_cond_B_D.detach()

        affline_scale_log_info["affline_emb_B_D"] = affline_emb_B_D.detach()
        affline_emb_B_D = self.affline_norm(affline_emb_B_D)

        # for logging purpose
        self.affline_scale_log_info = affline_scale_log_info
        self.affline_emb = affline_emb_B_D
        self.crossattn_emb = crossattn_emb
        self.crossattn_mask = crossattn_mask

        if self.use_cross_attn_mask:
            crossattn_mask = crossattn_mask[:, None, None, :].to(dtype=torch.bool)  # [B, 1, 1, length]
        else:
            crossattn_mask = None

        if self.blocks["block0"].x_format == "THWBD":
            x = rearrange(x_B_T_H_W_D, "B T H W D -> T H W B D")
            if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None:
                extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = rearrange(
                    extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D, "B T H W D -> T H W B D"
                )
            crossattn_emb = rearrange(crossattn_emb, "B M D -> M B D")

            if crossattn_mask:
                crossattn_mask = rearrange(crossattn_mask, "B M -> M B")

            if self.sequence_parallel:
                tp_group = parallel_state.get_tensor_model_parallel_group()
                # Sequence parallel requires the input tensor to be scattered along the first dimension.
                assert self.block_config == "FA-CA-MLP"  # Only support this block config for now
                T, H, W, B, D = x.shape
                # variable name x_T_H_W_B_D is no longer valid. x is reshaped to THW*1*1*b*D and will be reshaped back in FinalLayer
                x = x.view(T * H * W, 1, 1, B, D)
                assert x.shape[0] % parallel_state.get_tensor_model_parallel_world_size() == 0
                x = scatter_along_first_dim(x, tp_group)

                if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None:
                    extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.view(
                        T * H * W, 1, 1, B, D
                    )
                    extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = scatter_along_first_dim(
                        extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D, tp_group
                    )

        elif self.blocks["block0"].x_format == "BTHWD":
            x = x_B_T_H_W_D
        else:
            raise ValueError(f"Unknown x_format {self.blocks[0].x_format}")
        output = {
            "x": x,
            "affline_emb_B_D": affline_emb_B_D,
            "crossattn_emb": crossattn_emb,
            "crossattn_mask": crossattn_mask,
            "rope_emb_L_1_1_D": rope_emb_L_1_1_D,
            "adaln_lora_B_3D": adaln_lora_B_3D,
            "original_shape": original_shape,
            "extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D": extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
        }
        return output

    def forward_blocks_regular(
        self,
        x,
        affline_emb_B_D,
        crossattn_emb,
        crossattn_mask,
        rope_emb_L_1_1_D,
        adaln_lora_B_3D,
        extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
        feature_indices,
        original_shape,
        x_ctrl,
        return_features_early,
    ):
        features = []
        for name, block in self.blocks.items():
            assert (
                self.blocks["block0"].x_format == block.x_format
            ), f"First block has x_format {self.blocks[0].x_format}, got {block.x_format}"
            x = block(
                x,
                affline_emb_B_D,
                crossattn_emb,
                crossattn_mask,
                rope_emb_L_1_1_D=rope_emb_L_1_1_D,
                adaln_lora_B_3D=adaln_lora_B_3D,
                extra_per_block_pos_emb=extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
            )

            # Extract features
            block_idx = int(name.split("block")[-1])
            if block_idx in feature_indices:
                B, C, T, H, W = original_shape
                H = H // self.patch_spatial
                W = W // self.patch_spatial
                T = T // self.patch_temporal
                if self.sequence_parallel:
                    x_feat = gather_along_first_dim(x, parallel_state.get_tensor_model_parallel_group())
                    x_B_T_H_W_D = rearrange(x_feat, "(T H W) 1 1 B D -> B T H W D", T=T, H=H, W=W)
                else:
                    x_feat = x
                    if self.blocks["block0"].x_format == "THWBD":
                        x_B_T_H_W_D = rearrange(x_feat, "T H W B D -> B T H W D", T=T, H=H, W=W)
                    elif self.blocks["block0"].x_format == "BTHWD":
                        x_B_T_H_W_D = x_feat
                    else:
                        raise ValueError(f"Unknown x_format {self.blocks[-1].x_format}")

                features.append(x_B_T_H_W_D)

            if x_ctrl is not None and name in x_ctrl:
                x = x + x_ctrl[name]
            # If we have all of the features, we can exit early
            if return_features_early and len(features) == len(feature_indices):
                return features

        if self.blocks["block0"].x_format == "THWBD":
            x_B_T_H_W_D = rearrange(x, "T H W B D -> B T H W D")
        elif self.blocks["block0"].x_format == "BTHWD":
            x_B_T_H_W_D = x
        else:
            raise ValueError(f"Unknown x_format {self.blocks[-1].x_format}")

        x_B_D_T_H_W = self.decoder_head(
            x_B_T_H_W_D=x_B_T_H_W_D,
            emb_B_D=affline_emb_B_D,
            crossattn_emb=None,
            origin_shape=original_shape,
            crossattn_mask=None,
            adaln_lora_B_3D=adaln_lora_B_3D,
        )

        if len(feature_indices) == 0:
            # no features requested, return only the model output
            return x_B_D_T_H_W
        else:
            # score and features； score, features
            return x_B_D_T_H_W, features

    def forward_blocks_memory_save(
        self,
        x,
        affline_emb_B_D,
        crossattn_emb,
        crossattn_mask,
        rope_emb_L_1_1_D,
        adaln_lora_B_3D,
        extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
        feature_indices,
        original_shape,
        x_ctrl,
        return_features_early,
    ):
        x_before_gate = 0
        x_skip = rearrange(x, "T H W B D -> (T H W) B D")
        assert self.blocks["block0"].x_format == "THWBD"
        if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None:
            extra_per_block_pos_emb = rearrange(extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D, "T H W B D -> (T H W) B D")
        else:
            extra_per_block_pos_emb = None
        gate_L_B_D = 1.0

        features = []
        for name, block in self.blocks.items():
            gate_L_B_D, x_before_gate, x_skip = block(
                x_before_gate,
                x_skip,
                gate_L_B_D,
                affline_emb_B_D,
                crossattn_emb,
                crossattn_mask,
                rope_emb_L_1_1_D=rope_emb_L_1_1_D,
                adaln_lora_B_3D=adaln_lora_B_3D,
                extra_per_block_pos_emb=extra_per_block_pos_emb,
            )

            # Extract features.
            # Convert the block index in the memory save mode to the block index in the regular mode.
            block_idx = int(name.split("block")[-1]) - 1
            if block_idx in feature_indices:
                B, C, T_before_patchify, H_before_patchify, W_before_patchify = original_shape
                H = H_before_patchify // self.patch_spatial
                W = W_before_patchify // self.patch_spatial
                T = T_before_patchify // self.patch_temporal
                if self.sequence_parallel:
                    x_feat = gather_along_first_dim(x_skip, parallel_state.get_tensor_model_parallel_group())
                    x_B_T_H_W_D = rearrange(x_feat, "(T H W) 1 1 B D -> B T H W D", T=T, H=H, W=W)
                else:
                    x_feat = x_skip
                    x_B_T_H_W_D = rearrange(x_feat, "(T H W) B D -> B T H W D", T=T, H=H, W=W)

                features.append(x_B_T_H_W_D)

            new_name = f"block{block_idx}"
            if x_ctrl is not None and new_name in x_ctrl:
                x_ctrl_ = x_ctrl[new_name]
                x_ctrl_ = rearrange(x_ctrl_, "T H W B D -> (T H W) B D")
                x_skip = x_skip + x_ctrl_
            # If we have all of the features, we can exit early
            if return_features_early and len(features) == len(feature_indices):
                return features

        x_THW_B_D_before_gate = x_before_gate
        x_THW_B_D_skip = x_skip

        B, C, T_before_patchify, H_before_patchify, W_before_patchify = original_shape
        x_BT_HW_D_before_gate = rearrange(
            x_THW_B_D_before_gate,
            "(T H W) B D -> (B T) (H W) D",
            T=T_before_patchify // self.patch_temporal,
            H=H_before_patchify // self.patch_spatial,
            W=W_before_patchify // self.patch_spatial,
        )
        x_BT_HW_D_skip = rearrange(
            x_THW_B_D_skip,
            "(T H W) B D -> (B T) (H W) D",
            T=T_before_patchify // self.patch_temporal,
            H=H_before_patchify // self.patch_spatial,
            W=W_before_patchify // self.patch_spatial,
        )

        x_BT_HW_D = self.final_layer.forward_with_memory_save(
            x_BT_HW_D_before_gate=x_BT_HW_D_before_gate,
            x_BT_HW_D_skip=x_BT_HW_D_skip,
            gate_L_B_D=gate_L_B_D,
            emb_B_D=affline_emb_B_D,
            adaln_lora_B_3D=adaln_lora_B_3D,
        )

        # This is to ensure x_BT_HW_D has the correct shape because
        # when we merge T, H, W into one dimension, x_BT_HW_D has shape (B * T * H * W, 1*1, D).
        x_BT_HW_D = x_BT_HW_D.view(
            B * T_before_patchify // self.patch_temporal,
            H_before_patchify // self.patch_spatial * W_before_patchify // self.patch_spatial,
            -1,
        )
        x_B_D_T_H_W = rearrange(
            x_BT_HW_D,
            "(B T) (H W) (p1 p2 t C) -> B C (T t) (H p1) (W p2)",
            p1=self.patch_spatial,
            p2=self.patch_spatial,
            H=H_before_patchify // self.patch_spatial,
            W=W_before_patchify // self.patch_spatial,
            t=self.patch_temporal,
            B=B,
        )
        if len(feature_indices) == 0:
            # no features requested, return only the model output
            return x_B_D_T_H_W
        else:
            # score and features； score, features
            return x_B_D_T_H_W, features

    def forward(
        self,
        x: torch.Tensor,
        timesteps: torch.Tensor,
        crossattn_emb: torch.Tensor,
        crossattn_mask: Optional[torch.Tensor] = None,
        fps: Optional[torch.Tensor] = None,
        image_size: Optional[torch.Tensor] = None,
        padding_mask: Optional[torch.Tensor] = None,
        scalar_feature: Optional[torch.Tensor] = None,
        data_type: Optional[DataType] = DataType.VIDEO,
        x_ctrl: Optional[dict] = None,
        latent_condition: Optional[torch.Tensor] = None,
        latent_condition_sigma: Optional[torch.Tensor] = None,
        feature_indices: Optional[Container[int]] = None,
        return_features_early: bool = False,
        condition_video_augment_sigma: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> torch.Tensor | List[torch.Tensor] | Tuple[torch.Tensor, List[torch.Tensor]]:
        """
        Args:
            x: (B, C, T, H, W) tensor of spatial-temp inputs
            timesteps: (B, ) tensor of timesteps
            crossattn_emb: (B, N, D) tensor of cross-attention embeddings
            crossattn_mask: (B, N) tensor of cross-attention masks
            feature_indices: A set of feature indices (a set of integers) decides which blocks
                to extract features from. If the set is non-empty, then features will be returned.
                By default, feature_indices=None means extract no features.
            return_features_early: If true, the forward pass returns the features once the set is complete.
                This means the forward pass will not finish completely and no final output is returned.
            condition_video_augment_sigma: (B,) used in lvg(long video generation), we add noise with this sigma to augment condition input, the lvg model will condition on the condition_video_augment_sigma value;
                we need forward_before_blocks pass to the forward_before_blocks function.
        """
        if feature_indices is None:
            feature_indices = {}
        if return_features_early and len(feature_indices) == 0:
            # Exit immediately if user requested this.
            return []

        inputs = self.forward_before_blocks(
            x=x,
            timesteps=timesteps,
            crossattn_emb=crossattn_emb,
            crossattn_mask=crossattn_mask,
            fps=fps,
            image_size=image_size,
            padding_mask=padding_mask,
            scalar_feature=scalar_feature,
            data_type=data_type,
            latent_condition=latent_condition,
            latent_condition_sigma=latent_condition_sigma,
            condition_video_augment_sigma=condition_video_augment_sigma,
            **kwargs,
        )
        x, affline_emb_B_D, crossattn_emb, crossattn_mask, rope_emb_L_1_1_D, adaln_lora_B_3D, original_shape = (
            inputs["x"],
            inputs["affline_emb_B_D"],
            inputs["crossattn_emb"],
            inputs["crossattn_mask"],
            inputs["rope_emb_L_1_1_D"],
            inputs["adaln_lora_B_3D"],
            inputs["original_shape"],
        )
        extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = inputs["extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D"]
        if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None:
            assert (
                x.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape
            ), f"{x.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape} {original_shape}"

        if self.use_memory_save:
            return self.forward_blocks_memory_save(
                x,
                affline_emb_B_D,
                crossattn_emb,
                crossattn_mask,
                rope_emb_L_1_1_D,
                adaln_lora_B_3D,
                extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
                feature_indices,
                original_shape,
                x_ctrl,
                return_features_early,
            )

        return self.forward_blocks_regular(
            x,
            affline_emb_B_D,
            crossattn_emb,
            crossattn_mask,
            rope_emb_L_1_1_D,
            adaln_lora_B_3D,
            extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
            feature_indices,
            original_shape,
            x_ctrl,
            return_features_early,
        )

    @property
    def fsdp_wrap_block_cls(self):
        return DITBuildingBlock

    def enable_context_parallel(self, cp_group: ProcessGroup):
        cp_ranks = get_process_group_ranks(cp_group)
        cp_size = len(cp_ranks)
        # Set these attributes for spliting the data after embedding.
        self.cp_group = cp_group
        # Set these attributes for computing the loss.
        self.cp_size = cp_size

        self.pos_embedder.enable_context_parallel(cp_group)
        if self.extra_per_block_abs_pos_emb:
            self.extra_pos_embedder.enable_context_parallel(cp_group)
        # Loop through the model to set up context parallel.
        for block in self.blocks.values():
            for layer in block.blocks:
                if layer.block_type in ["mlp", "ff"]:
                    continue
                elif layer.block_type in ["cross_attn", "ca"]:
                    continue
                else:
                    layer.block.attn.attn_op.set_context_parallel_group(cp_group, cp_ranks, torch.cuda.Stream())

        log.debug(f"[CP] Enable context parallelism with size {cp_size}")

    def disable_context_parallel(self):
        self.cp_group = None
        self.cp_size = None

        self.pos_embedder.disable_context_parallel()
        if self.extra_per_block_abs_pos_emb:
            self.extra_pos_embedder.disable_context_parallel()

        # Loop through the model to disable context parallel.
        for block in self.blocks.values():
            for layer in block.blocks:
                if layer.block_type in ["mlp", "ff"]:
                    continue
                elif layer.block_type in ["cross_attn", "ca"]:
                    continue
                else:
                    layer.block.attn.attn_op.cp_group = None
                    layer.block.attn.attn_op.cp_ranks = None
                    layer.block.attn.attn_op.cp_stream = None

        log.debug("[CP] Disable context parallelism.")

    def enable_sequence_parallel(self):
        self._set_sequence_parallel(True)

    def disable_sequence_parallel(self):
        self._set_sequence_parallel(False)

    def _set_sequence_parallel(self, status: bool):
        self.sequence_parallel = status
        self.final_layer.sequence_parallel = status
        for block in self.blocks.values():
            for layer in block.blocks:
                if layer.block_type in ["full_attn", "fa", "cross_attn", "ca"]:
                    layer.block.attn.to_q[0].sequence_parallel = status
                    layer.block.attn.to_k[0].sequence_parallel = status
                    layer.block.attn.to_v[0].sequence_parallel = status
                    layer.block.attn.to_out[0].sequence_parallel = status
                    layer.block.attn.attn_op.sequence_parallel = status
                elif layer.block_type in ["mlp", "ff"]:
                    layer.block.layer1.sequence_parallel = status
                    layer.block.layer2.sequence_parallel = status
                else:
                    raise ValueError(f"Unknown block type {layer.block_type}")

    @property
    def is_context_parallel_enabled(self):
        return self.cp_group is not None