wan/modules/causal_model.py

# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
# Modified by Applied Intuition, Inc. in 2026.
# SPDX-License-Identifier: Apache-2.0

import functools
from wan.modules.attention import attention
import math
from wan.modules.model import (
    WanRMSNorm,
    rope_apply,
    rope_apply_given_freqs,
    WanLayerNorm,
    WAN_CROSSATTENTION_CLASSES,
    rope_params,
    MLPProj,
    sinusoidal_embedding_1d
)
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
from diffusers.configuration_utils import ConfigMixin, register_to_config
from torch.nn.attention.flex_attention import BlockMask
from diffusers.models.modeling_utils import ModelMixin
import torch.nn as nn
import torch
import math
import torch.distributed as dist
import time
import copy
from einops import rearrange

import os as _os
try:
    import torch._dynamo as _dynamo
    _dynamo.config.recompile_limit = int(
        _os.environ.get("TORCHDYNAMO_RECOMPILE_LIMIT", "64"))
    _dynamo.config.accumulated_recompile_limit = int(
        _os.environ.get("TORCHDYNAMO_ACCUMULATED_RECOMPILE_LIMIT", "4096"))
    print(
        "[causal_model] torch._dynamo recompile_limit="
        f"{_dynamo.config.recompile_limit} accumulated="
        f"{_dynamo.config.accumulated_recompile_limit}")
except Exception as _e:
    print(f"[causal_model] could not set torch._dynamo limits: {_e}")

if _os.environ.get("DISABLE_FLEX_COMPILE") == "1":
    # Env-gated escape hatch: keep flex_attention in eager mode. Removes
    # torch.compile + Triton autotune non-determinism. Used to isolate whether
    # compile-layer variability explains cross-run output divergence.
    print("[causal_model] DISABLE_FLEX_COMPILE=1 — flex_attention stays eager")
else:
    _flex_compile_mode = _os.environ.get("FLEX_COMPILE_MODE", "max-autotune-no-cudagraphs")
    print(f"[causal_model] flex_attention torch.compile mode={_flex_compile_mode}")
    flex_attention = torch.compile(
        flex_attention, dynamic=False, mode=_flex_compile_mode)

if _os.environ.get("REMIND_DMD_CHUNK_FLASH", "0") == "1":
    print("[causal_model] DMD chunk-causal attention uses chunked FlashAttention")


_VALID_CC_ROPE_MODES = (
    "standard",          # Wan original t/h/w RoPE; no camera/ProPE modules
    "dual_prope",        # full dual-attention PRoPE, no QK phase
    "cc_basic",          # QK phase (all D/2 slots), no PRoPE output residual
    "cc_output",         # QK phase (all D/2 slots) + P·x_std + prope_proj
    "cc_dual_channel",   # QK phase (last cc_phase_slots slots only), no PRoPE
    "cc_dual_output",    # QK phase (last cc_phase_slots slots only) + P·x_std + prope_proj
    "prope_residual",    # single-attention, P·x_std + prope_proj, NO QK phase
    "cc_value",          # QK phase (all D/2 slots) + V-side P_inv residual via value_proj
    "cc_full",           # QK phase (all D/2 slots) + V-side (value_proj) + O-side (prope_proj) residuals
)


def _camera_pose_features(viewmats: torch.Tensor,
                          Ks: torch.Tensor = None) -> torch.Tensor:
    """Build a compact per-frame pose descriptor from viewmats (+ Ks).

    Args:
        viewmats: [B, F, 4, 4] c2w, already translation-normalized upstream
                   (see CausalWanModel.forward).
        Ks:       [B, F, 3, 3] intrinsics, already focal-normalized upstream.
                   Provides focal/zoom signal (None for datasets without Ks).

    Returns:
        pose_c: [B, F, P] float32 descriptor. P depends on whether Ks is given:
          translation (3) + R flattened (9) + log-focal (2 if Ks else 0)  →  14 or 12.
    """
    assert viewmats.dim() == 4 and viewmats.shape[-2:] == (4, 4)
    B_, F_ = viewmats.shape[:2]
    t = viewmats[..., :3, 3].float()                           # [B, F, 3]
    R = viewmats[..., :3, :3].float().reshape(B_, F_, 9)       # [B, F, 9]
    feats = [t, R]
    if Ks is not None:
        # log-focal: Ks already ~O(1) after normalization, clamp for safety.
        fx = Ks[..., 0, 0].float().clamp(min=1e-3)
        fy = Ks[..., 1, 1].float().clamp(min=1e-3)
        lf = torch.stack([fx.log(), fy.log()], dim=-1)          # [B, F, 2]
        feats.append(lf)
    return torch.cat(feats, dim=-1)                             # [B, F, P]


class CameraPhaseMLP(nn.Module):
    """Zero-init linear layer: pose_c [B, F, P] → phase phi [B, F, D_half].

    Returns a complex phasor tensor `exp(j * phi)` ready to multiply against
    the per-token `freqs_x` of shape [F*H*W, 1, D_half]. When mask_range is
    set, phase is produced only for slots in [lo, hi); the others output 0
    phase → `exp(j*0) = 1` → those freq slots are untouched.

    At step 0 (weight=0) phase is identically 0, `exp(0)=1`, so attention
    output is bit-exact with a model that had no CameraPhaseMLP.
    """
    def __init__(self,
                 pose_dim: int,
                 head_dim_half: int,
                 mask_range: tuple = None):
        super().__init__()
        self.pose_dim = pose_dim
        self.head_dim_half = head_dim_half
        self.mask_range = mask_range  # None = all slots, else (lo, hi)
        # zero-init so phase_delta=0 ⇒ phasor=1+0j ⇒ freqs unchanged
        out_dim = head_dim_half if mask_range is None else (mask_range[1] - mask_range[0])
        self.proj = nn.Linear(pose_dim, out_dim, bias=False)
        nn.init.zeros_(self.proj.weight)

    def forward(self, pose_c: torch.Tensor) -> torch.Tensor:
        # pose_c: [B, F, P]  →  phi: [B, F, D_half] (float32 for stability of
        # cos/sin; the resulting complex phasor is multiplied into freqs which
        # are also float32/complex64 in rope_apply_given_freqs).
        B_, F_, _ = pose_c.shape
        w = self.proj.weight                                    # dtype-owner
        phi_active = torch.nn.functional.linear(
            pose_c.to(w.dtype), w).float()                      # [B, F, out_dim]
        if self.mask_range is None:
            phi = phi_active
        else:
            phi = pose_c.new_zeros(B_, F_, self.head_dim_half, dtype=torch.float32)
            lo, hi = self.mask_range
            phi[..., lo:hi] = phi_active
        return torch.polar(torch.ones_like(phi), phi)           # complex [B, F, D_half]


def _attention_with_weights(q, k, v):
    """Manual scaled_dot_product_attention that also returns the softmax weights.

    Used only when KV-cache importance probing is enabled. Slower than the
    fused kernel (no flash path) but needed because F.scaled_dot_product_attention
    does not expose attention weights.

    Shapes: q,k,v in [B, L, H, D]; returns (out [B, L_q, H, D], weights [B, H, L_q, L_kv]).
    Computes in float32 for numerical stability then casts back.
    """
    q_ = q.permute(0, 2, 1, 3)                             # [B, H, L_q, D]
    k_ = k.permute(0, 2, 1, 3)
    v_ = v.permute(0, 2, 1, 3)
    scale = q_.shape[-1] ** -0.5
    scores = torch.einsum('bhld,bhmd->bhlm', q_.float(), k_.float()) * scale
    attn = torch.softmax(scores, dim=-1)                   # [B, H, L_q, L_kv], float32
    out = torch.einsum('bhlm,bhmd->bhld', attn, v_.float()).to(v_.dtype)
    return out.permute(0, 2, 1, 3), attn


def _attention_with_frame_importance(
        q, k, v, *, hist_len: int, tokens_per_frame: int,
        query_chunk_size: int = 256):
    """Exact attention plus per-history-frame weights with bounded memory."""
    q_ = q.permute(0, 2, 1, 3)
    k_ = k.permute(0, 2, 1, 3)
    v_ = v.permute(0, 2, 1, 3)
    batch, heads, query_len, head_dim = q_.shape
    num_hist = (
        hist_len // tokens_per_frame
        if tokens_per_frame > 0 and hist_len % tokens_per_frame == 0
        else 0
    )
    importance_sum = torch.zeros(
        heads, num_hist, device=q.device, dtype=torch.float32)
    outputs = []
    scale = head_dim ** -0.5
    chunk = max(1, int(query_chunk_size))
    for start in range(0, query_len, chunk):
        q_part = q_[:, :, start:start + chunk]
        scores = torch.einsum(
            'bhld,bhmd->bhlm', q_part.float(), k_.float()) * scale
        attn = torch.softmax(scores, dim=-1)
        outputs.append(torch.einsum(
            'bhlm,bhmd->bhld', attn, v_.float()).to(v_.dtype))
        if num_hist:
            hist = attn[..., :hist_len].reshape(
                batch, heads, q_part.shape[2], num_hist, tokens_per_frame)
            importance_sum += hist.sum(dim=-1).sum(dim=(0, 2))
        del scores, attn
    output = torch.cat(outputs, dim=2).permute(0, 2, 1, 3)
    importance = importance_sum / max(1, batch * query_len)
    return output, importance


class CausalWanSelfAttention(nn.Module):

    def __init__(self,
                 dim,
                 num_heads,
                 qk_norm=True,
                 eps=1e-6,
                 cc_rope_mode: str = "dual_prope",
                 cc_pose_dim: int = 14,
                 cc_phase_slots: int = 16):
        assert dim % num_heads == 0
        super().__init__()
        assert cc_rope_mode in _VALID_CC_ROPE_MODES, (
            f"cc_rope_mode={cc_rope_mode!r} must be one of {_VALID_CC_ROPE_MODES}")
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.qk_norm = qk_norm
        self.eps = eps
        self.cc_rope_mode = cc_rope_mode
        self.fused_projections = False

        # layers (standard self-attention)
        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 = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()

        # ProPE dual-path output projection (zero-init → initial no-op).
        # See HY-WorldPlay's img_attn_prope_proj.
        # Only instantiated where it's actually read in forward:
        #   - standard        : no ProPE / camera residual modules at all
        #   - dual_prope      : gate for PRoPE 2nd-attention output (existing path)
        #   - cc_output       : gate for P·x_std (single-attention, camera-projected)
        #   - cc_dual_output  : same as cc_output but with masked-range QK phase
        #   - prope_residual  : gate for P·x_std WITHOUT any QK phase (isolated
        #                       residual-only branch for clean ablation vs cc_output)
        #   - cc_full         : gate for P·x_std (paired with V-side value_proj)
        # `cc_basic` / `cc_dual_channel` / `cc_value` do NOT use this module, so
        # we leave it as None to save ~30 × (dim² + dim) params and the
        # corresponding AdamW optimizer state (~1.6 GB on the 1.3B model). Any
        # stale prope_proj.* keys from pretrained checkpoints become harmless
        # "unexpected_keys" on load (strict=False).
        if cc_rope_mode in ("dual_prope", "cc_output", "cc_dual_output", "prope_residual", "cc_full"):
            self.prope_proj = nn.Linear(dim, dim)
            nn.init.zeros_(self.prope_proj.weight)
            if self.prope_proj.bias is not None:
                nn.init.zeros_(self.prope_proj.bias)
        else:
            self.prope_proj = None

        # cc_value / cc_full: zero-init Linear that mixes (P_inv · V) back into
        # V before attention, providing a "value-side" camera-geometry residual
        # that mirrors cc_output's "x_std-side" residual. Step-0 bit-exact
        # because the residual contributes 0 when value_proj.weight = bias = 0.
        # cc_full uses BOTH value_proj (V-side) and prope_proj (O-side) to
        # close the dual_prope geometric loop (V → world → attn → P → query
        # frame), but with zero-init residuals instead of replacements.
        if cc_rope_mode in ("cc_value", "cc_full"):
            self.value_proj = nn.Linear(dim, dim)
            nn.init.zeros_(self.value_proj.weight)
            if self.value_proj.bias is not None:
                nn.init.zeros_(self.value_proj.bias)
        else:
            self.value_proj = None

        # Camera-conditioned RoPE (CC-RoPE): zero-init per-freq phase delta.
        head_dim_half = self.head_dim // 2
        self.camera_phase_mlp = None
        if cc_rope_mode in ("cc_basic", "cc_output", "cc_value", "cc_full"):
            # full-band: all freq slots receive camera phase
            self.camera_phase_mlp = CameraPhaseMLP(
                pose_dim=cc_pose_dim,
                head_dim_half=head_dim_half,
                mask_range=None,
            )
        elif cc_rope_mode in ("cc_dual_channel", "cc_dual_output"):
            # M-RoPE style: camera phase only on the LAST `cc_phase_slots` freqs.
            # cc_dual_channel: masked phase, no PRoPE output residual.
            # cc_dual_output : masked phase + PRoPE output residual (P·x_std).
            lo = max(0, head_dim_half - cc_phase_slots)
            self.camera_phase_mlp = CameraPhaseMLP(
                pose_dim=cc_pose_dim,
                head_dim_half=head_dim_half,
                mask_range=(lo, head_dim_half),
            )

        # ── KV cache importance probing (off by default; no training impact) ──
        # When enabled, the cache-read branch of forward() swaps the fused
        # attention kernel for a manual path that exposes softmax weights, then
        # accumulates per-history-frame mean weight across denoising steps.
        self._log_attn: bool = False
        self._log_tpf: int = 0          # tokens per frame (set by pipeline)
        self._log_chunk_size: int = 0   # chunks-of-frames granularity (chunk_mask path)
        self._attn_sum = None           # [H, num_hist_frames]      — kv_cache path
        self._attn_count: int = 0
        self._attn_chunk_matrix = None  # [H, T_frames_Q, num_chunks_K] — chunk_mask path

    @staticmethod
    def _chunk_flash_attention(
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        tokens_per_chunk: int,
    ) -> torch.Tensor:
        """Exact chunk-causal attention using dense FlashAttention calls."""
        total_tokens = q.shape[1]
        if (
            k.shape[1] != total_tokens
            or v.shape[1] != total_tokens
            or tokens_per_chunk <= 0
            or total_tokens % tokens_per_chunk != 0
        ):
            raise ValueError(
                "chunk flash expects equal, unpadded q/k/v sequences divisible "
                f"by tokens_per_chunk: q={tuple(q.shape)} k={tuple(k.shape)} "
                f"v={tuple(v.shape)} tokens_per_chunk={tokens_per_chunk}"
            )

        outputs = []
        for start in range(0, total_tokens, tokens_per_chunk):
            end = start + tokens_per_chunk
            outputs.append(attention(q[:, start:end], k[:, :end], v[:, :end]))
        return torch.cat(outputs, dim=1)

    def _clean_history_chunk_flash_attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        half_tokens: int,
        tokens_per_chunk: int,
    ) -> torch.Tensor:
        """Structured clean-history attention via FlashAttention.

        The dual-stream clean-history mask is regular at chunk granularity:
        clean chunk i attends to clean chunks <= i, and noisy chunk i attends
        to clean chunks < i plus its own noisy chunk. Splitting those blocks
        keeps the exact attention support while avoiding generic FlexAttention
        backward, which is very slow for this 65k-token mask on A100.
        """
        total_tokens = half_tokens * 2
        if (
            q.shape[1] != total_tokens
            or k.shape[1] != total_tokens
            or v.shape[1] != total_tokens
            or half_tokens <= 0
            or tokens_per_chunk <= 0
            or half_tokens % tokens_per_chunk != 0
        ):
            raise ValueError(
                "clean-history chunk flash expects unpadded [clean,noisy] "
                f"tokens: q={tuple(q.shape)} k={tuple(k.shape)} "
                f"half_tokens={half_tokens} tokens_per_chunk={tokens_per_chunk}"
            )

        clean_out = []
        noisy_out = []
        for start in range(0, half_tokens, tokens_per_chunk):
            end = start + tokens_per_chunk
            clean_out.append(attention(q[:, start:end], k[:, :end], v[:, :end]))

            noisy_start = half_tokens + start
            noisy_end = half_tokens + end
            if start > 0:
                k_noisy = torch.cat(
                    [k[:, :start], k[:, noisy_start:noisy_end]], dim=1)
                v_noisy = torch.cat(
                    [v[:, :start], v[:, noisy_start:noisy_end]], dim=1)
            else:
                k_noisy = k[:, noisy_start:noisy_end]
                v_noisy = v[:, noisy_start:noisy_end]
            noisy_out.append(
                attention(q[:, noisy_start:noisy_end], k_noisy, v_noisy))

        return torch.cat(clean_out + noisy_out, dim=1)

    def enable_attn_log(self, tokens_per_frame: int, enable: bool = True,
                         chunk_size: int = 0):
        """Toggle attention logging. Resets accumulators.

        `tokens_per_frame`: size of one video frame in token units (required
            for both the kv_cache-streaming and chunk_mask logging paths).
        `chunk_size`: only used by the chunk_mask path (single-step training-
            style forward).  Set to the same chunk_size the forward was called
            with so the K-side can be aggregated to per-chunk columns.
        """
        self._log_attn = bool(enable)
        self._log_tpf = int(tokens_per_frame)
        self._log_chunk_size = int(chunk_size)
        self._attn_sum = None
        self._attn_count = 0
        self._attn_chunk_matrix = None

    def get_attn_log(self):
        """Return mean per-history-frame attention [H, num_hist] or None.
        (kv_cache-streaming inference path only.)
        """
        if self._attn_sum is None or self._attn_count == 0:
            return None
        return self._attn_sum / self._attn_count

    def get_attn_chunk_matrix(self):
        """Return per-Q-frame → per-K-chunk attention mass for the chunk_mask
        (single-step full-sequence) forward path.  Shape [H, T_frames_Q,
        num_chunks_K] on CPU fp32, or None if never captured.
        """
        return self._attn_chunk_matrix

    def _manual_chunk_mask_attention(self, q, k, v, tpf: int, cs: int):
        """Chunk-causal self-attention with softmax-weight bookkeeping.

        Used ONLY by the `_log_attn` probe path in the chunk_mask branch of
        `forward()`.  Replaces `flex_attention` with an explicit per-Q-frame
        loop so that we can aggregate softmax mass into a
        `[H, T_frames_Q, num_chunks_K]` matrix (cached on
        `self._attn_chunk_matrix`) and visualise it as a KV-importance heatmap.

        Memory is kept manageable by looping over Q frames: peak intermediate
        tensor is `[B, H, tpf, L_kv]`, not the full `[B, H, L, L]`.

        Args:
          q, k, v : [B, L, H, D] (RoPE already applied to q, k by caller).
          tpf     : tokens per latent frame.
          cs      : chunk_size (frames per chunk).  K is aggregated at this
                     granularity; Q stays per-frame for finer rows.
        Returns:
          x_std  : [B, L, H, D] attention output (same shape/semantics as
                    flex_attention with the chunk-causal block mask).
        """
        import math
        B, L, H, D = q.shape
        assert L % (tpf * cs) == 0, \
            f"L={L} not divisible by tpf*cs={tpf*cs}"
        T_frames = L // tpf
        num_chunks = T_frames // cs
        scale = 1.0 / math.sqrt(D)

        q_h = q.permute(0, 2, 1, 3).contiguous()      # [B, H, L, D]
        k_h = k.permute(0, 2, 1, 3).contiguous()
        v_h = v.permute(0, 2, 1, 3).contiguous()

        attn_mat = torch.zeros(H, T_frames, num_chunks,
                                device=q_h.device, dtype=torch.float32)
        out_per_frame = []
        for qf in range(T_frames):
            qs, qe = qf * tpf, (qf + 1) * tpf
            qc = qf // cs
            ke = (qc + 1) * cs * tpf                  # causal K extent
            q_i = q_h[:, :, qs:qe, :]                 # [B, H, tpf, D]
            k_i = k_h[:, :, :ke, :]                   # [B, H, ke, D]
            v_i = v_h[:, :, :ke, :]

            # bf16/fp16 matmul is plenty for softmax; accumulate in fp32.
            logits = torch.matmul(q_i, k_i.transpose(-2, -1)) * scale
            weights = logits.softmax(dim=-1)          # [B, H, tpf, ke]

            out_i = torch.matmul(weights.to(v_i.dtype), v_i)  # [B, H, tpf, D]
            out_per_frame.append(out_i)

            # K-chunk aggregation: group ke=visible*cs*tpf into (visible, cs*tpf) and sum.
            visible = qc + 1
            w_ck = weights.float().view(
                B, H, tpf, visible, cs * tpf,
            ).sum(dim=-1).mean(dim=(0, 2))             # [H, visible]
            attn_mat[:, qf, :visible] = w_ck

        x_std_h = torch.cat(out_per_frame, dim=2)     # [B, H, L, D]
        x_std = x_std_h.permute(0, 2, 1, 3)           # [B, L, H, D]
        self._attn_chunk_matrix = attn_mat.detach().cpu()
        return x_std

    def _accumulate_attn(self, attn_weights: torch.Tensor, hist_len: int):
        """attn_weights: [B, H, L_q, L_kv]. hist_len tokens at the front of L_kv
        belong to the history cache; the tail is current-chunk self-attention."""
        if hist_len <= 0 or self._log_tpf <= 0:
            return
        num_hist = hist_len // self._log_tpf
        if num_hist == 0 or hist_len % self._log_tpf != 0:
            return
        B, H, L_q, _ = attn_weights.shape
        hist = attn_weights[:, :, :, :hist_len].reshape(
            B, H, L_q, num_hist, self._log_tpf).sum(dim=-1)   # [B, H, L_q, num_hist]
        pf = hist.mean(dim=(0, 2)).detach()                    # [H, num_hist]
        self._accumulate_attn_summary(pf)

    def _accumulate_attn_summary(self, pf: torch.Tensor):
        if pf is None or pf.numel() == 0:
            return
        pf = pf.detach()
        if self._attn_sum is None or self._attn_sum.shape != pf.shape:
            self._attn_sum = pf.clone()
            self._attn_count = 1
        else:
            self._attn_sum = self._attn_sum + pf
            self._attn_count += 1

    def forward(
        self,
        x,
        seq_lens,
        freqs,
        kv_cache=None,
        kv_size=(0,0),
        viewmats=None,
        Ks=None,
        pose_c=None,
        tokens_per_frame: int = None,
        chunk_mask=None,
        prope_temporal_dim: int = 0,
        prope_freqs_offset: int = 0,
        prope_freqs_positions=None,
        sequence_parallel: bool = False,
        sequence_parallel_clean_history: bool = False,
        chunk_flash: bool = False,
        chunk_flash_tokens_per_chunk: int = 0,
        clean_history_chunk_flash: bool = False,
        clean_history_half_tokens: int = 0,
        clean_history_tokens_per_chunk: int = 0,
    ):
        r"""
        Args:
            x(Tensor): Shape [B, L, C]
            seq_lens(Tensor): Shape [B]
            freqs(Tensor): RoPE freqs [L, 1, C/num_heads] (already per-token)
            kv_cache: dict with keys k, v (and optionally k_prope, v_prope).
                     Used for inference (chunk-by-chunk streaming).
            kv_size: (offset, length). length<0 means context pass (write),
                     length=0 means no cache, length>0 means read [offset:offset+length]
            viewmats: [B, num_frames_in_x, 4, 4] — camera-to-world per frame of x.
                      If None and cc_rope_mode='dual_prope', ProPE path is skipped.
            Ks: [B, num_frames_in_x, 3, 3] intrinsics.
            pose_c: [B, F, P] per-frame pose descriptor (CC-RoPE modes only).
                    Produced by _camera_pose_features in CausalWanModel.forward.
            tokens_per_frame: int, needed by CC-RoPE to broadcast per-frame phase
                               to per-token freqs.
            chunk_mask: flex_attention BlockMask for full-sequence training mode.
                        When provided, skips KV cache entirely; uses flex_attention
                        with sparse block-causal mask (chunk i sees chunks 0..i).
                        Built by CausalWanModel._get_chunk_block_mask (cached).
        """
        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim

        def qkv_fn(x):
            q = self.norm_q(self.q(x)).view(b, s, n, d)
            k = self.norm_k(self.k(x)).view(b, s, n, d)
            v = self.v(x).view(b, s, n, d)
            return q, k, v

        q, k, v = qkv_fn(x)

        # ─── Build effective RoPE freqs (shared or camera-modulated) ───
        # cc_basic / cc_dual_channel / cc_output: multiply per-frame phasor
        # into the base freqs so a SINGLE attention captures camera-relative
        # position. Zero-init MLP ⇒ phasor=1 ⇒ freqs_eff ≡ freqs, i.e. step-0
        # output is bit-exact with "pure standard RoPE".
        freqs_eff = freqs
        if self.camera_phase_mlp is not None and pose_c is not None:
            assert tokens_per_frame is not None, \
                "tokens_per_frame is required when cc_rope_mode uses camera phase"
            phasor = self.camera_phase_mlp(pose_c)     # [B, F, D_half] complex
            B_, F_, Dh = phasor.shape
            # Expand per-frame phasor to per-token:
            # [B, F, 1, D_half] → [B, F, tpf, D_half] → [B, F*tpf, 1, D_half]
            phasor = phasor.unsqueeze(2).expand(B_, F_, tokens_per_frame, Dh)
            phasor = phasor.reshape(B_, F_ * tokens_per_frame, 1, Dh)
            # freqs: [L, 1, D_half] complex; broadcast to [1, L, 1, D_half]
            base = freqs.unsqueeze(0) if freqs.dim() == 3 else freqs
            freqs_eff = base * phasor                   # [B, L, 1, D_half]

        # ─── Standard RoPE path (always runs; with CC-RoPE freqs if enabled) ──
        roped_query = rope_apply_given_freqs(q, freqs_eff).type_as(v)
        roped_key   = rope_apply_given_freqs(k, freqs_eff).type_as(v)

        # ─── ProPE path (optional) ────────────────────────────────────
        # dual_prope: full dual-attention PRoPE (legacy).
        # cc_output / cc_dual_output / prope_residual: single-attention +
        #             post-hoc P-projection of x_std (apply_fn_o; skip Q/K/V
        #             transforms). They differ only in whether the QK phase
        #             MLP is active (full / masked / disabled).
        # cc_value  : single-attention + pre-hoc P_inv-projection of V mixed
        #             via zero-init Linear residual (V → V + value_proj(P_inv·V)).
        # cc_full   : single-attention + BOTH cc_value (V-side P_inv residual)
        #             AND cc_output (O-side P residual). Closes the dual_prope
        #             geometric loop (V→world→attn→P→query frame) via two
        #             zero-init Linear residuals.
        # cc_basic / cc_dual_channel: no PRoPE ops at all.
        prope_apply_fn_o = None
        q_prope = k_prope = v_prope = None
        if self.cc_rope_mode == "dual_prope":
            use_prope = (viewmats is not None)
            if use_prope:
                from prope.camera_rope import prope_qkv
                q_prope, k_prope, v_prope, prope_apply_fn_o = prope_qkv(
                    q.permute(0, 2, 1, 3),
                    k.permute(0, 2, 1, 3),
                    v.permute(0, 2, 1, 3),
                    viewmats=viewmats,
                    Ks=Ks,
                    temporal_dim=prope_temporal_dim,
                    freqs_offset=prope_freqs_offset,
                    freqs_positions=prope_freqs_positions,
                )
                q_prope = q_prope.permute(0, 2, 1, 3).type_as(v)
                k_prope = k_prope.permute(0, 2, 1, 3).type_as(v)
                v_prope = v_prope.permute(0, 2, 1, 3).type_as(v)
            use_prope_dual = use_prope
        elif self.cc_rope_mode in ("cc_output", "cc_dual_output", "prope_residual"):
            # Build only the output-side P transform (no 2nd attention).
            # head_dim must be divisible by 4 for the 4×4 block-diagonal matmul.
            # These three modes share the identical output-side PRoPE residual
            # (x_std → P → prope_proj). They differ only in the QK phase MLP:
            #   cc_output      : full-band phase (all D/2 freq slots)
            #   cc_dual_output : masked phase (last cc_phase_slots only)
            #   prope_residual : NO phase at all (camera_phase_mlp is None)
            # `prope_residual` serves as the clean "residual-only" ablation
            # cell vs cc_basic (phase-only) and cc_output (phase + residual).
            use_prope_dual = False
            if viewmats is not None and self.head_dim % 4 == 0:
                from prope.camera_rope import _prepare_apply_fns_all_dim
                _, _, prope_apply_fn_o = _prepare_apply_fns_all_dim(
                    head_dim=self.head_dim,
                    viewmats=viewmats,
                    Ks=Ks,
                    patches_x=None, patches_y=None,
                    image_width=None, image_height=None,
                )
        elif self.cc_rope_mode in ("cc_value", "cc_full"):
            # Pre-attention V-side residual:
            #   V ← V + value_proj(P_inv · V)
            # `value_proj` is zero-init ⇒ residual=0 at step 0 ⇒ bit-exact.
            # cc_value: V-side residual only (no output-side prope_proj).
            # cc_full : V-side residual AND output-side P-residual (handled
            #           by also setting prope_apply_fn_o below).
            use_prope_dual = False
            if viewmats is not None and self.head_dim % 4 == 0 and self.value_proj is not None:
                from prope.camera_rope import _prepare_apply_fns_all_dim
                _, prope_apply_fn_kv, _apply_fn_o_local = _prepare_apply_fns_all_dim(
                    head_dim=self.head_dim,
                    viewmats=viewmats,
                    Ks=Ks,
                    patches_x=None, patches_y=None,
                    image_width=None, image_height=None,
                )
                # v: [B, L, H, d] → [B, H, L, d] for tiled block-diag matmul.
                v_hnd = v.permute(0, 2, 1, 3).contiguous()        # [B, H, L, d]
                v_p_hnd = prope_apply_fn_kv(v_hnd)                # [B, H, L, d]
                v_p_flat = v_p_hnd.permute(0, 2, 1, 3).flatten(2) # [B, L, dim]
                v_residual = self.value_proj(v_p_flat)            # [B, L, dim]
                v_aug_flat = v.flatten(2) + v_residual            # [B, L, dim]
                v = v_aug_flat.view(b, s, n, d).type_as(v)        # [B, L, H, d]
                # cc_full: also feed apply_fn_o for the output-side residual.
                if self.cc_rope_mode == "cc_full":
                    prope_apply_fn_o = _apply_fn_o_local
        else:
            use_prope_dual = False

        # Back-compat alias for the branches below (training + inference).
        use_prope = use_prope_dual

        # ─── Training mode: full-sequence with chunk-causal mask ────
        # chunk_mask is a flex_attention BlockMask (sparse chunk-causal),
        # built once per sequence shape by CausalWanModel._get_chunk_block_mask.
        if chunk_mask is not None:
            def _flex(q_, k_, v_, block_mask):
                if sequence_parallel:
                    from wan.distributed.ulysses import distributed_flex_attention
                    return distributed_flex_attention(
                        q_, k_, v_, block_mask,
                        is_clean_history=sequence_parallel_clean_history,
                        clean_history_half_tokens=int(clean_history_half_tokens),
                        clean_history_tokens_per_chunk=int(
                            clean_history_tokens_per_chunk),
                        clean_history_chunk_flash=(
                            bool(sequence_parallel_clean_history)
                            and _os.environ.get(
                                "REMIND_TF_CHUNK_FLASH_SP", "1") != "0"
                        ))
                if clean_history_chunk_flash:
                    return self._clean_history_chunk_flash_attention(
                        q_, k_, v_,
                        int(clean_history_half_tokens),
                        int(clean_history_tokens_per_chunk),
                    )
                if chunk_flash:
                    return self._chunk_flash_attention(
                        q_, k_, v_, int(chunk_flash_tokens_per_chunk)
                    )
                # q,k,v: [B, L, H, d] → [B, H, L, d]
                q_h = q_.permute(0, 2, 1, 3).contiguous()
                k_h = k_.permute(0, 2, 1, 3).contiguous()
                v_h = v_.permute(0, 2, 1, 3).contiguous()
                q_len = q_h.shape[2]
                kv_len = k_h.shape[2]
                target_q_len = int(block_mask.shape[-2])
                target_kv_len = int(block_mask.shape[-1])
                q_pad = target_q_len - q_len
                kv_pad = target_kv_len - kv_len
                if q_pad < 0 or kv_pad < 0:
                    raise ValueError(
                        f"BlockMask shape {block_mask.shape} is smaller than "
                        f"q/k lengths {(q_len, kv_len)}")
                if q_pad > 0:
                    q_h = torch.cat([
                        q_h,
                        q_h.new_zeros(q_h.shape[0], q_h.shape[1], q_pad, q_h.shape[3]),
                    ], dim=2)
                if kv_pad > 0:
                    k_h = torch.cat([
                        k_h,
                        k_h.new_zeros(k_h.shape[0], k_h.shape[1], kv_pad, k_h.shape[3]),
                    ], dim=2)
                    v_h = torch.cat([
                        v_h,
                        v_h.new_zeros(v_h.shape[0], v_h.shape[1], kv_pad, v_h.shape[3]),
                    ], dim=2)
                o_h = flex_attention(q_h, k_h, v_h, block_mask=block_mask)
                if q_pad > 0:
                    o_h = o_h[:, :, :q_len]
                return o_h.permute(0, 2, 1, 3)  # back to [B, L, H, d]

            # Attention-importance probe (validation only — skips the fused
            # flex kernel and runs manual SDPA per Q frame so we can aggregate
            # softmax mass by K-chunk).  Falls back to flex if the required
            # metadata isn't set.
            if (self._log_attn and self._log_tpf > 0
                    and self._log_chunk_size > 0):
                x_std = self._manual_chunk_mask_attention(
                    roped_query, roped_key, v,
                    tpf=self._log_tpf, cs=self._log_chunk_size,
                )
                # Skip ProPE during probe (not needed for the heatmap; avoids
                # another expensive manual attention).  Falls back to zero for
                # the residual output so outputs in _combine_out aren't NaN.
                x_p = None if not use_prope else _flex(
                    q_prope, k_prope, v_prope, chunk_mask)
                return self._combine_out(x_std, x_p, use_prope, prope_apply_fn_o)

            x_std = _flex(roped_query, roped_key, v, chunk_mask)
            if use_prope:
                x_p = _flex(q_prope, k_prope, v_prope, chunk_mask)
            else:
                x_p = None
            return self._combine_out(x_std, x_p, use_prope, prope_apply_fn_o)

        # ─── Inference mode: KV cache streaming ─────────────────────
        if kv_cache is None:
            raise RuntimeError(
                "CausalWanSelfAttention: neither chunk_mask (training) nor "
                "kv_cache (inference) provided. If this is a training call, "
                "make sure chunk_size is reaching CausalWanModel.forward "
                "(beware of DDP wrapping hiding the forward signature)."
            )

        if kv_size[1] < 0:
            # Context pass: write to cache
            len_x = roped_query.shape[1]
            kv_cache["k"][:, kv_size[0]:kv_size[0]+len_x] = roped_key
            kv_cache["v"][:, kv_size[0]:kv_size[0]+len_x] = v
            if use_prope:
                # Store already-ProPE-transformed K/V (in world space via P_inv)
                kv_cache["k_prope"][:, kv_size[0]:kv_size[0]+len_x] = k_prope
                kv_cache["v_prope"][:, kv_size[0]:kv_size[0]+len_x] = v_prope
            # Attention on just current tokens (for residual)
            x_std = attention(roped_query, roped_key, v)
            if use_prope:
                x_p = attention(q_prope, k_prope, v_prope)
        else:
            if kv_size[1] == 0:
                # No cache, attend to current only
                x_std = attention(roped_query, roped_key, v)
                if use_prope:
                    x_p = attention(q_prope, k_prope, v_prope)
            else:
                # Read cache + current, run attention
                cache_k = kv_cache["k"][:, kv_size[0]:kv_size[0]+kv_size[1]]
                cache_v = kv_cache["v"][:, kv_size[0]:kv_size[0]+kv_size[1]]
                K_full = torch.cat([cache_k, roped_key], dim=1)
                V_full = torch.cat([cache_v, v], dim=1)
                # std-RoPE attention stays fused (positional-only; not the diagnostic
                # signal). When ProPE is inactive we fall back to logging std
                # attention so the heatmap isn't empty, but the preferred signal
                # lives on the ProPE branch below.
                if self._log_attn and not use_prope:
                    x_std, _attn_summary = _attention_with_frame_importance(
                        roped_query, K_full, V_full,
                        hist_len=kv_size[1], tokens_per_frame=self._log_tpf)
                    self._accumulate_attn_summary(_attn_summary)
                    del _attn_summary
                else:
                    x_std = attention(roped_query, K_full, V_full)
                if use_prope:
                    cache_kp = kv_cache["k_prope"][:, kv_size[0]:kv_size[0]+kv_size[1]]
                    cache_vp = kv_cache["v_prope"][:, kv_size[0]:kv_size[0]+kv_size[1]]
                    Kp_full = torch.cat([cache_kp, k_prope], dim=1)
                    Vp_full = torch.cat([cache_vp, v_prope], dim=1)
                    # ProPE Q/K are CAMERA-GEOMETRY-TRANSFORMED (Q' = P_iᵀ·Q,
                    # K' = P_j⁻¹·K). Their softmax weights encode how much the
                    # current-chunk camera geometry aligns with each history
                    # frame's camera — exactly the "which frame does the model
                    # rely on for memory" signal we want to log.
                    if self._log_attn:
                        x_p, _attn_summary = _attention_with_frame_importance(
                            q_prope, Kp_full, Vp_full,
                            hist_len=kv_size[1], tokens_per_frame=self._log_tpf)
                        self._accumulate_attn_summary(_attn_summary)
                        del _attn_summary
                    else:
                        x_p = attention(q_prope, Kp_full, Vp_full)

        if not use_prope:
            x_p = None
        return self._combine_out(x_std, x_p, use_prope, prope_apply_fn_o)

    def _combine_out(self, x_std, x_p, use_prope_dual, prope_apply_fn_o):
        """Fold std-RoPE (+ optional PRoPE) outputs into final [B, L, C].

        Three code paths:
          * dual_prope (use_prope_dual=True):
              `out = o(x_std) + prope_proj(P @ x_p_from_2nd_attn)`
          * cc_output / cc_dual_output / prope_residual / cc_full
            (prope_apply_fn_o is not None, use_prope_dual=False):
              `out = o(x_std) + prope_proj(P @ x_std)`  — camera geometry on x_std
          * cc_basic / cc_dual_channel / cc_value / no-camera: `out = o(x_std)`
            (cc_value's V-side residual is folded into x_std upstream;
             cc_full also adds the O-side P·x_std residual via prope_apply_fn_o.)

        All paths are bit-exact at step 0 when prope_proj / value_proj are
        zero-init (which WanDiffusionWrapper enforces for cc_* modes after
        checkpoint load).
        """
        x_std_flat = x_std.flatten(2)                         # [B, L, dim]
        out = self.o(x_std_flat)

        if use_prope_dual:
            # Existing dual-PRoPE: 2nd-attention output → P → prope_proj
            x_p_hnd = x_p.permute(0, 2, 1, 3)                 # [B, heads, L, d]
            x_p_hnd = prope_apply_fn_o(x_p_hnd)
            x_p_flat = x_p_hnd.permute(0, 2, 1, 3).flatten(2) # [B, L, dim]
            out = out + self.prope_proj(x_p_flat)
        elif prope_apply_fn_o is not None:
            # cc_output / cc_dual_output / prope_residual / cc_full: feed
            # P(x_std) through prope_proj (no 2nd attention). Identical
            # output-side math across these; they differ only in how (or
            # whether) the QK phase MLP writes into freqs_eff upstream and
            # whether V was pre-augmented (cc_full only).
            x_std_hnd = x_std.permute(0, 2, 1, 3)             # [B, heads, L, d]
            x_std_hnd = prope_apply_fn_o(x_std_hnd)
            x_p_flat = x_std_hnd.permute(0, 2, 1, 3).flatten(2)
            out = out + self.prope_proj(x_p_flat)

        return out


class CausalWanAttentionBlock(nn.Module):

    def __init__(self,
                 cross_attn_type,
                 dim,
                 ffn_dim,
                 num_heads,
                 qk_norm=True,
                 cross_attn_norm=False,
                 eps=1e-6,
                 cc_rope_mode: str = "dual_prope",
                 cc_pose_dim: int = 14,
                 cc_phase_slots: int = 16):
        super().__init__()
        self.dim = dim
        self.ffn_dim = ffn_dim
        self.num_heads = num_heads
        self.qk_norm = qk_norm
        self.cross_attn_norm = cross_attn_norm
        self.eps = eps

        # layers
        self.norm1 = WanLayerNorm(dim, eps)
        self.self_attn = CausalWanSelfAttention(
            dim, num_heads, qk_norm, eps,
            cc_rope_mode=cc_rope_mode,
            cc_pose_dim=cc_pose_dim,
            cc_phase_slots=cc_phase_slots,
        )
        self.norm3 = WanLayerNorm(
            dim, eps,
            elementwise_affine=True) if cross_attn_norm else nn.Identity()
        self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,
                                                                      num_heads,
                                                                      (-1, -1),
                                                                      qk_norm,
                                                                      eps)
        self.norm2 = WanLayerNorm(dim, eps)
        self.ffn = nn.Sequential(
            nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
            nn.Linear(ffn_dim, dim))

        # modulation
        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)

    def forward(
        self,
        x,
        e,
        seq_lens,
        freqs_x,
        context,
        context_lens,
        crossattn_cache=None,
        kv_cache=None,
        kv_size=(0,0),
        viewmats=None,
        Ks=None,
        pose_c=None,
        tokens_per_frame: int = None,
        chunk_mask=None,
        prope_temporal_dim: int = 0,
        prope_freqs_offset: int = 0,
        prope_freqs_positions=None,
        sequence_parallel: bool = False,
        sequence_parallel_clean_history: bool = False,
        chunk_flash: bool = False,
        chunk_flash_tokens_per_chunk: int = 0,
        clean_history_chunk_flash: bool = False,
        clean_history_half_tokens: int = 0,
        clean_history_tokens_per_chunk: int = 0,
        context_streams: int = 1,
    ):
        r"""
        Args:
            x(Tensor): Shape [B, L, C]
            e(Tensor): Shape [B, F, 6, C] — per-frame modulation.
                Tokens in x are assumed to be contiguous per-frame:
                L = F * tokens_per_frame, with tokens grouped by frame.
            viewmats/Ks: optional camera info for ProPE. If None, ProPE is skipped.
            chunk_mask: optional [L, L] mask → training mode (full-seq + chunk-causal).
            prope_temporal_dim: dims of head_dim allocated to temporal RoPE
                (0 = pure spatial ProPE). Passed through to self_attn.
            prope_freqs_offset: frame-index offset for temporal RoPE (matches
                the standard RoPE's freqs_offset in compact-RoPE inference).
        """
        B_, L_, D_ = x.shape
        F_ = e.shape[1]
        assert L_ % F_ == 0, f"L ({L_}) not divisible by F ({F_})"
        tpf = L_ // F_

        # modulation [1, 6, C] → broadcast with [B, F, 6, C] via unsqueeze(1)
        e_mod = self.modulation.unsqueeze(1) + e           # [B, F, 6, C]
        e_list = e_mod.chunk(6, dim=2)                      # 6 × [B, F, 1, C]

        def _mod(x_, scale, shift):
            # x_: [B, L, C]; scale/shift: [B, F, 1, C] → broadcast across tpf
            return (x_.view(B_, F_, tpf, D_) * (1 + scale) + shift).view(B_, L_, D_)

        def _res(x_, y_, gate):
            return (
                x_.view(B_, F_, tpf, D_) + y_.view(B_, F_, tpf, D_) * gate
            ).view(B_, L_, D_)

        y = self.self_attn(_mod(self.norm1(x), e_list[1], e_list[0]),
                           seq_lens, freqs_x,
                           kv_cache=kv_cache, kv_size=kv_size,
                           viewmats=viewmats, Ks=Ks,
                           pose_c=pose_c,
                           tokens_per_frame=tpf,
                           chunk_mask=chunk_mask,
                           prope_temporal_dim=prope_temporal_dim,
                           prope_freqs_offset=prope_freqs_offset,
                           prope_freqs_positions=prope_freqs_positions,
                           sequence_parallel=sequence_parallel,
                           sequence_parallel_clean_history=sequence_parallel_clean_history,
                           chunk_flash=chunk_flash,
                           chunk_flash_tokens_per_chunk=chunk_flash_tokens_per_chunk,
                           clean_history_chunk_flash=clean_history_chunk_flash,
                           clean_history_half_tokens=clean_history_half_tokens,
                           clean_history_tokens_per_chunk=clean_history_tokens_per_chunk)

        x = _res(x, y, e_list[2])

        # A 4D context carries one independent text prompt per latent chunk:
        # [B, N_chunk, L_text, C]. Vectorize cross-attention over
        # B * stream * chunk while leaving self-attention and its KV cache
        # untouched.
        if context.ndim == 4:
            B_ctx, num_contexts, text_len, context_dim = context.shape
            streams = int(context_streams)
            if B_ctx != B_ or streams <= 0 or F_ % streams != 0:
                raise ValueError(
                    "invalid chunk-text context layout: "
                    f"x_frames={F_}, context={tuple(context.shape)}, "
                    f"streams={streams}")
            frames_per_stream = F_ // streams
            if frames_per_stream % num_contexts != 0:
                raise ValueError(
                    "chunk-text contexts must evenly partition each stream: "
                    f"frames={frames_per_stream}, contexts={num_contexts}")
            frames_per_context = frames_per_stream // num_contexts
            tokens_per_context = frames_per_context * tpf
            x_grouped = self.norm3(x).view(
                B_, streams, num_contexts, tokens_per_context, D_
            ).reshape(B_ * streams * num_contexts, tokens_per_context, D_)
            context_grouped = context.unsqueeze(1).expand(
                B_, streams, num_contexts, text_len, context_dim,
            ).reshape(B_ * streams * num_contexts, text_len, context_dim)
            grouped_cache = (
                crossattn_cache if streams * num_contexts == 1 else None)
            cross_out = self.cross_attn(
                x_grouped, context_grouped, None,
                crossattn_cache=grouped_cache,
            ).view(B_, streams, num_contexts, tokens_per_context, D_)
            x = x + cross_out.reshape(B_, L_, D_)
        else:
            x = x + self.cross_attn(
                self.norm3(x), context, context_lens,
                crossattn_cache=crossattn_cache)

        y = self.ffn(_mod(self.norm2(x), e_list[4], e_list[3]))
        x = _res(x, y, e_list[5])

        return x


class CausalHead(nn.Module):
    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
        super().__init__()
        self.dim = dim
        self.out_dim = out_dim
        self.patch_size = patch_size
        self.eps = eps

        # layers
        out_dim = math.prod(patch_size) * out_dim
        self.norm = WanLayerNorm(dim, eps)
        self.head = nn.Linear(dim, out_dim)

        # modulation
        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)

    def forward(self, x, e):
        r"""
        Args:
            x(Tensor): Shape [B, L, C] where L = F * tokens_per_frame
            e(Tensor): Shape [B, F, C] — per-frame time embedding.
        """
        B_, L_, D_ = x.shape
        F_ = e.shape[1]
        assert L_ % F_ == 0, f"L ({L_}) not divisible by F ({F_})"
        tpf = L_ // F_
        # modulation [1, 2, C] + e.unsqueeze(2) [B, F, 1, C] → [B, F, 2, C]
        ss = self.modulation.unsqueeze(1) + e.unsqueeze(2)
        shift, scale = ss.chunk(2, dim=2)  # each [B, F, 1, C]
        x_fr = self.norm(x).view(B_, F_, tpf, D_)
        x_fr = x_fr * (1 + scale) + shift
        return self.head(x_fr.view(B_, L_, D_))


class CausalWanModel(ModelMixin, ConfigMixin):
    r"""
    Wan diffusion backbone supporting both text-to-video and image-to-video.
    """

    ignore_for_config = [
        'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim'
    ]
    _no_split_modules = ['WanAttentionBlock']
    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(self,
                 model_type='t2v',
                 patch_size=(1, 2, 2),
                 text_len=512,
                 in_dim=16,
                 dim=2048,
                 ffn_dim=8192,
                 freq_dim=256,
                 text_dim=4096,
                 out_dim=16,
                 num_heads=16,
                 num_layers=32,
                 qk_norm=True,
                 cross_attn_norm=True,
                 eps=1e-6,
                 # Spatiotemporal ProPE: split head_dim into spatial (ProPE) + temporal (RoPE).
                 # 0 = pure spatial (original ProPE). Must be even and leave head_dim - k
                 # divisible by 4. For head_dim=128, valid choices: 0, 8, 16, 24, 32...
                 prope_temporal_dim=0,
                 # Camera/pose-aware RoPE variants. 'standard' is the plain
                 # Wan original t/h/w RoPE control with no camera/ProPE modules.
                 # 'dual_prope' keeps legacy 2-attention path. 'cc_basic',
                 # 'cc_output', 'cc_dual_channel' fuse camera info into RoPE,
                 # running a SINGLE attention per block. All three are bit-exact
                 # identity at step 0 (zero-init MLP + prope_proj).
                 cc_rope_mode: str = "dual_prope",
                 cc_phase_slots: int = 16,
                 degradation_control_dim: int = 0,
                 degradation_control_hidden_dim: int = 256):
        r"""
        Initialize the diffusion model backbone.

        Args:
            model_type (`str`, *optional*, defaults to 't2v'):
                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
            text_len (`int`, *optional*, defaults to 512):
                Fixed length for text embeddings
            in_dim (`int`, *optional*, defaults to 16):
                Input video channels (C_in)
            dim (`int`, *optional*, defaults to 2048):
                Hidden dimension of the transformer
            ffn_dim (`int`, *optional*, defaults to 8192):
                Intermediate dimension in feed-forward network
            freq_dim (`int`, *optional*, defaults to 256):
                Dimension for sinusoidal time embeddings
            text_dim (`int`, *optional*, defaults to 4096):
                Input dimension for text embeddings
            out_dim (`int`, *optional*, defaults to 16):
                Output video channels (C_out)
            num_heads (`int`, *optional*, defaults to 16):
                Number of attention heads
            num_layers (`int`, *optional*, defaults to 32):
                Number of transformer blocks
            qk_norm (`bool`, *optional*, defaults to True):
                Enable query/key normalization
            cross_attn_norm (`bool`, *optional*, defaults to False):
                Enable cross-attention normalization
            eps (`float`, *optional*, defaults to 1e-6):
                Epsilon value for normalization layers
        """

        super().__init__()

        assert model_type in ['t2v', 'i2v', 'ti2v']
        self.model_type = model_type

        self.patch_size = patch_size
        self.text_len = text_len
        self.in_dim = in_dim
        self.dim = dim
        self.ffn_dim = ffn_dim
        self.freq_dim = freq_dim
        self.text_dim = text_dim
        self.out_dim = out_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.qk_norm = qk_norm
        self.cross_attn_norm = cross_attn_norm
        self.eps = eps
        self.prope_temporal_dim = prope_temporal_dim
        assert cc_rope_mode in _VALID_CC_ROPE_MODES, (
            f"cc_rope_mode={cc_rope_mode!r} must be one of {_VALID_CC_ROPE_MODES}")
        self.cc_rope_mode = cc_rope_mode
        self.cc_phase_slots = cc_phase_slots
        self.degradation_control_dim = int(degradation_control_dim)
        self.degradation_control_hidden_dim = int(
            degradation_control_hidden_dim)
        # Pose descriptor size: 3 (t) + 9 (R flat) + 2 (log fx, log fy) = 14.
        # Ks is optional; when absent, _camera_pose_features drops the 2 log-focal
        # entries (→ 12). We size the MLP for the MAX (14) and pad with zeros at
        # runtime if Ks is missing.
        self.cc_pose_dim = 14

        # embeddings
        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.degradation_control_embedding = None
        if self.degradation_control_dim > 0:
            self.degradation_control_embedding = nn.Sequential(
                nn.Linear(
                    self.degradation_control_dim,
                    self.degradation_control_hidden_dim,
                    bias=False,
                ),
                nn.SiLU(),
                nn.Linear(
                    self.degradation_control_hidden_dim, dim, bias=False),
            )
            self.reset_degradation_control_parameters()

        # blocks
        cross_attn_type = 'i2v_cross_attn' if model_type == 'i2v' else 't2v_cross_attn'
        self.blocks = nn.ModuleList([
            CausalWanAttentionBlock(
                cross_attn_type, dim, ffn_dim, num_heads, qk_norm,
                cross_attn_norm, eps,
                cc_rope_mode=cc_rope_mode,
                cc_pose_dim=self.cc_pose_dim,
                cc_phase_slots=cc_phase_slots,
            )
            for _ in range(num_layers)
        ])

        # head
        self.head = CausalHead(dim, out_dim, patch_size, eps)

        # buffers (don't use register_buffer otherwise dtype will be changed in to())
        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0

        self.gradient_checkpointing = False

    def reset_degradation_control_parameters(self):
        """Initialize a missing control adapter as an exact identity branch."""
        if self.degradation_control_embedding is None:
            return
        first = self.degradation_control_embedding[0]
        last = self.degradation_control_embedding[2]
        first.reset_parameters()
        last.reset_parameters()
        with torch.no_grad():
            last.weight.zero_()



    def get_transformer_module(self):
        return {type(self.blocks[0])}

    def init_freqs(self,device):
        d = self.dim // self.num_heads
        self.freqs = torch.cat([
            rope_params(1024, d - 4 * (d // 6)),
            rope_params(1024, 2 * (d // 6)),
            rope_params(1024, 2 * (d // 6))
        ], dim=1)
        self.freqs = self.freqs.to(device)

    def _get_chunk_block_mask(self, total_tokens, tokens_per_chunk, device):
        """
        Build (or fetch from cache) a flex_attention BlockMask for chunk-causal
        attention. Keyed on (total_tokens, tokens_per_chunk, device) so we only
        compile once per sequence shape.
        """
        if not hasattr(self, "_block_mask_cache"):
            self._block_mask_cache = {}
        padded_total_tokens = math.ceil(total_tokens / 128) * 128
        key = (total_tokens, tokens_per_chunk, str(device))
        cached = self._block_mask_cache.get(key)
        if cached is not None:
            return cached

        def chunk_causal(b, h, q_idx, kv_idx):
            # chunk index of query vs key; q's chunk must be >= kv's chunk
            valid = (q_idx < total_tokens) & (kv_idx < total_tokens)
            causal = (q_idx // tokens_per_chunk) >= (kv_idx // tokens_per_chunk)
            return (valid & causal) | (q_idx == kv_idx)

        block_mask = create_block_mask(
            chunk_causal,
            B=None, H=None,
            Q_LEN=padded_total_tokens, KV_LEN=padded_total_tokens,
            device=device,
            _compile=True,
        )
        self._block_mask_cache[key] = block_mask
        return block_mask

    def _get_clean_history_chunk_block_mask(self, half_tokens, tokens_per_chunk, device):
        """
        Build a BlockMask for a doubled [near_clean, noisy] training sequence.

        Layout:
          - tokens [0, half_tokens) are near-clean GT history.
          - tokens [half_tokens, 2 * half_tokens) are the noisy denoising stream.

        A noisy chunk can see only its own noisy chunk plus previous
        near-clean chunks.  A near-clean chunk can see near-clean chunks up to
        itself.  This mirrors Self-Forcing's clean-history mask.
        """
        if not hasattr(self, "_block_mask_cache"):
            self._block_mask_cache = {}
        total_tokens = half_tokens * 2
        padded_total_tokens = math.ceil(total_tokens / 128) * 128
        key = ("clean_history_clean_noisy", half_tokens, tokens_per_chunk, str(device))
        cached = self._block_mask_cache.get(key)
        if cached is not None:
            return cached

        def clean_history_mask(b, h, q_idx, kv_idx):
            # Keep this as pure index math so create_block_mask can compile it.
            # Semantics match the previous lookup-table mask:
            # clean chunk i -> clean chunks [0..i]
            # noisy chunk i -> noisy chunk i + clean chunks [0..i-1]
            is_clean_q = q_idx < half_tokens
            clean_chunk_end = ((q_idx // tokens_per_chunk) + 1) * tokens_per_chunk
            clean_mask = (
                is_clean_q
                & (kv_idx < clean_chunk_end)
                & (kv_idx < half_tokens)
            )

            is_noisy_q = (q_idx >= half_tokens) & (q_idx < total_tokens)
            noisy_rel_idx = q_idx - half_tokens
            noisy_chunk_idx = noisy_rel_idx // tokens_per_chunk
            noisy_start = half_tokens + noisy_chunk_idx * tokens_per_chunk
            noisy_end = noisy_start + tokens_per_chunk
            noisy_self_chunk = (
                is_noisy_q
                & (kv_idx >= noisy_start)
                & (kv_idx < noisy_end)
            )
            clean_context_end = noisy_chunk_idx * tokens_per_chunk
            noisy_prev_clean = (
                is_noisy_q
                & (kv_idx < clean_context_end)
            )
            return clean_mask | noisy_self_chunk | noisy_prev_clean | (q_idx == kv_idx)

        compile_mask = _os.environ.get("REMIND_TF_MASK_COMPILE", "0") != "0"
        block_mask = create_block_mask(
            clean_history_mask,
            B=None, H=None,
            Q_LEN=padded_total_tokens, KV_LEN=padded_total_tokens,
            device=device,
            _compile=compile_mask,
        )
        self._block_mask_cache[key] = block_mask
        return block_mask

    def _set_gradient_checkpointing(self, value=False):
        self.gradient_checkpointing = value

    def forward(
        self,
        x,
        t,
        context,
        seq_len,
        y=None,
        kv_cache: dict = None,
        crossattn_cache: dict = None,
        kv_size=(0,0),
        image_latent_input: torch.Tensor = None,
        render_latent_input: torch.Tensor = None,
        freqs_offset: int = 0,
        freqs_positions: torch.Tensor = None,
        viewmats: torch.Tensor = None,     # [B, F, 4, 4] per-frame c2w, optional
        Ks: torch.Tensor = None,           # [B, F, 3, 3] per-frame intrinsics
        chunk_size: int = None,            # latent frames per chunk (training mode)
        clean_history_clean_x: torch.Tensor = None,  # [B, C, F, H, W]
        clean_history_clean_t: torch.Tensor = None,  # [B, F] or [B, 1]
        degradation_control: torch.Tensor = None,  # [B, F, control_dim]
    ):
        r"""
        Run the diffusion model with kv caching.
        See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details.
        This function will be run for num_frame times.
        Process the latent frames one by one (1560 tokens each)

        Args:
            x (List[Tensor]):
                List of input video tensors, each with shape [C_in, F, H, W]
            t (Tensor):
                Diffusion timesteps tensor of shape [B]
            context (List[Tensor]):
                List of text embeddings each with shape [L, C]
            seq_len (`int`):
                Maximum sequence length for positional encoding
            y (List[Tensor], *optional*):
                Conditional video inputs for image-to-video mode, same shape as x

        Returns:
            List[Tensor]:
                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
        """
        # params
        device = self.patch_embedding.weight.device
        if hasattr(self, 'freqs'):
            if self.freqs.device != device:
                self.freqs = self.freqs.to(device)
        else:
            self.init_freqs(device)

        clean_history = clean_history_clean_x is not None
        f, h, w = x.shape[2:]
        orig_f = f
        if clean_history:
            assert kv_cache is None, "clean_history_clean_x is training-only"
            assert clean_history_clean_x.shape[2:] == x.shape[2:], (
                f"clean_history_clean_x shape {tuple(clean_history_clean_x.shape)} "
                f"must match x shape {tuple(x.shape)}")

        sp_pad_frames = 0
        if (chunk_size is not None and self.training and torch.is_grad_enabled()):
            try:
                from wan.distributed.util import get_sequence_parallel_world_size
                sp_size_for_padding = get_sequence_parallel_world_size()
            except Exception:
                sp_size_for_padding = 1
            if sp_size_for_padding > 1:
                align = sp_size_for_padding
                if chunk_size is not None and chunk_size > 0:
                    align = align * chunk_size // math.gcd(align, chunk_size)
                padded_f = math.ceil(f / align) * align
                sp_pad_frames = padded_f - f

        def _pad_latent_frames(tensor):
            if tensor is None or sp_pad_frames <= 0:
                return tensor
            pad = tensor.new_zeros(
                tensor.shape[0], tensor.shape[1], sp_pad_frames,
                *tensor.shape[3:])
            return torch.cat([tensor, pad], dim=2)

        def _pad_frame_tensor(tensor):
            if tensor is None or sp_pad_frames <= 0:
                return tensor
            pad = tensor[:, -1:].expand(
                tensor.shape[0], sp_pad_frames, *tensor.shape[2:])
            return torch.cat([tensor, pad], dim=1)

        if sp_pad_frames > 0:
            x = _pad_latent_frames(x)
            clean_history_clean_x = _pad_latent_frames(clean_history_clean_x)
            image_latent_input = _pad_latent_frames(image_latent_input)
            render_latent_input = _pad_latent_frames(render_latent_input)
            viewmats = _pad_frame_tensor(viewmats)
            Ks = _pad_frame_tensor(Ks)
            degradation_control = _pad_frame_tensor(degradation_control)
            if freqs_positions is not None:
                freqs_positions = freqs_positions.to(device=device, dtype=torch.long)
                extra = freqs_positions[-1] + torch.arange(
                    1, sp_pad_frames + 1, device=device, dtype=torch.long)
                freqs_positions = torch.cat([freqs_positions, extra], dim=0)
            f = padded_f
        h = h//2
        w = w//2

        c = self.dim // self.num_heads // 2
        freqs = self.freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)

        # Compute freqs_x once (same for all branches).  Most callers use the
        # compact contiguous convention `freqs_offset + arange(F)`.  Ref-cache
        # training/validation can pass explicit non-contiguous positive frame
        # positions, e.g. reference at 0 and current video at +40 latent frames
        # (equivalent to reference at t=-40 because RoPE attention only uses
        # relative phase differences).
        prope_freqs_positions = None
        if freqs_positions is None:
            base_freqs_positions = torch.arange(
                freqs_offset, freqs_offset + f, device=device, dtype=torch.long)
        else:
            freqs_positions = freqs_positions.to(device=device, dtype=torch.long)
            assert freqs_positions.dim() == 1 and freqs_positions.numel() == f, (
                f"freqs_positions must be [F={f}], got {tuple(freqs_positions.shape)}")
            assert int(freqs_positions.min().item()) >= 0, \
                "freqs_positions must be non-negative; shift positions instead"
            assert int(freqs_positions.max().item()) < freqs[0].shape[0], (
                f"freqs_positions max={int(freqs_positions.max().item())} exceeds "
                f"RoPE table length {freqs[0].shape[0]}")
            base_freqs_positions = freqs_positions
        if clean_history:
            effective_freqs_positions = torch.cat(
                [base_freqs_positions, base_freqs_positions], dim=0)
        else:
            effective_freqs_positions = base_freqs_positions
        assert int(effective_freqs_positions.max().item()) < freqs[0].shape[0], (
            f"freqs_positions max={int(effective_freqs_positions.max().item())} exceeds "
            f"RoPE table length {freqs[0].shape[0]}")
        temporal_freqs = freqs[0].index_select(0, effective_freqs_positions)
        if freqs_positions is not None or clean_history:
            prope_freqs_positions = effective_freqs_positions
        f_attn = f * 2 if clean_history else f
        freqs_x = torch.cat([
            temporal_freqs.view(f_attn, 1, 1, -1).expand(f_attn, h, w, -1),
            freqs[1][:h].view(1, h, 1, -1).expand(f_attn, h, w, -1),
            freqs[2][:w].view(1, 1, w, -1).expand(f_attn, h, w, -1)
        ], dim=-1).reshape(f_attn*h*w, 1, -1)

        # Input-channel handling.
        # If the model was configured with in_dim != x.shape[1], expand x with
        # render_latent_input (legacy v2v/t2v path) to reach self.in_dim.
        # If x already matches self.in_dim, pass through directly — this is the
        # "pure text + pose + RoPE/ProPE" regime where no render/mask is used.
        if x.shape[1] != self.in_dim:
            if render_latent_input is None:
                # Legacy t2v: pad to self.in_dim with zeros
                pad_ch = self.in_dim - x.shape[1]
                assert pad_ch > 0, \
                    f"x has {x.shape[1]} channels, model in_dim is {self.in_dim}"
                x = torch.cat([x, x.new_zeros(x.shape[0], pad_ch, *x.shape[2:])], dim=1)
            elif kv_size[1] >= 0:
                # Legacy v2v: concat render_latent_input
                x = torch.cat([x, render_latent_input], dim=1)
        assert x.shape[1] == self.in_dim, \
            f"x channels ({x.shape[1]}) != model in_dim ({self.in_dim})"
        if clean_history:
            if clean_history_clean_x.shape[1] != self.in_dim:
                pad_ch = self.in_dim - clean_history_clean_x.shape[1]
                assert pad_ch > 0 and render_latent_input is None, (
                    "clean_history_clean_x currently supports pure-I2V/in_dim "
                    f"padding only, got clean channels={clean_history_clean_x.shape[1]} "
                    f"model in_dim={self.in_dim}")
                clean_history_clean_x = torch.cat([
                    clean_history_clean_x,
                    clean_history_clean_x.new_zeros(
                        clean_history_clean_x.shape[0], pad_ch,
                        *clean_history_clean_x.shape[2:])
                ], dim=1)
            assert clean_history_clean_x.shape[1] == self.in_dim, (
                f"teacher clean channels ({clean_history_clean_x.shape[1]}) "
                f"!= model in_dim ({self.in_dim})")


        # embeddings
        x_noisy = [self.patch_embedding(u.unsqueeze(0)) for u in x]
        x_clean = (
            [self.patch_embedding(u.unsqueeze(0)) for u in clean_history_clean_x]
            if clean_history else None
        )
        grid_sizes = torch.stack([
            torch.as_tensor(u.shape[2:], dtype=torch.long, device=u.device)
            for u in x_noisy
        ])
        x_noisy = [u.flatten(2).transpose(1, 2) for u in x_noisy]
        seq_lens = torch.as_tensor(
            [u.size(1) for u in x_noisy], dtype=torch.long, device=x_noisy[0].device)
        assert seq_lens.max() <= seq_len
        if clean_history:
            x_clean = [u.flatten(2).transpose(1, 2) for u in x_clean]
            seq_lens_clean = torch.as_tensor(
                [u.size(1) for u in x_clean], dtype=torch.long, device=x_noisy[0].device)
            assert torch.equal(seq_lens_clean, seq_lens), (
                f"teacher clean token lengths {seq_lens_clean.tolist()} "
                f"!= noisy token lengths {seq_lens.tolist()}")
            x = torch.cat([
                torch.cat([u_clean, u_noisy], dim=1)
                for u_noisy, u_clean in zip(x_noisy, x_clean)
            ])
        else:
            x = torch.cat(x_noisy)

        # Per-frame timesteps. Accept t of shape [B], [B, 1] (uniform), or [B, F]
        # (per-frame, used for WorldPlay memory mode where past chunks have high noise).
        if t.dim() == 1:
            t_bf_noisy = t.unsqueeze(1).expand(-1, f).contiguous()
        elif t.shape[1] == 1:
            t_bf_noisy = t.expand(-1, f).contiguous()
        else:
            if t.shape[1] == orig_f and sp_pad_frames > 0:
                t = torch.cat(
                    [t, t[:, -1:].expand(-1, sp_pad_frames)],
                    dim=1)
            assert t.shape[1] == f, (
                f"timestep frames ({t.shape[1]}) must match input frames ({f})")
            t_bf_noisy = t
        if clean_history:
            if clean_history_clean_t is None:
                t_bf_clean = torch.zeros_like(t_bf_noisy)
            elif clean_history_clean_t.dim() == 1:
                t_bf_clean = clean_history_clean_t.unsqueeze(1).expand(-1, f).contiguous()
            elif clean_history_clean_t.shape[1] == 1:
                t_bf_clean = clean_history_clean_t.expand(-1, f).contiguous()
            else:
                if (clean_history_clean_t.shape[1] == orig_f
                        and sp_pad_frames > 0):
                    clean_history_clean_t = torch.cat(
                        [clean_history_clean_t,
                         clean_history_clean_t[:, -1:].expand(
                             -1, sp_pad_frames)],
                        dim=1)
                assert clean_history_clean_t.shape[1] == f, (
                    f"teacher clean timestep frames ({clean_history_clean_t.shape[1]}) "
                    f"must match input frames ({f})")
                t_bf_clean = clean_history_clean_t
            t_bf = torch.cat([t_bf_clean.to(t_bf_noisy.device), t_bf_noisy], dim=1)
        else:
            t_bf = t_bf_noisy

        B_sz = t_bf.shape[0]
        t_flat = t_bf.reshape(-1)  # [B*F]
        e_flat = self.time_embedding(
            sinusoidal_embedding_1d(self.freq_dim, t_flat).type_as(x)
        )  # [B*F, dim]
        if self.degradation_control_embedding is not None:
            if degradation_control is None:
                degradation_control = e_flat.new_zeros(
                    B_sz, f, self.degradation_control_dim)
            else:
                assert degradation_control.ndim == 3, (
                    "degradation_control must have shape [B, F, D], got "
                    f"{tuple(degradation_control.shape)}")
                assert degradation_control.shape[:2] == (B_sz, f), (
                    "degradation_control batch/frames must match input: "
                    f"got {tuple(degradation_control.shape[:2])}, "
                    f"expected {(B_sz, f)}")
                assert degradation_control.shape[2] == self.degradation_control_dim, (
                    "degradation_control feature dim mismatch: "
                    f"got {degradation_control.shape[2]}, "
                    f"expected {self.degradation_control_dim}")
                degradation_control = degradation_control.to(
                    device=e_flat.device, dtype=e_flat.dtype)
            if clean_history:
                degradation_control = torch.cat(
                    [degradation_control, degradation_control], dim=1)
            e_flat = e_flat + self.degradation_control_embedding(
                degradation_control.reshape(-1, self.degradation_control_dim))
        elif degradation_control is not None:
            raise ValueError(
                "degradation_control was provided but the model was created "
                "with degradation_control_dim=0")
        e0_flat = self.time_projection(e_flat).unflatten(1, (6, self.dim))  # [B*F, 6, dim]

        # Reshape to per-frame representation. Blocks expand to per-token internally.
        e = e_flat.view(B_sz, f_attn, self.dim)           # [B, F_attn, dim]
        e0 = e0_flat.view(B_sz, f_attn, 6, self.dim)      # [B, F_attn, 6, dim]
        e_head = e[:, f:] if clean_history else e

        if isinstance(context, torch.Tensor) and context.ndim == 4:
            B_context, N_context, L_context, D_context = context.shape
            if L_context > self.text_len:
                context = context[:, :, :self.text_len]
                L_context = self.text_len
            elif L_context < self.text_len:
                pad = context.new_zeros(
                    B_context, N_context,
                    self.text_len - L_context, D_context)
                context = torch.cat([context, pad], dim=2)
                L_context = self.text_len
            context = self.text_embedding(
                context.reshape(
                    B_context * N_context, L_context, D_context)
            ).reshape(B_context, N_context, L_context, self.dim)
        else:
            context = self.text_embedding(torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]))

        # Keep viewmats/Ks as [B, F, 4, 4] / [B, F, 3, 3] (one per frame).
        # prope_qkv handles the frame→token mapping natively: it treats
        # `cameras = viewmats.shape[1]` as the group count and applies each
        # camera matrix to `seqlen // cameras = h*w` consecutive tokens via
        # a batched einsum. Expanding to per-token was ~1560× redundant in
        # memory and forced the einsum into 32760 tiny per-camera matmuls
        # instead of 21 batched ones.
        viewmats_expanded = None
        Ks_expanded = None
        if viewmats is not None:
            assert viewmats.shape[1] == f, \
                f"viewmats frames ({viewmats.shape[1]}) != input frames ({f})"
            viewmats_expanded = viewmats.to(device=x.device, dtype=torch.float32)
            if Ks is not None:
                Ks_expanded = Ks.to(device=x.device, dtype=torch.float32)
                # Normalize intrinsics: pixel-unit focal lengths (400-1200) cause
                # ProPE projection matrices with O(1000) entries → bf16 overflow and
                # gradient explosion. Dividing by image width brings entries to O(1).
                # image_w = latent_w * patch * vae_spatial.
                # Wan2.1 VAE uses stride 8; Wan2.2-TI2V-5B native VAE uses
                # stride 16 and 48-channel latents.
                vae_spatial = 16.0 if self.head.out_dim == 48 else 8.0
                img_w = w * 2.0 * vae_spatial
                Ks_expanded = Ks_expanded.clone()
                Ks_expanded[..., 0, :] = Ks_expanded[..., 0, :] / img_w
                Ks_expanded[..., 1, :] = Ks_expanded[..., 1, :] / img_w
            # Normalize c2w translations per-scene: center cameras at mean position
            # and scale to unit sphere. DL3DV scene scales vary wildly (1m indoor to
            # 100m+ outdoor); without this, P = K @ c2w can have huge entries.
            viewmats_expanded = viewmats_expanded.clone()
            t_all = viewmats_expanded[:, :, :3, 3]          # [B, F, 3]
            t_center = t_all.mean(dim=1, keepdim=True)       # [B, 1, 3]
            t_scale = (t_all - t_center).norm(dim=-1).max(dim=1, keepdim=True).values.unsqueeze(-1)  # [B, 1, 1]
            t_scale = t_scale.clamp(min=0.01)                 # avoid div-by-zero
            viewmats_expanded[:, :, :3, 3] = (t_all - t_center) / t_scale
            if clean_history:
                viewmats_expanded = torch.cat(
                    [viewmats_expanded, viewmats_expanded], dim=1)
                if Ks_expanded is not None:
                    Ks_expanded = torch.cat([Ks_expanded, Ks_expanded], dim=1)

        # ── CC-RoPE: build per-frame pose descriptor once; broadcast to all blocks ──
        # Always emit shape [B, F, cc_pose_dim=14]. If Ks is None, the last 2 entries
        # (log-focal) are zero-padded so the Linear(14 → D_half) always lines up.
        pose_c = None
        if self.cc_rope_mode in {
            "cc_basic", "cc_output", "cc_value", "cc_full",
            "cc_dual_channel", "cc_dual_output",
        } and viewmats_expanded is not None:
            B_pose = viewmats_expanded.shape[0]
            F_pose = viewmats_expanded.shape[1]
            p = _camera_pose_features(viewmats_expanded, Ks_expanded)
            if p.shape[-1] < self.cc_pose_dim:
                pad = viewmats_expanded.new_zeros(
                    B_pose, F_pose, self.cc_pose_dim - p.shape[-1], dtype=torch.float32)
                p = torch.cat([p, pad], dim=-1)
            pose_c = p                                           # [B, F, 14] float32

        # Build chunk-wise causal BlockMask (flex_attention) if training.
        # Mask semantics: tokens in chunk i can attend to tokens in chunks [0..i];
        # within a chunk, full attention.
        #
        # BlockMask is sparse in chunk-block granularity — flex_attention skips
        # upper-triangular blocks entirely rather than multiplying by zero. This
        # is 2-3× faster than dense SDPA with a bool attn_mask (which falls
        # back to the math backend because flash requires is_causal/None).
        chunk_block_mask = None
        clean_history_half_tokens = 0
        clean_history_tokens_per_chunk = 0
        if chunk_size is not None:
            tokens_per_frame = h * w
            tokens_per_chunk = chunk_size * tokens_per_frame
            num_chunks = f // chunk_size
            half_tokens = num_chunks * tokens_per_chunk
            assert half_tokens == f * tokens_per_frame, \
                f"f ({f}) must be divisible by chunk_size ({chunk_size})"

            if clean_history:
                clean_history_half_tokens = half_tokens
                clean_history_tokens_per_chunk = tokens_per_chunk
                chunk_block_mask = self._get_clean_history_chunk_block_mask(
                    half_tokens, tokens_per_chunk, x.device)
            else:
                chunk_block_mask = self._get_chunk_block_mask(
                    half_tokens, tokens_per_chunk, x.device)

        sequence_parallel_enabled = False
        sequence_parallel_clean_history = False
        grid_sizes_for_unpatchify = grid_sizes
        if (chunk_block_mask is not None and self.training
                and torch.is_grad_enabled()):
            try:
                from wan.distributed.util import (
                    get_sequence_parallel_rank,
                    get_sequence_parallel_world_size,
                )
                sp_size = get_sequence_parallel_world_size()
            except Exception:
                sp_size = 1
            if sp_size > 1:
                sp_rank = get_sequence_parallel_rank()
                if f % sp_size != 0:
                    raise ValueError(
                        f"sequence_parallel_size={sp_size} requires latent "
                        f"frame count f={f} to be divisible by sp_size")
                tokens_per_frame = h * w
                local_f = f // sp_size
                fs = sp_rank * local_f
                fe = fs + local_f
                half_tokens = f * tokens_per_frame
                local_half_tokens = local_f * tokens_per_frame

                def _slice_frames(tensor):
                    if tensor is None:
                        return None
                    if clean_history:
                        return torch.cat(
                            [tensor[:, fs:fe], tensor[:, f + fs:f + fe]],
                            dim=1)
                    return tensor[:, fs:fe]

                def _slice_token_freqs(tensor):
                    if clean_history:
                        return torch.cat(
                            [tensor[fs * tokens_per_frame:fe * tokens_per_frame],
                             tensor[half_tokens + fs * tokens_per_frame:
                                    half_tokens + fe * tokens_per_frame]],
                            dim=0)
                    return tensor[fs * tokens_per_frame:fe * tokens_per_frame]

                if clean_history:
                    clean_slice = slice(fs * tokens_per_frame, fe * tokens_per_frame)
                    noisy_slice = slice(
                        half_tokens + fs * tokens_per_frame,
                        half_tokens + fe * tokens_per_frame)
                    x = torch.cat([x[:, clean_slice], x[:, noisy_slice]], dim=1)
                    seq_lens = torch.full_like(seq_lens, local_half_tokens)
                    sequence_parallel_clean_history = True
                else:
                    x = x[:, fs * tokens_per_frame:fe * tokens_per_frame]
                    seq_lens = torch.full_like(seq_lens, local_half_tokens)

                freqs_x = _slice_token_freqs(freqs_x)
                e = _slice_frames(e)
                e0 = _slice_frames(e0)
                e_head = e[:, local_f:] if clean_history else e
                viewmats_expanded = _slice_frames(viewmats_expanded)
                Ks_expanded = _slice_frames(Ks_expanded)
                pose_c = _slice_frames(pose_c)
                if context.ndim == 4:
                    if fs % chunk_size != 0 or fe % chunk_size != 0:
                        raise ValueError(
                            "sequence-parallel frame shards must align to "
                            "chunk-text prompt boundaries")
                    context = context[:, fs // chunk_size:fe // chunk_size]
                if prope_freqs_positions is not None:
                    if clean_history:
                        prope_freqs_positions = torch.cat(
                            [prope_freqs_positions[fs:fe],
                             prope_freqs_positions[f + fs:f + fe]],
                            dim=0)
                    else:
                        prope_freqs_positions = prope_freqs_positions[fs:fe]
                grid_sizes_for_unpatchify = grid_sizes.clone()
                grid_sizes_for_unpatchify[:, 0] = local_f
                sequence_parallel_enabled = True

        # arguments
        kwargs = dict(
            e=e0,
            seq_lens=seq_lens,
            freqs_x=freqs_x,
            context=context,
            context_lens=None,
            kv_size=kv_size,
            viewmats=viewmats_expanded,
            Ks=Ks_expanded,
            pose_c=pose_c,                   # [B, F, 14] or None
            tokens_per_frame=h * w,          # per-token broadcast for CC-RoPE
            chunk_mask=chunk_block_mask,   # BlockMask (flex) or None
            # Spatiotemporal ProPE: per-frame temporal position = freqs_offset + frame_idx
            prope_temporal_dim=self.prope_temporal_dim,
            prope_freqs_offset=freqs_offset,
            prope_freqs_positions=prope_freqs_positions,
            sequence_parallel=sequence_parallel_enabled,
            sequence_parallel_clean_history=sequence_parallel_clean_history,
            chunk_flash=(
                not bool(clean_history)
                and bool(chunk_block_mask is not None)
                and not bool(sequence_parallel_enabled)
                and _os.environ.get("REMIND_DMD_CHUNK_FLASH", "0") == "1"
            ),
            chunk_flash_tokens_per_chunk=(
                tokens_per_chunk if chunk_block_mask is not None else 0
            ),
            clean_history_chunk_flash=(
                bool(clean_history)
                and bool(chunk_block_mask is not None)
                and not bool(sequence_parallel_enabled)
                and _os.environ.get("REMIND_TF_CHUNK_FLASH", "1") != "0"
            ),
            clean_history_half_tokens=clean_history_half_tokens,
            clean_history_tokens_per_chunk=clean_history_tokens_per_chunk,
            context_streams=(2 if clean_history else 1),
        )

        def create_custom_forward(module):
            def custom_forward(*inputs, **kwargs):
                return module(*inputs, **kwargs)
            return custom_forward

        for block_index, block in enumerate(self.blocks):
            kwargs['kv_cache'] = kv_cache[block_index] if kv_cache is not None else None
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                x= torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    x, **kwargs,
                    use_reentrant=False,
                )
            else:
                x= block(x, **kwargs)

        if clean_history:
            x = x[:, seq_lens[0]:]

        x = self.head(x, e_head)
        x = self.unpatchify(x, grid_sizes_for_unpatchify)
        x = torch.stack(x)
        if sequence_parallel_enabled:
            from wan.distributed.util import gather_forward_with_grad
            x = gather_forward_with_grad(x, dim=2)
        if sp_pad_frames > 0:
            x = x[:, :, :orig_f]

        return x

    def unpatchify(self, x, grid_sizes):
        r"""
        Reconstruct video tensors from patch embeddings.

        Args:
            x (List[Tensor]):
                List of patchified features, each with shape [L, C_out * prod(patch_size)]
            grid_sizes (Tensor):
                Original spatial-temporal grid dimensions before patching,
                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)

        Returns:
            List[Tensor]:
                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
        """

        c = self.out_dim
        out = []
        for u, v in zip(x, grid_sizes.tolist()):
            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
            u = torch.einsum('fhwpqrc->cfphqwr', u)
            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
            out.append(u)
        return out