# Any4D DiT built on top of Wan 2.1 Fun 1.3B InP as the base video diffusion model.
#
# This file bundles the Wan 2.1 DiT source directly (stripped of DiffSynth's
# dance/music/camera extras) so we do not depend on diffsynth at runtime.
#
# The contract of Any4DWanDiT.forward() matches custom/any4d/a4d_network.py:Any4DDiT.forward()
# so the rest of the Any4D pipeline (pack_streams_from_entries, denoise,
# assemble_network_inputs, …) can call it unmodified.
#
# Wan 2.1 Fun 1.3B InP hyperparameters (canonical):
#     dim=1536, ffn_dim=8960, num_heads=12, num_layers=30,
#     patch_size=(1, 2, 2), in_dim=36 (16 RGB + 4 mask + 16 ref), out_dim=16,
#     text_dim=4096 (UMT5-XXL), freq_dim=256, eps=1e-6, has_image_input=True.
#
# Original source: https://github.com/modelscope/DiffSynth-Studio
#     diffsynth/models/wan_video_dit.py (Apache-2.0)

import math
import os
from typing import Any, Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

from custom.any4d.a4d_config import Any4DConfig
from imaginaire.utils import log


# ================================================================================
# ============================ Wan 2.1 DiT (inline) ==============================
# ================================================================================

# Set A4D_WAN_FORCE_SDPA=1 to skip FlashAttention and use torch SDPA. Needed on GPUs
# without a compiled FlashAttention kernel image (e.g. Blackwell sm_120 with an FA3
# build compiled for Hopper sm_90, which otherwise hard-exits the process with
# "no kernel image is available for execution on the device"). Default is unchanged.
_FORCE_SDPA = os.environ.get('A4D_WAN_FORCE_SDPA', '0') == '1'

try:
    import flash_attn_interface
    FLASH_ATTN_3_AVAILABLE = not _FORCE_SDPA
except ModuleNotFoundError:
    FLASH_ATTN_3_AVAILABLE = False

try:
    import flash_attn
    FLASH_ATTN_2_AVAILABLE = not _FORCE_SDPA
except ModuleNotFoundError:
    FLASH_ATTN_2_AVAILABLE = False


def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int):
    if FLASH_ATTN_3_AVAILABLE:
        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
        x = flash_attn_interface.flash_attn_func(q, k, v)
        if isinstance(x, tuple):
            x = x[0]
        return rearrange(x, "b s n d -> b s (n d)", n=num_heads)
    if FLASH_ATTN_2_AVAILABLE:
        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
        x = flash_attn.flash_attn_func(q, k, v)
        return rearrange(x, "b s n d -> b s (n d)", n=num_heads)
    q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
    k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
    v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
    x = F.scaled_dot_product_attention(q, k, v)
    return rearrange(x, "b n s d -> b s (n d)", n=num_heads)


def modulate(x, shift, scale):
    return x * (1 + scale) + shift


def sinusoidal_embedding_1d(dim, position):
    sinusoid = torch.outer(
        position.type(torch.float64),
        torch.pow(10000,
                  -torch.arange(dim // 2, dtype=torch.float64, device=position.device).div(dim // 2)),
    )
    return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1).to(position.dtype)


def _precompute_freqs_cis_1d(dim: int, end: int = 1024, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].double() / dim))
    freqs = torch.outer(torch.arange(end, device=freqs.device), freqs)
    return torch.polar(torch.ones_like(freqs), freqs)  # complex64


def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0):
    f = _precompute_freqs_cis_1d(dim - 2 * (dim // 3), end, theta)
    h = _precompute_freqs_cis_1d(dim // 3, end, theta)
    w = _precompute_freqs_cis_1d(dim // 3, end, theta)
    return f, h, w


def rope_apply(x, freqs, num_heads):
    x = rearrange(x, "b s (n d) -> b s n d", n=num_heads)
    x_out = torch.view_as_complex(
        x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2))
    x_out = torch.view_as_real(x_out * freqs).flatten(2)
    return x_out.to(x.dtype)


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
        self.normalized_shape = (dim,)

    def forward(self, x):
        dtype = x.dtype
        return (x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
                ).to(dtype) * self.weight


class AttentionModule(nn.Module):
    def __init__(self, num_heads):
        super().__init__()
        self.num_heads = num_heads

    def forward(self, q, k, v):
        return flash_attention(q, k, v, num_heads=self.num_heads)


class SelfAttention(nn.Module):
    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.q = nn.Linear(dim, dim)
        self.k = nn.Linear(dim, dim)
        self.v = nn.Linear(dim, dim)
        self.o = nn.Linear(dim, dim)
        self.norm_q = RMSNorm(dim, eps=eps)
        self.norm_k = RMSNorm(dim, eps=eps)
        self.attn = AttentionModule(num_heads)

    def forward(self, x, freqs):
        q = self.norm_q(self.q(x))
        k = self.norm_k(self.k(x))
        v = self.v(x)
        q = rope_apply(q, freqs, self.num_heads)
        k = rope_apply(k, freqs, self.num_heads)
        return self.o(self.attn(q, k, v))


class CrossAttention(nn.Module):
    """Cross-attention to text embeddings, with optional CLIP-image sub-branch (used by InP)."""

    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.q = nn.Linear(dim, dim)
        self.k = nn.Linear(dim, dim)
        self.v = nn.Linear(dim, dim)
        self.o = nn.Linear(dim, dim)
        self.norm_q = RMSNorm(dim, eps=eps)
        self.norm_k = RMSNorm(dim, eps=eps)
        self.has_image_input = has_image_input
        if has_image_input:
            self.k_img = nn.Linear(dim, dim)
            self.v_img = nn.Linear(dim, dim)
            self.norm_k_img = RMSNorm(dim, eps=eps)
        self.attn = AttentionModule(num_heads)

    def forward(self, x: torch.Tensor, y: torch.Tensor, img_context_len: int = 0):
        if self.has_image_input and img_context_len > 0:
            img = y[:, :img_context_len]
            ctx = y[:, img_context_len:]
        else:
            ctx = y
            img = None
        q = self.norm_q(self.q(x))
        k = self.norm_k(self.k(ctx))
        v = self.v(ctx)
        x_out = self.attn(q, k, v)
        if img is not None:
            k_img = self.norm_k_img(self.k_img(img))
            v_img = self.v_img(img)
            x_out = x_out + flash_attention(q, k_img, v_img, num_heads=self.num_heads)
        return self.o(x_out)


class GateModule(nn.Module):
    def forward(self, x, gate, residual):
        return x + gate * residual


class DiTBlock(nn.Module):
    def __init__(self, has_image_input: bool, dim: int, num_heads: int,
                 ffn_dim: int, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.ffn_dim = ffn_dim
        self.self_attn = SelfAttention(dim, num_heads, eps)
        self.cross_attn = CrossAttention(dim, num_heads, eps, has_image_input=has_image_input)
        self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
        self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
        self.norm3 = nn.LayerNorm(dim, eps=eps)
        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim),
                                 nn.GELU(approximate='tanh'),
                                 nn.Linear(ffn_dim, dim))
        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim ** 0.5)
        self.gate = GateModule()

    def forward(self, x, context, t_mod, freqs, img_context_len: int = 0):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1)
        input_x = modulate(self.norm1(x), shift_msa, scale_msa)
        x = self.gate(x, gate_msa, self.self_attn(input_x, freqs))
        x = x + self.cross_attn(self.norm3(x), context, img_context_len=img_context_len)
        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
        x = self.gate(x, gate_mlp, self.ffn(input_x))
        return x


class MLP(nn.Module):
    def __init__(self, in_dim, out_dim, has_pos_emb=False):
        super().__init__()
        self.proj = nn.Sequential(
            nn.LayerNorm(in_dim),
            nn.Linear(in_dim, in_dim),
            nn.GELU(),
            nn.Linear(in_dim, out_dim),
            nn.LayerNorm(out_dim),
        )
        self.has_pos_emb = has_pos_emb
        if has_pos_emb:
            self.emb_pos = nn.Parameter(torch.zeros((1, 514, 1280)))

    def forward(self, x):
        if self.has_pos_emb:
            x = x + self.emb_pos.to(dtype=x.dtype, device=x.device)
        return self.proj(x)


class Head(nn.Module):
    def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.out_dim = out_dim
        self.patch_size = patch_size
        self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
        self.head = nn.Linear(dim, out_dim * math.prod(patch_size))
        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim ** 0.5)

    def forward(self, x, t_mod):
        # t_mod arrives as the time embedding (B, D); add a singleton mod-param axis so it
        # broadcasts against self.modulation (1, 2, D) for ANY batch size. Without this,
        # batch=1 worked only by accidental broadcast while batch>1 raised
        # "size of tensor a (2) must match b (B) at dim 1" (Wan's original Head unsqueezes here).
        if t_mod.dim() == 2:
            t_mod = t_mod.unsqueeze(1)  # (B, D) -> (B, 1, D)
        shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(2, dim=1)
        return self.head(self.norm(x) * (1 + scale) + shift)


class WanModel(nn.Module):
    """Wan 2.1 DiT (I2V / T2V variant, stripped of DiffSynth dance/music/camera extras)."""

    def __init__(
        self,
        dim: int,
        in_dim: int,
        ffn_dim: int,
        out_dim: int,
        text_dim: int,
        freq_dim: int,
        eps: float,
        patch_size: Tuple[int, int, int],
        num_heads: int,
        num_layers: int,
        has_image_input: bool,
        has_image_pos_emb: bool = False,
    ):
        super().__init__()
        self.dim = dim
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.ffn_dim = ffn_dim
        self.freq_dim = freq_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.has_image_input = has_image_input
        self.patch_size = patch_size
        self.eps = eps

        self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
        self.text_embedding = nn.Sequential(
            nn.Linear(text_dim, dim),
            nn.GELU(approximate='tanh'),
            nn.Linear(dim, dim),
        )
        self.time_embedding = nn.Sequential(
            nn.Linear(freq_dim, dim),
            nn.SiLU(),
            nn.Linear(dim, dim),
        )
        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
        self.blocks = nn.ModuleList([
            DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps) for _ in range(num_layers)
        ])
        self.head = Head(dim, out_dim, patch_size, eps)

        head_dim = dim // num_heads
        self.freqs = precompute_freqs_cis_3d(head_dim)

        if has_image_input:
            # CLIP image feature projector (clip_feature_dim = 1280 for open-clip ViT-H).
            self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb)

    def patchify(self, x: torch.Tensor):
        x = self.patch_embedding(x)
        f, h, w = x.shape[2], x.shape[3], x.shape[4]
        x = rearrange(x, 'b c f h w -> b (f h w) c')
        return x, (f, h, w)

    def unpatchify(self, x: torch.Tensor, grid_size):
        return rearrange(
            x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)',
            f=grid_size[0], h=grid_size[1], w=grid_size[2],
            x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2],
        )


# ================================================================================
# ============================= Any4D wrapper ====================================
# ================================================================================


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.out_channels = out_channels
        if proj_mode == 'per_token':
            self.proj = nn.ModuleList([
                nn.Sequential(
                    nn.Linear(in_channels, hidden_channels),
                    nn.GELU(approximate='tanh'),
                    nn.Linear(hidden_channels, out_channels),
                ) for _ in range(num_tokens)
            ])
        elif proj_mode == 'shared':
            self.proj = nn.Sequential(
                nn.Linear(in_channels, hidden_channels),
                nn.GELU(approximate='tanh'),
                nn.Linear(hidden_channels, out_channels),
            )
            self.id_embs = nn.Parameter(torch.randn(num_tokens, out_channels))
        else:
            raise ValueError(f'Invalid proj_mode: {proj_mode}')

    def forward(self, x):
        B = x.shape[0]
        assert x.shape == (B, self.num_tokens, self.in_channels)
        if self.proj_mode == 'per_token':
            return torch.stack([p(x[:, i, :]) for i, p in enumerate(self.proj)], dim=-2)
        return self.proj(x) + self.id_embs.view(1, self.num_tokens, self.out_channels)


class Any4DWanDiT(nn.Module):
    """Any4D DiT with Wan 2.1 Fun 1.3B InP as the base video diffusion backbone.

    Forward signature matches `custom/any4d/a4d_network.py:Any4DDiT.forward()`:
        x_in_streams: dict of stream_name -> (B, C, T, H, W) or (B, T, D).
        timesteps: dict of stream_name -> (B, 1, T, 1, 1) or (B, T, 1).
        crossattn_emb: (B, Nc, text_dim) — text embeddings.
    Returns y_out_streams dict with the same keys.
    """

    # Wan 2.1 Fun 1.3B InP canonical kwargs.
    WAN_1_3B_INP_KWARGS: Dict[str, Any] = dict(
        dim=1536,
        in_dim=36,          # 16 RGB latent + 4 input-frame mask + 16 reference latent.
        ffn_dim=8960,
        out_dim=16,         # predicts 16-channel latent.
        text_dim=4096,      # UMT5-XXL.
        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,
    ):
        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 / optimizer / checkpointer treat it like any other
        # weight. Initialized from a seeded semi-orthogonal matrix (reproducible start);
        # P is ALWAYS learned — there is no fixed-P path. Built only in K-subspace mode
        # (0 < K < C); full-channel PPD (K >= C) doesn't use a projection.
        self.sn_P_raw = None
        if getattr(any4d_config, 'use_structured_noise', False):
            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)
                _raw0 = _P0.t().contiguous()                                  # (C, K)
                _pinit = getattr(any4d_config, 'structured_noise_p_init_path', '') or ''
                _per_mod = bool(getattr(any4d_config, 'structured_noise_per_modality_p', False))
                _frozen = bool(getattr(any4d_config, 'structured_noise_freeze_p', False))

                def _reg_buffer(name, t):
                    if name in self.__dict__:                 # drop placeholder/None attr
                        del self.__dict__[name]
                    self.register_buffer(name, t.contiguous(), persistent=True)

                if _per_mod:
                    # PER-MODALITY P (variant 4: pretrain-then-freeze, per-modality). Load a dict
                    # {modality -> (C, K)} of pre-optimized projections. rgb -> sn_P_raw (the noise
                    # projection P_rgb, so single-P consumers still work); each map -> sn_P_raw_{m}
                    # (the phase-source projection for that modality). All registered as FSDP-safe
                    # buffers (per-modality is currently freeze-only).
                    from custom.wan.structured_noise import SN_MODALITIES
                    assert _frozen, 'per-modality P is supported only with structured_noise_freeze_p=True'
                    assert _pinit, 'per-modality P requires structured_noise_p_init_path (a dict {mod:(C,K)})'
                    _pdict = torch.load(_pinit, map_location='cpu')
                    assert isinstance(_pdict, dict), \
                        f'per-modality p_init must be a dict {{mod:(C,K)}}, got {type(_pdict)}'
                    for _m in SN_MODALITIES:
                        assert _m in _pdict, f'per-modality P dict missing {_m!r}: have {list(_pdict)}'
                        _t = _pdict[_m].float()
                        assert tuple(_t.shape) == tuple(_raw0.shape), \
                            f'per-modality P[{_m}] shape {tuple(_t.shape)} != {tuple(_raw0.shape)}'
                        _reg_buffer('sn_P_raw' if _m == 'rgb' else f'sn_P_raw_{_m}', _t)
                    log.info(f'[sn_P_raw] PER-MODALITY frozen buffers for {SN_MODALITIES} from {_pinit}')
                else:
                    # Single shared P. Optionally init from a pre-optimized (C,K) tensor (variant 3).
                    if _pinit:
                        _pre = torch.load(_pinit, map_location='cpu').float()
                        assert tuple(_pre.shape) == tuple(_raw0.shape), \
                            f'pre-optimized P shape {tuple(_pre.shape)} != {tuple(_raw0.shape)}'
                        _raw0 = _pre.contiguous()
                        log.info(f'[sn_P_raw] initialized from pre-optimized P: {_pinit}')
                    # FROZEN mode: register as a BUFFER (not a Parameter). A buffer is absent from
                    # model.parameters() so the optimizer can never update it, and — critically —
                    # FSDP does not flatten buffers into trainable FlatParameters (which silently
                    # re-enables grad on a requires_grad=False Parameter mixed with trainable ones,
                    # the bug that let the "frozen" P drift in the first syncfz run). The buffer is
                    # still saved in state_dict / copied into the EMA module, so eval loads it.
                    if _frozen:
                        _reg_buffer('sn_P_raw', _raw0)
                        log.info('[sn_P_raw] FROZEN as buffer — pretrain-then-freeze mode (FSDP-safe)')
                    else:
                        self.sn_P_raw = nn.Parameter(_raw0)                  # (C, K), trainable

        wan_kwargs = dict(self.WAN_1_3B_INP_KWARGS)
        if wan_model_kwargs:
            wan_kwargs.update(wan_model_kwargs)
        self.wan = WanModel(**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']

        num_views = any4d_config.num_views

        # ---- Shared Wan InP patch embed & head (pretrained, used by all views) ----
        # Any4D data pipeline produces 18-ch per-view tensors (rgb latent + input_mask +
        # output_mask). We reshape them into Wan's 36-ch InP layout (rgb + 4-ch mask +
        # 16-ch reference) inside `_to_wan_36ch` just before this Conv3d, so the
        # pretrained Wan patch_embedding is used as-is across all views.
        # Fun-InP (in_dim=36) sets has_image_input=True; the pure T2V subclass
        # (Any4DWanT2VDiT, in_dim=16) sets it False. The CLIP image cross-attn branch
        # is dead code in this integration anyway (forward always passes img_context_len=0),
        # so both are fine — the flag only controls whether WanModel builds img_emb/k_img/v_img.
        if not self.has_image_input:
            log.info('Any4DWanDiT: has_image_input=False (pure T2V; CLIP image branch disabled)')
        self.x_embedder = self.wan.patch_embedding          # alias — kept for forward loop
        self.final_existing = self.wan.head                 # alias — same, for all views
        self.x_embedder_newviews = nn.ModuleList()          # unused: all views share x_embedder
        self.final_layer_newviews = nn.ModuleList()         # unused: all views share final_existing

        # ---- View embeddings (positional marker per view) ----
        # v0 stays at zero so v0 tokens match Wan's native training distribution exactly.
        # v1, v2 get a small random signature so Wan's self-attention can distinguish
        # them from v0 at step 0 (otherwise all three views are permutation-symmetric
        # and the attention sees the joint token sequence as homogeneous, which
        # empirically collapses target-frame generation toward the latent origin).
        view_emb_std = any4d_config.view_emb_std
        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 ----
        # action_embedder projects trajectory tokens into a timestep-level offset.
        # Final Linear is zero-initialized so the trajectory adds zero to t_emb at init
        # → Wan's time conditioning is unperturbed. The model learns non-zero output during
        # training. This is the ControlNet zero-conv trick applied to the adaln stream.
        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)
            # Zero the sattn output projection so sattn stream is a no-op at init.
            self._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)
            # xattn tokens attach to the *text* context, so they live in text_dim space.
            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)
            # Zero the xattn in-projection so appended context tokens are zero at init.
            self._zero_last_linear(self.xattn_inproj)

        # Sample-counter stub used by VidarModel._update_train_stats.
        self.accum_video_sample_counter = 0

    # ------------------------------------------------------------------
    # Stubs expected by the Cosmos-Predict2 pipeline helpers (context-parallel / FSDP
    # coordination). Any4D's Wan path is single-GPU (or DDP), so these are no-ops.
    # ------------------------------------------------------------------
    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):
        """Per-block FSDP sharding of the Wan transformer stack.
        """
        if mesh is None:
            return self
        from torch.distributed._composable.fsdp import fully_shard as _fully_shard
        blocks = self.wan.blocks
        for i, block in enumerate(blocks):
            reshard_after_forward = i < len(blocks) - 1
            _fully_shard(block, mesh=mesh, reshard_after_forward=reshard_after_forward)
        return self

    @staticmethod
    def _zero_last_linear(module: nn.Module):
        """Zero the weight and bias of the *last* nn.Linear inside `module`."""
        last_lin = None
        for m in module.modules():
            if isinstance(m, nn.Linear):
                last_lin = m
        if last_lin is None:
            return
        with torch.no_grad():
            nn.init.zeros_(last_lin.weight)
            if last_lin.bias is not None:
                nn.init.zeros_(last_lin.bias)

    # ------------------------------------------------------------------ helpers

    def _to_wan_36ch(self, x_v_34: torch.Tensor) -> torch.Tensor:
        """Reshape Any4D's 34-ch per-view tensor into Wan 2.1 Fun InP's 36-ch input layout.

        Input layout (Any4D stream after pack_streams + assemble_network_inputs):
            x_v_34[:, 0:16]  = noisy rgb VAE latent (cs_in * noisy target at all frames,
                               because Any4DWanPipeline overrides rgb's input_mask to 0
                               and output_mask to 1 — see a4d_pipe_wan.denoise)
            x_v_34[:, 16:17] = input_mask value  (1 at given frames, 0 elsewhere; read-through)
            x_v_34[:, 17:18] = output_mask value (0 at given frames, 1 elsewhere; unused here,
                               kept for Any4D bookkeeping / back-compat with loss masks)
            x_v_34[:, 18:34] = Wan reference latent — VAE-encoded pixel video with real
                               pixels at every given frame and gray (0 in [-1, 1]) elsewhere;
                               built by a4d_vae_wan._build_wan_ref_latent from the per-view
                               input_mask. Generalizes across variable k in the contiguous
                               `[1]*k + [0]*(T-k)` conditioning pattern. Read-through:
                               assemble gives us the clean x0 values since the ref entry's
                               input_mask defaults to 1.

        Output layout (Wan Fun InP's pretrained patch_embedding input):
            out[:, 0:16]  = noisy_slot  (xt, actively denoised)
            out[:, 16:20] = mask_slot   (1 at given frames, broadcast 1→4 ch)
            out[:, 20:36] = ref_slot    (clean Wan reference latent)

        Pure tensor slicing + broadcast — no learned parameters.
        """
        noisy_slot = x_v_34[:, 0:16]
        imask      = x_v_34[:, 16:17]
        ref_slot   = x_v_34[:, 18:34]
        mask_slot  = imask.expand(-1, 4, -1, -1, -1)           # (B, 4, T, H, W)
        return torch.cat([noisy_slot, mask_slot, ref_slot], dim=1)

    # ------------------------------------------------------------------ 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]:
        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 ----------------
        ref_key = f'v{V_used[0]}'
        t_flat = timesteps[ref_key].flatten(1).mean(dim=1)  # (B,)
        t_emb = sinusoidal_embedding_1d(self.wan.freq_dim, t_flat).to(crossattn_emb.dtype)
        t_emb = self.wan.time_embedding(t_emb)  # (B, D)
        if cfg.has_adaln_stream:
            action = x_in_streams['adaln']  # (B, T, C)
            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))  # (B, 6, D)

        # ---------------- cross-attention context ----------------
        # Project text from text_dim -> D.
        context = self.wan.text_embedding(crossattn_emb)  # (B, Nc, D)
        img_context_len = 0

        # xattn stream: treat as extra tokens appended to text, before text projection.
        if cfg.has_xattn_stream:
            xattn_in = x_in_streams['xattn']                   # (B, Nx, Cx)
            xattn_text = self.xattn_inproj(xattn_in)           # (B, Nx, text_dim)
            extra_ctx = self.wan.text_embedding(xattn_text)    # (B, Nx, D)
            context = torch.cat([context, extra_ctx], dim=1)

        # ---------------- per-view patch embedding ----------------
        # All views share the same Wan pretrained patch_embedding (36 input channels).
        # Per-view differentiation comes from the view_embs added after patch_embed.
        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_18 = x_in_streams[f'v{v_idx}']              # (B, 34, T, H, W) after rgb{N}_ref entry
            x_v_36 = self._to_wan_36ch(x_v_18)              # (B, 36, T, H, W)
            tokens = self.x_embedder(x_v_36)                # (B, D, Tt, Hh, Ww)
            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 stream: extra self-attention tokens appended to the main sequence.
        sattn_slot = None
        if cfg.has_sattn_stream:
            sattn_in = x_in_streams['sattn']               # (B, Ns, Cs)
            sattn_emb = self.sattn_inproj(sattn_in)        # (B, Ns, D)
            sattn_slot = (sum(seq_lens), sattn_emb.shape[1])
            token_parts.append(sattn_emb)

        x_tokens = torch.cat(token_parts, dim=1)  # (B, Sum_L, D)

        # ---------------- RoPE frequencies for the extended sequence ----------------
        freqs = self._build_freqs(grid_shapes, sattn_slot, x_tokens.device)

        # ---------------- Wan DiT blocks ----------------
        x = x_tokens
        for block in self.wan.blocks:
            x = block(x, context, t_mod, freqs, img_context_len=img_context_len)

        # ---------------- split back & per-view head ----------------
        # All views share Wan's pretrained head (out_dim = 16 rgb latent channels).
        # The pipeline expects net_output to match the per-view input stream in channel
        # count (= cfg.num_video_in_channels[v_idx], typically 34 = rgb + 2 masks + 16 ref).
        # Wan's head only predicts the rgb slot (16 ch); we zero-pad the remaining
        # channels so downstream arithmetic (cs_skip*yt + cs_out*net_out) preserves the
        # mask/ref slots unchanged, which is correct since they are conditioning-only
        # inputs and don't participate in the loss (supervise_mask is 0 for those).
        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]
            y_v = self.final_existing(tokens_v, t_emb.to(tokens_v.dtype))  # (B, L, pt*ph*pw*16)
            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

    def _build_freqs(
        self,
        grid_shapes: List[Tuple[int, int, int]],
        sattn_slot: Optional[Tuple[int, int]],
        device: torch.device,
    ) -> torch.Tensor:
        """Construct per-token RoPE complex freqs for the concatenated view sequence + sattn pad.
        Output shape: (Sum_L, 1, head_dim / 2) complex — matches `rope_apply`'s expected format.
        """
        parts = []
        f_freqs, h_freqs, w_freqs = self.wan.freqs  # complex precomputed
        f_freqs, h_freqs, w_freqs = f_freqs.to(device), h_freqs.to(device), w_freqs.to(device)

        for (Tt, Hh, Ww) in grid_shapes:
            f_f = f_freqs[:Tt].view(Tt, 1, 1, -1).expand(Tt, Hh, Ww, -1)
            h_f = h_freqs[:Hh].view(1, Hh, 1, -1).expand(Tt, Hh, Ww, -1)
            w_f = w_freqs[:Ww].view(1, 1, Ww, -1).expand(Tt, Hh, Ww, -1)
            parts.append(torch.cat([f_f, h_f, w_f], dim=-1).reshape(Tt * Hh * Ww, 1, -1))

        if sattn_slot is not None:
            _, s_len = sattn_slot
            # Trivial RoPE (all freqs = identity / zero-angle complex) for sattn tokens.
            zero_complex = torch.complex(
                torch.ones(s_len, 1, parts[0].shape[-1], device=device),
                torch.zeros(s_len, 1, parts[0].shape[-1], device=device),
            )
            parts.append(zero_complex)

        return torch.cat(parts, dim=0)


# Re-export WanModel so downstream code (pipelines, checkpoint surgery) can import it from here
# rather than from diffsynth.
__all__ = ['Any4DWanDiT', 'WanModel', 'DiTBlock', 'Head']