memo/models/transformer_3d.py

from dataclasses import dataclass
from typing import Optional

import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models import ModelMixin
from diffusers.utils import BaseOutput
from einops import rearrange, repeat
from torch import nn

from memo.models.attention import JointAudioTemporalBasicTransformerBlock, TemporalBasicTransformerBlock


def create_custom_forward(module, return_dict=None):
    def custom_forward(*inputs):
        if return_dict is not None:
            return module(*inputs, return_dict=return_dict)

        return module(*inputs)

    return custom_forward


@dataclass
class Transformer3DModelOutput(BaseOutput):
    sample: torch.FloatTensor


class Transformer3DModel(ModelMixin, ConfigMixin):
    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(
        self,
        num_attention_heads: int = 16,
        attention_head_dim: int = 88,
        in_channels: Optional[int] = None,
        num_layers: int = 1,
        dropout: float = 0.0,
        norm_num_groups: int = 32,
        cross_attention_dim: Optional[int] = None,
        attention_bias: bool = False,
        activation_fn: str = "geglu",
        use_linear_projection: bool = False,
        only_cross_attention: bool = False,
        upcast_attention: bool = False,
        unet_use_cross_frame_attention=None,
        unet_use_temporal_attention=None,
        use_audio_module=False,
        depth=0,
        unet_block_name=None,
        emo_drop_rate=0.3,
        is_final_block=False,
    ):
        super().__init__()
        self.use_linear_projection = use_linear_projection
        self.num_attention_heads = num_attention_heads
        self.attention_head_dim = attention_head_dim
        inner_dim = num_attention_heads * attention_head_dim
        self.use_audio_module = use_audio_module
        # Define input layers
        self.in_channels = in_channels
        self.is_final_block = is_final_block

        self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
        if use_linear_projection:
            self.proj_in = nn.Linear(in_channels, inner_dim)
        else:
            self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)

        if use_audio_module:
            self.transformer_blocks = nn.ModuleList(
                [
                    JointAudioTemporalBasicTransformerBlock(
                        dim=inner_dim,
                        num_attention_heads=num_attention_heads,
                        attention_head_dim=attention_head_dim,
                        dropout=dropout,
                        cross_attention_dim=cross_attention_dim,
                        activation_fn=activation_fn,
                        attention_bias=attention_bias,
                        only_cross_attention=only_cross_attention,
                        upcast_attention=upcast_attention,
                        unet_use_cross_frame_attention=unet_use_cross_frame_attention,
                        unet_use_temporal_attention=unet_use_temporal_attention,
                        depth=depth,
                        unet_block_name=unet_block_name,
                        use_ada_layer_norm=True,
                        emo_drop_rate=emo_drop_rate,
                        is_final_block=(is_final_block and d == num_layers - 1),
                    )
                    for d in range(num_layers)
                ]
            )
        else:
            self.transformer_blocks = nn.ModuleList(
                [
                    TemporalBasicTransformerBlock(
                        inner_dim,
                        num_attention_heads,
                        attention_head_dim,
                        dropout=dropout,
                        cross_attention_dim=cross_attention_dim,
                        activation_fn=activation_fn,
                        attention_bias=attention_bias,
                        only_cross_attention=only_cross_attention,
                        upcast_attention=upcast_attention,
                    )
                    for _ in range(num_layers)
                ]
            )

        # 4. Define output layers
        if use_linear_projection:
            self.proj_out = nn.Linear(in_channels, inner_dim)
        else:
            self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)

        self.gradient_checkpointing = False

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

    def forward(
        self,
        hidden_states,
        ref_img_feature=None,
        encoder_hidden_states=None,
        attention_mask=None,
        timestep=None,
        emotion=None,
        uc_mask=None,
        return_dict: bool = True,
    ):
        # Input
        assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
        video_length = hidden_states.shape[2]
        hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")

        if self.use_audio_module:
            if encoder_hidden_states.dim() == 4:
                encoder_hidden_states = rearrange(
                    encoder_hidden_states,
                    "bs f margin dim -> (bs f) margin dim",
                )
        else:
            if encoder_hidden_states.shape[0] != hidden_states.shape[0]:
                encoder_hidden_states = repeat(encoder_hidden_states, "b n c -> (b f) n c", f=video_length)

        batch, _, height, weight = hidden_states.shape
        residual = hidden_states
        if self.use_audio_module:
            residual_audio = encoder_hidden_states

        hidden_states = self.norm(hidden_states)
        if not self.use_linear_projection:
            hidden_states = self.proj_in(hidden_states)
            inner_dim = hidden_states.shape[1]
            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
        else:
            inner_dim = hidden_states.shape[1]
            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
            hidden_states = self.proj_in(hidden_states)

        # Blocks
        for block in self.transformer_blocks:
            if self.training and self.gradient_checkpointing:
                if isinstance(block, TemporalBasicTransformerBlock):
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(block),
                        hidden_states,
                        ref_img_feature,
                        None,  # attention_mask
                        encoder_hidden_states,
                        timestep,
                        None,  # cross_attention_kwargs
                        video_length,
                        uc_mask,
                    )
                elif isinstance(block, JointAudioTemporalBasicTransformerBlock):
                    (
                        hidden_states,
                        encoder_hidden_states,
                    ) = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(block),
                        hidden_states,
                        encoder_hidden_states,
                        attention_mask,
                        emotion,
                    )
                else:
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(block),
                        hidden_states,
                        encoder_hidden_states,
                        timestep,
                        attention_mask,
                        video_length,
                    )
            else:
                if isinstance(block, TemporalBasicTransformerBlock):
                    hidden_states = block(
                        hidden_states=hidden_states,
                        ref_img_feature=ref_img_feature,
                        encoder_hidden_states=encoder_hidden_states,
                        timestep=timestep,
                        video_length=video_length,
                        uc_mask=uc_mask,
                    )
                elif isinstance(block, JointAudioTemporalBasicTransformerBlock):
                    hidden_states, encoder_hidden_states = block(
                        hidden_states,  # shape [2, 4096, 320]
                        encoder_hidden_states=encoder_hidden_states,  # shape [2, 20, 640]
                        attention_mask=attention_mask,
                        emotion=emotion,
                    )
                else:
                    hidden_states = block(
                        hidden_states,  # shape [2, 4096, 320]
                        encoder_hidden_states=encoder_hidden_states,  # shape [2, 20, 640]
                        attention_mask=attention_mask,
                        timestep=timestep,
                        video_length=video_length,
                    )

        # Output
        if not self.use_linear_projection:
            hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
            hidden_states = self.proj_out(hidden_states)
        else:
            hidden_states = self.proj_out(hidden_states)
            hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()

        output = hidden_states + residual

        if self.use_audio_module and not self.is_final_block:
            audio_output = encoder_hidden_states + residual_audio
        else:
            audio_output = None

        output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
        if not return_dict:
            if self.use_audio_module:
                return output, audio_output
            else:
                return output

        if self.use_audio_module:
            return output, audio_output
        else:
            return output