""" 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_proj.weight) if use_cam_pose: self.action_mlp = MLP_CamPose(dim) else: self.action_mlp = MLP_Action(dim) self.self_attn = SSEGDN( 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 # 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) -> torch.Tensor: # x: (B, N, D) full flattened spatial-temporal tokens; SSE-GLA attends # across all N = T*P tokens (no spatial/temporal factorization). def _call(xx, pe): wg = self.noise_write_gate(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) 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 o_b, _ = _call(x_b, pe_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) x = self.gate(x, gate_msa, self._run_cgla(input_x, pose_emb)) # 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"]