# Any4D DiT built on top of Wan 2.1 Fun-V1.1-1.3B-Control. # # Differences from the InP variant in ``a4d_network_wan.py``: # - patch_embedding in_dim = 48 (not 36) # Layout per DiffSynth's WanVideoUnit_FunControl (wan_video.py:542-557): # [0:16] noisy RGB latent # [16:32] control latent (VAE-encoded; canny / depth / seg / ... — any 16-ch # pseudo-RGB pretrained Wan-VAE can encode) # [32:48] y latent (zeros for pure text+control mode; would carry InP ref # latent for I2V mode) # No 4-channel mask slot (Fun-Control's mask is implicit / outside the DiT input). # - has_ref_conv = True # Separate Conv2d that maps a 16-ch reference latent into a (B, dim, H/2, W/2) # feature, added token-side to the first frame. Only used when a reference image # is provided. For pure text+control -> video training, no reference image is # passed; the ref_conv weights are loaded from the checkpoint but unused. # - has_image_input = True # CLIP image embeddings are still consumed via cross-attention (k_img/v_img branch). # For pure T2V, we pass zero CLIP context (img_context_len=0), which routes # around the k_img/v_img branch. import math from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.utils.checkpoint from einops import rearrange from custom.any4d.a4d_config import Any4DConfig from custom.wan.a4d_network_wan import ( Any4DWanDiT, LowDimProjLayer, WanModel, sinusoidal_embedding_1d, ) from imaginaire.utils import log class WanModelControl(WanModel): """Wan 2.1 DiT (Fun-V1.1-Control variant) with an extra ref_conv module. Inherits everything from ``WanModel`` (patch_embedding, blocks, head, ...) and only adds the ``ref_conv`` Conv2d that the Control checkpoint ships with. ``ref_conv`` is loaded so the checkpoint's keys all land on a parameter — but its activation is wired in via ``Any4DWanControlDiT.forward`` only when a reference image is supplied (currently never in this integration). """ def __init__(self, *args, has_ref_conv: bool = True, ref_conv_in_dim: int = 16, **kwargs): super().__init__(*args, **kwargs) self.has_ref_conv = has_ref_conv if has_ref_conv: # Matches diffsynth's WanModel: nn.Conv2d(16, dim, (2, 2), stride=(2, 2)). self.ref_conv = nn.Conv2d( ref_conv_in_dim, self.dim, kernel_size=(2, 2), stride=(2, 2), ) class Any4DWanControlDiT(nn.Module): """Any4D DiT with Wan 2.1 Fun-V1.1-1.3B-Control as the base video diffusion model. Forward signature matches ``Any4DWanDiT.forward()`` so the pipeline / loss / visualization code is unchanged. The only differences from ``Any4DWanDiT``: * Wan kwargs flip ``in_dim`` 36 -> 48 and add ``has_ref_conv=True``. * ``_to_wan_48ch`` replaces ``_to_wan_36ch``, mapping the 50-ch per-view stream (rgb + masks + ref + edge_control) into the 48-ch DiT input layout (noisy / control / ref). """ # Wan 2.1 Fun-V1.1-1.3B-Control canonical kwargs (verified from the published # checkpoint's config.json): # dim=1536, in_dim=48, ffn_dim=8960, out_dim=16, text_dim=4096, freq_dim=256, # eps=1e-6, patch_size=[1,2,2], num_heads=12, num_layers=30, # has_image_input=True (model_type=i2v), has_ref_conv=True (in_dim_ref_conv=16), # add_control_adapter=False (in_dim_control_adapter=24 is unused in this ckpt). WAN_1_3B_CONTROL_KWARGS: Dict[str, Any] = dict( dim=1536, in_dim=48, ffn_dim=8960, out_dim=16, text_dim=4096, freq_dim=256, eps=1e-6, patch_size=(1, 2, 2), num_heads=12, num_layers=30, has_image_input=True, has_image_pos_emb=False, ) def __init__( self, any4d_config: Any4DConfig, wan_model_kwargs: Optional[dict] = None, view_emb_std: float = 0.06, ): super().__init__() from custom.any4d import a4d_surgery any4d_config = a4d_surgery.validate_config(any4d_config) self.any4d_config = any4d_config # LEARNABLE orthonormal projection for PPD structured noise: a (C, K) parameter, # re-orthonormalized via QR each step in Any4DWanModel.compute_loss. Kept inside # the DiT so FSDP grad-sync, the optimizer, and the checkpointer treat it like any # other DiT weight. Initialized from a seeded semi-orthogonal matrix (reproducible # start). P is ALWAYS learned — there is no fixed-P path. self.sn_P_raw = None if getattr(any4d_config, 'use_structured_noise', False): # The Control DiT supports only setup 1: a seeded, jointly-learned, SHARED P. The # pretrain-then-freeze / per-modality-P options are setup 2 (wan_variant='t2v', handled # by the base Any4DWanDiT) — fail loudly rather than silently training with a random P. for _flag in ('structured_noise_p_init_path', 'structured_noise_freeze_p', 'structured_noise_per_modality_p'): assert not getattr(any4d_config, _flag, None), ( f'{_flag} is set but wan_variant=control supports only a seeded, jointly-learned ' f'shared P (setup 1); use wan_variant=t2v for a pretrained/frozen/per-modality P.') from custom.any4d.a4d_logistics import rgb_latent_span from custom.wan.structured_noise import _semi_orthogonal_P _rs, _re = rgb_latent_span(any4d_config.video_entries[0]) _C = _re - _rs _K = int(getattr(any4d_config, 'structured_noise_channels', 16)) if 0 < _K < _C: _P0 = _semi_orthogonal_P(_C, _K, int(getattr(any4d_config, 'structured_noise_seed', 0)), torch.device('cpu'), torch.float32) # (K, C) self.sn_P_raw = nn.Parameter(_P0.t().contiguous()) # (C, K), trainable wan_kwargs = dict(self.WAN_1_3B_CONTROL_KWARGS) if wan_model_kwargs: wan_kwargs.update(wan_model_kwargs) self.wan = WanModelControl(has_ref_conv=True, **wan_kwargs) self.model_channels = wan_kwargs['dim'] self.patch_size = tuple(wan_kwargs['patch_size']) self.num_heads = wan_kwargs['num_heads'] self.head_dim = self.model_channels // self.num_heads self.text_dim = wan_kwargs['text_dim'] self.has_image_input = wan_kwargs['has_image_input'] # Trade recompute for activation memory across the Wan DiT block stack. Required # for many-view JOINT cross-view training at full resolution (the concatenated # per-view token sequence makes the self-attention activations dominate VRAM). self.use_grad_ckpt = bool(getattr(any4d_config, 'wan_gradient_checkpointing', False)) self.grad_ckpt_period = max(1, int(getattr(any4d_config, 'wan_grad_ckpt_period', 1))) # Selective Activation Checkpointing (SAC), same scheme as the cosmos DiT # (cosmos_predict2/networks/selective_activation_checkpoint.py): wrap blocks with # PyTorch checkpoint_wrapper + a per-op context_fn. mode='mm_only' SAVES matmuls + # attention (SDPA/flash) and RECOMPUTES only cheap ops (norms/GELU/elementwise) -> # NO attention recompute -> much faster than full-block recompute, at higher VRAM # (which we have). mode='block_wise' = full-block recompute (context_fn=None). self.sac_mode = str(getattr(any4d_config, 'wan_sac_mode', 'mm_only')) if self.use_grad_ckpt: from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( CheckpointImpl, checkpoint_wrapper as _ckpt_wrap) _ctx_fn = None if self.sac_mode == 'mm_only': from cosmos_predict2.networks.selective_activation_checkpoint import mm_only_context_fn _ctx_fn = mm_only_context_fn for _bi in range(len(self.wan.blocks)): if _bi % self.grad_ckpt_period == 0: self.wan.blocks[_bi] = _ckpt_wrap( self.wan.blocks[_bi], checkpoint_impl=CheckpointImpl.NO_REENTRANT, context_fn=_ctx_fn, preserve_rng_state=False, ) if _ctx_fn is not None else _ckpt_wrap( self.wan.blocks[_bi], checkpoint_impl=CheckpointImpl.NO_REENTRANT, preserve_rng_state=False, ) log.info(f"[Wan] SAC enabled: mode={self.sac_mode}, every {self.grad_ckpt_period} " f"block(s), {len(self.wan.blocks)} total") num_views = any4d_config.num_views self.x_embedder = self.wan.patch_embedding self.final_existing = self.wan.head self.x_embedder_newviews = nn.ModuleList() self.final_layer_newviews = nn.ModuleList() # Per-view positional markers — same convention as Any4DWanDiT. self.view_emb_existing = nn.Parameter(torch.zeros(1, self.model_channels)) self.view_embs_newviews = nn.ParameterList([ nn.Parameter(torch.randn(1, self.model_channels) * view_emb_std) for _ in range(num_views - 1) ]) self._view_emb_init_std = view_emb_std # Low-dim streams (adaln / sattn / xattn) — copy of Any4DWanDiT logic. if any4d_config.has_adaln_stream: Ta = any4d_config.num_lowdim_adaln_tokens Ca = any4d_config.num_lowdim_adaln_channels self.action_embedder = nn.Sequential( nn.Linear(Ta * Ca, self.model_channels), nn.SiLU(), nn.Linear(self.model_channels, self.model_channels), ) with torch.no_grad(): nn.init.zeros_(self.action_embedder[-1].weight) if self.action_embedder[-1].bias is not None: nn.init.zeros_(self.action_embedder[-1].bias) else: self.action_embedder = None if any4d_config.has_sattn_stream: h = min(any4d_config.lowdim_hidden_channels, any4d_config.num_lowdim_sattn_channels, self.model_channels // 4) self.sattn_inproj = LowDimProjLayer( any4d_config.lowdim_sattn_proj_mode, any4d_config.num_lowdim_sattn_tokens, any4d_config.num_lowdim_sattn_channels, h, self.model_channels) self.sattn_outproj = LowDimProjLayer( any4d_config.lowdim_sattn_proj_mode, any4d_config.num_lowdim_sattn_tokens, self.model_channels, h, any4d_config.num_lowdim_sattn_channels) Any4DWanDiT._zero_last_linear(self.sattn_outproj) if any4d_config.has_xattn_stream: h = min(any4d_config.lowdim_hidden_channels, any4d_config.num_lowdim_xattn_channels, self.model_channels // 4) self.xattn_inproj = LowDimProjLayer( any4d_config.lowdim_xattn_proj_mode, any4d_config.num_lowdim_xattn_tokens, any4d_config.num_lowdim_xattn_channels, h, self.text_dim) Any4DWanDiT._zero_last_linear(self.xattn_inproj) self.accum_video_sample_counter = 0 # --- Stubs expected by the trainer / context-parallel machinery --- def disable_context_parallel(self): pass def enable_context_parallel(self, *a, **kw): pass @property def is_context_parallel_enabled(self) -> bool: return False def fully_shard(self, mesh=None, **kwargs): return self def _to_wan_48ch(self, x_v_50: torch.Tensor) -> torch.Tensor: """Reshape Any4D's 50-ch per-view stream into Wan Fun-Control's 48-ch DiT input. Per-view stream layout (set by video_entries in the experiment config): [0:16] noisy RGB latent (xt / target) [16:17] input_mask (Any4D bookkeeping, dropped) [17:18] output_mask (Any4D bookkeeping, dropped) [18:34] ref latent — the Wan Fun-Control y/image-conditioning slot. For a view-synth ANCHOR view this carries the anchor's OWN full clean VAE latent (reference-image broadcast, built by a4d_vae_wan._build_wan_anchor_ref_latent); for every other view and for ALL legacy/forecast Control runs it is ZEROS. [34:50] edge_control latent (VAE-encoded canny pseudo-RGB; future modalities like _depth_control / _seg_control would land here too) Wan Fun-Control DiT input layout (matches DiffSynth's WanVideoUnit_FunControl: ``y = torch.cat([control_latents, zeros], dim=1)`` then ``cat([x, y], dim=1)``): [0:16] noisy RGB latent [16:32] control latent [32:48] y / reference-image latent (zeros in pure text+control mode; the clean anchor latent in view-synth mode) We now PASS THE REF SLOT THROUGH to the pretrained image-conditioning ('y') channel [32:48]. Fun-Control's y channel was trained with either all-zeros (text+control mode) or a VAE-encoded REFERENCE IMAGE broadcast (image-conditioned mode). a4d_vae_wan fills the [18:34] slot with the latter recipe ONLY for the view-synth anchor view (its own clean latent), and ZEROS otherwise — so legacy / forecast Control runs (anchor absent) keep y=0 and stay on-distribution, while the anchor view feeds a clean reference image into the I2V path. The target view then attends to the anchor's clean ref tokens via the 2-view joint attention. """ noisy_slot = x_v_50[:, 0:16] control_slot = x_v_50[:, 34:50] # y / reference-image channel: anchor's clean latent when present, else zeros # (the ref slot is guaranteed zero for non-anchor views and legacy runs). ref_slot = x_v_50[:, 18:34] return torch.cat([noisy_slot, control_slot, ref_slot], dim=1) 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]: cfg = self.any4d_config V_max = cfg.num_views V_used = [v for v in range(V_max) if f'v{v}' in x_in_streams] assert cfg.video_concat_mode == 'view', \ f'Only video_concat_mode="view" is supported, got "{cfg.video_concat_mode}"' # Time conditioning (matches Any4DWanDiT). ref_key = f'v{V_used[0]}' t_flat = timesteps[ref_key].flatten(1).mean(dim=1) t_emb = sinusoidal_embedding_1d(self.wan.freq_dim, t_flat).to(crossattn_emb.dtype) t_emb = self.wan.time_embedding(t_emb) if cfg.has_adaln_stream: action = x_in_streams['adaln'] action_flat = rearrange(action, 'b t c -> b (t c)') t_emb = t_emb + self.action_embedder(action_flat) t_mod = self.wan.time_projection(t_emb).unflatten(1, (6, self.model_channels)) # Cross-attention context = T5 text only. # has_image_input=True but img_context_len=0 routes around the k_img/v_img branch, # which keeps Fun-Control in pure-T2V mode (no CLIP image conditioning). context = self.wan.text_embedding(crossattn_emb) img_context_len = 0 if cfg.has_xattn_stream: xattn_in = x_in_streams['xattn'] xattn_text = self.xattn_inproj(xattn_in) extra_ctx = self.wan.text_embedding(xattn_text) context = torch.cat([context, extra_ctx], dim=1) # Per-view patch embedding. view_embs = [self.view_emb_existing] + list(self.view_embs_newviews) token_parts: List[torch.Tensor] = [] grid_shapes: List[Tuple[int, int, int]] = [] seq_lens: List[int] = [] for v_idx in V_used: x_v_50 = x_in_streams[f'v{v_idx}'] x_v_48 = self._to_wan_48ch(x_v_50) tokens = self.x_embedder(x_v_48) B, D, Tt, Hh, Ww = tokens.shape tokens = rearrange(tokens, 'b d t h w -> b (t h w) d') tokens = tokens + view_embs[v_idx].view(1, 1, D) token_parts.append(tokens) grid_shapes.append((Tt, Hh, Ww)) seq_lens.append(Tt * Hh * Ww) sattn_slot = None if cfg.has_sattn_stream: sattn_in = x_in_streams['sattn'] sattn_emb = self.sattn_inproj(sattn_in) sattn_slot = (sum(seq_lens), sattn_emb.shape[1]) token_parts.append(sattn_emb) x_tokens = torch.cat(token_parts, dim=1) # Reuse parent's RoPE freq builder. freqs = Any4DWanDiT._build_freqs(self, grid_shapes, sattn_slot, x_tokens.device) # Wan DiT blocks. DiTBlock.forward signature (see a4d_network_wan.py): # forward(self, x, context, t_mod, freqs, img_context_len: int = 0) # checkpoint() forwards *args + **kwargs straight to the block. x = x_tokens for block in self.wan.blocks: # Blocks are pre-wrapped with the SAC/checkpoint_wrapper in __init__ when # gradient checkpointing is enabled, so a plain call self-checkpoints (recompute, # per the mm_only/block_wise policy, happens only during backward). x = block(x, context, t_mod, freqs, img_context_len=img_context_len) # Per-view head; pad to per-view input channel count so cs_skip*yt + cs_out*net_out # preserves the conditioning slots (control + ref). Cout_pred = 16 pt, ph, pw = self.patch_size y_out_streams: Dict[str, torch.Tensor] = {} offset = 0 for v_pos, v_idx in enumerate(V_used): L_v = seq_lens[v_pos] tokens_v = x[:, offset:offset + L_v] offset += L_v Tt, Hh, Ww = grid_shapes[v_pos] # Head expects t_mod shape (B, 2, dim) so its `(modulation + t_mod).chunk(2, dim=1)` # works. Passing t_emb (B, dim) broadcasts correctly only when B == 1 # (right-align makes it (1, 1, dim)); for B > 1, the right-align becomes # (1, B, dim) and collides with modulation's (1, 2, dim). Unsqueeze to # (B, 1, dim) so broadcast yields (B, 2, dim) regardless of B. y_v = self.final_existing(tokens_v, t_emb.unsqueeze(1).to(tokens_v.dtype)) y_v = rearrange( y_v, 'b (t h w) (pt ph pw c) -> b c (t pt) (h ph) (w pw)', t=Tt, h=Hh, w=Ww, pt=pt, ph=ph, pw=pw, c=Cout_pred, ) C_total = cfg.num_video_in_channels[v_idx] if C_total > Cout_pred: B_, _, T_, H_, W_ = y_v.shape pad = torch.zeros(B_, C_total - Cout_pred, T_, H_, W_, device=y_v.device, dtype=y_v.dtype) y_v = torch.cat([y_v, pad], dim=1) y_out_streams[f'v{v_idx}'] = y_v if sattn_slot is not None: s_off, s_len = sattn_slot sattn_out = x[:, s_off:s_off + s_len] y_out_streams['sattn'] = self.sattn_outproj(sattn_out) return y_out_streams __all__ = ['Any4DWanControlDiT', 'WanModelControl']