src/model/spacetime.py

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

import sys
sys.path.append('./')

from einops import rearrange, repeat
from model.dit import *
from diffusers.models.embeddings import LabelEmbedding 

class PointEmbed(nn.Module):
    def __init__(self, hidden_dim=96, dim=512):
        super().__init__()

        assert hidden_dim % 6 == 0

        self.embedding_dim = hidden_dim
        e = torch.pow(2, torch.arange(self.embedding_dim // 6)).float() * np.pi
        e = torch.stack([
            torch.cat([e, torch.zeros(self.embedding_dim // 6),
                        torch.zeros(self.embedding_dim // 6)]),
            torch.cat([torch.zeros(self.embedding_dim // 6), e,
                        torch.zeros(self.embedding_dim // 6)]),
            torch.cat([torch.zeros(self.embedding_dim // 6),
                        torch.zeros(self.embedding_dim // 6), e]),
        ])
        self.register_buffer('basis', e)  # 3 x 16

        self.mlp = nn.Linear(self.embedding_dim+3, dim)

    @staticmethod
    def embed(input, basis):
        projections = torch.einsum(
            'bnd,de->bne', input, basis)
        embeddings = torch.cat([projections.sin(), projections.cos()], dim=2)
        return embeddings
    
    def forward(self, input):
        # input: B x N x 3
        embed = self.mlp(torch.cat([self.embed(input, self.basis), input], dim=2)) # B x N x C
        return embed

class AdaLayerNorm(nn.Module):
    r"""
    Norm layer modified to incorporate timestep embeddings.

    Parameters:
        embedding_dim (`int`): The size of each embedding vector.
        num_embeddings (`int`, *optional*): The size of the embeddings dictionary.
        output_dim (`int`, *optional*):
        norm_elementwise_affine (`bool`, defaults to `False):
        norm_eps (`bool`, defaults to `False`):
        chunk_dim (`int`, defaults to `0`):
    """

    def __init__(
        self,
        embedding_dim: int,
        num_embeddings: Optional[int] = None,
        output_dim: Optional[int] = None,
        norm_elementwise_affine: bool = False,
        norm_eps: float = 1e-5,
        chunk_dim: int = 0,
    ):
        super().__init__()

        self.chunk_dim = chunk_dim
        output_dim = output_dim or embedding_dim * 2

        if num_embeddings is not None:
            self.emb = nn.Embedding(num_embeddings, embedding_dim)
        else:
            self.emb = None

        self.silu = nn.SiLU()
        self.linear = nn.Linear(embedding_dim, output_dim)
        self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)

    def forward(
        self, x: torch.Tensor, timestep: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        if self.emb is not None:
            temb = self.emb(timestep)

        temb = self.linear(self.silu(temb))

        if self.chunk_dim == 1:
            # This is a bit weird why we have the order of "shift, scale" here and "scale, shift" in the
            # other if-branch. This branch is specific to CogVideoX for now.
            shift, scale = temb.chunk(2, dim=1)
            shift = shift[:, None, :]
            scale = scale[:, None, :]
        else:
            scale, shift = temb.chunk(2, dim=0)

        x = self.norm(x) * (1 + scale) + shift
        return x


@maybe_allow_in_graph
class SpatialTemporalTransformerBlock(nn.Module):
    r"""
    Transformer block used in [CogVideoX](https://github.com/THUDM/CogVideo) model.

    Parameters:
        dim (`int`):
            The number of channels in the input and output.
        num_attention_heads (`int`):
            The number of heads to use for multi-head attention.
        attention_head_dim (`int`):
            The number of channels in each head.
        time_embed_dim (`int`):
            The number of channels in timestep embedding.
        dropout (`float`, defaults to `0.0`):
            The dropout probability to use.
        activation_fn (`str`, defaults to `"gelu-approximate"`):
            Activation function to be used in feed-forward.
        attention_bias (`bool`, defaults to `False`):
            Whether or not to use bias in attention projection layers.
        qk_norm (`bool`, defaults to `True`):
            Whether or not to use normalization after query and key projections in Attention.
        norm_elementwise_affine (`bool`, defaults to `True`):
            Whether to use learnable elementwise affine parameters for normalization.
        norm_eps (`float`, defaults to `1e-5`):
            Epsilon value for normalization layers.
        final_dropout (`bool` defaults to `False`):
            Whether to apply a final dropout after the last feed-forward layer.
        ff_inner_dim (`int`, *optional*, defaults to `None`):
            Custom hidden dimension of Feed-forward layer. If not provided, `4 * dim` is used.
        ff_bias (`bool`, defaults to `True`):
            Whether or not to use bias in Feed-forward layer.
        attention_out_bias (`bool`, defaults to `True`):
            Whether or not to use bias in Attention output projection layer.
    """

    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        time_embed_dim: int,
        dropout: float = 0.0,
        activation_fn: str = "gelu-approximate",
        attention_bias: bool = False,
        qk_norm: bool = True,
        norm_elementwise_affine: bool = True,
        norm_eps: float = 1e-5,
        final_dropout: bool = True,
        ff_inner_dim: Optional[int] = None,
        ff_bias: bool = True,
        attention_out_bias: bool = True,
    ):
        super().__init__()

        # 1. Self Attention
        self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)

        self.attn1 = Attention(
            query_dim=dim,
            dim_head=attention_head_dim,
            heads=num_attention_heads,
            qk_norm="layer_norm" if qk_norm else None,
            eps=1e-6,
            bias=attention_bias,
            out_bias=attention_out_bias,
            processor=CogVideoXAttnProcessor2_0(),
        )

        self.norm_temp = AdaLayerNorm(dim, chunk_dim=1)
        self.attn_temp = Attention(
            query_dim=dim,
            dim_head=attention_head_dim,
            heads=num_attention_heads,
            qk_norm="layer_norm" if qk_norm else None,
            eps=1e-6,
            bias=attention_bias,
            out_bias=attention_out_bias,
        )

        # 2. Feed Forward
        self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)

        self.ff = FeedForward(
            dim,
            dropout=dropout,
            activation_fn=activation_fn,
            final_dropout=final_dropout,
            inner_dim=ff_inner_dim,
            bias=ff_bias,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ) -> torch.Tensor:
        temb_in = temb
        text_seq_length = encoder_hidden_states.size(1)

        B, F, N, C = hidden_states.shape
        hidden_states = hidden_states.reshape(-1, N, C)
        if encoder_hidden_states.shape[0] != B * F:
            encoder_hidden_states = encoder_hidden_states.repeat_interleave(F, 0)
        temb = temb_in.repeat_interleave(F, 0)

        # norm & modulate
        norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
            hidden_states, encoder_hidden_states, temb
        )

        # Spatial Attention
        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
            hidden_states=norm_hidden_states,
            encoder_hidden_states=norm_encoder_hidden_states,
            image_rotary_emb=image_rotary_emb,
        )

        hidden_states = hidden_states + gate_msa * attn_hidden_states
        encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states

        # norm & modulate
        norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(
            hidden_states, encoder_hidden_states, temb
        )

        # feed-forward
        norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
        ff_output = self.ff(norm_hidden_states)

        hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:]
        encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length]

        ## Time Attention
        hidden_states = rearrange(hidden_states, '(b f) n c -> (b n) f c', f=F)
        temb = temb_in.repeat_interleave(N, 0)
        norm_hidden_states = self.norm_temp(hidden_states, temb=temb)
        hidden_states = self.attn_temp(norm_hidden_states) + hidden_states
        hidden_states = rearrange(hidden_states, '(b n) f c -> b f n c', n=N)

        # hidden_states = rearrange(hidden_states, '(b f) n c -> b f n c', f=F)

        return hidden_states, encoder_hidden_states

@maybe_allow_in_graph
class SpatialOnlyTransformerBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        time_embed_dim: int,
        dropout: float = 0.0,
        activation_fn: str = "gelu-approximate",
        attention_bias: bool = False,
        qk_norm: bool = True,
        norm_elementwise_affine: bool = True,
        norm_eps: float = 1e-5,
        final_dropout: bool = True,
        ff_inner_dim: Optional[int] = None,
        ff_bias: bool = True,
        attention_out_bias: bool = True,
    ):
        super().__init__()

        # 1. Self Attention
        self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)

        self.attn1 = Attention(
            query_dim=dim,
            dim_head=attention_head_dim,
            heads=num_attention_heads,
            qk_norm="layer_norm" if qk_norm else None,
            eps=1e-6,
            bias=attention_bias,
            out_bias=attention_out_bias,
            processor=CogVideoXAttnProcessor2_0(),
        )

        # 2. Feed Forward
        self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)

        self.ff = FeedForward(
            dim,
            dropout=dropout,
            activation_fn=activation_fn,
            final_dropout=final_dropout,
            inner_dim=ff_inner_dim,
            bias=ff_bias,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ) -> torch.Tensor:
        temb_in = temb
        text_seq_length = encoder_hidden_states.size(1)

        B, F, N, C = hidden_states.shape
        hidden_states = hidden_states.reshape(-1, N, C)
        if encoder_hidden_states.shape[0] != B * F:
            encoder_hidden_states = encoder_hidden_states.repeat_interleave(F, 0)
        temb = temb_in.repeat_interleave(F, 0)

        # norm & modulate
        norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
            hidden_states, encoder_hidden_states, temb
        )

        # Spatial Attention
        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
            hidden_states=norm_hidden_states,
            encoder_hidden_states=norm_encoder_hidden_states,
            image_rotary_emb=image_rotary_emb,
        )

        hidden_states = hidden_states + gate_msa * attn_hidden_states
        encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states

        # norm & modulate
        norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(
            hidden_states, encoder_hidden_states, temb
        )

        # feed-forward
        norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
        ff_output = self.ff(norm_hidden_states)

        hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:]
        encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length]

        hidden_states = rearrange(hidden_states, '(b f) n c -> b f n c', f=F)

        return hidden_states, encoder_hidden_states

@maybe_allow_in_graph
class TemporalOnlyTransformerBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        time_embed_dim: int,
        dropout: float = 0.0,
        activation_fn: str = "gelu-approximate",
        attention_bias: bool = False,
        qk_norm: bool = True,
        norm_elementwise_affine: bool = True,
        norm_eps: float = 1e-5,
        final_dropout: bool = True,
        ff_inner_dim: Optional[int] = None,
        ff_bias: bool = True,
        attention_out_bias: bool = True,
    ):
        super().__init__()

        # 1. Self Attention
        # self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)

        # self.attn1 = Attention(
        #     query_dim=dim,
        #     dim_head=attention_head_dim,
        #     heads=num_attention_heads,
        #     qk_norm="layer_norm" if qk_norm else None,
        #     eps=1e-6,
        #     bias=attention_bias,
        #     out_bias=attention_out_bias,
        #     processor=CogVideoXAttnProcessor2_0(),
        # )

        self.norm_temp = AdaLayerNorm(dim, chunk_dim=1)
        self.attn_temp = Attention(
            query_dim=dim,
            dim_head=attention_head_dim,
            heads=num_attention_heads,
            qk_norm="layer_norm" if qk_norm else None,
            eps=1e-6,
            bias=attention_bias,
            out_bias=attention_out_bias,
        )

        # 2. Feed Forward
        self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)

        self.ff = FeedForward(
            dim,
            dropout=dropout,
            activation_fn=activation_fn,
            final_dropout=final_dropout,
            inner_dim=ff_inner_dim,
            bias=ff_bias,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ) -> torch.Tensor:
        temb_in = temb
        text_seq_length = encoder_hidden_states.size(1)

        B, F, N, C = hidden_states.shape
        hidden_states = hidden_states.reshape(-1, N, C)
        if encoder_hidden_states.shape[0] != B * F:
            encoder_hidden_states = encoder_hidden_states.repeat_interleave(F, 0)
        temb = temb_in.repeat_interleave(F, 0)

        # # norm & modulate
        # norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
        #     hidden_states, encoder_hidden_states, temb
        # )

        # # Spatial Attention
        # attn_hidden_states, attn_encoder_hidden_states = self.attn1(
        #     hidden_states=norm_hidden_states,
        #     encoder_hidden_states=norm_encoder_hidden_states,
        #     image_rotary_emb=image_rotary_emb,
        # )

        # hidden_states = hidden_states + gate_msa * attn_hidden_states
        # encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states

        # norm & modulate
        norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(
            hidden_states, encoder_hidden_states, temb
        )

        # feed-forward
        norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
        ff_output = self.ff(norm_hidden_states)

        hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:]
        encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length]

        ## Time Attention
        hidden_states = rearrange(hidden_states, '(b f) n c -> (b n) f c', f=F)
        temb = temb_in.repeat_interleave(N, 0)
        norm_hidden_states = self.norm_temp(hidden_states, temb=temb)
        hidden_states = self.attn_temp(norm_hidden_states) + hidden_states
        hidden_states = rearrange(hidden_states, '(b n) f c -> b f n c', n=N)

        # hidden_states = rearrange(hidden_states, '(b f) n c -> b f n c', f=F)

        return hidden_states, encoder_hidden_states


@maybe_allow_in_graph
class SpatialTemporalTransformerBlockv2(nn.Module):
    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        time_embed_dim: int,
        dropout: float = 0.0,
        activation_fn: str = "gelu-approximate",
        attention_bias: bool = False,
        qk_norm: bool = True,
        norm_elementwise_affine: bool = True,
        norm_eps: float = 1e-5,
        final_dropout: bool = True,
        ff_inner_dim: Optional[int] = None,
        ff_bias: bool = True,
        attention_out_bias: bool = True,
    ):
        super().__init__()

        # 1. Self Attention
        self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)

        self.attn1 = Attention(
            query_dim=dim,
            dim_head=attention_head_dim,
            heads=num_attention_heads,
            qk_norm="layer_norm" if qk_norm else None,
            eps=1e-6,
            bias=attention_bias,
            out_bias=attention_out_bias,
            processor=CogVideoXAttnProcessor2_0(),
        )

        self.norm_temp = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
        self.attn_temp = Attention(
            query_dim=dim,
            dim_head=attention_head_dim,
            heads=num_attention_heads,
            qk_norm="layer_norm" if qk_norm else None,
            eps=1e-6,
            bias=attention_bias,
            out_bias=attention_out_bias,
            processor=CogVideoXAttnProcessor2_0(),
        )

        # 2. Feed Forward
        self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)

        self.ff = FeedForward(
            dim,
            dropout=dropout,
            activation_fn=activation_fn,
            final_dropout=final_dropout,
            inner_dim=ff_inner_dim,
            bias=ff_bias,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        encoder_hidden_states_time: torch.Tensor,
        temb: torch.Tensor,
        indices: Optional[torch.Tensor] = None,
        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ) -> torch.Tensor:
        temb_in = temb
        text_seq_length = encoder_hidden_states.size(1)

        B, F, N, C = hidden_states.shape
        hidden_states = hidden_states.reshape(-1, N, C)
        if encoder_hidden_states.shape[0] != B * F:
            encoder_hidden_states = encoder_hidden_states.repeat_interleave(F, 0)
        temb = temb_in.repeat_interleave(F, 0)

        # norm & modulate
        norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
            hidden_states, encoder_hidden_states, temb
        )

        # Spatial Attention
        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
            hidden_states=norm_hidden_states,
            encoder_hidden_states=norm_encoder_hidden_states,
            indices=indices,
            image_rotary_emb=image_rotary_emb,
        )

        hidden_states = hidden_states + gate_msa * attn_hidden_states
        encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states

        # norm & modulate
        norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(
            hidden_states, encoder_hidden_states, temb
        )

        # feed-forward
        norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
        ff_output = self.ff(norm_hidden_states)

        hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:]
        encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length]

        ## Time Attention
        hidden_states = rearrange(hidden_states, '(b f) n c -> (b n) f c', f=F)
        temb = temb_in.repeat_interleave(N, 0)
        norm_hidden_states, norm_encoder_hidden_states_time, gate_msa, enc_gate_msa = self.norm_temp(
            hidden_states, encoder_hidden_states_time, temb
        )
        attn_hidden_states, attn_encoder_hidden_states_time = self.attn_temp(
            hidden_states=norm_hidden_states,
            encoder_hidden_states=norm_encoder_hidden_states_time
        )
        hidden_states = hidden_states + gate_msa * attn_hidden_states
        encoder_hidden_states_time = encoder_hidden_states_time + enc_gate_msa * attn_encoder_hidden_states_time
        hidden_states = rearrange(hidden_states, '(b n) f c -> b f n c', n=N)

        return hidden_states, encoder_hidden_states, encoder_hidden_states_time

class SpaitalTemporalTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin):
    """
    A Transformer model for video-like data in [CogVideoX](https://github.com/THUDM/CogVideo).

    Parameters:
        num_attention_heads (`int`, defaults to `30`):
            The number of heads to use for multi-head attention.
        attention_head_dim (`int`, defaults to `64`):
            The number of channels in each head.
        in_channels (`int`, defaults to `16`):
            The number of channels in the input.
        out_channels (`int`, *optional*, defaults to `16`):
            The number of channels in the output.
        flip_sin_to_cos (`bool`, defaults to `True`):
            Whether to flip the sin to cos in the time embedding.
        time_embed_dim (`int`, defaults to `512`):
            Output dimension of timestep embeddings.
        ofs_embed_dim (`int`, defaults to `512`):
            Output dimension of "ofs" embeddings used in CogVideoX-5b-I2B in version 1.5
        text_embed_dim (`int`, defaults to `4096`):
            Input dimension of text embeddings from the text encoder.
        num_layers (`int`, defaults to `30`):
            The number of layers of Transformer blocks to use.
        dropout (`float`, defaults to `0.0`):
            The dropout probability to use.
        attention_bias (`bool`, defaults to `True`):
            Whether to use bias in the attention projection layers.
        sample_width (`int`, defaults to `90`):
            The width of the input latents.
        sample_height (`int`, defaults to `60`):
            The height of the input latents.
        sample_frames (`int`, defaults to `49`):
            The number of frames in the input latents. Note that this parameter was incorrectly initialized to 49
            instead of 13 because CogVideoX processed 13 latent frames at once in its default and recommended settings,
            but cannot be changed to the correct value to ensure backwards compatibility. To create a transformer with
            K latent frames, the correct value to pass here would be: ((K - 1) * temporal_compression_ratio + 1).
        patch_size (`int`, defaults to `2`):
            The size of the patches to use in the patch embedding layer.
        temporal_compression_ratio (`int`, defaults to `4`):
            The compression ratio across the temporal dimension. See documentation for `sample_frames`.
        max_text_seq_length (`int`, defaults to `226`):
            The maximum sequence length of the input text embeddings.
        activation_fn (`str`, defaults to `"gelu-approximate"`):
            Activation function to use in feed-forward.
        timestep_activation_fn (`str`, defaults to `"silu"`):
            Activation function to use when generating the timestep embeddings.
        norm_elementwise_affine (`bool`, defaults to `True`):
            Whether to use elementwise affine in normalization layers.
        norm_eps (`float`, defaults to `1e-5`):
            The epsilon value to use in normalization layers.
        spatial_interpolation_scale (`float`, defaults to `1.875`):
            Scaling factor to apply in 3D positional embeddings across spatial dimensions.
        temporal_interpolation_scale (`float`, defaults to `1.0`):
            Scaling factor to apply in 3D positional embeddings across temporal dimensions.
    """

    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(
        self,
        num_attention_heads: int = 8,
        attention_head_dim: int = 64,
        in_channels: int = 3,
        out_channels: Optional[int] = 3,
        flip_sin_to_cos: bool = True,
        freq_shift: int = 0,
        time_embed_dim: int = 512,
        ofs_embed_dim: Optional[int] = None,
        text_embed_dim: int = 4096,
        num_layers: int = 8,
        dropout: float = 0.0,
        attention_bias: bool = True,
        sample_points: int = 2048,
        sample_frames: int = 48,
        patch_size: int = 1,
        patch_size_t: Optional[int] = None,
        temporal_compression_ratio: int = 4,
        max_text_seq_length: int = 226,
        activation_fn: str = "gelu-approximate",
        timestep_activation_fn: str = "silu",
        norm_elementwise_affine: bool = True,
        norm_eps: float = 1e-5,
        spatial_interpolation_scale: float = 1.875,
        temporal_interpolation_scale: float = 1.0,
        use_positional_embeddings: bool = True,
        use_learned_positional_embeddings: bool = False,
        patch_bias: bool = True,
        cond_seq_length: int = 4,
        cond_seq_length_t: int = 2,
        transformer_block: str = "SpatialTemporalTransformerBlock",
        num_classes: int = 0,
        class_dropout_prob: float = 0.0,
    ):
        super().__init__()
        inner_dim = num_attention_heads * attention_head_dim

        if use_positional_embeddings and use_learned_positional_embeddings:
            raise ValueError(
                "There are no CogVideoX checkpoints available with disable rotary embeddings and learned positional "
                "embeddings. If you're using a custom model and/or believe this should be supported, please open an "
                "issue at https://github.com/huggingface/diffusers/issues."
            )

        self.embedding_dropout = nn.Dropout(dropout)

        # 2. Time embeddings and ofs embedding(Only CogVideoX1.5-5B I2V have)

        self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift)
        self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn)

        self.ofs_proj = None
        self.ofs_embedding = None
        if ofs_embed_dim:
            self.ofs_proj = Timesteps(ofs_embed_dim, flip_sin_to_cos, freq_shift)
            self.ofs_embedding = TimestepEmbedding(
                ofs_embed_dim, ofs_embed_dim, timestep_activation_fn
            )  # same as time embeddings, for ofs

        self.class_embedder = None
        if num_classes > 0:
            self.class_embedder = LabelEmbedding(num_classes, time_embed_dim, class_dropout_prob)

        self.transformer_block = transformer_block
        if transformer_block == "SpatialTemporalTransformerBlock":
            TransformerBlock = SpatialTemporalTransformerBlock
        elif transformer_block == "SpatialTemporalTransformerBlockv2":
            TransformerBlock = SpatialTemporalTransformerBlockv2
        elif transformer_block == "SpatialTemporalTransformerBlockv3":
            TransformerBlock = SpatialTemporalTransformerBlockv3
        elif transformer_block == "SpatialOnlyTransformerBlock":
            TransformerBlock = SpatialOnlyTransformerBlock
        elif transformer_block == "TemporalOnlyTransformerBlock":
            TransformerBlock = TemporalOnlyTransformerBlock

        # 3. Define spatio-temporal transformers blocks
        self.transformer_blocks = nn.ModuleList(
            [
                TransformerBlock(
                    dim=inner_dim,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    time_embed_dim=time_embed_dim,
                    dropout=dropout,
                    activation_fn=activation_fn,
                    attention_bias=attention_bias,
                    norm_elementwise_affine=norm_elementwise_affine,
                    norm_eps=norm_eps,
                )
                for _ in range(num_layers)
            ]
        )
        self.norm_final = nn.LayerNorm(inner_dim, norm_eps, norm_elementwise_affine)

        # 4. Output blocks
        self.norm_out = AdaLayerNorm(
            embedding_dim=time_embed_dim,
            output_dim=2 * inner_dim,
            norm_elementwise_affine=norm_elementwise_affine,
            norm_eps=norm_eps,
            chunk_dim=1,
        )

        if patch_size_t is None:
            # For CogVideox 1.0
            output_dim = patch_size * patch_size * out_channels
        else:
            # For CogVideoX 1.5
            output_dim = patch_size * patch_size * patch_size_t * out_channels

        self.proj_out = nn.Linear(inner_dim, output_dim)

        self.gradient_checkpointing = False
        
        if use_positional_embeddings or use_learned_positional_embeddings:
            self.embed_dim = num_attention_heads * attention_head_dim
            self.cond_seq_length = cond_seq_length
            self.cond_seq_length_t = cond_seq_length_t
            persistent = use_learned_positional_embeddings
            pos_embedding = self._get_positional_embeddings(sample_points, sample_frames)
            self.register_buffer("pos_embedding", pos_embedding, persistent=persistent)

    def _get_positional_embeddings(self, points: int, frames: int, device: Optional[torch.device] = None) -> torch.Tensor:
        pos_embedding = get_3d_sincos_pos_embed(
            self.embed_dim,
            points,
            frames,
            device=device,
            output_type="pt",
        )
        pos_embedding = pos_embedding.flatten(0, 1)
        joint_pos_embedding = pos_embedding.new_zeros(
            1, self.cond_seq_length + points * frames, self.embed_dim, requires_grad=False
        )
        joint_pos_embedding.data[:, self.cond_seq_length:].copy_(pos_embedding)
        return joint_pos_embedding
    
    def _set_gradient_checkpointing(self, module, value=False):
        self.gradient_checkpointing = value

    @property
    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
    def attn_processors(self) -> Dict[str, AttentionProcessor]:
        r"""
        Returns:
            `dict` of attention processors: A dictionary containing all attention processors used in the model with
            indexed by its weight name.
        """
        # set recursively
        processors = {}

        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
            if hasattr(module, "get_processor"):
                processors[f"{name}.processor"] = module.get_processor()

            for sub_name, child in module.named_children():
                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)

            return processors

        for name, module in self.named_children():
            fn_recursive_add_processors(name, module, processors)

        return processors

    def forward(
        self,
        hidden_states: torch.Tensor, # [batch_size]
        encoder_hidden_states: torch.Tensor, # [batch_size]
        timestep: Union[int, float, torch.LongTensor],
        timestep_cond: Optional[torch.Tensor] = None,
        class_labels: Optional[torch.Tensor] = None,
        force_drop_ids: Optional[torch.Tensor] = None,
        ofs: Optional[Union[int, float, torch.LongTensor]] = None,
        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        attention_kwargs: Optional[Dict[str, Any]] = None,
        indices: Optional[torch.LongTensor] = None,
        return_dict: bool = True,
    ):
        if attention_kwargs is not None:
            attention_kwargs = attention_kwargs.copy()
            lora_scale = attention_kwargs.pop("scale", 1.0)
        else:
            lora_scale = 1.0

        if USE_PEFT_BACKEND:
            # weight the lora layers by setting `lora_scale` for each PEFT layer
            scale_lora_layers(self, lora_scale)
        else:
            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
                logger.warning(
                    "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
                )
        
        # 1. Time embedding
        timesteps = timestep
        t_emb = self.time_proj(timesteps)

        # timesteps does not contain any weights and will always return f32 tensors
        # but time_embedding might actually be running in fp16. so we need to cast here.
        # there might be better ways to encapsulate this.
        t_emb = t_emb.to(dtype=hidden_states.dtype)
        emb = self.time_embedding(t_emb, timestep_cond)

        # TODO: check force drop id shape
        if self.class_embedder is not None:
            assert class_labels is not None
            class_labels = self.class_embedder(class_labels, force_drop_ids=force_drop_ids) # (N, D)
            emb = emb + class_labels

        if self.ofs_embedding is not None:
            ofs_emb = self.ofs_proj(ofs)
            ofs_emb = ofs_emb.to(dtype=hidden_states.dtype)
            ofs_emb = self.ofs_embedding(ofs_emb)
            emb = emb + ofs_emb
        
        B, F, N, C = hidden_states.shape
        full_seq = torch.cat([encoder_hidden_states, hidden_states.reshape(B, F*N, -1)], axis=1)

        # 2. Patch embedding
        pos_embedding = self.pos_embedding
        pos_embedding = pos_embedding.to(dtype=full_seq.dtype)
        hidden_states = full_seq + pos_embedding
        
        hidden_states = self.embedding_dropout(hidden_states)
        encoder_hidden_states = hidden_states[:, :self.cond_seq_length]
        hidden_states = hidden_states[:, self.cond_seq_length:].reshape(B, F, N, C)

        if self.transformer_block not in ["SpatialTemporalTransformerBlock", 'TemporalOnlyTransformerBlock', 'SpatialOnlyTransformerBlock']:
            encoder_hidden_states_time = hidden_states[:, :self.cond_seq_length_t]
            encoder_hidden_states_time = rearrange(encoder_hidden_states_time, 'b f n c -> (b n) f c')
            hidden_states = hidden_states[:, self.cond_seq_length_t:]

        # 3. Transformer blocks
        for i, block in enumerate(self.transformer_blocks):
            if torch.is_grad_enabled() and self.gradient_checkpointing:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)
                    return custom_forward

                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
                if self.transformer_block in ["SpatialTemporalTransformerBlock", 'TemporalOnlyTransformerBlock', 'SpatialOnlyTransformerBlock']:
                    hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(block),
                        hidden_states,
                        encoder_hidden_states,
                        emb,
                        image_rotary_emb,
                        **ckpt_kwargs,
                    )
                else:
                    hidden_states, encoder_hidden_states, encoder_hidden_states_time = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(block),
                        hidden_states,
                        encoder_hidden_states,
                        encoder_hidden_states_time,
                        emb,
                        image_rotary_emb,
                        indices=indices,
                        **ckpt_kwargs,
                    )
            else:
                if self.transformer_block in ["SpatialTemporalTransformerBlock", 'TemporalOnlyTransformerBlock', 'SpatialOnlyTransformerBlock']:
                    hidden_states, encoder_hidden_states = block(
                        hidden_states=hidden_states,
                        encoder_hidden_states=encoder_hidden_states,
                        temb=emb,
                        image_rotary_emb=image_rotary_emb,
                    )
                else:
                    hidden_states, encoder_hidden_states, encoder_hidden_states_time = block(
                        hidden_states=hidden_states,
                        encoder_hidden_states=encoder_hidden_states,
                        encoder_hidden_states_time=encoder_hidden_states_time,
                        temb=emb,
                        indices=indices,
                        image_rotary_emb=image_rotary_emb,
                    )
        hidden_states = rearrange(hidden_states, 'b f n c -> b (f n) c')
        # 4. Final block
        hidden_states = self.norm_final(hidden_states)
        hidden_states = self.norm_out(hidden_states, temb=emb)
        output = self.proj_out(hidden_states)

        if USE_PEFT_BACKEND:
            # remove `lora_scale` from each PEFT layer
            unscale_lora_layers(self, lora_scale)

        return output

class MDM_ST(nn.Module):
    def __init__(self, n_points, n_frame, n_feats, model_config):
        super().__init__()
        print('use new model')

        self.n_points = n_points
        self.n_feats = n_feats
        self.latent_dim = model_config.latent_dim
        self.cond_frame = 1 if model_config.frame_cond else 0
        self.frame_cond = model_config.frame_cond

        if model_config.get('point_embed', True):
            self.input_encoder = PointEmbed(dim=self.latent_dim)
        else:
            print('not using point embedding')
            self.input_encoder = nn.Linear(n_feats, self.latent_dim)
        self.mask_cond = model_config.get('mask_cond', False)
        if self.mask_cond:
            print('Use mask condition')
            self.mask_encoder = nn.Linear(1, self.latent_dim)
            self.cond_frame += 1
        self.pred_offset = model_config.get('pred_offset', True)
        self.num_neighbors = model_config.get('num_neighbors', 0)
        self.max_num_forces = model_config.get('max_num_forces', 1)
        self.model_config = model_config

        self.cond_seq_length = 2

        self.E_cond_encoder = nn.Linear(1, self.latent_dim)
        self.nu_cond_encoder = nn.Linear(1, self.latent_dim)
        self.force_as_token = model_config.get('force_as_token', True)
        self.force_as_latent = model_config.get('force_as_latent', False)

        if self.force_as_latent:
            self.input_encoder = nn.Linear(n_feats + 4 * self.max_num_forces, self.latent_dim)
        elif self.force_as_token:
            self.cond_seq_length += self.max_num_forces * 2
            self.force_cond_encoder = nn.Linear(3, self.latent_dim)
            self.drag_point_encoder = nn.Linear(3, self.latent_dim)
        else:
            self.cond_seq_length += 2
            self.force_cond_encoder = nn.Linear(3, self.latent_dim)
            self.drag_point_encoder = nn.Linear(3, self.latent_dim)

        self.gravity_emb = model_config.get('gravity_emb', False)
        if self.gravity_emb:
            self.gravity_embedding = nn.Embedding(2, self.latent_dim)
            self.cond_seq_length += 1

        if self.model_config.floor_cond:
            self.floor_encoder = nn.Linear(1, self.latent_dim)
            self.cond_seq_length += 1
        
        if self.model_config.coeff_cond:
            self.coeff_encoder = nn.Linear(1, self.latent_dim)
            self.cond_seq_length += 1

        self.num_mat = model_config.get('num_mat', 0)
        if model_config.class_token:
            self.class_embedding = nn.Embedding(model_config.num_mat, self.latent_dim)
            self.cond_seq_length += 1

        self.class_dropout_prob = model_config.get('class_dropout_prob', 0.0)
        self.dit = SpaitalTemporalTransformer(sample_points=n_points, sample_frames=n_frame+self.cond_frame, in_channels=n_feats,
            num_layers=model_config.n_layers, num_attention_heads=self.latent_dim // 64, time_embed_dim=self.latent_dim, cond_seq_length=self.cond_seq_length, cond_seq_length_t=self.cond_frame, transformer_block=model_config.transformer_block, num_classes=self.num_mat, class_dropout_prob=self.class_dropout_prob)
        
        self._init_weights()

    def _init_weights(self):
        if self.gravity_emb:
            nn.init.normal_(self.gravity_embedding.weight, mean=0.0, std=0.1)
    
    def enable_gradient_checkpointing(self):
        self.dit._set_gradient_checkpointing(True)

    def forward(self, x, timesteps, init_pc, force, E, nu, drag_mask, drag_point, floor_height, gravity_label=None, coeff=None, y=None, null_emb=None):
        
        """
        x: [batch_size, frame, n_points, n_feats], denoted x_t in the paper
        timesteps: [batch_size] (int)
        """
        
        bs, n_frame, n_points, n_feats = x.shape
        
        init_pc = init_pc.reshape(bs, n_points, n_feats)
        force = force.unsqueeze(1) if force.ndim == 2 else force
        drag_point = drag_point.unsqueeze(1) if drag_point.ndim == 2 else drag_point
        E = E.unsqueeze(1)
        nu = nu.unsqueeze(1)

        if self.num_neighbors > 0:
            rel_dist = torch.cdist(init_pc, init_pc)
            dist, indices = rel_dist.topk(self.num_neighbors, largest = False)
            indices = indices.repeat_interleave(n_frame, 0)
            # indices = torch.cat([indices, torch.tensor([2048, 2049, 2050, 2051])[None, None].repeat(bs*n_frame, n_points, 1).to(indices.device)], axis=2)
        else:
            indices = None
        
        if self.force_as_token:
            force_emb = self.force_cond_encoder(force) + self.gravity_embedding(gravity_label) if self.gravity_emb else self.force_cond_encoder(force)
            encoder_hidden_states = torch.cat([self.E_cond_encoder(E), self.nu_cond_encoder(nu)], axis=1)
            # force_info = torch.cat([force, drag_point], dim=-1) # (B, n_forces, 7)
            # force_tokens = self.force_cond_encoder(force_info)
            encoder_hidden_states = torch.cat([encoder_hidden_states, force_emb, self.drag_point_encoder(drag_point[..., :3])], axis=1)
        elif self.force_as_latent:
            encoder_hidden_states = torch.cat([self.E_cond_encoder(E), self.nu_cond_encoder(nu)], axis=1)
            force = force.unsqueeze(1).repeat(1, n_points, 1, 1) # (B, n_points, n_forces, 3)
            all_force = torch.cat([force, drag_mask.permute(0, 2, 1, 3)], dim=-1).reshape(bs, n_points, -1) # (B, n_points, n_forces, 4)
        else:
            encoder_hidden_states = torch.cat([self.force_cond_encoder(force), self.E_cond_encoder(E),
                self.nu_cond_encoder(nu), self.drag_point_encoder(drag_point[..., :3])], axis=1) 
        if self.gravity_emb:
            encoder_hidden_states = torch.cat([encoder_hidden_states, self.gravity_embedding(gravity_label)], axis=1)
        if self.model_config.class_token:
            class_labels = y.unsqueeze(1)
            class_labels = self.class_embedding(class_labels)
            encoder_hidden_states = torch.cat([encoder_hidden_states, class_labels], axis=1)
        if self.model_config.floor_cond:
            floor_height = floor_height.unsqueeze(1) if floor_height is not None else None
            encoder_hidden_states = torch.cat([encoder_hidden_states, self.floor_encoder(floor_height)], axis=1)
        if self.model_config.coeff_cond:
            coeff = coeff.unsqueeze(1) if coeff is not None else None
            encoder_hidden_states = torch.cat([encoder_hidden_states, self.coeff_encoder(coeff)], axis=1)
        if null_emb is not None:
            encoder_hidden_states = encoder_hidden_states * null_emb
        if self.frame_cond:
            x = torch.cat([init_pc.unsqueeze(1), x], axis=1) # Condition on first frame
        if self.force_as_latent:
            all_force = all_force.unsqueeze(1).repeat(1, x.shape[1], 1, 1) # (B, n_frame, n_points, n_forces*4)
            x = torch.cat([x, all_force], dim=-1) # (B, n_frame, n_points, n_feats+n_forces * 4)
            n_feats = x.shape[-1]
        hidden_states = self.input_encoder(x.reshape(-1, n_points,
            n_feats)).reshape(bs, -1, n_points, self.latent_dim)
        if self.mask_cond:
            mask = self.mask_encoder(drag_mask[:, :1])
            hidden_states = torch.cat([mask, hidden_states], axis=1)
        if self.model_config.transformer_block in ["SpatialTemporalTransformerBlock", "TemporalOnlyTransformerBlock", "SpatialOnlyTransformerBlock"]:
            output = self.dit(hidden_states, encoder_hidden_states, timesteps, class_labels=y).reshape(bs, -1, n_points, 3)[:, self.cond_frame:]
        else:
            output = self.dit(hidden_states, encoder_hidden_states, timesteps, indices=indices).reshape(bs, -1, n_points, 3)
        output = output + init_pc.unsqueeze(1) if self.pred_offset else output
            
        return output

if __name__ == "__main__":

    # Diffusion
    from omegaconf import OmegaConf
    from options import TestingConfig
    cfg_path = '../traj-diff/configs/eval.yaml'
    config_path = 'model_config.yaml'
    device = 'cuda'
    schema = OmegaConf.structured(TestingConfig)
    cfg = OmegaConf.load(cfg_path)
    cfg = OmegaConf.merge(schema, cfg)
    n_training_frames = cfg.train_dataset.n_training_frames
    n_frames_interval = cfg.train_dataset.n_frames_interval

    point_num = 2048
    frame_num = 24
    x = torch.randn(1, frame_num, point_num, 3).to(device).to(torch.float16)
    timesteps = torch.tensor([999]).int().to(device).to(torch.float16)
    init_pc = torch.randn(1, 1, point_num, 3).to(device).to(torch.float16)
    force = torch.randn(1, 3).to(device).to(torch.float16)
    E = torch.randn(1, 1).to(device).to(torch.float16)
    nu = torch.randn(1, 1).to(device).to(torch.float16)

    x = nn.Parameter(x)

    with torch.enable_grad():
    # with torch.no_grad():
        t_total = 0 
        for i in range(100):
            model = MDM_ST(point_num, frame_num, 3, cfg.model_config).to(device).to(torch.float16)
            model.train()
            import time
            t0 = time.time()
            output = model(x, timesteps, init_pc, force, E, nu, None, force, torch.zeros_like(E), torch.ones_like(E), None)
            loss = output.sum()
            loss.backward()
            t1 = time.time()
            if i > 10:
                t_total += t1 - t0
            print(t1 - t0)

    print("Average time: ", t_total / 90)