code/diffsynth/models/memory/u_vit_cgla_blocks.py

"""
Wan ``DiTBlock`` with CGLA (SSE-GLA) replacing the self-attention.

A Wan DiTBlock is: self-attention -> cross-attention(text) -> FFN, each with its
own AdaLN modulation. Here the **self-attention is swapped for a camera-guided
linear attention** (DFOT ``SSEGLA``, :class:`fla.layers.sse.SSEGLA`) that runs on
the full flattened spatial-temporal token sequence ``(B, T*P, D)`` — *no*
spatial/temporal factorization (the dfot intra-frame softmax + inter-frame split
is removed). The cross-attention and FFN are Wan's, unchanged, so a Wan
checkpoint loads them.

Weight-loading (Wan 2.1 / 2.2 -> this block):
  * ``cross_attn`` / ``norm1`` / ``norm2`` / ``norm3`` / ``ffn`` /
    ``modulation`` / ``gate`` : Wan ``DiTBlock`` submodules (load fully).
  * ``self_attn.{q,k,v,o}_proj`` : the SSE-GLA's q/k/v/o projections have the
    same shapes as Wan's ``self_attn.{q,k,v,o}`` (``Linear(dim, dim)``, since
    ``key_dim == value_dim == dim``), so Wan's softmax-attn q/k/v/o weights
    initialise them. Use :func:`remap_wan_to_cgla` to rename
    ``self_attn.{q,k,v,o}`` -> ``self_attn.{q,k,v,o}_proj`` before loading.
  * Wan's ``self_attn.norm_q`` / ``norm_k`` (qk-RMSNorms) do NOT load — the
    SSE-GLA path has no qk-norm (they are "unexpected").
  * The rest of the SSE-GLA params (gates / sparse routing / LoRA deltas /
    ``pose_*`` injection / ``noise_write_gate``) are new (trained from scratch).

Three camera-pose-as-PE variants via ``use_pose_rope`` (handled inside the real
``SSEGLA``; see ``fla/layers/sse.py``):

- ``use_pose_rope=False`` (CGLA): pose into sparse stream's q2/k2/gk2/eta.
- ``use_pose_rope=True`` (PRoPE): the above PLUS dfot ``PoseRoPE`` — q/k rotated
  by learned per-token camera-pose angles (relative camera PE).
- UCPE (``mechanism="ucpe"``): PRoPE PLUS an absolute-orientation
  ``cam_encoder`` (``Linear(pose_dim -> dim)``, zero-init, added to x) built
  into this block.

This block IS the DiT block (a drop-in for ``wan_video_dit.DiTBlock``); the Wan
pipeline calls ``block(x, context, t_mod, freqs, actions)`` where ``actions`` is
the per-frame 12-dim RT camera pose, broadcast here to per-token ``pose_emb``.

``fla`` (and triton / rotary_embedding_torch) are used directly. The vendored
``flash-linear-attention`` at the Echo-Memory repo root is put on
``sys.path`` below.
"""

from __future__ import annotations

import os
import sys

_REPO_ROOT = os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
)
_FLA_PATH = os.path.join(_REPO_ROOT, "flash-linear-attention")
if os.path.isdir(_FLA_PATH) and _FLA_PATH not in sys.path:
    sys.path.insert(0, _FLA_PATH)

import torch
import torch.nn as nn

from fla.layers.sse import SSEGLA, SSEGDN

# Wan DiT components — same shapes a Wan checkpoint expects, for weight loading.
from diffsynth.models.wan_video_dit import SelfAttention, CrossAttention, GateModule, modulate


# ── CGLA block — SSE-GLA replacing Wan self-attention, + cross-attn + FFN ───────


class MLP_Action(nn.Module):
    def __init__(self, out_dim, sliding_window_size=3, r=4):
        super().__init__()
        self.proj_action = nn.Linear(r * sliding_window_size * 10, out_dim)
        nn.init.zeros_(self.proj_action.weight)
        nn.init.zeros_(self.proj_action.bias)
        self.sliding_window_size = sliding_window_size
        self.r = r

    def forward(self, x):
        bs, nr, act_dim = x.shape
        r = self.r
        n = nr // r
        actions = x.reshape(bs, n, r, act_dim)
        actions = F.pad(actions, (0, 0, 0, 0, self.sliding_window_size - 1, 1), mode="replicate")
        action_windows = []
        for i in range(self.sliding_window_size):
            action_windows.append(actions[:, i:i + n + 1])
        actions = torch.cat(action_windows, dim=2)
        actions = actions.reshape(bs, n + 1, -1)
        actions = self.proj_action(actions)
        return actions


class MLP_CamPose(nn.Module):
    def __init__(self, out_dim, pose_dim=12):
        super().__init__()
        self.proj = nn.Linear(pose_dim, out_dim)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)

    def forward(self, x):
        return self.proj(x)


class CGLATransformerBlock(nn.Module):
    """A Wan ``DiTBlock`` whose self-attention is CGLA (SSE-GLA).

    This IS the DiT block (a drop-in replacement for ``wan_video_dit.DiTBlock`` /
    ``DiTBlock_w_Action``): self-attn -> cross-attn(text) -> FFN, each with its own
    AdaLN modulation. The only change vs a Wan DiTBlock is that the softmax
    self-attention is swapped for a camera-guided **linear** attention (DFOT
    ``SSEGLA``) that attends across the full flattened spatial-temporal token
    sequence ``(B, T*P, D)`` (no spatial/temporal factorization); cross-attn and
    FFN are Wan's, unchanged, so a Wan checkpoint loads them.

    Forward signature matches the Wan pipeline's block call
    ``block(x, context, t_mod, freqs, actions)``: ``actions`` is the per-frame
    12-dim RT camera pose, broadcast here to per-token ``pose_emb`` and fed to
    the SSE-GLA's sparse stream (q2/k2/gk2/eta). ``freqs`` (Wan 3D-RoPE) is
    accepted for signature compatibility but unused — the SSE-GLA has no
    positional RoPE (the camera pose is the signal).

    Weight-loading (Wan 2.1 / 2.2 -> this block), via :func:`remap_wan_to_cgla`:
      * ``cross_attn`` / ``norm1`` / ``norm2`` / ``norm3`` / ``ffn`` /
        ``modulation`` / ``gate`` : Wan ``DiTBlock`` submodules (load fully).
      * ``self_attn.{q,k,v,o}_proj`` : SSE-GLA q/k/v/o projections, same shapes
        as Wan's ``self_attn.{q,k,v,o}`` (``Linear(dim, dim)``); Wan's softmax
        q/k/v/o initialise them. Wan's ``self_attn.norm_q`` / ``norm_k`` do NOT
        load (SSE-GLA has no qk-norm) — they are "unexpected".
      * The rest of the SSE-GLA params (gates / sparse routing / LoRA deltas /
        ``pose_*`` / ``noise_write_gate`` / ``cgla_gate``) are new (trained).

    Three camera-pose-as-PE variants via ``mechanism``:
      * ``"cgla"`` : ``use_pose_rope=False`` — pose into sparse q2/k2/gk2/eta.
      * ``"prope"``: ``use_pose_rope=True`` — above PLUS dfot ``PoseRoPE``
        (q/k rotated by learned per-token camera-pose angles => relative pose).
      * ``"ucpe"`` : ``"prope"`` PLUS an absolute-orientation ``cam_encoder``
        (``Linear(pose_dim -> dim)``, zero-init, added to x).

    Identity at init: the CGLA self-attn residual is scaled by
    ``tanh(self.cgla_gate)`` (zero-init), so at step 0 the self-attn contributes
    exactly 0 and the block reduces to Wan's cross-attn + FFN (the frozen Wan
    backbone is undisturbed). The SSE-GLA's pose-injection "up" projections are
    additionally zero-init by dfot design.
    """

    def __init__(
        self,
        has_image_input: bool,
        dim: int,
        num_heads: int,
        ffn_dim: int,
        eps: float = 1e-6,
        # SSE-GLA shape (dfot config; key_dim == dim so Wan q/k/v load).
        head_dim: int = 128,
        num_sparse_partition: int = 4,
        num_writer: int = 1,
        num_reader: int = 1,
        pose_dim: int = 12,
        pose_bottleneck: int = 64,
        gate_logit_normalizer: int = 16,
        gate_low_rank_dim: int = 16,
        use_pose_rope: bool | None = None,
        use_pose_gate_mod: bool = False,
        layer_idx: int = 0,
        # Camera-pose-as-PE variant: cgla | prope | ucpe.
        mechanism: str = "cgla",
        bidirectional: bool = True,
        # Accepted for call-site compatibility; emb_dim unused (FiLM emb = x).
        add_action_attn=False,
        action_use_temporal_attention: bool = False, 
        use_cam_pose: bool = False,
        emb_dim: int = 1024,
        num_patches: int = None,
        temporal_length: int = None,
        dropout: float = 0.0,
        rope=None,
        mode: str = "chunk",
        sse_implementation: str = "mask",
        expand_v: float = 1.0,
        **_legacy,
    ):
        super().__init__()
        self.dim = int(dim)
        self.num_heads = int(num_heads)
        self.ffn_dim = int(ffn_dim)
        self.has_image_input = bool(has_image_input)
        self.emb_dim = int(emb_dim)
        self.pose_dim = int(dim)
        self.bidirectional = bool(bidirectional)

        mechanism = str(mechanism or "cgla").lower()
        assert mechanism in ("cgla", "prope", "ucpe"), (
            f"mechanism must be one of cgla/prope/ucpe, got {mechanism!r}"
        )
        self.mechanism = mechanism
        if use_pose_rope is None:
            use_pose_rope = mechanism in ("prope", "ucpe")
        self.use_abs_orientation = (mechanism == "ucpe")

        action_use_temporal_attention = False # HardCODE(jiakuihu)
        assert add_action_attn and (not action_use_temporal_attention)
        if add_action_attn:
            # ── self_attn = CGLA (SSE-GLA), on the full flattened token sequence ──
            self.self_attn_with_action = SelfAttention(dim, num_heads, eps)
            nn.init.zeros_(self.self_attn_with_action.o.weight)
            nn.init.zeros_(self.self_attn_with_action.o.bias)
        if use_cam_pose:
            self.action_mlp = MLP_CamPose(dim)
        else:
            self.action_mlp = MLP_Action(dim)

        import os as _os
        _SSEBackend = SSEGLA if _os.environ.get("CGLA_BACKEND", "ssegdn").lower() == "ssegla" else SSEGDN
        self.self_attn = _SSEBackend(
            hidden_size=self.dim,
            num_heads=self.num_heads,
            head_dim=head_dim,
            expand_v=expand_v,                 # value_dim == dim -> o_proj: Linear(dim, dim)
            mode=mode,
            num_sparse_partition=num_sparse_partition,
            num_writer=num_writer,
            num_reader=num_reader,
            sse_implementation=sse_implementation,
            gate_logit_normalizer=gate_logit_normalizer,
            gate_low_rank_dim=gate_low_rank_dim,
            pose_dim=self.pose_dim,
            # pose_bottleneck=pose_bottleneck,
            rope=None,                         # no positional RoPE (pose is the signal)
            use_pose_rope=use_pose_rope,
            use_pose_gate_mod=use_pose_gate_mod,
            layer_idx=layer_idx,
        )
        self.self_attn.num_heads = self.num_heads
        self.self_attn.head_dim = self.dim // self.num_heads
        # Audit #8 (P1): the comments promise a zero-init tanh(cgla_gate) identity
        # gate on the CGLA residual, but it was never implemented — the branch fed
        # the residual stream through the non-zero AdaLN gate from step 0, and the
        # nan-locator shows the residual stream overflowing by layer 28 (q1 inf).
        # Zero-init scalar gate => identity at step 0; the gate opens by gradient.
        self.cgla_gate = nn.Parameter(torch.zeros(1))

        # Write gate for the GLA state (dfot noise_write_gate). Derived from x
        # (the block has no separate noise embedding); sigmoid(bias=5) ~= 0.99
        # at init. The SSE-GLA's own o_proj (loaded from Wan's self_attn.o) is the
        # attention output projection — no separate out_proj is needed.
        self.noise_write_gate = nn.Linear(self.dim, 1, bias=True)
        nn.init.zeros_(self.noise_write_gate.weight)
        nn.init.constant_(self.noise_write_gate.bias, 5.0)

        # UCPE "Absolute Orientation Encoding": a zero-init cam_encoder adds the
        # per-token camera pose to x before the block. Zero-init => identity at 0.
        if self.use_abs_orientation:
            self.cam_encoder = nn.Linear(self.pose_dim, self.dim, bias=True)
            nn.init.zeros_(self.cam_encoder.weight)
            nn.init.zeros_(self.cam_encoder.bias)

        # ── Wan DiTBlock submodules (load fully from a Wan checkpoint) ──
        self.cross_attn = CrossAttention(
            self.dim, self.num_heads, eps, has_image_input=self.has_image_input
        )
        self.norm1 = nn.LayerNorm(self.dim, eps=eps, elementwise_affine=False)
        self.norm2 = nn.LayerNorm(self.dim, eps=eps, elementwise_affine=False)
        self.norm3 = nn.LayerNorm(self.dim, eps=eps)
        self.ffn = nn.Sequential(
            nn.Linear(self.dim, self.ffn_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(self.ffn_dim, self.dim),
        )
        self.modulation = nn.Parameter(torch.randn(1, 6, self.dim) / self.dim**0.5)
        self.gate = GateModule()
        self.action_use_temporal_attention = action_use_temporal_attention

        self._aux_loss = None

    # def load_state_dict(self, state_dict, strict=True):
    #     # Remap Wan's self_attn.{q,k,v,o} -> SSE-GLA's {q,k,v,o}_proj so Wan's
    #     # softmax self-attention weights initialise the linear-attention q/k/v/o
    #     # projections (same shapes). For DiT-level loads, apply
    #     # ``remap_wan_to_cgla`` to the full state_dict first.
    #     return super().load_state_dict(remap_wan_to_cgla(state_dict), strict=strict)

    def _pose_emb_from_actions(self, actions, x: torch.Tensor):
        """Broadcast per-frame RT ``(B, F, pose_dim)`` -> per-token ``(B, N, pose_dim)``.

        ``actions`` may be a list, ``(F, pose_dim)``, or ``(B, F, pose_dim)``.
        Returns None if the pose cannot be aligned to the token grid (vanilla GLA).
        """
        if actions is None:
            return None
        if not torch.is_tensor(actions):
            actions = torch.tensor(actions, device=x.device, dtype=x.dtype)
        else:
            actions = actions.to(device=x.device, dtype=x.dtype)
        if actions.ndim == 2:                       # (F, pose_dim) -> (1, F, pose_dim)
            actions = actions.unsqueeze(0)
        if actions.ndim != 3 or actions.shape[-1] != self.pose_dim:
            return None
        b, n, _ = x.shape
        f = actions.shape[1]
        if f <= 1 or n % f != 0:
            return None
        if actions.shape[0] != b:
            actions = actions.expand(b, *actions.shape[1:])
        s = n // f
        return actions.unsqueeze(2).expand(b, f, s, self.pose_dim).reshape(
            b, n, self.pose_dim
        )

    def _run_cgla(self, x: torch.Tensor, pose_emb, t_feat=None) -> torch.Tensor:
        # x: (B, N, D) full flattened spatial-temporal tokens; SSE-GLA attends
        # across all N = T*P tokens (no spatial/temporal factorization).
        import os as _os
        _gate_from_t = t_feat is not None and _os.environ.get("CGLA_WRITE_GATE", "x") == "t"

        def _call(xx, pe, tf):
            # audit #7: DFOT gates writes by NOISE LEVEL; opt-in via CGLA_WRITE_GATE=t
            wg = self.noise_write_gate(tf if _gate_from_t else xx).sigmoid()  # (B, N, 1)
            o, info, _ = self.self_attn(
                xx,
                attention_mask=None,
                pose_emb=pe,
                write_gate=wg,
                use_cache=False,
            )
            return o, info

        o, info = _call(x, pose_emb, t_feat)
        aux = info[1] if info is not None and len(info) > 1 else None
        self._aux_loss = aux if torch.is_tensor(aux) else torch.zeros(
            (), device=x.device, dtype=x.dtype
        )
        # GLA is causal; for the two-chunk target-prefix / context-suffix layout,
        # run on the flipped sequence too and average so target tokens can
        # retrieve the later context (mirrors dfot's BiJointCGLATransformerBlock).
        if self.bidirectional:
            x_b = torch.flip(x, dims=[1])
            pe_b = torch.flip(pose_emb, dims=[1]) if pose_emb is not None else None
            tf_b = torch.flip(t_feat, dims=[1]) if t_feat is not None else None
            o_b, _ = _call(x_b, pe_b, tf_b)
            o_b = torch.flip(o_b, dims=[1])
            o = (o + o_b) * 0.5
        return o

    def forward(self, x, context, t_mod, freqs, actions=None):
        # Wan pipeline calls block(x, context, t_mod, freqs, actions); actions is
        # the per-frame 12-dim RT camera pose -> per-token pose_emb for the SSE-GLA.
        original_x = x
        
        actions = self.action_mlp(actions.to(x.dtype)).to(x.dtype)
        bs, num_frames, dim = actions.shape
        actions = actions.reshape(bs, num_frames, 1, dim)
        x = x.reshape(bs, num_frames, -1, dim)
        pose_emb = actions.repeat(1, 1, x.shape[2], 1).flatten(1, 2)
        if self.use_abs_orientation and actions is not None:
            x = x + self.cam_encoder(actions)
        else:
            x = x + actions

        if hasattr(self, "self_attn_with_action"):
            if not self.action_use_temporal_attention:
                x = x.reshape(bs, -1, dim)
                x = original_x + self.self_attn_with_action(x, freqs)
            else:
                from einops import rearrange
                x = rearrange(x, "b f p d -> (b p) f d")
                attn_out = self.self_attn_with_action(x)
                attn_out = rearrange(attn_out, "(b p) f d -> b f p d", b=bs)
                x = original_x + attn_out.reshape(bs, -1, dim)
        else:
            x = x.reshape(bs, -1, dim)

        # Full Wan DiTBlock path: CGLA self-attn -> cross-attn -> FFN,
        # each with AdaLN modulation. The self-attn residual is gated by
        # tanh(cgla_gate) (zero-init => identity at step 0).
        has_seq = len(t_mod.shape) == 4
        chunk_dim = 2 if has_seq else 1
        (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=chunk_dim)
        if has_seq:
            shift_msa = shift_msa.squeeze(2); scale_msa = scale_msa.squeeze(2)
            gate_msa = gate_msa.squeeze(2)
            shift_mlp = shift_mlp.squeeze(2); scale_mlp = scale_mlp.squeeze(2)
            gate_mlp = gate_mlp.squeeze(2)

        # 1. CGLA self-attention (linear attention on flattened tokens).
        input_x = modulate(self.norm1(x), shift_msa, scale_msa)
        t_feat = None
        import os as _os
        if _os.environ.get("CGLA_WRITE_GATE", "x") == "t":
            # noise-level feature: (modulation + t_mod) pooled over the 6 AdaLN
            # slots; per-frame when t_mod is sequential, broadcast to tokens.
            _tm = self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod
            _tf = _tm.mean(dim=2) if has_seq else _tm.mean(dim=1, keepdim=True)
            _n_tok = input_x.shape[1]
            _per = max(1, _n_tok // _tf.shape[1])
            t_feat = _tf.repeat_interleave(_per, dim=1)
            if t_feat.shape[1] != _n_tok:
                t_feat = t_feat[:, :_n_tok]
        x = self.gate(x, gate_msa, torch.tanh(self.cgla_gate) * self._run_cgla(input_x, pose_emb, t_feat))
        # 2. Cross-attention to text (if context provided).
        if context is not None:
            x = x + self.cross_attn(self.norm3(x), context)
        # 3. FFN.
        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
        x = self.gate(x, gate_mlp, self.ffn(input_x))
        return x


def remap_wan_to_cgla(state_dict):
    """Remap a Wan DiT checkpoint's self-attn keys to CGLA (SSE-GLA) names.

    Wan's ``self_attn.{q,k,v,o}`` (softmax attention projections, shape
    ``Linear(dim, dim)``) initialise the SSE-GLA's ``{q,k,v,o}_proj``
    projections (same shapes, since ``key_dim == value_dim == dim``). This
    remaps every ``*.self_attn.{q,k,v,o}.*`` key to ``*.self_attn.{q,k,v,o}_proj.*``
    so a Wan checkpoint loads the linear-attention projections. All other keys
    (cross_attn / ffn / norm / modulation / ...) pass through unchanged — they
    already match the CGLA block's Wan submodules.
    """
    out = {}
    for name, param in state_dict.items():
        if ".self_attn.q." in name:
            name = name.replace(".self_attn.q.", ".self_attn.q_proj.", 1)
        elif ".self_attn.k." in name:
            name = name.replace(".self_attn.k.", ".self_attn.k_proj.", 1)
        elif ".self_attn.v." in name:
            name = name.replace(".self_attn.v.", ".self_attn.v_proj.", 1)
        elif ".self_attn.o." in name:
            name = name.replace(".self_attn.o.", ".self_attn.o_proj.", 1)
        out[name] = param
    return out


__all__ = ["CGLATransformerBlock", "remap_wan_to_cgla"]