# 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.rope_dim = kwargs['rope_dim'] # 128 # NOTE(bvh): dead param from the deprecated sattn stream, removed so it never reaches optimizer. # Old ckpts still carry the key => now also stripped on load (see a4d_surgery + checkpointer). # self.sattn_rope = torch.nn.Parameter(torch.randn(1, 1, 1, self.rope_dim)) self.block_gate_fix = kwargs['block_gate_fix'] self.isolate_stub_streams = kwargs['isolate_stub_streams'] 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 combined with unequal number of views across ranks, but I solved this by: 1) disabling sharding for some specific modules (see also text2image_dit.py:MiniTrainDIT.fully_shard()). 2) creating tiny stub streams for missing views (see also a4d_logistics.py:pack_streams_from_entries()). To clarify, this means that the list of streams should be the exact same across all ranks, which then avoids FSDP deadlocks, but does not affect forward pass calculation/logic as long as isolate_stub_streams = True. - 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' 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 # NOTE(bvh): ^ this has a DIFFERENT meaning here compared to Any4DDiT! V_x = list(x_B_T_H_W_D) # Make a shallow copy V_normalized_x = [] V_gate_next = [] V_shapes = [] V_seqlens = [0] if self.isolate_stub_streams: # NOTE(bvh): stub streams/tokens should be skipped from the consequential operations # like attention calculations etc. # this is an attempt at fixing training artefacts (Mar - May 2026) # this assumes (T, H, W) = (1, 2, 2) and not more V_real = [v for v in range(V_used) if V_x[v].shape[1:4].numel() > 4] else: V_real = list(range(V_used)) # 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 cat_rope_emb_L_1_1_D = torch.cat([rope_emb_L_1_1_D[v] for v in V_real]) # (V*T*H*W, 1, 1, D) f32 # TODO shape # ================================ PER VIEW # for v in range(V_used): for v in V_real: 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 # bookmark(bvh): SAC if sattn_emb_in is not None: raise RuntimeError('this is deprecated in Any4D for the time being, ' 'since it complicates the SAC op_count strategy') 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): for v in V_real: 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): for v in V_real: 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 kwargs['isolate_stub_streams'] = any4d_config.fsdp_stub_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) ]) # NOTE(yams_any4d): V4Head — volumetric action head on pre-final-layer DiT # features (custom/any4d/v4_head.py). Lives on the DiT so FSDP / optimizer / # checkpointer pick it up; params are re-unfrozen after LoRA injection. if getattr(any4d_config, 'v4head_enabled', False): from custom.any4d.v4_head import V4Head _t_out = getattr(any4d_config, 'num_lowdim_adaln_tokens', -1) or -1 if _t_out <= 0: # not yet resolved by surgery at init time _t_out = max(v[1] for v in any4d_config.lowdim_adaln_entries.values()) self.v4head = V4Head( prep_path=any4d_config.v4head_prep_path, d_in=self.model_channels, pred_size=any4d_config.v4head_pred_size, n_z=any4d_config.v4head_n_z, t_out=_t_out, temporal_mode=getattr(any4d_config, 'v4head_temporal_mode', 'per_step'), use_wrist=getattr(any4d_config, 'v4head_use_wrist', True), n_lat_frames=getattr(any4d_config, 'v4head_n_lat_frames', 11), input_channel_norm=getattr(any4d_config, 'v4head_input_channel_norm', False), ) else: self.v4head = None self._v4head_feats = None # stashed per forward when v4head is enabled # ========================================================= # ======== 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. ''' # DEBUG # print(f'@@@@ Any4DDiT.forward() BEGIN: rank={torch.distributed.get_rank()}, iteration={iteration}\n' + \ # '\n'.join([f'{k}: {v.shape} {v.device}' for k, v in x_in_streams.items()])) 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': , # '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): ^ this has a DIFFERENT meaning here compared to Any4DBlock! # NOTE(bvh): ^ in theory, this system should work for non-consecutive views as well, # but in practice we typically construct batches with the subset 0...k-1 for now. # Still, we should support the general case (so do not conflate v_idx and v_pos). # NOTE(bvh): other video_concat_mode values (time, vert, horz, spat) are now deprecated assert any4d_config.video_concat_mode == 'view', \ f'Only video_concat_mode="view" is supported, got "{any4d_config.video_concat_mode}"' # ================================================ # ======== 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. # NOTE(bvh): v_idx = view index (for module lookup), v_pos = position in list. # Dicts are keyed by v_idx; lists are position-ordered (for block interface). x_proj_dict = {} # v_idx -> (B, T, H, W, D) bf16 timesteps_per_view = {} # v_idx -> (B, T) bf16 for v_idx in V_used: x_in_v = x_in_streams[f'v{v_idx}'] # (B, C, T, H, W) timesteps_v = timesteps[f'v{v_idx}'][:, 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_idx](x_in_v) # (B, T, H, W, D) bf16 if v_idx >= 1: x_proj_v += self.view_embs_newviews[v_idx - 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_dict[v_idx] = x_proj_v timesteps_per_view[v_idx] = timesteps_v # Process positional encodings, diffusion timestep encodings, and adaptive layernorm. # Each view gets its own pos/timestep embeddings. rope_emb = {} # v_idx -> (T*H*W, 1, 1, D) f32 t_embedding = {} # v_idx -> (B, T, D) bf16 t_emb_norm = {} # v_idx -> (B, T, D) bf16 adaln_lora = {} # v_idx -> (B, T, D*3) bf16 # legacy_network_behavior=1 pins all per-view t_embedders to V_used[-1] (replicates the # pre-cfc1f99e (2026-03-31) loop-leak bug for old-ckpt inference: _legacy_t_idx = V_used[-1] if any4d_config.legacy_network_behavior == 1 else None for v_idx in V_used: rope_emb[v_idx] = self.pos_embedder(x_proj_dict[v_idx], fps=fps) # (T*H*W, 1, 1, D) f32 _t_idx = _legacy_t_idx if _legacy_t_idx is not None else v_idx (t_emb_v, adaln_v) = my_view_t_embedders[_t_idx](timesteps_per_view[v_idx]) # ^ (B, T, D) bf16, (B, T, D*3) bf16 if any4d_config.has_adaln_stream: t_emb_v = t_emb_v + action_emb_B_D # (B, T, D) bf16 adaln_v = adaln_v + action_emb_B_3D # (B, T, D*3) bf16 t_embedding[v_idx] = t_emb_v # (B, T, D) bf16 t_emb_norm[v_idx] = self.t_embedding_norm(t_emb_v) # (B, T, D) bf16 adaln_lora[v_idx] = adaln_v # (B, T, D*3) bf16 # =============================================== # ======== Handle low-dimensional inputs ======== # =============================================== # DEPRECATED(bvh): sattn_stream will be removed; predictions should go via video streams only. 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 ======== # ========================================== # Convert dicts to position-ordered lists for block interface. x_block_list = [x_proj_dict[v_idx] for v_idx in V_used] # list of (B, T, H, W, D) bf16 rope_emb_list = [rope_emb[v_idx] for v_idx in V_used] # list of (T*H*W, 1, 1, D) f32 t_emb_norm_list = [t_emb_norm[v_idx] for v_idx in V_used] # list of (B, T, D) bf16 adaln_lora_list = [adaln_lora[v_idx] for v_idx in V_used] # list of (B, T, D*3) bf16 # len = 28 for block in self.blocks: (x_block_list, sattn_out) = block( x_B_T_H_W_D=x_block_list, # list of (B, T, H, W, D) bf16 emb_B_T_D=t_emb_norm_list, # list of (B, T, D) bf16 crossattn_emb=crossattn_emb, # (B, NC, D) bf16 rope_emb_L_1_1_D=rope_emb_list, # list of (T*H*W, 1, 1, 128) f32 adaln_lora_B_T_3D=adaln_lora_list, # 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_list = x_block_list # list of (B, T, H, W, D) bf16 # y_block_list: [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 ======== # ================================================= # Convert block output (position-ordered list) back to v_idx-keyed dict. y_block_dict = {v_idx: y_block_list[v_pos] for v_pos, v_idx in enumerate(V_used)} # ^ v_idx -> (B, T, H, W, D) bf16 # NOTE(yams_any4d): stash pre-final-layer features (NOT detached — the V4Head # loss backprops through them into the LoRA adapters) for a4d_model._v4head_compute. if self.v4head is not None: self._v4head_feats = {v_idx: y_block_dict[v_idx] for v_idx in V_used} # Apply per-view output projection (final layer). 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_idx in V_used: y_proj_v = my_final_layers[v_idx]( y_block_dict[v_idx], t_embedding[v_idx], adaln_lora_B_T_3D=adaln_lora[v_idx]) y_out_streams[f'v{v_idx}'] = self.unpatchify(y_proj_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} # DEBUG # print(f'@@@@ Any4DDiT.forward() END: rank={torch.distributed.get_rank()}, iteration={iteration}\n' + \ # '\n'.join([f'{k}: {v.shape} {v.device}' for k, v in y_out_streams.items()])) return y_out_streams