# Created by BVH, Jun - Aug 2025.
# The modified architecture (diffusion transformer and transformer blocks).

from einops import rearrange
from typing import Dict, List, Optional, Tuple

import torch
from torch import nn

from cosmos_predict2.models.action_video2world_dit import ActionConditionedMinimalV1LVGDiT, Mlp
from cosmos_predict2.models.text2image_dit import PatchEmbed, FinalLayer, Block, Timesteps, TimestepEmbedding
from cosmos_predict2.models.video2world_model import DataType
from custom.any4d.a4d_config import Any4DConfig
from imaginaire.utils.graph import create_cuda_graph
# from custom.any4d import a4d_surgery  # avoid due to circular import

from imaginaire.utils import log


class LowDimProjLayer(nn.Module):
    def __init__(
            self,
            proj_mode: str,
            num_tokens: int,
            in_channels: int,
            hidden_channels: int,
            out_channels: int
        ):
        super().__init__()

        self.proj_mode = proj_mode
        self.num_tokens = num_tokens
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels

        if proj_mode == 'per_token':
            self.proj = nn.ModuleList([
                Mlp(in_channels, hidden_channels, out_channels)
                for _ in range(num_tokens)
            ])
        
        elif proj_mode == 'shared':
            self.proj = Mlp(in_channels, hidden_channels, out_channels)
            self.id_embs = nn.Parameter(torch.randn(num_tokens, out_channels))
        
        else:
            raise ValueError(f'Invalid projection mode: {proj_mode}')

    def forward(self, x: torch.Tensor):
        '''
        :param x: (B, num_tokens, in_channels)
        :return y: (B, num_tokens, out_channels)
        '''
        B = x.shape[0]
        assert x.shape == (B, self.num_tokens, self.in_channels), \
            f'LowDimProjLayer input shape must be (B, {self.num_tokens}, {self.in_channels}), but got {x.shape}'
        
        if self.proj_mode == 'per_token':
            y = torch.stack([proj(x[:, i, :])
                             for i, proj in enumerate(self.proj)], dim=-2)
        
        elif self.proj_mode == 'shared':
            x2 = self.proj(x)
            id_embs_bc = self.id_embs.view((1, self.num_tokens, self.out_channels))
            y = x2 + id_embs_bc
        
        else:
            raise ValueError(f'Invalid projection mode: {self.proj_mode}')

        assert y.shape == (B, self.num_tokens, self.out_channels), \
            f'LowDimProjLayer output shape must be (B, {self.num_tokens}, {self.out_channels}), but got {y.shape}'
        return y


class Any4DBlock(Block):
    def __init__(
        self,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.force_any4d = False
        self.force_vanilla = False
        self.return_list = True
        self.block_gate_fix = kwargs['block_gate_fix']
        self.rope_dim = kwargs['rope_dim']  # 128
        self.sattn_rope = torch.nn.Parameter(torch.randn(1, 1, 1, self.rope_dim))
    
    def forward(
        self, 
        x_B_T_H_W_D: list[torch.Tensor],  # list of (B, T, H, W, D) bf16
        emb_B_T_D: list[torch.Tensor],  # list of (B, T, D) bf16; affected by both timestep and action
        crossattn_emb: torch.Tensor,  # (B, NC, D) bf16
        rope_emb_L_1_1_D: list[torch.Tensor],  # list of (T*H*W, 1, 1, D) f32
        adaln_lora_B_T_3D: list[torch.Tensor],  # list of (B, T, D*3) bf16; affected by both timestep and action
        extra_per_block_pos_emb: Optional[torch.Tensor] = None,  # None
        sattn_emb_in: Optional[torch.Tensor] = None,  # (B, NS, D) bf16
        sattn_t_emb_norm: Optional[torch.Tensor] = None,  # (B, NS, D) bf16
        sattn_adaln_lora: Optional[torch.Tensor] = None,  # (B, NS, D*3) bf16
        xattn_emb: Optional[torch.Tensor] = None,  # (B, NX, D) bf16
        **leftover,
    ) -> list[torch.Tensor]:
        '''
        NOTES / WARNINGS (bvh):
        - We should periodically and carefully check updates by NVIDIA to
        text2image_dit.py:Block.forward(), since this implementation mirrors / extends it
        to multi-view / heterogeneous data.
        - If V_used is variable across batches (and thus across GPUs / FSDP ranks),
        this used to give problems because of FSDP sharding, but I solved this by disabling
        sharding for some specific modules.
        - Selective activation checkpointing (SAC) might also be problematic when we modify
        this, because the vanilla code (predict2_2B_720_context_fn) assumes op count 16 by default.
        The solution for now is to switch to the stateless mm_only mode,
        which will use more VRAM, but is also simpler, faster, and more robust.
        '''
        correct_format = (isinstance(x_B_T_H_W_D, list))
        if self.force_any4d:
            if correct_format:
                pass  # resume as normal
            
            else:
                log.warning('Any4DBlock called with force_any4d = True but non-list / homogeneous data (probably a mistake)? '
                            'Correcting format...')
                x_B_T_H_W_D = [x_B_T_H_W_D]
                emb_B_T_D = [emb_B_T_D]
                adaln_lora_B_T_3D = [adaln_lora_B_T_3D]
                rope_emb_L_1_1_D = [rope_emb_L_1_1_D]
                # resume as normal

        elif self.force_vanilla:
            # force vanilla behavior (also single-view / homogeneous tensors only)
            log.warning('Any4DBlock called with force_vanilla = True. '
                        'Reverting to default Block behavior...')
            
            if not correct_format:
                if len(x_B_T_H_W_D) != 1:
                    log.critical('force_vanilla = True but multiple views were given (probably a mistake)?')
                
                x_B_T_H_W_D = x_B_T_H_W_D[0]
                emb_B_T_D = emb_B_T_D[0]
                adaln_lora_B_T_3D = adaln_lora_B_T_3D[0]
                rope_emb_L_1_1_D = rope_emb_L_1_1_D[0]
            
            retval = super().forward(
                x_B_T_H_W_D=x_B_T_H_W_D,
                emb_B_T_D=emb_B_T_D,
                crossattn_emb=crossattn_emb,
                rope_emb_L_1_1_D=rope_emb_L_1_1_D,
                adaln_lora_B_T_3D=adaln_lora_B_T_3D
            )
            
            if self.return_list:
                return [retval]
            else:
                return retval

        elif not correct_format:
            # fallback to vanilla behavior (single-view / homogeneous tensors only)
            raise ValueError('Any4DBlock called with non-list / homogeneous data (probably a mistake)?')

        # BVH shape notes (any4d):
        # V_x_B_T_H_W_D: [tensor[2, 7, 8, 10, 2048] bf16 n=2293760 (4.4Mb) x∈[-1.984, 2.078] μ=0.001 σ=0.093 grad RegisterPostBackwardFunctionBackward cuda:0,
        #                 tensor[2, 7, 8, 10, 2048] bf16 n=2293760 (4.4Mb) x∈[-1.086, 1.086] μ=0.001 σ=0.066 grad RegisterPostBackwardFunctionBackward cuda:0]
        # V_emb_B_T_D: [tensor[2, 7, 2048] bf16 n=28672 (56Kb) x∈[-0.002, 0.102] μ=0.043 σ=0.042 grad RegisterPostBackwardFunctionBackward cuda:0,
        #               tensor[2, 7, 2048] bf16 n=28672 (56Kb) x∈[-0.002, 0.102] μ=0.043 σ=0.042 grad RegisterPostBackwardFunctionBackward cuda:0]
        # crossattn_emb: tensor[2, 512, 1024] bf16 n=1048576 (2Mb) x∈[-0.648, 0.578] μ=1.317e-05 σ=0.018 cuda:0
        # V_rope_emb_L_1_1_D: [tensor[560, 1, 1, 128] n=71680 (0.3Mb) x∈[0., 9.000] μ=0.458 σ=1.138 cuda:0,
        #                      tensor[560, 1, 1, 128] n=71680 (0.3Mb) x∈[0., 9.000] μ=0.458 σ=1.138 cuda:0]
        # V_adaln_lora_B_T_3D: [tensor[2, 7, 6144] bf16 n=86016 (0.2Mb) x∈[-9.812, 6.344] μ=0.676 σ=1.281 grad RegisterPostBackwardFunctionBackward cuda:0,
        #                       tensor[2, 7, 6144] bf16 n=86016 (0.2Mb) x∈[-9.812, 6.344] μ=0.676 σ=1.281 grad RegisterPostBackwardFunctionBackward cuda:0]

        assert extra_per_block_pos_emb is None, 'not supported in Any4D, see text2image_dit.py'
        assert self.use_adaln_lora, 'not supported in Any4D, see text2image_dit.py'

        cat_rope_emb_L_1_1_D = torch.cat(rope_emb_L_1_1_D, dim=0)  # (V*T*H*W, 1, 1, D) f32
        # cat_rope_emb_L_1_1_D: tensor[1120, 1, 1, 128] n=143360 (0.5Mb) x∈[0., 9.000] μ=0.458 σ=1.138 cuda:0

        def _fn(_x_B_T_H_W_D, _norm_layer, _scale_B_T_1_1_D, _shift_B_T_1_1_D):
            return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D

        V_used = len(x_B_T_H_W_D)  # int
        V_x = list(x_B_T_H_W_D)  # Make a shallow copy
        V_normalized_x = []
        V_gate_next = []
        V_shapes = []
        V_seqlens = [0]

        # ================================ PER VIEW
        
        for v in range(V_used):
            shift_self_attn_B_T_D, scale_self_attn_B_T_D, gate_self_attn_B_T_D = (
                self.adaln_modulation_self_attn(emb_B_T_D[v]) + adaln_lora_B_T_3D[v]
            ).chunk(3, dim=-1)

            shift_self_attn_B_T_1_1_D = rearrange(shift_self_attn_B_T_D, "b t d -> b t 1 1 d")
            scale_self_attn_B_T_1_1_D = rearrange(scale_self_attn_B_T_D, "b t d -> b t 1 1 d")
            gate_self_attn_B_T_1_1_D = rearrange(gate_self_attn_B_T_D, "b t d -> b t 1 1 d")
            # shift_self_attn_B_T_1_1_D: tensor[2, 7, 1, 1, 2048] bf16 n=28672 (56Kb) x∈[-5.750, 4.938] μ=0.001 σ=0.248 grad ViewBackward0 cuda:0
            # scale_self_attn_B_T_1_1_D: tensor[2, 7, 1, 1, 2048] bf16 n=28672 (56Kb) x∈[-5.344, 8.250] μ=-0.238 σ=1.805 grad ViewBackward0 cuda:0
            # gate_self_attn_B_T_1_1_D: tensor[2, 7, 1, 1, 2048] bf16 n=28672 (56Kb) x∈[-10.250, 2.406] μ=-0.010 σ=0.527 grad ViewBackward0 cuda:0
            
            normalized_x_B_T_H_W_D = _fn(
                V_x[v],
                self.layer_norm_self_attn,
                scale_self_attn_B_T_1_1_D,
                shift_self_attn_B_T_1_1_D
            )

            normalized_x_B_L_D = rearrange(normalized_x_B_T_H_W_D, "b t h w d -> b (t h w) d")
            (T, H, W) = V_x[v].shape[1:4]
            
            V_normalized_x.append(normalized_x_B_L_D)
            V_shapes.append((T, H, W))
            V_seqlens.append(V_seqlens[v] + T*H*W)
            V_gate_next.append(gate_self_attn_B_T_1_1_D)

        # ================================ LOWDIM

        cat_normalized_x_B_L_D = torch.cat(V_normalized_x, dim=1)  # (B, V*T*H*W, D) bf16
        NV = cat_normalized_x_B_L_D.shape[1]
        # cat_normalized_x_B_L_D: tensor[2, 1120, 2048] bf16 n=4587520 (8.8Mb) x∈[-37.000, 39.750] μ=-0.013 σ=1.609 grad CatBackward0 cuda:0

        if sattn_emb_in is not None:
            sattn_emb = sattn_emb_in  # (B, NS, D) bf16
            NS = sattn_emb.shape[1]

            (shift_esa, scale_esa, gate_esa) = (
                self.adaln_modulation_self_attn(sattn_t_emb_norm) + sattn_adaln_lora
            ).chunk(3, dim=-1)
            # ^ 3x (B, NS, D) bf16

            sattn_emb2 = _fn(sattn_emb, self.layer_norm_self_attn, scale_esa, shift_esa)
            # ^ (B, NS, D) bf16
            
            aug_x = torch.cat([cat_normalized_x_B_L_D, sattn_emb2], dim=1)
            # ^ (B, V*T*H*W + NS, D) bf16
            
            aug_rope = torch.cat([cat_rope_emb_L_1_1_D, self.sattn_rope.repeat(NS, 1, 1, 1)], dim=0)
            # ^ (V*T*H*W + NS, 1, 1, D) f32
        
        else:
            aug_x = cat_normalized_x_B_L_D  # (B, V*T*H*W, D) bf16
            aug_rope = cat_rope_emb_L_1_1_D  # (V*T*H*W, 1, 1, D) f32
        
        # ================================ SHARED
        
        # bookmark(bvh): any4d dit_self_attn
        aug_res = self.self_attn(
            aug_x,
            None,
            rope_emb=aug_rope,
        )
        cat_result_B_L_D = aug_res[:, 0:NV]
        # cat_result_B_L_D: tensor[2, 1120, 2048] bf16 n=4587520 (8.8Mb) x∈[-1.500, 0.984] μ=-0.003 σ=0.205 grad UnsafeViewBackward0 cuda:0

        # ================================ LOWDIM

        if sattn_emb_in is not None:
            sattn_emb3 = aug_res[:, NV:NV+NS]
            sattn_emb4 = sattn_emb + gate_esa * sattn_emb3
        
        # ================================ PER VIEW
        
        for v in range(V_used):
            shift_cross_attn_B_T_D, scale_cross_attn_B_T_D, gate_cross_attn_B_T_D = (
                self.adaln_modulation_cross_attn(emb_B_T_D[v]) + adaln_lora_B_T_3D[v]
            ).chunk(3, dim=-1)
            
            shift_cross_attn_B_T_1_1_D = rearrange(shift_cross_attn_B_T_D, "b t d -> b t 1 1 d")
            scale_cross_attn_B_T_1_1_D = rearrange(scale_cross_attn_B_T_D, "b t d -> b t 1 1 d")
            gate_cross_attn_B_T_1_1_D = rearrange(gate_cross_attn_B_T_D, "b t d -> b t 1 1 d")

            result_B_L_D = cat_result_B_L_D[:, V_seqlens[v]:V_seqlens[v+1]]  # (B, T*H*W, D) bf16
            
            (T, H, W) = V_shapes[v]
            result_B_T_H_W_D = rearrange(result_B_L_D, "b (t h w) d -> b t h w d", t=T, h=H, w=W)
            
            # this is gate_self_attn_B_T_1_1_D
            V_x[v] = V_x[v] + V_gate_next[v] * result_B_T_H_W_D
            
            # NOTE(bvh): Unfortunately we cannot use _x_fn() here since things are too entangled.
            normalized_x_B_T_H_W_D = _fn(
                V_x[v],
                self.layer_norm_cross_attn,
                scale_cross_attn_B_T_1_1_D,
                shift_cross_attn_B_T_1_1_D
            )

            normalized_x_B_L_D = rearrange(normalized_x_B_T_H_W_D, "b t h w d -> b (t h w) d")
            
            V_normalized_x[v] = normalized_x_B_L_D
            V_gate_next[v] = gate_cross_attn_B_T_1_1_D

        # ================================ LOWDIM

        cat_normalized_x_B_L_D = torch.cat(V_normalized_x, dim=1)  # (B, V*T*H*W, D) bf16
        # cat_normalized_x_B_L_D: tensor[2, 1120, 2048] bf16 n=4587520 (8.8Mb) x∈[-20.875, 52.250] μ=-0.023 σ=1.727 grad CatBackward0 cuda:0

        if sattn_emb_in is not None:
            (shift_esa, scale_esa, gate_esa) = (
                self.adaln_modulation_cross_attn(sattn_t_emb_norm) + sattn_adaln_lora
            ).chunk(3, dim=-1)
            # ^ 3x (B, NS, D) bf16

            sattn_emb5 = _fn(sattn_emb4, self.layer_norm_cross_attn, scale_esa, shift_esa)
            # ^ (B, NS, D) bf16

            aug_x = torch.cat([cat_normalized_x_B_L_D, sattn_emb5], dim=1)
            # ^ (B, V*T*H*W + NS, D) bf16
        
            # NOTE: aug_rope already calculated above

        else:
            aug_x = cat_normalized_x_B_L_D
            # ^ (B, V*T*H*W, D) bf16

        if xattn_emb is not None:
            aug_crossattn = torch.cat([crossattn_emb, xattn_emb], dim=1)
            # ^ (B, NC + NX, D) bf16
        
        else:
            aug_crossattn = crossattn_emb
            # ^ (B, NC, D) bf16

        # ================================ SHARED

        # bookmark(bvh): any4d dit_cross_attn
        aug_res = self.cross_attn(
            aug_x,
            aug_crossattn,
            rope_emb=aug_rope,
        )
        cat_result_B_L_D = aug_res[:, 0:NV]
        # cat_result_B_L_D: tensor[2, 1120, 2048] bf16 n=4587520 (8.8Mb) x∈[-0.030, 0.019] μ=4.387e-05 σ=0.002 grad UnsafeViewBackward0 cuda:0
        
        # ================================ LOWDIM

        if sattn_emb_in is not None:
            sattn_emb6 = aug_res[:, NV:NV+NS]
            sattn_emb7 = sattn_emb4 + gate_esa * sattn_emb6

        # ================================ PER VIEW

        for v in range(V_used):
            shift_mlp_B_T_D, scale_mlp_B_T_D, gate_mlp_B_T_D = (
                self.adaln_modulation_mlp(emb_B_T_D[v]) + adaln_lora_B_T_3D[v]
            ).chunk(3, dim=-1)
        
            shift_mlp_B_T_1_1_D = rearrange(shift_mlp_B_T_D, "b t d -> b t 1 1 d")
            scale_mlp_B_T_1_1_D = rearrange(scale_mlp_B_T_D, "b t d -> b t 1 1 d")
            gate_mlp_B_T_1_1_D = rearrange(gate_mlp_B_T_D, "b t d -> b t 1 1 d")

            result_B_L_D = cat_result_B_L_D[:, V_seqlens[v]:V_seqlens[v+1]]  # (B, T*H*W, D) bf16
            
            (T, H, W) = V_shapes[v]
            result_B_T_H_W_D = rearrange(result_B_L_D, "b (t h w) d -> b t h w d", t=T, h=H, w=W)
            
            # new / more consistent notation:
            if self.block_gate_fix:
                # NOTE: this should be enabled for all new training / resume runs
                # this is gate_cross_attn_B_T_1_1_D
                V_x[v] = V_x[v] + V_gate_next[v] * result_B_T_H_W_D
            else:
                # NOTE: this was accidentally the default until 12/4
                V_x[v] = V_x[v] + gate_mlp_B_T_1_1_D * result_B_T_H_W_D
            
            # old notation / like vanilla Block:
            # NOTE: this was accidentally the default until 12/4
            # V_x[v] = result_B_T_H_W_D * gate_mlp_B_T_1_1_D + V_x[v]
            
            normalized_x_B_T_H_W_D = _fn(
                V_x[v],
                self.layer_norm_mlp,
                scale_mlp_B_T_1_1_D,
                shift_mlp_B_T_1_1_D
            )
            
            # NOTE(bvh): No need to do this step jointly across views,
            # since mlp = GPT2FeedForward which has linear layers only.
            result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D)
            
            V_x[v] = V_x[v] + gate_mlp_B_T_1_1_D * result_B_T_H_W_D  #@IgnoreException

        # ================================ LOWDIM
        
        if sattn_emb_in is not None:
            (shift_esa, scale_esa, gate_esa) = (
                self.adaln_modulation_mlp(sattn_t_emb_norm) + sattn_adaln_lora
            ).chunk(3, dim=-1)
            # ^ 3x (B, NS, D) bf16

            sattn_emb8 = _fn(sattn_emb7, self.layer_norm_mlp, scale_esa, shift_esa)
            sattn_emb9 = self.mlp(sattn_emb8)
            sattn_emb10 = sattn_emb7 + gate_esa * sattn_emb9
            
            # Return sattn output (avoid in-place updates for checkpoint compatibility)
            sattn_out = sattn_emb10
        
        else:
            sattn_out = None

        # ================================ SHARED

        # BVH shape notes (any4d):
        # V_x: [tensor[2, 7, 8, 10, 2048] bf16 n=2293760 (4.4Mb) x∈[-2.234, 7.625] μ=0.005 σ=0.262 grad AddBackward0 cuda:0,
        #       tensor[2, 7, 8, 10, 2048] bf16 n=2293760 (4.4Mb) x∈[-5.344, 9.250] μ=-0.000 σ=0.350 grad AddBackward0 cuda:0]

        if self.return_list:
            return (V_x, sattn_out)
        
        else:
            if len(V_x) != 1:
                log.critical('return_list = False but multiple views were returned (probably a mistake)?')
            
            return (V_x[0], sattn_out)


class Any4DDiT(ActionConditionedMinimalV1LVGDiT):
    '''
    NOTE(bvh): We inherit from Cosmos-Predict2-2B-Action-Conditioned-Sample instead of
    Cosmos-Predict2-2B-Video2World, in case we want to condition on actions.
    (Caveat: it will be hard to use the pretrained representation because of the
    fixed 7 * 12 action dimension.)
    '''
    def __init__(
        self,
        *args,
        # in_channels_newviews: list[int] = None,
        # out_channels_newviews: list[int] = None,
        any4d_config: Any4DConfig = None,
        **kwargs,
    ):
        # Prepare surgery
        from custom.any4d import a4d_surgery
        any4d_config = a4d_surgery.validate_config(any4d_config)

        num_views = any4d_config.num_views
        in_channels_existing = any4d_config.num_video_in_channels[0]
        out_channels_existing = any4d_config.num_video_out_channels[0]

        kwargs['concat_padding_mask'] = True
        kwargs['in_channels_override'] = in_channels_existing  # extend pretrained PatchEmbed
        kwargs['out_channels_override'] = out_channels_existing  # extend pretrained FinalLayer

        # kwargs['action_dim'] = any4d_config.num_lowdim_adaln_channels
        (Ta, Ca) = (any4d_config.num_lowdim_adaln_tokens, any4d_config.num_lowdim_adaln_channels)
        kwargs['action_dim'] = Ta * Ca

        kwargs['block_cls'] = Any4DBlock  # accept heterogeneous inputs (multi-view)
        kwargs['block_gate_fix'] = any4d_config.block_gate_fix

        super().__init__(*args, **kwargs)
        # ^ this initializes x_embedder, final_layer,
        # action_embedder_B_D, and action_embedder_B_3D.

        self.any4d_config = any4d_config
    
        # ==========================================================
        # ======== Handle high-dimensional inputs & outputs ========
        # ==========================================================
        
        in_channels_newviews = any4d_config.num_video_in_channels[1:]
        out_channels_newviews = any4d_config.num_video_out_channels[1:]
        
        # NOTE(bvh): This should match cosmos_predict2/models/text2image_dit.py:MiniTrainDIT.build_patch_embed().
        self.x_embedder_newviews = nn.ModuleList([
            PatchEmbed(  # all embedders (both existing & new) have bias = False
                spatial_patch_size=self.patch_spatial,  # = 2
                temporal_patch_size=self.patch_temporal,  # = 1
                in_channels=in_channels_newviews[view_idx],  # = depends on surgery
                out_channels=self.model_channels  # = 2048 (2B) / 5120 (14B)
            ) for view_idx in range(num_views - 1)
        ])

        # NOTE(bvh): This unique viewpoint identifier acts as a kind of positional encoding,
        # and is always randomly initialized and added to all tokens in each video stream
        # after input projection.
        self.view_embs_newviews = nn.ParameterList([
            nn.Parameter(torch.randn(1, self.model_channels) * any4d_config.view_emb_std)
            for view_idx in range(num_views - 1)
        ])

        self.t_embedder_newviews = nn.ModuleList([
            nn.Sequential(
                Timesteps(self.model_channels),
                TimestepEmbedding(self.model_channels, self.model_channels, use_adaln_lora=self.use_adaln_lora),
            ) for view_idx in range(num_views - 1)
        ])

        # NOTE(bvh): This should match cosmos_predict2/models/text2image_dit.py:MiniTrainDIT.__init__().
        self.final_layer_newviews = nn.ModuleList([
            FinalLayer(
                hidden_size=self.model_channels,  # = 2048 (2B) / 5120 (14B)
                spatial_patch_size=self.patch_spatial,  # = 2
                temporal_patch_size=self.patch_temporal,  # = 1
                out_channels=out_channels_newviews[view_idx],  # = depends on surgery
                use_adaln_lora=self.use_adaln_lora,  # = True
                adaln_lora_dim=self.adaln_lora_dim,  # = 256
            ) for view_idx in range(num_views - 1)
        ])
    
        # =========================================================
        # ======== Handle low-dimensional inputs & outputs ========
        # =========================================================

        if any4d_config.has_sattn_stream:
            sattn_hidden_channels = min(any4d_config.lowdim_hidden_channels,
                                        any4d_config.num_lowdim_sattn_channels,
                                        self.model_channels // 4)
            self.sattn_inproj = LowDimProjLayer(
                proj_mode=any4d_config.lowdim_sattn_proj_mode,
                num_tokens=any4d_config.num_lowdim_sattn_tokens,
                in_channels=any4d_config.num_lowdim_sattn_channels,
                hidden_channels=sattn_hidden_channels,
                out_channels=self.model_channels,
            )
            self.sattn_outproj = LowDimProjLayer(
                proj_mode=any4d_config.lowdim_sattn_proj_mode,
                num_tokens=any4d_config.num_lowdim_sattn_tokens,
                in_channels=self.model_channels,
                hidden_channels=sattn_hidden_channels,
                out_channels=any4d_config.num_lowdim_sattn_channels,
            )
            self.t_embedder_sattn = nn.Sequential(
                Timesteps(self.model_channels),
                TimestepEmbedding(self.model_channels, self.model_channels, use_adaln_lora=self.use_adaln_lora),
            )

        if any4d_config.has_xattn_stream:
            xattn_hidden_channels = min(any4d_config.lowdim_hidden_channels,
                                        any4d_config.num_lowdim_xattn_channels,
                                        self.model_channels // 4)
            self.xattn_inproj = LowDimProjLayer(
                proj_mode=any4d_config.lowdim_xattn_proj_mode,
                num_tokens=any4d_config.num_lowdim_xattn_tokens,
                in_channels=any4d_config.num_lowdim_xattn_channels,
                hidden_channels=xattn_hidden_channels,
                out_channels=self.model_channels,
            )

        # NOTE(bvh): lowdim_adaln_entries already taken care of by action_dim above.

        pass


    # NOTE(bvh): This method is adapted from:
    # - cosmos_predict2/models/action_video2world_dit.py:ActionConditionedMinimalV1LVGDiT.forward()
    # which itself was adapted from:
    # - cosmos_predict2/models/video2world_dit.py:MinimalV1LVGDiT.forward()
    # - cosmos_predict2/models/text2image_dit.py:MiniTrainDIT.forward().
    def forward(
        self,
        x_in_streams: dict[str, torch.Tensor],
        timesteps: dict[str, torch.Tensor],
        crossattn_emb: torch.Tensor,
        fps: Optional[torch.Tensor] = None,
        use_cuda_graphs: bool = False,
        iteration: int = -1,
        **leftover,
    ) -> Dict[str, torch.Tensor]:
        '''
        :param x_in_streams: dict mapping stream name to (B, C, T, H, W) or (B, T, D) tensor.
        :param timesteps: dict mapping stream name to (B, 1/T, 1) or (B, 1, 1/T, 1, 1) tensor.
        :param crossattn_emb: (B, NC, D) tensor.
        :param fps: (B) tensor.
        :param use_cuda_graphs: bool.
        :return y_out_streams: dict mapping stream name to (B, C, T, H, W) or (B, T, D) tensor.
        '''
        assert not(self.extra_per_block_abs_pos_emb), 'not supported in Any4D, see text2image_dit.py'
        assert 'rope' in self.pos_emb_cls, 'not supported in Any4D, see text2image_dit.py'
        assert not(use_cuda_graphs), 'not supported in Any4D, see text2image_dit.py'
        
        # BVH shape notes (any4d):
        # x_in_streams: {
        #   'v0': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-3.609, 3.422] μ=-0.108 σ=0.910 cuda:0,
        #   'xattn': tensor[2, 30, 14] bf16 n=840 (1.6Kb) all_zeros cuda:0,
        #   'adaln': tensor[2, 7, 21] bf16 n=294 x∈[-0.992, 1.000] μ=0.114 σ=0.516 cuda:0,
        #   'v1': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-3.656, 3.703] μ=-0.102 σ=0.926 cuda:0,
        #   'sattn': tensor[2, 25, 22] bf16 n=1100 (2.1Kb) all_zeros cuda:0}
        # timesteps: {
        #   'v0': tensor[2, 1, 7, 1, 1] n=14 x∈[9.999e-05, 9.999e-05] μ=9.999e-05 σ=0. cuda:0,
        #   'v1': tensor[2, 1, 7, 1, 1] n=14 x∈[0.193, 0.198] μ=0.195 σ=0.002 cuda:0,
        #   'sattn': tensor[2, 25, 1] n=50 x∈[9.999e-05, 9.999e-05] μ=9.999e-05 σ=0. cuda:0}
        # crossattn_emb: tensor[2, 512, 1024] bf16 n=1048576 (2Mb) x∈[-0.621, 0.570] μ=1.049e-05 σ=0.013 cuda:0
        # fps: tensor[2] bf16 μ=10.000 σ=0. cuda:0 [10.000, 10.000]
        # use_cuda_graphs: False
        # iteration: 1+
        # leftover: {
        #   'data_type': <DataType.VIDEO: 'video'>,
        #   'padding_mask': tensor[2, 1, 128, 160] bf16 n=40960 (80Kb) all_zeros cuda:0,
        #   'use_video_condition': tensor[1] bool cuda:0 [True],
        #   'gt_frames': None,
        #   'condition_video_input_mask_B_C_T_H_W': None}
        
        any4d_config = self.any4d_config
        V_max = any4d_config.num_views
        V_used = [v for v in range(V_max) if f'v{v}' in x_in_streams]  # list of int
        # NOTE(bvh): ^ in theory, this should work for non-consecutive views as well,
        # but in practice we typically construct batches with the subset 0...k-1 for now.

        if any4d_config.video_concat_mode.startswith('spat:'):
            (NH, NW) = map(int, any4d_config.video_concat_mode.split(':')[1].split(','))
            assert NH * NW == len(V_used), 'Number of spatial tiles must match number of views'
        
        # ================================================
        # ======== Handle high-dimensional inputs ========
        # ================================================
        
        if any4d_config.video_proj_mode == 'per_view':
            my_x_embedders = [self.x_embedder, *self.x_embedder_newviews]
            my_final_layers = [self.final_layer, *self.final_layer_newviews]
        elif any4d_config.video_proj_mode == 'shared':
            my_x_embedders = [self.x_embedder] * V_max
            my_final_layers = [self.final_layer] * V_max
        else:
            raise ValueError(f'Invalid video_proj_mode: {any4d_config.video_proj_mode}')

        if any4d_config.view_timestep_mode == 'shared':
            my_view_t_embedders = [self.t_embedder] * V_max
        elif any4d_config.view_timestep_mode == 'per_view':
            my_view_t_embedders = [self.t_embedder, *self.t_embedder_newviews]
        else:
            raise ValueError(f'Invalid view_timestep_mode: {any4d_config.view_timestep_mode}')

        if any4d_config.has_adaln_stream:
            action = x_in_streams['adaln']  # (B, T, D).
            
            # project adaln to action embedding
            assert action is not None, "adaln must be provided"
            action_flat = rearrange(action, "b t d -> b 1 (t d)")
            action_emb_B_D = self.action_embedder_B_D(action_flat)  # (B, 1, D) bf16
            action_emb_B_3D = self.action_embedder_B_3D(action_flat)  # (B, 1, D*3) bf16

        # First encode all viewpoint contents (high-dim video streams) separately.
        x_proj_all = []
        timesteps_all = []

        for v in V_used:
            x_in_v = x_in_streams[f'v{v}']  # (B, C, T, H, W)
            timesteps_v = timesteps[f'v{v}'][:, 0, :, 0, 0]  # (B, 1/T)

            # This is taken out of text2image_dit.py:MiniTrainDIT.prepare_embedded_sequence(),
            # because we have to distinguish input projection from positional embedding:
            x_proj_v = my_x_embedders[v](x_in_v)  # (B, T, H, W, D) bf16

            if any4d_config.video_concat_mode == 'view' and v >= 1:
                x_proj_v += self.view_embs_newviews[v - 1]  # (B, T, H, W, D) bf16

            T = x_proj_v.shape[1]
            if timesteps_v.shape[1] == 1:  # convert (B, 1) to (B, T) as needed
                timesteps_v = timesteps_v.repeat(1, T)  # (B, T) bf16

            x_proj_all.append(x_proj_v)  # list of (B, T, H, W, D) bf16
            timesteps_all.append(timesteps_v)  # list of (B, T, 1) bf16
        
        # Combine video streams beforehand, or keep separate, as desired.
        # NOTE(bvh): Unless we are doing time or view concatenation,
        # all noise levels (= timesteps) are expected / required to be the same.
        if any4d_config.video_concat_mode == 'time':
            cat_dims = [x_proj_v.shape[1] for x_proj_v in x_proj_all]
            cmb_x_proj_all = [torch.cat(x_proj_all, dim=1)]  # (B, V*T, H, W, D) bf16
            cmb_timesteps = [torch.cat(timesteps_all, dim=1)]  # (B, V*T) bf16

        elif any4d_config.video_concat_mode == 'vert':
            cat_dims = [x_proj_v.shape[2] for x_proj_v in x_proj_all]
            cmb_x_proj_all = [torch.cat(x_proj_all, dim=2)]  # (B, T, V*H, W, D) bf16
            cmb_timesteps = [timesteps_all[0]]  # (B, T) bf16

        elif any4d_config.video_concat_mode == 'horz':
            cat_dims = [x_proj_v.shape[3] for x_proj_v in x_proj_all]
            cmb_x_proj_all = [torch.cat(x_proj_all, dim=3)]  # (B, T, H, V*W, D) bf16
            cmb_timesteps = [timesteps_all[0]]  # (B, T) bf16

        elif any4d_config.video_concat_mode.startswith('spat:'):
            # cat_dims not needed since all dimensions have to be equal anyway
            x_proj_stack = torch.stack(x_proj_all, dim=0)  # (V, B, T, H, W, D) bf16
            cmb_x_proj_all = [rearrange(x_proj_stack,
                '(nh nw) b t h w d -> b t (nh h) (nw w) d', nh=NH, nw=NW)]  # (B, T, NH*H, NW*W, D) bf16
            cmb_timesteps = [timesteps_all[0]]  # (B, T) bf16

        elif any4d_config.video_concat_mode == 'view':
            # Pass-through since already handled above & below
            cmb_x_proj_all = x_proj_all  # list of (B, T, H, W, D) bf16
            cmb_timesteps = timesteps_all  # list of (B, T) bf16

        else:
            raise ValueError(f'Invalid video_concat_mode: {any4d_config.video_concat_mode}')
        
        # Now process positional encodings, diffusion timestep encodings, and adaptive layernorm.
        rope_emb_all = []
        t_embedding_all = []
        t_emb_norm_all = []
        adaln_lora_all = []

        for (x_proj_v, timesteps_v) in zip(cmb_x_proj_all, cmb_timesteps):
            # This is taken out of text2image_dit.py:MiniTrainDIT.prepare_embedded_sequence(),
            # because we have to distinguish input projection from positional embedding:
            rope_emb_v = self.pos_embedder(x_proj_v, fps=fps)  # (T*H*W, 1, 1, D) f32
            # NOTE(bvh): annoyingly, ^ this does NOT depend on fps anymore in new cosmos;
            # TODO(bvh): maybe make configurable, and check effect of batch dimension
            # ^ I meant make WHAT configurable ??

            (t_embedding_v, adaln_lora_v) = my_view_t_embedders[v](timesteps_v)
            # ^ (B, T, D) bf16, (B, T, D*3) bf16

            if any4d_config.has_adaln_stream:
                # add action embedding to the timestep embedding and adaln_lora
                t_embedding_v = t_embedding_v + action_emb_B_D  # (B, T, D) bf16
                adaln_lora_v = adaln_lora_v + action_emb_B_3D  # (B, T, D*3) bf16

            t_emb_norm_v = self.t_embedding_norm(t_embedding_v)  # (B, T, D) bf16

            rope_emb_all.append(rope_emb_v)  # list of (T*H*W, 1, 1, D) f32
            t_embedding_all.append(t_embedding_v)  # list of (B, T, D) bf16
            t_emb_norm_all.append(t_emb_norm_v)  # list of (B, T, D) bf16
            adaln_lora_all.append(adaln_lora_v)  # list of (B, T, D*3) bf16

        # ===============================================
        # ======== Handle low-dimensional inputs ========
        # ===============================================

        if any4d_config.has_sattn_stream:
            sattn_in = x_in_streams['sattn']  # (B, NS, D) bf16
            sattn_emb_in = self.sattn_inproj(sattn_in)  # (B, NS, D) bf16
            NS = sattn_emb_in.shape[1]

            timesteps_sattn = timesteps[f'sattn'][:, :, 0]  # (B, 1/T/NS/?)
            if any4d_config.lowdim_timestep_mode == 'per_stream':
                my_sattn_t_embedder = self.t_embedder_sattn
            else:
                my_sattn_t_embedder = my_view_t_embedders[0]
            
            if timesteps_sattn.shape[1] == 1:  # convert (B, 1) to (B, NS) as needed
                timesteps_sattn = timesteps_sattn.repeat(1, NS)  # (B, NS) bf16
            
            (sattn_t_embedding, sattn_adaln_lora) = my_sattn_t_embedder(timesteps_sattn)
            # ^ (B, NS, D) bf16, (B, NS, D*3) bf16

            if any4d_config.has_adaln_stream:
                sattn_t_embedding = sattn_t_embedding + action_emb_B_D  # (B, NS, D) bf16
                sattn_adaln_lora = sattn_adaln_lora + action_emb_B_3D  # (B, NS, D*3) bf16

            sattn_t_emb_norm = self.t_embedding_norm(sattn_t_embedding)  # (B, NS, D) bf16
            sattn_emb = sattn_emb_in  # (B, NS, D) bf16 - this flows through blocks

        else:
            sattn_emb = None
            sattn_t_emb_norm = None
            sattn_adaln_lora = None
        
        if any4d_config.has_xattn_stream:
            xattn_in = x_in_streams['xattn']  # (B, NX, D) bf16
            xattn_emb = self.xattn_inproj(xattn_in)  # (B, NX, D) bf16
        
        else:
            xattn_emb = None
    
        # ==========================================
        # ======== Apply transformer blocks ========
        # ==========================================

        blocks = self.blocks
        x_block_all = cmb_x_proj_all
        # x_block_all: [tensor[2, 7, 8, 10, 2048] bf16 n=2293760 (4.4Mb) x∈[-1.984, 2.078] μ=0.001 σ=0.093 grad UnsafeViewBackward0 cuda:0,
        #               tensor[2, 7, 8, 10, 2048] bf16 n=2293760 (4.4Mb) x∈[-1.086, 1.086] μ=0.001 σ=0.066 grad AddBackward0 cuda:0]
        
        # len = 28
        for block in blocks:
            (x_block_all, sattn_out) = block(
                x_B_T_H_W_D=x_block_all,  # list of (B, T, H, W, D) bf16
                emb_B_T_D=t_emb_norm_all,  # list of (B, T, D) bf16
                crossattn_emb=crossattn_emb,  # (B, NC, D) bf16
                rope_emb_L_1_1_D=rope_emb_all,  # list of (T*H*W, 1, 1, 128) f32
                adaln_lora_B_T_3D=adaln_lora_all,  # list of (B, T, D*3) bf16
                sattn_emb_in=sattn_emb,  # (B, NS, D) bf16
                sattn_t_emb_norm=sattn_t_emb_norm,  # (B, NS, D) bf16
                sattn_adaln_lora=sattn_adaln_lora,  # (B, NS, D*3) bf16
                xattn_emb=xattn_emb,  # (B, NX, D) bf16
            )
            # Update sattn_emb for the next block
            if sattn_out is not None:
                sattn_emb = sattn_out

        y_block_all = x_block_all  # list of (B, T, H, W, D) bf16
        # y_block_all: [tensor[2, 7, 8, 10, 2048] bf16 n=2293760 (4.4Mb) x∈[-85.000, 270.000] μ=-0.578 σ=11.125 grad AddBackward0 cuda:0,
        #               tensor[2, 7, 8, 10, 2048] bf16 n=2293760 (4.4Mb) x∈[-82.000, 196.000] μ=-0.320 σ=12.312 grad AddBackward0 cuda:0]

        # ================================================
        # ======== Handle low-dimensional outputs ========
        # ================================================
        
        if any4d_config.has_sattn_stream:
            sattn_emb_out = sattn_emb  # (B, NS, D) bf16
        
        else:
            sattn_emb_out = None

        # =================================================
        # ======== Handle high-dimensional outputs ========
        # =================================================

        # Separate video streams depending on how they were combined before.
        sep_y_block_all = []
        offset = 0

        if any4d_config.video_concat_mode == 'time':
            for v in range(len(V_used)):
                sep_y_block_all.append(y_block_all[:, offset:offset + cat_dims[v]])
                offset += cat_dims[v]

        elif any4d_config.video_concat_mode == 'vert':
            for v in range(len(V_used)):
                sep_y_block_all.append(y_block_all[:, :, offset:offset + cat_dims[v]])
                offset += cat_dims[v]

        elif any4d_config.video_concat_mode == 'horz':
            for v in range(len(V_used)):
                sep_y_block_all.append(y_block_all[:, :, :, offset:offset + cat_dims[v]])
                offset += cat_dims[v]

        elif any4d_config.video_concat_mode.startswith('spat:'):
            sep_y_block_all = rearrange(y_block_all,
                'v b t (nh h) (nw w) d -> (nh nw) b t h w d', nh=NH, nw=NW)  # (NH*NW, B, T, H, W, D) bf16
            sep_y_block_all = list(sep_y_block_all)  # list of (B, T, H, W, D) bf16

        elif any4d_config.video_concat_mode == 'view':
            # Pass-through since already in correct format
            sep_y_block_all = y_block_all  # list of (B, T, H, W, D) bf16

        y_proj_all = []
        for (v, x_block_v) in enumerate(sep_y_block_all):
            y_proj_v = my_final_layers[v](x_block_v, t_embedding_all[v], adaln_lora_B_T_3D=adaln_lora_all[v])
            y_proj_v = self.unpatchify(y_proj_v)
            y_proj_all.append(y_proj_v)

        # y_proj_all: [tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-336.000, 282.000] μ=-1.836 σ=30.625 grad UnsafeViewBackward0 cuda:0,
        #              tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-332.000, 466.000] μ=4.781 σ=41.250 grad UnsafeViewBackward0 cuda:0]

        y_out_streams = dict()

        if any4d_config.has_sattn_stream:
            sattn_out = self.sattn_outproj(sattn_emb_out)
            y_out_streams['sattn'] = sattn_out  # (B, NS, D) bf16
        
        for v in V_used:
            y_out_streams[f'v{v}'] = y_proj_all[v]  # (B, T, H, W, D) bf16
        
        # BVH shape notes (any4d):
        # y_out_streams: {
        #   'v0': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-146.000, 322.000] μ=7.531 σ=39.250 grad UnsafeViewBackward0 cuda:0,
        #   'v1': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-404.000, 364.000] μ=0.350 σ=41.250 grad UnsafeViewBackward0 cuda:0}
        
        return y_out_streams