diff --git a/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py b/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py new file mode 100644 index 000000000..54ee1f936 --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py @@ -0,0 +1,253 @@ +# Copyright 2024 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Inference-side helpers for the bidirectional fused GDN path. + +Precision knob: env var ``FUSED_GDN_PRECISION`` or ``PRECISION_OVERRIDE``: + 0=IEEE fp32 dots, 1=TF32, 2=bf16 TC + fp32 state [default], 3=bf16 TC + bf16 state. +""" + +# ruff: noqa: E501 + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +# ===================================================================== +# GPU-adaptive kernel config +# ===================================================================== + + +def _get_kernel_config() -> dict: + """Return optimal kernel parameters for the current GPU. + + STATE_FP32 (fp32 state_prev) needs ~128KB SRAM (H100 228KB), vs ~96KB for + bf16 state_prev (fits GB10's 101KB). + """ + if not torch.cuda.is_available(): + return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 4, "STATE_FP32": False} + smem = torch.cuda.get_device_properties(0).shared_memory_per_multiprocessor + state_fp32 = smem >= 150 * 1024 # H100 (228KB) yes, GB10 (101KB) no + return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 8, "STATE_FP32": state_fp32} + + +_KCFG = None + + +def _kcfg(): + global _KCFG + if _KCFG is None: + _KCFG = _get_kernel_config() + return _KCFG + + +# precision=0 → IEEE fp32 dots + fp32 state (DOT_PRECISION=2, STATE_FP32=1) +# precision=1 → TF32 dots + fp32 state (DOT_PRECISION=1, STATE_FP32=1) +# precision=2 → bf16 dots + fp32 state (DOT_PRECISION=0, STATE_FP32=1) [default] +# precision=3 → bf16 dots + bf16 state (DOT_PRECISION=0, STATE_FP32=0) +def _precision_params(precision: int) -> tuple: + if precision == 0: + return 2, True + elif precision == 1: + return 1, True + elif precision == 3: + return 0, False + else: # default + return 0, True + + +_env_prec = os.environ.get("FUSED_GDN_PRECISION", None) +PRECISION_OVERRIDE: int | None = int(_env_prec) if _env_prec is not None else None + + +def _resolve_launch_config() -> tuple: + """Returns (prec, dot_prec, state_fp32, num_warps). + + Uses ``PRECISION_OVERRIDE`` when set, else ``_kcfg()`` (per-GPU SRAM). + num_warps clamped to 4 when dots run on fp32 operands (more registers). + """ + cfg = _kcfg() + prec = PRECISION_OVERRIDE if PRECISION_OVERRIDE is not None else 2 + dot_prec, state_fp32 = _precision_params(prec) + if PRECISION_OVERRIDE is None: + state_fp32 = cfg["STATE_FP32"] + nw = cfg["num_warps"] + if dot_prec >= 1: + nw = min(nw, 4) + return prec, dot_prec, state_fp32, nw + + +def prepare_rope_tables( + rotary_emb, N: int, D: int, device +) -> tuple[torch.Tensor, torch.Tensor]: + """Complex rotary_emb `(1, 1, N, D//2)` → expanded (N, D) cos/sin tables. + + Encodes the interleaved-pair rotation + y[2i] = x[2i]*cos[i] - x[2i+1]*sin[i] + y[2i+1] = x[2i]*sin[i] + x[2i+1]*cos[i] + as y[d] = x[d]*cos_exp[d] + x[d^1]*sin_exp[d] + where sin_exp[2i] = -sin[i], sin_exp[2i+1] = +sin[i]. + + Returns (cos_exp, sin_exp) both (N, D) float32, contiguous. + """ + if rotary_emb is None: + return ( + torch.ones(N, D, device=device, dtype=torch.float32), + torch.zeros(N, D, device=device, dtype=torch.float32), + ) + freqs = rotary_emb.squeeze(0).squeeze(0) # (N, D//2) complex + cos_half = freqs.real.float() + sin_half = freqs.imag.float() + rope_cos = cos_half.repeat_interleave(2, dim=-1) + rope_sin = torch.stack([-sin_half, sin_half], dim=-1).reshape(N, D) + return rope_cos.contiguous(), rope_sin.contiguous() + + +def _precompute_inv_rms( + qkv: torch.Tensor, idx: int, C: int, eps: float = 1e-5 +) -> torch.Tensor: + """Compute 1/RMS for one component of QKV over the full C = H*D channel dim. + + qkv: (B, N, 3, H, D); idx: 0=Q, 1=K, 2=V; C: H*D. Returns (B, N) float32. + """ + raw = qkv[:, :, idx].float() # (B, N, H, D) + sq_sum = (raw * raw).sum(dim=(-2, -1)) # (B, N) + return torch.rsqrt(sq_sum / C + eps) + + +# ===================================================================== +# Fused single-pass Q+K inverse-RMS Triton kernel +# ===================================================================== +# Single Triton launch that reads each `(b, n)` row of `qkv` once and emits +# both `q_inv_rms[b, n]` and `k_inv_rms[b, n]`. Replaces two separate PyTorch +# scans (cast→square→sum→rsqrt) over `qkv[:, :, 0]` and `qkv[:, :, 1]`. +# +# Layout assumed: `qkv` is (B, N, 3, H, D) contiguous, so the C = H*D channels +# for a given (b, n, qkv_idx) live in a contiguous memory span. + + +@triton.jit +def _fused_qk_inv_rms_kernel( + qkv_ptr, # *T_in (B, N, 3, H, D), contiguous + q_inv_rms_ptr, # *float32 (B, N) + k_inv_rms_ptr, # *float32 (B, N) + N: tl.constexpr, + C: tl.constexpr, # H * D + eps, + BLOCK_C: tl.constexpr, +): + bn_id = tl.program_id(0) + qkv_row_stride = 3 * C + row_base = bn_id * qkv_row_stride + q_base = row_base + k_base = row_base + C + + offs = tl.arange(0, BLOCK_C) + mask = offs < C + + q_vals = tl.load(qkv_ptr + q_base + offs, mask=mask, other=0.0).to(tl.float32) + k_vals = tl.load(qkv_ptr + k_base + offs, mask=mask, other=0.0).to(tl.float32) + + q_sq = tl.sum(q_vals * q_vals, axis=0) + k_sq = tl.sum(k_vals * k_vals, axis=0) + + inv_c = 1.0 / C + q_inv = tl.rsqrt(q_sq * inv_c + eps) + k_inv = tl.rsqrt(k_sq * inv_c + eps) + + tl.store(q_inv_rms_ptr + bn_id, q_inv) + tl.store(k_inv_rms_ptr + bn_id, k_inv) + + +def fused_qk_inv_rms( + qkv: torch.Tensor, + eps: float = 1e-5, +) -> tuple[torch.Tensor, torch.Tensor]: + """Single-pass Triton fused Q+K inverse-RMS. + + Replaces two ``_precompute_inv_rms`` scans with one launch that reads each + ``(b, n)`` row of ``qkv`` exactly once. + qkv: (B, N, 3, H, D) contiguous. Returns (q_inv_rms, k_inv_rms), each (B, N) float32. + """ + assert qkv.is_contiguous(), "qkv must be contiguous (B, N, 3, H, D)" + assert ( + qkv.dim() == 5 and qkv.shape[2] == 3 + ), f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}" + B, N, _, H, D = qkv.shape + C = H * D + q_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device) + k_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device) + BLOCK_C = triton.next_power_of_2(C) + _fused_qk_inv_rms_kernel[(B * N,)]( + qkv, + q_inv_rms, + k_inv_rms, + N=N, + C=C, + eps=eps, + BLOCK_C=BLOCK_C, + ) + return q_inv_rms, k_inv_rms + + +# ===================================================================== +# Bidirectional GDN entry point (delegates to chunkwise) +# ===================================================================== + + +def fused_bigdn_func( + qkv: torch.Tensor, # (B, N, 3, H, D) + q_inv_rms: torch.Tensor, # (B, N) float32 + k_inv_rms: torch.Tensor, # (B, N) float32 + q_norm_weight: torch.Tensor, # (C,) float32 + k_norm_weight: torch.Tensor, # (C,) float32 + rope_cos: torch.Tensor, # (N, D) float32 + rope_sin: torch.Tensor, # (N, D) float32 + beta: torch.Tensor, # (B, H, F, S) + decay: torch.Tensor, # (B, H, F) + F: int, + S: int, + k_scale: float, + eps: float = 1e-6, +) -> torch.Tensor: + """Bidirectional fused GDN. Returns ``(B, N, H, D)``. + + Thin entry point kept for call-site stability; delegates to + :func:`fused_bigdn_bidi_chunkwise` from ``fused_gdn_chunkwise``. + """ + from sglang.jit_kernel.diffusion.triton.sana_wm_gdn_chunkwise import ( + fused_bigdn_bidi_chunkwise, + ) + + return fused_bigdn_bidi_chunkwise( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + beta, + decay, + F=F, + S=S, + k_scale=k_scale, + eps=eps, + ) diff --git a/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn_chunkwise.py b/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn_chunkwise.py new file mode 100644 index 000000000..70e5c0b8f --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn_chunkwise.py @@ -0,0 +1,1734 @@ +# Copyright 2024 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +# ruff: noqa: E501 + +""" +Fused GDN — Chunkwise-parallel forward (v2). + +V2 changes vs v1: + 1. Phase A split into two kernels along GDN data streams (KV and Z — gating + sub-paths, not CUDA streams; both launch on the same CUDA stream). Z is + lighter (no V/Cos/Sin loads, no K_pair flip), enabling 2 blocks/SM + resident on H100 for latency hiding. + 2. Phase A stores (I - P_kv) / (I - P_z) instead of P_kv/P_z so Phase B's + MMA `(I-P_kv) @ M` folds the identity-add in (no separate elementwise pass). + +BiGDN inference path: QK_NORM=1, USE_PRECOMPUTED_RMS=1, SAVE_STATE=0. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + +_CAM_IDENTITY_CACHE: dict = {} + +# Per-architecture launch config (auto-selected via compute capability). +# Empirically tuned at production config (B=1..8, T=11, S=920, H=20, D=112). +# Two effects drive BLOCK_S: +# 1. Precision: fp32 operand fragments are 2× bf16. BLOCK_S=64 + fp32 → register +# spills (40-100× slower); BLOCK_S=32 + fp32 → no spills. fp32 forces BLOCK_S=32. +# 2. Arch (bf16): A100 (192 KB SRAM) prefers BLOCK_S=32; H100/GB200 (228 KB) +# tolerate BLOCK_S=64. +# +# 8 tuned knobs across 3 phases: Phase A (nw, BS); Phase B (nw, use_acc, ns); +# Phase C (nw, BS, ns). Values from empirical sweeps (T6 A100/H100 2026-04-19; +# Blackwell-DC 2026-04-20; Spark GB10 in 5da52db6 / 3ad104d0). Phase B persistent +# M[128,128] fp32 = 64 KB → nw controls register spread; Phase C loaded M = 64 KB +# → BS controls transient SMEM. New arch: pick closest bucket, then override in +# _CHUNKWISE_SHAPE_OVERRIDES once a targeted sweep lands. + + +@dataclass(frozen=True) +class _PhaseCfg: + nw: int # num_warps + BS: int = 0 # BLOCK_S (Phase A/C only; 0 = N/A for Phase B) + ns: int = 1 # num_stages + use_acc: bool = False # Phase B only: fold A_f via MMA accumulator + + +@dataclass(frozen=True) +class _ChunkwiseCfg: + A: _PhaseCfg + B: _PhaseCfg + C: _PhaseCfg + + def as_tuple(self) -> tuple: + """Flatten to the 8-tuple the legacy API returns.""" + return ( + self.A.nw, + self.A.BS, + self.B.nw, + self.B.ns, + self.B.use_acc, + self.C.nw, + self.C.BS, + self.C.ns, + ) + + +# ────────────────────────────────────────────────────────────────── +# Primary tuning table: (arch_key, prec_key) → _ChunkwiseCfg. +# Arch keys: +# "ampere" sm_80 A100 (164 KB SRAM, no WGMMA) +# "hopper" sm_90 H100 (228 KB SRAM, WGMMA) +# "blackwell_dc" sm_100 B200 / GB200 (228 KB SRAM, WGMMA v2) +# "blackwell_spark" sm_120+ with < 150 KB SRAM 5090 / GB10 (~102 KB SRAM) +# Prec keys: +# "bf16" dot_prec == 0 (bf16 TC, half-size operand fragments) +# "fp32" dot_prec >= 1 (TF32 TC or IEEE Markidis 3-pass; same launch shape) +# ────────────────────────────────────────────────────────────────── +_CHUNKWISE_TUNING: dict[tuple[str, str], _ChunkwiseCfg] = { + # A100: smaller SRAM than Hopper, no WGMMA → bigger CTAs hide MMA latency. + # Phase B fp32 needs nw=32 to spread persistent M across warps (no acc-fusion + # available pre-Hopper, so ns=2 fills the MMA pipeline slot instead). + ("ampere", "bf16"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=32), + B=_PhaseCfg(nw=8, use_acc=False, ns=1), + C=_PhaseCfg(nw=4, BS=32, ns=1), # nw=4 bf16 C: 27% faster than nw=8 per T6 + ), + ("ampere", "fp32"): _ChunkwiseCfg( + # 2026-04-30 retune: Phase A nw=16 BS=32 is 8-13× faster than the legacy + # nw=8 across all F; Phase A was the A100 sink/rolling bottleneck. + A=_PhaseCfg(nw=16, BS=32), + B=_PhaseCfg(nw=32, use_acc=False, ns=2), # ns=2 fills pipe (no acc-fusion) + C=_PhaseCfg( + nw=16, BS=32, ns=1 + ), # 2026-04-30 retune: nw=16 BS=32 is 2.8x faster (was nw=8 BS=16) + ), + # Hopper (H100): WGMMA + 228 KB SRAM → big tiles win at bf16. + # Phase B fp32 uses acc-fusion (MMA accumulator folds A_f in one op, +12%). + ("hopper", "bf16"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=64), + B=_PhaseCfg(nw=4, use_acc=False, ns=1), # small CTAs pack better on WGMMA + C=_PhaseCfg(nw=8, BS=32, ns=1), + ), + ("hopper", "fp32"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=32), # fp32 operand 2× bigger → half BS + B=_PhaseCfg( + nw=32, use_acc=False, ns=1 + ), # 2026-04-29 retune: acc_fusion=False is 3x faster post precision-gate fix + C=_PhaseCfg( + nw=16, BS=32, ns=1 + ), # 2026-04-30 retune: nw=16 BS=32 is 1.7x faster (was nw=8 BS=16) + ), + # Blackwell-DC (B200 / GB200): 228 KB SRAM + improved WGMMA codegen. + # bf16 likes small CTAs (nw=4); fp32 stays at nw=8 (nw=4 + BS=64 fp32 = 92× regression). + ("blackwell_dc", "bf16"): _ChunkwiseCfg( + A=_PhaseCfg(nw=4, BS=64), + B=_PhaseCfg(nw=4, use_acc=False, ns=1), + C=_PhaseCfg(nw=8, BS=64, ns=1), # 228 KB SRAM leaves room for BS=64 bf16 + ), + ("blackwell_dc", "fp32"): _ChunkwiseCfg( + A=_PhaseCfg( + nw=8, BS=128 + ), # 2026-04-30 retune: nw=8 BS=128 ~5% faster at production F=3-6 (sweep across F=3,5,6,11) + B=_PhaseCfg( + nw=32, use_acc=False, ns=3 + ), # 2026-04-29 retune: 14x faster (was nw=8 acc=True 17ms; now nw=32 ns=3 acc=False 1.23ms) + C=_PhaseCfg( + nw=4, BS=64, ns=1 + ), # 2026-04-30 retune: nw=4 BS=64 is 3-5x faster than old nw=8 BS=16 (sweep 2026-04-30) + ), + # Blackwell-Spark (5090 / GB10, ~102 KB SRAM): small-chip SRAM penalty, no + # Blackwell-DC WGMMA-v2 register-spread benefit. Behaves like Hopper at fp32 + # (Phase B wants nw=32, not nw=8 like DC). BS one step smaller than DC; Phase A + # bf16 wants nw=8 (nw=4 was 22× slower, 2026-04-20). Phase B nw=32 fp32: + # 1.84×/2.65× (GB10/5090) over prior nw=8 (sweep 2026-04-24, F=11 S=920). + ("blackwell_spark", "bf16"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=32), + B=_PhaseCfg( + nw=8, use_acc=False, ns=1 + ), # nw=8 (not 4) at bf16: ~5% across F=3,6,11 + # C.nw=4 BS=32: ~3.5% faster than nw=8 (Phase C is bandwidth-bound, fewer + # warps schedule better on small SRAM). BS=64 bf16 OOMs Spark SRAM. + C=_PhaseCfg(nw=4, BS=32, ns=1), + ), + ("blackwell_spark", "fp32"): _ChunkwiseCfg( + A=_PhaseCfg(nw=8, BS=16), # fp32 operand 2× bigger → BS=16 (half of DC's 32) + # nw=16 OOMs the 102 KB SRAM cap at TF32 (needs 131 KB); nw=8 fits and is + # within noise. The D-tile path (auto-enabled on spark, see + # `_pick_phase_b_d_splits`) is ~2.6× faster at TF32 / ~13% at IEEE — these + # baseline params only apply when PHASE_B_D_SPLITS=1 is forced. + B=_PhaseCfg(nw=8, use_acc=False, ns=1), + C=_PhaseCfg(nw=8, BS=16, ns=1), # binding constraint: M.fp32 64 KB + Q stage + ), +} + + +# Shape-aware override table (empty by default). Keyed by +# (arch_key, prec_key, shape_hint); exact-match, values are full `_ChunkwiseCfg` +# (no partial overrides). Strictly additive — base table is the fallback. +# Populate only when a targeted sweep shows a shape regresses with the arch config. +_CHUNKWISE_SHAPE_OVERRIDES: dict[tuple[str, str, str], _ChunkwiseCfg] = {} + + +# Per-(cap, dot_prec) exact overrides (pins a specific GPU model if the arch +# bucket is wrong for it). Also empty by default. +_ARCH_OVERRIDES: dict = {} + + +def _arch_key(cap: tuple) -> str: + """Map compute capability → named arch bucket in `_CHUNKWISE_TUNING`. + + Blackwell (cap[0] >= 10) splits into "blackwell_dc"/"blackwell_spark" by SRAM + size (≥150 KB vs less). Unknown archs / no CUDA → conservative "ampere". + """ + if cap[0] == 8: + return "ampere" + if cap[0] == 9: + return "hopper" + if cap[0] >= 10: + has_big_sram = True + if torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + smem = getattr(props, "shared_memory_per_multiprocessor", 228 * 1024) + has_big_sram = smem >= 150 * 1024 + return "blackwell_dc" if has_big_sram else "blackwell_spark" + return "ampere" + + +def _prec_key(dot_prec: int) -> str: + return "fp32" if dot_prec >= 1 else "bf16" + + +def _auto_config(dot_prec: int, cap: tuple, shape_hint: str | None = None) -> tuple: + """Look up chunkwise launch params. Resolution order: shape override → + per-(arch, prec) table → ("ampere", prec) fallback. (`_ARCH_OVERRIDES` is + applied by `_get_arch_config`, not here.) + + Returns the legacy 8-tuple `(a_nw, a_BS, b_nw, b_ns, b_use_acc, c_nw, c_BS, c_ns)`. + """ + arch = _arch_key(cap) + prec = _prec_key(dot_prec) + + if shape_hint is not None: + cfg = _CHUNKWISE_SHAPE_OVERRIDES.get((arch, prec, shape_hint)) + if cfg is not None: + return cfg.as_tuple() + + cfg = _CHUNKWISE_TUNING.get((arch, prec)) or _CHUNKWISE_TUNING[("ampere", prec)] + return cfg.as_tuple() + + +def _get_arch_config( + dot_precision: int = 0, + shape_hint: str | None = None, + device: torch.device | int | None = None, +): + """Returns (a_warps, a_BLOCK_S, b_warps, b_stages, b_use_acc_fusion, + c_warps, c_BLOCK_S, c_stages). + + dot_precision: 0=bf16 TC, 1=TF32 TC, 2=IEEE fp32. + device: capability source; pass ``qkv.device`` in multi-GPU single-process + setups so the right tuning bucket is chosen (defaults to current device). + """ + if not torch.cuda.is_available(): + cap = (9, 0) # assume modern when querying from CPU + else: + if device is None: + dev_idx = torch.cuda.current_device() + elif isinstance(device, int): + dev_idx = device + else: + dev_idx = ( + device.index + if device.index is not None + else torch.cuda.current_device() + ) + cap = torch.cuda.get_device_capability(dev_idx) + key = (cap, dot_precision) + if key in _ARCH_OVERRIDES: + return _ARCH_OVERRIDES[key] + return _auto_config(dot_precision, cap, shape_hint) + + +# ════════════════════════════════════════════════════════════════ +# Phase A — split into KV and Z kernels +# ════════════════════════════════════════════════════════════════ + + +@triton.jit +def _phase_a_kv_kernel( + qkv_ptr, + stride_b: tl.constexpr, + stride_n: tl.constexpr, + stride_3: tl.constexpr, + stride_h: tl.constexpr, + stride_d: tl.constexpr, + beta_ptr, + k_inv_rms_ptr, + k_norm_w_ptr, + rope_cos_ptr, + rope_sin_ptr, + I_minus_P_kv_ptr, # output: (I - K_rot^T diag(β) K_rot) + A_ptr, # output: K_rot^T diag(β) V + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + K_SCALE, + NORM_EPS: tl.constexpr, + DOT_PRECISION: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, + SKIP_RELU: tl.constexpr = False, +): + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid = tl.program_id(0) + pid_b = pid // (H * F) + pid_hf = pid % (H * F) + pid_h = pid_hf // F + pid_f = pid_hf % F + bh = pid_b * H + pid_h + N: tl.constexpr = F * S + + qkv_bh = qkv_ptr + pid_b * stride_b + pid_h * stride_h + beta_bhf = beta_ptr + bh * (F * S) + pid_f * S + I_P_kv_bhf = ( + I_minus_P_kv_ptr + bh * F * BLOCK_D * BLOCK_D + pid_f * BLOCK_D * BLOCK_D + ) + A_bhf = A_ptr + bh * F * BLOCK_D * BLOCK_D + pid_f * BLOCK_D * BLOCK_D + + offs_d = tl.arange(0, BLOCK_D) + mask_d = offs_d < D + offs_d_pair = offs_d ^ 1 + mask_d_pair = offs_d_pair < D + + nw_offset = pid_h * D + k_nw = tl.load(k_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to( + tl.float32 + ) + k_nw_pair = tl.load( + k_norm_w_ptr + nw_offset + offs_d_pair, mask=mask_d_pair, other=0.0 + ).to(tl.float32) + + # fp32 accumulators avoid bf16 round-off compounding across the loop. + P_kv_acc = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + A_acc = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + + k_scale = K_SCALE + n_base = pid_f * S + + for s0 in range(0, S, BLOCK_S): + offs_s = s0 + tl.arange(0, BLOCK_S) + mask_s = offs_s < S + mask_sd = mask_s[:, None] & mask_d[None, :] + n_idx = n_base + offs_s + + k_ptrs = ( + qkv_bh + + n_idx[:, None] * stride_n + + 1 * stride_3 + + offs_d[None, :] * stride_d + ) + v_ptrs = ( + qkv_bh + + n_idx[:, None] * stride_n + + 2 * stride_3 + + offs_d[None, :] * stride_d + ) + K_raw = tl.load(k_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + V_raw = tl.load(v_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + beta_t = tl.load(beta_bhf + offs_s, mask=mask_s, other=0.0).to(tl.float32) + + k_inv_rms = tl.load( + k_inv_rms_ptr + pid_b * N + n_idx, mask=mask_s, other=1.0 + ).to(tl.float32) + K_normed = K_raw * k_inv_rms[:, None] * k_nw[None, :] + if SKIP_RELU: + K = K_normed * k_scale + else: + K = tl.where(K_normed > 0, K_normed, 0.0) * k_scale + + K_pair_raw = tl.reshape( + tl.flip(tl.reshape(K_raw, (BLOCK_S, BLOCK_D // 2, 2)), dim=2), + (BLOCK_S, BLOCK_D), + ) + K_pair_normed = K_pair_raw * k_inv_rms[:, None] * k_nw_pair[None, :] + if SKIP_RELU: + K_pair = K_pair_normed * k_scale + else: + K_pair = tl.where(K_pair_normed > 0, K_pair_normed, 0.0) * k_scale + + rope_ptrs = n_idx[:, None] * D + offs_d[None, :] + Cos = tl.load(rope_cos_ptr + rope_ptrs, mask=mask_sd, other=1.0).to(tl.float32) + Sin = tl.load(rope_sin_ptr + rope_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + K_rot = K * Cos + K_pair * Sin + + beta_Krot = beta_t[:, None] * K_rot + beta_V = beta_t[:, None] * V_raw + + K_rot_T = tl.trans(K_rot) + P_kv_acc += tl.dot( + K_rot_T.to(dot_dtype), + beta_Krot.to(dot_dtype), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + A_acc += tl.dot( + K_rot_T.to(dot_dtype), + beta_V.to(dot_dtype), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + + # Padded positions are 0 by construction (K_rot is 0 outside D). + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + diag_in_range = ( + (offs_d[:, None] == offs_d[None, :]) & mask_d[:, None] & mask_d[None, :] + ) + I_minus_P_kv = tl.where(diag_in_range, 1.0 - P_kv_acc, -P_kv_acc) + if DOT_PRECISION >= 1: + tl.store(I_P_kv_bhf + offs_dd, I_minus_P_kv) + tl.store(A_bhf + offs_dd, A_acc) + else: + tl.store(I_P_kv_bhf + offs_dd, I_minus_P_kv.to(tl.bfloat16)) + tl.store(A_bhf + offs_dd, A_acc.to(tl.bfloat16)) + + +@triton.jit +def _phase_a_z_kernel( + qkv_ptr, + stride_b: tl.constexpr, + stride_n: tl.constexpr, + stride_3: tl.constexpr, + stride_h: tl.constexpr, + stride_d: tl.constexpr, + beta_ptr, + k_inv_rms_ptr, + k_norm_w_ptr, + I_minus_P_z_ptr, # output: (I - K^T diag(β) K) + B_ptr, # output: K^T β + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + K_SCALE, + NORM_EPS: tl.constexpr, + DOT_PRECISION: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Z stream: uses K (no RoPE). Cheaper than KV — no V load, no RoPE, no K_pair.""" + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid = tl.program_id(0) + pid_b = pid // (H * F) + pid_hf = pid % (H * F) + pid_h = pid_hf // F + pid_f = pid_hf % F + bh = pid_b * H + pid_h + N: tl.constexpr = F * S + + qkv_bh = qkv_ptr + pid_b * stride_b + pid_h * stride_h + beta_bhf = beta_ptr + bh * (F * S) + pid_f * S + I_P_z_bhf = I_minus_P_z_ptr + bh * F * BLOCK_D * BLOCK_D + pid_f * BLOCK_D * BLOCK_D + B_bhf = B_ptr + bh * F * BLOCK_D + pid_f * BLOCK_D + + offs_d = tl.arange(0, BLOCK_D) + mask_d = offs_d < D + + nw_offset = pid_h * D + k_nw = tl.load(k_norm_w_ptr + nw_offset + offs_d, mask=mask_d, other=0.0).to( + tl.float32 + ) + + P_z_acc = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + B_acc = tl.zeros([BLOCK_D], dtype=tl.float32) + + k_scale = K_SCALE + n_base = pid_f * S + + for s0 in range(0, S, BLOCK_S): + offs_s = s0 + tl.arange(0, BLOCK_S) + mask_s = offs_s < S + mask_sd = mask_s[:, None] & mask_d[None, :] + n_idx = n_base + offs_s + + k_ptrs = ( + qkv_bh + + n_idx[:, None] * stride_n + + 1 * stride_3 + + offs_d[None, :] * stride_d + ) + K_raw = tl.load(k_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + beta_t = tl.load(beta_bhf + offs_s, mask=mask_s, other=0.0).to(tl.float32) + + k_inv_rms = tl.load( + k_inv_rms_ptr + pid_b * N + n_idx, mask=mask_s, other=1.0 + ).to(tl.float32) + K_normed = K_raw * k_inv_rms[:, None] * k_nw[None, :] + K = tl.where(K_normed > 0, K_normed, 0.0) * k_scale + + beta_K = beta_t[:, None] * K + + K_T = tl.trans(K) + P_z_acc += tl.dot( + K_T.to(dot_dtype), + beta_K.to(dot_dtype), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + B_acc += tl.sum(beta_K, axis=0) + + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + diag_in_range = ( + (offs_d[:, None] == offs_d[None, :]) & mask_d[:, None] & mask_d[None, :] + ) + I_minus_P_z = tl.where(diag_in_range, 1.0 - P_z_acc, -P_z_acc) + + if DOT_PRECISION >= 1: + tl.store(I_P_z_bhf + offs_dd, I_minus_P_z) + else: + tl.store(I_P_z_bhf + offs_dd, I_minus_P_z.to(tl.bfloat16)) + # B stays fp32 (vector, ~0.5 KB, negligible HBM cost). + tl.store(B_bhf + offs_d, B_acc) + + +def phase_a( + qkv: torch.Tensor, + beta: torch.Tensor, + q_inv_rms: torch.Tensor, + k_inv_rms: torch.Tensor, + q_norm_w: torch.Tensor, + k_norm_w: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + F: int, + S: int, + k_scale: float = 1.0, + norm_eps: float = 1e-5, + num_warps: int | None = None, + num_stages: int = 1, + BLOCK_S: int | None = None, + dot_precision: int = 0, + skip_relu: bool = False, + skip_z: bool = False, +): + """Compute (I-P_kv), A, (I-P_z), B for all (B, H, F) via 2 kernels (KV + Z). + + `skip_relu=True`: pure linear K-stream prep (no ReLU). Used by the camera + branch where K is already ReLU'd then rotated by UCPE+RoPE — re-applying ReLU + on rotated values would clobber legitimate negatives. + + `skip_z=True`: skip the Z kernel, return placeholder I_P_z/B_z. Used by + NUM_ONLY (camera) callers that never consume the denominator scan. + """ + if num_warps is None or BLOCK_S is None: + a_w, a_bs, *_ = _get_arch_config(dot_precision, device=qkv.device) + if num_warps is None: + num_warps = a_w + if BLOCK_S is None: + BLOCK_S = a_bs + B, N, three, H, D = qkv.shape + assert three == 3 and N == F * S + BLOCK_D = triton.next_power_of_2(D) + BH = B * H + + # FAIR-COMPARE PATCH: keep fp32 inter-phase bridge at P0/P1 to match pytorch/fused + bridge_dtype = torch.float32 if dot_precision >= 1 else torch.bfloat16 + I_P_kv = torch.empty(BH, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=bridge_dtype) + A = torch.empty(BH, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=bridge_dtype) + + beta_c = beta.contiguous() + grid = (BH * F,) + + _phase_a_kv_kernel[grid]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + beta_c, + k_inv_rms, + k_norm_w, + rope_cos, + rope_sin, + I_P_kv, + A, + H=H, + F=F, + S=S, + D=D, + K_SCALE=k_scale, + NORM_EPS=norm_eps, + DOT_PRECISION=dot_precision, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + SKIP_RELU=skip_relu, + num_warps=num_warps, + num_stages=num_stages, + ) + + if skip_z: + # NUM_ONLY (camera) callers don't consume the Z scan; placeholders let + # Phase B skip all Z loads/stores too. + I_P_z = torch.empty(1, device=qkv.device, dtype=bridge_dtype) + B_z = torch.empty(1, device=qkv.device, dtype=torch.float32) + return I_P_kv, A, I_P_z, B_z + + I_P_z = torch.empty(BH, F, BLOCK_D, BLOCK_D, device=qkv.device, dtype=bridge_dtype) + # B stays fp32 — small vector (~0.5 KB/frame), no benefit to downcast. + B_z = torch.empty(BH, F, BLOCK_D, device=qkv.device, dtype=torch.float32) + + _phase_a_z_kernel[grid]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + beta_c, + k_inv_rms, + k_norm_w, + I_P_z, + B_z, + H=H, + F=F, + S=S, + D=D, + K_SCALE=k_scale, + NORM_EPS=norm_eps, + DOT_PRECISION=dot_precision, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + num_warps=num_warps, + num_stages=num_stages, + ) + return I_P_kv, A, I_P_z, B_z + + +# ════════════════════════════════════════════════════════════════ +# Phase B — serial scan, uses pre-stored (I - P) so MMA folds in M +# ════════════════════════════════════════════════════════════════ + + +@triton.jit +def _phase_b_kernel( + I_P_kv_ptr, + A_ptr, + I_P_z_ptr, + B_ptr, + decay_ptr, + M_fwd_ptr, + z_fwd_ptr, + M_rev_ptr, + z_rev_ptr, + init_state_kv_ptr, # (BH, BLOCK_D, BLOCK_D) — read when LOAD_INIT_STATE=1 + init_state_z_ptr, # (BH, BLOCK_D) + final_state_kv_ptr, # (BH, BLOCK_D, BLOCK_D) — written when SAVE_FINAL_STATE=1 + final_state_z_ptr, # (BH, BLOCK_D) + BH: tl.constexpr, + F: tl.constexpr, + BLOCK_D: tl.constexpr, + DOT_PRECISION: tl.constexpr, + USE_ACC_FUSION: tl.constexpr, + LOAD_INIT_STATE: tl.constexpr, # forward scan seeded with init state (vs zeros) + SAVE_FINAL_STATE: tl.constexpr, # write M_{F-1} of forward scan to final_state_* + DIRECTION: tl.constexpr, # 0=both, 1=fwd-only, 2=rev-only + COMBINED_HISTORY: tl.constexpr, # rev read-add-stores into M_fwd_ptr → M_hist[f] + # = M_fwd[f] + M_rev[f]; skips the F-1 zero-write (rev value there is zero). + # DIRECTION=0 only. Saves one Phase C launch + one M-shaped buffer downstream. + SKIP_Z: tl.constexpr, +): + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid = tl.program_id(0) + bh = pid + + offs_d = tl.arange(0, BLOCK_D) + offs_dd = offs_d[:, None] * BLOCK_D + offs_d[None, :] + + # ── Forward scan (skip when DIRECTION=2 i.e. rev-only) ── + if DIRECTION != 2: + if LOAD_INIT_STATE: + M = tl.load(init_state_kv_ptr + bh * BLOCK_D * BLOCK_D + offs_dd).to( + tl.float32 + ) + if not SKIP_Z: + z = tl.load(init_state_z_ptr + bh * BLOCK_D + offs_d).to(tl.float32) + else: + M = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + if not SKIP_Z: + z = tl.zeros([BLOCK_D], dtype=tl.float32) + for f in range(F): + I_P_kv_f = tl.load( + I_P_kv_ptr + + bh * F * BLOCK_D * BLOCK_D + + f * BLOCK_D * BLOCK_D + + offs_dd + ) + A_f = tl.load( + A_ptr + bh * F * BLOCK_D * BLOCK_D + f * BLOCK_D * BLOCK_D + offs_dd + ) + g_f = tl.load(decay_ptr + bh * F + f).to(tl.float32) + + # M = g · (I - P_kv) M + A_f + if USE_ACC_FUSION: + # Pre-scale (I-P) by g and fold A_f via the MMA accumulator → + # A_f + g·(I-P)·M in one MMA, no separate M_temp tensor. + I_P_scaled = I_P_kv_f.to(tl.float32) * g_f + M = tl.dot( + I_P_scaled.to(dot_dtype), + M.to(dot_dtype), + acc=A_f.to(tl.float32), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + else: + M_temp = tl.dot( + I_P_kv_f.to(dot_dtype), + M.to(dot_dtype), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + M = g_f * M_temp + A_f + + tl.store( + M_fwd_ptr + + bh * F * BLOCK_D * BLOCK_D + + f * BLOCK_D * BLOCK_D + + offs_dd, + M, + ) + if not SKIP_Z: + I_P_z_f = tl.load( + I_P_z_ptr + + bh * F * BLOCK_D * BLOCK_D + + f * BLOCK_D * BLOCK_D + + offs_dd + ) + B_f = tl.load(B_ptr + bh * F * BLOCK_D + f * BLOCK_D + offs_d) + # z = g · (I - P_z) z + B_f + z_temp = tl.sum(I_P_z_f * z[None, :], axis=1) + z = g_f * z_temp + B_f + tl.store(z_fwd_ptr + bh * F * BLOCK_D + f * BLOCK_D + offs_d, z) + + # Save terminal forward state for state-cached inference (autoregressive). + if SAVE_FINAL_STATE: + tl.store(final_state_kv_ptr + bh * BLOCK_D * BLOCK_D + offs_dd, M) + if not SKIP_Z: + tl.store(final_state_z_ptr + bh * BLOCK_D + offs_d, z) + + # ── Reverse scan (skip when DIRECTION=1 i.e. fwd-only) ── + if DIRECTION != 1: + M = tl.zeros([BLOCK_D, BLOCK_D], dtype=tl.float32) + if not SKIP_Z: + z = tl.zeros([BLOCK_D], dtype=tl.float32) + # COMBINED_HISTORY: skip the F-1 zero-write so M_hist[F-1] keeps the fwd + # value (rev value there is zero by construction). + if not COMBINED_HISTORY: + tl.store( + M_rev_ptr + + bh * F * BLOCK_D * BLOCK_D + + (F - 1) * BLOCK_D * BLOCK_D + + offs_dd, + M, + ) + if not SKIP_Z: + tl.store(z_rev_ptr + bh * F * BLOCK_D + (F - 1) * BLOCK_D + offs_d, z) + for f_iter in range(F - 1): + f_src = F - 1 - f_iter + f_dst = f_src - 1 + I_P_kv_f = tl.load( + I_P_kv_ptr + + bh * F * BLOCK_D * BLOCK_D + + f_src * BLOCK_D * BLOCK_D + + offs_dd + ) + A_f = tl.load( + A_ptr + bh * F * BLOCK_D * BLOCK_D + f_src * BLOCK_D * BLOCK_D + offs_dd + ) + g_f = tl.load(decay_ptr + bh * F + f_src).to(tl.float32) + + if USE_ACC_FUSION: + I_P_scaled = I_P_kv_f.to(tl.float32) * g_f + M = tl.dot( + I_P_scaled.to(dot_dtype), + M.to(dot_dtype), + acc=A_f.to(tl.float32), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + else: + M_temp = tl.dot( + I_P_kv_f.to(dot_dtype), + M.to(dot_dtype), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + M = g_f * M_temp + A_f + + if not SKIP_Z: + I_P_z_f = tl.load( + I_P_z_ptr + + bh * F * BLOCK_D * BLOCK_D + + f_src * BLOCK_D * BLOCK_D + + offs_dd + ) + B_f = tl.load(B_ptr + bh * F * BLOCK_D + f_src * BLOCK_D + offs_d) + z_temp = tl.sum(I_P_z_f * z[None, :], axis=1) + z = g_f * z_temp + B_f + + if COMBINED_HISTORY: + # Read-add-store the rev contribution into the fwd buffer slot + # (fwd just wrote M_fwd[f_dst]; stays in L1/L2). + M_addr = ( + M_fwd_ptr + + bh * F * BLOCK_D * BLOCK_D + + f_dst * BLOCK_D * BLOCK_D + + offs_dd + ) + tl.store(M_addr, tl.load(M_addr) + M) + if not SKIP_Z: + z_addr = z_fwd_ptr + bh * F * BLOCK_D + f_dst * BLOCK_D + offs_d + tl.store(z_addr, tl.load(z_addr) + z) + else: + tl.store( + M_rev_ptr + + bh * F * BLOCK_D * BLOCK_D + + f_dst * BLOCK_D * BLOCK_D + + offs_dd, + M, + ) + if not SKIP_Z: + tl.store(z_rev_ptr + bh * F * BLOCK_D + f_dst * BLOCK_D + offs_d, z) + + +def phase_b_triton( + I_P_kv, + A, + I_P_z, + B, + decay, + F, + num_warps=None, + num_stages=None, + use_acc_fusion=None, + dot_precision=0, + init_state_kv=None, + init_state_z=None, + return_final_state=False, + direction=0, + combined_history=False, + skip_z=False, +): + """Phase B serial-F scan over (B*H,). + + Forward scan can be seeded with `init_state_kv`/`init_state_z` (autoregressive + chunk > 0) and write terminal `M_{F-1}`/`z_{F-1}` when `return_final_state=True`. + + `direction`: 0=both, 1=forward-only, 2=reverse-only (single-direction state + cache). `combined_history` (direction=0 only): rev read-add-stores into the + fwd buffer so it becomes M_hist[f]=M_fwd[f]+M_rev[f], letting Phase C run once + (it's linear: `Q @ (M_fwd + M_rev) = Q @ M_fwd + Q @ M_rev`); M_rev/z_rev are + then placeholders. `skip_z`: skip the Z/denominator recurrence (camera + num-only scans, Phase C with num_only=True). + + Returns (M_fwd, z_fwd, M_rev, z_rev) (+ (final_kv, final_z) when + return_final_state). Skipped-direction outputs are 1-element placeholders the + kernel never touches — callers must discard the slot they didn't ask for. + Reverse scan always seeds from zero (upstream bidi convention: only forward + state is cached). + """ + BH = I_P_kv.shape[0] + _, _, BLOCK_D, _ = A.shape # A is always full [BH, F, BLOCK_D, BLOCK_D] + device, fdtype = I_P_kv.device, torch.float32 + + if num_warps is None or num_stages is None or use_acc_fusion is None: + _, _, b_w, b_s, b_acc, *_ = _get_arch_config(dot_precision, device=device) + if num_warps is None: + num_warps = b_w + if num_stages is None: + num_stages = b_s + if use_acc_fusion is None: + use_acc_fusion = b_acc + + if combined_history and direction != 0: + raise ValueError("combined_history=True requires direction=0 (bidi)") + + # Kernel is DIRECTION-gated (constexpr), so inactive buffers can be 1-element + # placeholders — frees ~4× M_fwd-shaped allocs per single-direction call. + decay_flat = decay.reshape(BH, F).contiguous().float() + + load_init = init_state_kv is not None + dummy = torch.empty(1, device=device, dtype=fdtype) + full_M = lambda: torch.empty(BH, F, BLOCK_D, BLOCK_D, device=device, dtype=fdtype) + full_z = lambda: torch.empty(BH, F, BLOCK_D, device=device, dtype=fdtype) + M_fwd = dummy if direction == 2 else full_M() + z_fwd = dummy if (direction == 2 or skip_z) else full_z() + # Combined-history reuses M_fwd/z_fwd as M_hist/z_hist; rev outputs are + # placeholders even though DIRECTION!=1. + M_rev = dummy if (direction == 1 or combined_history) else full_M() + z_rev = dummy if (direction == 1 or combined_history or skip_z) else full_z() + if load_init: + init_kv = init_state_kv.contiguous().view(BH, BLOCK_D, BLOCK_D) + init_z = dummy if skip_z else init_state_z.contiguous().view(BH, BLOCK_D) + else: + init_kv = dummy + init_z = dummy + + if return_final_state: + final_kv = torch.empty(BH, BLOCK_D, BLOCK_D, device=device, dtype=fdtype) + final_z = ( + dummy if skip_z else torch.empty(BH, BLOCK_D, device=device, dtype=fdtype) + ) + else: + final_kv = dummy + final_z = dummy + + d_splits, nw_override, ns_override, acc_override = _pick_phase_b_d_splits( + BLOCK_D, dot_precision=dot_precision + ) + if d_splits > 1: + D_TILE = BLOCK_D // d_splits + # D-tile-specific tuning if available, else baseline. + nw_use = nw_override if nw_override is not None else num_warps + ns_use = ns_override if ns_override is not None else num_stages + acc_use = acc_override if acc_override is not None else use_acc_fusion + _phase_b_dtile_kernel[(BH, d_splits)]( + I_P_kv, + A, + I_P_z, + B, + decay_flat, + M_fwd, + z_fwd, + M_rev, + z_rev, + init_kv, + init_z, + final_kv, + final_z, + BH=BH, + F=F, + BLOCK_D=BLOCK_D, + D_TILE=D_TILE, + DOT_PRECISION=dot_precision, + USE_ACC_FUSION=acc_use, + LOAD_INIT_STATE=1 if load_init else 0, + SAVE_FINAL_STATE=1 if return_final_state else 0, + DIRECTION=direction, + COMBINED_HISTORY=1 if combined_history else 0, + SKIP_Z=1 if skip_z else 0, + num_warps=nw_use, + num_stages=ns_use, + ) + else: + _phase_b_kernel[(BH,)]( + I_P_kv, + A, + I_P_z, + B, + decay_flat, + M_fwd, + z_fwd, + M_rev, + z_rev, + init_kv, + init_z, + final_kv, + final_z, + BH=BH, + F=F, + BLOCK_D=BLOCK_D, + DOT_PRECISION=dot_precision, + USE_ACC_FUSION=use_acc_fusion, + LOAD_INIT_STATE=1 if load_init else 0, + SAVE_FINAL_STATE=1 if return_final_state else 0, + DIRECTION=direction, + COMBINED_HISTORY=1 if combined_history else 0, + SKIP_Z=1 if skip_z else 0, + num_warps=num_warps, + num_stages=num_stages, + ) + if return_final_state: + return M_fwd, z_fwd, M_rev, z_rev, final_kv, final_z + return M_fwd, z_fwd, M_rev, z_rev + + +# Phase B D-tile — j-axis split for grid parallelism (#118). +# Same recurrence as _phase_b_kernel but each program owns a D_TILE-wide slice of +# M's output column dim. Grid: (BH, d_splits). M_new[*, j_tile] depends only on +# M_prev[*, j_tile] and full (I-P_kv) — independent across j-tiles. z is +# unsplittable; only `pid_d == 0` updates/writes z. +@triton.jit +def _phase_b_dtile_kernel( + I_P_kv_ptr, + A_ptr, + I_P_z_ptr, + B_ptr, + decay_ptr, + M_fwd_ptr, + z_fwd_ptr, + M_rev_ptr, + z_rev_ptr, + init_state_kv_ptr, + init_state_z_ptr, + final_state_kv_ptr, + final_state_z_ptr, + BH: tl.constexpr, + F: tl.constexpr, + BLOCK_D: tl.constexpr, + D_TILE: tl.constexpr, + DOT_PRECISION: tl.constexpr, + USE_ACC_FUSION: tl.constexpr, + LOAD_INIT_STATE: tl.constexpr, + SAVE_FINAL_STATE: tl.constexpr, + DIRECTION: tl.constexpr, + COMBINED_HISTORY: tl.constexpr, + SKIP_Z: tl.constexpr, +): + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid_bh = tl.program_id(0) + pid_d = tl.program_id(1) + bh = pid_bh + + offs_d_full = tl.arange(0, BLOCK_D) + offs_d_tile = pid_d * D_TILE + tl.arange(0, D_TILE) + offs_dd_full = offs_d_full[:, None] * BLOCK_D + offs_d_full[None, :] + offs_dd_tile = offs_d_full[:, None] * BLOCK_D + offs_d_tile[None, :] + + is_lead = pid_d == 0 + + if DIRECTION != 2: + if LOAD_INIT_STATE: + M = tl.load(init_state_kv_ptr + bh * BLOCK_D * BLOCK_D + offs_dd_tile).to( + tl.float32 + ) + else: + M = tl.zeros([BLOCK_D, D_TILE], dtype=tl.float32) + if not SKIP_Z: + z = tl.zeros([BLOCK_D], dtype=tl.float32) + if is_lead and LOAD_INIT_STATE: + z = tl.load(init_state_z_ptr + bh * BLOCK_D + offs_d_full).to( + tl.float32 + ) + + for f in range(F): + I_P_kv_f = tl.load( + I_P_kv_ptr + + bh * F * BLOCK_D * BLOCK_D + + f * BLOCK_D * BLOCK_D + + offs_dd_full + ) + A_f = tl.load( + A_ptr + + bh * F * BLOCK_D * BLOCK_D + + f * BLOCK_D * BLOCK_D + + offs_dd_tile + ) + g_f = tl.load(decay_ptr + bh * F + f).to(tl.float32) + + if USE_ACC_FUSION: + I_P_scaled = I_P_kv_f.to(tl.float32) * g_f + M = tl.dot( + I_P_scaled.to(dot_dtype), + M.to(dot_dtype), + acc=A_f.to(tl.float32), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + else: + M_temp = tl.dot( + I_P_kv_f.to(dot_dtype), + M.to(dot_dtype), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + M = g_f * M_temp + A_f + + tl.store( + M_fwd_ptr + + bh * F * BLOCK_D * BLOCK_D + + f * BLOCK_D * BLOCK_D + + offs_dd_tile, + M, + ) + + if is_lead and not SKIP_Z: + I_P_z_f = tl.load( + I_P_z_ptr + + bh * F * BLOCK_D * BLOCK_D + + f * BLOCK_D * BLOCK_D + + offs_dd_full + ) + B_f = tl.load(B_ptr + bh * F * BLOCK_D + f * BLOCK_D + offs_d_full) + z_temp = tl.sum(I_P_z_f * z[None, :], axis=1) + z = g_f * z_temp + B_f + tl.store(z_fwd_ptr + bh * F * BLOCK_D + f * BLOCK_D + offs_d_full, z) + + if SAVE_FINAL_STATE: + tl.store(final_state_kv_ptr + bh * BLOCK_D * BLOCK_D + offs_dd_tile, M) + if is_lead and not SKIP_Z: + tl.store(final_state_z_ptr + bh * BLOCK_D + offs_d_full, z) + + if DIRECTION != 1: + M = tl.zeros([BLOCK_D, D_TILE], dtype=tl.float32) + if not SKIP_Z: + z = tl.zeros([BLOCK_D], dtype=tl.float32) + + if not COMBINED_HISTORY: + tl.store( + M_rev_ptr + + bh * F * BLOCK_D * BLOCK_D + + (F - 1) * BLOCK_D * BLOCK_D + + offs_dd_tile, + M, + ) + if is_lead and not SKIP_Z: + tl.store( + z_rev_ptr + bh * F * BLOCK_D + (F - 1) * BLOCK_D + offs_d_full, z + ) + + for f_iter in range(F - 1): + f_src = F - 1 - f_iter + f_dst = f_src - 1 + I_P_kv_f = tl.load( + I_P_kv_ptr + + bh * F * BLOCK_D * BLOCK_D + + f_src * BLOCK_D * BLOCK_D + + offs_dd_full + ) + A_f = tl.load( + A_ptr + + bh * F * BLOCK_D * BLOCK_D + + f_src * BLOCK_D * BLOCK_D + + offs_dd_tile + ) + g_f = tl.load(decay_ptr + bh * F + f_src).to(tl.float32) + + if USE_ACC_FUSION: + I_P_scaled = I_P_kv_f.to(tl.float32) * g_f + M = tl.dot( + I_P_scaled.to(dot_dtype), + M.to(dot_dtype), + acc=A_f.to(tl.float32), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + else: + M_temp = tl.dot( + I_P_kv_f.to(dot_dtype), + M.to(dot_dtype), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + M = g_f * M_temp + A_f + + if is_lead and not SKIP_Z: + I_P_z_f = tl.load( + I_P_z_ptr + + bh * F * BLOCK_D * BLOCK_D + + f_src * BLOCK_D * BLOCK_D + + offs_dd_full + ) + B_f = tl.load(B_ptr + bh * F * BLOCK_D + f_src * BLOCK_D + offs_d_full) + z_temp = tl.sum(I_P_z_f * z[None, :], axis=1) + z = g_f * z_temp + B_f + + if COMBINED_HISTORY: + M_addr = ( + M_fwd_ptr + + bh * F * BLOCK_D * BLOCK_D + + f_dst * BLOCK_D * BLOCK_D + + offs_dd_tile + ) + tl.store(M_addr, tl.load(M_addr) + M) + if is_lead and not SKIP_Z: + z_addr = ( + z_fwd_ptr + bh * F * BLOCK_D + f_dst * BLOCK_D + offs_d_full + ) + tl.store(z_addr, tl.load(z_addr) + z) + else: + tl.store( + M_rev_ptr + + bh * F * BLOCK_D * BLOCK_D + + f_dst * BLOCK_D * BLOCK_D + + offs_dd_tile, + M, + ) + if is_lead and not SKIP_Z: + tl.store( + z_rev_ptr + bh * F * BLOCK_D + f_dst * BLOCK_D + offs_d_full, z + ) + + +_PHASE_B_DTILE_ARCH_CACHE: dict = {} # (dev, dot_prec) -> (d_splits, nw, ns, acc) + + +def _pick_phase_b_d_splits(BLOCK_D: int, dot_precision: int = 0): + """Returns (d_splits, nw_override, ns_override, acc_override). + + `d_splits=1` → baseline `_phase_b_kernel` with `_CHUNKWISE_TUNING` config. + `d_splits>1` → `_phase_b_dtile_kernel` with nw/ns/acc overrides. + Per-arch optima from a 96-config sweep (2026-04-29, T=11 B=1 IEEE) and a + multi-arch bf16/TF32 sweep (2026-05-06, F=11 S=920); see inline cfg comments. + Env overrides: PHASE_B_D_SPLITS, PHASE_B_DTILE_NW/NS/ACC (1=True / 0=False). + """ + import os + + env_d = os.environ.get("PHASE_B_D_SPLITS", None) + if env_d is not None: + d = int(env_d) + if d < 1 or BLOCK_D % d != 0: + return (1, None, None, None) + nw = int(os.environ.get("PHASE_B_DTILE_NW", "0")) or None + ns = int(os.environ.get("PHASE_B_DTILE_NS", "0")) or None + acc_env = os.environ.get("PHASE_B_DTILE_ACC", None) + acc = bool(int(acc_env)) if acc_env is not None else None + return (d, nw, ns, acc) + try: + import torch + + if not torch.cuda.is_available(): + return (1, None, None, None) + dev = torch.cuda.current_device() + cache_key = (dev, dot_precision) + if cache_key not in _PHASE_B_DTILE_ARCH_CACHE: + cap = torch.cuda.get_device_capability(dev) + major, minor = cap[0], cap[1] + if dot_precision == 2: + # IEEE fp32: D-tile dominates baseline on every arch. + if major == 8 and minor == 0: + cfg = (4, 32, 1, True) # A100 + elif major == 9: + cfg = (4, 32, 1, True) # H100 (Hopper) + elif major == 8 and minor == 9: + cfg = (8, 4, 1, False) # Ada (assume Blackwell-like) + elif major >= 10: + cfg = (8, 4, 1, False) # GB200/B200, 5090, GB10 + else: + cfg = (1, None, None, None) # unknown — baseline + else: + # bf16/TF32: cap-specific dispatch (per-cap D-tile-vs-baseline win): + # sm_80 A100 / sm_90 H100 / sm_100 GB200: D-tile WIN ~10-12% — (4,8,2,F). + # sm_120 5090: D-tile WIN 2.6×(P1)/1.13×(P2) — (8,8,1,F); TF32 baseline OOMs. + # sm_121 GB10: baseline WINS (4% over D-tile). Despite matching sm_120 SRAM, + # the baseline fits all configs up to nw=16 ns=2 here (Triton/codegen + # difference between consumer-Blackwell variants) and saturates the chip. + if major == 8 and minor == 0: + cfg = (4, 8, 2, False) # A100 + elif major == 9: + cfg = (4, 8, 2, False) # H100 + elif major == 10: + cfg = (4, 8, 2, False) # GB200 / B200 + elif major == 12 and minor == 0: + cfg = (8, 8, 1, False) # 5090 + elif major == 12 and minor == 1: + cfg = (1, None, None, None) # GB10 — baseline wins + else: + cfg = (1, None, None, None) # Ada, unknown + _PHASE_B_DTILE_ARCH_CACHE[cache_key] = cfg + return _PHASE_B_DTILE_ARCH_CACHE[cache_key] + except Exception: + return (1, None, None, None) + + +# Phase C — Pass 2 output (per (B, H, F)). Same as v1. + + +@triton.jit +def _phase_c_kernel( + qkv_ptr, + stride_b: tl.constexpr, + stride_n: tl.constexpr, + stride_3: tl.constexpr, + stride_h: tl.constexpr, + stride_d: tl.constexpr, + q_inv_rms_ptr, + q_norm_w_ptr, + rope_cos_ptr, + rope_sin_ptr, + M_ptr, + z_ptr, + num_ptr, + den_ptr, + H: tl.constexpr, + F: tl.constexpr, + S: tl.constexpr, + D: tl.constexpr, + NORM_EPS: tl.constexpr, + DOT_PRECISION: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, + ACCUMULATE: tl.constexpr = False, + SKIP_LAST_F: tl.constexpr = False, + SKIP_RELU: tl.constexpr = False, + NUM_ONLY: tl.constexpr = False, +): + if DOT_PRECISION >= 1: + dot_dtype = tl.float32 + else: + dot_dtype = tl.bfloat16 + dot_ip: tl.constexpr = "ieee" if DOT_PRECISION == 2 else "tf32" + + pid = tl.program_id(0) + pid_b = pid // (H * F) + pid_hf = pid % (H * F) + pid_h = pid_hf // F + pid_f = pid_hf % F + bh = pid_b * H + pid_h + N: tl.constexpr = F * S + + # SKIP_LAST_F (reverse-accumulate callers): M_rev[F-1]/z_rev[F-1] are exactly + # zero (rev scan inits to zero, write loop fills only f 0, Q_normed, 0.0) + Q_pair = tl.where(Q_pair_normed > 0, Q_pair_normed, 0.0) + + rope_ptrs = n_idx[:, None] * D + offs_d[None, :] + Cos = tl.load(rope_cos_ptr + rope_ptrs, mask=mask_sd, other=1.0).to(tl.float32) + Sin = tl.load(rope_sin_ptr + rope_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + Q_rot = Q * Cos + Q_pair * Sin + + num = tl.dot( + Q_rot.to(dot_dtype), + M_f.to(dot_dtype), + out_dtype=tl.float32, + input_precision=dot_ip, + ) + if not NUM_ONLY: + den = tl.sum(Q * z_f[None, :], axis=1) + + num_ptrs = num_bh + n_idx[:, None] * (H * D) + offs_d[None, :] + if not NUM_ONLY: + den_ptrs = den_bh + n_idx + if ACCUMULATE: + # Reverse-direction Phase C: add onto forward's already-written buffer + # instead of allocating a separate one. + prev_num = tl.load(num_ptrs, mask=mask_sd, other=0.0).to(tl.float32) + num = num + prev_num + if not NUM_ONLY: + prev_den = tl.load(den_ptrs, mask=mask_s, other=0.0).to(tl.float32) + den = den + prev_den + if DOT_PRECISION >= 1: + tl.store(num_ptrs, num, mask=mask_sd) + if not NUM_ONLY: + tl.store(den_ptrs, den, mask=mask_s) + else: + tl.store(num_ptrs, num.to(tl.bfloat16), mask=mask_sd) + if not NUM_ONLY: + tl.store(den_ptrs, den.to(tl.bfloat16), mask=mask_s) + + +def phase_c( + qkv, + q_inv_rms, + q_norm_w, + rope_cos, + rope_sin, + M, + z, + F, + S, + num_warps=None, + num_stages=None, + BLOCK_S=None, + dot_precision=0, + num_out=None, + den_out=None, + accumulate=False, + skip_last_frame=False, + skip_relu: bool = False, + num_only: bool = False, +): + """Phase C Pass-2 output. + + ``accumulate=True``: add into caller's ``num_out``/``den_out`` (fuses + reverse-direction output into the forward buffer — saves ~45 MB at B=1 bf16). + ``skip_last_frame=True``: early-return f=F-1 programs; valid only for the + reverse-accumulate call, where M[F-1]/z[F-1] are guaranteed zero. + ``skip_relu=True``: as in Phase A KV — Q is already ReLU'd then rotated by + UCPE+RoPE; re-applying ReLU on rotated Q would clobber legitimate negatives. + ``num_only=True``: skip denominator entirely (writes only ``num_out``; + ``den_out`` may be None). Camera branch has no Z scan. + """ + if num_warps is None or num_stages is None or BLOCK_S is None: + *_, c_w, c_bs, c_s = _get_arch_config(dot_precision, device=qkv.device) + if num_warps is None: + num_warps = c_w + if num_stages is None: + num_stages = c_s + if BLOCK_S is None: + BLOCK_S = c_bs + B, N, three, H, D = qkv.shape + BLOCK_D = triton.next_power_of_2(D) + if num_out is None: + num_out = torch.empty( + B, + N, + H, + D, + device=qkv.device, + dtype=(torch.float32 if dot_precision >= 1 else torch.bfloat16), + ) + if den_out is None and not num_only: + den_out = torch.empty( + B, + H, + N, + device=qkv.device, + dtype=(torch.float32 if dot_precision >= 1 else torch.bfloat16), + ) + elif num_only and den_out is None: + # 1-element placeholder; kernel guards den loads/stores under NUM_ONLY. + den_out = torch.empty( + 1, + device=qkv.device, + dtype=(torch.float32 if dot_precision >= 1 else torch.bfloat16), + ) + + _phase_c_kernel[(B * H * F,)]( + qkv, + qkv.stride(0), + qkv.stride(1), + qkv.stride(2), + qkv.stride(3), + qkv.stride(4), + q_inv_rms, + q_norm_w, + rope_cos, + rope_sin, + M, + z, + num_out, + den_out, + H=H, + F=F, + S=S, + D=D, + NORM_EPS=1e-5, + DOT_PRECISION=dot_precision, + BLOCK_D=BLOCK_D, + BLOCK_S=BLOCK_S, + ACCUMULATE=1 if accumulate else 0, + SKIP_LAST_F=skip_last_frame, + SKIP_RELU=skip_relu, + NUM_ONLY=num_only, + num_warps=num_warps, + num_stages=num_stages, + ) + return num_out, den_out + + +def fused_bigdn_bidi_chunkwise( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_w, + k_norm_w, + rope_cos, + rope_sin, + beta, + decay, + F, + S, + k_scale=1.0, + eps=1e-6, + norm_eps=1e-5, + dot_precision=0, + init_state_kv=None, + init_state_z=None, + return_final_state=False, +): + """Bidi chunkwise GDN forward, optionally state-cached for autoregressive + sampling (chunk 0 = full bidi + state save; chunks > 0 seed forward from saved + state). Reverse always seeds from zero (upstream convention). + + Pipeline: Phase A once → Phase B direction=0, combined_history=True (rev summed + into fwd buffer in-kernel so M_hist[f]=M_fwd[f]+M_rev[f]) → Phase C ONCE on + M_hist. Exact via Phase C linearity `Q @ (M_fwd + M_rev) = Q @ M_fwd + Q @ M_rev`. + Replaces the prior 2× Phase B + 2× Phase C pattern. + """ + I_P_kv, A, I_P_z, B_z = phase_a( + qkv, + beta, + q_inv_rms, + k_inv_rms, + q_norm_w, + k_norm_w, + rope_cos, + rope_sin, + F=F, + S=S, + k_scale=k_scale, + norm_eps=norm_eps, + dot_precision=dot_precision, + ) + + if return_final_state: + M_hist, z_hist, _, _, final_kv, final_z = phase_b_triton( + I_P_kv, + A, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=0, + init_state_kv=init_state_kv, + init_state_z=init_state_z, + return_final_state=True, + combined_history=True, + ) + else: + M_hist, z_hist, _, _ = phase_b_triton( + I_P_kv, + A, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=0, + init_state_kv=init_state_kv, + init_state_z=init_state_z, + combined_history=True, + ) + num_out, den_out = phase_c( + qkv, + q_inv_rms, + q_norm_w, + rope_cos, + rope_sin, + M_hist, + z_hist, + F=F, + S=S, + dot_precision=dot_precision, + accumulate=False, + ) + del M_hist, z_hist, I_P_kv, A, I_P_z, B_z + + total_den = den_out.float().permute(0, 2, 1).unsqueeze(-1) # (B, N, H, 1) + out = (num_out.float() / (total_den + eps)).to(qkv.dtype) + del num_out, den_out, total_den + if return_final_state: + B = qkv.shape[0] + H = qkv.shape[3] + D = qkv.shape[4] + BLOCK_D = final_kv.shape[1] + state_kv = ( + final_kv.view(B, H, BLOCK_D, BLOCK_D)[:, :, :D, :D] + .transpose(-1, -2) + .contiguous() + ) + state_z = final_z.view(B, H, BLOCK_D)[:, :, :D].unsqueeze(-1).contiguous() + return out, state_kv, state_z + return out + + +def _default_dot_prec() -> int: + try: + from sglang.jit_kernel.diffusion.triton.sana_wm_gdn import ( + _resolve_launch_config, + ) + + _, dot_prec, _, _ = _resolve_launch_config() + return dot_prec + except Exception: + return 0 + + +def _cam_identity_tables( + *, + B: int, + N: int, + H: int, + D: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + device_index = device.index if device.type == "cuda" else None + key = (device.type, device_index, B, N, H * D, D) + cached = _CAM_IDENTITY_CACHE.get(key) + if cached is not None: + return cached + + ones_inv_rms = torch.ones(B, N, device=device, dtype=torch.float32) + ones_nw = torch.ones(H * D, device=device, dtype=torch.float32) + ones_cos = torch.ones(N, D, device=device, dtype=torch.float32) + zeros_sin = torch.zeros(N, D, device=device, dtype=torch.float32) + cached = (ones_inv_rms, ones_nw, ones_cos, zeros_sin) + _CAM_IDENTITY_CACHE[key] = cached + return cached + + +def cam_scan_bidi_chunkwise( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + dot_precision: int | None = None, +) -> torch.Tensor: + """Bidirectional camera scan for SANA-WM's numerator-only branch. + + q, k, v: camera-prepared ``(B, H, D, N)`` fp32; beta: ``(B, H, F, S)`` fp32; + decay: ``(B, H, F)`` fp32. Returns ``(B, H, D, N)`` fp32. + """ + assert ( + q.shape == k.shape == v.shape + ), f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}" + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + assert beta.is_contiguous() and decay.is_contiguous() + assert ( + q.dtype == torch.float32 + ), f"cam_scan_bidi_chunkwise requires fp32 q/k/v, got {q.dtype}" + + B, H, D, N = q.shape + F = beta.shape[2] + assert N % F == 0 + S = N // F + assert beta.shape == (B, H, F, S) + assert decay.shape == (B, H, F) + + if dot_precision is None: + dot_precision = _default_dot_prec() + + qkv = torch.empty(B, N, 3, H, D, device=q.device, dtype=q.dtype) + qkv[:, :, 0].copy_(q.permute(0, 3, 1, 2)) + qkv[:, :, 1].copy_(k.permute(0, 3, 1, 2)) + qkv[:, :, 2].copy_(v.permute(0, 3, 1, 2)) + + ones_inv_rms, ones_nw, ones_cos, zeros_sin = _cam_identity_tables( + B=B, N=N, H=H, D=D, device=q.device + ) + I_P_kv, A, I_P_z, B_z = phase_a( + qkv, + beta, + ones_inv_rms, + ones_inv_rms, + ones_nw, + ones_nw, + ones_cos, + zeros_sin, + F=F, + S=S, + k_scale=1.0, + norm_eps=1e-5, + dot_precision=dot_precision, + skip_relu=True, + skip_z=True, + ) + M_hist, z_hist, _, _ = phase_b_triton( + I_P_kv, + A, + I_P_z, + B_z, + decay, + F=F, + dot_precision=dot_precision, + direction=0, + combined_history=True, + skip_z=True, + ) + num_out, _ = phase_c( + qkv, + ones_inv_rms, + ones_nw, + ones_cos, + zeros_sin, + M_hist, + z_hist, + F=F, + S=S, + dot_precision=dot_precision, + skip_relu=True, + num_only=True, + ) + return num_out.permute(0, 2, 3, 1).contiguous().to(torch.float32) diff --git a/python/sglang/multimodal_gen/apps/realtime_webui/app.js b/python/sglang/multimodal_gen/apps/realtime_webui/app.js index 86a8be0c7..ade87daf6 100644 --- a/python/sglang/multimodal_gen/apps/realtime_webui/app.js +++ b/python/sglang/multimodal_gen/apps/realtime_webui/app.js @@ -431,10 +431,7 @@ async function decodeFrameBatch(header, data) { } function isWorkerDecodableContentType(contentType) { - return ( - isWorkerDecodableRawContentType(contentType) || - isEncodedPreviewContentType(contentType) - ); + return isWorkerDecodableRawContentType(contentType); } function isWorkerDecodableRawContentType(contentType) { @@ -1238,7 +1235,7 @@ async function payloadToArrayBuffer(data) { return data.arrayBuffer(); } -function drawFrame(image) { +function drawFrame(image, { close = true, markRendered = true } = {}) { const sourceWidth = image.width; const sourceHeight = image.height; let drawSource = image; @@ -1258,9 +1255,9 @@ function drawFrame(image) { ctx.imageSmoothingEnabled = true; ctx.imageSmoothingQuality = "high"; ctx.drawImage(drawSource, 0, 0, sourceWidth, sourceHeight); - renderedPreviewFrames += 1; + if (markRendered) renderedPreviewFrames += 1; setPreviewState("live"); - if (!(image instanceof ImageData)) image.close?.(); + if (close && !(image instanceof ImageData)) image.close?.(); } function renderLoop(now) { @@ -1528,11 +1525,11 @@ function receive(data, epoch) { const payload = message.payload; delete message.payload; enqueueDecodeBatch(message, payload, epoch); - setStatus("Live", "live"); + if (!renderedPreviewFrames) setStatus("Receiving", "live"); return; } pendingHeader = message; - if (pendingHeader) setStatus("Live", "live"); + if (pendingHeader && !renderedPreviewFrames) setStatus("Receiving", "live"); return; } const header = pendingHeader; @@ -1557,6 +1554,9 @@ async function decodeAndEnqueueFrameBatch(header, data, epoch) { return; } const now = performance.now(); + if (!renderedPreviewFrames && decodedFrames.length) { + drawFrame(decodedFrames[0].image, { close: false, markRendered: false }); + } // record source frames before preview playback can hold or drop for latency recordDecodedFrameBatch(decodedFrames); const enqueueResult = playbackController.enqueueDecodedFrames(header, decodedFrames, now); diff --git a/python/sglang/multimodal_gen/apps/realtime_webui/index.html b/python/sglang/multimodal_gen/apps/realtime_webui/index.html index f81db4b36..3489af8f3 100644 --- a/python/sglang/multimodal_gen/apps/realtime_webui/index.html +++ b/python/sglang/multimodal_gen/apps/realtime_webui/index.html @@ -163,6 +163,6 @@ - + diff --git a/python/sglang/multimodal_gen/configs/models/dits/sana_wm.py b/python/sglang/multimodal_gen/configs/models/dits/sana_wm.py new file mode 100644 index 000000000..69a7daf90 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/sana_wm.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig +from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_transformer_blocks + + +@dataclass +class SanaWMArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field( + default_factory=lambda: [is_blocks_or_transformer_blocks] + ) + + # --- Core dims (upstream: depth=20, hidden=2240, heads=20, linear_head_dim=112) --- + patch_size: int = 1 + in_channels: int = 128 # LTX-2 VAE latent channels + out_channels: int = 128 + num_layers: int = 20 + + # Patch embedder uses (1, patch_size, patch_size) — temporal patch is always 1. + patch_size_t: int = 1 + + num_attention_heads: int = 20 + attention_head_dim: int = 112 # = linear_head_dim + linear_head_dim: int = 112 + + # --- Cross-attention (text conditioning) --- + # In upstream, cross-attn uses num_heads (=20) with head_dim = hidden/num_heads = 112. + num_cross_attention_heads: int = 20 + cross_attention_head_dim: int = 112 + cross_attention_dim: int = 2240 # query dim used inside MultiHeadCrossAttention + cross_norm: bool = True + + # Gemma-2-2b-it hidden size (input to y_embedder.y_proj). + caption_channels: int = 2304 + model_max_length: int = 300 + y_norm: bool = True + y_norm_scale_factor: float = 0.01 + y_norm_eps: float = 1e-5 + + mlp_ratio: float = 3.0 + qk_norm: bool = True + norm_eps: float = 1e-6 + timestep_norm_scale_factor: float = 1.0 + + # --- Hybrid GDN/Softmax attention --- + # softmax_every_n=4 => blocks where (i+1)%4 == 0 use softmax main branch, + # i.e. block indices {3, 7, 11, 15, 19}. + softmax_every_n: int = 4 + + # --- GDN ShortConvolution params --- + conv_kernel_size: int = 4 + k_conv_only: bool = True + chunk_gdn_chunk_size: int = 21 + update_rule: str = "torch_chunk" # main branch update rule + cam_update_rule: str = "torch_chunk" # camera branch update rule + # main GDN scan backend: "auto" uses the SANA-WM Triton fast path on + # supported CUDA inference runs, otherwise falls back to the torch scan. + gdn_backend: str = "auto" + + # --- Camera conditioning --- + cam_attn_compress: int = 1 # cam_dim == in_dim + init_cam_from_base: bool = True + use_chunk_plucker_post_attn: bool = True + use_chunk_plucker_input: bool = False + chunk_plucker_channels: int = 48 # 8 orig frames × 6D Plücker + chunk_plucker_post_attn_blocks: int = 20 + + chunk_split_strategy: str = "first_chunk_plus_one" + chunk_size: int = 10 + # Upstream currently forwards chunk metadata through the softmax blocks but + # does not apply a chunk-causal mask there. Keep this disabled by default + # for checkpoint-output parity; it can be enabled for experiments. + use_chunked_softmax_attention: bool = False + + # --- Temporal FFN (GLUMBConvTemp) --- + ffn_type: str = "GLUMBConvTemp" + t_kernel_size: int = 3 + mlp_acts: tuple = field(default_factory=lambda: ("silu", "silu", None)) + + # --- Position embedding --- + pos_embed_type: str = "wan_rope" + + # --- VAE coupling (LTX-2) --- + vae_temporal_stride: int = 8 # original-frames per latent frame + vae_spatial_stride: int = 32 # pixels per latent token (per spatial axis) + + sample_size: int = 32 # legacy, unused + guidance_embeds: bool = False + class_dropout_prob: float = 0.0 + + # the released checkpoints store raw upstream parameter names; streaming + # also keeps an unused all-zero pos_embed while the native model uses RoPE + param_names_mapping: dict = field( + default_factory=lambda: { + "^pos_embed$": "", + } + ) + + def __post_init__(self): + super().__post_init__() + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + + +@dataclass +class SanaWMConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=SanaWMArchConfig) + prefix: str = "SanaWM" diff --git a/python/sglang/multimodal_gen/configs/models/dits/sana_wm_refiner.py b/python/sglang/multimodal_gen/configs/models/dits/sana_wm_refiner.py new file mode 100644 index 000000000..99cef259f --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/sana_wm_refiner.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig +from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_transformer_blocks + + +@dataclass +class SanaWMRefinerArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field( + default_factory=lambda: [is_blocks_or_transformer_blocks] + ) + + # Core dims + in_channels: int = 128 + out_channels: int = 128 + patch_size: int = 1 + patch_size_t: int = 1 + num_layers: int = 28 + num_attention_heads: int = 32 + attention_head_dim: int = 64 + cross_attention_dim: int = 4096 + caption_channels: int = 4096 + + qk_norm: bool = True + norm_eps: float = 1e-6 + apply_gated_attention: bool = False + + timestep_scale_multiplier: float = 1000.0 + rope_type: str = "interleaved" + + # RoPE coord generation + sampling_rate: int = 16000 + hop_length: int = 160 + scale_factors: tuple = (8, 32, 32) + base_num_frames: int = 20 + base_height: int = 2048 + base_width: int = 2048 + causal_offset: int = 1 + + # Map Diffusers-style param keys to sglang's LTX-2 primitive naming. + # The refiner reuses LTX2Attention / LTX2FeedForward, so it inherits the + # same naming differences vs Diffusers: + # * ff.net.0.proj / ff.net.2 -> proj_in / proj_out (LTX2FeedForward) + # * norm_q / norm_k -> q_norm / k_norm (LTX2Attention) + # Keep this aligned with LTX2ArchConfig.param_names_mapping in ltx_2.py. + param_names_mapping: dict = field( + default_factory=lambda: { + r"^(transformer_blocks\.\d+\.ff)\.net\.0\.proj\.(.*)$": r"\1.proj_in.\2", + r"^(transformer_blocks\.\d+\.ff)\.net\.2\.(.*)$": r"\1.proj_out.\2", + r"(.*)\.norm_q\.(.*)$": r"\1.q_norm.\2", + r"(.*)\.norm_k\.(.*)$": r"\1.k_norm.\2", + } + ) + + def __post_init__(self): + super().__post_init__() + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + + +@dataclass +class SanaWMRefinerConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=SanaWMRefinerArchConfig) + prefix: str = "SanaWMRefiner" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 70a177157..aee206d52 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -317,6 +317,12 @@ class PipelineConfig: def prepare_calculated_size(self, image): return self.calculate_condition_image_size(image, image.width, image.height) + def preprocess_realtime_condition_image(self, batch, _vae_image_processor) -> bool: + """Realtime hook: optionally preprocess the first-frame condition image + in-place. Return True if handled (skip the standard path), False to fall + back to the normal condition-image preprocessing. Default: not handled.""" + return False + def prepare_image_processor_kwargs(self, batch, neg=False): return {} @@ -744,6 +750,31 @@ class PipelineConfig: help="Override the selected pipeline config's resolution setting. Only applies to pipelines that define a resolution field.", ) + # SANA-WM streaming knobs. default=None so they only apply to pipeline + # configs that define these fields (e.g. SanaWMPipelineConfig); other + # configs are left untouched by update_config_from_args. + parser.add_argument( + f"--{prefix_with_dot}streaming", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}streaming", + default=None, + help="SANA-WM: enable chunk-causal streaming (forward_long) generation.", + ) + parser.add_argument( + f"--{prefix_with_dot}refiner-chunked", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}refiner_chunked", + default=None, + help="SANA-WM: chunk-wise streaming refiner (vs whole-clip dense refiner).", + ) + parser.add_argument( + f"--{prefix_with_dot}num-frame-per-block", + type=int, + dest=f"{prefix_with_dot.replace('-', '_')}num_frame_per_block", + default=None, + help="SANA-WM: latent frames per streaming chunk (default 3).", + ) + # DiT configuration parser.add_argument( f"--{prefix_with_dot}dit-precision", @@ -941,6 +972,30 @@ class PipelineConfig: ) # 1.5. Adjust pipeline config for fine-tuned VAE if needed pipeline_config_cls = model_info.pipeline_config_cls + # If an explicit pipeline_class_name refines the model-default config + # (e.g. SanaWMRealtimePipeline -> SanaWMRealtimeConfig, a subclass of + # the model-resolved SanaWMPipelineConfig), prefer the pipeline's own + # config so realtime-only wiring (the /v1/realtime_video adapter) is + # selected. Only applies when the explicit config strictly subclasses + # the model default, so non-realtime pipelines are unaffected. + if pipeline_class_name: + explicit_config_classes = get_pipeline_config_classes( + pipeline_class_name + ) + if explicit_config_classes is not None: + explicit_config_cls = explicit_config_classes[0] + if ( + isinstance(explicit_config_cls, type) + and isinstance(pipeline_config_cls, type) + and explicit_config_cls is not pipeline_config_cls + and issubclass(explicit_config_cls, pipeline_config_cls) + ): + logger.info( + f"Refining pipeline config {pipeline_config_cls.__name__} " + f"-> {explicit_config_cls.__name__} for explicit " + f"pipeline_class_name={pipeline_class_name}" + ) + pipeline_config_cls = explicit_config_cls vae_path = kwargs.get(prefix_with_dot + "vae_path") or kwargs.get("vae_path") if vae_path is None: component_paths = kwargs.get( diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py index 254946865..edde43fb7 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -24,6 +24,11 @@ from sglang.multimodal_gen.runtime.distributed import ( get_sp_world_size, ) +# Distilled 3-step stage-2 sigma schedule. Single source of truth shared by +# LTX2TwoStagePipeline and the SANA-WM refiner stages (matches NVlabs +# `inference_sana_wm.py` / the LTX-2 distilled refiner). +STAGE_2_DISTILLED_SIGMA_VALUES: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0) + def pack_text_embeds( text_hidden_states: torch.Tensor, diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py index f8580e652..2c16ffbd2 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py @@ -24,3 +24,4 @@ class ModelDeploymentConfig: fsdp_auto_min_available_memory_gb: float | None = None fsdp_auto_requires_cfg: bool = True fsdp_auto_requires_default_parallelism: bool = True + auto_enable_cfg_parallel: bool = True diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py new file mode 100644 index 000000000..386cc4234 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 + + +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits.sana_wm import SanaWMConfig +from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput +from sglang.multimodal_gen.configs.models.encoders.base import EncoderConfig +from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config +from sglang.multimodal_gen.configs.models.vaes.ltx_video import LTXVideoVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + PipelineConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import ( + ModelDeploymentConfig, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +def sana_wm_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: + """Extract Gemma-2 last hidden state as text conditioning (same as SANA T2I).""" + return outputs.last_hidden_state + + +SANA_WM_CHI_PROMPT: tuple[str, ...] = ( + 'Given a user prompt, generate an "Enhanced prompt" that provides detailed ' + "visual descriptions suitable for image generation. Evaluate the level of " + "detail in the user prompt:", + "- If the prompt is simple, focus on adding specifics about colors, shapes, " + "sizes, textures, and spatial relationships to create vivid and concrete scenes.", + "- If the prompt is already detailed, refine and enhance the existing details " + "slightly without overcomplicating.", + "Here are examples of how to transform or refine prompts:", + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled " + "up in a round shape, sleeping peacefully on a warm sunny windowsill, " + "surrounded by pots of blooming red flowers.", + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene " + "at dusk, featuring glowing street lamps, a diverse crowd of people in " + "colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", + "Please generate only the enhanced description for the prompt below and " + "avoid including any additional commentary or evaluations:", + "User Prompt: ", +) + + +@dataclass +class SanaWMPipelineConfig(PipelineConfig): + """Pipeline config for the SANA-WM TI2V world model (text + first-frame image -> video, + optional 6-DoF camera trajectory, optional Stage-2 LTX-2 refiner).""" + + task_type: ModelTaskType = ModelTaskType.TI2V + + # SanaWMBeforeDenoisingStage._splice_first_frame handles condition-image + # resize + VAE-encode itself, so bypass the framework's generic TI2V + # preprocessing in InputValidationStage. Without this, the framework path + # reads `vae_config.arch_config.scale_factor_spatial` which LTXVideoVAEArchConfig + # does not expose (it uses `spatial_compression_ratio` instead). LTX-2 sets + # the same flag for the same reason -- both are TI2V on LTXVideoVAEConfig. + skip_input_image_preprocess: bool = True + + # --- Guidance --- + # SANA-WM uses standard CFG via guidance_scale; no embedded guidance token. + should_use_guidance: bool = False + + enable_autocast: bool = False + + # --- Streaming self-forcing (S1c) --- + # When ``streaming`` is set, the pipeline uses the autoregressive + # SanaWMStreamingDenoisingStage (forward_long, chunk-by-chunk) instead of the + # one-shot bidirectional denoise. ``num_frame_per_block`` is the streaming + # chunk size in LATENT frames (distinct from the DiT's intra-attention + # ``arch.chunk_size``). ``denoising_step_list`` must end in 0. + streaming: bool = False + num_frame_per_block: int = 3 + num_cached_blocks: int = 2 + sink_token: bool = True + denoising_step_list: tuple = (1000, 960, 889, 727, 0) + streaming_cfg_scale: float = 1.0 + # Streaming refiner (S2b): chunked LTX-2 sink/current refiner. + sink_size: int = 1 + refiner_block_size: int = 3 + refiner_kv_max_frames: int = 11 + refiner_seed: int = 42 + # True -> chunked streaming refiner (low-latency, causal); False -> whole-clip + # dense refiner (global context, max quality, non-streaming). + refiner_chunked: bool = True + + # --- DiT --- + dit_config: DiTConfig = field(default_factory=SanaWMConfig) + + # --- VAE: LTX-2 (128ch, 8× temporal, 32× spatial) --- + vae_config: VAEConfig = field(default_factory=LTXVideoVAEConfig) + vae_precision: str = "bf16" + # Match NVlabs SANA-WM inference: long videos must use LTX-2 spatial + # tiling plus framewise temporal decode to avoid oversized Conv3d pads. + vae_tiling: bool = True + vae_sp: bool = False # no VAE SP for now + vae_framewise_encoding: bool = True + vae_framewise_decoding: bool = True + vae_tile_sample_min_num_frames: int = 96 + vae_tile_sample_stride_num_frames: int = 64 + + # Load both encoder and decoder (need encoder for first-frame conditioning) + def __post_init__(self): + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + self.vae_config.use_tiling = self.vae_tiling + self.vae_config.use_temporal_tiling = self.vae_framewise_decoding + self.vae_config.tile_sample_min_num_frames = self.vae_tile_sample_min_num_frames + self.vae_config.tile_sample_stride_num_frames = ( + self.vae_tile_sample_stride_num_frames + ) + self.vae_config.blend_num_frames = ( + self.vae_tile_sample_min_num_frames - self.vae_tile_sample_stride_num_frames + ) + + # --- Text encoder: Gemma-2-2b-it (single encoder, same as SANA T2I) --- + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (Gemma2Config(),) + ) + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",)) + text_encoder_extra_args: list[dict] = field( + default_factory=lambda: [ + { + # Match NVlabs SANA-WM prompt encoding: positive and negative + # branches must have the same token dimension for CFG concat. + "padding": "max_length", + "return_attention_mask": True, + } + ] + ) + chi_prompt: tuple[str, ...] = SANA_WM_CHI_PROMPT + preprocess_text_funcs: tuple[Callable | None, ...] = field( + default_factory=lambda: (None,) + ) + postprocess_text_funcs: tuple[Callable, ...] = field( + default_factory=lambda: (sana_wm_postprocess_text,) + ) + + # --- Scheduler --- + # linear_flow training schedule from the released config.yaml. + flow_shift: float = 9.95 + # Official NVlabs/Sana inference resolves inference_flow_shift first and + # falls back to flow_shift only when it is absent. + inference_flow_shift: float | None = 9.8 + + # --- Video shape --- + # VAE strides: (temporal, spatial_h, spatial_w) + vae_stride: tuple = (8, 32, 32) # LTX-2 VAE temporal=8, spatial=32 + + # --- Camera conditioning --- + camera_conditioning: bool = True # set False to disable camera branch + + # --- Deployment --- + def get_model_deployment_config(self) -> ModelDeploymentConfig: + return ModelDeploymentConfig( + auto_dit_layerwise_offload=True, + # Conservative auto-FSDP gate for the 720p world-model path. Users + # can still force FSDP explicitly on smaller cards. + fsdp_auto_min_available_memory_gb=60, + ) + + # --- Latent shape --- + def prepare_latent_shape(self, batch, batch_size: int, num_frames: int): + """ + Returns 5D latent shape: (B, 128, T_latent, H_sp, W_sp). + T_latent = ceil((num_frames - 1) / temporal_stride) + 1 + """ + t_stride = self.vae_stride[0] + h_stride = self.vae_stride[1] + w_stride = self.vae_stride[2] if len(self.vae_stride) > 2 else h_stride + + if batch.height % h_stride != 0 or batch.width % w_stride != 0: + raise ValueError( + "SANA-WM height/width must be divisible by the LTX-2 spatial " + f"stride ({h_stride}, {w_stride}); got " + f"height={batch.height}, width={batch.width}." + ) + + T_latent = (num_frames - 1) // t_stride + 1 + H_sp = batch.height // h_stride + W_sp = batch.width // w_stride + z_dim = self.vae_config.arch_config.latent_channels # 128 + + return (batch_size, z_dim, T_latent, H_sp, W_sp) + + def adjust_num_frames(self, num_frames: int) -> int: + """Ensure (num_frames - 1) is divisible by VAE temporal stride.""" + t_stride = self.vae_stride[0] + if (num_frames - 1) % t_stride != 0: + adjusted = ((num_frames - 1) // t_stride) * t_stride + 1 + logger.warning( + f"num_frames - 1 must be divisible by temporal stride {t_stride}. " + f"Rounding {num_frames} → {adjusted}." + ) + return adjusted + return num_frames + + # --- Text embedding accessors --- + def get_pos_prompt_embeds(self, batch): + return batch.prompt_embeds[0] + + def get_neg_prompt_embeds(self, batch): + return batch.negative_prompt_embeds[0] + + # --- Conditioning kwargs for DenoisingStage --- + def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): + """Build positive conditioning kwargs passed to SanaWMTransformer3DModel.forward. + + The DiT forward signature consumes: + * encoder_hidden_states -- Gemma-2 embeddings (set by DenoisingStage) + * timestep -- diffusion step (set by DenoisingStage) + * encoder_attention_mask -- text padding mask + * camera_conditions -- (B, T_lat, 20) latent-frame UCPE raymap + * chunk_plucker -- (B, 48, T_lat, H, W) packed Plücker raymap + """ + out = {} + + m = batch.prompt_attention_mask + if isinstance(m, (list, tuple)): + out["encoder_attention_mask"] = m[0] if m else None + elif m is not None: + out["encoder_attention_mask"] = m + + # Camera conditioning (built by SanaWMBeforeDenoisingStage) + if hasattr(batch, "extra") and batch.extra: + cc = batch.extra.get("camera_conditions", None) + cp = batch.extra.get("chunk_plucker", None) + if cc is not None: + out["camera_conditions"] = cc + if cp is not None: + out["chunk_plucker"] = cp + + return out + + def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): + """Build negative conditioning kwargs for CFG. + + Camera/plucker are structural video conditions, not text conditions. + NVlabs SANA-WM duplicates them for both CFG branches and only swaps + text embeddings/masks. + """ + out = {} + m = batch.negative_attention_mask + if isinstance(m, (list, tuple)): + out["encoder_attention_mask"] = m[0] if m else None + elif m is not None: + out["encoder_attention_mask"] = m + if hasattr(batch, "extra") and batch.extra: + cc = batch.extra.get("camera_conditions", None) + cp = batch.extra.get("chunk_plucker", None) + if cc is not None: + out["camera_conditions"] = cc + if cp is not None: + out["chunk_plucker"] = cp + return out + + # --- Post-processing --- + def post_denoising_loop(self, latents: torch.Tensor, batch) -> torch.Tensor: + """No token un-packing needed; 5D latents are already spatial.""" + return latents + + def shard_latents_for_sp(self, batch, latents): + # SANA-WM uses frame-wise GDN recurrent scan and a temporal depth-wise + # conv (GLUMBConvTemp, t_kernel=3) that both span across frames. Splitting + # the latent along T would truncate the GDN hidden state and drop the + # GLUMBConvTemp halo at rank boundaries, producing silent wrong outputs. + # Camera/Plücker tensors are also indexed in lockstep with T and would + # need matching shards. Disable SP until a halo-exchange-aware impl lands. + return latents, False + + def gather_latents_for_sp(self, latents): + return latents + + def get_decode_scale_and_shift(self, device, dtype, vae): + """Invert the LTX-2 latent normalization used before denoising. + + SANA-WM uses the LTX-2 VAE. Upstream encodes as + ``(z - latents_mean) * scaling_factor / latents_std`` and decodes by + applying the inverse transform. + """ + latents_mean = getattr(vae, "latents_mean", None) + latents_std = getattr(vae, "latents_std", None) + + scaling_factor = ( + getattr(getattr(vae, "config", None), "scaling_factor", None) + or getattr(vae, "scaling_factor", None) + or getattr(self.vae_config.arch_config, "scaling_factor", None) + or 1.0 + ) + if isinstance(scaling_factor, (int, float)) and float(scaling_factor) == 0.0: + scaling_factor = 1.0 + + if isinstance(latents_mean, torch.Tensor) and isinstance( + latents_std, torch.Tensor + ): + latents_mean = latents_mean.to(device=device, dtype=dtype).view( + 1, -1, 1, 1, 1 + ) + latents_std = latents_std.to(device=device, dtype=dtype).view( + 1, -1, 1, 1, 1 + ) + sf = torch.tensor(float(scaling_factor), device=device, dtype=dtype).view( + 1, 1, 1, 1, 1 + ) + return sf / latents_std, latents_mean + + sf = torch.tensor(float(scaling_factor), device=device, dtype=dtype).view( + 1, 1, 1, 1, 1 + ) + return sf, None + + +class SanaWMRealtimeConfig(SanaWMPipelineConfig): + """Realtime alias of the SANA-WM pipeline config. + + Same numerics/fields as SanaWMPipelineConfig (our correct streaming pipeline); + exists so the realtime-serving registry can key adapters on a realtime config + class (matching the upstream/LingBot-World pattern) without renaming the base. + """ + + def get_model_deployment_config(self) -> ModelDeploymentConfig: + return ModelDeploymentConfig( + auto_dit_layerwise_offload=True, + auto_disable_component_offload_min_available_memory_gb=120, + auto_disable_component_offload_components=("dit",), + auto_enable_cfg_parallel=False, + ) diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 7cbfc101c..f15dee3e3 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -702,7 +702,7 @@ class SamplingParams: raise user_kwargs = dict(kwargs) - user_kwargs.pop("diffusers_kwargs", None) + diffusers_kwargs = user_kwargs.pop("diffusers_kwargs", None) user_sampling_params = type(sampling_params)(*args, **user_kwargs) # TODO: refactor @@ -710,6 +710,8 @@ class SamplingParams: user_sampling_params, explicit_fields=set(user_kwargs.keys()) ) sampling_params._explicit_fields = set(user_kwargs.keys()) + if diffusers_kwargs is not None: + sampling_params.diffusers_kwargs = diffusers_kwargs sampling_params._adjust(server_args) sampling_params._validate_with_pipeline_config(server_args.pipeline_config) @@ -945,6 +947,36 @@ class SamplingParams: '--image-path "img1.png" "img2.png"' ), ) + add_argument( + "--action", + type=str, + help=( + "SANA-WM WASD/IJKL action DSL, e.g. " + "'w-80,jw-40,w-40,lw-60,w-100'. Model-specific fields are " + "ignored by other pipelines." + ), + ) + add_argument( + "--translation-speed", + "--translation_speed", + type=float, + dest="translation_speed", + help="SANA-WM action DSL per-frame translation speed.", + ) + add_argument( + "--rotation-speed-deg", + "--rotation_speed_deg", + type=float, + dest="rotation_speed_deg", + help="SANA-WM action DSL per-frame rotation speed in degrees.", + ) + add_argument( + "--pitch-limit-deg", + "--pitch_limit_deg", + type=float, + dest="pitch_limit_deg", + help="SANA-WM action DSL absolute pitch clamp in degrees.", + ) add_argument( "--moba-config-path", type=str, diff --git a/python/sglang/multimodal_gen/configs/sample/sana_wm.py b/python/sglang/multimodal_gen/configs/sample/sana_wm.py new file mode 100644 index 000000000..87024f83f --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/sana_wm.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sampling parameters for SANA-WM TI2V world model generation.""" + +from dataclasses import dataclass +from typing import Any, Optional, Sequence, Union + +from sglang.multimodal_gen.configs.sample.sampling_params import ( + DataType, + SamplingParams, +) + +# Type alias for the camera tensor inputs. Accept torch.Tensor, numpy arrays, +# or nested Python lists — coerced to torch.Tensor downstream in the stage. +CameraTensorLike = Union[Any, Sequence[Sequence[Sequence[float]]]] + + +@dataclass +class SanaWMSamplingParams(SamplingParams): + """Default sampling parameters for SANA-WM 720p (704×1280) 16fps video. + + Frame counts must satisfy (num_frames - 1) % 8 == 0. + + Optional camera conditioning: + camera_to_world: (T, 4, 4) extrinsics, one per output frame. + intrinsics: (T, 3, 3) pinhole intrinsics, one per output frame. + action: WASD/IJKL action DSL (e.g. "w-80,jw-40,w-40"), rolled + out to camera_to_world before the camera branch. + Omitted camera_to_world -> static identity camera. Omitted intrinsics -> + centered heuristic; pass explicit intrinsics for closest NVlabs parity. + """ + + data_type: DataType = DataType.VIDEO + + # Resolution: 720p landscape (LTX-2 VAE requires multiples of 32) + height: int = 704 + width: int = 1280 + + # 49 = (49-1)/8 = 6 latent frames → ~3 seconds at 16fps + num_frames: int = 49 + + # SANA-WM is trained at 16fps (override base default of 24). + fps: int = 16 + + num_inference_steps: int = 20 + + guidance_scale: float = 4.5 + + # NVlabs' SANA-WM inference defaults to an empty negative prompt. + negative_prompt: str = "" + + # --- Camera trajectory (6-DoF) — optional --- + camera_to_world: Optional[CameraTensorLike] = None + intrinsics: Optional[CameraTensorLike] = None + action: Optional[str] = None + translation_speed: float = ( + 0.04 # match official streaming (STREAMING_TRANSLATION_SPEED) + ) + rotation_speed_deg: float = 1.2 + pitch_limit_deg: float = 85.0 + + def build_request_extra(self) -> dict[str, Any]: + extra = super().build_request_extra() + if self.action is not None and self.camera_to_world is not None: + raise ValueError( + "SANA-WM accepts either action or camera_to_world, not both." + ) + if self.camera_to_world is not None: + extra["camera_to_world"] = self.camera_to_world + if self.intrinsics is not None: + extra["intrinsics"] = self.intrinsics + if self.action is not None: + extra["action"] = self.action + extra["translation_speed"] = self.translation_speed + extra["rotation_speed_deg"] = self.rotation_speed_deg + extra["pitch_limit_deg"] = self.pitch_limit_deg + return extra diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 0871eaf5a..22f7d7cf6 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -77,6 +77,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImagePipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import ( StableDiffusion3PipelineConfig, ) @@ -131,6 +132,7 @@ from sglang.multimodal_gen.configs.sample.qwenimage import ( QwenImageSamplingParams, ) from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams +from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams from sglang.multimodal_gen.configs.sample.stablediffusion3 import ( StableDiffusion3SamplingParams, ) @@ -977,6 +979,20 @@ def _register_configs(): ], ) + # SANA-WM (register BEFORE generic SANA T2I to prevent "sana" detector false-match) + register_configs( + sampling_param_cls=SanaWMSamplingParams, + pipeline_config_cls=SanaWMPipelineConfig, + hf_model_paths=[ + "Efficient-Large-Model/SANA-WM_bidirectional", + "Efficient-Large-Model/SANA-WM_streaming", + ], + model_detectors=[ + # Match "sana-wm" or "sana_wm" but NOT plain T2I "sana" checkpoints. + lambda hf_id: ("sana-wm" in hf_id.lower() or "sana_wm" in hf_id.lower()), + ], + ) + # Cosmos3 — single checkpoint serves T2V, I2V, and T2I. Mode is dispatched # per-request inside the pipeline from ``num_frames`` and ``image_path``. # Both Nano (8B) and Super (32B) share the same pipeline; arch dimensions @@ -1005,7 +1021,13 @@ def _register_configs(): "Efficient-Large-Model/Sana_1600M_512px_diffusers", "Efficient-Large-Model/Sana_600M_512px_diffusers", ], - model_detectors=[lambda hf_id: "sana" in hf_id.lower()], + model_detectors=[ + lambda hf_id: ( + "sana" in hf_id.lower() + and "sana-wm" not in hf_id.lower() + and "sana_wm" not in hf_id.lower() + ) + ], ) # FireRed-Image-Edit diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index 24e218bd6..3802954d2 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -73,6 +73,22 @@ async def _wait_until_http_ready(server_args: ServerArgs) -> None: raise RuntimeError(f"HTTP server did not become ready at {health_url}") +def _is_realtime_serving(server_args: ServerArgs) -> bool: + """A realtime pipeline establishes per-session state over the WebSocket, so + the synthetic server-warmup request (which has no session) cannot run — it + would fail in the realtime stage and abort startup. Detect it via the + realtime-adapter registry and skip server warmup.""" + try: + from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.registry import ( + get_realtime_model_adapter, + ) + + get_realtime_model_adapter(server_args) + return True + except Exception: + return False + + async def _run_server_warmup_after_http_ready( server_args: ServerArgs, warmup_done: asyncio.Event ) -> None: @@ -81,6 +97,7 @@ async def _run_server_warmup_after_http_ready( not server_args.warmup or not server_args.server_warmup or server_args.warmup_resolutions is not None + or _is_realtime_serving(server_args) ): warmup_done.set() return diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py index 6514c6fdb..a15b46b3a 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py @@ -4,7 +4,6 @@ from __future__ import annotations import os import tempfile -from collections import deque from typing import TYPE_CHECKING, Any from fastapi import WebSocket @@ -28,12 +27,13 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( from sglang.multimodal_gen.runtime.entrypoints.utils import ( prepare_request, ) +from sglang.multimodal_gen.runtime.realtime.camera_controls import ( + RealtimeCameraControlState, +) from sglang.multimodal_gen.runtime.realtime.condition_events import ( ConditionEvent, ConditionEventQueue, ControlSignal, - ControlStateSamplingQueue, - ControlStateTransition, ) from sglang.multimodal_gen.runtime.server_args import get_global_server_args @@ -53,22 +53,18 @@ LINGBOT_REALTIME_DEFAULT_NUM_INFERENCE_STEPS = 4 LINGBOT_REALTIME_MIN_CONDITION_CHUNKS = 2 -class LingBotWorldRealtimeState: +class LingBotWorldRealtimeState(RealtimeCameraControlState): def __init__(self): - self.events = ConditionEventQueue(max_events={"prompt": 1}) - self.camera_state = ControlStateSamplingQueue( - default_item=[], + super().__init__( min_pulse_items=1, + script_maxlen=512, max_transitions=512, ) - self.camera_script_queue: deque[ControlSignal] = deque(maxlen=512) - self.latest_sampled_event_id: int | None = None + self.events = ConditionEventQueue(max_events={"prompt": 1}) def clear(self) -> None: + super().clear() self.events.clear() - self.camera_state.clear() - self.camera_script_queue.clear() - self.latest_sampled_event_id = None def receive_prompt(self, prompt: str, *, event_id: int | None = None) -> None: self.events.push( @@ -82,125 +78,17 @@ class LingBotWorldRealtimeState: ) ) - def receive_camera_script( - self, - camera_actions: list[list[str]], - *, - event_id: int | None = None, - ) -> None: - self.camera_script_queue.clear() - self.camera_state.clear() - for actions in camera_actions: - self.camera_script_queue.append( - ControlSignal( - kind="camera_actions", - payload=list(actions), - seq_id=event_id, - ) - ) - - def receive_camera_state_transitions( - self, - transitions: list[ControlStateTransition], - ) -> None: - self.camera_script_queue.clear() - self.camera_state.push_many(transitions) - - def receive_camera_actions( - self, - camera_actions: list[list[str]], - *, - event_id: int | None = None, - ) -> None: - self.receive_camera_script(camera_actions, event_id=event_id) - - def receive_camera_state( - self, - actions: list[str], - *, - event_id: int | None = None, - timestamp_ms: int | None = None, - ) -> None: - self.receive_camera_state_transitions( - [ - ControlStateTransition( - payload=list(actions), - seq_id=event_id, - timestamp_ms=timestamp_ms, - ) - ] - ) - - def _sample_camera_script(self, chunk_size: int) -> list[list[str]]: - chunk: list[list[str]] = [] - latest_event_id = self.latest_sampled_event_id - while self.camera_script_queue and len(chunk) < chunk_size: - signal = self.camera_script_queue.popleft() - chunk.append(list(signal.payload)) - latest_event_id = signal.seq_id - while len(chunk) < chunk_size: - chunk.append([]) - self.latest_sampled_event_id = latest_event_id - return chunk - - def _camera_state_transition( - self, - actions: list[str], - *, - event_id: int | None, - timestamp_ms: int | None, - ) -> ControlStateTransition: - return ControlStateTransition( - payload=list(actions), - seq_id=event_id, - timestamp_ms=timestamp_ms, - ) - - def _camera_transitions_from_event_payload( - self, - payload: dict[str, Any], - *, - event_id: int | None, - ) -> list[ControlStateTransition]: - transitions = payload.get("transitions") - if not isinstance(transitions, list): - raise ValueError("camera_actions state payload requires transitions") - result = [] - for transition in transitions: - if not isinstance(transition, dict): - raise ValueError("camera_actions transition must be a map") - actions = transition.get("actions") - if not isinstance(actions, list): - raise ValueError("camera_actions transition actions must be a list") - timestamp_ms = transition.get("client_ts_ms") - if timestamp_ms is not None: - timestamp_ms = int(timestamp_ms) - result.append( - self._camera_state_transition( - list(actions), - event_id=event_id, - timestamp_ms=timestamp_ms, - ) - ) - return result - def receive_camera_event_payload( self, payload: Any, *, event_id: int | None, ) -> str: - if isinstance(payload, dict) and payload.get("mode") == "state": - transitions = self._camera_transitions_from_event_payload( - payload, - event_id=event_id, - ) - self.receive_camera_state_transitions(transitions) - return f"kind=camera_actions, mode=state, transitions={len(transitions)}" - - camera_actions = LingBotWorldRealtimeAdapter._validate_camera_actions(payload) - self.receive_camera_script(camera_actions, event_id=event_id) - return f"kind=camera_actions, mode=script, frames={len(camera_actions)}" + return super().receive_camera_event_payload( + payload, + event_id=event_id, + validate_camera_actions=LingBotWorldRealtimeAdapter._validate_camera_actions, + ) def sample_prompt(self) -> str: prompt = self.events.pop_latest("prompt") @@ -209,20 +97,6 @@ class LingBotWorldRealtimeState: self.latest_sampled_event_id = self.events.last_sampled_seq_id("prompt") return prompt - def sample_camera_actions(self, chunk_size: int) -> list[list[str]] | None: - """samples a sequence of camera actions for the chunk with chunk_size frames - - Args: - chunk_size: number of frames - """ - if self.camera_script_queue: - return self._sample_camera_script(chunk_size) - action_list = self.camera_state.sample_chunk(chunk_size) - if action_list is None: - return None - self.latest_sampled_event_id = self.camera_state.latest_sampled_seq_id() - return [list(actions) for actions in action_list] - def has_prompt(self) -> bool: return self.events.has_events("prompt") diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/sana_wm_realtime_adapter.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/sana_wm_realtime_adapter.py new file mode 100644 index 000000000..3ba2967aa --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/sana_wm_realtime_adapter.py @@ -0,0 +1,359 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import os +import tempfile +from typing import TYPE_CHECKING, Any + +from fastapi import WebSocket + +from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( + RealtimeEvent, + RealtimeVideoGenerationsRequest, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import ( + RealtimeChunkInputs, + RealtimeModelAdapter, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_output_adapter import ( + RawRGBRealtimeOutputAdapter, + RealtimeFrameSendStats, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( + build_sampling_params, + save_image_to_path, +) +from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.base import ( + normalize_sana_wm_camera_actions, + parse_sana_wm_action_string, + snap_sana_wm_num_frames, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.self_forcing import ( + SanaWMSelfForcingSampler, +) +from sglang.multimodal_gen.runtime.realtime.camera_controls import ( + RealtimeCameraControlState, +) +from sglang.multimodal_gen.runtime.server_args import get_global_server_args + +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import ( + GenerateSession, + RealtimeChunkContext, + ) + from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import ( + OutputBatch, + Req, + ) + from sglang.multimodal_gen.runtime.server_args import ServerArgs + + +SANA_WM_DEFAULT_SIZE = "1280x704" +SANA_WM_DEFAULT_NUM_FRAMES = 1081 +SANA_WM_DEFAULT_FPS = 16 +SANA_WM_DEFAULT_STEPS = 4 +SANA_WM_DEFAULT_GUIDANCE = 1.0 +SANA_WM_CONTROL_PULSE_FRAMES = 8 + + +def _normalize_sana_wm_state_actions(actions: list[Any]) -> list[str]: + return [str(action).lower() for action in actions] + + +class SanaWMRealtimeAdapterState(RealtimeCameraControlState): + def __init__(self): + super().__init__( + min_pulse_items=SANA_WM_CONTROL_PULSE_FRAMES, + script_maxlen=2048, + max_transitions=512, + normalize_state_actions=_normalize_sana_wm_state_actions, + ) + self.base_condition_inputs: dict[str, Any] = {} + + def clear(self) -> None: + super().clear() + self.base_condition_inputs.clear() + + def receive_camera_event_payload( + self, + payload: Any, + *, + event_id: int | None, + ) -> str: + return super().receive_camera_event_payload( + payload, + event_id=event_id, + validate_camera_actions=SanaWMRealtimeAdapter._validate_camera_actions, + ) + + +class SanaWMRealtimeAdapter(RealtimeModelAdapter): + name = "sana_wm_realtime" + + def __init__(self): + self.output_adapter = RawRGBRealtimeOutputAdapter() + + def create_state(self) -> SanaWMRealtimeAdapterState: + return SanaWMRealtimeAdapterState() + + def _state(self, session: GenerateSession) -> SanaWMRealtimeAdapterState: + state = session.adapter_state + if not isinstance(state, SanaWMRealtimeAdapterState): + raise TypeError("SANA-WM realtime adapter state is not initialized") + return state + + @staticmethod + def _validate_camera_actions(payload: Any) -> list[list[str]]: + return normalize_sana_wm_camera_actions( + payload, error_label="camera_actions event payload" + ) + + @staticmethod + def _raw_frame_count(result: OutputBatch) -> int | None: + if result.raw_frame_batches is None: + return None + return sum(len(frames) for frames in result.raw_frame_batches) + + async def on_init( + self, + session: GenerateSession, + request: RealtimeVideoGenerationsRequest, + ) -> None: + if request.first_frame is None: + raise ValueError("SANA-WM realtime requires first_frame") + + request.size = request.size or SANA_WM_DEFAULT_SIZE + if request.num_frames is not None: + request.num_frames = int(request.num_frames) + else: + # Open-ended session: keep num_frames unset so prepare_next_request + # samples uniform action chunks (no front-loaded segmentation), and + # flag the stage explicitly via condition_inputs — + # build_sampling_params strips None fields, so the per-chunk batch + # would otherwise carry the SamplingParams default num_frames. + request.condition_inputs = { + **(request.condition_inputs or {}), + "sana_wm_open_ended": True, + } + request.fps = int(request.fps or SANA_WM_DEFAULT_FPS) + request.num_inference_steps = int( + request.num_inference_steps or SANA_WM_DEFAULT_STEPS + ) + request.guidance_scale = float( + request.guidance_scale or SANA_WM_DEFAULT_GUIDANCE + ) + if request.negative_prompt is None: + request.negative_prompt = "" + if request.generator_device is None: + request.generator_device = "cuda" + + state = self._state(session) + condition_inputs = dict(request.condition_inputs or {}) + camera_actions = condition_inputs.pop("camera_actions", None) + action = condition_inputs.pop("action", None) + if camera_actions is not None and action is not None: + raise ValueError("pass only one of camera_actions or action") + if camera_actions is not None: + state.receive_camera_event_payload(camera_actions, event_id=None) + if action is not None: + if not isinstance(action, str) or not action: + raise ValueError("action condition input must be a non-empty string") + state.receive_camera_script( + parse_sana_wm_action_string(action), event_id=None + ) + state.base_condition_inputs = condition_inputs + + server_args = get_global_server_args() + if server_args.input_save_path is not None: + uploads_dir = server_args.input_save_path + os.makedirs(uploads_dir, exist_ok=True) + else: + if session.input_temp_dir is None: + session.input_temp_dir = tempfile.mkdtemp(prefix="sglang_input_") + uploads_dir = session.input_temp_dir + + if isinstance( + request.first_frame, str + ) and request.first_frame.lower().startswith(("http://", "https://")): + suffix = os.path.splitext(request.first_frame.split("?", 1)[0])[1] + digest = hashlib.sha256(request.first_frame.encode("utf-8")).hexdigest()[ + :16 + ] + target_path = os.path.join(uploads_dir, f"realtime_ref_{digest}{suffix}") + if os.path.exists(target_path): + request.first_frame = target_path + return + else: + target_path = os.path.join(uploads_dir, f"{session.id}_first_frame") + image_path = await save_image_to_path(request.first_frame, target_path) + request.first_frame = image_path + + async def wait_for_next_chunk(self, session: GenerateSession) -> None: + del session + + def ingest_event( + self, + session: GenerateSession, + event: RealtimeEvent, + ) -> str: + state = self._state(session) + if event.kind == "camera_actions": + return state.receive_camera_event_payload( + event.payload, + event_id=event.event_id, + ) + if event.kind == "action": + if not isinstance(event.payload, str) or not event.payload: + raise ValueError("action event payload must be a non-empty string") + camera_actions = parse_sana_wm_action_string(event.payload) + state.receive_camera_script(camera_actions, event_id=event.event_id) + return f"kind=action, frames={len(camera_actions)}" + raise ValueError(f"unsupported event kind: {event.kind}") + + def _sample_chunk_inputs( + self, + session: GenerateSession, + chunk: RealtimeChunkContext, + action_chunk_size: int, + ) -> RealtimeChunkInputs: + state = self._state(session) + request = session.request + if request is None: + raise ValueError("realtime request is not initialized") + + condition_inputs = dict(state.base_condition_inputs) if chunk.index == 0 else {} + camera_actions = state.sample_camera_actions(action_chunk_size) + if camera_actions is not None: + condition_inputs["camera_actions"] = camera_actions + return RealtimeChunkInputs( + prompt=request.prompt, + condition_inputs=condition_inputs, + ) + + def _build_sampling_params( + self, + session: GenerateSession, + chunk: RealtimeChunkContext, + chunk_inputs: RealtimeChunkInputs, + chunk_size: int, + ): + request = session.request + if request is None: + raise ValueError("realtime request is not initialized") + + return build_sampling_params( + chunk.request_id, + prompt=chunk_inputs.prompt, + size=request.size, + num_frames=request.num_frames, + fps=request.fps, + image_path=request.first_frame, + output_file_name=chunk.request_id, + save_output=False, + seed=request.seed, + generator_device=request.generator_device, + num_inference_steps=request.num_inference_steps, + guidance_scale=request.guidance_scale, + guidance_scale_2=request.guidance_scale_2, + negative_prompt=request.negative_prompt, + enable_teacache=request.enable_teacache, + enable_frame_interpolation=request.enable_frame_interpolation, + frame_interpolation_exp=request.frame_interpolation_exp, + frame_interpolation_scale=request.frame_interpolation_scale, + frame_interpolation_model_path=request.frame_interpolation_model_path, + enable_upscaling=request.enable_upscaling, + upscaling_model_path=request.upscaling_model_path, + upscaling_scale=request.upscaling_scale, + diffusers_kwargs=request.diffusers_kwargs, + profile=request.profile, + num_profiled_timesteps=request.num_profiled_timesteps, + profile_all_stages=request.profile_all_stages, + perf_dump_path=request.perf_dump_path, + output_path=request.output_path, + output_compression=request.output_compression, + output_quality=request.output_quality, + condition_inputs=chunk_inputs.condition_inputs, + realtime_chunk_size=chunk_size, + ) + + def prepare_next_request( + self, + session: GenerateSession, + server_args: ServerArgs, + chunk: RealtimeChunkContext, + ) -> Req: + arch_config = server_args.pipeline_config.dit_config.arch_config + chunk_size = int(getattr(arch_config, "num_frames_per_block", 3)) + temporal_compression = int( + server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio + ) + # Match action sampling to the latent span used by the batch path. Chunk + # 0 may carry a front-loaded remainder, so a fixed nfpb*tc action count + # would read static-padded camera poses and drift from batch output. + action_chunk_size = chunk_size * temporal_compression + req_num_frames = ( + session.request.num_frames if session.request is not None else None + ) + if req_num_frames is not None: + snapped = snap_sana_wm_num_frames( + int(req_num_frames), stride=temporal_compression + ) + latent_t = (snapped - 1) // temporal_compression + 1 + segments = SanaWMSelfForcingSampler.create_autoregressive_segments( + latent_t, chunk_size + ) + idx = int(chunk.index) + if 0 <= idx and idx + 1 < len(segments): + action_chunk_size = ( + segments[idx + 1] - segments[idx] + ) * temporal_compression + chunk_inputs = self._sample_chunk_inputs(session, chunk, action_chunk_size) + sampling_params = self._build_sampling_params( + session, + chunk, + chunk_inputs, + chunk_size, + ) + batch = prepare_request( + server_args=server_args, + sampling_params=sampling_params, + ) + batch.session = session.realtime_session + batch.realtime_session_id = session.id + batch.return_raw_frames = True + batch.block_idx = chunk.index + batch.realtime_event_id = self._state(session).latest_sampled_event_id + if session.request is not None: + # Forward the full transport config like the LingBot adapter does — + # the shared RawRGB output adapter / realtime_video_api consume + # preview width + pacing too; dropping them silently disabled both + # features for SANA-WM sessions. + batch.realtime_output_format = session.request.realtime_output_format + batch.realtime_preview_max_width = ( + session.request.realtime_preview_max_width + ) + batch.realtime_output_pacing = bool(session.request.realtime_output_pacing) + return batch + + async def send_output( + self, + ws: WebSocket, + session: GenerateSession, + result: OutputBatch, + batch: Req, + ) -> RealtimeFrameSendStats: + return await self.output_adapter.send(ws, session, result, batch) + + def on_chunk_complete(self, session: GenerateSession, result: OutputBatch) -> None: + if session.request is not None and self._raw_frame_count(result) == 0: + session.request.max_chunks = session.generate_chunk_cnt + 1 + session.generate_chunk_completed() + + def dispose(self, session: GenerateSession) -> None: + state = session.adapter_state + if isinstance(state, SanaWMRealtimeAdapterState): + state.clear() + self.output_adapter.reset() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py index decc09363..349a3b879 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py @@ -124,6 +124,7 @@ def _frame_shape_from_metadata( RAW_RGB_FRAMES_PER_WS_MESSAGE = 16 +ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE = 6 FRAME_BATCH_PACK_OFFLOAD_BYTES = 64 * 1024 WEBP_DEFAULT_QUALITY = 90 JPEG_DEFAULT_QUALITY = 95 @@ -139,12 +140,15 @@ class _TransportPayload: metadata: dict[str, int | str | bool | list[int]] -def _split_frame_batch(frames: list[bytes]) -> list[list[bytes]]: +def _split_frame_batch( + frames: list[bytes], + frames_per_message: int = RAW_RGB_FRAMES_PER_WS_MESSAGE, +) -> list[list[bytes]]: if not frames: return [frames] return [ - frames[i : i + RAW_RGB_FRAMES_PER_WS_MESSAGE] - for i in range(0, len(frames), RAW_RGB_FRAMES_PER_WS_MESSAGE) + frames[i : i + frames_per_message] + for i in range(0, len(frames), frames_per_message) ] @@ -495,7 +499,7 @@ class RawRGBRealtimeOutputAdapter: stats = empty_frame_send_stats(content_type) for frames in frame_batches: split_batches = ( - [frames] + _split_frame_batch(frames, ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE) if _is_encoded_preview_transport( content_type=content_type, output_format=output_format, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py index 44181d52e..851358542 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py @@ -31,14 +31,24 @@ def _register_builtin_realtime_model_adapters() -> None: from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import ( LingBotWorldCausalDMDConfig, ) + from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import ( + SanaWMRealtimeConfig, + ) from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.adapters.lingbot_world_realtime_adapter import ( LingBotWorldRealtimeAdapter, ) + from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.adapters.sana_wm_realtime_adapter import ( + SanaWMRealtimeAdapter, + ) register_realtime_model_adapter( LingBotWorldCausalDMDConfig, LingBotWorldRealtimeAdapter, ) + register_realtime_model_adapter( + SanaWMRealtimeConfig, + SanaWMRealtimeAdapter, + ) _BUILTIN_ADAPTERS_REGISTERED = True diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm.py new file mode 100644 index 000000000..387324dec --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm.py @@ -0,0 +1,906 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +from typing import Callable, List, Optional, Tuple + +import torch +import torch.nn as nn + +from sglang.multimodal_gen.configs.models.dits.sana_wm import SanaWMConfig +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT + +# Re-exported for back-compat: callers import these names from this module path. +from sglang.multimodal_gen.runtime.models.dits.sana_wm_components import ( # noqa: F401 + _CACHE_TYPE_CONCAT, + _CACHE_TYPE_STATE, + _INT32_SAFE_CONV_ELEMENTS, + _NUM_STREAM_CACHE_SLOTS, + _SLOT_CAM_K, + _SLOT_CAM_V, + _SLOT_FFN_TCONV, + _SLOT_K, + _SLOT_SHORTCONV, + _SLOT_TYPE_FLAG, + _SLOT_V, + BidirectionalGDNUCPESinglePathLiteLA, + CaptionEmbedder, + GLUMBConvTemp, + MultiHeadCrossAttention, + PatchEmbedMS3D, + T2IFinalLayer, + TimestepEmbedder, + WanRotaryPosEmbed, + _apply_block_diagonal, + _apply_complex_rope, + _apply_ray_projmat, + _apply_rotary_emb_bhnd, + _apply_rotary_emb_dn, + _bidirectional_short_conv, + _build_ucpe_apply_fns, + _compute_fov_from_focal, + _compute_frame_gates, + _ConvLayer, + _downscale_to_reference_rms, + _flip_and_shift, + _gdn_chunk_scan_forward, + _gdn_scan_bidirectional, + _gdn_scan_cached, + _gdn_scan_forward, + _gdn_scan_forward_stateful, + _invert_SE3, + _log_sana_wm_triton_cam_gdn_fallback, + _log_sana_wm_triton_gdn_fallback, + _RMSNorm, + _sana_wm_chunk_boundaries_for_attention, + _sana_wm_chunk_index_from_chunk_size, + _sana_wm_chunked_attention, + _sana_wm_normalize_chunk_index, + _sana_wm_padded_scale, + _sana_wm_sdpa, + _ShortConvolution, + _single_path_delta_chunk_scan_forward, + _single_path_delta_scan_bidirectional, + _single_path_delta_scan_cached, + _single_path_delta_scan_forward, + _single_path_delta_scan_forward_stateful, + _sinusoidal_timestep_embedding, + _slice_rope_for_cam, + _slice_rope_to_current_chunk, + _temporal_short_conv_cached, + _tensor_cache_key, + _UpstreamMlp, + compute_chunk_plucker, + process_camera_conditions_ucpe, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm import ( + parity_probe, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class SanaWMBlock(nn.Module): + """One transformer block of SANA-WM.""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + head_dim: int, + mlp_ratio: float, + t_kernel_size: int, + qk_norm: bool, + cross_norm: bool, + conv_kernel_size: int, + k_conv_only: bool, + softmax_main: bool, + use_chunk_plucker_post_attn: bool, + chunk_size: Optional[int] = None, + chunk_split_strategy: str = "uniform", + update_rule: str = "torch_chunk", + cam_update_rule: str = "torch_chunk", + chunk_gdn_chunk_size: int = 21, + use_chunked_softmax_attention: bool = False, + gdn_backend: str = "auto", + ) -> None: + super().__init__() + self.softmax_main = softmax_main + self.chunk_size = chunk_size + self.chunk_split_strategy = chunk_split_strategy + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + + self.attn = BidirectionalGDNUCPESinglePathLiteLA( + in_dim=hidden_size, + heads=num_heads, + head_dim=head_dim, + qk_norm=qk_norm, + conv_kernel_size=conv_kernel_size, + k_conv_only=k_conv_only, + softmax_main=softmax_main, + update_rule=update_rule, + cam_update_rule=cam_update_rule, + chunk_gdn_chunk_size=chunk_gdn_chunk_size, + use_chunked_softmax_attention=use_chunked_softmax_attention, + gdn_backend=gdn_backend, + ) + + self.cross_attn = MultiHeadCrossAttention( + d_model=hidden_size, + num_heads=num_heads, + qk_norm=cross_norm, + ) + + self.mlp = GLUMBConvTemp( + in_features=hidden_size, + hidden_features=int(hidden_size * mlp_ratio), + t_kernel_size=t_kernel_size, + ) + + self.scale_shift_table = nn.Parameter( + torch.randn(6, hidden_size) / hidden_size**0.5 + ) + + if use_chunk_plucker_post_attn: + self.plucker_proj = nn.Linear(hidden_size, hidden_size, bias=True) + nn.init.zeros_(self.plucker_proj.weight) + nn.init.zeros_(self.plucker_proj.bias) + else: + self.plucker_proj = None + + @staticmethod + def _modulate( + x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor + ) -> torch.Tensor: + return x * (1 + scale) + shift + + @staticmethod + def _reshape_framewise_modulation( + x: torch.Tensor, + num_frames: int, + ) -> tuple[torch.Tensor, int]: + B, N, C = x.shape + tokens_per_frame = N // num_frames + return x.reshape(B, num_frames, tokens_per_frame, C), tokens_per_frame + + def forward( + self, + x: torch.Tensor, # (B, N, D) + y: torch.Tensor, # (B, L, D) text embeds + t: torch.Tensor, # (B, 6*D) AdaLN-single + HW: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor], + prope_fns: Optional[Tuple[Callable, Callable, Callable]], + plucker_emb: Optional[torch.Tensor], + mask: Optional[torch.Tensor], + chunk_size: Optional[int] = None, + chunk_split_strategy: Optional[str] = None, + chunk_index: Optional[List[int]] = None, + ) -> torch.Tensor: + B = x.shape[0] + if t.dim() == 2: + num_frames = None + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + t.reshape(B, 6, -1) + ).chunk(6, dim=1) + else: + num_frames = t.reshape(B, -1, 6, t.shape[-1] // 6).shape[1] + t = t.reshape(B, num_frames, 6, -1) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None, None, :, :] + t + ).chunk(6, dim=2) + + # Self-attention with UCPE camera branch + if num_frames is None: + x_in = self._modulate(self.norm1(x), shift_msa, scale_msa) + else: + x_norm, tokens_per_frame = self._reshape_framewise_modulation( + self.norm1(x), num_frames + ) + x_in = self._modulate(x_norm, shift_msa, scale_msa).reshape_as(x) + attn_out = self.attn( + x_in, + HW=HW, + rotary_emb=rotary_emb, + prope_fns=prope_fns, + chunk_size=self.chunk_size if chunk_size is None else chunk_size, + chunk_split_strategy=( + self.chunk_split_strategy + if chunk_split_strategy is None + else chunk_split_strategy + ), + chunk_index=chunk_index, + ) + if num_frames is None: + x = x + gate_msa * attn_out + else: + attn_out = attn_out.reshape(B, num_frames, tokens_per_frame, -1) + x = x + (gate_msa * attn_out).reshape_as(x) + + # Plücker post-attn injection (zero-init linear) + if self.plucker_proj is not None and plucker_emb is not None: + x = x + self.plucker_proj(plucker_emb) + + # Cross-attention + x = x + self.cross_attn(x, y, mask=mask) + + # FFN + if num_frames is None: + x_in = self._modulate(self.norm2(x), shift_mlp, scale_mlp) + x = x + gate_mlp * self.mlp(x_in, HW=HW) + else: + x_norm, tokens_per_frame = self._reshape_framewise_modulation( + self.norm2(x), num_frames + ) + x_in = self._modulate(x_norm, shift_mlp, scale_mlp).reshape_as(x) + mlp_out = self.mlp(x_in, HW=HW).reshape(B, num_frames, tokens_per_frame, -1) + x = x + (gate_mlp * mlp_out).reshape_as(x) + return x + + def forward_long( + self, + x: torch.Tensor, + y: torch.Tensor, + t: torch.Tensor, + HW: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor], + prope_fns: Optional[Tuple[Callable, Callable, Callable]], + plucker_emb: Optional[torch.Tensor], + mask: Optional[torch.Tensor], + *, + kv_cache: list, + save_kv_cache: bool, + ) -> Tuple[torch.Tensor, list]: + """Streaming counterpart of ``forward``: threads the per-block 10-slot ``kv_cache`` through cached attention + FFN.""" + B = x.shape[0] + if t.dim() == 2: + num_frames = None + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + t.reshape(B, 6, -1) + ).chunk(6, dim=1) + else: + num_frames = t.reshape(B, -1, 6, t.shape[-1] // 6).shape[1] + t = t.reshape(B, num_frames, 6, -1) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None, None, :, :] + t + ).chunk(6, dim=2) + + if num_frames is None: + x_in = self._modulate(self.norm1(x), shift_msa, scale_msa) + else: + x_norm, tokens_per_frame = self._reshape_framewise_modulation( + self.norm1(x), num_frames + ) + x_in = self._modulate(x_norm, shift_msa, scale_msa).reshape_as(x) + + attn_out, kv_cache = self.attn.forward_long( + x_in, + HW=HW, + rotary_emb=rotary_emb, + prope_fns=prope_fns, + kv_cache=kv_cache, + save_kv_cache=save_kv_cache, + ) + if num_frames is None: + x = x + gate_msa * attn_out + else: + attn_out = attn_out.reshape(B, num_frames, tokens_per_frame, -1) + x = x + (gate_msa * attn_out).reshape_as(x) + + if self.plucker_proj is not None and plucker_emb is not None: + x = x + self.plucker_proj(plucker_emb) + + x = x + self.cross_attn(x, y, mask=mask) + + if num_frames is None: + x_in = self._modulate(self.norm2(x), shift_mlp, scale_mlp) + else: + x_norm, tokens_per_frame = self._reshape_framewise_modulation( + self.norm2(x), num_frames + ) + x_in = self._modulate(x_norm, shift_mlp, scale_mlp).reshape_as(x) + + # GLUMBConvTemp returns a tuple whenever the streaming path is active + # (ffn_tail set OR save requested); branch on tuple-ness, never on + # save_kv_cache alone (a read-only pass with a populated slot 9 still + # returns a tuple). + mlp_out = self.mlp( + x_in, + HW=HW, + ffn_tail=kv_cache[_SLOT_FFN_TCONV], + save_ffn_tail=save_kv_cache, + ) + if isinstance(mlp_out, tuple): + mlp_out, ffn_tail = mlp_out + if save_kv_cache: + kv_cache[_SLOT_FFN_TCONV] = ffn_tail + if num_frames is None: + x = x + gate_mlp * mlp_out + else: + mlp_out = mlp_out.reshape(B, num_frames, tokens_per_frame, -1) + x = x + (gate_mlp * mlp_out).reshape_as(x) + return x, kv_cache + + +class SanaWMTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): + """SANA-WM 2.6B TI2V world model. + + Forward inputs: + hidden_states: (B, C, T, H, W) 128-ch LTX-2 latent + encoder_hidden_states: (B, L, 2304) Gemma-2 embeddings + timestep: (B,) + encoder_attention_mask: (B, L) optional bool + camera_conditions: (B, T, 20) latent-frame raymap: + 16 c2w + (fx,fy,cx,cy) + chunk_plucker: (B, 48, T, H, W) optional, computed + from camera_conditions + if absent. + + Returns: ``(B, C, T, H, W)`` predicted velocity / noise. + """ + + _fsdp_shard_conditions = SanaWMConfig()._fsdp_shard_conditions + _compile_conditions = SanaWMConfig()._compile_conditions + _supported_attention_backends = SanaWMConfig()._supported_attention_backends + param_names_mapping = SanaWMConfig().param_names_mapping + reverse_param_names_mapping = SanaWMConfig().reverse_param_names_mapping + lora_param_names_mapping: dict = {} + + def __init__(self, config: SanaWMConfig, hf_config=None, **kwargs) -> None: + super().__init__(config, hf_config=hf_config or {}, **kwargs) + arch = config.arch_config + + self.patch_size = (arch.patch_size_t, arch.patch_size, arch.patch_size) + self.inner_dim = arch.num_attention_heads * arch.attention_head_dim + self.hidden_size = self.inner_dim + self.num_attention_heads = arch.num_attention_heads + self.attention_head_dim = arch.attention_head_dim + self.out_channels = arch.out_channels + self.num_channels_latents = arch.num_channels_latents + self.vae_temporal_stride = arch.vae_temporal_stride + self.timestep_norm_scale_factor = getattr( + arch, "timestep_norm_scale_factor", 1.0 + ) + + # --- Embedders --- + self.x_embedder = PatchEmbedMS3D( + self.patch_size, + arch.in_channels, + self.inner_dim, + bias=True, + ) + + self.t_embedder = TimestepEmbedder(self.inner_dim, frequency_embedding_size=256) + self.t_block = nn.Sequential( + nn.SiLU(), + nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True), + ) + + self.y_embedder = CaptionEmbedder( + in_channels=arch.caption_channels, + hidden_size=self.inner_dim, + token_num=arch.model_max_length, + ) + self.y_norm = bool(getattr(arch, "y_norm", True)) + self.attention_y_norm = _RMSNorm( + self.inner_dim, + scale_factor=getattr(arch, "y_norm_scale_factor", 1.0), + eps=getattr(arch, "y_norm_eps", 1e-5), + ) + + # 3-channel raymap embedder -- kept for state_dict compatibility but + # only invoked when ``use_chunk_plucker_post_attn`` is False. + # When ``True`` (the case for the released checkpoint) the absmap + # path is skipped entirely. + self.raymap_embedder = PatchEmbedMS3D( + self.patch_size, + 3, + self.inner_dim, + bias=True, + ) + # 48-channel plucker embedder (chunk-packed) + if arch.use_chunk_plucker_post_attn or arch.use_chunk_plucker_input: + self.plucker_embedder = PatchEmbedMS3D( + self.patch_size, + arch.chunk_plucker_channels, + self.inner_dim, + bias=True, + ) + nn.init.zeros_(self.plucker_embedder.proj.weight) + nn.init.zeros_(self.plucker_embedder.proj.bias) + else: + self.plucker_embedder = None + self.use_chunk_plucker_post_attn = arch.use_chunk_plucker_post_attn + self.use_chunk_plucker_input = arch.use_chunk_plucker_input + self.chunk_size = getattr(arch, "chunk_size", None) + self.chunk_split_strategy = getattr(arch, "chunk_split_strategy", "uniform") + + # --- RoPE --- + self.rope = WanRotaryPosEmbed( + attention_head_dim=arch.linear_head_dim, + patch_size=self.patch_size, + max_seq_len=1024, + ) + + # --- Transformer blocks --- + depth = arch.num_layers + self.softmax_every_n = arch.softmax_every_n + softmax_idx = set( + i + for i in range(depth) + if arch.softmax_every_n > 0 and (i + 1) % arch.softmax_every_n == 0 + ) + self.softmax_block_indices = tuple(sorted(softmax_idx)) + + self.blocks = nn.ModuleList( + [ + SanaWMBlock( + hidden_size=self.inner_dim, + num_heads=arch.num_attention_heads, + head_dim=arch.linear_head_dim, + mlp_ratio=arch.mlp_ratio, + t_kernel_size=arch.t_kernel_size, + qk_norm=arch.qk_norm, + cross_norm=arch.cross_norm, + conv_kernel_size=arch.conv_kernel_size, + k_conv_only=arch.k_conv_only, + softmax_main=(i in softmax_idx), + use_chunk_plucker_post_attn=( + arch.use_chunk_plucker_post_attn + and ( + arch.chunk_plucker_post_attn_blocks < 0 + or i < arch.chunk_plucker_post_attn_blocks + ) + ), + chunk_size=self.chunk_size, + chunk_split_strategy=self.chunk_split_strategy, + update_rule=getattr(arch, "update_rule", "torch_chunk"), + cam_update_rule=getattr(arch, "cam_update_rule", "torch_chunk"), + chunk_gdn_chunk_size=getattr(arch, "chunk_gdn_chunk_size", 21), + use_chunked_softmax_attention=getattr( + arch, "use_chunked_softmax_attention", False + ), + gdn_backend=getattr(arch, "gdn_backend", "auto"), + ) + for i in range(depth) + ] + ) + + self.final_layer = T2IFinalLayer( + self.inner_dim, self.patch_size, self.out_channels + ) + + # Cache RoPE freqs per shape -- avoids recomputation across denoising + # steps with constant latent shapes. + self._freqs_cache: dict = {} + self._ucpe_apply_fns_cache: Optional[ + Tuple[Tuple, torch.Tensor, Tuple[Callable, Callable, Callable]] + ] = None + self._plucker_emb_cache: Optional[Tuple[Tuple, torch.Tensor, torch.Tensor]] = ( + None + ) + + # FSDP shard targets + self.layer_names = ["blocks"] + + def post_load_weights(self) -> None: + # FSDP loader initializes the model on meta and only materializes + # tensors that appear in the checkpoint. WanRotaryPosEmbed._freqs is a + # derived, non-persistent constant, so recompute it deterministically. + for module in self.modules(): + if isinstance(module, WanRotaryPosEmbed): + if module._freqs.is_meta: + module._init_freqs_buffer() + + # ------------------------------------------------------------------ # + # Forward + # ------------------------------------------------------------------ # + + def _get_freqs(self, T: int, H: int, W: int, device: torch.device) -> torch.Tensor: + key = (T, H, W, str(device)) + if key not in self._freqs_cache: + self._freqs_cache[key] = self.rope((T, H, W), device) + return self._freqs_cache[key] + + def _get_freqs_window( + self, + start: int, + end: int, + H: int, + W: int, + device: torch.device, + frame_index: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """RoPE freqs for a streaming chunk at GLOBAL frame positions. + + ``frame_index`` (per-token global positions) overrides ``(start, end)``; + the count branch is cached, the tensor branch is computed fresh. + """ + if frame_index is not None: + return self.rope((end - start, H, W), device, frame_index=frame_index) + key = ("win", int(start), int(end), H, W, str(device)) + if key not in self._freqs_cache: + self._freqs_cache[key] = self.rope(((int(start), int(end)), H, W), device) + return self._freqs_cache[key] + + def _get_ucpe_apply_fns( + self, + camera_conditions: torch.Tensor, + *, + HW: Tuple[int, int, int], + freqs: torch.Tensor, + ) -> Tuple[Callable, Callable, Callable]: + head_dim = self.attention_head_dim + if torch.is_grad_enabled(): + raymats = process_camera_conditions_ucpe( + camera_conditions, + HW=HW, + patch_size=self.patch_size, + ) + raymats_flat = raymats.reshape(camera_conditions.shape[0], -1, 4, 4) + return _build_ucpe_apply_fns(head_dim, raymats_flat, freqs) + + key = ( + "ucpe", + HW, + self.patch_size, + head_dim, + _tensor_cache_key(camera_conditions), + _tensor_cache_key(freqs), + ) + cached = self._ucpe_apply_fns_cache + if cached is not None and cached[0] == key: + return cached[2] + + raymats = process_camera_conditions_ucpe( + camera_conditions, + HW=HW, + patch_size=self.patch_size, + ) + raymats_flat = raymats.reshape(camera_conditions.shape[0], -1, 4, 4) + prope_fns = _build_ucpe_apply_fns(head_dim, raymats_flat, freqs) + self._ucpe_apply_fns_cache = (key, camera_conditions, prope_fns) + return prope_fns + + def _get_plucker_emb( + self, + chunk_plucker: torch.Tensor, + *, + latent_token_count: int, + ) -> torch.Tensor: + if self.plucker_embedder is None: + raise ValueError("SANA-WM plucker_embedder is not initialized.") + + weight = self.plucker_embedder.proj.weight + bias = self.plucker_embedder.proj.bias + key = ( + "plucker_emb", + latent_token_count, + self.patch_size, + _tensor_cache_key(chunk_plucker), + _tensor_cache_key(weight), + None if bias is None else _tensor_cache_key(bias), + ) + if not torch.is_grad_enabled(): + cached = self._plucker_emb_cache + if cached is not None and cached[0] == key: + return cached[2] + + plucker_emb = self.plucker_embedder(chunk_plucker.to(weight.dtype)) + if plucker_emb.shape[1] != latent_token_count: + raise ValueError( + f"plucker_emb token count {plucker_emb.shape[1]} != " + f"latent token count {latent_token_count}; " + "expected chunk_plucker shape (B, 48, T, H, W)." + ) + + if not torch.is_grad_enabled(): + self._plucker_emb_cache = (key, chunk_plucker, plucker_emb) + return plucker_emb + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + timestep: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + camera_conditions: Optional[torch.Tensor] = None, + chunk_plucker: Optional[torch.Tensor] = None, + guidance: Optional[torch.Tensor] = None, # kept for compat + **kwargs, + ) -> torch.Tensor: + if encoder_hidden_states is None: + raise ValueError("SANA-WM forward requires encoder_hidden_states.") + if timestep is None: + raise ValueError("SANA-WM forward requires timestep.") + + B, C, T_raw, H_raw, W_raw = hidden_states.shape + p_t, p_h, p_w = self.patch_size + T = T_raw // p_t + H = H_raw // p_h + W = W_raw // p_w + chunk_size = kwargs.get("chunk_size", self.chunk_size) + chunk_split_strategy = kwargs.get( + "chunk_split_strategy", self.chunk_split_strategy + ) + chunk_index = kwargs.get("chunk_index", None) + + # Patch embed: (B, C, T, H, W) -> (B, T*H*W, D) + x = self.x_embedder(hidden_states.to(dtype=self.x_embedder.proj.weight.dtype)) + + # Timestep AdaLN-single. SANA-WM's LTX sampler passes per-frame + # timesteps shaped (B, 1, T) so the clean first-frame condition can stay + # at timestep 0 while remaining latent frames denoise. Keep the scalar + # path for generic scheduler compatibility. + if self.timestep_norm_scale_factor != 1.0: + timestep_for_embed = ( + timestep.float() / self.timestep_norm_scale_factor + ).to(torch.float32) + else: + timestep_for_embed = timestep.long().to(torch.float32) + + if timestep_for_embed.dim() == 1: + t_emb = self.t_embedder(timestep_for_embed) # (B, D) + t6 = self.t_block(t_emb) # (B, 6D) + else: + timestep_shape = tuple(timestep_for_embed.shape) + t_flat = self.t_embedder(timestep_for_embed.flatten()) + t6_flat = self.t_block(t_flat) + t_emb = t_flat.unflatten(0, timestep_shape) + t6 = t6_flat.unflatten(0, timestep_shape) + + if isinstance(encoder_attention_mask, (list, tuple)): + encoder_attention_mask = encoder_attention_mask[0] + y = encoder_hidden_states + if y.dim() == 3: + y = y.unsqueeze(1) + y = self.y_embedder(y).squeeze(1) # (B, L, D) + if y.shape[0] != B: + y = y.expand(B, -1, -1).contiguous() + if self.y_norm: + y = self.attention_y_norm(y) + if encoder_attention_mask is not None and encoder_attention_mask.shape[0] != B: + encoder_attention_mask = encoder_attention_mask.expand(B, -1).contiguous() + + freqs = self._get_freqs(T, H, W, x.device) + + # Camera conditioning: UCPE prope_fns + Plücker + prope_fns = None + if camera_conditions is not None: + if camera_conditions.shape[1] != T: + raise ValueError( + "SANA-WM camera_conditions must be sampled at latent " + f"frames: got {camera_conditions.shape[1]} frames, " + f"expected T={T}." + ) + prope_fns = self._get_ucpe_apply_fns( + camera_conditions, + HW=(T, H, W), + freqs=freqs, + ) + + # Plücker post-attn embedding (shared across all blocks) + plucker_emb = None + needs_plucker_emb = ( + chunk_plucker is not None + and self.plucker_embedder is not None + and (self.use_chunk_plucker_post_attn or self.use_chunk_plucker_input) + ) + if needs_plucker_emb: + plucker_emb = self._get_plucker_emb( + chunk_plucker, + latent_token_count=x.shape[1], + ) # (B, T*H*W, D) + + if self.use_chunk_plucker_input and plucker_emb is not None: + x = x + plucker_emb + + if not self.use_chunk_plucker_post_attn: + plucker_emb = None + + # --- 6. Transformer blocks --- + HW = (T, H, W) + for block in self.blocks: + x = block( + x, + y=y, + t=t6, + HW=HW, + rotary_emb=freqs, + prope_fns=prope_fns, + plucker_emb=plucker_emb, + mask=encoder_attention_mask, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + chunk_index=chunk_index, + ) + + x = self.final_layer(x, t_emb) # (B, N, p_t*p_h*p_w*C_out) + + # Un-patch + x = x.reshape(B, T, H, W, p_t, p_h, p_w, self.out_channels) + x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).contiguous() + x = x.reshape(B, self.out_channels, T * p_t, H * p_h, W * p_w) + return x + + def forward_long( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + timestep: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + camera_conditions: Optional[torch.Tensor] = None, + chunk_plucker: Optional[torch.Tensor] = None, + *, + kv_cache: Optional[list] = None, + save_kv_cache: bool = True, + start_f: Optional[int] = None, + end_f: Optional[int] = None, + frame_index: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, list]: + """Streaming autoregressive forward over a chunk of latent frames. + + RoPE / camera / plücker are windowed to the chunk's GLOBAL frame range + ``[start_f, end_f)``; a per-block 10-slot ``kv_cache`` carries recurrent + state / concat-windows across chunks. Returns ``(out, new_cache)``. + """ + if encoder_hidden_states is None: + raise ValueError("SANA-WM forward_long requires encoder_hidden_states.") + if timestep is None: + raise ValueError("SANA-WM forward_long requires timestep.") + + if kv_cache is None: + kv_cache = [[None] * _NUM_STREAM_CACHE_SLOTS for _ in self.blocks] + + B, C, T_raw, H_raw, W_raw = hidden_states.shape + p_t, p_h, p_w = self.patch_size + T = T_raw // p_t + H = H_raw // p_h + W = W_raw // p_w + start = 0 if start_f is None else int(start_f) + end = start + T if end_f is None else int(end_f) + + x = self.x_embedder(hidden_states.to(dtype=self.x_embedder.proj.weight.dtype)) + + # Timestep AdaLN-single: force the framewise (B, 1, T) path so blocks + # always apply per-frame modulation. + if timestep.dim() == 1: + timestep = timestep[:, None, None].expand(-1, 1, T) + elif timestep.dim() == 2: + timestep = timestep[:, None, :] + if self.timestep_norm_scale_factor != 1.0: + timestep_for_embed = ( + timestep.float() / self.timestep_norm_scale_factor + ).to(torch.float32) + else: + timestep_for_embed = timestep.long().to(torch.float32) + timestep_shape = tuple(timestep_for_embed.shape) + t_flat = self.t_embedder(timestep_for_embed.flatten()) + t6_flat = self.t_block(t_flat) + t_emb = t_flat.unflatten(0, timestep_shape) + t6 = t6_flat.unflatten(0, timestep_shape) + + if isinstance(encoder_attention_mask, (list, tuple)): + encoder_attention_mask = encoder_attention_mask[0] + y = encoder_hidden_states + if y.dim() == 3: + y = y.unsqueeze(1) + y = self.y_embedder(y).squeeze(1) + if y.shape[0] != B: + y = y.expand(B, -1, -1).contiguous() + if self.y_norm: + y = self.attention_y_norm(y) + if encoder_attention_mask is not None and encoder_attention_mask.shape[0] != B: + encoder_attention_mask = encoder_attention_mask.expand(B, -1).contiguous() + + # RoPE windowed to global frame positions [start, end) + freqs = self._get_freqs_window( + start, end, H, W, x.device, frame_index=frame_index + ) + + # Camera conditioning: slice to the chunk, co-windowed w/ freqs + prope_fns = None + if camera_conditions is not None: + if camera_conditions.shape[1] != T: + # .contiguous(): canonical layout regardless of how the caller + # built the full-length tensor, so batch and realtime windows are + # kernel-level identical (slice offset/stride changes the reduction + # order otherwise — measured 1e-7 seeds amplifying to %-level drift + # through the bf16 block stack). + camera_conditions = camera_conditions[:, start:end].contiguous() + if camera_conditions.shape[0] != B: + camera_conditions = camera_conditions.repeat( + B // camera_conditions.shape[0], 1, 1 + ) + prope_fns = self._get_ucpe_apply_fns( + camera_conditions, HW=(T, H, W), freqs=freqs + ) + + # Plücker post-attn / input embedding, sliced to the chunk. + if chunk_plucker is not None and chunk_plucker.shape[2] != T: + chunk_plucker = chunk_plucker[ + :, :, start:end + ].contiguous() # see camera note + if chunk_plucker is not None and chunk_plucker.shape[0] != B: + chunk_plucker = chunk_plucker.repeat( + B // chunk_plucker.shape[0], 1, 1, 1, 1 + ) + plucker_emb = None + needs_plucker_emb = ( + chunk_plucker is not None + and self.plucker_embedder is not None + and (self.use_chunk_plucker_post_attn or self.use_chunk_plucker_input) + ) + if needs_plucker_emb: + plucker_emb = self._get_plucker_emb( + chunk_plucker, latent_token_count=x.shape[1] + ) + if self.use_chunk_plucker_input and plucker_emb is not None: + x = x + plucker_emb + if not self.use_chunk_plucker_post_attn: + plucker_emb = None + + # parity harness (env-gated, no-op in prod): on the FIRST sink-path call + # (frame_index not None), checksum the pre-block tensors and x after every + # block to localize where the two execution paths first diverge. + _probe_path = os.environ.get(parity_probe.ENV_BLOCK_PROBE) + _probe = None + if ( + _probe_path + and frame_index is not None + and not getattr(self, "_block_probe_done", False) + ): + _ck = parity_probe.checksum + _probe = { + "x_embed": _ck(x), + "t6": _ck(t6), + "y": _ck(y), + "freqs": ( + ( + tuple(freqs.shape), + float(freqs.real.detach().double().sum().item()), + float(freqs.imag.detach().double().sum().item()), + ) + if freqs is not None + else None + ), + "plucker_emb": _ck(plucker_emb), + "frame_index": frame_index.tolist(), + } + + HW = (T, H, W) + new_cache = [] + for i, block in enumerate(self.blocks): + x, block_cache = block.forward_long( + x, + y, + t6, + HW, + freqs, + prope_fns, + plucker_emb, + encoder_attention_mask, + kv_cache=kv_cache[i], + save_kv_cache=save_kv_cache, + ) + new_cache.append(block_cache) + if _probe is not None: + _probe[f"x_after_block_{i:02d}"] = parity_probe.checksum(x) + if _probe is not None: + torch.save(_probe, _probe_path) + self._block_probe_done = True + + x = self.final_layer(x, t_emb) + x = x.reshape(B, T, H, W, p_t, p_h, p_w, self.out_channels) + x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).contiguous() + x = x.reshape(B, self.out_channels, T * p_t, H * p_h, W * p_w) + return x, new_cache + + +EntryClass = SanaWMTransformer3DModel diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py new file mode 100644 index 000000000..a4cd0f720 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py @@ -0,0 +1,3037 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SANA-WM DiT building blocks (components). + +Reusable primitives, RoPE / UCPE camera geometry, GDN scan kernels, embedders, +FFN, and the SANA-WM attention modules. +""" + +import math +from typing import Callable, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.embeddings import get_1d_rotary_pos_embed + +from sglang.multimodal_gen.runtime.layers.attention import LocalAttention +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +_SANA_WM_TRITON_GDN_DISABLED_REASON: Optional[str] = None +_SANA_WM_TRITON_GDN_FALLBACK_LOGGED = False +_SANA_WM_TRITON_CAM_GDN_DISABLED_REASON: Optional[str] = None +_SANA_WM_TRITON_CAM_GDN_FALLBACK_LOGGED = False + +# Streaming per-block cache: a 10-slot list per block (mirrors the reference +# SANA-WM forward_long). GDN blocks store recurrent state (0/1 main, 2 cam), +# the K short-conv prefix (4), and a STATE type flag (6); softmax blocks store +# a K/V concat-window (0/1 main, 2/3 cam) and a CONCAT type flag (6); all blocks +# store the FFN temporal-conv tail (9). Slots 5/7/8 are unused. +_NUM_STREAM_CACHE_SLOTS = 10 +_SLOT_K = 0 +_SLOT_V = 1 +_SLOT_CAM_K = 2 +_SLOT_CAM_V = 3 +_SLOT_SHORTCONV = 4 +_SLOT_TYPE_FLAG = 6 +_SLOT_FFN_TCONV = 9 +_CACHE_TYPE_CONCAT = 0.0 +_CACHE_TYPE_STATE = 1.0 + + +def _tensor_cache_key(tensor: torch.Tensor) -> Tuple: + # Inference-mode tensors (e.g. an un-laundered camera tensor on the cfg=1.0 + # streaming path) raise on ._version access; they are never mutated, so 0 is safe. + try: + version = int(tensor._version) + except (RuntimeError, AttributeError): + version = 0 + return ( + tuple(tensor.shape), + tuple(tensor.stride()), + str(tensor.device), + tensor.dtype, + tensor.data_ptr(), + version, + ) + + +def _log_sana_wm_triton_gdn_fallback(reason: str) -> None: + global _SANA_WM_TRITON_GDN_FALLBACK_LOGGED + if not _SANA_WM_TRITON_GDN_FALLBACK_LOGGED: + logger.warning( + "SANA-WM Triton GDN fast path is unavailable; falling back to torch " + "GDN scan. reason=%s", + reason, + ) + _SANA_WM_TRITON_GDN_FALLBACK_LOGGED = True + + +def _log_sana_wm_triton_cam_gdn_fallback(reason: str) -> None: + global _SANA_WM_TRITON_CAM_GDN_FALLBACK_LOGGED + if not _SANA_WM_TRITON_CAM_GDN_FALLBACK_LOGGED: + logger.warning( + "SANA-WM Triton camera GDN fast path is unavailable; falling back " + "to torch camera scan. reason=%s", + reason, + ) + _SANA_WM_TRITON_CAM_GDN_FALLBACK_LOGGED = True + + +# --------------------------------------------------------------------------- +# Small primitives (RMSNorm, ShortConvolution). Parameter names/shapes match +# the released checkpoint so it loads cleanly. +# --------------------------------------------------------------------------- + + +class _RMSNorm(nn.Module): + """RMSNorm with signature ``RMSNorm(dim, scale_factor=1.0, eps=1e-6)``. + + Parameter name: ``weight`` (shape ``(dim,)``). + """ + + def __init__( + self, + dim: int, + scale_factor: float = 1.0, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim) * scale_factor) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Upstream Sana RMSNorm does both normalization and weight multiply in + # fp32 before casting back to the input dtype. + x_in = x + x32 = x.float() + rms = x32.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt() + return (x32 * rms * self.weight.to(dtype=x32.dtype)).type_as(x_in) + + +class _ShortConvolution(nn.Module): + """Depth-wise causal Conv1d along the temporal axis. + + Weight shape ``(hidden_size, 1, K)`` (groups=hidden_size). Input is + ``(B, T, C)`` with causal padding of ``K-1`` on the left. + """ + + def __init__(self, hidden_size: int, kernel_size: int) -> None: + super().__init__() + self.hidden_size = hidden_size + self.kernel_size = kernel_size + self.weight = nn.Parameter(torch.zeros(hidden_size, 1, kernel_size)) + # identity init: last tap = 1. The released SANA-WM checkpoint has no + # ShortConvolution bias, so keep this module bias-free. + with torch.no_grad(): + self.weight[:, 0, -1] = 1.0 + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, None]: + # x: (B, T, C) -> (B, C, T) for conv1d + x_bct = x.transpose(1, 2) + x_pad = F.pad(x_bct, (self.kernel_size - 1, 0)) + y = F.conv1d(x_pad, self.weight, bias=None, groups=self.hidden_size) + return y.transpose(1, 2), None + + +def _bidirectional_short_conv( + x: torch.Tensor, # (B*S, T, C) + conv: _ShortConvolution, +) -> torch.Tensor: + """Forward + backward causal pass minus the shared center tap (symmetric/non-causal filter).""" + y_fwd, _ = conv(x) + y_bwd, _ = conv(x.flip(1)) + y_bwd = y_bwd.flip(1) + w_center = conv.weight[:, 0, -1] # (C,) + center = x * w_center.view(1, 1, -1) + return (y_fwd + y_bwd - center).to(x.dtype) + + +def _temporal_short_conv_cached( + x: torch.Tensor, # (B, N=T*S, C) + conv: "_ShortConvolution", + HW: Tuple[int, int, int], + *, + prefix: Optional[torch.Tensor] = None, + save_prefix: bool = True, + bidirectional: bool = True, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Chunk-causal short conv on K for streaming `forward_long` (cache slot 4). + + The FORWARD (causal) pass prepends the previous chunk's last ``kernel-1`` + frames so it stays continuous across chunks; the BACKWARD pass (when + ``bidirectional``) is recomputed intra-chunk. The last ``kernel-1`` input + frames are returned as the next chunk's ``prefix``. A single chunk with no + prefix reduces to ``_bidirectional_short_conv``. + """ + B, N, C = x.shape + T, H, W = HW + S = H * W + pad = conv.kernel_size - 1 + xt = x.view(B, T, S, C).permute(0, 2, 1, 3).contiguous().reshape(B * S, T, C) + + if prefix is not None: + prefix = prefix.to(device=xt.device, dtype=xt.dtype) + y_fwd, _ = conv(torch.cat([prefix, xt], dim=1)) + y_fwd = y_fwd[:, -T:] + else: + y_fwd, _ = conv(xt) + + if bidirectional: + y_bwd, _ = conv(xt.flip(1)) + y_bwd = y_bwd.flip(1) + center = xt * conv.weight[:, 0, -1].view(1, 1, -1) + y = (y_fwd + y_bwd - center).to(xt.dtype) + else: + y = y_fwd.to(xt.dtype) + + new_prefix = xt[:, -pad:].detach().clone() if (save_prefix and pad > 0) else prefix + y = y.reshape(B, S, T, C).permute(0, 2, 1, 3).reshape(B, N, C) + return y, new_prefix + + +# --------------------------------------------------------------------------- +# 3D RoPE +# --------------------------------------------------------------------------- + + +class WanRotaryPosEmbed(nn.Module): + """3D rotary position embeddings split across (t, h, w) head dims. + + Returns complex ``freqs`` of shape ``(1, 1, T*H*W, D/2)``. + """ + + def __init__( + self, + attention_head_dim: int, + patch_size: Tuple[int, int, int], + max_seq_len: int = 1024, + theta: float = 10000.0, + ) -> None: + super().__init__() + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + self.theta = theta + self._init_freqs_buffer() + + def _init_freqs_buffer(self) -> None: + # This is a persistent=False buffer, so it is not in the checkpoint and + # stays on meta after FSDP weight load; post_load_weights re-runs this + # to rematerialize it. + h_dim = w_dim = 2 * (self.attention_head_dim // 6) + t_dim = self.attention_head_dim - h_dim - w_dim + + freqs = [] + for dim in [t_dim, h_dim, w_dim]: + freq = get_1d_rotary_pos_embed( + dim, + self.max_seq_len, + self.theta, + use_real=False, + repeat_interleave_real=False, + freqs_dtype=torch.float64, + ) + freqs.append(freq) + self.register_buffer("_freqs", torch.cat(freqs, dim=1), persistent=False) + + def forward( + self, + fhw, + device: torch.device, + frame_index: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # `fhw[0]` is either an int frame count (dense; positions 0..T-1) or a + # `(start, end)` tuple selecting GLOBAL frame positions for a streaming + # chunk. `frame_index`, when given, supplies per-token global frame + # positions directly and overrides `fhw[0]`, so chunk-by-chunk RoPE + # stays at absolute positions. + fspec, pph, ppw = fhw + freqs = self._freqs.to(device) + d = self.attention_head_dim + t_size = d // 2 - 2 * (d // 6) + h_size = d // 6 + w_size = d // 6 + ft, fh, fw = freqs.split_with_sizes([t_size, h_size, w_size], dim=1) + + if frame_index is not None: + f_pos = frame_index.to(device=device, dtype=torch.long) + elif isinstance(fspec, (tuple, list)): + f_pos = torch.arange( + int(fspec[0]), int(fspec[1]), device=device, dtype=torch.long + ) + else: + f_pos = torch.arange(int(fspec), device=device, dtype=torch.long) + ppf = int(f_pos.shape[0]) + + freqs_t = ft[f_pos].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_h = fh[:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_w = fw[:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + out = torch.cat([freqs_t, freqs_h, freqs_w], dim=-1) + return out.reshape(1, 1, ppf * pph * ppw, -1) + + +def _slice_rope_to_current_chunk( + rotary_emb: Optional[torch.Tensor], current_tokens: int +) -> Optional[torch.Tensor]: + """Defensive no-op: keep the trailing ``current_tokens`` of a RoPE table. + + ``forward_long`` builds ``freqs`` windowed to exactly the chunk, so this is + a no-op there; it only trims if a caller hands in a wider table.""" + if rotary_emb is None or rotary_emb.shape[-2] == current_tokens: + return rotary_emb + return rotary_emb[..., -current_tokens:, :] + + +def _apply_rotary_emb_dn( + hidden_states: torch.Tensor, freqs: torch.Tensor +) -> torch.Tensor: + """Apply complex RoPE to a tensor of shape ``(B, H, D, N)`` (GDN layout). + + ``freqs`` is complex with shape ``(1, 1, N, D/2)``. + """ + # (B, H, D, N) -> (B, H, N, D) + x = hidden_states.permute(0, 1, 3, 2).to(torch.float64).contiguous() + x_c = torch.view_as_complex(x.unflatten(-1, (-1, 2))) + y = torch.view_as_real(x_c * freqs).flatten(-2, -1) + return y.permute(0, 1, 3, 2).type_as(hidden_states) + + +def _apply_rotary_emb_bhnd( + hidden_states: torch.Tensor, freqs: torch.Tensor +) -> torch.Tensor: + """Apply complex RoPE to ``(B, H, N, D)`` (softmax attention layout).""" + x = hidden_states.to(torch.float64).contiguous() + x_c = torch.view_as_complex(x.unflatten(-1, (-1, 2))) + y = torch.view_as_real(x_c * freqs).flatten(-2, -1) + return y.type_as(hidden_states) + + +# --------------------------------------------------------------------------- +# UCPE block-diagonal apply primitives +# --------------------------------------------------------------------------- + + +def _apply_ray_projmat(feats: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor: + """Per-token 4x4 projmat applied to channels grouped by 4. + + feats: (B, H, N, D), matrix: (B, N, 4, 4). + """ + B, Hh, N, D = feats.shape + return torch.einsum( + "bnij,bhnkj->bhnki", + matrix, + feats.reshape(B, Hh, N, -1, 4), + ).reshape(feats.shape) + + +def _apply_complex_rope( + hidden_states: torch.Tensor, freqs: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + if inverse: + freqs = freqs.conj() + x_real = hidden_states.to(torch.float64) + if x_real.stride(-1) != 1: + x_real = x_real.contiguous() + x_c = torch.view_as_complex(x_real.unflatten(-1, (-1, 2))) + return torch.view_as_real(x_c * freqs).flatten(-2, -1).type_as(hidden_states) + + +def _apply_block_diagonal( + feats: torch.Tensor, + func_size_pairs: List[Tuple[Callable[[torch.Tensor], torch.Tensor], int]], +) -> torch.Tensor: + funcs, block_sizes = zip(*func_size_pairs) + assert feats.shape[-1] == sum(block_sizes), (feats.shape, block_sizes) + x_blocks = torch.split(feats, list(block_sizes), dim=-1) + return torch.cat([f(b) for f, b in zip(funcs, x_blocks)], dim=-1) + + +def _sana_wm_chunk_index_from_chunk_size( + T: int, + chunk_size: int, + strategy: str = "uniform", +) -> list[int]: + """Return temporal chunk start indices.""" + if chunk_size <= 0: + raise ValueError(f"chunk_size must be > 0, got {chunk_size}.") + if T <= 0: + raise ValueError(f"T must be > 0, got {T}.") + + strategy = "uniform" if strategy is None else str(strategy).lower() + + if strategy in ("uniform", "default"): + indices = list(range(0, T, chunk_size)) + if len(indices) > 1 and (T - indices[-1]) < chunk_size: + indices.pop() + return indices + + if strategy in ("first_frame", "first_frame_alone", "first_frame_only"): + if T <= 1: + return [0] + indices = [0] + list(range(1, T, chunk_size)) + if len(indices) > 2 and (T - indices[-1]) < chunk_size: + indices.pop() + return indices + + if strategy in ("first_plus_one", "first_chunk_plus_one"): + if T <= chunk_size + 1: + return [0] + indices = [0] + list(range(chunk_size + 1, T, chunk_size)) + if len(indices) > 1 and (T - indices[-1]) < chunk_size: + indices.pop() + return indices + + raise ValueError( + f"Unknown chunk_split_strategy '{strategy}'. Supported: " + "uniform, first_frame, first_plus_one." + ) + + +def _sana_wm_normalize_chunk_index( + chunk_index: Optional[List[int]], + T: int, + chunk_size: Optional[int] = None, + chunk_split_strategy: str = "uniform", +) -> list[int]: + if chunk_index is not None: + normalized = [int(idx) for idx in chunk_index] + if not normalized or normalized[0] != 0: + normalized = [0] + [idx for idx in normalized if idx > 0] + normalized = [idx for idx in normalized if idx < T] + if not normalized: + normalized = [0] + else: + if chunk_size is None: + raise ValueError("Either chunk_index or chunk_size must be provided.") + normalized = _sana_wm_chunk_index_from_chunk_size( + T, + int(chunk_size), + strategy=chunk_split_strategy, + ) + + if normalized[-1] != T: + normalized.append(T) + if any(end <= start for start, end in zip(normalized[:-1], normalized[1:])): + raise ValueError(f"chunk_index must be strictly increasing, got {normalized}.") + return normalized + + +def _sana_wm_chunk_boundaries_for_attention( + HW: Tuple[int, int, int], + chunk_size: Optional[int], + chunk_split_strategy: str, + chunk_index: Optional[List[int]], +) -> Optional[list[int]]: + T, _, _ = HW + if chunk_index is None and (chunk_size is None or int(chunk_size) >= T): + return None + + boundaries = _sana_wm_normalize_chunk_index( + chunk_index, + T, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + ) + if boundaries == [0, T]: + return None + return boundaries + + +def _sana_wm_sdpa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + softmax_scale: float, +) -> torch.Tensor: + q_sdpa = q.transpose(1, 2) + k_sdpa = k.transpose(1, 2) + v_sdpa = v.transpose(1, 2) + + dtype_orig = q_sdpa.dtype + head_dim = q_sdpa.shape[-1] + need_pad = head_dim not in (32, 64, 128, 256) and head_dim < 256 + if need_pad: + pad_to = 128 if head_dim <= 128 else 256 + pad_size = pad_to - head_dim + q_sdpa = F.pad(q_sdpa, (0, pad_size)) + k_sdpa = F.pad(k_sdpa, (0, pad_size)) + v_sdpa = F.pad(v_sdpa, (0, pad_size)) + + # CUDA SDPA cannot use flash kernels for fp32; cast to bf16 on that path + # and cast the output back to the caller dtype. + if q_sdpa.device.type == "cuda" and q_sdpa.dtype == torch.float32: + q_sdpa = q_sdpa.bfloat16() + k_sdpa = k_sdpa.bfloat16() + v_sdpa = v_sdpa.bfloat16() + + out = F.scaled_dot_product_attention( + q_sdpa, + k_sdpa, + v_sdpa, + dropout_p=0.0, + is_causal=False, + scale=softmax_scale, + ) + if need_pad: + out = out[..., :head_dim] + return out.transpose(1, 2).to(dtype_orig) + + +def _sana_wm_padded_scale(head_dim: int) -> float: + """Softmax scale for the head-padding SDPA path: pad the head dim to 128 + (or 256) and let SDPA use the *padded*-dim default scale. For SANA-WM's + head_dim=112 this is 1/sqrt(128), NOT 1/sqrt(112).""" + if head_dim in (32, 64, 128, 256) or head_dim >= 256: + return head_dim**-0.5 + pad_to = 128 if head_dim <= 128 else 256 + return pad_to**-0.5 + + +def _sana_wm_chunked_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + HW: Tuple[int, int, int], + chunk_size: Optional[int], + chunk_split_strategy: str, + chunk_index: Optional[List[int]], + softmax_scale: float, +) -> Optional[torch.Tensor]: + """Exact chunk-causal softmax attention without materializing an NxN mask.""" + boundaries = _sana_wm_chunk_boundaries_for_attention( + HW, + chunk_size, + chunk_split_strategy, + chunk_index, + ) + if boundaries is None: + return None + + _, H_sp, W_sp = HW + tokens_per_frame = H_sp * W_sp + out_chunks = [] + for start_frame, end_frame in zip(boundaries[:-1], boundaries[1:]): + query_start = start_frame * tokens_per_frame + query_end = end_frame * tokens_per_frame + kv_end = end_frame * tokens_per_frame + out_chunks.append( + _sana_wm_sdpa( + q[:, query_start:query_end], + k[:, :kv_end], + v[:, :kv_end], + softmax_scale=softmax_scale, + ) + ) + return torch.cat(out_chunks, dim=1) + + +def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: + Rinv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = Rinv + out[..., :3, 3] = -torch.einsum("...ij,...j->...i", Rinv, transforms[..., :3, 3]) + out[..., 3, 3] = 1.0 + return out + + +def _slice_rope_for_cam( + rotary_emb: Optional[torch.Tensor], + head_dim: int, + rope_dim: int, +) -> Optional[torch.Tensor]: + """Re-slice WanRotaryPosEmbed output to a smaller rope_dim.""" + if rotary_emb is None: + return None + orig_t_size = head_dim // 2 - 2 * (head_dim // 6) + orig_h_size = head_dim // 6 + new_t_size = rope_dim // 2 - 2 * (rope_dim // 6) + new_h_size = rope_dim // 6 + new_w_size = rope_dim // 6 + t_part = rotary_emb[..., :new_t_size] + h_part = rotary_emb[..., orig_t_size : orig_t_size + new_h_size] + w_part = rotary_emb[ + ..., orig_t_size + orig_h_size : orig_t_size + orig_h_size + new_w_size + ] + return torch.cat([t_part, h_part, w_part], dim=-1) + + +def _build_ucpe_apply_fns( + head_dim: int, + raymats: torch.Tensor, # (B, N, 4, 4) -- ray<-world + rotary_emb: Optional[torch.Tensor], +) -> Tuple[Callable, Callable, Callable]: + """Build the (apply_q, apply_kv, apply_o) callables used in the camera + branch. Splits the head_dim in two: half for 4x4 projmat tiling, half + for complex-RoPE rotation. + """ + P = raymats + P_T = P.transpose(-1, -2) + P_inv = _invert_SE3(P) + + rotary_emb_cam = _slice_rope_for_cam(rotary_emb, head_dim, head_dim // 2) + if rotary_emb_cam is not None: + + def rope_fn(x: torch.Tensor) -> torch.Tensor: + return _apply_complex_rope(x, rotary_emb_cam, inverse=False) + + def rope_fn_inv(x: torch.Tensor) -> torch.Tensor: + return _apply_complex_rope(x, rotary_emb_cam, inverse=True) + + else: + + def rope_fn(x: torch.Tensor) -> torch.Tensor: + return x + + def rope_fn_inv(x: torch.Tensor) -> torch.Tensor: + return x + + half = head_dim // 2 + + def ray_proj_t(y: torch.Tensor) -> torch.Tensor: + return _apply_ray_projmat(y, P_T) + + def ray_proj_inv(y: torch.Tensor) -> torch.Tensor: + return _apply_ray_projmat(y, P_inv) + + def ray_proj(y: torch.Tensor) -> torch.Tensor: + return _apply_ray_projmat(y, P) + + def apply_q(x: torch.Tensor) -> torch.Tensor: + return _apply_block_diagonal(x, [(ray_proj_t, half), (rope_fn, half)]) + + def apply_kv(x: torch.Tensor) -> torch.Tensor: + return _apply_block_diagonal(x, [(ray_proj_inv, half), (rope_fn, half)]) + + def apply_o(x: torch.Tensor) -> torch.Tensor: + return _apply_block_diagonal(x, [(ray_proj, half), (rope_fn_inv, half)]) + + return apply_q, apply_kv, apply_o + + +def _compute_fov_from_focal(focal: torch.Tensor, image_size: int) -> torch.Tensor: + """fov = 2 * atan(image_size / (2 * focal))""" + return 2.0 * torch.atan(image_size / (2.0 * focal.clamp(min=1e-6))) + + +def _unproject_grid( + x_fov: torch.Tensor, + y_fov: torch.Tensor, + H: int, + W: int, + cx: torch.Tensor, + cy: torch.Tensor, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Compute camera-space unit ray directions for each latent token. + + Returns (B, F, H, W, 3). + """ + B, F_dim = x_fov.shape + # Upstream `create_grid` uses integer pixel coordinates [0, W-1] / + # [0, H-1] rather than half-pixel centers. + u = torch.arange(W, device=device, dtype=dtype) + v = torch.arange(H, device=device, dtype=dtype) + u = u.view(1, 1, 1, W).expand(B, F_dim, H, W) + v = v.view(1, 1, H, 1).expand(B, F_dim, H, W) + cx_e = cx.view(B, F_dim, 1, 1) + cy_e = cy.view(B, F_dim, 1, 1) + tan_x = torch.tan(x_fov / 2.0).view(B, F_dim, 1, 1) + tan_y = torch.tan(y_fov / 2.0).view(B, F_dim, 1, 1) + # Map pixel (u, v) -> (x_dir, y_dir) on the z=1 plane. + dx = (u - cx_e) / max(W, 1) * 2.0 * tan_x + dy = (v - cy_e) / max(H, 1) * 2.0 * tan_y + dz = torch.ones_like(dx) + d = torch.stack([dx, dy, dz], dim=-1) # (B, F, H, W, 3) + return F.normalize(d, dim=-1) + + +def process_camera_conditions_ucpe( + camera_conditions: torch.Tensor, # (B, F, 20) + HW: Tuple[int, int, int], + patch_size: Tuple[int, int, int] = (1, 1, 1), +) -> torch.Tensor: + """Convert ``(B, F, 20)`` flat camera conditions into ``raymats``. + + Layout: first 16 = c2w 4x4 flatten, last 4 = ``(fx, fy, cx, cy)``. + Returns ``raymats`` of shape ``(B, F, H, W, 4, 4)`` (ray<-world). + """ + B, F_dim, _ = camera_conditions.shape + _, H, W = HW + device = camera_conditions.device + dtype = camera_conditions.dtype + + c2w_flat = camera_conditions[..., :16] + C_to_W = c2w_flat.view(B, F_dim, 4, 4) + + fx = camera_conditions[..., 16] + fy = camera_conditions[..., 17] + cx = camera_conditions[..., 18] + cy = camera_conditions[..., 19] + + image_h = H * patch_size[1] + image_w = W * patch_size[2] + x_fov = _compute_fov_from_focal(fx, image_w) + y_fov = _compute_fov_from_focal(fy, image_h) + + # cx/cy are in pixel units -- scale to latent grid like upstream + # (``cx / patch_size[2]``). + cx_lat = cx / float(patch_size[2]) + cy_lat = cy / float(patch_size[1]) + + d_cam = _unproject_grid( + x_fov, y_fov, H, W, cx_lat, cy_lat, device, dtype + ) # (B,F,H,W,3) + + # Build per-token "ray<-world" 4x4 following upstream `world_to_ray_mats`. + R_c2w = C_to_W[..., :3, :3] # (B, F, 3, 3) + t_c2w = C_to_W[..., :3, 3] # (B, F, 3) + d_world = torch.einsum("bfij,bfhwj->bfhwi", R_c2w, d_cam) + z_ray = F.normalize(d_world, dim=-1, eps=1e-6) + cam_y = R_c2w[..., :, 1].view(B, F_dim, 1, 1, 3).expand(B, F_dim, H, W, 3) + x_ray = F.normalize(torch.cross(cam_y, z_ray, dim=-1), dim=-1, eps=1e-6) + y_ray = F.normalize(torch.cross(z_ray, x_ray, dim=-1), dim=-1, eps=1e-6) + R_ray_to_world = torch.stack([x_ray, y_ray, z_ray], dim=-1) + + # P = ray<-world = inverse of [R_ray_to_world | t_c2w] + R_w_to_ray = R_ray_to_world.transpose(-1, -2) + t_w_to_ray = -torch.einsum( + "bfhwij,bfj->bfhwi", + R_w_to_ray, + t_c2w, + ) + raymats = torch.zeros(B, F_dim, H, W, 4, 4, device=device, dtype=dtype) + raymats[..., :3, :3] = R_w_to_ray + raymats[..., :3, 3] = t_w_to_ray + raymats[..., 3, 3] = 1.0 + invalid = torch.isnan(d_world).any(dim=-1) + if bool(invalid.any()): + eye = torch.eye(4, device=device, dtype=dtype) + raymats[invalid] = eye + return raymats + + +def compute_chunk_plucker( + camera_conditions: torch.Tensor, # (B, F_orig, 20) + HW: Tuple[int, int, int], # latent (T, H, W) + vae_temporal_stride: int = 8, + patch_size: Tuple[int, int, int] = (1, 1, 1), +) -> torch.Tensor: + """Compute the 48-channel packed Plücker raymap consumed by + ``plucker_embedder``. + + Official SANA-WM centers each chunk on the latent-frame timestamp: + ``0, stride, 2*stride, ...``. For timestamp 0 the chunk is clamped to the + first ``stride`` frames; for later timestamps it uses the preceding + ``stride - 1`` frames plus the current frame. Short tail chunks are padded + by repeating the final available frame. + + Each latent frame packs ``vae_temporal_stride`` original-frame Plücker + coords ``[d, o x d]`` (6D each) into 48 channels. Output shape is + ``(B, 48, T, H, W)`` for direct consumption by Conv3d. + """ + B, F_orig, _ = camera_conditions.shape + T, H, W = HW + device = camera_conditions.device + dtype = camera_conditions.dtype + + c2w = camera_conditions[..., :16].view(B, F_orig, 4, 4) + fx = camera_conditions[..., 16] + fy = camera_conditions[..., 17] + cx = camera_conditions[..., 18] + cy = camera_conditions[..., 19] + + image_h = H * patch_size[1] + image_w = W * patch_size[2] + x_fov = _compute_fov_from_focal(fx, image_w) + y_fov = _compute_fov_from_focal(fy, image_h) + + d_cam = _unproject_grid( + x_fov, + y_fov, + H, + W, + cx / float(patch_size[2]), + cy / float(patch_size[1]), + device, + dtype, + ) # (B, F_orig, H, W, 3) + + R = c2w[..., :3, :3] + o = c2w[..., :3, 3] # (B, F_orig, 3) + d_world = F.normalize(torch.einsum("bfij,bfhwj->bfhwi", R, d_cam), dim=-1) + o_exp = o.view(B, F_orig, 1, 1, 3).expand_as(d_world) + moment = torch.cross(o_exp, d_world, dim=-1) + plucker = torch.cat([d_world, moment], dim=-1) # (B, F_orig, H, W, 6) + + time_indices = torch.arange( + 0, F_orig, vae_temporal_stride, device=device, dtype=torch.long + ) + if time_indices.numel() < T: + pad = T - int(time_indices.numel()) + last = ( + time_indices[-1:] + if time_indices.numel() + else torch.zeros(1, device=device, dtype=torch.long) + ) + time_indices = torch.cat([time_indices, last.repeat(pad)], dim=0) + time_indices = time_indices[:T] + + chunks = [] + for time_index in time_indices.tolist(): + start = max(0, int(time_index) - vae_temporal_stride + 1) + end = min(start + vae_temporal_stride, F_orig) + chunk = plucker[:, start:end] + if chunk.shape[1] < vae_temporal_stride: + pad = vae_temporal_stride - chunk.shape[1] + pad_chunk = chunk[:, -1:].repeat(1, pad, 1, 1, 1) + chunk = torch.cat([chunk, pad_chunk], dim=1) + chunks.append(chunk) + + plucker = torch.stack(chunks, dim=1) # (B, T, stride, H, W, 6) + plucker = plucker.permute(0, 1, 3, 4, 2, 5).reshape( + B, T, H, W, vae_temporal_stride * 6 + ) + # (B, 48, T, H, W) for Conv3d + plucker = plucker.permute(0, 4, 1, 2, 3).contiguous() + return plucker + + +# --------------------------------------------------------------------------- +# Patch embedding / caption embedder / final layer -- upstream key names +# --------------------------------------------------------------------------- + + +class PatchEmbedMS3D(nn.Module): + """3D patch embedder used by SANA-WM for ``x_embedder``, + ``raymap_embedder``, and ``plucker_embedder``. + + Parameter names: ``proj.weight`` / ``proj.bias``. + """ + + def __init__( + self, + patch_size: Tuple[int, int, int], + in_chans: int, + embed_dim: int, + kernel_size: Optional[Tuple[int, int, int]] = None, + bias: bool = True, + ) -> None: + super().__init__() + kernel_size = kernel_size or patch_size + assert patch_size[0] == 1, "Temporal patch must be 1 for SANA-WM." + self.patch_size = patch_size + self.proj = nn.Conv3d( + in_chans, embed_dim, kernel_size=kernel_size, stride=patch_size, bias=bias + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj(x) # (B, D, T, H, W) + return x.flatten(2).transpose(1, 2) # (B, T*H*W, D) + + +class _UpstreamMlp(nn.Module): + """timm-style Mlp used by ``y_embedder.y_proj`` (fc1/fc2 with bias=True).""" + + def __init__( + self, in_features: int, hidden_features: int, out_features: int + ) -> None: + super().__init__() + self.fc1 = nn.Linear(in_features, hidden_features, bias=True) + self.act = nn.GELU(approximate="tanh") + self.fc2 = nn.Linear(hidden_features, out_features, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc2(self.act(self.fc1(x))) + + +class CaptionEmbedder(nn.Module): + """Upstream ``CaptionEmbedder``: projects text embeddings to hidden size, plus a learned null-caption table for CFG. + + Inference-only forward: ``y`` may be ``(B, 1, L, in_channels)`` or + ``(B, L, in_channels)``; the null table is not applied (no inference-time + dropout). Returns ``(B, 1, L, hidden_size)``. + """ + + def __init__( + self, in_channels: int, hidden_size: int, token_num: int = 300 + ) -> None: + super().__init__() + self.y_proj = _UpstreamMlp(in_channels, hidden_size, hidden_size) + # buffer in upstream -- registered with nn.Parameter wrapper but as buffer. + self.register_buffer( + "y_embedding", + torch.randn(token_num, in_channels) / in_channels**0.5, + persistent=True, + ) + + def forward(self, y: torch.Tensor) -> torch.Tensor: + if y.dim() == 3: + y = y.unsqueeze(1) + return self.y_proj(y) + + +class T2IFinalLayer(nn.Module): + """Output AdaLN + Linear, with the scale_shift_table living *inside* the + module (upstream stores it as ``final_layer.scale_shift_table``). + """ + + def __init__( + self, hidden_size: int, patch_size: Tuple[int, int, int], out_channels: int + ) -> None: + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear( + hidden_size, math.prod(patch_size) * out_channels, bias=True + ) + self.scale_shift_table = nn.Parameter( + torch.randn(2, hidden_size) / hidden_size**0.5 + ) + self.out_channels = out_channels + + def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + # t can be either (B, D) for scalar timesteps or (B, 1, T, D) for + # SANA-WM's first-frame-conditioned flow_euler_ltx sampler. + if t.dim() == 2: + shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1) + x = self.norm_final(x) * (1 + scale) + shift + else: + B, N, D = x.shape + t = t.reshape(B, -1, D) + num_frames = t.shape[1] + tokens_per_frame = N // num_frames + shift, scale = ( + self.scale_shift_table[None, None, :, :] + t[:, :, None, :] + ).chunk(2, dim=2) + x = self.norm_final(x).reshape(B, num_frames, tokens_per_frame, D) + x = x * (1 + scale) + shift + x = x.reshape(B, N, D) + return self.linear(x) + + +# --------------------------------------------------------------------------- +# Timestep embedder (upstream sana_blocks.TimestepEmbedder, key names +# ``t_embedder.mlp.0/2.{weight,bias}``) +# --------------------------------------------------------------------------- + + +def _sinusoidal_timestep_embedding( + t: torch.Tensor, dim: int, max_period: float = 10000.0 +) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(half, dtype=torch.float32, device=t.device) + / half + ) + args = t.float()[:, None] * freqs[None] + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1) + return emb + + +class TimestepEmbedder(nn.Module): + """Upstream ``TimestepEmbedder``. Stored as ``mlp.0/2`` (Linear, SiLU, + Linear) -- we mirror that exact layout.""" + + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256) -> None: + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + t_freq = _sinusoidal_timestep_embedding(t, self.frequency_embedding_size) + # first linear's weight dtype matches upstream + return self.mlp(t_freq.to(self.mlp[0].weight.dtype)) + + +# --------------------------------------------------------------------------- +# GLUMBConvTemp -- upstream basic_modules.GLUMBConvTemp. +# +# Stored sub-modules (all referenced by their checkpoint key prefix): +# * inverted_conv.conv Conv2d(in, hidden*2, 1) +# * depth_conv.conv Conv2d(hidden*2, hidden*2, 3, groups=hidden*2) +# * point_conv.conv Conv2d(hidden, out, 1, bias=False) +# * t_conv Conv2d(out, out, kernel=(t_k, 1), padding=(t_pad, 0), bias=False) +# +# Upstream wraps each Conv2d in a ``ConvLayer`` that owns the conv as +# ``.conv``; that's why the key has the extra ``.conv`` segment. +# --------------------------------------------------------------------------- + +_INT32_SAFE_CONV_ELEMENTS = 1 << 30 + + +class _ConvLayer(nn.Module): + """Thin wrapper -- has a single ``conv`` member to match upstream + ``inverted_conv.conv`` / ``depth_conv.conv`` / ``point_conv.conv`` keys. + """ + + def __init__(self, conv: nn.Module, act: Optional[nn.Module] = None) -> None: + super().__init__() + self.conv = conv + self.act = act + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv(x) + if self.act is not None: + x = self.act(x) + return x + + +class GLUMBConvTemp(nn.Module): + """Spatial GLU MBConv + additive temporal Conv2d (zero-init). + + Operates on ``(B, T*H*W, C)``; reshape to ``(B*T, C, H, W)`` for spatial + convs, then ``(B, C, T, H*W)`` for ``t_conv`` (a 2D depth-wise temporal + conv), residually added. + """ + + def __init__( + self, in_features: int, hidden_features: int, t_kernel_size: int = 3 + ) -> None: + super().__init__() + self.inverted_conv = _ConvLayer( + nn.Conv2d(in_features, hidden_features * 2, 1, 1, 0, bias=True), + act=nn.SiLU(inplace=False), + ) + self.depth_conv = _ConvLayer( + nn.Conv2d( + hidden_features * 2, + hidden_features * 2, + 3, + 1, + 1, + groups=hidden_features * 2, + bias=True, + ), + act=None, + ) + self.point_conv = _ConvLayer( + nn.Conv2d(hidden_features, in_features, 1, 1, 0, bias=False), + act=None, + ) + self.glu_act = nn.SiLU(inplace=False) + + t_padding = t_kernel_size // 2 + self.t_conv = nn.Conv2d( + in_features, + in_features, + kernel_size=(t_kernel_size, 1), + stride=1, + padding=(t_padding, 0), + bias=False, + ) + nn.init.zeros_(self.t_conv.weight) + + def _apply_spatial(self, x: torch.Tensor) -> torch.Tensor: + x = self.inverted_conv(x) + x = self.depth_conv(x) + a, g = x.chunk(2, dim=1) + return self.point_conv(a * self.glu_act(g)) + + def _apply_spatial_autochunked(self, x: torch.Tensor) -> torch.Tensor: + """Avoid oversized Conv2d calls on long videos while keeping short path fused.""" + BT, _, H, W = x.shape + elements_per_bt = self.inverted_conv.conv.out_channels * H * W + max_bt = max(1, _INT32_SAFE_CONV_ELEMENTS // elements_per_bt) + if BT <= max_bt: + return self._apply_spatial(x) + return torch.cat( + [ + self._apply_spatial(x[start : start + max_bt]) + for start in range(0, BT, max_bt) + ], + dim=0, + ) + + def forward( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + *, + ffn_tail: Optional[torch.Tensor] = None, + save_ffn_tail: bool = False, + ): + B, N, C = x.shape + T, H, W = HW + assert N == T * H * W, f"GLUMBConvTemp: N={N} != T*H*W={T * H * W}" + + # Spatial path is frame-local, so identical chunked or whole. + x_sp = x.reshape(B * T, H, W, C).permute(0, 3, 1, 2).contiguous() + x_sp = self._apply_spatial_autochunked(x_sp) # (B*T, C, H, W) + + x_t = ( + x_sp.view(B, T, C, H * W).permute(0, 2, 1, 3).contiguous() + ) # (B, C, T, S=H*W) + + if ffn_tail is None and not save_ffn_tail: + # Dense (bidirectional / symmetric-padding) path. + x_out = x_t + self.t_conv(x_t) + return x_out.permute(0, 2, 3, 1).reshape(B, N, C) + + # Streaming causal path (cache slot 9): prepend the previous chunk's last + # `pad` frames as real left context, then drop them from the output. The + # right edge still sees the conv's own zero-pad, so a chunk matches the + # whole-sequence pass for every frame whose right context lies in-chunk + # (always true for the final chunk). Mirrors the reference GLUMBConvTemp. + pad = self.t_conv.kernel_size[0] // 2 + conv_in = x_t + padded_size = 0 + if ffn_tail is not None: + prefix = ffn_tail.to(device=x_t.device, dtype=x_t.dtype) + conv_in = torch.cat([prefix[:, :, -pad:], x_t], dim=2) + padded_size = conv_in.shape[2] - x_t.shape[2] + new_tail = ( + x_t[:, :, -pad:].detach().clone() + if (save_ffn_tail and pad > 0) + else ffn_tail + ) + tconv_out = self.t_conv(conv_in)[:, :, padded_size:] + x_out = x_t + tconv_out + return x_out.permute(0, 2, 3, 1).reshape(B, N, C), new_tail + + +# --------------------------------------------------------------------------- +# Frame-gate and DeltaNet update rule. The recurrent and chunk-parallel forms +# are numerically equivalent. +# --------------------------------------------------------------------------- + + +def _compute_frame_gates( + x: torch.Tensor, # (B, N, C) + HW: Tuple[int, int, int], + heads: int, + beta_proj: nn.Linear, + gate_proj: nn.Linear, + dt_bias: torch.Tensor, + A_log: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Frame-level beta/decay gates. + + Returns: + beta: (B, H, T, S) in (0, 1) via sigmoid + decay: (B, H, T) in (0, 1) via exp(-A * softplus(.)) + """ + B, N, C = x.shape + T, H, W = HW + S = H * W + beta = beta_proj(x).sigmoid().reshape(B, T, S, heads).permute(0, 3, 1, 2) + x_frame = x.reshape(B, T, S, C).mean(dim=2) + a_out = gate_proj(x_frame).float() + dt = dt_bias.float().view(1, 1, -1) + A_val = A_log.float().exp().view(1, 1, -1) + decay = (-A_val * F.softplus(a_out + dt)).exp().transpose(1, 2) # (B, H, T) + return beta, decay + + +def _gdn_scan_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_rot: torch.Tensor, + k_rot: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + eps: float = 1e-6, + return_components: bool = False, +) -> torch.Tensor: + """Causal recurrent GDN scan over T. Tensors are in (B, H, D, N=T*S) layout.""" + B, H, D, N = q.shape + T = beta.shape[2] + S = N // T + + def fold(x): + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) # (B, H, T, D, S) + + q, k, v = fold(q), fold(k), fold(v) + q_rot, k_rot = fold(q_rot), fold(k_rot) + if beta.ndim == 4: + beta_e = beta.unsqueeze(3) # (B, H, T, 1, S) + else: + beta_e = beta.view(B, H, T, 1, 1) + decay_e = decay.view(B, H, T, 1, 1) + + state_kv = torch.zeros(B, H, D, D, device=q.device, dtype=q.dtype) + state_z = torch.zeros(B, H, D, 1, device=q.device, dtype=q.dtype) + num_list, den_list = [], [] + target_z = 1.0 + for t in range(T): + qt, kt, vt = q[:, :, t], k[:, :, t], v[:, :, t] + qrt, krt = q_rot[:, :, t], k_rot[:, :, t] + bt, gt = beta_e[:, :, t], decay_e[:, :, t] + state_kv = state_kv * gt + state_z = state_z * gt + v_pred = torch.matmul(state_kv, krt) + delta_v = (vt - v_pred) * bt + state_kv = state_kv + torch.matmul(delta_v, krt.transpose(-1, -2)) + z_pred = torch.matmul(state_z.transpose(-1, -2), kt) + delta_z = (target_z - z_pred) * bt + state_z = state_z + torch.matmul(kt, delta_z.transpose(-1, -2)) + num_list.append(torch.matmul(state_kv, qrt)) # (B, H, D, S) + den_list.append(torch.matmul(state_z.transpose(-1, -2), qt)) # (B, H, 1, S) + + num_stacked = torch.stack(num_list, dim=2) # (B, H, T, D, S) + den_stacked = torch.stack(den_list, dim=2) # (B, H, T, 1, S) + + def restore(tensor, d_out): + return tensor.permute(0, 1, 3, 2, 4).reshape(B, H, d_out, N) + + num = restore(num_stacked, D) + den = restore(den_stacked, 1) + if return_components: + return num, den + return num / (den + eps) + + +# --------------------------------------------------------------------------- +# Streaming (autoregressive) state-carrying scans for the `forward_long` path: +# a forward-only recurrence that SEEDS its state from a prior chunk and RETURNS +# the final state, so a long video can be generated chunk-by-chunk while +# remaining numerically identical to the monolithic forward scan. +# --------------------------------------------------------------------------- + + +def _gdn_scan_forward_stateful( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_rot: torch.Tensor, + k_rot: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + init_state_kv: Optional[torch.Tensor] = None, + init_state_z: Optional[torch.Tensor] = None, + eps: float = 1e-6, + return_components: bool = False, + return_state: bool = False, +): + """Main-branch GDN causal scan that carries KV/Z state across chunks. + + The state is seeded from + ``init_state_kv``/``init_state_z`` (None → zeros, i.e. the first chunk) and, + when ``return_state`` is set, the final ``(state_kv, state_z)`` is returned so + the next chunk continues the recurrence. Tensors are ``(B, H, D, N=T*S)``. + """ + B, H, D, N = q.shape + T = beta.shape[2] + S = N // T + + def fold(x): + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + q, k, v = fold(q), fold(k), fold(v) + q_rot, k_rot = fold(q_rot), fold(k_rot) + beta_e = beta.unsqueeze(3) if beta.ndim == 4 else beta.view(B, H, T, 1, 1) + decay_e = decay.view(B, H, T, 1, 1) + + state_kv = ( + torch.zeros(B, H, D, D, device=q.device, dtype=q.dtype) + if init_state_kv is None + else init_state_kv.to(device=q.device, dtype=q.dtype).clone() + ) + state_z = ( + torch.zeros(B, H, D, 1, device=q.device, dtype=q.dtype) + if init_state_z is None + else init_state_z.to(device=q.device, dtype=q.dtype).clone() + ) + num_list, den_list = [], [] + target_z = 1.0 + for t in range(T): + qt, kt, vt = q[:, :, t], k[:, :, t], v[:, :, t] + qrt, krt = q_rot[:, :, t], k_rot[:, :, t] + bt, gt = beta_e[:, :, t], decay_e[:, :, t] + state_kv = state_kv * gt + state_z = state_z * gt + delta_v = (vt - torch.matmul(state_kv, krt)) * bt + state_kv = state_kv + torch.matmul(delta_v, krt.transpose(-1, -2)) + delta_z = (target_z - torch.matmul(state_z.transpose(-1, -2), kt)) * bt + state_z = state_z + torch.matmul(kt, delta_z.transpose(-1, -2)) + num_list.append(torch.matmul(state_kv, qrt)) + den_list.append(torch.matmul(state_z.transpose(-1, -2), qt)) + + num = torch.stack(num_list, dim=2).permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + den = torch.stack(den_list, dim=2).permute(0, 1, 3, 2, 4).reshape(B, H, 1, N) + out = (num, den) if return_components else num / (den + eps) + if return_state: + return out, (state_kv, state_z) + return out + + +def _single_path_delta_scan_forward_stateful( + q_rot: torch.Tensor, + k_rot: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + init_state_kv: Optional[torch.Tensor] = None, + return_state: bool = False, +): + """Camera-branch (numerator-only) delta-rule scan that carries state across + chunks via a seedable / returnable ``state_kv`` for chunked autoregressive + generation. + """ + B, H, D, N = q_rot.shape + T = beta.shape[2] + S = N // T + + def fold(x): + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + q_rot, k_rot, v = fold(q_rot), fold(k_rot), fold(v) + beta_e = beta.unsqueeze(3) if beta.ndim == 4 else beta.view(B, H, T, 1, 1) + decay_e = decay.view(B, H, T, 1, 1) + + state_kv = ( + torch.zeros(B, H, D, D, device=q_rot.device, dtype=q_rot.dtype) + if init_state_kv is None + else init_state_kv.to(device=q_rot.device, dtype=q_rot.dtype).clone() + ) + out_list = [] + for t in range(T): + qrt, krt, vt = q_rot[:, :, t], k_rot[:, :, t], v[:, :, t] + bt, gt = beta_e[:, :, t], decay_e[:, :, t] + state_kv = state_kv * gt + delta_v = (vt - torch.matmul(state_kv, krt)) * bt + state_kv = state_kv + torch.matmul(delta_v, krt.transpose(-1, -2)) + out_list.append(torch.matmul(state_kv, qrt)) + + out = torch.stack(out_list, dim=2).permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + if return_state: + return out, state_kv + return out + + +def _gdn_chunk_scan_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_rot: torch.Tensor, + k_rot: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + chunk_size: Optional[int] = 21, + eps: float = 1e-6, + return_components: bool = False, +) -> torch.Tensor: + """Chunk-scan form of SANA GDN. + + Computes W/U per chunk instead of materializing all temporal transitions at + once, keeping peak memory closer to the recurrent path on long videos. + """ + B, H, D, N = q.shape + T = beta.shape[2] + S = N // T + + def fold(x: torch.Tensor) -> torch.Tensor: + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + q, k, v = fold(q), fold(k), fold(v) + q_rot, k_rot = fold(q_rot), fold(k_rot) + if beta.ndim == 4: + beta_e = beta.unsqueeze(3) + else: + beta_e = beta.view(B, H, T, 1, 1) + decay_e = decay.view(B, H, T, 1, 1) + + if chunk_size is None or int(chunk_size) <= 0: + boundaries = [0, T] + else: + boundaries = _sana_wm_normalize_chunk_index( + None, + T, + chunk_size=int(chunk_size), + chunk_split_strategy="uniform", + ) + + eye = torch.eye(D, device=q.device, dtype=q.dtype).view(1, 1, 1, D, D) + state_kv = torch.zeros(B, H, D, D, device=q.device, dtype=q.dtype) + state_z = torch.zeros(B, H, D, 1, device=q.device, dtype=q.dtype) + num_chunks, den_chunks = [], [] + + for start, end in zip(boundaries[:-1], boundaries[1:]): + q_c = q[:, :, start:end] + k_c = k[:, :, start:end] + v_c = v[:, :, start:end] + q_rot_c = q_rot[:, :, start:end] + k_rot_c = k_rot[:, :, start:end] + beta_c = beta_e[:, :, start:end] + decay_c = decay_e[:, :, start:end] + + k_rot_beta = k_rot_c * beta_c + w_kv = decay_c * (eye - torch.matmul(k_rot_beta, k_rot_c.transpose(-1, -2))) + u_kv = torch.matmul(v_c * beta_c, k_rot_c.transpose(-1, -2)) + + k_beta = k_c * beta_c + w_z = decay_c * (eye - torch.matmul(k_beta, k_c.transpose(-1, -2))) + u_z = k_beta.sum(dim=-1, keepdim=True) + + state_kv_frames, state_z_frames = [], [] + for offset in range(end - start): + state_kv = torch.matmul(state_kv, w_kv[:, :, offset]) + u_kv[:, :, offset] + state_z = torch.matmul(w_z[:, :, offset], state_z) + u_z[:, :, offset] + state_kv_frames.append(state_kv) + state_z_frames.append(state_z) + + state_kv_all = torch.stack(state_kv_frames, dim=2) + state_z_all = torch.stack(state_z_frames, dim=2) + num_chunks.append(torch.matmul(state_kv_all, q_rot_c)) + den_chunks.append(torch.matmul(state_z_all.transpose(-1, -2), q_c)) + + num = torch.cat(num_chunks, dim=2) + den = torch.cat(den_chunks, dim=2) + + def restore(tensor: torch.Tensor, d_out: int) -> torch.Tensor: + return tensor.permute(0, 1, 3, 2, 4).reshape(B, H, d_out, N) + + num = restore(num, D) + den = restore(den, 1) + if return_components: + return num, den + return num / (den + eps) + + +def _flip_and_shift(x: torch.Tensor, dim: int, shift_val: float = 0.0) -> torch.Tensor: + """Flip along ``dim`` then shift by one with ``shift_val`` filling the head. + Used for the backward pass of bidirectional GDN.""" + x_flipped = x.flip(dim) + idx = [slice(None)] * x.ndim + idx[dim] = slice(0, 1) + head = torch.full_like(x_flipped[tuple(idx)], shift_val) + idx[dim] = slice(0, -1) + return torch.cat([head, x_flipped[tuple(idx)]], dim=dim) + + +def _gdn_scan_bidirectional( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_rot: torch.Tensor, + k_rot: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + HW: Tuple[int, int, int], + chunk_size: Optional[int] = None, + eps: float = 1e-6, +) -> torch.Tensor: + """Bidirectional GDN: forward (inclusive) + backward (exclusive) scan, summed in numerator/denominator space.""" + + def run_scan( + q_in: torch.Tensor, + k_in: torch.Tensor, + v_in: torch.Tensor, + q_rot_in: torch.Tensor, + k_rot_in: torch.Tensor, + beta_in: torch.Tensor, + decay_in: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if chunk_size is None: + return _gdn_scan_forward( + q_in, + k_in, + v_in, + q_rot_in, + k_rot_in, + beta_in, + decay_in, + eps=eps, + return_components=True, + ) + return _gdn_chunk_scan_forward( + q_in, + k_in, + v_in, + q_rot_in, + k_rot_in, + beta_in, + decay_in, + chunk_size=chunk_size, + eps=eps, + return_components=True, + ) + + num_fwd, den_fwd = run_scan(q, k, v, q_rot, k_rot, beta, decay) + + # Backward pass: flip Q, flip+shift K/V/k_rot/beta and shift decay-by-1. + B, H, D, N = q.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + + def to_time(x, d): + return x.view(B, H, d, T, S).permute(0, 1, 3, 2, 4) + + def from_time(x, d): + return x.permute(0, 1, 3, 2, 4).reshape(B, H, d, N) + + q_t = to_time(q, D) + k_t = to_time(k, D) + v_t = to_time(v, D) + q_rot_t = to_time(q_rot, D) + k_rot_t = to_time(k_rot, D) + + q_bwd = torch.flip(q_t, dims=[2]) + q_rot_bwd = torch.flip(q_rot_t, dims=[2]) + k_bwd = _flip_and_shift(k_t, dim=2, shift_val=0.0) + v_bwd = _flip_and_shift(v_t, dim=2, shift_val=0.0) + k_rot_bwd = _flip_and_shift(k_rot_t, dim=2, shift_val=0.0) + beta_bwd = _flip_and_shift(beta, dim=2, shift_val=0.0) + decay_bwd = _flip_and_shift(decay, dim=2, shift_val=1.0) + + q_bwd_f = from_time(q_bwd, D) + k_bwd_f = from_time(k_bwd, D) + v_bwd_f = from_time(v_bwd, D) + q_rot_bwd_f = from_time(q_rot_bwd, D) + k_rot_bwd_f = from_time(k_rot_bwd, D) + + num_bwd_flipped, den_bwd_flipped = run_scan( + q_bwd_f, + k_bwd_f, + v_bwd_f, + q_rot_bwd_f, + k_rot_bwd_f, + beta_bwd, + decay_bwd, + ) + + def flip_back(tensor): + d_actual = tensor.shape[2] + t_struct = tensor.view(B, H, d_actual, T, S) + return torch.flip(t_struct, dims=[3]).reshape(B, H, d_actual, N) + + num_bwd = flip_back(num_bwd_flipped) + den_bwd = flip_back(den_bwd_flipped) + return (num_fwd + num_bwd) / (den_fwd + den_bwd + eps) + + +def _single_path_delta_scan_forward( + q_rot: torch.Tensor, + k_rot: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, +) -> torch.Tensor: + """Numerator-only camera delta-rule recurrence. + + This is intentionally separate from ``_gdn_scan_forward``: the SANA-WM + camera branch does not use the GDN denominator path; reusing the + main-branch GDN recurrence here changes the latent distribution + substantially once camera conditioning is enabled. + """ + B, H, D, N = q_rot.shape + T = beta.shape[2] + S = N // T + + def fold(x: torch.Tensor) -> torch.Tensor: + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + q_rot = fold(q_rot) + k_rot = fold(k_rot) + v = fold(v) + + if beta.ndim == 4: + beta_e = beta.unsqueeze(3) + else: + beta_e = beta.view(B, H, T, 1, 1) + decay_e = decay.view(B, H, T, 1, 1) + + state_kv = torch.zeros(B, H, D, D, device=q_rot.device, dtype=q_rot.dtype) + out_list = [] + for t in range(T): + qrt = q_rot[:, :, t] + krt = k_rot[:, :, t] + vt = v[:, :, t] + bt = beta_e[:, :, t] + gt = decay_e[:, :, t] + + state_kv = state_kv * gt + v_pred = torch.matmul(state_kv, krt) + delta_v = (vt - v_pred) * bt + state_kv = state_kv + torch.matmul(delta_v, krt.transpose(-1, -2)) + out_list.append(torch.matmul(state_kv, qrt)) + + out = torch.stack(out_list, dim=2) + return out.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + + +def _single_path_delta_chunk_scan_forward( + q_rot: torch.Tensor, + k_rot: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + chunk_size: Optional[int] = 21, +) -> torch.Tensor: + """Chunk-scan form of the camera single-path delta rule.""" + B, H, D, N = q_rot.shape + T = beta.shape[2] + S = N // T + + def fold(x: torch.Tensor) -> torch.Tensor: + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + q_rot = fold(q_rot) + k_rot = fold(k_rot) + v = fold(v) + if beta.ndim == 4: + beta_e = beta.unsqueeze(3) + else: + beta_e = beta.view(B, H, T, 1, 1) + decay_e = decay.view(B, H, T, 1, 1) + + if chunk_size is None or int(chunk_size) <= 0: + boundaries = [0, T] + else: + boundaries = _sana_wm_normalize_chunk_index( + None, + T, + chunk_size=int(chunk_size), + chunk_split_strategy="uniform", + ) + + eye = torch.eye(D, device=q_rot.device, dtype=q_rot.dtype).view(1, 1, 1, D, D) + state_kv = torch.zeros(B, H, D, D, device=q_rot.device, dtype=q_rot.dtype) + out_chunks = [] + + for start, end in zip(boundaries[:-1], boundaries[1:]): + q_rot_c = q_rot[:, :, start:end] + k_rot_c = k_rot[:, :, start:end] + v_c = v[:, :, start:end] + beta_c = beta_e[:, :, start:end] + decay_c = decay_e[:, :, start:end] + + k_rot_beta = k_rot_c * beta_c + w_kv = decay_c * (eye - torch.matmul(k_rot_beta, k_rot_c.transpose(-1, -2))) + u_kv = torch.matmul(v_c * beta_c, k_rot_c.transpose(-1, -2)) + + state_frames = [] + for offset in range(end - start): + state_kv = torch.matmul(state_kv, w_kv[:, :, offset]) + u_kv[:, :, offset] + state_frames.append(state_kv) + + state_all = torch.stack(state_frames, dim=2) + out_chunks.append(torch.matmul(state_all, q_rot_c)) + + out = torch.cat(out_chunks, dim=2) + return out.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + + +def _single_path_delta_scan_bidirectional( + q_rot: torch.Tensor, + k_rot: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + HW: Tuple[int, int, int], + chunk_size: Optional[int] = None, +) -> torch.Tensor: + """Bidirectional single-path camera scan: inclusive forward + exclusive backward (``flip_and_shift`` on K/V, beta, decay).""" + scan_forward = ( + _single_path_delta_chunk_scan_forward + if chunk_size is not None + else _single_path_delta_scan_forward + ) + if chunk_size is None: + out_fwd = scan_forward(q_rot, k_rot, v, beta, decay) + else: + out_fwd = scan_forward(q_rot, k_rot, v, beta, decay, chunk_size=chunk_size) + + B, H, D, N = q_rot.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + + def to_time(x: torch.Tensor) -> torch.Tensor: + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + def from_time(x: torch.Tensor) -> torch.Tensor: + return x.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + + q_rot_t = to_time(q_rot) + k_rot_t = to_time(k_rot) + v_t = to_time(v) + + q_rot_bwd = torch.flip(q_rot_t, dims=[2]) + k_rot_bwd = _flip_and_shift(k_rot_t, dim=2, shift_val=0.0) + v_bwd = _flip_and_shift(v_t, dim=2, shift_val=0.0) + beta_bwd = _flip_and_shift(beta, dim=2, shift_val=0.0) + decay_bwd = _flip_and_shift(decay, dim=2, shift_val=1.0) + + if chunk_size is None: + out_bwd_flipped = scan_forward( + from_time(q_rot_bwd), + from_time(k_rot_bwd), + from_time(v_bwd), + beta_bwd, + decay_bwd, + ) + else: + out_bwd_flipped = scan_forward( + from_time(q_rot_bwd), + from_time(k_rot_bwd), + from_time(v_bwd), + beta_bwd, + decay_bwd, + chunk_size=chunk_size, + ) + out_bwd = torch.flip( + out_bwd_flipped.view(B, H, D, T, S), + dims=[3], + ).reshape(B, H, D, N) + return out_fwd + out_bwd + + +# --------------------------------------------------------------------------- +# Chunk-causal cached scans for the streaming `forward_long` path: the FORWARD +# (inclusive) pass carries recurrent state across chunks; the BACKWARD +# (exclusive) pass is recomputed intra-chunk and stateless. A single chunk with +# no carried state reduces exactly to the bidirectional scans, while a sequence +# processed chunk-by-chunk stays continuous in the forward direction. +# --------------------------------------------------------------------------- + + +def _gdn_scan_cached( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_rot: torch.Tensor, + k_rot: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + init_state_kv: Optional[torch.Tensor] = None, + init_state_z: Optional[torch.Tensor] = None, + eps: float = 1e-6, +) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Chunk-causal main-branch GDN scan for streaming `forward_long`. + + The forward pass seeds/returns ``(state_kv, state_z)`` so chunks stay + continuous; the backward pass is intra-chunk and stateless. Returns + ``(out, (state_kv, state_z))``. + """ + (num_fwd, den_fwd), (state_kv, state_z) = _gdn_scan_forward_stateful( + q, + k, + v, + q_rot, + k_rot, + beta, + decay, + init_state_kv=init_state_kv, + init_state_z=init_state_z, + eps=eps, + return_components=True, + return_state=True, + ) + + B, H, D, N = q.shape + T = beta.shape[2] + S = N // T + + def to_time(x, d): + return x.view(B, H, d, T, S).permute(0, 1, 3, 2, 4) + + def from_time(x, d): + return x.permute(0, 1, 3, 2, 4).reshape(B, H, d, N) + + q_bwd = from_time(torch.flip(to_time(q, D), dims=[2]), D) + q_rot_bwd = from_time(torch.flip(to_time(q_rot, D), dims=[2]), D) + k_bwd = from_time(_flip_and_shift(to_time(k, D), dim=2, shift_val=0.0), D) + v_bwd = from_time(_flip_and_shift(to_time(v, D), dim=2, shift_val=0.0), D) + k_rot_bwd = from_time(_flip_and_shift(to_time(k_rot, D), dim=2, shift_val=0.0), D) + beta_bwd = _flip_and_shift(beta, dim=2, shift_val=0.0) + decay_bwd = _flip_and_shift(decay, dim=2, shift_val=1.0) + + num_bwd_flipped, den_bwd_flipped = _gdn_scan_forward( + q_bwd, + k_bwd, + v_bwd, + q_rot_bwd, + k_rot_bwd, + beta_bwd, + decay_bwd, + eps=eps, + return_components=True, + ) + + def flip_back(tensor, d): + return torch.flip(tensor.view(B, H, d, T, S), dims=[3]).reshape(B, H, d, N) + + num_bwd = flip_back(num_bwd_flipped, D) + den_bwd = flip_back(den_bwd_flipped, 1) + out = (num_fwd + num_bwd) / (den_fwd + den_bwd + eps) + return out, (state_kv, state_z) + + +def _single_path_delta_scan_cached( + q_rot: torch.Tensor, + k_rot: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + init_state_kv: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Chunk-causal camera single-path delta scan for streaming `forward_long`. + + The forward pass carries ``state_kv`` across chunks; the backward pass is + intra-chunk/stateless. Returns ``(out, state_kv)``. + """ + out_fwd, state_kv = _single_path_delta_scan_forward_stateful( + q_rot, + k_rot, + v, + beta, + decay, + init_state_kv=init_state_kv, + return_state=True, + ) + + B, H, D, N = q_rot.shape + T = beta.shape[2] + S = N // T + + def to_time(x): + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + def from_time(x): + return x.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + + q_rot_bwd = from_time(torch.flip(to_time(q_rot), dims=[2])) + k_rot_bwd = from_time(_flip_and_shift(to_time(k_rot), dim=2, shift_val=0.0)) + v_bwd = from_time(_flip_and_shift(to_time(v), dim=2, shift_val=0.0)) + beta_bwd = _flip_and_shift(beta, dim=2, shift_val=0.0) + decay_bwd = _flip_and_shift(decay, dim=2, shift_val=1.0) + + out_bwd_flipped = _single_path_delta_scan_forward( + q_rot_bwd, + k_rot_bwd, + v_bwd, + beta_bwd, + decay_bwd, + ) + out_bwd = torch.flip(out_bwd_flipped.view(B, H, D, T, S), dims=[3]).reshape( + B, H, D, N + ) + return out_fwd + out_bwd, state_kv + + +def _downscale_to_reference_rms( + ref: torch.Tensor, + transformed: torch.Tensor, + eps: float = 1e-6, +) -> torch.Tensor: + """Clamp UCPE-transformed channel RMS to the pre-transform envelope. + + The UCPE matrices include translations that can inflate transformed Q/K/V + magnitudes; downscale per (batch, head, token) before the camera recurrence + so inference stays on the training distribution. + """ + ref_rms = ref.square().mean(dim=2, keepdim=True).add(eps).sqrt() + transformed_rms = transformed.square().mean(dim=2, keepdim=True).add(eps).sqrt() + scale = (ref_rms / transformed_rms.clamp_min(eps)).clamp(max=1.0) + return transformed * scale + + +# --------------------------------------------------------------------------- +# SANA-WM attention block: main GDN (or softmax) + UCPE camera branch. The cam +# branch shares ``proj`` / ``output_gate`` with the main branch. +# --------------------------------------------------------------------------- + + +class BidirectionalGDNUCPESinglePathLiteLA(nn.Module): + """Bidirectional GDN main branch + UCPE camera branch (single-path + output: ``main + out_proj_cam(cam_raw)`` then shared output gate + + shared output projection). + """ + + def __init__( + self, + in_dim: int, + heads: int, + head_dim: int, + qk_norm: bool = True, + conv_kernel_size: int = 4, + k_conv_only: bool = True, + eps: float = 1e-8, + softmax_main: bool = False, + update_rule: str = "torch_chunk", + cam_update_rule: str = "torch_chunk", + chunk_gdn_chunk_size: int = 21, + use_chunked_softmax_attention: bool = False, + gdn_backend: str = "auto", + ) -> None: + super().__init__() + out_dim = heads * head_dim + assert ( + out_dim == in_dim + ), f"in_dim ({in_dim}) must equal heads*head_dim ({out_dim})" + self.in_dim = in_dim + self.out_dim = out_dim + self.heads = heads + self.dim = head_dim + self.eps = eps + self.softmax_main = softmax_main + self.update_rule = update_rule + self.cam_update_rule = cam_update_rule + self.chunk_gdn_chunk_size = chunk_gdn_chunk_size + self.use_chunked_softmax_attention = use_chunked_softmax_attention + self.gdn_backend = gdn_backend + if self.update_rule not in ("torch_chunk", "torch_recurrent"): + raise ValueError(f"Unsupported SANA-WM update_rule: {self.update_rule}") + if self.cam_update_rule not in ("torch_chunk", "torch_recurrent"): + raise ValueError( + f"Unsupported SANA-WM cam_update_rule: {self.cam_update_rule}" + ) + if self.gdn_backend not in ("auto", "torch", "triton"): + raise ValueError( + "Unsupported SANA-WM gdn_backend: " + f"{self.gdn_backend}. Expected one of auto, torch, triton." + ) + + # Fused QKV + output proj (proj shared with cam branch). + self.qkv = nn.Linear(in_dim, 3 * out_dim, bias=False) + self.proj = nn.Linear(out_dim, out_dim, bias=True) + + if qk_norm: + self.q_norm = _RMSNorm(in_dim, eps=1e-5) + self.k_norm = _RMSNorm(in_dim, eps=1e-5) + self.q_norm_cam = _RMSNorm(in_dim, eps=1e-5) + self.k_norm_cam = _RMSNorm(in_dim, eps=1e-5) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + self.q_norm_cam = nn.Identity() + self.k_norm_cam = nn.Identity() + + # Also held by the softmax variant for state_dict compat. + self.beta_proj = nn.Linear(in_dim, heads, bias=True) + self.gate_proj = nn.Linear(in_dim, heads, bias=True) + self.A_log = nn.Parameter(torch.log(torch.empty(heads).uniform_(0, 16))) + self.dt_bias = nn.Parameter(torch.full((heads,), -5.0)) + self.register_buffer("recall_gate", torch.zeros(1)) + self.output_gate = nn.Linear(in_dim, out_dim, bias=True) + + # Short convs on K (k_conv_only=True). The upstream softmax variant + # ``BidirectionalSoftmaxUCPESinglePathLiteLA`` does NOT carry these + # short convs in either the main or the cam branch, so the released + # checkpoint has no conv_k/conv_k_cam tensors for softmax blocks + # (indices {3, 7, 11, 15, 19} when softmax_every_n=4). Creating them + # here would leave the loader staring at missing checkpoint keys it + # cannot synthesize. Skip creation when softmax_main=True. + if conv_kernel_size > 0 and not softmax_main: + self.conv_k = _ShortConvolution(out_dim, conv_kernel_size) + if not k_conv_only: + self.conv_q = _ShortConvolution(out_dim, conv_kernel_size) + self.conv_v = _ShortConvolution(out_dim, conv_kernel_size) + else: + self.conv_q = None + self.conv_v = None + else: + self.conv_k = self.conv_q = self.conv_v = None + self.conv_kernel_size = conv_kernel_size + self.k_conv_only = k_conv_only + + # Camera branch: separate QKV + zero-init output proj + separate K + # ShortConvolution. Both branches share ``proj``. + self.q_proj_cam = nn.Linear(in_dim, out_dim, bias=True) + self.k_proj_cam = nn.Linear(in_dim, out_dim, bias=True) + self.v_proj_cam = nn.Linear(in_dim, out_dim, bias=True) + self.out_proj_cam = nn.Linear(out_dim, out_dim, bias=True) + nn.init.zeros_(self.out_proj_cam.weight) + nn.init.zeros_(self.out_proj_cam.bias) + if conv_kernel_size > 0 and not softmax_main: + self.conv_k_cam = _ShortConvolution(out_dim, conv_kernel_size) + if not k_conv_only: + self.conv_q_cam = _ShortConvolution(out_dim, conv_kernel_size) + self.conv_v_cam = _ShortConvolution(out_dim, conv_kernel_size) + else: + self.conv_q_cam = None + self.conv_v_cam = None + else: + self.conv_k_cam = self.conv_q_cam = self.conv_v_cam = None + + # Softmax-variant blocks route attention through SGLang's pluggable + # backend (FA3 / FlashInfer / Triton / SDPA). GDN blocks compute + # attention via the scan path and don't go through this. + if softmax_main: + self.softmax_attn = LocalAttention( + num_heads=heads, + head_size=head_dim, + ) + self._triton_rope_tables_cache: Optional[ + Tuple[Tuple, Tuple[torch.Tensor, torch.Tensor]] + ] = None + self._triton_norm_weights_cache: Optional[ + Tuple[Tuple, Tuple[torch.Tensor, torch.Tensor]] + ] = None + self._cam_qkv_params_cache: Optional[ + Tuple[Tuple, Tuple[torch.Tensor, torch.Tensor]] + ] = None + + # ------------------------------------------------------------------ # + + @staticmethod + def _temporal_short_conv( + x: torch.Tensor, # (B, N, C) + conv: _ShortConvolution, + HW: Tuple[int, int, int], + bidirectional: bool = True, + ) -> torch.Tensor: + B, N, C = x.shape + T, H, W = HW + S = H * W + # (B, T, S, C) -> (B*S, T, C): T onto the time axis for the conv. + y = x.view(B, T, S, C).permute(0, 2, 1, 3).contiguous().reshape(B * S, T, C) + if bidirectional: + y = _bidirectional_short_conv(y, conv) + else: + y, _ = conv(y) + # back to (B, T*S, C) + return y.reshape(B, S, T, C).permute(0, 2, 1, 3).reshape(B, N, C) + + # ------------------------------------------------------------------ # + + def _get_cam_qkv_params(self) -> Tuple[torch.Tensor, torch.Tensor]: + weights = ( + self.q_proj_cam.weight, + self.k_proj_cam.weight, + self.v_proj_cam.weight, + ) + biases = ( + self.q_proj_cam.bias, + self.k_proj_cam.bias, + self.v_proj_cam.bias, + ) + if torch.is_grad_enabled(): + return torch.cat(weights, dim=0), torch.cat(biases, dim=0) + + key = ( + "cam_qkv_params", + tuple(_tensor_cache_key(weight) for weight in weights), + tuple(_tensor_cache_key(bias) for bias in biases), + ) + cached = self._cam_qkv_params_cache + if cached is not None and cached[0] == key: + return cached[1] + + params = ( + torch.cat(weights, dim=0).contiguous(), + torch.cat(biases, dim=0).contiguous(), + ) + self._cam_qkv_params_cache = (key, params) + return params + + def _cam_qkv( + self, x: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + qkv_weight, qkv_bias = self._get_cam_qkv_params() + return F.linear(x, qkv_weight, qkv_bias).chunk(3, dim=-1) + + # ------------------------------------------------------------------ # + + def _triton_gdn_unavailable_reason( + self, + qkv: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + HW: Tuple[int, int, int], + ) -> Optional[str]: + if self.gdn_backend == "torch": + return "gdn_backend=torch" + if _SANA_WM_TRITON_GDN_DISABLED_REASON is not None: + return _SANA_WM_TRITON_GDN_DISABLED_REASON + if self.training or torch.is_grad_enabled(): + return "requires eval/inference mode" + if not qkv.is_cuda: + return "requires CUDA tensor" + if qkv.dtype not in (torch.float16, torch.bfloat16): + return f"requires fp16/bf16 qkv, got {qkv.dtype}" + if not qkv.is_contiguous(): + return "qkv must be contiguous" + if beta.ndim != 4: + return f"requires beta shape (B, H, T, S), got {tuple(beta.shape)}" + + B, N, three, heads, head_dim = qkv.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + if three != 3: + return f"requires qkv third dim=3, got {three}" + if N != T * S: + return f"requires N=T*S, got N={N}, T*S={T * S}" + if head_dim > 128: + return f"requires head_dim <= 128, got {head_dim}" + if beta.shape != (B, heads, T, S): + return f"requires beta shape {(B, heads, T, S)}, got {tuple(beta.shape)}" + if decay.shape != (B, heads, T): + return f"requires decay shape {(B, heads, T)}, got {tuple(decay.shape)}" + if not hasattr(self.q_norm, "weight") or not hasattr(self.k_norm, "weight"): + return "requires learned q/k RMSNorm weights" + return None + + def _get_triton_norm_weights(self) -> Tuple[torch.Tensor, torch.Tensor]: + q_weight = self.q_norm.weight + k_weight = self.k_norm.weight + key = ( + "triton_norm_weights", + _tensor_cache_key(q_weight), + _tensor_cache_key(k_weight), + ) + cached = self._triton_norm_weights_cache + if cached is not None and cached[0] == key: + return cached[1] + + weights = ( + q_weight.float().contiguous(), + k_weight.float().contiguous(), + ) + self._triton_norm_weights_cache = (key, weights) + return weights + + def _get_triton_rope_tables( + self, + prepare_rope_tables: Callable, + rotary_emb: Optional[torch.Tensor], + *, + N: int, + head_dim: int, + device: torch.device, + ) -> Tuple[torch.Tensor, torch.Tensor]: + key = ( + "triton_rope_tables", + N, + head_dim, + str(device), + None if rotary_emb is None else _tensor_cache_key(rotary_emb), + ) + cached = self._triton_rope_tables_cache + if cached is not None and cached[0] == key: + return cached[1] + + tables = prepare_rope_tables(rotary_emb, N, head_dim, device) + self._triton_rope_tables_cache = (key, tables) + return tables + + def _maybe_main_branch_triton_gdn( + self, + qkv: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + HW: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor], + ) -> Optional[torch.Tensor]: + global _SANA_WM_TRITON_GDN_DISABLED_REASON + + reason = self._triton_gdn_unavailable_reason(qkv, beta, decay, HW) + if reason is not None: + if self.gdn_backend == "triton": + raise RuntimeError(f"SANA-WM Triton GDN backend unavailable: {reason}") + return None + + try: + from sglang.jit_kernel.diffusion.triton.sana_wm_gdn import ( + fused_bigdn_func, + fused_qk_inv_rms, + prepare_rope_tables, + ) + + B, N, _, heads, head_dim = qkv.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + q_norm_weight, k_norm_weight = self._get_triton_norm_weights() + norm_eps = float(getattr(self.q_norm, "eps", 1e-5)) + q_inv_rms, k_inv_rms = fused_qk_inv_rms(qkv, eps=norm_eps) + rope_cos, rope_sin = self._get_triton_rope_tables( + prepare_rope_tables, + rotary_emb, + N=N, + head_dim=head_dim, + device=qkv.device, + ) + out = fused_bigdn_func( + qkv, + q_inv_rms, + k_inv_rms, + q_norm_weight=q_norm_weight, + k_norm_weight=k_norm_weight, + rope_cos=rope_cos, + rope_sin=rope_sin, + beta=beta.contiguous(), + decay=decay.contiguous(), + F=T, + S=S, + k_scale=(head_dim**-0.5) * (S**-0.5), + eps=self.eps, + ) + return out.reshape(B, N, heads * head_dim) + except Exception as exc: + if self.gdn_backend == "triton": + raise + _SANA_WM_TRITON_GDN_DISABLED_REASON = str(exc) + _log_sana_wm_triton_gdn_fallback(str(exc)) + return None + + # ------------------------------------------------------------------ # + + def _triton_cam_gdn_unavailable_reason( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + HW: Tuple[int, int, int], + ) -> Optional[str]: + if self.gdn_backend == "torch": + return "gdn_backend=torch" + if _SANA_WM_TRITON_CAM_GDN_DISABLED_REASON is not None: + return _SANA_WM_TRITON_CAM_GDN_DISABLED_REASON + if self.training or torch.is_grad_enabled(): + return "requires eval/inference mode" + if not q.is_cuda: + return "requires CUDA tensor" + if ( + q.dtype != torch.float32 + or k.dtype != torch.float32 + or v.dtype != torch.float32 + ): + return f"requires fp32 q/k/v, got {q.dtype}/{k.dtype}/{v.dtype}" + if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous(): + return "q/k/v must be contiguous" + if beta.ndim not in (3, 4): + return f"requires beta rank 3 or 4, got {tuple(beta.shape)}" + + B, heads, head_dim, N = q.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + if k.shape != q.shape or v.shape != q.shape: + return f"q/k/v shape mismatch: {q.shape}/{k.shape}/{v.shape}" + if N != T * S: + return f"requires N=T*S, got N={N}, T*S={T * S}" + if beta.ndim == 3 and beta.shape != (B, heads, T): + return f"requires beta shape {(B, heads, T)}, got {tuple(beta.shape)}" + if beta.ndim == 4 and beta.shape != (B, heads, T, S): + return ( + f"requires beta shape {(B, heads, T, S)}, " f"got {tuple(beta.shape)}" + ) + if decay.shape != (B, heads, T): + return f"requires decay shape {(B, heads, T)}, got {tuple(decay.shape)}" + if head_dim > 128: + return f"requires head_dim <= 128, got {head_dim}" + return None + + def _maybe_cam_branch_triton_scan( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + HW: Tuple[int, int, int], + ) -> Optional[torch.Tensor]: + global _SANA_WM_TRITON_CAM_GDN_DISABLED_REASON + + precheck_reason = None + if self.gdn_backend == "torch": + precheck_reason = "gdn_backend=torch" + elif _SANA_WM_TRITON_CAM_GDN_DISABLED_REASON is not None: + precheck_reason = _SANA_WM_TRITON_CAM_GDN_DISABLED_REASON + elif self.training or torch.is_grad_enabled(): + precheck_reason = "requires eval/inference mode" + elif not q.is_cuda: + precheck_reason = "requires CUDA tensor" + + if precheck_reason is not None: + if self.gdn_backend == "triton": + raise RuntimeError( + "SANA-WM Triton camera GDN backend unavailable: " + f"{precheck_reason}" + ) + return None + + q = q.float().contiguous() + k = k.float().contiguous() + v = v.float().contiguous() + reason = self._triton_cam_gdn_unavailable_reason(q, k, v, beta, decay, HW) + if reason is not None: + if self.gdn_backend == "triton": + raise RuntimeError( + f"SANA-WM Triton camera GDN backend unavailable: {reason}" + ) + return None + + try: + from sglang.jit_kernel.diffusion.triton.sana_wm_gdn_chunkwise import ( + cam_scan_bidi_chunkwise, + ) + + B, heads, _, _ = q.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + if beta.ndim == 3: + beta_in = beta.unsqueeze(-1).expand(B, heads, T, S).contiguous() + else: + beta_in = beta.contiguous() + out = cam_scan_bidi_chunkwise( + q, + k, + v, + beta_in.float(), + decay.float().contiguous(), + ) + return out + except Exception as exc: + if self.gdn_backend == "triton": + raise + _SANA_WM_TRITON_CAM_GDN_DISABLED_REASON = str(exc) + _log_sana_wm_triton_cam_gdn_fallback(str(exc)) + return None + + # ------------------------------------------------------------------ # + + def _main_branch_gdn( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return ``(out_raw, beta, decay)`` where ``out_raw`` is the GDN + scan result before output gate / proj. + """ + B, N, C = x.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + + qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.dim).contiguous() + q, k, v = qkv.unbind(2) + + if self.conv_k is not None: + k = self._temporal_short_conv( + k.reshape(B, N, C), self.conv_k, HW, bidirectional=True + ) + k = k.reshape(B, N, self.heads, self.dim) + qkv = torch.stack((q, k, v), dim=2).contiguous() + + beta, decay = _compute_frame_gates( + x, + HW, + self.heads, + self.beta_proj, + self.gate_proj, + self.dt_bias, + self.A_log, + ) + + triton_out = self._maybe_main_branch_triton_gdn( + qkv, + beta, + decay, + HW, + rotary_emb, + ) + if triton_out is not None: + return triton_out, beta, decay + + q, k, v = qkv.unbind(2) + + q = self.q_norm(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + + # ReLU kernel + key scale + q = F.relu(q) + k = F.relu(k) + k_scale = (self.dim**-0.5) * (S**-0.5) + k = k * k_scale + + # (B, H, D, N) + q = q.permute(0, 2, 3, 1) + k = k.permute(0, 2, 3, 1) + v = v.permute(0, 2, 3, 1) + + if rotary_emb is not None: + q_rot = _apply_rotary_emb_dn(q, rotary_emb) + k_rot = _apply_rotary_emb_dn(k, rotary_emb) + else: + q_rot = q + k_rot = k + + # fp32 scan for stability + dtype = q.dtype + scan_chunk_size = ( + self.chunk_gdn_chunk_size if self.update_rule == "torch_chunk" else None + ) + out = _gdn_scan_bidirectional( + q.float(), + k.float(), + v.float(), + q_rot.float(), + k_rot.float(), + beta.float(), + decay.float(), + HW=HW, + chunk_size=scan_chunk_size, + eps=self.eps, + ).to(dtype) + + out = out.permute(0, 3, 1, 2).reshape(B, N, C) + return out, beta, decay + + def _main_branch_softmax( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor], + chunk_size: Optional[int] = None, + chunk_split_strategy: str = "uniform", + chunk_index: Optional[List[int]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Softmax variant of the main branch, via SGLang's pluggable attention backend. + + Returns ``(out_raw, beta, decay)`` so the cam branch can reuse the shared + gates -- like upstream ``BidirectionalSoftmaxUCPESinglePathLiteLA``. + """ + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.dim) + q, k, v = qkv.unbind(2) + if self.conv_k is not None: + k = self._temporal_short_conv( + k.reshape(B, N, C), self.conv_k, HW, bidirectional=True + ) + k = k.reshape(B, N, self.heads, self.dim) + q = self.q_norm(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + # RoPE primitives use (B, H, N, D); LocalAttention takes (B, N, H, D). + q = q.permute(0, 2, 1, 3) # (B, H, N, D) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + if rotary_emb is not None: + q = _apply_rotary_emb_bhnd(q, rotary_emb) + k = _apply_rotary_emb_bhnd(k, rotary_emb) + q_in = q.transpose(1, 2).contiguous() + k_in = k.transpose(1, 2).contiguous() + v_in = v.transpose(1, 2).contiguous() + out = None + if self.use_chunked_softmax_attention: + out = _sana_wm_chunked_attention( + q_in, + k_in, + v_in, + HW=HW, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + chunk_index=chunk_index, + softmax_scale=self.softmax_attn.softmax_scale, + ) + if out is None: + out = self.softmax_attn(q_in, k_in, v_in) + out = out.reshape(B, N, C) + + # Gates are needed by the cam branch and also exist in the softmax + # variant's state dict. + beta, decay = _compute_frame_gates( + x, + HW, + self.heads, + self.beta_proj, + self.gate_proj, + self.dt_bias, + self.A_log, + ) + return out, beta, decay + + # ------------------------------------------------------------------ # + + def _cam_branch( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + apply_q: Callable, + apply_kv: Callable, + apply_o: Callable, + beta: torch.Tensor, + decay: torch.Tensor, + ) -> torch.Tensor: + B, N, C = x.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + + q, k, v = self._cam_qkv(x) + q = q.reshape(B, N, self.heads, self.dim) + k = k.reshape(B, N, self.heads, self.dim) + v = v.reshape(B, N, self.heads, self.dim) + + if self.conv_k_cam is not None: + k = self._temporal_short_conv( + k.reshape(B, N, C), self.conv_k_cam, HW, bidirectional=True + ) + k = k.reshape(B, N, self.heads, self.dim) + + q = self.q_norm_cam(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm_cam(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + + q = F.relu(q) + k = F.relu(k) + k_scale = (self.dim**-0.5) * (S**-0.5) + k = k * k_scale + + # (B, H, N, D) for UCPE apply, then back to (B, H, D, N) for the scan. + q_bhnd = q.permute(0, 2, 1, 3) + k_bhnd = k.permute(0, 2, 1, 3) + v_bhnd = v.permute(0, 2, 1, 3) + + q_proj = apply_q(q_bhnd) + kv_proj = apply_kv(torch.cat([k_bhnd, v_bhnd], dim=1)) + k_proj, v_proj = torch.chunk(kv_proj, chunks=2, dim=1) + + q_pre_dn = q_bhnd.permute(0, 1, 3, 2) + q_dn = q_proj.permute(0, 1, 3, 2) + k_pre_dn = k_bhnd.permute(0, 1, 3, 2) + k_dn = k_proj.permute(0, 1, 3, 2) + v_pre_dn = v_bhnd.permute(0, 1, 3, 2) + v_dn = v_proj.permute(0, 1, 3, 2) + + q_dn = _downscale_to_reference_rms(q_pre_dn, q_dn) + k_dn = _downscale_to_reference_rms(k_pre_dn, k_dn) + v_dn = _downscale_to_reference_rms(v_pre_dn, v_dn) + + pre_ucpe_k_norm = torch.linalg.vector_norm( + k_pre_dn.float(), dim=2, keepdim=True + ).clamp_min(1e-6) + post_ucpe_k_norm = torch.linalg.vector_norm( + k_dn.float(), dim=2, keepdim=True + ).clamp_min(1e-6) + inflation_sq = (post_ucpe_k_norm / pre_ucpe_k_norm) ** 2 + frame_inflation_sq = inflation_sq.view(B, self.heads, T, S).mean(dim=-1) + if beta.ndim == 3: + beta = beta / frame_inflation_sq.clamp_min(1.0) + elif beta.ndim == 4: + beta = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0) + + dtype = q_dn.dtype + out = self._maybe_cam_branch_triton_scan(q_dn, k_dn, v_dn, beta, decay, HW) + if out is None: + scan_chunk_size = ( + self.chunk_gdn_chunk_size + if self.cam_update_rule == "torch_chunk" + else None + ) + out = _single_path_delta_scan_bidirectional( + q_dn.float(), + k_dn.float(), + v_dn.float(), + beta.float(), + decay.float(), + HW=HW, + chunk_size=scan_chunk_size, + ) + out = out.to(dtype) + out_bhnd = out.permute(0, 1, 3, 2) + out_bhnd = apply_o(out_bhnd) + return out_bhnd.permute(0, 2, 1, 3).reshape(B, N, C) + + def _cam_branch_softmax( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + apply_q: Callable, + apply_kv: Callable, + apply_o: Callable, + chunk_size: Optional[int] = None, + chunk_split_strategy: str = "uniform", + chunk_index: Optional[List[int]] = None, + ) -> torch.Tensor: + B, N, C = x.shape + + q, k, v = self._cam_qkv(x) + q = q.reshape(B, N, self.heads, self.dim) + k = k.reshape(B, N, self.heads, self.dim) + v = v.reshape(B, N, self.heads, self.dim) + + if self.conv_k_cam is not None: + k = self._temporal_short_conv( + k.reshape(B, N, C), self.conv_k_cam, HW, bidirectional=True + ) + k = k.reshape(B, N, self.heads, self.dim) + + q = self.q_norm_cam(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm_cam(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + + q_bhnd = q.permute(0, 2, 1, 3) + k_bhnd = k.permute(0, 2, 1, 3) + v_bhnd = v.permute(0, 2, 1, 3) + + q_proj = apply_q(q_bhnd) + kv_proj = apply_kv(torch.cat([k_bhnd, v_bhnd], dim=1)) + k_proj, v_proj = torch.chunk(kv_proj, chunks=2, dim=1) + + q_pre_dn = q_bhnd.permute(0, 1, 3, 2) + k_pre_dn = k_bhnd.permute(0, 1, 3, 2) + v_pre_dn = v_bhnd.permute(0, 1, 3, 2) + q_dn = _downscale_to_reference_rms(q_pre_dn, q_proj.permute(0, 1, 3, 2)) + k_dn = _downscale_to_reference_rms(k_pre_dn, k_proj.permute(0, 1, 3, 2)) + v_dn = _downscale_to_reference_rms(v_pre_dn, v_proj.permute(0, 1, 3, 2)) + + q_in = q_dn.permute(0, 3, 1, 2).contiguous() + k_in = k_dn.permute(0, 3, 1, 2).contiguous() + v_in = v_dn.permute(0, 3, 1, 2).contiguous() + out = None + if self.use_chunked_softmax_attention: + out = _sana_wm_chunked_attention( + q_in, + k_in, + v_in, + HW=HW, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + chunk_index=chunk_index, + softmax_scale=self.softmax_attn.softmax_scale, + ) + if out is None: + out = self.softmax_attn(q_in, k_in, v_in) # (B, N, H, D) + + out_bhnd = out.transpose(1, 2).contiguous() + out_bhnd = apply_o(out_bhnd) + return out_bhnd.transpose(1, 2).reshape(B, N, C) + + # ------------------------------------------------------------------ # + + def forward( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor] = None, + prope_fns: Optional[Tuple[Callable, Callable, Callable]] = None, + chunk_size: Optional[int] = None, + chunk_split_strategy: str = "uniform", + chunk_index: Optional[List[int]] = None, + ) -> torch.Tensor: + if self.softmax_main: + main_raw, beta, decay = self._main_branch_softmax( + x, + HW, + rotary_emb, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + chunk_index=chunk_index, + ) + else: + main_raw, beta, decay = self._main_branch_gdn(x, HW, rotary_emb) + + if prope_fns is not None: + apply_q, apply_kv, apply_o = prope_fns + if self.softmax_main: + cam_raw = self._cam_branch_softmax( + x, + HW, + apply_q, + apply_kv, + apply_o, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + chunk_index=chunk_index, + ) + else: + cam_raw = self._cam_branch( + x, HW, apply_q, apply_kv, apply_o, beta, decay + ) + combined = main_raw + self.out_proj_cam(cam_raw) + else: + combined = main_raw + + # Shared output gate + shared output projection. The SiLU gate is + # evaluated in fp32 and multiplied before casting for proj. + gate = F.silu(self.output_gate(x).to(torch.float32)) + combined = combined * gate + return self.proj(combined.to(self.proj.weight.dtype)) + + # ------------------------------------------------------------------ # + # Streaming `forward_long`: chunk-causal cached branch methods + dispatcher. + # Each takes a per-block 10-slot `kv_cache` list (mutated in place) + + # `save_kv_cache`. GDN/cam scans carry recurrent state (slots 0/1/2); + # softmax blocks use a concat-window (slots 0/1 main, 2/3 cam); short-conv + # prefix (slot 4); type flag (slot 6). A single chunk with an empty cache + # reduces to the dense bidirectional path. + # ------------------------------------------------------------------ # + + def _main_branch_gdn_cached( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor], + kv_cache: list, + save_kv_cache: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, N, C = x.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + + qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.dim).contiguous() + q, k, v = qkv.unbind(2) + + # Cached forward prefix + intra-chunk backward short conv on K. + if self.conv_k is not None: + k, new_prefix = _temporal_short_conv_cached( + k.reshape(B, N, C), + self.conv_k, + HW, + prefix=kv_cache[_SLOT_SHORTCONV], + save_prefix=save_kv_cache, + bidirectional=True, + ) + k = k.reshape(B, N, self.heads, self.dim) + if save_kv_cache: + kv_cache[_SLOT_SHORTCONV] = new_prefix + + beta, decay = _compute_frame_gates( + x, HW, self.heads, self.beta_proj, self.gate_proj, self.dt_bias, self.A_log + ) + + # Bypass the Triton fast path: it carries no state. + q = self.q_norm(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + q = F.relu(q) + k = F.relu(k) + k = k * ((self.dim**-0.5) * (S**-0.5)) + + q = q.permute(0, 2, 3, 1) + k = k.permute(0, 2, 3, 1) + v = v.permute(0, 2, 3, 1) + + rotary_emb = _slice_rope_to_current_chunk(rotary_emb, N) + if rotary_emb is not None: + q_rot = _apply_rotary_emb_dn(q, rotary_emb) + k_rot = _apply_rotary_emb_dn(k, rotary_emb) + else: + q_rot = q + k_rot = k + + dtype = q.dtype + out, (state_kv, state_z) = _gdn_scan_cached( + q.float(), + k.float(), + v.float(), + q_rot.float(), + k_rot.float(), + beta.float(), + decay.float(), + init_state_kv=kv_cache[_SLOT_K], + init_state_z=kv_cache[_SLOT_V], + eps=self.eps, + ) + out = out.to(dtype) + if save_kv_cache: + kv_cache[_SLOT_K] = state_kv.detach().clone() + kv_cache[_SLOT_V] = state_z.detach().clone() + kv_cache[_SLOT_TYPE_FLAG] = torch.tensor( + [_CACHE_TYPE_STATE], device=x.device + ) + + out = out.permute(0, 3, 1, 2).reshape(B, N, C) + return out, beta, decay + + def _main_branch_softmax_cached( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor], + kv_cache: list, + save_kv_cache: bool, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.dim) + q, k, v = qkv.unbind(2) + if self.conv_k is not None: # softmax blocks have conv_k=None + k, new_prefix = _temporal_short_conv_cached( + k.reshape(B, N, C), + self.conv_k, + HW, + prefix=kv_cache[_SLOT_SHORTCONV], + save_prefix=save_kv_cache, + bidirectional=True, + ) + k = k.reshape(B, N, self.heads, self.dim) + if save_kv_cache: + kv_cache[_SLOT_SHORTCONV] = new_prefix + q = self.q_norm(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + q = q.permute(0, 2, 1, 3) # (B, H, N, D) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + rotary_emb = _slice_rope_to_current_chunk(rotary_emb, N) + if rotary_emb is not None: + q = _apply_rotary_emb_bhnd(q, rotary_emb) + k = _apply_rotary_emb_bhnd(k, rotary_emb) + q_in = q.transpose(1, 2).contiguous() # (B, N_cur, H, D) + k_in = k.transpose(1, 2).contiguous() + v_in = v.transpose(1, 2).contiguous() + + # Concat-window: store the CURRENT chunk K/V (RoPE'd at absolute + # positions), then prepend the cached prefix. Q stays current-chunk. + cached_k = kv_cache[_SLOT_K] + cached_v = kv_cache[_SLOT_V] + if save_kv_cache: + kv_cache[_SLOT_K] = k_in.detach().clone() + kv_cache[_SLOT_V] = v_in.detach().clone() + kv_cache[_SLOT_TYPE_FLAG] = torch.tensor( + [_CACHE_TYPE_CONCAT], device=x.device + ) + if cached_k is not None: + k_in = torch.cat([cached_k.to(k_in.dtype), k_in], dim=1) + v_in = torch.cat([cached_v.to(v_in.dtype), v_in], dim=1) + + out = _sana_wm_sdpa( + q_in, k_in, v_in, softmax_scale=_sana_wm_padded_scale(self.dim) + ) + out = out.reshape(B, N, C) + + beta, decay = _compute_frame_gates( + x, HW, self.heads, self.beta_proj, self.gate_proj, self.dt_bias, self.A_log + ) + return out, beta, decay + + def _cam_branch_cached( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + apply_q: Callable, + apply_kv: Callable, + apply_o: Callable, + beta: torch.Tensor, + decay: torch.Tensor, + kv_cache: list, + save_kv_cache: bool, + ) -> torch.Tensor: + B, N, C = x.shape + T, H_sp, W_sp = HW + S = H_sp * W_sp + + q, k, v = self._cam_qkv(x) + q = q.reshape(B, N, self.heads, self.dim) + k = k.reshape(B, N, self.heads, self.dim) + v = v.reshape(B, N, self.heads, self.dim) + + # Cam K short conv stays bidirectional/uncached; slot 4 is main-K only. + if self.conv_k_cam is not None: + k = self._temporal_short_conv( + k.reshape(B, N, C), self.conv_k_cam, HW, bidirectional=True + ) + k = k.reshape(B, N, self.heads, self.dim) + + q = self.q_norm_cam(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm_cam(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + q = F.relu(q) + k = F.relu(k) + k = k * ((self.dim**-0.5) * (S**-0.5)) + + q_bhnd = q.permute(0, 2, 1, 3) + k_bhnd = k.permute(0, 2, 1, 3) + v_bhnd = v.permute(0, 2, 1, 3) + + q_proj = apply_q(q_bhnd) + kv_proj = apply_kv(torch.cat([k_bhnd, v_bhnd], dim=1)) + k_proj, v_proj = torch.chunk(kv_proj, chunks=2, dim=1) + + q_pre_dn = q_bhnd.permute(0, 1, 3, 2) + q_dn = q_proj.permute(0, 1, 3, 2) + k_pre_dn = k_bhnd.permute(0, 1, 3, 2) + k_dn = k_proj.permute(0, 1, 3, 2) + v_pre_dn = v_bhnd.permute(0, 1, 3, 2) + v_dn = v_proj.permute(0, 1, 3, 2) + + # No RMS downscale here: full post-UCPE q/k/v feed the scan; inflation + # is computed from full post-UCPE K vs pre-UCPE K and absorbed only into + # beta. + pre_ucpe_k_norm = torch.linalg.vector_norm( + k_pre_dn.float(), dim=2, keepdim=True + ).clamp_min(1e-6) + post_ucpe_k_norm = torch.linalg.vector_norm( + k_dn.float(), dim=2, keepdim=True + ).clamp_min(1e-6) + inflation_sq = (post_ucpe_k_norm / pre_ucpe_k_norm) ** 2 + frame_inflation_sq = inflation_sq.view(B, self.heads, T, S).mean(dim=-1) + if beta.ndim == 3: + beta = beta / frame_inflation_sq.clamp_min(1.0) + elif beta.ndim == 4: + beta = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0) + + dtype = q_dn.dtype + # Bypass the Triton cam scan: it carries no state. + out, cam_state = _single_path_delta_scan_cached( + q_dn.float(), + k_dn.float(), + v_dn.float(), + beta.float(), + decay.float(), + init_state_kv=kv_cache[_SLOT_CAM_K], + ) + out = out.to(dtype) + if save_kv_cache: + kv_cache[_SLOT_CAM_K] = cam_state.detach().clone() + + out_bhnd = out.permute(0, 1, 3, 2) + out_bhnd = apply_o(out_bhnd) + return out_bhnd.permute(0, 2, 1, 3).reshape(B, N, C) + + def _cam_branch_softmax_cached( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + apply_q: Callable, + apply_kv: Callable, + apply_o: Callable, + kv_cache: list, + save_kv_cache: bool, + ) -> torch.Tensor: + B, N, C = x.shape + + q, k, v = self._cam_qkv(x) + q = q.reshape(B, N, self.heads, self.dim) + k = k.reshape(B, N, self.heads, self.dim) + v = v.reshape(B, N, self.heads, self.dim) + + if self.conv_k_cam is not None: + k = self._temporal_short_conv( + k.reshape(B, N, C), self.conv_k_cam, HW, bidirectional=True + ) + k = k.reshape(B, N, self.heads, self.dim) + + q = self.q_norm_cam(q.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + k = self.k_norm_cam(k.reshape(B, N, C)).reshape(B, N, self.heads, self.dim) + + q_bhnd = q.permute(0, 2, 1, 3) + k_bhnd = k.permute(0, 2, 1, 3) + v_bhnd = v.permute(0, 2, 1, 3) + + q_proj = apply_q(q_bhnd) + kv_proj = apply_kv(torch.cat([k_bhnd, v_bhnd], dim=1)) + k_proj, v_proj = torch.chunk(kv_proj, chunks=2, dim=1) + + # No RMS downscale here: full post-UCPE q/k/v feed SDPA directly. + q_dn = q_proj.permute(0, 1, 3, 2) + k_dn = k_proj.permute(0, 1, 3, 2) + v_dn = v_proj.permute(0, 1, 3, 2) + + q_in = q_dn.permute(0, 3, 1, 2).contiguous() # (B, N_cur, H, D) + k_in = k_dn.permute(0, 3, 1, 2).contiguous() + v_in = v_dn.permute(0, 3, 1, 2).contiguous() + + cached_cam_k = kv_cache[_SLOT_CAM_K] + cached_cam_v = kv_cache[_SLOT_CAM_V] + if save_kv_cache: + kv_cache[_SLOT_CAM_K] = k_in.detach().clone() + kv_cache[_SLOT_CAM_V] = v_in.detach().clone() + if cached_cam_k is not None: + k_in = torch.cat([cached_cam_k.to(k_in.dtype), k_in], dim=1) + v_in = torch.cat([cached_cam_v.to(v_in.dtype), v_in], dim=1) + + out = _sana_wm_sdpa( + q_in, k_in, v_in, softmax_scale=_sana_wm_padded_scale(self.dim) + ) # (B, N_cur, H, D) + out_bhnd = out.transpose(1, 2).contiguous() + out_bhnd = apply_o(out_bhnd) + return out_bhnd.transpose(1, 2).reshape(B, N, C) + + def forward_long( + self, + x: torch.Tensor, + HW: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor] = None, + prope_fns: Optional[Tuple[Callable, Callable, Callable]] = None, + *, + kv_cache: list, + save_kv_cache: bool, + ) -> Tuple[torch.Tensor, list]: + if self.softmax_main: + main_raw, beta, decay = self._main_branch_softmax_cached( + x, HW, rotary_emb, kv_cache, save_kv_cache + ) + else: + main_raw, beta, decay = self._main_branch_gdn_cached( + x, HW, rotary_emb, kv_cache, save_kv_cache + ) + + if prope_fns is not None: + apply_q, apply_kv, apply_o = prope_fns + if self.softmax_main: + cam_raw = self._cam_branch_softmax_cached( + x, HW, apply_q, apply_kv, apply_o, kv_cache, save_kv_cache + ) + else: + cam_raw = self._cam_branch_cached( + x, + HW, + apply_q, + apply_kv, + apply_o, + beta, + decay, + kv_cache, + save_kv_cache, + ) + combined = main_raw + self.out_proj_cam(cam_raw) + else: + combined = main_raw + + gate = F.silu(self.output_gate(x).to(torch.float32)) + combined = combined * gate + return self.proj(combined.to(self.proj.weight.dtype)), kv_cache + + +# --------------------------------------------------------------------------- +# Cross-attention (text conditioning). Stored as +# ``cross_attn.{q_linear, kv_linear, proj, q_norm, k_norm}``. +# --------------------------------------------------------------------------- + + +class MultiHeadCrossAttention(nn.Module): + def __init__(self, d_model: int, num_heads: int, qk_norm: bool = True) -> None: + super().__init__() + assert d_model % num_heads == 0 + self.d_model = d_model + self.num_heads = num_heads + self.head_dim = d_model // num_heads + + self.q_linear = nn.Linear(d_model, d_model, bias=True) + self.kv_linear = nn.Linear(d_model, d_model * 2, bias=True) + self.proj = nn.Linear(d_model, d_model, bias=True) + if qk_norm: + self.q_norm = _RMSNorm(d_model, eps=1e-6) + self.k_norm = _RMSNorm(d_model, eps=1e-6) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + # Padding-mask path falls back to SDPA internally; the unmasked path + # can pick FA3 / FlashInfer / etc. + self.attn = LocalAttention( + num_heads=num_heads, + head_size=self.head_dim, + ) + + def forward( + self, + x: torch.Tensor, # (B, N, D) + cond: torch.Tensor, # (B, L, D) + mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + B, N, D = x.shape + + q = self.q_linear(x) + kv = self.kv_linear(cond).view(B, -1, 2, D) + k, v = kv.unbind(2) + # LocalAttention takes (B, N, H, D). + q = self.q_norm(q).view(B, N, self.num_heads, self.head_dim) + k = self.k_norm(k).view(B, -1, self.num_heads, self.head_dim) + v = v.view(B, -1, self.num_heads, self.head_dim) + + attn_mask = mask.bool() if mask is not None else None + out = self.attn(q, k, v, attn_mask=attn_mask) # (B, N, H, D) + out = out.reshape(B, N, D) + return self.proj(out) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py new file mode 100644 index 000000000..8d9b353b9 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any, Optional + +import torch +import torch.nn as nn + +from sglang.multimodal_gen.configs.models.dits.sana_wm_refiner import ( + SanaWMRefinerArchConfig, + SanaWMRefinerConfig, +) +from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear +from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.models.dits.ltx_2 import ( + LTX2AdaLayerNormSingle, + LTX2Attention, + LTX2AudioVideoRotaryPosEmbed, + LTX2FeedForward, + LTX2TextProjection, +) + + +def pack_latents( + latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1 +) -> torch.Tensor: + """Pack a 5D latent (B, C, T, H, W) into a 3D token sequence (B, L, in_dim).""" + B, _, T, H, W = latents.shape + pT = T // patch_size_t + pH = H // patch_size + pW = W // patch_size + latents = latents.reshape(B, -1, pT, patch_size_t, pH, patch_size, pW, patch_size) + return latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + + +def unpack_latents( + tokens: torch.Tensor, + num_frames: int, + height: int, + width: int, + patch_size: int = 1, + patch_size_t: int = 1, +) -> torch.Tensor: + """Inverse of `pack_latents`: (B, L, out_dim) -> (B, C, T, H, W).""" + B = tokens.size(0) + tokens = tokens.reshape( + B, + num_frames // patch_size_t, + height // patch_size, + width // patch_size, + -1, + patch_size_t, + patch_size, + patch_size, + ) + return ( + tokens.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + ) + + +def _slice_rope( + rope: tuple[torch.Tensor, torch.Tensor], start: int, end: Optional[int] = None +) -> tuple[torch.Tensor, torch.Tensor]: + """Slice along token axis for either interleaved (rank-3) or split (rank-4).""" + cos, sin = rope + end_ = end if end is not None else cos.shape[-2 if cos.ndim == 4 else 1] + if cos.ndim == 3: + return cos[:, start:end_], sin[:, start:end_] + if cos.ndim == 4: + return cos[:, :, start:end_, :], sin[:, :, start:end_, :] + raise ValueError(f"Unexpected RoPE rank: {cos.ndim}") + + +def _streaming_self_attention( + attn: LTX2Attention, + hidden_states: torch.Tensor, + video_rotary_emb: tuple[torch.Tensor, torch.Tensor], + n_context_tokens: int, +) -> torch.Tensor: + """Streaming SLA: context attends to context only, current attends to context+current. + + Mirrors NVlabs `inference_sana_wm.py::_streaming_self_attention`. + """ + seq_len = hidden_states.shape[1] + if n_context_tokens <= 0 or n_context_tokens >= seq_len: + return attn(hidden_states, context=None, pe=video_rotary_emb) + + ctx_rope = _slice_rope(video_rotary_emb, 0, n_context_tokens) + out_ctx = attn( + hidden_states[:, :n_context_tokens], + context=None, + pe=ctx_rope, + ) + + cur_rope = _slice_rope(video_rotary_emb, n_context_tokens, seq_len) + out_cur = attn( + hidden_states[:, n_context_tokens:], + context=hidden_states, + pe=cur_rope, + k_pe=video_rotary_emb, + ) + return torch.cat([out_ctx, out_cur], dim=1) + + +class SanaWMRefinerBlock(nn.Module): + """Video-only LTX-2 transformer block. + + Diffusers-compatible layout: `norm1 -> attn1 (self) -> norm2 -> attn2 (cross) -> norm3 -> ff`, + each modulated via per-block `scale_shift_table` + token-wise `temb`. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + cross_attention_dim: int, + qk_norm: bool = True, + norm_eps: float = 1e-6, + apply_gated_attention: bool = False, + prefix: str = "", + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + self.dim = int(dim) + + self.norm1 = nn.RMSNorm(self.dim, eps=norm_eps, elementwise_affine=False) + self.attn1 = LTX2Attention( + query_dim=self.dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + norm_eps=norm_eps, + qk_norm=qk_norm, + apply_gated_attention=apply_gated_attention, + prefix=f"{prefix}.attn1", + quant_config=quant_config, + ) + + self.norm2 = nn.RMSNorm(self.dim, eps=norm_eps, elementwise_affine=False) + self.attn2 = LTX2Attention( + query_dim=self.dim, + context_dim=cross_attention_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + norm_eps=norm_eps, + qk_norm=qk_norm, + use_local_attention=True, + apply_gated_attention=apply_gated_attention, + prefix=f"{prefix}.attn2", + quant_config=quant_config, + ) + + self.norm3 = nn.RMSNorm(self.dim, eps=norm_eps, elementwise_affine=False) + self.ff = LTX2FeedForward(self.dim, dim_out=self.dim, quant_config=quant_config) + + self.scale_shift_table = nn.Parameter(torch.randn(6, self.dim) / self.dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + video_rotary_emb: tuple[torch.Tensor, torch.Tensor], + encoder_attention_mask: Optional[torch.Tensor] = None, + n_context_tokens: int = 0, + ) -> torch.Tensor: + B = hidden_states.size(0) + T = temb.size(1) + D = self.dim + ada = self.scale_shift_table[None, None].to( + device=temb.device, dtype=temb.dtype + ) + temb.reshape(B, T, 6, D) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada.unbind( + dim=2 + ) + + normed = self.norm1(hidden_states) * (1 + scale_msa) + shift_msa + attn_out = _streaming_self_attention( + self.attn1, + normed, + video_rotary_emb, + n_context_tokens=n_context_tokens, + ) + hidden_states = hidden_states + attn_out * gate_msa + + normed = self.norm2(hidden_states) + ca_out = self.attn2( + normed, + context=encoder_hidden_states, + mask=encoder_attention_mask, + pe=None, + ) + hidden_states = hidden_states + ca_out + + normed = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp + hidden_states = hidden_states + self.ff(normed) * gate_mlp + return hidden_states + + +class SanaWMLTX2VideoRefiner(CachableDiT, LayerwiseOffloadableModuleMixin): + """SANA-WM stage-2 LTX-2 video-only refiner. + + Loads Diffusers-format refiner weights from `/refiner/transformer/`. + Audio params present in the checkpoint are silently dropped by the loader's + `strict=False` state_dict load. + """ + + _fsdp_shard_conditions = SanaWMRefinerArchConfig()._fsdp_shard_conditions + _compile_conditions = SanaWMRefinerArchConfig()._compile_conditions + _supported_attention_backends = ( + SanaWMRefinerArchConfig()._supported_attention_backends + ) + param_names_mapping = SanaWMRefinerArchConfig().param_names_mapping + reverse_param_names_mapping: dict = {} + lora_param_names_mapping: dict = {} + + def __init__( + self, + config: SanaWMRefinerConfig, + hf_config: dict[str, Any], + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__(config, hf_config=hf_config) + arch = config.arch_config + + self.in_channels = int(arch.in_channels) + self.out_channels = int(arch.out_channels) + self.patch_size = int(arch.patch_size) + self.patch_size_t = int(arch.patch_size_t) + self.hidden_size = int(arch.hidden_size) + self.num_attention_heads = int(arch.num_attention_heads) + self.num_channels_latents = int(arch.num_channels_latents) + self.attention_head_dim = int(arch.attention_head_dim) + self.timestep_scale_multiplier = float(arch.timestep_scale_multiplier) + self.rope_type = str(arch.rope_type) + + in_dim = ( + self.in_channels * self.patch_size_t * self.patch_size * self.patch_size + ) + out_dim = ( + self.out_channels * self.patch_size_t * self.patch_size * self.patch_size + ) + + self.proj_in = ColumnParallelLinear( + in_dim, + self.hidden_size, + bias=True, + gather_output=True, + quant_config=quant_config, + ) + + self.time_embed = LTX2AdaLayerNormSingle( + self.hidden_size, embedding_coefficient=6 + ) + self.caption_projection = LTX2TextProjection( + in_features=int(arch.caption_channels), + hidden_size=self.hidden_size, + out_features=self.hidden_size, + act_fn="gelu_tanh", + ) + + self.transformer_blocks = nn.ModuleList( + [ + SanaWMRefinerBlock( + dim=self.hidden_size, + num_attention_heads=self.num_attention_heads, + attention_head_dim=self.attention_head_dim, + cross_attention_dim=int(arch.cross_attention_dim), + qk_norm=bool(arch.qk_norm), + norm_eps=float(arch.norm_eps), + apply_gated_attention=bool(arch.apply_gated_attention), + prefix=f"transformer_blocks.{i}", + quant_config=quant_config, + ) + for i in range(int(arch.num_layers)) + ] + ) + + self.scale_shift_table = nn.Parameter( + torch.randn(2, self.hidden_size) / self.hidden_size**0.5 + ) + self.norm_out = nn.LayerNorm( + self.hidden_size, eps=float(arch.norm_eps), elementwise_affine=False + ) + self.proj_out = ColumnParallelLinear( + self.hidden_size, + out_dim, + bias=True, + gather_output=True, + quant_config=quant_config, + ) + + # LTX2AudioVideoRotaryPosEmbed expects `dim` to be the *total* hidden + # size (num_heads * head_dim), not the per-head dim. It internally + # reshapes cos/sin to (B, T, num_heads, head_dim/2). Passing + # `attention_head_dim` here would size the RoPE to head_dim/num_heads + # and produce a (1, num_heads, L, 2) cos/sin that won't match + # LTX2Attention's q/k. See LTX2Transformer3DAVModel.__init__ in + # ltx_2.py for the canonical convention (`dim=self.hidden_size`). + self.rope = LTX2AudioVideoRotaryPosEmbed( + dim=self.hidden_size, + patch_size=self.patch_size, + patch_size_t=self.patch_size_t, + base_num_frames=int(arch.base_num_frames), + base_height=int(arch.base_height), + base_width=int(arch.base_width), + sampling_rate=int(arch.sampling_rate), + hop_length=int(arch.hop_length), + scale_factors=tuple(arch.scale_factors), + causal_offset=int(arch.causal_offset), + modality="video", + rope_type=self.rope_type, + num_attention_heads=self.num_attention_heads, + ) + + self.layer_names = ["transformer_blocks"] + + def _scale_timestep_for_adaln(self, timestep: torch.Tensor) -> torch.Tensor: + return timestep * self.timestep_scale_multiplier + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states_image=None, + encoder_attention_mask: Optional[torch.Tensor] = None, + num_frames: Optional[int] = None, + height: Optional[int] = None, + width: Optional[int] = None, + fps: float = 24.0, + n_context_tokens: int = 0, + guidance=None, + **kwargs, + ) -> torch.Tensor: + # Accept either packed (B, L, in_dim) or raw 5D (B, C, T, H, W). + if hidden_states.dim() == 5: + B_, _, T_, H_, W_ = hidden_states.shape + if num_frames is None: + num_frames = T_ + if height is None: + height = H_ + if width is None: + width = W_ + hidden_states = pack_latents( + hidden_states, + patch_size=self.patch_size, + patch_size_t=self.patch_size_t, + ) + packed_input = True + else: + if num_frames is None or height is None or width is None: + raise ValueError( + "num_frames/height/width are required when hidden_states is pre-packed." + ) + packed_input = False + + B = hidden_states.size(0) + + video_coords = self.rope.prepare_video_coords( + batch_size=B, + num_frames=num_frames, + height=height, + width=width, + device=hidden_states.device, + fps=fps, + ) + video_rotary_emb = self.rope( + video_coords, + device=hidden_states.device, + out_dtype=hidden_states.dtype, + ) + + hidden_states, _ = self.proj_in(hidden_states) + + scaled_t = self._scale_timestep_for_adaln(timestep) + temb, embedded_timestep = self.time_embed( + scaled_t.flatten(), hidden_dtype=hidden_states.dtype + ) + if timestep.dim() >= 2: + temb = temb.view(B, -1, temb.size(-1)) + embedded_timestep = embedded_timestep.view( + B, -1, embedded_timestep.size(-1) + ) + else: + temb = temb.view(B, 1, temb.size(-1)) + embedded_timestep = embedded_timestep.view(B, 1, embedded_timestep.size(-1)) + + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(B, -1, self.hidden_size) + + for block in self.transformer_blocks: + hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + video_rotary_emb=video_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + n_context_tokens=n_context_tokens, + ) + + scale_shift_values = self.scale_shift_table[None, None].to( + device=hidden_states.device, dtype=hidden_states.dtype + ) + embedded_timestep[:, :, None].to(hidden_states.dtype) + shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] + hidden_states = self.norm_out(hidden_states) * (1 + scale) + shift + hidden_states, _ = self.proj_out(hidden_states) + + if packed_input: + return hidden_states + return unpack_latents( + hidden_states, + num_frames=num_frames, + height=height, + width=width, + patch_size=self.patch_size, + patch_size_t=self.patch_size_t, + ) + + +EntryClass = SanaWMLTX2VideoRefiner diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py index 9ddaea282..207362e86 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py @@ -87,14 +87,47 @@ class LTX2VideoCausalConv3d(nn.Module): padding_mode=spatial_padding_mode, ) - def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor: + def forward( + self, + hidden_states: torch.Tensor, + causal: bool = True, + conv_cache: Optional[dict] = None, + cache_key: Optional[str] = None, + ) -> torch.Tensor: time_kernel_size = self.kernel_size[0] if causal: - pad_left = hidden_states[:, :, :1, :, :].repeat( - (1, 1, time_kernel_size - 1, 1, 1) - ) - hidden_states = torch.concatenate([pad_left, hidden_states], dim=2) + if ( + conv_cache is not None + and cache_key is not None + and time_kernel_size > 1 + ): + # Streaming: prepend the previous chunk's last (k-1) frames of the + # PADDED conv input (first chunk: replicate frame 0, the "sink"), + # then store the last (k-1) frames of THIS padded input for the next + # chunk. Storing the padded tail (not the raw input) keeps exactly + # k-1 frames even for short chunks. Left-pad k-1 with temporal pad 0 / + # stride 1 keeps the output T invariant. + prev = conv_cache.get(cache_key) + if prev is None: + pad_left = hidden_states[:, :, :1, :, :].repeat( + (1, 1, time_kernel_size - 1, 1, 1) + ) + else: + pad_left = prev.to( + device=hidden_states.device, dtype=hidden_states.dtype + ) + hidden_states = torch.concatenate([pad_left, hidden_states], dim=2) + conv_cache[cache_key] = ( + hidden_states[:, :, -(time_kernel_size - 1) :, :, :] + .detach() + .clone() + ) + else: + pad_left = hidden_states[:, :, :1, :, :].repeat( + (1, 1, time_kernel_size - 1, 1, 1) + ) + hidden_states = torch.concatenate([pad_left, hidden_states], dim=2) else: pad_left = hidden_states[:, :, :1, :, :].repeat( (1, 1, (time_kernel_size - 1) // 2, 1, 1) @@ -200,6 +233,8 @@ class LTX2VideoResnetBlock3d(nn.Module): temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, causal: bool = True, + conv_cache: Optional[dict] = None, + cache_key: Optional[str] = None, ) -> torch.Tensor: hidden_states = inputs @@ -214,7 +249,12 @@ class LTX2VideoResnetBlock3d(nn.Module): hidden_states = hidden_states * (1 + scale_1) + shift_1 hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.conv1(hidden_states, causal=causal) + hidden_states = self.conv1( + hidden_states, + causal=causal, + conv_cache=conv_cache, + cache_key=None if cache_key is None else f"{cache_key}.conv1", + ) if self.per_channel_scale1 is not None: spatial_shape = hidden_states.shape[-2:] @@ -236,7 +276,12 @@ class LTX2VideoResnetBlock3d(nn.Module): hidden_states = self.nonlinearity(hidden_states) hidden_states = self.dropout(hidden_states) - hidden_states = self.conv2(hidden_states, causal=causal) + hidden_states = self.conv2( + hidden_states, + causal=causal, + conv_cache=conv_cache, + cache_key=None if cache_key is None else f"{cache_key}.conv2", + ) if self.per_channel_scale2 is not None: spatial_shape = hidden_states.shape[-2:] @@ -343,9 +388,27 @@ class LTXVideoUpsampler3d(nn.Module): spatial_padding_mode=spatial_padding_mode, ) - def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor: + def forward( + self, + hidden_states: torch.Tensor, + causal: bool = True, + conv_cache: Optional[dict] = None, + cache_key: Optional[str] = None, + ) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = hidden_states.shape + # The temporal pixel-shuffle expands 1 latent frame -> stride[0] sample + # frames, then drops the first (stride[0]-1) frames as the causal "anchor". + # Under chunking that drop must happen ONLY at the true clip start, else a + # frame is lost at every chunk boundary. trim_start is read once and applied + # to BOTH the residual and shuffle branches (mirrors the reference). + trim_start = self.stride[0] - 1 + if conv_cache is not None and cache_key is not None: + if conv_cache.get(f"{cache_key}.trim_applied", False): + trim_start = 0 + else: + conv_cache[f"{cache_key}.trim_applied"] = True + if self.residual: residual = hidden_states.reshape( batch_size, @@ -367,9 +430,11 @@ class LTXVideoUpsampler3d(nn.Module): self.stride[0] * self.stride[1] * self.stride[2] ) // self.upscale_factor residual = residual.repeat(1, repeats, 1, 1, 1) - residual = residual[:, :, self.stride[0] - 1 :] + residual = residual[:, :, trim_start:] - hidden_states = self.conv(hidden_states, causal=causal) + hidden_states = self.conv( + hidden_states, causal=causal, conv_cache=conv_cache, cache_key=cache_key + ) hidden_states = hidden_states.reshape( batch_size, -1, @@ -386,7 +451,7 @@ class LTXVideoUpsampler3d(nn.Module): .flatten(4, 5) .flatten(2, 3) ) - hidden_states = hidden_states[:, :, self.stride[0] - 1 :] + hidden_states = hidden_states[:, :, trim_start:] if self.residual: hidden_states = hidden_states + residual @@ -597,6 +662,8 @@ class LTX2VideoMidBlock3d(nn.Module): temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, causal: bool = True, + conv_cache: Optional[dict] = None, + cache_key: Optional[str] = None, ) -> torch.Tensor: r"""Forward method of the `LTXMidBlock3D` class.""" @@ -616,7 +683,14 @@ class LTX2VideoMidBlock3d(nn.Module): resnet, hidden_states, temb, generator, causal ) else: - hidden_states = resnet(hidden_states, temb, generator, causal=causal) + hidden_states = resnet( + hidden_states, + temb, + generator, + causal=causal, + conv_cache=conv_cache, + cache_key=None if cache_key is None else f"{cache_key}.resnets.{i}", + ) return hidden_states @@ -782,9 +856,18 @@ class LTX2VideoUpBlock3d(nn.Module): temb: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, causal: bool = True, + conv_cache: Optional[dict] = None, + cache_key: Optional[str] = None, ) -> torch.Tensor: if self.conv_in is not None: - hidden_states = self.conv_in(hidden_states, temb, generator, causal=causal) + hidden_states = self.conv_in( + hidden_states, + temb, + generator, + causal=causal, + conv_cache=conv_cache, + cache_key=None if cache_key is None else f"{cache_key}.conv_in", + ) if self.time_embedder is not None: temb = self.time_embedder( @@ -797,8 +880,15 @@ class LTX2VideoUpBlock3d(nn.Module): temb = temb.view(hidden_states.size(0), -1, 1, 1, 1) if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, causal=causal) + for j, upsampler in enumerate(self.upsamplers): + hidden_states = upsampler( + hidden_states, + causal=causal, + conv_cache=conv_cache, + cache_key=( + None if cache_key is None else f"{cache_key}.upsamplers.{j}" + ), + ) for i, resnet in enumerate(self.resnets): if torch.is_grad_enabled() and self.gradient_checkpointing: @@ -806,7 +896,14 @@ class LTX2VideoUpBlock3d(nn.Module): resnet, hidden_states, temb, generator, causal ) else: - hidden_states = resnet(hidden_states, temb, generator, causal=causal) + hidden_states = resnet( + hidden_states, + temb, + generator, + causal=causal, + conv_cache=conv_cache, + cache_key=None if cache_key is None else f"{cache_key}.resnets.{i}", + ) return hidden_states @@ -1116,10 +1213,17 @@ class LTX2VideoDecoder3d(nn.Module): hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, causal: Optional[bool] = None, + conv_cache: Optional[dict] = None, ) -> torch.Tensor: causal = causal or self.is_causal + _ck = (lambda k: k) if conv_cache is not None else (lambda k: None) - hidden_states = self.conv_in(hidden_states, causal=causal) + hidden_states = self.conv_in( + hidden_states, + causal=causal, + conv_cache=conv_cache, + cache_key=_ck("conv_in"), + ) if self.timestep_scale_multiplier is not None: temb = temb * self.timestep_scale_multiplier @@ -1134,10 +1238,22 @@ class LTX2VideoDecoder3d(nn.Module): up_block, hidden_states, temb, None, causal ) else: - hidden_states = self.mid_block(hidden_states, temb, causal=causal) + hidden_states = self.mid_block( + hidden_states, + temb, + causal=causal, + conv_cache=conv_cache, + cache_key=_ck("mid_block"), + ) - for up_block in self.up_blocks: - hidden_states = up_block(hidden_states, temb, causal=causal) + for i, up_block in enumerate(self.up_blocks): + hidden_states = up_block( + hidden_states, + temb, + causal=causal, + conv_cache=conv_cache, + cache_key=_ck(f"up_blocks.{i}"), + ) hidden_states = self.norm_out(hidden_states) @@ -1155,7 +1271,12 @@ class LTX2VideoDecoder3d(nn.Module): hidden_states = hidden_states * (1 + scale) + shift hidden_states = self.conv_act(hidden_states) - hidden_states = self.conv_out(hidden_states, causal=causal) + hidden_states = self.conv_out( + hidden_states, + causal=causal, + conv_cache=conv_cache, + cache_key=_ck("conv_out"), + ) p = self.patch_size p_t = self.patch_size_t @@ -1990,4 +2111,38 @@ class AutoencoderKLLTX2Video(ParallelTiledVAE): return dec -EntryClass = AutoencoderKLLTX2Video +class AutoencoderKLCausalLTX2Video(AutoencoderKLLTX2Video): + """Streaming causal LTX-2 VAE. + + Same weights / architecture as ``AutoencoderKLLTX2Video`` but loaded with + ``decoder_causal=True`` (the checkpoint config.json flips it automatically), + plus a chunk-by-chunk ``decode_chunk`` that threads a per-conv ``conv_cache`` + across chunks so frames decode causally as the streaming denoise produces + them -- the ``decode_per_frame_with_cache`` equivalent at chunk granularity. + The streaming VAE config.json's ``_class_name`` is + ``AutoencoderKLCausalLTX2Video``, so this name must be registered. + """ + + @staticmethod + def reset_decoder_cache() -> dict: + """Fresh decoder cache (per-conv feat tails + upsampler trim flags).""" + return {} + + def decode_chunk( + self, + z_chunk: torch.Tensor, + conv_cache: dict, + *, + temb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Causally decode one chunk of latents, carrying ``conv_cache`` across calls. + + ``z_chunk`` is ``(B, C_latent, T_chunk, h, w)`` and must already be + de-normalized (the stage applies the latents scale/shift). Returns the + pixel chunk ``(B, 3, T_out, h*sf, w*sf)`` and mutates ``conv_cache`` in + place. Bypasses the tiled/framewise ``decode`` dispatch on purpose. + """ + return self.decoder(z_chunk, temb, causal=True, conv_cache=conv_cache) + + +EntryClass = [AutoencoderKLLTX2Video, AutoencoderKLCausalLTX2Video] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index 52651762f..e8b27be0a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -5,6 +5,9 @@ import numpy as np import torch from diffusers import FlowMatchEulerDiscreteScheduler +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + STAGE_2_DISTILLED_SIGMA_VALUES as _SHARED_STAGE_2_DISTILLED_SIGMA_VALUES, +) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( LTX2PipelineConfig, is_ltx23_native_variant, @@ -779,7 +782,7 @@ class LTX2TwoStageResidencyController: class LTX2TwoStagePipeline(_BaseLTX2Pipeline): pipeline_name = "LTX2TwoStagePipeline" - STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0] + STAGE_2_DISTILLED_SIGMA_VALUES = list(_SHARED_STAGE_2_DISTILLED_SIGMA_VALUES) STAGE_1_DISTILLED_LORA_STRENGTH = 0.0 STAGE_2_DISTILLED_LORA_STRENGTH = 1.0 STAGE_1_DENOISING_SAMPLER_NAME = "euler" diff --git a/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_pipeline.py new file mode 100644 index 000000000..97e08d0ba --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_pipeline.py @@ -0,0 +1,322 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os + +from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig +from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams +from sglang.multimodal_gen.runtime.loader.utils import get_memory_usage_of_component +from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + InputValidationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm import ( + SanaWMBeforeDenoisingStage, + SanaWMDecodingStage, + SanaWMDenoisingStage, + SanaWMTextEncodingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.refiner import ( + OfficialDiffusersLTX2RefinerModule, + OfficialGemma3TextEncoderModule, + SanaWMLTX2RefinerStage, + SanaWMRefinerDecodingStage, + default_sana_wm_refiner_dtype, + sana_wm_skip_refiner_enabled, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import ( + SanaWMStreamingDecodingStage, + SanaWMStreamingDenoisingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming_refiner import ( + SanaWMStreamingRefinerStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +# Stage-2 refiner sub-modules live under `/refiner/...`, not at the +# model root. They're loaded manually in `initialize_pipeline` rather than via +# `_required_config_modules`, because the framework verifier resolves every +# required module key as a literal top-level subdir of the materialized model. + + +logger = init_logger(__name__) + + +class SanaWMPipeline(LoRAPipeline, ComposedPipelineBase): + """SANA-WM TI2V pipeline (single-stage).""" + + pipeline_name = "SanaWMPipeline" + pipeline_config_cls = SanaWMPipelineConfig + sampling_params_cls = SanaWMSamplingParams + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + ] + + @staticmethod + def _validate_parallelism_args(server_args: ServerArgs) -> None: + tp_size = getattr(server_args, "tp_size", 1) or 1 + if tp_size != 1: + raise ValueError( + "SANA-WM does not support tensor parallelism yet. " + "Use --num-gpus with FSDP/CFG parallelism instead of " + f"--tp-size {tp_size}." + ) + + sp_degree = getattr(server_args, "sp_degree", 1) or 1 + if sp_degree != 1: + raise ValueError( + "SANA-WM does not support temporal sequence parallelism yet. " + "Stage-1 GDN/GLUMBConvTemp span frames and require halo/state " + "exchange before latents can be sharded. Use --num-gpus with " + "FSDP/CFG parallelism instead of " + f"--sp-degree {sp_degree}." + ) + + def create_pipeline_stages(self, server_args: ServerArgs): + self._validate_parallelism_args(server_args) + self.add_stage(InputValidationStage()) + + self.add_stage( + SanaWMTextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + "prompt_encoding_stage", + ) + + self.add_stage( + SanaWMBeforeDenoisingStage( + vae=self.get_module("vae"), + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + pipeline_config=server_args.pipeline_config, + ), + "sana_wm_before_denoising", + ) + + if getattr(server_args.pipeline_config, "streaming", False): + DenoiseStage = SanaWMStreamingDenoisingStage + else: + DenoiseStage = SanaWMDenoisingStage + self.add_stage( + DenoiseStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + # Subclasses (e.g. SanaWMTwoStagePipeline) insert latent-domain stages + # between denoising and VAE decoding. + self._maybe_add_refiner_stage(server_args) + + self._add_decoding_stage(server_args) + + def _add_decoding_stage(self, server_args: ServerArgs = None) -> None: + if server_args is not None and getattr( + server_args.pipeline_config, "streaming", False + ): + DecodeStage = SanaWMStreamingDecodingStage + else: + DecodeStage = SanaWMDecodingStage + self.add_stage( + DecodeStage( + vae=self.get_module("vae"), + pipeline=self, + component_name="vae", + ), + "decoding_stage", + ) + + def _maybe_add_refiner_stage(self, server_args: ServerArgs) -> None: + """Hook for subclasses; single-stage pipeline is a no-op.""" + return None + + +class SanaWMTwoStagePipeline(SanaWMPipeline): + """SANA-WM two-stage pipeline: SANA-WM DiT + LTX-2 latent refiner. + + Stage-1 produces a coarse 720p latent; the LTX-2 refiner runs 3 Euler steps + on it before VAE decode, matching the NVlabs ``inference_sana_wm.py`` default. + """ + + pipeline_name = "SanaWMTwoStagePipeline" + + # Stage-2 refiner sub-modules and their on-disk layout. Loaded through the + # official Diffusers/Transformers classes because NVlabs' reference refiner + # is a narrow video-only wrapper around those modules. + _REFINER_SUB_MODULES: tuple[tuple[str, str], ...] = ( + ("transformer_2", "refiner/transformer"), + ("connectors", "refiner/connectors"), + ("text_encoder_2", "refiner/text_encoder"), + # The refiner Gemma-3 ships its tokenizer files alongside the encoder. + ("tokenizer_2", "refiner/text_encoder"), + ) + + def initialize_pipeline(self, server_args: ServerArgs) -> None: + super().initialize_pipeline(server_args) + if sana_wm_skip_refiner_enabled(): + logger.info( + "SANA-WM refiner component loading skipped by " + "SGLANG_SANA_WM_SKIP_REFINER." + ) + return + self._load_refiner_modules(server_args) + + def _resolve_refiner_paths(self, server_args: ServerArgs) -> tuple[str, str]: + component_paths = getattr(server_args, "component_paths", {}) or {} + refiner_root = component_paths.get( + "refiner", os.path.join(self.model_path, "refiner") + ) + refiner_gemma_root = component_paths.get( + "refiner_text_encoder", + component_paths.get( + "text_encoder_2", os.path.join(refiner_root, "text_encoder") + ), + ) + return refiner_root, refiner_gemma_root + + def _resolve_refiner_component_path( + self, server_args: ServerArgs, module_name: str, subpath: str + ) -> str: + component_paths = getattr(server_args, "component_paths", {}) or {} + if module_name in component_paths: + return self._resolve_component_path(server_args, module_name, subpath) + + if ( + "refiner" not in component_paths + and "refiner_text_encoder" not in component_paths + ): + return self._resolve_component_path(server_args, module_name, subpath) + + refiner_root, refiner_gemma_root = self._resolve_refiner_paths(server_args) + if module_name in ("text_encoder_2", "tokenizer_2"): + return refiner_gemma_root + + rel_subpath = subpath.removeprefix("refiner/") + return os.path.join(refiner_root, rel_subpath) + + def _load_refiner_modules(self, server_args: ServerArgs) -> None: + for module_name, subpath in self._REFINER_SUB_MODULES: + component_path = self._resolve_refiner_component_path( + server_args, module_name, subpath + ) + logger.info( + "SANA-WM loading refiner component %s from %s", + module_name, + component_path, + ) + module, memory_usage = self._load_official_refiner_component( + module_name, + component_path, + server_args, + ) + self.modules[module_name] = module + self.memory_usages[module_name] = memory_usage + + @staticmethod + def _load_official_refiner_component( + module_name: str, + component_path: str, + server_args: ServerArgs, + ): + """Load SANA-WM refiner modules through the same libraries as NVlabs. + + The upstream wrapper (``diffusion/refiner/diffusers_ltx2_refiner.py``) + keeps the LTX-2 transformer/connectors as Diffusers modules and only + customizes the video-only forward surface; use that path for the + quality-critical stage-2 refiner instead of the experimental native port. + """ + + dtype = default_sana_wm_refiner_dtype(server_args) + if module_name == "transformer_2": + from diffusers.models.transformers.transformer_ltx2 import ( + LTX2VideoTransformer3DModel, + ) + + module = LTX2VideoTransformer3DModel.from_pretrained( + component_path, + torch_dtype=dtype, + ).eval() + module = OfficialDiffusersLTX2RefinerModule(module) + elif module_name == "connectors": + from diffusers.pipelines.ltx2 import LTX2TextConnectors + + module = LTX2TextConnectors.from_pretrained( + component_path, + torch_dtype=dtype, + ).eval() + elif module_name == "text_encoder_2": + from transformers import Gemma3ForConditionalGeneration + + module = Gemma3ForConditionalGeneration.from_pretrained( + component_path, + torch_dtype=dtype, + low_cpu_mem_usage=True, + ).eval() + module = OfficialGemma3TextEncoderModule(module) + elif module_name == "tokenizer_2": + from transformers import AutoTokenizer + + module = AutoTokenizer.from_pretrained(component_path) + else: + raise ValueError(f"Unsupported SANA-WM refiner component: {module_name}") + + memory_usage = get_memory_usage_of_component(module) + logger.info( + "Loaded %s: %s (official native version). model size: %s GB", + module_name, + module.__class__.__name__, + memory_usage if memory_usage is not None else "NA", + ) + return module, memory_usage or 0.0 + + def _maybe_add_refiner_stage(self, server_args: ServerArgs) -> None: + if sana_wm_skip_refiner_enabled(): + return + pc = server_args.pipeline_config + common = dict( + transformer=self.get_module("transformer_2"), + connectors=self.get_module("connectors"), + text_encoder=self.get_module("text_encoder_2"), + tokenizer=self.get_module("tokenizer_2"), + dtype=default_sana_wm_refiner_dtype(server_args), + ) + if getattr(pc, "streaming", False) and getattr(pc, "refiner_chunked", True): + stage = SanaWMStreamingRefinerStage( + **common, + block_size=int(getattr(pc, "refiner_block_size", 3)), + kv_max_frames=int(getattr(pc, "refiner_kv_max_frames", 11)), + sink_size=int(getattr(pc, "sink_size", 1)), + seed=int(getattr(pc, "refiner_seed", 42)), + ) + else: + stage = SanaWMLTX2RefinerStage(**common) + self.add_stage(stage, "sana_wm_refiner") + + def _add_decoding_stage(self, server_args: ServerArgs = None) -> None: + # Streaming and skip-refiner both route to the base decode + # (SanaWMStreamingDecodingStage / dense decode); otherwise dense refiner-decode. + streaming = server_args is not None and getattr( + server_args.pipeline_config, "streaming", False + ) + if streaming or sana_wm_skip_refiner_enabled(): + return super()._add_decoding_stage(server_args) + self.add_stage( + SanaWMRefinerDecodingStage( + vae=self.get_module("vae"), + pipeline=self, + component_name="vae", + ), + "decoding_stage", + ) + + +EntryClass = [SanaWMPipeline, SanaWMTwoStagePipeline] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_realtime_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_realtime_pipeline.py new file mode 100644 index 000000000..31f2ea6b2 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_realtime_pipeline.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import ( + SanaWMRealtimeConfig, +) +from sglang.multimodal_gen.runtime.pipelines.sana_wm_pipeline import ( + SanaWMTwoStagePipeline, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + RealtimeInputValidationStage, + RealtimeTextEncodingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm import ( + SanaWMTextEncodingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_chain import ( + SanaWMCameraCondStage, + SanaWMCausalDecodeChainStage, + SanaWMChunkedRefinerChainStage, + SanaWMCondFrameEncodeStage, + SanaWMRealtimeLatentPrepStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.refiner import ( + default_sana_wm_refiner_dtype, + sana_wm_skip_refiner_enabled, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import ( + SanaWMStreamingDenoisingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming_refiner import ( + SanaWMStreamingRefinerStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model + +DEFAULT_SANA_WM_TEXT_ENCODER = "Efficient-Large-Model/gemma-2-2b-it" + + +class SanaWMRealtimeTextEncodingStage( + RealtimeTextEncodingStage, SanaWMTextEncodingStage +): + """Realtime text encoding using SANA-WM prompt processing. + + MRO contract: ``RealtimeTextEncodingStage.forward`` (per-session cache) calls + ``super().forward``, which must resolve to ``SanaWMTextEncodingStage.forward``. + This preserves the chi prompt prefix and official prompt window used by the + batch path. + """ + + +class SanaWMRealtimePipeline(SanaWMTwoStagePipeline): + """SANA-WM realtime interactive pipeline. + + Extends the two-stage pipeline to inherit refiner sub-module loading (``transformer_2`` / + ``connectors`` / ``text_encoder_2`` / ``tokenizer_2``). The streaming refiner stage built + here is purely a carrier of those modules handed to ``SanaWMRealtimeStage``, not added to + the stage list (the realtime stage drives the incremental stage-1 session + chunked refiner + runner per user action). + """ + + pipeline_name = "SanaWMRealtimePipeline" + is_video_pipeline = True + # Must be the realtime config so get_realtime_model_adapter() resolves the + # SANA-WM adapter (the realtime registry keys on SanaWMRealtimeConfig). + pipeline_config_cls = SanaWMRealtimeConfig + + def _resolve_component_path( + self, server_args: ServerArgs, module_name: str, load_module_name: str + ) -> str: + if ( + module_name in {"text_encoder", "tokenizer"} + and module_name not in server_args.component_paths + ): + return maybe_download_model(DEFAULT_SANA_WM_TEXT_ENCODER) + return super()._resolve_component_path( + server_args, + module_name, + load_module_name, + ) + + def _build_realtime_refiner_stage(self, server_args: ServerArgs): + """Build the chunked streaming refiner carrier when refiner modules exist.""" + if sana_wm_skip_refiner_enabled(): + return None + if self.get_module("transformer_2") is None: + return None + + pc = server_args.pipeline_config + return SanaWMStreamingRefinerStage( + transformer=self.get_module("transformer_2"), + connectors=self.get_module("connectors"), + text_encoder=self.get_module("text_encoder_2"), + tokenizer=self.get_module("tokenizer_2"), + dtype=default_sana_wm_refiner_dtype(server_args), + block_size=int(getattr(pc, "refiner_block_size", 3)), + kv_max_frames=int(getattr(pc, "refiner_kv_max_frames", 11)), + sink_size=int(getattr(pc, "sink_size", 1)), + seed=int(getattr(pc, "refiner_seed", 42)), + ) + + def create_pipeline_stages(self, server_args: ServerArgs): + refiner_stage = self._build_realtime_refiner_stage(server_args) + common = dict( + transformer=self.get_module("transformer"), + vae=self.get_module("vae"), + model_path=self.model_path, + ) + self.add_stage(RealtimeInputValidationStage()) + self.add_stage( + SanaWMRealtimeTextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ) + ) + self.add_stage(SanaWMCondFrameEncodeStage(**common)) + self.add_stage( + SanaWMRealtimeLatentPrepStage( + use_refiner=refiner_stage is not None, **common + ) + ) + self.add_stage(SanaWMCameraCondStage(**common)) + self.add_stage( + SanaWMStreamingDenoisingStage( + transformer=self.get_module("transformer"), keep_resident=True + ) + ) + if refiner_stage is not None: + self.add_stage( + SanaWMChunkedRefinerChainStage(refiner_stage=refiner_stage, **common) + ) + self.add_stage(SanaWMCausalDecodeChainStage(**common)) + + +EntryClass = SanaWMRealtimePipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 66a765aef..a700d18f8 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -436,4 +436,5 @@ class OutputBatch: self.rollout_trajectory_data = None self.trajectory_decoded = None self.output_file_paths = None + self.raw_frame_batches = None self.noise_pred = None diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py index e1157c507..1936eaddf 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py @@ -65,6 +65,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av i from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( LTX2DenoisingStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_diffusion import ( + RealtimeDiffusionStage, + RealtimeStageComponent, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_input_validation import ( RealtimeInputValidationStage, ) @@ -95,6 +99,8 @@ __all__ = [ "PipelineStage", "InputValidationStage", "RealtimeInputValidationStage", + "RealtimeDiffusionStage", + "RealtimeStageComponent", "TimestepPreparationStage", "DMDTimestepPreparationStage", "LatentPreparationStage", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index a14d173c0..1a130cd87 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -31,6 +31,33 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) +def scale_and_shift_latents(latents: torch.Tensor, server_args, vae) -> torch.Tensor: + """De-normalize latents before VAE decode (single shared implementation). + + Used by DecodingStage.scale_and_shift and by realtime stages that decode + outside a DecodingStage instance. + """ + scaling_factor, shift_factor = ( + server_args.pipeline_config.get_decode_scale_and_shift( + latents.device, latents.dtype, vae + ) + ) + + # 1. scale + if isinstance(scaling_factor, torch.Tensor): + latents = latents / scaling_factor.to(latents.device, latents.dtype) + else: + latents = latents / scaling_factor + + # 2. apply shifting if needed + if shift_factor is not None: + if isinstance(shift_factor, torch.Tensor): + latents = latents + shift_factor.to(latents.device, latents.dtype) + else: + latents = latents + shift_factor + return latents + + def _ensure_tensor_decode_output(decode_output): """ Ensure VAE decode output is a tensor. @@ -106,25 +133,7 @@ class DecodingStage(PipelineStage): return result def scale_and_shift(self, latents: torch.Tensor, server_args): - scaling_factor, shift_factor = ( - server_args.pipeline_config.get_decode_scale_and_shift( - latents.device, latents.dtype, self.vae - ) - ) - - # 1. scale - if isinstance(scaling_factor, torch.Tensor): - latents = latents / scaling_factor.to(latents.device, latents.dtype) - else: - latents = latents / scaling_factor - - # 2. apply shifting if needed - if shift_factor is not None: - if isinstance(shift_factor, torch.Tensor): - latents += shift_factor.to(latents.device, latents.dtype) - else: - latents += shift_factor - return latents + return scale_and_shift_latents(latents, server_args, self.vae) @torch.no_grad() def decode( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/__init__.py new file mode 100644 index 000000000..bfa9f95dc --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/__init__.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SANA-WM pipeline stages (package). + +The realtime serving framework drives ``SanaWMRealtimeStage`` over the +``/v1/realtime_video`` WebSocket. Base stages + helpers are re-exported here +for back-compat. +""" + +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.base import ( + SanaWMBeforeDenoisingStage, + SanaWMDecodingStage, + SanaWMDenoisingStage, + SanaWMTextEncodingStage, + _align_sana_wm_cfg_text_conditions, + configure_sana_wm_ltx2_vae_for_long_video, + parse_sana_wm_action_string, + sana_wm_action_to_camera_to_world, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_stage import ( + SanaWMRealtimeStage, +) + +__all__ = [ + "SanaWMBeforeDenoisingStage", + "SanaWMDecodingStage", + "SanaWMDenoisingStage", + "SanaWMTextEncodingStage", + "_align_sana_wm_cfg_text_conditions", + "configure_sana_wm_ltx2_vae_for_long_video", + "parse_sana_wm_action_string", + "sana_wm_action_to_camera_to_world", + "SanaWMRealtimeStage", +] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py new file mode 100644 index 000000000..f4282654b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py @@ -0,0 +1,2325 @@ +# SPDX-License-Identifier: Apache-2.0 + +import inspect +import math +import os +import time +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +import torchvision.transforms.functional as TF +from diffusers.utils.torch_utils import randn_tensor +from PIL import Image + +from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.distributed.communication_op import ( + cfg_model_parallel_all_reduce, +) +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, +) +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + ComponentUse, +) +from sglang.multimodal_gen.runtime.models.dits.sana_wm_components import ( + compute_chunk_plucker, +) +from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import ( + get_or_create_request_scheduler, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( + PipelineStage, + StageParallelismType, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + DenoisingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import ( + TextEncodingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + +SANA_WM_TARGET_HEIGHT = 704 +SANA_WM_TARGET_WIDTH = 1280 +_SANA_WM_DIAGNOSTICS_ENVS = ( + "SGLANG_SANA_WM_DIAGNOSTICS", + "SGLANG_SANA_WM_LOG_TENSOR_STATS", +) + +_SANA_WM_DEFAULT_VAE_TILE_MIN_FRAMES = 96 +_SANA_WM_DEFAULT_VAE_TILE_STRIDE_FRAMES = 64 +_SANA_WM_BIDIRECTIONAL_DEFAULT_TRANSLATION_SPEED = 0.05 +_SANA_WM_DEFAULT_TRANSLATION_SPEED = ( + 0.04 # match official streaming (STREAMING_TRANSLATION_SPEED) +) +_SANA_WM_DEFAULT_ROTATION_SPEED_DEG = 1.2 +_SANA_WM_DEFAULT_PITCH_LIMIT_DEG = 85.0 +_SANA_WM_CONDITION_IMAGE_PREPROCESS_KEY = "sana_wm_condition_image_preprocess" +_SANA_WM_ALLOWED_ACTION_KEYS = frozenset("wasdijkl") + + +def sana_wm_pil_to_model_tensor( + image: Image.Image, *, device: torch.device, dtype: torch.dtype +) -> torch.Tensor: + arr = np.asarray(image, dtype=np.float32) / 255.0 + tensor = torch.from_numpy(arr).permute(2, 0, 1) + return (tensor * 2.0 - 1.0).unsqueeze(0).unsqueeze(2).to(device=device, dtype=dtype) + + +def _sana_wm_rot_x(angle_rad: float) -> np.ndarray: + c, s = np.cos(angle_rad), np.sin(angle_rad) + return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=np.float64) + + +def _sana_wm_rot_y(angle_rad: float) -> np.ndarray: + c, s = np.cos(angle_rad), np.sin(angle_rad) + return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64) + + +def sana_wm_compute_resize_crop_geometry( + src_w: int, src_h: int, target_h: int, target_w: int +) -> tuple[int, int, int, int]: + scale = max(target_h / float(src_h), target_w / float(src_w)) + resized_w = max(target_w, int(round(src_w * scale))) + resized_h = max(target_h, int(round(src_h * scale))) + left = (resized_w - target_w) // 2 + top = (resized_h - target_h) // 2 + return resized_w, resized_h, left, top + + +def normalize_sana_wm_camera_actions( + payload: Any, + *, + allow_none: bool = False, + error_label: str = "camera_actions", +) -> list[list[str]]: + if payload is None and allow_none: + return [] + if not isinstance(payload, list): + raise ValueError(f"{error_label} must be list[list[str]]") + out: list[list[str]] = [] + for frame_actions in payload: + if not isinstance(frame_actions, list): + raise ValueError(f"{error_label} must be list[list[str]]") + out.append([str(key).lower() for key in frame_actions]) + return out + + +def parse_sana_wm_action_string(action: str) -> list[list[str]]: + cleaned = "".join(action.replace("，", ",").split()) + if not cleaned: + raise ValueError("action string is empty") + + per_frame: list[list[str]] = [] + for segment in cleaned.split(","): + if not segment or "-" not in segment: + raise ValueError( + f"invalid action segment {segment!r}; expected '-'" + ) + keys_part, duration = segment.rsplit("-", 1) + if not duration.isdigit() or int(duration) <= 0: + raise ValueError(f"invalid duration in action segment {segment!r}") + + if keys_part.lower() == "none": + keys: list[str] = [] + else: + bad = sorted( + { + char + for char in keys_part.lower() + if char not in _SANA_WM_ALLOWED_ACTION_KEYS + } + ) + if bad: + raise ValueError( + f"unknown action keys {bad}; allowed keys are " + f"{sorted(_SANA_WM_ALLOWED_ACTION_KEYS)}" + ) + keys = sorted(set(keys_part.lower())) + # Fresh list per frame: repeated frames must NOT alias one list object + # (callers could mutate a frame and silently edit all its repeats). + per_frame.extend([list(keys) for _ in range(int(duration))]) + return per_frame + + +def sana_wm_action_to_camera_to_world_array( + action: str, + *, + translation_speed: float = _SANA_WM_BIDIRECTIONAL_DEFAULT_TRANSLATION_SPEED, + rotation_speed_deg: float = _SANA_WM_DEFAULT_ROTATION_SPEED_DEG, + pitch_limit_deg: float = _SANA_WM_DEFAULT_PITCH_LIMIT_DEG, + strafe_yaw_coupling: float = 0.4, +) -> np.ndarray: + per_frame = parse_sana_wm_action_string(action) + rotate_rad = math.radians(rotation_speed_deg) + pitch_limit_rad = math.radians(pitch_limit_deg) + current = np.eye(4, dtype=np.float64) + poses = [current.copy()] + current_pitch = 0.0 + + for keys in per_frame: + held = set(keys) + rotation = current[:3, :3] + translation = current[:3, 3] + + pitch_delta = (rotate_rad if "i" in held else 0.0) - ( + rotate_rad if "k" in held else 0.0 + ) + next_pitch = current_pitch + pitch_delta + if -pitch_limit_rad <= next_pitch <= pitch_limit_rad: + current_pitch = next_pitch + else: + pitch_delta = 0.0 + + yaw_delta = (rotate_rad if "l" in held else 0.0) - ( + rotate_rad if "j" in held else 0.0 + ) + strafe_yaw = (rotate_rad if "d" in held else 0.0) - ( + rotate_rad if "a" in held else 0.0 + ) + yaw_delta += strafe_yaw_coupling * strafe_yaw + rotation = _sana_wm_rot_y(yaw_delta) @ rotation @ _sana_wm_rot_x(pitch_delta) + + forward = rotation[:, 2].copy() + right = rotation[:, 0].copy() + forward[1] = 0.0 + right[1] = 0.0 + forward_norm = float(np.linalg.norm(forward)) + right_norm = float(np.linalg.norm(right)) + if forward_norm > 0.0: + forward /= forward_norm + 1e-6 + if right_norm > 0.0: + right /= right_norm + 1e-6 + + move = np.zeros(3, dtype=np.float64) + if "w" in held: + move += forward * translation_speed + if "s" in held: + move -= forward * translation_speed + if "d" in held: + move += right * translation_speed + if "a" in held: + move -= right * translation_speed + + current = np.eye(4, dtype=np.float64) + current[:3, :3] = rotation + current[:3, 3] = translation + move + poses.append(current.copy()) + + return np.stack(poses, axis=0).astype(np.float32) + + +def sana_wm_load_camera(path: Path) -> np.ndarray: + c2w = np.load(path).astype(np.float32) + if c2w.ndim != 3 or c2w.shape[1:] != (4, 4): + raise ValueError( + f"camera trajectory must have shape (F, 4, 4); got {c2w.shape}" + ) + return c2w + + +def sana_wm_load_intrinsics(path: Path, num_frames: int) -> np.ndarray: + arr = np.load(path).astype(np.float32) + if arr.shape == (4,): + return np.broadcast_to(arr, (num_frames, 4)).copy() + if arr.shape == (3, 3): + vec = np.array([arr[0, 0], arr[1, 1], arr[0, 2], arr[1, 2]], dtype=np.float32) + return np.broadcast_to(vec, (num_frames, 4)).copy() + if arr.ndim == 3 and arr.shape[1:] == (3, 3) and arr.shape[0] >= num_frames: + arr = arr[:num_frames] + return np.stack( + [arr[:, 0, 0], arr[:, 1, 1], arr[:, 0, 2], arr[:, 1, 2]], axis=1 + ) + raise ValueError( + f"unsupported intrinsics shape {arr.shape}; expected (4,), (3,3), or (F,3,3)" + ) + + +def snap_sana_wm_num_frames( + num_frames: int, stride: int = 8, upper_bound: int | None = None +) -> int: + if num_frames < 1: + return 1 + if (num_frames - 1) % stride == 0: + return min(num_frames, upper_bound) if upper_bound is not None else num_frames + + floor_candidate = num_frames - ((num_frames - 1) % stride) + ceil_candidate = floor_candidate + stride + snapped = ( + floor_candidate + if num_frames - floor_candidate < ceil_candidate - num_frames + else ceil_candidate + ) + if upper_bound is not None and snapped > upper_bound: + snapped = floor_candidate + return max(snapped, 1) + + +def sana_wm_resize_and_center_crop( + image: Image.Image, + target_h: int = SANA_WM_TARGET_HEIGHT, + target_w: int = SANA_WM_TARGET_WIDTH, +) -> tuple[Image.Image, tuple[int, int], tuple[int, int], tuple[int, int]]: + src_w, src_h = image.size + resized_w, resized_h, left, top = sana_wm_compute_resize_crop_geometry( + src_w, src_h, target_h, target_w + ) + resized = image.resize((resized_w, resized_h), Image.LANCZOS) + crop = resized.crop((left, top, left + target_w, top + target_h)) + return crop, (src_w, src_h), (resized_w, resized_h), (left, top) + + +def sana_wm_vae_scaling_factor(vae, pipeline_config) -> float: + """Resolve the VAE scaling factor (batch-reference precedence). + + vae.config.scaling_factor -> vae.scaling_factor -> + pipeline_config.vae_config.arch_config.scaling_factor -> 1.0; zero is + treated as unset. Single source for both the batch and realtime + first-frame encode paths. + """ + scaling_factor = ( + getattr(getattr(vae, "config", None), "scaling_factor", None) + or getattr(vae, "scaling_factor", None) + or getattr( + getattr(getattr(pipeline_config, "vae_config", None), "arch_config", None), + "scaling_factor", + None, + ) + or 1.0 + ) + if isinstance(scaling_factor, torch.Tensor): + return float(scaling_factor.item()) + scaling_factor = float(scaling_factor) + return 1.0 if scaling_factor == 0.0 else scaling_factor + + +def sana_wm_normalize_vae_latents( + vae, z: torch.Tensor, pipeline_config +) -> torch.Tensor: + """Normalize freshly-encoded VAE latents (batch-reference semantics). + + ``(z - latents_mean) * scaling_factor / latents_std`` when the VAE carries + mean/std buffers (the LTX-2 causal VAE does); legacy shift-then-scale + otherwise. ``z`` is expected in float32 (drifting this was parity root + cause #2). + """ + latents_mean = getattr(vae, "latents_mean", None) + latents_std = getattr(vae, "latents_std", None) + scaling_factor = sana_wm_vae_scaling_factor(vae, pipeline_config) + if sana_wm_diagnostics_enabled(): + logger.info( + "[SANA-WM diagnostics] VAE encode normalization: " + "has_latents_mean_std=%s scaling_factor=%.6g", + isinstance(latents_mean, torch.Tensor) + and isinstance(latents_std, torch.Tensor), + scaling_factor, + ) + if isinstance(latents_mean, torch.Tensor) and isinstance(latents_std, torch.Tensor): + latents_mean = latents_mean.to(device=z.device, dtype=z.dtype).view( + 1, -1, 1, 1, 1 + ) + latents_std = latents_std.to(device=z.device, dtype=z.dtype).view( + 1, -1, 1, 1, 1 + ) + return (z - latents_mean) * scaling_factor / latents_std + + # Legacy VAE convention: encode applies shift before scaling. + shift_factor = getattr(vae, "shift_factor", None) + if shift_factor is not None: + z = z - ( + shift_factor.to(z.device, z.dtype) + if isinstance(shift_factor, torch.Tensor) + else shift_factor + ) + return z * scaling_factor + + +def sana_wm_action_to_camera_to_world( + action: str, + *, + translation_speed: float = _SANA_WM_DEFAULT_TRANSLATION_SPEED, + rotation_speed_deg: float = _SANA_WM_DEFAULT_ROTATION_SPEED_DEG, + pitch_limit_deg: float = _SANA_WM_DEFAULT_PITCH_LIMIT_DEG, +) -> torch.Tensor: + """Roll out upstream SANA-WM action DSL to a camera-to-world trajectory. + + Coordinate convention is OpenCV: +X right, +Y down, +Z forward. Returned + shape is ``(N+1, 4, 4)``, float32. + + Uses the same SANA-WM kinematics as the realtime path, including the + strafe->yaw coupling ``yaw += 0.4 * (d - a)`` from the reference. + """ + return torch.from_numpy( + sana_wm_action_to_camera_to_world_array( + action, + translation_speed=translation_speed, + rotation_speed_deg=rotation_speed_deg, + pitch_limit_deg=pitch_limit_deg, + ) + ) + + +def _resolve_sana_wm_vae_frame_tile_value( + pipeline_config: SanaWMPipelineConfig, + direct_attr: str, + vae_config_attr: str, + default: int, +) -> int: + direct_value = getattr(pipeline_config, direct_attr, default) + if direct_value != default: + return int(direct_value) + + vae_config = getattr(pipeline_config, "vae_config", None) + nested_value = getattr(vae_config, vae_config_attr, None) + return int(nested_value or default) + + +def sana_wm_diagnostics_enabled() -> bool: + """Whether to emit detailed SANA-WM tensor-quality diagnostics.""" + return any( + os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on", "debug"} + for name in _SANA_WM_DIAGNOSTICS_ENVS + ) + + +def log_sana_wm_tensor_stats(label: str, tensor: torch.Tensor | None) -> None: + """Log compact tensor statistics for reference-quality alignment. + + Enable with ``SGLANG_SANA_WM_DIAGNOSTICS=1``. The fingerprint is a strided + sum over at most ~4096 values, for spotting deterministic drift. + """ + if not sana_wm_diagnostics_enabled(): + return + if tensor is None: + logger.info("[SANA-WM diagnostics] %s: None", label) + return + if not isinstance(tensor, torch.Tensor): + logger.info( + "[SANA-WM diagnostics] %s: non-tensor type=%s", + label, + type(tensor).__name__, + ) + return + + with torch.no_grad(): + data = tensor.detach() + if data.numel() == 0: + logger.info( + "[SANA-WM diagnostics] %s: shape=%s dtype=%s device=%s empty", + label, + tuple(data.shape), + data.dtype, + data.device, + ) + return + + stats = data.float() + finite = torch.isfinite(stats) + finite_ratio = float(finite.float().mean().item()) + finite_stats = stats[finite] if bool(finite.any().item()) else stats.reshape(-1) + flat = finite_stats.reshape(-1) + stride = max(1, flat.numel() // 4096) + fingerprint = float(flat[::stride].sum().item()) + std = float(finite_stats.std(unbiased=False).item()) + logger.info( + "[SANA-WM diagnostics] %s: shape=%s dtype=%s device=%s " + "finite=%.6f min=%.6g max=%.6g mean=%.6g std=%.6g " + "l2=%.6g fingerprint=%.6g", + label, + tuple(data.shape), + data.dtype, + data.device, + finite_ratio, + float(finite_stats.min().item()), + float(finite_stats.max().item()), + float(finite_stats.mean().item()), + std, + float(torch.linalg.vector_norm(finite_stats).item()), + fingerprint, + ) + + +def configure_sana_wm_ltx2_vae_for_long_video( + vae: Any, + pipeline_config: SanaWMPipelineConfig, + *, + log_info: Any | None = None, +) -> None: + """Apply SANA-WM's upstream LTX-2 VAE tiling knobs. + + The LTX-2 VAE implements temporal tiled decode but, unlike the generic VAE + base, does not enable it from ``enable_tiling()`` alone. Without this, + 321-frame 720p decode can enter one large Conv3d path and hit PyTorch's + 32-bit index math limit. + """ + + min_frames = _resolve_sana_wm_vae_frame_tile_value( + pipeline_config, + "vae_tile_sample_min_num_frames", + "tile_sample_min_num_frames", + _SANA_WM_DEFAULT_VAE_TILE_MIN_FRAMES, + ) + stride_frames = _resolve_sana_wm_vae_frame_tile_value( + pipeline_config, + "vae_tile_sample_stride_num_frames", + "tile_sample_stride_num_frames", + _SANA_WM_DEFAULT_VAE_TILE_STRIDE_FRAMES, + ) + + use_tiling = bool(getattr(pipeline_config, "vae_tiling", True)) + if use_tiling and hasattr(vae, "enable_tiling"): + try: + vae.enable_tiling( + tile_sample_min_num_frames=min_frames, + tile_sample_stride_num_frames=stride_frames, + ) + except TypeError: + vae.enable_tiling() + + if hasattr(vae, "use_framewise_encoding"): + vae.use_framewise_encoding = bool( + getattr(pipeline_config, "vae_framewise_encoding", True) + ) + if hasattr(vae, "use_framewise_decoding"): + vae.use_framewise_decoding = bool( + getattr(pipeline_config, "vae_framewise_decoding", True) + ) + + if hasattr(vae, "tile_sample_min_num_frames"): + vae.tile_sample_min_num_frames = min_frames + if hasattr(vae, "tile_sample_stride_num_frames"): + vae.tile_sample_stride_num_frames = stride_frames + + if log_info is not None: + log_info( + "SANA-WM VAE tiling configured: spatial=%s, framewise_encode=%s, " + "framewise_decode=%s, tile_frames_min=%d, tile_frames_stride=%d", + getattr(vae, "use_tiling", use_tiling), + getattr(vae, "use_framewise_encoding", None), + getattr(vae, "use_framewise_decoding", None), + getattr(vae, "tile_sample_min_num_frames", min_frames), + getattr(vae, "tile_sample_stride_num_frames", stride_frames), + ) + + +class SanaWMDecodingStage(DecodingStage): + """Decode SANA-WM LTX-2 latents with upstream long-video VAE settings.""" + + @torch.no_grad() + def decode( + self, + latents: torch.Tensor, + server_args: ServerArgs, + *, + vae_dtype: torch.dtype, + ) -> torch.Tensor: + configure_sana_wm_ltx2_vae_for_long_video( + self.vae, + server_args.pipeline_config, + log_info=self.log_info, + ) + frames = super().decode(latents, server_args, vae_dtype=vae_dtype) + log_sana_wm_tensor_stats("decode.frames", frames) + return frames + + +def _first_tensor(value: Any) -> torch.Tensor | None: + if isinstance(value, (list, tuple)): + return value[0] if value else None + return value if isinstance(value, torch.Tensor) else None + + +def _to_device_dtype( + value: torch.Tensor | None, + *, + device: torch.device, + dtype: torch.dtype | None = None, +) -> torch.Tensor | None: + if value is None: + return None + if dtype is None: + return value.to(device=device) + return value.to(device=device, dtype=dtype) + + +def _cat_optional_tensors( + neg: torch.Tensor | None, + pos: torch.Tensor | None, +) -> torch.Tensor | None: + if neg is None and pos is None: + return None + if neg is None: + return pos + if pos is None: + return neg + return torch.cat([neg, pos], dim=0) + + +def _text_sequence_dim(tensor: torch.Tensor) -> int: + return -2 if tensor.ndim >= 3 else -1 + + +def _pad_text_sequence( + tensor: torch.Tensor | None, + target_length: int, +) -> torch.Tensor | None: + if tensor is None: + return None + seq_dim = _text_sequence_dim(tensor) + current_length = tensor.shape[seq_dim] + if current_length == target_length: + return tensor + if current_length > target_length: + index = [slice(None)] * tensor.ndim + index[seq_dim] = slice(0, target_length) + return tensor[tuple(index)] + + pad_shape = list(tensor.shape) + pad_shape[seq_dim] = target_length - current_length + padding = torch.zeros(pad_shape, device=tensor.device, dtype=tensor.dtype) + return torch.cat([tensor, padding], dim=seq_dim) + + +def _default_attention_mask_for_embeds(embeds: torch.Tensor) -> torch.Tensor: + return torch.ones( + (embeds.shape[0], embeds.shape[-2]), + device=embeds.device, + dtype=torch.long, + ) + + +def _align_sana_wm_cfg_text_conditions( + pos_embeds: torch.Tensor, + neg_embeds: torch.Tensor | None, + pos_mask: torch.Tensor | None, + neg_mask: torch.Tensor | None, +) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, +]: + if neg_embeds is None: + return pos_embeds, neg_embeds, pos_mask, neg_mask + + target_length = max(pos_embeds.shape[-2], neg_embeds.shape[-2]) + pos_embeds = _pad_text_sequence(pos_embeds, target_length) + neg_embeds = _pad_text_sequence(neg_embeds, target_length) + + if pos_mask is None: + pos_mask = _default_attention_mask_for_embeds(pos_embeds) + if neg_mask is None: + neg_mask = _default_attention_mask_for_embeds(neg_embeds) + pos_mask = _pad_text_sequence(pos_mask, target_length) + neg_mask = _pad_text_sequence(neg_mask, target_length) + return pos_embeds, neg_embeds, pos_mask, neg_mask + + +class SanaWMTextEncodingStage(TextEncodingStage): + """Gemma-2 text encoding that mirrors NVlabs SANA-WM inference. + + The official script prepends a long ``chi_prompt`` only to the positive + branch, tokenizes that longer string, then keeps token 0 and the last + 299 tokens. The negative branch remains a normal 300-token padded prompt. + Keeping this local avoids bending the shared TextEncodingStage around a + model-specific prompt-window contract. + """ + + @staticmethod + def _text_encoder_max_length(server_args: ServerArgs) -> int: + encoder_cfg = server_args.pipeline_config.text_encoder_configs[0] + arch_config = getattr(encoder_cfg, "arch_config", None) + return int(getattr(arch_config, "text_len", 300) or 300) + + @staticmethod + def _chi_prompt(server_args: ServerArgs) -> str: + parts = getattr(server_args.pipeline_config, "chi_prompt", ()) or () + return "\n".join(parts) + + @staticmethod + def _select_official_prompt_window( + tensor: torch.Tensor | None, + max_length: int, + ) -> torch.Tensor | None: + if tensor is None: + return None + seq_dim = _text_sequence_dim(tensor) + if tensor.shape[seq_dim] <= max_length: + return tensor + index = [slice(None)] * tensor.ndim + tail_start = tensor.shape[seq_dim] - max_length + 1 + select = torch.cat( + [ + torch.zeros(1, device=tensor.device, dtype=torch.long), + torch.arange( + tail_start, + tensor.shape[seq_dim], + device=tensor.device, + dtype=torch.long, + ), + ], + dim=0, + ) + index[seq_dim] = select + return tensor[tuple(index)] + + @staticmethod + def _seq_lens_from_masks(masks: list[torch.Tensor | None]) -> list[list[int]]: + seq_lens = [] + for mask in masks: + if mask is None: + seq_lens.append([]) + else: + seq_lens.append([int(x) for x in mask.long().sum(dim=-1).tolist()]) + return seq_lens + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if len(self.text_encoders) != 1: + raise ValueError( + "SANA-WM stage-1 expects exactly one Gemma-2 text encoder." + ) + assert batch.prompt is not None + + max_length = self._text_encoder_max_length(server_args) + chi_prompt = self._chi_prompt(server_args) + prompt_text = batch.prompt + if isinstance(prompt_text, str): + prompt_text = [prompt_text] + else: + prompt_text = list(prompt_text) + + tokenizer = self.tokenizers[0] + if chi_prompt: + prompt_text = [chi_prompt + text for text in prompt_text] + max_length_all = len(tokenizer.encode(chi_prompt)) + max_length - 2 + else: + max_length_all = max_length + + ( + prompt_embeds_list, + prompt_masks_list, + pooler_embeds_list, + prompt_embeds_masks_list, + _prompt_seq_lens_list, + ) = self.encode_text( + prompt_text, + server_args, + encoder_index=[0], + return_attention_mask=True, + max_length=max_length_all, + padding="max_length", + truncation=True, + ) + + prompt_embeds_list = [ + self._select_official_prompt_window(tensor, max_length) + for tensor in prompt_embeds_list + ] + prompt_masks_list = [ + self._select_official_prompt_window(tensor, max_length) + for tensor in prompt_masks_list + ] + prompt_embeds_masks_list = [ + self._select_official_prompt_window(tensor, max_length) + for tensor in prompt_embeds_masks_list + ] + prompt_seq_lens_list = self._seq_lens_from_masks(prompt_masks_list) + + if batch.do_classifier_free_guidance: + assert isinstance(batch.negative_prompt, str) + ( + neg_embeds_list, + neg_masks_list, + neg_pooler_embeds_list, + neg_embeds_masks_list, + _neg_seq_lens_list, + ) = self.encode_text( + batch.negative_prompt, + server_args, + encoder_index=[0], + return_attention_mask=True, + max_length=max_length, + padding="max_length", + truncation=True, + ) + neg_seq_lens_list = self._seq_lens_from_masks(neg_masks_list) + + self._append_positive_text_outputs( + batch, + prompt_embeds_list, + prompt_masks_list, + pooler_embeds_list, + prompt_embeds_masks_list, + prompt_seq_lens_list, + ) + + if batch.do_classifier_free_guidance: + self._append_negative_text_outputs( + batch, + prompt_embeds_list, + neg_embeds_list, + neg_masks_list, + neg_pooler_embeds_list, + neg_embeds_masks_list, + neg_seq_lens_list, + ) + + self.log_info( + "SANA-WM text encoded with chi_prompt=%s, prompt_window=%d, " + "positive_raw_window=%d", + "yes" if chi_prompt else "no", + max_length, + max_length_all, + ) + log_sana_wm_tensor_stats("text.prompt_embeds", prompt_embeds_list[0]) + if batch.do_classifier_free_guidance: + log_sana_wm_tensor_stats("text.negative_prompt_embeds", neg_embeds_list[0]) + + return batch + + +class SanaWMDenoisingStage(DenoisingStage): + """SANA-WM stage-1 sampler matching NVlabs ``flow_euler_ltx``. + + The generic denoising stage uses one scalar timestep for every latent token + and updates the whole tensor. Official SANA-WM inference uses per-frame + timesteps: the first-frame condition stays at timestep 0 and is not updated, + while the remaining latent frames denoise normally. + """ + + @property + def parallelism_type(self) -> StageParallelismType: + if self.server_args.enable_cfg_parallel: + return StageParallelismType.CFG_PARALLEL + return StageParallelismType.REPLICATED + + @staticmethod + def _combine_cfg_parallel_noise( + noise_pred: torch.Tensor, + guidance_scale: float, + cfg_rank: int, + ) -> torch.Tensor: + if cfg_rank == 0: + partial = guidance_scale * noise_pred + elif cfg_rank == 1: + partial = (1.0 - guidance_scale) * noise_pred + else: + partial = torch.zeros_like(noise_pred) + return cfg_model_parallel_all_reduce(partial) + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if batch.latents is None: + raise ValueError("SANA-WM denoising requires initialized latents.") + if batch.latents.ndim != 5: + raise ValueError( + "SANA-WM denoising expects 5D latents shaped (B, C, T, H, W), " + f"got {tuple(batch.latents.shape)}." + ) + + device = get_local_torch_device() + target_dtype = PRECISION_TO_TYPE.get( + getattr(server_args.pipeline_config, "dit_precision", "bf16"), + torch.bfloat16, + ) + scheduler = getattr( + batch, "scheduler", None + ) or get_or_create_request_scheduler(batch, self.scheduler) + timesteps = batch.timesteps + if timesteps is None: + raise ValueError("SANA-WM denoising requires prepared timesteps.") + + latents = batch.latents.to(device=device, dtype=target_dtype) + init_latents = latents.clone() + condition_mask = torch.zeros_like(latents) + condition_mask[:, :, :1] = 1 + + pos_embeds = _to_device_dtype( + _first_tensor(server_args.pipeline_config.get_pos_prompt_embeds(batch)), + device=device, + dtype=target_dtype, + ) + pos_mask = _to_device_dtype( + _first_tensor(batch.prompt_attention_mask), device=device + ) + if pos_embeds is None: + raise ValueError("SANA-WM denoising requires positive prompt embeds.") + + do_cfg = bool(batch.do_classifier_free_guidance) + neg_embeds = None + neg_mask = None + if do_cfg: + neg_embeds = _to_device_dtype( + _first_tensor(server_args.pipeline_config.get_neg_prompt_embeds(batch)), + device=device, + dtype=target_dtype, + ) + neg_mask = _to_device_dtype( + _first_tensor(batch.negative_attention_mask), device=device + ) + if neg_embeds is None: + raise ValueError("SANA-WM CFG requires negative prompt embeds.") + + pos_embeds, neg_embeds, pos_mask, neg_mask = ( + _align_sana_wm_cfg_text_conditions( + pos_embeds, neg_embeds, pos_mask, neg_mask + ) + ) + + extra = batch.extra or {} + diffusers_kwargs = extra.get("diffusers_kwargs", {}) + if not isinstance(diffusers_kwargs, dict): + diffusers_kwargs = {} + chunk_kwargs = {} + for key in ("chunk_index", "chunk_size", "chunk_split_strategy"): + value = extra.get(key, diffusers_kwargs.get(key)) + if value is not None: + chunk_kwargs[key] = value + camera_conditions = _to_device_dtype( + extra.get("camera_conditions"), device=device, dtype=target_dtype + ) + chunk_plucker = _to_device_dtype( + extra.get("chunk_plucker"), device=device, dtype=target_dtype + ) + + cfg_parallel = bool(server_args.enable_cfg_parallel and do_cfg) + cfg_rank = get_classifier_free_guidance_rank() if cfg_parallel else 0 + if cfg_parallel and get_classifier_free_guidance_world_size() > 2: + logger.warning_once( + "SANA-WM CFG parallel uses two guidance branches; extra CFG ranks " + "run dummy forwards and contribute zeros." + ) + + if cfg_parallel: + if cfg_rank == 1: + branch_embeds = neg_embeds + branch_mask = neg_mask + else: + branch_embeds = pos_embeds + branch_mask = pos_mask + model_kwargs = { + "encoder_hidden_states": branch_embeds, + "encoder_attention_mask": branch_mask, + "camera_conditions": camera_conditions, + "chunk_plucker": chunk_plucker, + } + else: + model_kwargs = { + "encoder_hidden_states": ( + torch.cat([neg_embeds, pos_embeds], dim=0) if do_cfg else pos_embeds + ), + "encoder_attention_mask": ( + _cat_optional_tensors(neg_mask, pos_mask) if do_cfg else pos_mask + ), + "camera_conditions": ( + torch.cat([camera_conditions, camera_conditions], dim=0) + if do_cfg and camera_conditions is not None + else camera_conditions + ), + "chunk_plucker": ( + torch.cat([chunk_plucker, chunk_plucker], dim=0) + if do_cfg and chunk_plucker is not None + else chunk_plucker + ), + } + model_kwargs.update(chunk_kwargs) + + condition_mask_input = ( + condition_mask + if cfg_parallel or not do_cfg + else torch.cat([condition_mask, condition_mask], dim=0) + ) + timestep_condition_limit = (1.0 - condition_mask_input.float()) * 1000.0 + + self.log_info( + "SANA-WM flow_euler_ltx denoising: latent=%s, steps=%d, cfg=%s, " + "cfg_parallel=%s, guidance_scale=%.4f, first_frame_locked=yes", + tuple(latents.shape), + len(timesteps), + do_cfg, + cfg_parallel, + float(getattr(batch, "guidance_scale", 1.0) or 1.0), + ) + log_sana_wm_tensor_stats("denoise.input_latents", latents) + + start_time = time.perf_counter() + with self.use_declared_component( + component_name="transformer", module=self.transformer + ) as transformer: + assert transformer is not None + self.transformer = transformer + + for step_idx, t in enumerate(self.progress_bar(timesteps)): + if cfg_parallel: + latent_model_input = latents + else: + latent_model_input = ( + torch.cat([latents, latents], dim=0) if do_cfg else latents + ) + + timestep = t.expand(condition_mask_input.shape).float() + timestep = torch.minimum(timestep, timestep_condition_limit) + model_timestep = timestep[:, :1, :, 0, 0] + + with set_forward_context( + current_timestep=step_idx, + attn_metadata=None, + forward_batch=batch, + ): + noise_pred = transformer( + hidden_states=latent_model_input.to(target_dtype), + timestep=model_timestep, + **model_kwargs, + ) + + if do_cfg: + guidance_scale = float(getattr(batch, "guidance_scale", 1.0) or 1.0) + if cfg_parallel: + noise_pred = self._combine_cfg_parallel_noise( + noise_pred, guidance_scale, cfg_rank + ) + else: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * ( + noise_pred_text - noise_pred_uncond + ) + timestep = timestep.chunk(2)[0] + + latents_dtype = latents.dtype + latents_shape = latents.shape + batch_size, channels, _, _, _ = latents_shape + scheduler_output = scheduler.step( + -noise_pred.reshape(batch_size, channels, -1).transpose(1, 2), + t, + latents.reshape(batch_size, channels, -1).transpose(1, 2), + per_token_timesteps=timestep.reshape(batch_size, channels, -1)[ + :, 0 + ], + return_dict=False, + )[0] + denoised_latents = scheduler_output.transpose(1, 2).reshape( + latents_shape + ) + + tokens_to_denoise = t.float() / 1000.0 - 1e-6 < (1.0 - condition_mask) + latents = torch.where(tokens_to_denoise, denoised_latents, latents) + if latents.dtype != latents_dtype: + latents = latents.to(latents_dtype) + + if sana_wm_diagnostics_enabled() and ( + step_idx == 0 or step_idx == len(timesteps) - 1 + ): + log_sana_wm_tensor_stats( + f"denoise.step_{step_idx}.noise_pred", noise_pred + ) + log_sana_wm_tensor_stats( + f"denoise.step_{step_idx}.latents", latents + ) + + log_sana_wm_tensor_stats("denoise.output_latents", latents) + unchanged = (latents[:, :, :1] - init_latents[:, :, :1]).abs().max().item() + self.log_info( + "SANA-WM flow_euler_ltx denoising finished in %.4f seconds; " + "first_frame_max_delta=%.6g", + time.perf_counter() - start_time, + float(unchanged), + ) + batch.latents = server_args.pipeline_config.post_denoising_loop(latents, batch) + return batch + + +class SanaWMBeforeDenoisingStage(PipelineStage): + """ + Monolithic pre-processing stage for SANA-WM TI2V inference. + + Must run after SanaWMTextEncodingStage, which populates batch.prompt_embeds. + """ + + def __init__( + self, + vae, + transformer, + scheduler, + pipeline_config: SanaWMPipelineConfig, + ): + super().__init__() + self.vae = vae + self.transformer = transformer + self.scheduler = scheduler + self.pipeline_config = pipeline_config + + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + if self.vae is None: + return [] + stage_name = self._component_stage_name(stage_name) + pipeline_config = getattr(server_args, "pipeline_config", self.pipeline_config) + vae_dtype = PRECISION_TO_TYPE[pipeline_config.vae_precision] + return [ + ComponentUse( + stage_name=stage_name, + component_name="vae", + target_dtype=vae_dtype, + ) + ] + + @torch.no_grad() + def _vae_encode_image( + self, + image: torch.Tensor, # (1, C, H, W) or (1, C, 1, H, W) in [0, 1] float + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Encode a single image frame through the VAE encoder.""" + vae = self.vae + configure_sana_wm_ltx2_vae_for_long_video(vae, self.pipeline_config) + vae_dtype = PRECISION_TO_TYPE.get( + self.pipeline_config.vae_precision, torch.bfloat16 + ) + + # Normalize image to [-1, 1] range expected by the VAE + if image.max() > 1.01: + image = image / 255.0 + image = (image * 2.0 - 1.0).to(device=device, dtype=vae_dtype) + + # Add temporal dim if absent: (B, C, H, W) -> (B, C, 1, H, W) + if image.dim() == 4: + image = image.unsqueeze(2) + + log_sana_wm_tensor_stats("first_frame.pixel_input_normalized", image) + with self.use_declared_component( + component_name="vae", + module=vae, + target_dtype=vae_dtype, + ) as active_vae: + vae = active_vae if active_vae is not None else vae + z = self._extract_vae_latents(vae.encode(image)).float() + + z = sana_wm_normalize_vae_latents(vae, z, self.pipeline_config) + + log_sana_wm_tensor_stats("first_frame.latent_normalized", z) + return z.to(dtype=dtype) # (1, 128, 1, H_sp, W_sp) + + @staticmethod + def _extract_vae_latents(encoded: Any) -> torch.Tensor: + """Return deterministic VAE latents from common Diffusers outputs.""" + latent_dist = getattr(encoded, "latent_dist", None) + if latent_dist is not None: + if hasattr(latent_dist, "mode"): + return latent_dist.mode() + mean = getattr(latent_dist, "mean", None) + if isinstance(mean, torch.Tensor): + return mean + if callable(mean): + return mean() + if hasattr(latent_dist, "sample"): + return latent_dist.sample() + + if isinstance(encoded, tuple) and encoded: + return SanaWMBeforeDenoisingStage._extract_vae_latents(encoded[0]) + if isinstance(encoded, torch.Tensor): + return encoded + raise TypeError( + "Unsupported VAE encode output for SANA-WM first-frame conditioning: " + f"{type(encoded).__name__}" + ) + + def _prepare_noise_latents( + self, + shape: tuple, + dtype: torch.dtype, + device: torch.device, + generator: ( + torch.Generator | list[torch.Generator] | tuple[torch.Generator, ...] + ), + ) -> torch.Tensor: + if isinstance(generator, (list, tuple)): + if not generator: + raise ValueError("SANA-WM generator list must not be empty.") + if len(generator) == 1: + return randn_tensor( + shape, generator=generator[0], device=device, dtype=dtype + ) + if len(generator) != shape[0]: + raise ValueError( + "SANA-WM generator list length must match latent batch size; " + f"got {len(generator)} generators for batch {shape[0]}." + ) + sample_shape = (1, *shape[1:]) + return torch.cat( + [ + randn_tensor( + sample_shape, + generator=sample_generator, + device=device, + dtype=dtype, + ) + for sample_generator in generator + ], + dim=0, + ) + return randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + @staticmethod + def _generator_from_seed( + seed: int | list[int] | tuple[int, ...] | None, + *, + batch_size: int, + device: torch.device, + ) -> torch.Generator | list[torch.Generator]: + if seed is None: + seed = 0 + if isinstance(seed, (list, tuple)): + if not seed: + raise ValueError("SANA-WM seed list must not be empty.") + if len(seed) == 1: + seed = seed[0] + elif len(seed) == batch_size: + return [ + torch.Generator(device=device).manual_seed(int(sample_seed)) + for sample_seed in seed + ] + else: + raise ValueError( + "SANA-WM seed list length must be 1 or match latent batch " + f"size; got {len(seed)} seeds for batch {batch_size}." + ) + return torch.Generator(device=device).manual_seed(int(seed)) + + @staticmethod + def _canonical_condition_image_tensor(image: torch.Tensor) -> torch.Tensor: + """Return image as NCHW RGB float tensor without changing its value range.""" + image = image.float() + if image.dim() == 5 and image.shape[2] == 1: + image = image.squeeze(2) + if image.dim() == 3: + if image.shape[0] in (1, 3, 4): + image = image.unsqueeze(0) + elif image.shape[-1] in (1, 3, 4): + image = image.permute(2, 0, 1).unsqueeze(0) + else: + raise ValueError( + "condition_image tensor must be CHW or HWC with 1, 3, " + f"or 4 channels; got {tuple(image.shape)}." + ) + elif image.dim() == 4: + if image.shape[1] in (1, 3, 4): + pass + elif image.shape[-1] in (1, 3, 4): + image = image.permute(0, 3, 1, 2) + else: + raise ValueError( + "condition_image tensor must be NCHW or NHWC with 1, 3, " + f"or 4 channels; got {tuple(image.shape)}." + ) + else: + raise ValueError( + "condition_image tensor must have shape CHW, HWC, NCHW, NHWC, " + f"or NCHW singleton-video; got {tuple(image.shape)}." + ) + + if image.shape[1] == 1: + image = image.expand(-1, 3, -1, -1) + elif image.shape[1] == 4: + image = image[:, :3] + elif image.shape[1] != 3: + raise ValueError( + f"condition_image must have 1, 3, or 4 channels; got {image.shape[1]}." + ) + return image.contiguous() + + @staticmethod + def _resize_center_crop_tensor( + image: torch.Tensor, + *, + target_h: int, + target_w: int, + ) -> tuple[torch.Tensor, dict[str, tuple[int, int]]]: + """Match official SANA-WM resize-then-center-crop preprocessing.""" + image = SanaWMBeforeDenoisingStage._canonical_condition_image_tensor(image) + src_h, src_w = int(image.shape[-2]), int(image.shape[-1]) + resized_w, resized_h, left, top = sana_wm_compute_resize_crop_geometry( + src_w, src_h, target_h, target_w + ) + if resized_h != src_h or resized_w != src_w: + image = F.interpolate( + image, + size=(resized_h, resized_w), + mode="bilinear", + align_corners=False, + ) + image = image[..., top : top + target_h, left : left + target_w].contiguous() + return image, { + "source_size": (src_w, src_h), + "resized_size": (resized_w, resized_h), + "crop_offset": (left, top), + "target_size": (target_w, target_h), + } + + @staticmethod + def _preprocess_condition_image( + condition_image: Any, + *, + target_h: int, + target_w: int, + ) -> tuple[torch.Tensor, dict[str, tuple[int, int]]]: + """Aspect-preserving resize + center crop, mirroring NVlabs/Sana.""" + if isinstance(condition_image, list): + if len(condition_image) == 0: + raise ValueError( + "condition_image list is empty; SANA-WM requires a first " + "frame conditioning image." + ) + condition_image = condition_image[0] + + if isinstance(condition_image, Image.Image): + image = condition_image.convert("RGB") + src_w, src_h = image.size + resized_w, resized_h, left, top = sana_wm_compute_resize_crop_geometry( + src_w, src_h, target_h, target_w + ) + resampling_enum = getattr(Image, "Resampling", None) + resampling = ( + resampling_enum.LANCZOS + if resampling_enum is not None + else Image.LANCZOS + ) + image = image.resize((resized_w, resized_h), resampling) + image = image.crop((left, top, left + target_w, top + target_h)) + return TF.to_tensor(image).unsqueeze(0), { + "source_size": (src_w, src_h), + "resized_size": (resized_w, resized_h), + "crop_offset": (left, top), + "target_size": (target_w, target_h), + } + + if isinstance(condition_image, torch.Tensor): + return SanaWMBeforeDenoisingStage._resize_center_crop_tensor( + condition_image, + target_h=target_h, + target_w=target_w, + ) + + raise TypeError( + "condition_image must be a PIL image, tensor, or non-empty list; " + f"got {type(condition_image).__name__}." + ) + + @staticmethod + def _transform_intrinsics_for_condition_image( + intrinsics_vec4: torch.Tensor, + preprocess_info: ( + dict[str, tuple[int, int]] | list[dict[str, tuple[int, int]]] | None + ), + ) -> torch.Tensor: + """Map source-image intrinsics into the cropped output pixel grid.""" + if not preprocess_info: + return intrinsics_vec4 + if isinstance(preprocess_info, list): + transform = ( + SanaWMBeforeDenoisingStage._transform_intrinsics_for_condition_image + ) + if len(preprocess_info) == 1: + return transform(intrinsics_vec4, preprocess_info[0]) + if len(preprocess_info) != intrinsics_vec4.shape[0]: + raise ValueError( + "SANA-WM condition-image preprocess metadata length must " + "match intrinsics batch size; got " + f"{len(preprocess_info)} metadata entries for batch " + f"{intrinsics_vec4.shape[0]}." + ) + return torch.cat( + [ + transform(intrinsics_vec4[index : index + 1], info) + for index, info in enumerate(preprocess_info) + ], + dim=0, + ) + src_w, src_h = preprocess_info["source_size"] + resized_w, resized_h = preprocess_info["resized_size"] + left, top = preprocess_info["crop_offset"] + sx = resized_w / float(src_w) + sy = resized_h / float(src_h) + out = intrinsics_vec4.clone() + out[..., 0] *= sx + out[..., 2] = out[..., 2] * sx - left + out[..., 1] *= sy + out[..., 3] = out[..., 3] * sy - top + return out + + @torch.no_grad() + def _splice_first_frame( + self, + latents: torch.Tensor, # (B, 128, T_lat, H_sp, W_sp) + condition_image, # PIL Image or torch.Tensor + dtype: torch.dtype, + device: torch.device, + batch: Req | None = None, + ) -> torch.Tensor: + """Replace latents[:, :, 0] with VAE-encoded first frame.""" + B, _C, _T_lat, H_sp, W_sp = latents.shape + target_h = H_sp * self.pipeline_config.vae_stride[1] # 32 + target_w = W_sp * self.pipeline_config.vae_stride[2] # 32 + condition_images = self._condition_images_for_batch(condition_image, B) + first_frame_latents = [] + preprocess_infos = [] + for image in condition_images: + img_tensor, preprocess_info = self._preprocess_condition_image( + image, + target_h=target_h, + target_w=target_w, + ) + preprocess_infos.append(preprocess_info) + first_frame_latents.append( + self._vae_encode_image(img_tensor, dtype, device) + ) + + if batch is not None: + if not hasattr(batch, "extra") or batch.extra is None: + batch.extra = {} + batch.extra[_SANA_WM_CONDITION_IMAGE_PREPROCESS_KEY] = ( + preprocess_infos[0] if len(preprocess_infos) == 1 else preprocess_infos + ) + self.log_info( + "First-frame condition image preprocessed: source=%s, resized=%s, " + "crop_offset=%s, target=%s.", + preprocess_infos[0]["source_size"], + preprocess_infos[0]["resized_size"], + preprocess_infos[0]["crop_offset"], + preprocess_infos[0]["target_size"], + ) + if len(preprocess_infos) > 1: + self.log_info( + "Processed %d batched first-frame images.", len(preprocess_infos) + ) + + first_frame_z = torch.cat(first_frame_latents, dim=0) + if first_frame_z.shape[0] == 1 and B > 1: + first_frame_z = first_frame_z.expand(B, -1, -1, -1, -1) + elif first_frame_z.shape[0] != B: + raise ValueError( + "SANA-WM first-frame latent batch does not match noise batch: " + f"{first_frame_z.shape[0]} vs {B}." + ) + + latents = latents.clone() + latents[:, :, 0:1] = first_frame_z + log_sana_wm_tensor_stats("latents.after_first_frame_splice", latents) + return latents + + @staticmethod + def _condition_images_for_batch(condition_image: Any, batch_size: int) -> list[Any]: + if isinstance(condition_image, list): + if not condition_image: + raise ValueError( + "condition_image list is empty; SANA-WM requires a first " + "frame conditioning image." + ) + if len(condition_image) == 1 or len(condition_image) == batch_size: + return list(condition_image) + raise ValueError( + "SANA-WM condition_image list must contain one image or one " + f"image per batch item; got {len(condition_image)} images for " + f"batch {batch_size}." + ) + return [condition_image] + + @staticmethod + def _pad_or_trim_frames(tensor: torch.Tensor, num_frames: int) -> torch.Tensor: + current = tensor.shape[1] + if current == num_frames: + return tensor + if current > num_frames: + return tensor[:, :num_frames] + if current == 0: + raise ValueError("camera trajectory must contain at least one frame") + pad = num_frames - current + last = tensor[:, -1:].repeat(1, pad, *([1] * (tensor.ndim - 2))) + return torch.cat([tensor, last], dim=1) + + @staticmethod + def _maybe_load_npy_tensor(value: Any, field_name: str) -> Any: + if isinstance(value, (str, os.PathLike)): + path = os.fspath(value) + if not path.endswith(".npy"): + raise ValueError( + f"{field_name} path must point to a .npy file, got {path!r}" + ) + return torch.from_numpy(np.load(path)) + return value + + @staticmethod + def _first_mapping_value(mapping: dict[str, Any], *keys: str) -> Any: + for key in keys: + if key in mapping and mapping[key] is not None: + return mapping[key] + return None + + @staticmethod + def _first_request_mapping_value( + extra: dict[str, Any], + diffusers_kwargs: dict[str, Any], + *keys: str, + ) -> Any: + for mapping in (extra, diffusers_kwargs): + for key in keys: + if key in mapping and mapping[key] is not None: + return mapping[key] + return None + + @staticmethod + def _request_float_value( + extra: dict[str, Any], + diffusers_kwargs: dict[str, Any], + *keys: str, + default: float, + ) -> float: + value = SanaWMBeforeDenoisingStage._first_request_mapping_value( + extra, diffusers_kwargs, *keys + ) + return default if value is None else float(value) + + @staticmethod + def _request_action_value( + extra: dict[str, Any], + diffusers_kwargs: dict[str, Any], + ) -> Any: + return SanaWMBeforeDenoisingStage._first_request_mapping_value( + extra, diffusers_kwargs, "action", "sana_wm_action" + ) + + @staticmethod + def _pad_or_trim_action_trajectory( + trajectory: torch.Tensor, + num_frames: int, + ) -> torch.Tensor: + current = trajectory.shape[0] + if current == num_frames: + return trajectory + if current > num_frames: + return trajectory[:num_frames] + pad = num_frames - current + return torch.cat([trajectory, trajectory[-1:].repeat(pad, 1, 1)], dim=0) + + @staticmethod + def _coerce_action_camera_to_world( + value: Any, + *, + batch_size: int, + num_frames: int, + translation_speed: float, + rotation_speed_deg: float, + pitch_limit_deg: float, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + if isinstance(value, str): + actions = [value] + elif isinstance(value, (list, tuple)) and all( + isinstance(x, str) for x in value + ): + actions = list(value) + else: + raise ValueError( + "SANA-WM action must be a string or a list of strings, " + f"got {type(value).__name__}." + ) + + trajectories = [ + SanaWMBeforeDenoisingStage._pad_or_trim_action_trajectory( + sana_wm_action_to_camera_to_world( + action, + translation_speed=translation_speed, + rotation_speed_deg=rotation_speed_deg, + pitch_limit_deg=pitch_limit_deg, + ), + num_frames, + ) + for action in actions + ] + camera = torch.stack(trajectories, dim=0).to(device=device, dtype=dtype) + if camera.shape[0] == 1 and batch_size > 1: + camera = camera.expand(batch_size, -1, -1, -1) + elif camera.shape[0] != batch_size: + raise ValueError( + f"SANA-WM action batch {camera.shape[0]} does not match {batch_size}" + ) + return camera + + @staticmethod + def _action_num_frames_for_request(batch: Req) -> int | None: + extra = getattr(batch, "extra", None) or {} + diffusers_kwargs = extra.get("diffusers_kwargs", {}) + if not isinstance(diffusers_kwargs, dict): + diffusers_kwargs = {} + action = SanaWMBeforeDenoisingStage._request_action_value( + extra, diffusers_kwargs + ) + if action is None: + return None + if isinstance(action, str): + return len(parse_sana_wm_action_string(action)) + 1 + if isinstance(action, (list, tuple)) and all( + isinstance(x, str) for x in action + ): + return max(len(parse_sana_wm_action_string(x)) + 1 for x in action) + raise ValueError( + "SANA-WM action must be a string or a list of strings, " + f"got {type(action).__name__}." + ) + + @staticmethod + def _coerce_camera_to_world( + value: Any, + *, + batch_size: int, + num_frames: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + value = SanaWMBeforeDenoisingStage._maybe_load_npy_tensor( + value, "camera_to_world" + ) + camera = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + if camera.dim() == 3: + camera = camera.unsqueeze(0) + if camera.dim() != 4 or camera.shape[-2:] != (4, 4): + raise ValueError( + "camera_to_world must have shape (F,4,4) or (B,F,4,4), " + f"got {tuple(camera.shape)}" + ) + camera = camera.to(device=device, dtype=dtype) + if camera.shape[0] == 1 and batch_size > 1: + camera = camera.expand(batch_size, -1, -1, -1) + elif camera.shape[0] != batch_size: + raise ValueError( + f"camera_to_world batch {camera.shape[0]} does not match {batch_size}" + ) + return SanaWMBeforeDenoisingStage._pad_or_trim_frames(camera, num_frames) + + @staticmethod + def _intrinsics_matrix_to_vec4(intrinsics: torch.Tensor) -> torch.Tensor: + return torch.stack( + [ + intrinsics[..., 0, 0], + intrinsics[..., 1, 1], + intrinsics[..., 0, 2], + intrinsics[..., 1, 2], + ], + dim=-1, + ) + + @staticmethod + def _coerce_intrinsics_vec4( + value: Any, + *, + batch_size: int, + num_frames: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + value = SanaWMBeforeDenoisingStage._maybe_load_npy_tensor(value, "intrinsics") + intrinsics = ( + value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + ) + intrinsics = intrinsics.to(device=device, dtype=dtype) + + if intrinsics.dim() == 1 and intrinsics.shape[0] == 4: + intrinsics = intrinsics.view(1, 1, 4) + elif intrinsics.dim() == 2 and intrinsics.shape == (3, 3): + intrinsics = SanaWMBeforeDenoisingStage._intrinsics_matrix_to_vec4( + intrinsics + ).view(1, 1, 4) + elif intrinsics.dim() == 2 and intrinsics.shape[-1] == 4: + intrinsics = intrinsics.unsqueeze(0) + elif intrinsics.dim() == 3 and intrinsics.shape[-2:] == (3, 3): + vec4 = SanaWMBeforeDenoisingStage._intrinsics_matrix_to_vec4(intrinsics) + if intrinsics.shape[0] >= num_frames: + intrinsics = vec4.unsqueeze(0) + elif intrinsics.shape[0] == batch_size: + intrinsics = vec4.unsqueeze(1) + else: + raise ValueError( + "intrinsics with shape (N,3,3) must use N>=num_frames " + f"or N=batch_size, got N={intrinsics.shape[0]}, " + f"num_frames={num_frames}, batch_size={batch_size}" + ) + elif intrinsics.dim() == 3 and intrinsics.shape[-1] == 4: + pass + elif intrinsics.dim() == 4 and intrinsics.shape[-2:] == (3, 3): + intrinsics = SanaWMBeforeDenoisingStage._intrinsics_matrix_to_vec4( + intrinsics + ) + else: + raise ValueError( + "intrinsics must have shape (4,), (F,4), (B,F,4), " + "(3,3), (F,3,3), or (B,F,3,3); " + f"got {tuple(intrinsics.shape)}" + ) + + if intrinsics.shape[0] == 1 and batch_size > 1: + intrinsics = intrinsics.expand(batch_size, -1, -1) + elif intrinsics.shape[0] != batch_size: + raise ValueError( + f"intrinsics batch {intrinsics.shape[0]} does not match {batch_size}" + ) + if intrinsics.shape[1] == 1 and num_frames > 1: + intrinsics = intrinsics.expand(-1, num_frames, -1) + return SanaWMBeforeDenoisingStage._pad_or_trim_frames(intrinsics, num_frames) + + @staticmethod + def _relative_camera_poses(camera_to_world: torch.Tensor) -> torch.Tensor: + input_dtype = camera_to_world.dtype + camera_to_world = camera_to_world.float() + first_inv = torch.linalg.inv(camera_to_world[:, :1]) + poses = torch.matmul(first_inv, camera_to_world) + eye = torch.eye( + 4, + device=camera_to_world.device, + dtype=camera_to_world.dtype, + ) + poses[:, 0] = eye + return poses.to(dtype=input_dtype) + + @staticmethod + def _scale_intrinsics_to_latent( + intrinsics_vec4: torch.Tensor, + *, + pixel_h: int, + pixel_w: int, + latent_h: int, + latent_w: int, + ) -> torch.Tensor: + intrinsics_latent = intrinsics_vec4.clone() + intrinsics_latent[..., [0, 2]] *= latent_w / float(pixel_w) + intrinsics_latent[..., [1, 3]] *= latent_h / float(pixel_h) + return intrinsics_latent + + @staticmethod + def _flatten_camera_conditions( + camera_to_world: torch.Tensor, + intrinsics_vec4: torch.Tensor, + ) -> torch.Tensor: + c2w_flat = camera_to_world.reshape( + camera_to_world.shape[0], + camera_to_world.shape[1], + 16, + ) + return torch.cat( + [c2w_flat, intrinsics_vec4], + dim=-1, + ) + + @staticmethod + def _latent_frame_camera_conditions( + camera_conditions: torch.Tensor, + *, + num_frames: int, + latent_frames: int, + vae_temporal_stride: int, + ) -> torch.Tensor: + time_indices = torch.arange( + 0, + num_frames, + vae_temporal_stride, + device=camera_conditions.device, + dtype=torch.long, + ) + if time_indices.numel() < latent_frames: + pad = latent_frames - int(time_indices.numel()) + time_indices = torch.cat( + [time_indices, time_indices[-1:].repeat(pad)], dim=0 + ) + time_indices = time_indices[:latent_frames] + return camera_conditions.index_select(1, time_indices) + + def _default_static_camera( + self, + *, + batch_size: int, + num_frames: int, + pixel_h: int, + pixel_w: int, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + camera_to_world = torch.eye(4, device=device, dtype=dtype).view(1, 1, 4, 4) + camera_to_world = camera_to_world.repeat(batch_size, num_frames, 1, 1) + focal = 0.8 * float(max(pixel_h, pixel_w)) + intrinsics = torch.tensor( + [focal, focal, pixel_w / 2.0, pixel_h / 2.0], + device=device, + dtype=dtype, + ).view(1, 1, 4) + intrinsics = intrinsics.repeat(batch_size, num_frames, 1) + return camera_to_world, intrinsics + + @staticmethod + def _has_explicit_camera_request(batch: Req) -> bool: + extra = getattr(batch, "extra", None) or {} + if any( + key in extra and extra[key] is not None + for key in ( + "camera_conditions", + "chunk_plucker", + "camera_to_world", + "intrinsics", + "action", + "sana_wm_action", + ) + ): + return True + diffusers_kwargs = extra.get("diffusers_kwargs", {}) + if not isinstance(diffusers_kwargs, dict): + return False + return any( + key in diffusers_kwargs and diffusers_kwargs[key] is not None + for key in ( + "camera_conditions", + "chunk_plucker", + "camera_to_world", + "camera_to_world_path", + "camera_path", + "intrinsics", + "intrinsics_path", + "action", + "sana_wm_action", + ) + ) + + def _build_camera_conditioning( + self, + batch: Req, + *, + batch_size: int, + num_frames: int, + latent_shape: tuple, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, str]: + if not hasattr(batch, "extra") or batch.extra is None: + batch.extra = {} + if not self.pipeline_config.camera_conditioning: + return None, None, "disabled" + + extra = batch.extra + T_lat = latent_shape[2] + sp_h = latent_shape[3] + sp_w = latent_shape[4] + vae_temporal_stride = self.pipeline_config.vae_stride[0] + camera_compute_dtype = torch.float32 + + diffusers_kwargs = extra.get("diffusers_kwargs", {}) + if not isinstance(diffusers_kwargs, dict): + diffusers_kwargs = {} + camera_conditions = self._first_request_mapping_value( + extra, diffusers_kwargs, "camera_conditions" + ) + chunk_plucker = self._first_request_mapping_value( + extra, diffusers_kwargs, "chunk_plucker" + ) + preprocess_info = extra.get(_SANA_WM_CONDITION_IMAGE_PREPROCESS_KEY) + action = self._request_action_value(extra, diffusers_kwargs) + arch = getattr( + getattr(self.pipeline_config, "dit_config", None), + "arch_config", + None, + ) + requires_chunk_plucker = bool( + getattr(arch, "use_chunk_plucker_post_attn", False) + or getattr(arch, "use_chunk_plucker_input", False) + ) + if action is not None and ( + camera_conditions is not None or chunk_plucker is not None + ): + raise ValueError( + "SANA-WM action cannot be combined with prepacked " + "camera_conditions/chunk_plucker." + ) + if camera_conditions is not None: + camera_conditions = ( + camera_conditions + if isinstance(camera_conditions, torch.Tensor) + else torch.as_tensor(camera_conditions) + ).to(device=device, dtype=camera_compute_dtype) + if camera_conditions.dim() == 2: + camera_conditions = camera_conditions.unsqueeze(0) + if camera_conditions.dim() != 3: + raise ValueError( + "camera_conditions must have shape (T,20) or (B,T,20), " + f"got {tuple(camera_conditions.shape)}" + ) + if camera_conditions.shape[0] == 1 and batch_size > 1: + camera_conditions = camera_conditions.expand(batch_size, -1, -1) + if camera_conditions.shape[0] != batch_size: + raise ValueError( + "camera_conditions batch dimension must be 1 or match " + f"request batch size {batch_size}, got " + f"{camera_conditions.shape[0]}." + ) + if camera_conditions.shape[-1] != 20: + raise ValueError( + "camera_conditions must have last dimension 20, got " + f"{tuple(camera_conditions.shape)}" + ) + if camera_conditions.shape[1] == T_lat: + source = "prepacked" + if chunk_plucker is None and requires_chunk_plucker: + raise ValueError( + "Prepacked latent-frame camera_conditions require " + "chunk_plucker for this SANA-WM checkpoint. Pass " + "chunk_plucker with shape (B,48,T,H,W), or pass " + "original-frame camera_conditions so SGLang can " + "derive chunk_plucker." + ) + else: + source = "prebuilt_original_frames" + original_camera_conditions = self._pad_or_trim_frames( + camera_conditions, num_frames + ) + camera_conditions = self._latent_frame_camera_conditions( + original_camera_conditions, + num_frames=num_frames, + latent_frames=T_lat, + vae_temporal_stride=vae_temporal_stride, + ) + if chunk_plucker is None: + chunk_plucker = compute_chunk_plucker( + camera_conditions=original_camera_conditions, + HW=(T_lat, sp_h, sp_w), + vae_temporal_stride=vae_temporal_stride, + patch_size=(1, 1, 1), + ) + else: + camera_to_world = extra.get("camera_to_world", None) + intrinsics = extra.get("intrinsics", None) + if camera_to_world is None: + camera_to_world = self._first_mapping_value( + diffusers_kwargs, + "camera_to_world", + "camera_to_world_path", + "camera_path", + ) + if intrinsics is None: + intrinsics = self._first_mapping_value( + diffusers_kwargs, + "intrinsics", + "intrinsics_path", + ) + if action is not None and camera_to_world is not None: + raise ValueError( + "SANA-WM action and camera_to_world/camera_path are " + "mutually exclusive." + ) + + if action is not None: + source = ( + "action" if intrinsics is not None else "action_default_intrinsics" + ) + translation_speed = self._request_float_value( + extra, + diffusers_kwargs, + "translation_speed", + "action_translation_speed", + "sana_wm_translation_speed", + default=_SANA_WM_DEFAULT_TRANSLATION_SPEED, + ) + rotation_speed_deg = self._request_float_value( + extra, + diffusers_kwargs, + "rotation_speed_deg", + "action_rotation_speed_deg", + "sana_wm_rotation_speed_deg", + default=_SANA_WM_DEFAULT_ROTATION_SPEED_DEG, + ) + pitch_limit_deg = self._request_float_value( + extra, + diffusers_kwargs, + "pitch_limit_deg", + "action_pitch_limit_deg", + "sana_wm_pitch_limit_deg", + default=_SANA_WM_DEFAULT_PITCH_LIMIT_DEG, + ) + camera_to_world = self._coerce_action_camera_to_world( + action, + batch_size=batch_size, + num_frames=num_frames, + translation_speed=translation_speed, + rotation_speed_deg=rotation_speed_deg, + pitch_limit_deg=pitch_limit_deg, + device=device, + dtype=camera_compute_dtype, + ) + self.log_info( + "SANA-WM action trajectory rolled out: frames=%d, " + "translation_speed=%.6g, rotation_speed_deg=%.6g, " + "pitch_limit_deg=%.6g", + camera_to_world.shape[1], + translation_speed, + rotation_speed_deg, + pitch_limit_deg, + ) + if intrinsics is None: + _, intrinsics_vec4 = self._default_static_camera( + batch_size=batch_size, + num_frames=num_frames, + pixel_h=batch.height, + pixel_w=batch.width, + device=device, + dtype=camera_compute_dtype, + ) + self.log_info( + "No intrinsics provided; using heuristic centered " + "intrinsics for the action trajectory." + ) + else: + intrinsics_vec4 = self._coerce_intrinsics_vec4( + intrinsics, + batch_size=batch_size, + num_frames=num_frames, + device=device, + dtype=camera_compute_dtype, + ) + intrinsics_vec4 = self._transform_intrinsics_for_condition_image( + intrinsics_vec4, + preprocess_info, + ) + elif camera_to_world is not None: + source = ( + "request" + if intrinsics is not None + else "request_default_intrinsics" + ) + camera_to_world = self._coerce_camera_to_world( + camera_to_world, + batch_size=batch_size, + num_frames=num_frames, + device=device, + dtype=camera_compute_dtype, + ) + if intrinsics is None: + _, intrinsics_vec4 = self._default_static_camera( + batch_size=batch_size, + num_frames=num_frames, + pixel_h=batch.height, + pixel_w=batch.width, + device=device, + dtype=camera_compute_dtype, + ) + self.log_info( + "No intrinsics provided; using heuristic centered " + "intrinsics for the request camera trajectory." + ) + else: + intrinsics_vec4 = self._coerce_intrinsics_vec4( + intrinsics, + batch_size=batch_size, + num_frames=num_frames, + device=device, + dtype=camera_compute_dtype, + ) + intrinsics_vec4 = self._transform_intrinsics_for_condition_image( + intrinsics_vec4, + preprocess_info, + ) + elif intrinsics is not None: + source = "default_static_request_intrinsics" + camera_to_world, _ = self._default_static_camera( + batch_size=batch_size, + num_frames=num_frames, + pixel_h=batch.height, + pixel_w=batch.width, + device=device, + dtype=camera_compute_dtype, + ) + intrinsics_vec4 = self._coerce_intrinsics_vec4( + intrinsics, + batch_size=batch_size, + num_frames=num_frames, + device=device, + dtype=camera_compute_dtype, + ) + intrinsics_vec4 = self._transform_intrinsics_for_condition_image( + intrinsics_vec4, + preprocess_info, + ) + self.log_info( + "No camera trajectory provided; using static identity " + "poses with request intrinsics." + ) + else: + source = "default_static" + self.log_info( + "No camera trajectory provided; using a static identity " + "camera with heuristic centered intrinsics. Pass " + "camera_to_world/intrinsics for camera-controlled output." + ) + camera_to_world, intrinsics_vec4 = self._default_static_camera( + batch_size=batch_size, + num_frames=num_frames, + pixel_h=batch.height, + pixel_w=batch.width, + device=device, + dtype=camera_compute_dtype, + ) + + camera_to_world = self._relative_camera_poses(camera_to_world) + intrinsics_vec4 = self._scale_intrinsics_to_latent( + intrinsics_vec4, + pixel_h=batch.height, + pixel_w=batch.width, + latent_h=sp_h, + latent_w=sp_w, + ) + original_camera_conditions = self._flatten_camera_conditions( + camera_to_world, intrinsics_vec4 + ) + camera_conditions = self._latent_frame_camera_conditions( + original_camera_conditions, + num_frames=num_frames, + latent_frames=T_lat, + vae_temporal_stride=vae_temporal_stride, + ) + if chunk_plucker is None: + chunk_plucker = compute_chunk_plucker( + camera_conditions=original_camera_conditions, + HW=(T_lat, sp_h, sp_w), + vae_temporal_stride=vae_temporal_stride, + patch_size=(1, 1, 1), + ) + + if chunk_plucker is not None: + chunk_plucker = ( + chunk_plucker + if isinstance(chunk_plucker, torch.Tensor) + else torch.as_tensor(chunk_plucker) + ).to(device=device, dtype=dtype) + if chunk_plucker.dim() == 4: + chunk_plucker = chunk_plucker.unsqueeze(0) + if chunk_plucker.shape[0] == 1 and batch_size > 1: + chunk_plucker = chunk_plucker.expand(batch_size, -1, -1, -1, -1) + if chunk_plucker.dim() != 5: + raise ValueError( + "chunk_plucker must have shape (48,T,H,W) or " + f"(B,48,T,H,W), got {tuple(chunk_plucker.shape)}" + ) + if chunk_plucker.shape[0] != batch_size: + raise ValueError( + "chunk_plucker batch dimension must be 1 or match " + f"request batch size {batch_size}, got " + f"{chunk_plucker.shape[0]}." + ) + expected_chunk_shape = (batch_size, 48, T_lat, sp_h, sp_w) + if tuple(chunk_plucker.shape) != expected_chunk_shape: + raise ValueError( + "chunk_plucker shape mismatch for SANA-WM: expected " + f"{expected_chunk_shape}, got {tuple(chunk_plucker.shape)}." + ) + + if camera_conditions is not None: + camera_conditions = camera_conditions.to(device=device, dtype=dtype) + + return camera_conditions, chunk_plucker, source + + def _prepare_timesteps( + self, + batch: Req, + server_args: ServerArgs, + device: torch.device, + ): + """Set up scheduler timesteps and populate batch.timesteps, .sigmas.""" + scheduler = get_or_create_request_scheduler(batch, self.scheduler) + num_inference_steps = batch.num_inference_steps + + flow_shift = getattr( + self.pipeline_config, + "inference_flow_shift", + None, + ) + if flow_shift is None: + flow_shift = getattr(self.pipeline_config, "flow_shift", 9.95) + kwargs = {} + + # diffusers FlowMatchEulerDiscreteScheduler supports mu/shift + sig_params = inspect.signature(scheduler.set_timesteps).parameters + if "shift" in sig_params: + kwargs["shift"] = flow_shift + elif "mu" in sig_params: + # Convert flow_shift to mu: mu ~= log(shift) + kwargs["mu"] = math.log(flow_shift) + + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + sigmas = scheduler.sigmas.tolist() + if sigmas: + self.log_info( + "FlowMatch timesteps prepared: steps=%d, flow_shift=%.4f, " + "sigma_start=%.6f, sigma_end=%.6f", + num_inference_steps, + flow_shift, + float(sigmas[0]), + float(sigmas[-1]), + ) + + batch.timesteps = timesteps + batch.sigmas = sigmas + batch.scheduler = scheduler + return batch + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + """Pre-process everything needed by DenoisingStage for SANA-WM. + + Expects batch to already have prompt_embeds set by SanaWMTextEncodingStage. + """ + device = get_local_torch_device() + dtype = PRECISION_TO_TYPE.get( + getattr(self.pipeline_config, "dit_precision", "bf16"), + torch.bfloat16, + ) + if not hasattr(batch, "extra") or batch.extra is None: + batch.extra = {} + + # Adjust num_frames to be compatible with VAE temporal stride. + requested_num_frames = batch.num_frames or 49 + action_num_frames = self._action_num_frames_for_request(batch) + if action_num_frames is not None and action_num_frames < requested_num_frames: + self.log_info( + "SANA-WM action trajectory has %d frames; capping requested " + "num_frames=%d before VAE stride adjustment.", + action_num_frames, + requested_num_frames, + ) + requested_num_frames = action_num_frames + num_frames = requested_num_frames + num_frames = self.pipeline_config.adjust_num_frames(num_frames) + batch.num_frames = num_frames + self.log_info( + "SANA-WM prepare: seed=%s, size=%dx%d, frames=%d, " + "vae_stride=%s, diagnostics=%s", + getattr(batch, "seed", None), + batch.width, + batch.height, + num_frames, + self.pipeline_config.vae_stride, + "on" if sana_wm_diagnostics_enabled() else "off", + ) + + batch_size = batch.batch_size or 1 + generator = getattr(batch, "generator", None) + if not isinstance(generator, (list, tuple, torch.Generator)): + generator = self._generator_from_seed( + getattr(batch, "seed", None), + batch_size=batch_size, + device=device, + ) + batch.generator = generator + + latent_shape = self.pipeline_config.prepare_latent_shape( + batch, batch_size, num_frames + ) + # latent_shape: (B, 128, T_latent, H_sp, W_sp) + latents = self._prepare_noise_latents(latent_shape, dtype, device, generator) + log_sana_wm_tensor_stats("latents.initial_noise", latents) + + batch.raw_latent_shape = latent_shape + + condition_image = getattr(batch, "condition_image", None) + if condition_image is not None: + try: + latents = self._splice_first_frame( + latents, condition_image, dtype, device, batch=batch + ) + self.log_info("First-frame spliced into noise latents.") + except Exception as e: + raise RuntimeError( + "SANA-WM first-frame conditioning failed; refusing to " + "continue with pure-noise latents because that produces " + "misleading low-quality output." + ) from e + else: + raise ValueError( + "SANA-WM is a TI2V world model and requires condition_image " + "for first-frame conditioning. Provide --image-path, " + "--condition-image, or the equivalent API image input." + ) + + batch.latents = latents + + # The released SANA-WM checkpoint is camera-conditioned. Official + # inference requires a camera trajectory or action DSL. If the SGLang + # request omits one, use a static identity trajectory so the UCPE path + # remains active instead of silently dropping all camera conditioning. + try: + camera_conditions, chunk_plucker, camera_source = ( + self._build_camera_conditioning( + batch, + batch_size=batch_size, + num_frames=num_frames, + latent_shape=latent_shape, + device=device, + dtype=dtype, + ) + ) + except Exception as e: + if self._has_explicit_camera_request(batch): + raise RuntimeError( + "SANA-WM camera conditioning failed for an explicitly " + "provided camera/intrinsics request." + ) from e + logger.warning( + "SANA-WM camera conditioning failed: %s. Disabling camera branch.", + e, + ) + camera_conditions, chunk_plucker, camera_source = None, None, "error" + + if camera_conditions is not None: + batch.extra["camera_conditions"] = camera_conditions + log_sana_wm_tensor_stats("camera_conditions", camera_conditions) + if chunk_plucker is not None: + batch.extra["chunk_plucker"] = chunk_plucker + log_sana_wm_tensor_stats("chunk_plucker", chunk_plucker) + self.log_info( + "SANA-WM camera conditioning: source=%s, raymap=%s, chunk_plucker=%s", + camera_source, + None if camera_conditions is None else tuple(camera_conditions.shape), + None if chunk_plucker is None else tuple(chunk_plucker.shape), + ) + + batch = self._prepare_timesteps(batch, server_args, device) + + # Ensure prompt_embeds is a list (DenoisingStage expects list[Tensor]). + if isinstance(batch.prompt_embeds, torch.Tensor): + batch.prompt_embeds = [batch.prompt_embeds] + if batch.negative_prompt_embeds is not None and isinstance( + batch.negative_prompt_embeds, torch.Tensor + ): + batch.negative_prompt_embeds = [batch.negative_prompt_embeds] + + batch.do_classifier_free_guidance = getattr(batch, "guidance_scale", 1.0) > 1.0 + + self.log_info( + "BeforeDenoisingStage done: latent=%s, T_lat=%d, H_sp=%d, W_sp=%d, " + "num_inference_steps=%d, camera=%s", + str(latent_shape), + latent_shape[2], + latent_shape[3], + latent_shape[4], + batch.num_inference_steps, + "yes" if batch.extra.get("camera_conditions") is not None else "no", + ) + return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/parity_probe.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/parity_probe.py new file mode 100644 index 000000000..b1f25f336 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/parity_probe.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Env-gated parity dump/probe helpers (no-op in production). + +Single home for the debug harness that localized the realtime<->batch parity +bugs (previously copy-pasted across streaming/realtime/refiner/sana_wm). + +Env vars: +- ``SANAWM_FORK_DUMP_DIR`` — batch-path dumps (init_noise, conditioning, + per-chunk stage-1, kv/refiner probes, weights fingerprint) +- ``SANAWM_RT_DUMP_DIR`` — realtime-path dumps (same set) +- ``SANAWM_BLOCK_PROBE`` — file path: per-block forward_long checksums on + the first sink-path call +- ``SANAWM_INJECT_DIR`` — batch-path injection inputs (read, not written) + +All checksums reduce in float64 (bf16/fp32 upcasts are lossless, so the +formula is bitwise-stable across both paths). +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import torch + +ENV_FORK_DUMP = "SANAWM_FORK_DUMP_DIR" +ENV_RT_DUMP = "SANAWM_RT_DUMP_DIR" +ENV_BLOCK_PROBE = "SANAWM_BLOCK_PROBE" +ENV_INJECT = "SANAWM_INJECT_DIR" + + +def probe_dir(*env_names: str) -> Path | None: + """First configured dump dir among ``env_names`` (created on demand).""" + for name in env_names: + value = os.environ.get(name) + if value: + path = Path(value) + path.mkdir(parents=True, exist_ok=True) + return path + return None + + +def dump_tensor(dirpath: Path | str | None, name: str, tensor) -> None: + """Save ``tensor`` as float32 CPU under ``dirpath/name.pt`` (None-safe).""" + if dirpath is None or tensor is None: + return + torch.save(tensor.detach().float().cpu(), Path(dirpath) / f"{name}.pt") + + +def dump_obj(dirpath: Path | str | None, name: str, obj) -> None: + """Save a picklable object (e.g. a checksum dict) under ``dirpath/name.pt``.""" + if dirpath is None or obj is None: + return + torch.save(obj, Path(dirpath) / f"{name}.pt") + + +def checksum(tensor) -> tuple[tuple[int, ...], float] | None: + """(shape, float64 sum) of a tensor; None passes through.""" + if tensor is None: + return None + return ( + tuple(tensor.shape), + float(tensor.detach().float().double().sum().item()), + ) + + +def kv_cache_checksums(chunk_kv, sink_num: int) -> dict: + """Per-block, per-slot checksums of an accumulated KV cache.""" + probe: dict = {"sink_num": sink_num} + for block_id, slots in enumerate(chunk_kv): + for slot_id, tensor in enumerate(slots): + if tensor is not None: + probe[f"b{block_id:02d}s{slot_id}"] = checksum(tensor) + return probe + + +def weights_fingerprint(module: torch.nn.Module) -> dict[str, float]: + """Per-parameter |.|-sum fingerprint (catches wrong/mutated weights).""" + return { + name: float(param.detach().float().abs().sum().item()) + for name, param in module.named_parameters() + } diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_chain.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_chain.py new file mode 100644 index 000000000..daed8f247 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_chain.py @@ -0,0 +1,556 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SANA-WM realtime stage chain. + +Per-tick pipeline: + + RealtimeInputValidationStage (framework) + SanaWMRealtimeTextEncodingStage (framework cache via MRO) + SanaWMCondFrameEncodeStage -> batch.image_latent (+ inputs snapshot) + SanaWMRealtimeLatentPrepStage -> batch.latents = this tick's pre-noised + chunk(s); batch.extra["sana_wm_chunk_plan"] + SanaWMCameraCondStage -> batch.extra camera_conditions/chunk_plucker + SanaWMStreamingDenoisingStage (session path; SanaWMStreamCacheState) + SanaWMChunkedRefinerChainStage -> batch.latents = refined buffer + SanaWMCausalDecodeChainStage -> OutputBatch (decodes past its frontier) + +Chain stages share ``SanaWMRealtimeStage`` for first-frame image handling and +causal VAE decode helpers; model-specific chunk planning stays in this module. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import torch + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.models.dits.sana_wm_components import ( + compute_chunk_plucker, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req +from sglang.multimodal_gen.runtime.realtime.causal_state import ( + RealtimeCausalDecodeState, +) +from sglang.multimodal_gen.runtime.realtime.session import BaseRealtimeState +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +from . import parity_probe +from .base import ( + _SANA_WM_DEFAULT_ROTATION_SPEED_DEG, + _SANA_WM_DEFAULT_TRANSLATION_SPEED, + SanaWMBeforeDenoisingStage, + configure_sana_wm_ltx2_vae_for_long_video, + snap_sana_wm_num_frames, +) +from .realtime_stage import ( + DEFAULT_REFINER_BLOCK_SIZE, + DEFAULT_REFINER_KV_MAX_FRAMES, + SANA_WM_HEIGHT, + SANA_WM_WIDTH, + SanaWMRealtimeStage, + _motion_param, +) +from .streaming import SanaWMStreamCacheState, SanaWMStreamingDenoisingStage +from .streaming_refiner import RefinerChunkRunner + + +# --------------------------------------------------------------------- # +# Per-session state blobs (one per owning stage) +# --------------------------------------------------------------------- # +class SanaWMSessionInputsState(BaseRealtimeState): + """Session input snapshot — written by the cond-encode stage, read by the + camera stage (crop geometry) and the refiner stage (prompt).""" + + def __init__(self): + super().__init__() + self.image = None + self.intrinsics_image = None + self.src_size = None + self.resized_size = None + self.crop_offset = None + self.target_height: int | None = None + self.target_width: int | None = None + self.prompt: str = "" + self.open_ended = False + self.latent_t: int | None = None # None => open-ended + self.num_frame_per_block = DEFAULT_REFINER_BLOCK_SIZE + self.sink_size = 1 + self.camera_actions: list[list[str]] = [] + self.max_camera_actions = 0 # unlimited (sessions stream past any horizon) + self.static_c2w: np.ndarray | None = None + self.intrinsics_raw: np.ndarray | None = None + self.translation_speed = _SANA_WM_DEFAULT_TRANSLATION_SPEED + self.rotation_speed_deg = _SANA_WM_DEFAULT_ROTATION_SPEED_DEG + self.first_latent: torch.Tensor | None = None + + def dispose(self): + super().dispose() + self.__init__() + + +class SanaWMNoiseState(BaseRealtimeState): + """Seeded noise discipline — full-horizon buffer for fixed N (sliced per + chunk, matching the offline draw bitwise), seeded fallback otherwise.""" + + def __init__(self): + super().__init__() + self.noise_buffer: torch.Tensor | None = None + self.generator: torch.Generator | None = None + self.segments: list[int] | None = None # front-loaded grid (fixed N) + + def dispose(self): + super().dispose() + self.noise_buffer = None + self.generator = None + self.segments = None + + +class SanaWMRefinerChainState(BaseRealtimeState): + def __init__(self): + super().__init__() + self.runner: RefinerChunkRunner | None = None + self.refined_full: torch.Tensor | None = None + self.next_ref_idx = 0 + self.block_size = DEFAULT_REFINER_BLOCK_SIZE + self.sink_size = 1 + + def dispose(self): + super().dispose() + self.runner = None + self.refined_full = None + self.next_ref_idx = 0 + self.block_size = DEFAULT_REFINER_BLOCK_SIZE + self.sink_size = 1 + + +# --------------------------------------------------------------------- # +# Chain stages +# --------------------------------------------------------------------- # +class SanaWMCondFrameEncodeStage(SanaWMRealtimeStage): + """Encode the first frame and write the session input snapshot.""" + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + session = self.require_session(batch, context="SANA-WM realtime chain") + self._pipeline_config = server_args.pipeline_config + device = get_local_torch_device() + weight_dtype = PRECISION_TO_TYPE.get( + getattr(server_args.pipeline_config, "dit_precision", "bf16"), + torch.bfloat16, + ) + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + + self.vae = self.vae.to(device=device, dtype=vae_dtype).eval() + configure_sana_wm_ltx2_vae_for_long_video(self.vae, server_args.pipeline_config) + + st = session.get_or_create_state(SanaWMSessionInputsState) + if batch.block_idx == 0 or st.image is None: + st.dispose() + ( + st.image, + st.intrinsics_image, + st.src_size, + st.resized_size, + st.crop_offset, + ) = self._prepare_image(batch) + st.target_width, st.target_height = st.image.size + st.prompt = str(batch.prompt) + st.open_ended = self._is_open_ended(batch) + if st.open_ended: + st.latent_t = None + else: + num_frames = snap_sana_wm_num_frames(int(batch.num_frames), stride=8) + batch.num_frames = num_frames + st.latent_t = (num_frames - 1) // 8 + 1 + st.num_frame_per_block = int( + batch.extra.get("sana_wm_num_frame_per_block", 3) + ) + st.sink_size = int(batch.extra.get("sana_wm_sink_size", 1)) + st.translation_speed = _motion_param( + batch, "translation_speed", _SANA_WM_DEFAULT_TRANSLATION_SPEED + ) + st.rotation_speed_deg = _motion_param( + batch, "rotation_speed_deg", _SANA_WM_DEFAULT_ROTATION_SPEED_DEG + ) + st.static_c2w = self._prepare_static_camera( + batch, + num_frames=( + ((st.latent_t - 1) * 8 + 1) + if st.latent_t is not None + else st.num_frame_per_block * 8 + 1 + ), + translation_speed=st.translation_speed, + rotation_speed_deg=st.rotation_speed_deg, + ) + st.first_latent = self._get_first_frame_latent( + batch, + st.image, + device=device, + vae_dtype=vae_dtype, + latent_dtype=weight_dtype, + ) + batch.image_latent = st.first_latent + return batch + + +class SanaWMRealtimeLatentPrepStage(SanaWMRealtimeStage): + """This tick's chunk PLAN + pre-noised latents. + + Fixed N: front-loaded segments + a full-horizon seeded buffer sliced per + chunk (bitwise-matching the offline draw); past the horizon / open-ended: + uniform chunks from the seeded fallback generator. The plan holds more than + one chunk only when the refiner block grid lags the stage-1 chunk grid (the + option-b boundary tick) so every tick emits frames.""" + + def __init__(self, *, use_refiner: bool, **kwargs) -> None: + super().__init__(**kwargs) + self.use_refiner = bool(use_refiner) + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + device = get_local_torch_device() + weight_dtype = PRECISION_TO_TYPE.get( + getattr(server_args.pipeline_config, "dit_precision", "bf16"), + torch.bfloat16, + ) + session = self.require_session(batch, context="SANA-WM realtime chain") + inputs = session.get_or_create_state(SanaWMSessionInputsState) + noise = session.get_or_create_state(SanaWMNoiseState) + cache = session.get_or_create_state(SanaWMStreamCacheState) + first_latent = batch.image_latent + if first_latent is None: + raise ValueError("cond-frame latent missing (run the encode stage first)") + + nfpb = int(inputs.num_frame_per_block) + if batch.block_idx == 0 or ( + noise.noise_buffer is None and noise.generator is None + ): + noise.dispose() + seed = batch.seed[0] if isinstance(batch.seed, list) else batch.seed + if seed is not None: + noise.generator = torch.Generator(device=device).manual_seed(int(seed)) + if inputs.latent_t is not None: + noise.segments = SanaWMStreamingDenoisingStage._autoregressive_segments( + int(inputs.latent_t), nfpb + ) + gen = ( + batch.generator[0] + if isinstance(batch.generator, list) + else batch.generator + ) or noise.generator + buf = torch.randn( + 1, + first_latent.shape[1], + int(inputs.latent_t), + first_latent.shape[-2], + first_latent.shape[-1], + device=device, + dtype=weight_dtype, + generator=gen, + ) + buf[:, :, :1] = first_latent + noise.noise_buffer = buf + parity_probe.dump_tensor( # parity harness (no-op in prod) + parity_probe.probe_dir(parity_probe.ENV_RT_DUMP), + "noise_buffer", + buf, + ) + + # Chunk plan: enough stage-1 chunks for >=1 refiner block. + stage1_cur = cache.chunk_indices[-1] + chunk_idx = cache.chunk_idx + + def _next_len(idx: int) -> int: + if noise.segments is not None and idx + 1 < len(noise.segments): + # Front-loaded grid (chunk 0's segment already includes the + # conditioning frame). + return noise.segments[idx + 1] - noise.segments[idx] + n = nfpb + if idx == 0: + # Uniform (open-ended) grid: chunk 0 = cond frame + nfpb new + # frames, matching the retired engine's step(n_frames=nfpb). + n += first_latent.shape[2] + return n + + plan: list[int] = [] + sim = stage1_cur + idx = chunk_idx + if self.use_refiner and chunk_idx > 0: + sink = int(inputs.sink_size) + block = int(inputs.num_frame_per_block) + refined_cur = sink + max(0, (sim - sink)) // block * block + target = refined_cur + block + while sim < target: + n = _next_len(idx) + plan.append(n) + sim += n + idx += 1 + else: + plan.append(_next_len(idx)) + + pieces = [] + pos = stage1_cur + for i, n in enumerate(plan): + is_chunk0 = (chunk_idx + i) == 0 + new_frames = n - (first_latent.shape[2] if is_chunk0 else 0) + lo = pos + (first_latent.shape[2] if is_chunk0 else 0) + hi = lo + new_frames + if noise.noise_buffer is not None and hi <= noise.noise_buffer.shape[2]: + draw = noise.noise_buffer[:, :, lo:hi].to(device, weight_dtype) + else: + draw = torch.randn( + first_latent.shape[0], + first_latent.shape[1], + new_frames, + first_latent.shape[-2], + first_latent.shape[-1], + device=device, + dtype=weight_dtype, + generator=noise.generator, + ) + pieces.append( + torch.cat([first_latent.to(device, weight_dtype), draw], dim=2) + if is_chunk0 + else draw + ) + pos = hi + batch.latents = torch.cat(pieces, dim=2) + batch.extra["sana_wm_chunk_plan"] = plan + return batch + + +class SanaWMCameraCondStage(SanaWMRealtimeStage): + """Camera conditioning windows — appends this tick's actions and rebuilds + the full-length raymap/chunk_plucker through the END of the planned chunks + (single-sourced through the batch helpers; parity RC#3).""" + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + device = get_local_torch_device() + weight_dtype = PRECISION_TO_TYPE.get( + getattr(server_args.pipeline_config, "dit_precision", "bf16"), + torch.bfloat16, + ) + session = self.require_session(batch, context="SANA-WM realtime chain") + inputs = session.get_or_create_state(SanaWMSessionInputsState) + cache = session.get_or_create_state(SanaWMStreamCacheState) + plan = list(batch.extra.get("sana_wm_chunk_plan") or []) + target_latent = cache.chunk_indices[-1] + sum(plan) + if cache.chunk_idx == 0 and plan: + target_latent = max(target_latent, plan[0]) + + self._append_realtime_camera_actions(batch, inputs) + camera, plucker = self._build_camera_windows( + batch, + inputs, + target_latent=target_latent, + device=device, + dtype=weight_dtype, + ) + batch.extra["camera_conditions"] = camera + batch.extra["chunk_plucker"] = plucker + return batch + + def _build_camera_windows( + self, + batch: Req, + inputs: SanaWMSessionInputsState, + *, + target_latent: int, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Port of the mega-stage's ``_update_camera_tensors`` (value-identical + ops; returns instead of writing legacy state). Coverage runs through the + END of the planned chunks; within a fixed horizon ``max()`` keeps it + EXACTLY at num_frames — bitwise-identical camera tensors.""" + if ( + inputs.src_size is None + or inputs.resized_size is None + or inputs.crop_offset is None + ): + raise ValueError("SANA-WM crop metadata is not initialized") + needed = (int(target_latent) - 1) * 8 + 1 + if inputs.open_ended: + num_frames = needed + else: + num_frames = max(int(batch.num_frames), needed) + condition_inputs = batch.condition_inputs or {} + if "translation_speed" in condition_inputs: + inputs.translation_speed = float(condition_inputs["translation_speed"]) + if "rotation_speed_deg" in condition_inputs: + inputs.rotation_speed_deg = float(condition_inputs["rotation_speed_deg"]) + c2w = self._camera_from_state( + inputs, + num_frames=num_frames, + translation_speed=inputs.translation_speed, + rotation_speed_deg=inputs.rotation_speed_deg, + ) + intrinsics_raw = self._prepare_intrinsics( + batch, inputs, num_frames=num_frames, device=device + ) + inputs.intrinsics_raw = intrinsics_raw + vae_time_stride = 8 + pixel_h = int(inputs.target_height or batch.height or SANA_WM_HEIGHT) + pixel_w = int(inputs.target_width or batch.width or SANA_WM_WIDTH) + latent_h = pixel_h // 32 + latent_w = pixel_w // 32 + latent_t = (num_frames - 1) // vae_time_stride + 1 + camera_to_world = ( + torch.from_numpy(np.asarray(c2w, dtype=np.float32)) + .unsqueeze(0) + .to(device=device, dtype=torch.float32) + ) + intrinsics_vec4 = ( + torch.from_numpy(np.asarray(intrinsics_raw, dtype=np.float32)) + .unsqueeze(0) + .to(device=device, dtype=torch.float32) + ) + intrinsics_vec4 = ( + SanaWMBeforeDenoisingStage._transform_intrinsics_for_condition_image( + intrinsics_vec4, + { + "source_size": inputs.src_size, + "resized_size": inputs.resized_size, + "crop_offset": inputs.crop_offset, + }, + ) + ) + rel_poses = SanaWMBeforeDenoisingStage._relative_camera_poses(camera_to_world) + intrinsics_latent = SanaWMBeforeDenoisingStage._scale_intrinsics_to_latent( + intrinsics_vec4, + pixel_h=pixel_h, + pixel_w=pixel_w, + latent_h=latent_h, + latent_w=latent_w, + ) + original_camera = SanaWMBeforeDenoisingStage._flatten_camera_conditions( + rel_poses, intrinsics_latent + ) + raymap = SanaWMBeforeDenoisingStage._latent_frame_camera_conditions( + original_camera, + num_frames=num_frames, + latent_frames=latent_t, + vae_temporal_stride=vae_time_stride, + ).to(device=device, dtype=dtype) + chunk_plucker = compute_chunk_plucker( + original_camera, + HW=(latent_t, latent_h, latent_w), + vae_temporal_stride=vae_time_stride, + patch_size=(1, 1, 1), + ).to(device=device, dtype=dtype) + return raymap, chunk_plucker + + +class SanaWMChunkedRefinerChainStage(SanaWMRealtimeStage): + """Chunked LTX-2 refiner on the uniform ``sink + i*block`` grid — refines + COMPLETE blocks only (option b) and hands the refined buffer downstream.""" + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + device = get_local_torch_device() + session = self.require_session(batch, context="SANA-WM realtime chain") + inputs = session.get_or_create_state(SanaWMSessionInputsState) + st = session.get_or_create_state(SanaWMRefinerChainState) + stage1 = batch.latents # growing stage-1 buffer from the denoise stage + if stage1 is None: + raise ValueError("refiner stage expects the stage-1 latent buffer") + + with set_forward_context( + current_timestep=batch.block_idx, attn_metadata=None, forward_batch=batch + ): + if st.runner is None: + st.block_size = int(inputs.num_frame_per_block) + st.sink_size = int(inputs.sink_size) + first_latent = batch.image_latent + st.refined_full = first_latent.detach().clone() + st.runner = self._build_chain_refiner_runner( + inputs, + st, + batch, + device=device, + spatial_shape=( + int(first_latent.shape[3]), + int(first_latent.shape[4]), + ), + ) + + frontier = st.refined_full.shape[2] + end_f = stage1.shape[2] + while True: + block_start = frontier + block_end = block_start + st.block_size # complete blocks only + if block_end > end_f: + break + sink_seed = ( + stage1[:, :, : st.sink_size] + if block_start == st.sink_size + else None + ) + refined = st.runner.refine_block( + block_idx=st.next_ref_idx, + clean_block=stage1[:, :, block_start:block_end].contiguous(), + block_start=block_start, + block_end=block_end, + sink_seed_frames=sink_seed, + ) + st.refined_full = torch.cat( + [st.refined_full, refined.to(st.refined_full.dtype)], dim=2 + ) + st.next_ref_idx += 1 + frontier = block_end + + batch.latents = st.refined_full + return batch + + def _build_chain_refiner_runner( + self, inputs, st, batch, *, device, spatial_shape + ) -> RefinerChunkRunner: + # Reuse the parity-validated builder via a state adapter that exposes + # prompt/sink/block/kv_max in the legacy state shape it expects. + legacy = SimpleNamespace( + prompt=inputs.prompt, + sink_size=st.sink_size, + refiner_block_size=st.block_size, + refiner_kv_max_frames=int( + batch.extra.get( + "sana_wm_refiner_kv_max_frames", DEFAULT_REFINER_KV_MAX_FRAMES + ) + ), + ) + return self._build_refiner_runner( + legacy, + batch, + device=device, + spatial_shape=spatial_shape, + seed=int(batch.extra.get("sana_wm_refiner_seed", batch.seed)), + fps=float(batch.fps), + ) + + +class SanaWMCausalDecodeChainStage(SanaWMRealtimeStage): + """Causal-VAE chunk decode past this session's decode frontier.""" + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: + device = get_local_torch_device() + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + + self.vae = self.vae.to(device=device, dtype=vae_dtype).eval() + configure_sana_wm_ltx2_vae_for_long_video(self.vae, server_args.pipeline_config) + session = self.require_session(batch, context="SANA-WM realtime chain") + st = session.get_or_create_state(RealtimeCausalDecodeState) + src = batch.latents + if src is None or src.shape[2] <= st.next_dec_idx: + return self._empty_output(batch) + with set_forward_context( + current_timestep=batch.block_idx, attn_metadata=None, forward_batch=batch + ): + frames = self._decode_chunk( + src[:, :, st.next_dec_idx :], st, server_args, vae_dtype=vae_dtype + ) + st.next_dec_idx = src.shape[2] + return OutputBatch(output=frames.to(torch.float32), metrics=batch.metrics) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_stage.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_stage.py new file mode 100644 index 000000000..ea97ee27a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_stage.py @@ -0,0 +1,448 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import os +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image + +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_diffusion import ( + RealtimeDiffusionStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +from .base import ( + SANA_WM_TARGET_HEIGHT, + SANA_WM_TARGET_WIDTH, + SanaWMBeforeDenoisingStage, + normalize_sana_wm_camera_actions, + sana_wm_action_to_camera_to_world_array, + sana_wm_load_camera, + sana_wm_load_intrinsics, + sana_wm_normalize_vae_latents, + sana_wm_pil_to_model_tensor, + sana_wm_resize_and_center_crop, +) +from .refiner import ( + STAGE_2_DISTILLED_SIGMA_VALUES, + SanaWMLTX2RefinerStage, + _unwrap_diffusers_ltx2_refiner, +) +from .streaming_refiner import ( + RefinerChunkRunner, + _RefinerCore, +) + +logger = init_logger(__name__) + +SANA_WM_HEIGHT = SANA_WM_TARGET_HEIGHT +SANA_WM_WIDTH = SANA_WM_TARGET_WIDTH +DEFAULT_REFINER_BLOCK_SIZE = 3 +DEFAULT_REFINER_KV_MAX_FRAMES = 11 + + +@contextmanager +def _deterministic_vae_encode_context(): + prev_benchmark = torch.backends.cudnn.benchmark + prev_deterministic = torch.backends.cudnn.deterministic + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + try: + yield + finally: + torch.backends.cudnn.benchmark = prev_benchmark + torch.backends.cudnn.deterministic = prev_deterministic + + +def _normalize_camera_actions(payload: Any) -> list[list[str]]: + return normalize_sana_wm_camera_actions(payload, allow_none=True) + + +def _actions_to_action_string(actions: list[list[str]]) -> str: + if not actions: + return "none-1" + + segments: list[str] = [] + current = tuple(sorted(set(actions[0]))) + count = 0 + for frame_actions in actions: + normalized = tuple(sorted(set(frame_actions))) + if normalized == current: + count += 1 + continue + key = "".join(current) if current else "none" + segments.append(f"{key}-{count}") + current = normalized + count = 1 + key = "".join(current) if current else "none" + segments.append(f"{key}-{count}") + return ",".join(segments) + + +def _normalize_intrinsics_array(arr: Any, num_frames: int) -> np.ndarray: + intrinsics = np.asarray(arr, dtype=np.float32) + if intrinsics.shape == (4,): + return np.broadcast_to(intrinsics, (num_frames, 4)).copy() + if intrinsics.shape == (3, 3): + vec = np.array( + [ + intrinsics[0, 0], + intrinsics[1, 1], + intrinsics[0, 2], + intrinsics[1, 2], + ], + dtype=np.float32, + ) + return np.broadcast_to(vec, (num_frames, 4)).copy() + if intrinsics.ndim == 2 and intrinsics.shape[1] == 4: + if intrinsics.shape[0] < num_frames: + pad = np.broadcast_to( + intrinsics[-1:], (num_frames - intrinsics.shape[0], 4) + ) + intrinsics = np.concatenate([intrinsics, pad], axis=0) + return intrinsics[:num_frames].copy() + if intrinsics.ndim == 3 and intrinsics.shape[1:] == (3, 3): + if intrinsics.shape[0] < num_frames: + pad = np.broadcast_to( + intrinsics[-1:], (num_frames - intrinsics.shape[0], 3, 3) + ) + intrinsics = np.concatenate([intrinsics, pad], axis=0) + intrinsics = intrinsics[:num_frames] + return np.stack( + [ + intrinsics[:, 0, 0], + intrinsics[:, 1, 1], + intrinsics[:, 0, 2], + intrinsics[:, 1, 2], + ], + axis=1, + ).astype(np.float32) + raise ValueError( + "intrinsics must have shape (4,), (3,3), (F,4), or (F,3,3), " + f"got {intrinsics.shape}" + ) + + +def _motion_param(batch: Req, name: str, default: float) -> float: + value = (batch.condition_inputs or {}).get(name) + if value is None: + value = batch.extra.get(f"sana_wm_{name}", default) + return float(value) + + +def _default_source_intrinsics(image: Image.Image, num_frames: int) -> np.ndarray: + width, height = image.size + focal = 0.8 * float(max(width, height)) + intrinsics = np.array( + [focal, focal, width / 2.0, height / 2.0], + dtype=np.float32, + ) + return np.broadcast_to(intrinsics, (num_frames, 4)).copy() + + +class SanaWMRealtimeStage(RealtimeDiffusionStage): + def __init__( + self, + *, + transformer: torch.nn.Module, + vae: torch.nn.Module, + model_path: str, + refiner_stage: SanaWMLTX2RefinerStage | None = None, + ) -> None: + super().__init__( + transformer=transformer, + vae=vae, + model_path=model_path, + default_height=SANA_WM_HEIGHT, + default_width=SANA_WM_WIDTH, + ) + # Injected by the pipeline; we reuse its transformer / text encoder to + # drive the RefinerChunkRunner. None -> stage-1-only output. + self.refiner_stage = refiner_stage + self.first_frame_latent_cache = None + + def _prepare_image( + self, batch: Req + ) -> tuple[ + Image.Image, Image.Image, tuple[int, int], tuple[int, int], tuple[int, int] + ]: + if isinstance(batch.condition_image, Image.Image): + original = batch.condition_image.convert("RGB") + elif batch.image_path is not None and isinstance(batch.image_path, str): + original = Image.open(batch.image_path).convert("RGB") + else: + raise ValueError("SANA-WM realtime requires a first-frame image") + target_h, target_w = self.target_pixel_size(batch) + cropped, src_size, resized_size, crop_offset = sana_wm_resize_and_center_crop( + original, target_h, target_w + ) + return cropped, original, src_size, resized_size, crop_offset + + def _prepare_static_camera( + self, + batch: Req, + *, + num_frames: int, + translation_speed: float, + rotation_speed_deg: float, + ) -> np.ndarray | None: + condition_inputs = batch.condition_inputs or {} + camera_path = condition_inputs.get("camera_path") + if camera_path is not None: + c2w = sana_wm_load_camera(Path(str(camera_path))) + elif condition_inputs.get("camera") is not None: + c2w = np.asarray(condition_inputs["camera"], dtype=np.float32) + elif condition_inputs.get("action") is not None: + c2w = sana_wm_action_to_camera_to_world_array( + str(condition_inputs["action"]), + translation_speed=translation_speed, + rotation_speed_deg=rotation_speed_deg, + ) + else: + return None + + if c2w.ndim != 3 or c2w.shape[1:] != (4, 4): + raise ValueError( + f"camera trajectory must have shape (F,4,4), got {c2w.shape}" + ) + if c2w.shape[0] < num_frames: + pad = np.broadcast_to(c2w[-1:], (num_frames - c2w.shape[0], 4, 4)) + c2w = np.concatenate([c2w, pad], axis=0) + return c2w[:num_frames].astype(np.float32) + + def _prepare_intrinsics( + self, + batch: Req, + state, + *, + num_frames: int, + device: torch.device, + ) -> np.ndarray: + condition_inputs = batch.condition_inputs or {} + if condition_inputs.get("intrinsics_path") is not None: + return sana_wm_load_intrinsics( + Path(str(condition_inputs["intrinsics_path"])), num_frames + ) + if condition_inputs.get("intrinsics") is not None: + return _normalize_intrinsics_array( + condition_inputs["intrinsics"], num_frames + ) + if state.intrinsics_raw is not None: + cached = state.intrinsics_raw + if cached.shape[0] < num_frames: + # Open-ended growth: hold the last row (fixed-horizon sessions + # cache the full-length array, so this never triggers). + pad = np.broadcast_to( + cached[-1:], (num_frames - cached.shape[0],) + cached.shape[1:] + ) + cached = np.concatenate([cached, pad], axis=0) + return cached + if state.intrinsics_image is None: + raise ValueError("SANA-WM image is not initialized") + logger.info("No intrinsics provided; using heuristic centered intrinsics.") + return _default_source_intrinsics(state.intrinsics_image, num_frames) + + def _append_realtime_camera_actions(self, batch: Req, state) -> None: + actions = _normalize_camera_actions( + (batch.condition_inputs or {}).get("camera_actions") + ) + if actions: + state.camera_actions.extend(actions) + if ( + state.max_camera_actions > 0 + and len(state.camera_actions) > state.max_camera_actions + ): + del state.camera_actions[state.max_camera_actions :] + + def _camera_from_state( + self, + state, + *, + num_frames: int, + translation_speed: float, + rotation_speed_deg: float, + ) -> np.ndarray: + if state.static_c2w is not None: + c2w = state.static_c2w + if c2w.shape[0] < num_frames: + # Open-ended growth: hold the last pose (fixed-horizon sessions + # always precompute the full trajectory, so this never triggers). + pad = np.broadcast_to(c2w[-1:], (num_frames - c2w.shape[0], 4, 4)) + c2w = np.concatenate([c2w, pad], axis=0) + return c2w[:num_frames] + + num_actions = max(0, num_frames - 1) + actions = list(state.camera_actions[:num_actions]) + if len(actions) < num_actions: + actions.extend([[] for _ in range(num_actions - len(actions))]) + return sana_wm_action_to_camera_to_world_array( + _actions_to_action_string(actions), + translation_speed=translation_speed, + rotation_speed_deg=rotation_speed_deg, + )[:num_frames] + + @torch.inference_mode() + def _encode_first_frame( + self, + image: Image.Image, + *, + device: torch.device, + vae_dtype: torch.dtype, + latent_dtype: torch.dtype, + ) -> torch.Tensor: + if hasattr(self.vae, "enable_tiling"): + self.vae.enable_tiling() + image_tensor = sana_wm_pil_to_model_tensor( + image, device=device, dtype=vae_dtype + ) + with _deterministic_vae_encode_context(): + # Must use the SAME shared core as the batch path's _vae_encode_image: + # drifting this was a parity root cause. + z = SanaWMBeforeDenoisingStage._extract_vae_latents( + self.vae.encode(image_tensor.to(device=device, dtype=vae_dtype)) + ).float() + z = sana_wm_normalize_vae_latents( + self.vae, z, getattr(self, "_pipeline_config", None) + ) + return z.to(device=device, dtype=latent_dtype) + + def _first_frame_cache_key( + self, + batch: Req, + *, + device: torch.device, + vae_dtype: torch.dtype, + latent_dtype: torch.dtype, + ) -> tuple | None: + if batch.image_path is None or not isinstance(batch.image_path, str): + return None + stat = os.stat(batch.image_path) + target_h, target_w = self.target_pixel_size(batch) + return ( + batch.image_path, + stat.st_mtime_ns, + stat.st_size, + target_h, + target_w, + device.type, + device.index, + vae_dtype, + latent_dtype, + ) + + @torch.inference_mode() + def _get_first_frame_latent( + self, + batch: Req, + image: Image.Image, + *, + device: torch.device, + vae_dtype: torch.dtype, + latent_dtype: torch.dtype, + ) -> torch.Tensor: + cache_key = self._first_frame_cache_key( + batch, + device=device, + vae_dtype=vae_dtype, + latent_dtype=latent_dtype, + ) + if ( + cache_key is not None + and self.first_frame_latent_cache is not None + and self.first_frame_latent_cache[0] == cache_key + ): + return self.first_frame_latent_cache[1] + + first_latent = self._encode_first_frame( + image, + device=device, + vae_dtype=vae_dtype, + latent_dtype=latent_dtype, + ) + if cache_key is not None: + self.first_frame_latent_cache = (cache_key, first_latent.detach()) + return first_latent + + @torch.inference_mode() + def _decode_chunk( + self, + latents: torch.Tensor, + state, + server_args: ServerArgs, + *, + vae_dtype: torch.dtype, + ) -> torch.Tensor: + if latents.shape[2] == 0: + vae_stride = getattr(server_args.pipeline_config, "vae_stride", (8, 32, 32)) + return torch.empty( + ( + latents.shape[0], + 3, + 0, + int(latents.shape[3]) * int(vae_stride[1]), + int(latents.shape[4]) * int(vae_stride[2]), + ), + dtype=torch.float32, + device=latents.device, + ) + conv_cache = self.ensure_causal_vae_conv_cache(state) + z = self.scale_and_shift_latents(latents.to(vae_dtype), server_args) + decoded = self.vae.decode_chunk(z, conv_cache) + return (decoded / 2 + 0.5).clamp(0, 1) + + def _build_refiner_runner( + self, + state, + batch: Req, + *, + device: torch.device, + spatial_shape: tuple[int, int], + seed: int, + fps: float, + ) -> RefinerChunkRunner: + """Build the chunked LTX-2 refiner runner that carries sink/history KV across blocks.""" + rs = self.refiner_stage + if rs is None: + raise RuntimeError("SANA-WM realtime refiner stage is not initialized") + # Keep refiner sub-modules resident: the realtime stage runs the refiner + # directly, outside the offline stage's use_declared_component context. + for _mod in (rs.text_encoder, rs.connectors, rs.transformer): + if _mod is not None: + _mod.to(device) + prompt_embeds, prompt_mask = rs._encode_prompt(state.prompt, device) + unwrapped = _unwrap_diffusers_ltx2_refiner(rs.transformer) + core = _RefinerCore(unwrapped, device, rs.dtype) + sigmas = torch.tensor( + STAGE_2_DISTILLED_SIGMA_VALUES, dtype=torch.float32, device=device + ) + return RefinerChunkRunner( + core, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_mask, + fps=fps, + sigmas=sigmas, + source_sink_frames=state.sink_size, + block_size=state.refiner_block_size, + kv_max_frames=state.refiner_kv_max_frames, + seed=seed, + spatial_shape=spatial_shape, + ) + + @staticmethod + def _is_open_ended(batch: Req) -> bool: + """Open-ended = the init request carried no num_frames. + + The adapter flags this via condition_inputs (build_sampling_params strips + None fields, so batch.num_frames would otherwise carry the SamplingParams + default and be indistinguishable from an explicit request).""" + condition_inputs = batch.condition_inputs or {} + if bool(condition_inputs.get("sana_wm_open_ended")): + return True + num_frames = getattr(batch, "num_frames", None) + return num_frames is None or int(num_frames) <= 0 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/refiner.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/refiner.py new file mode 100644 index 000000000..31e1f7213 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/refiner.py @@ -0,0 +1,799 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +from typing import Any + +import torch +from torch import nn + +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + STAGE_2_DISTILLED_SIGMA_VALUES as _STAGE_2_DISTILLED_SIGMA_VALUES, +) +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + pack_text_embeds, +) +from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_classifier_free_guidance_rank, +) +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + ComponentUse, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) +from sglang.multimodal_gen.runtime.models.dits.sana_wm_refiner_transformer import ( + pack_latents, + unpack_latents, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( + PipelineStage, + StageParallelismType, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +from .base import ( + SanaWMDecodingStage, + log_sana_wm_tensor_stats, + sana_wm_diagnostics_enabled, +) + +logger = init_logger(__name__) + +# Distilled 3-step sigma schedule, matches NVlabs `inference_sana_wm.py`. +# Canonical value lives in the LTX-2 pipeline config (shared with LTX2TwoStagePipeline). +STAGE_2_DISTILLED_SIGMA_VALUES: tuple[float, ...] = _STAGE_2_DISTILLED_SIGMA_VALUES + +# Default Gemma-3 token budget for the refiner prompt encoder. +_REFINER_TEXT_MAX_LENGTH = 1024 + + +class _OfficialLayerwiseModule(nn.Module, LayerwiseOffloadableModuleMixin): + def __init__(self, module: nn.Module) -> None: + super().__init__() + self.module = module + self.layerwise_offload_managers = [] + + def forward(self, *args, **kwargs): + return self.module(*args, **kwargs) + + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError as exc: + wrapped = self.__dict__.get("module") + if wrapped is not None: + return getattr(wrapped, name) + modules = self.__dict__.get("_modules", {}) + wrapped = modules.get("module") + if wrapped is not None: + return getattr(wrapped, name) + raise exc + + +class OfficialDiffusersLTX2RefinerModule(_OfficialLayerwiseModule): + """Thin offload wrapper around Diffusers' official LTX-2 refiner module.""" + + layer_names = ["module.transformer_blocks"] + + +class OfficialGemma3TextEncoderModule(_OfficialLayerwiseModule): + """Thin offload wrapper around HF Gemma-3 used by the official refiner.""" + + layer_names = [ + "module.language_model.layers", + "module.model.language_model.layers", + ] + + +def _truthy_flag(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return False + + +def sana_wm_skip_refiner_enabled(batch: Req | None = None) -> bool: + if os.getenv("SGLANG_SANA_WM_SKIP_REFINER", "").strip().lower() in { + "1", + "true", + "yes", + "on", + }: + return True + if batch is None: + return False + extra = getattr(batch, "extra", None) or {} + diffusers_kwargs = extra.get("diffusers_kwargs", {}) + if not isinstance(diffusers_kwargs, dict): + diffusers_kwargs = {} + return any( + _truthy_flag(value) + for value in ( + extra.get("skip_refiner"), + extra.get("sana_wm_skip_refiner"), + diffusers_kwargs.get("skip_refiner"), + diffusers_kwargs.get("sana_wm_skip_refiner"), + ) + ) + + +def default_sana_wm_refiner_dtype(server_args: ServerArgs) -> torch.dtype: + precision = getattr(server_args.pipeline_config, "dit_precision", "bf16") + return PRECISION_TO_TYPE.get(precision, torch.bfloat16) + + +def _is_current_cfg_main_rank() -> bool: + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return True + try: + return get_classifier_free_guidance_rank() == 0 + except AssertionError: + return True + + +def _pack_text_embeds( + text_hidden_states: torch.Tensor, + sequence_lengths: torch.Tensor, + *, + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, +) -> torch.Tensor: + """Gemma-3 masked min-max text-embed pooling for the LTX-2 refiner. + + Delegates to framework-wide ``pack_text_embeds`` (verified BITWISE-identical, + incl. bf16 and both padding sides) — a private fork would drift, like the + realtime<->batch parity bugs. + """ + return pack_text_embeds( + text_hidden_states, + sequence_lengths, + padding_side=padding_side, + scale_factor=scale_factor, + eps=eps, + ) + + +def _refiner_config_value(transformer: nn.Module, name: str) -> Any: + transformer = _unwrap_diffusers_ltx2_refiner(transformer) + config = getattr(transformer, "config", None) + if config is not None: + if isinstance(config, dict) and name in config: + return config[name] + if hasattr(config, name): + return getattr(config, name) + return getattr(transformer, name) + + +def _unwrap_diffusers_ltx2_refiner(transformer: nn.Module) -> nn.Module: + if isinstance(transformer, OfficialDiffusersLTX2RefinerModule): + return transformer.module + return transformer + + +def _uses_diffusers_ltx2_refiner(transformer: nn.Module) -> bool: + transformer = _unwrap_diffusers_ltx2_refiner(transformer) + return transformer.__class__.__name__ == "LTX2VideoTransformer3DModel" + + +def _as_additive_attention_mask( + attention_mask: torch.Tensor | None, + dtype: torch.dtype, +) -> torch.Tensor | None: + if attention_mask is None: + return None + if attention_mask.ndim == 2: + return ((1 - attention_mask.to(dtype)) * -10000.0).unsqueeze(1) + return attention_mask.to(dtype) + + +def _forward_diffusers_video_only( + transformer: nn.Module, + *, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_attention_mask: torch.Tensor | None, + num_frames: int, + height: int, + width: int, + fps: float, + n_context_tokens: int, +) -> torch.Tensor: + """Official SANA-WM LTX-2 video-only forward adapted to injected modules.""" + + batch_size = hidden_states.size(0) + encoder_attention_mask = _as_additive_attention_mask( + encoder_attention_mask, hidden_states.dtype + ) + + video_coords = transformer.rope.prepare_video_coords( + batch_size, num_frames, height, width, hidden_states.device, fps=fps + ) + video_rotary_emb = transformer.rope(video_coords, device=hidden_states.device) + + hidden_states = transformer.proj_in(hidden_states) + temb, embedded_timestep = transformer.time_embed( + timestep.flatten(), + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(batch_size, -1, temb.size(-1)) + embedded_timestep = embedded_timestep.view( + batch_size, -1, embedded_timestep.size(-1) + ) + + encoder_hidden_states = transformer.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view( + batch_size, -1, hidden_states.size(-1) + ) + + for block in transformer.transformer_blocks: + hidden_states = _forward_diffusers_video_block( + block=block, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + video_rotary_emb=video_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + n_context_tokens=n_context_tokens, + ) + + scale_shift_values = ( + transformer.scale_shift_table[None, None] + embedded_timestep[:, :, None] + ) + shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] + hidden_states = transformer.norm_out(hidden_states) + hidden_states = hidden_states * (1 + scale) + shift + return transformer.proj_out(hidden_states) + + +def _forward_diffusers_video_block( + *, + block: nn.Module, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + video_rotary_emb: tuple[torch.Tensor, torch.Tensor], + encoder_attention_mask: torch.Tensor | None, + n_context_tokens: int, +) -> torch.Tensor: + batch_size = hidden_states.size(0) + + norm_hidden_states = block.norm1(hidden_states) + num_ada_params = block.scale_shift_table.shape[0] + ada_values = block.scale_shift_table[None, None].to(temb.device) + temb.reshape( + batch_size, temb.size(1), num_ada_params, -1 + ) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind( + dim=2 + ) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + + attn_hidden_states = _streaming_diffusers_self_attention( + attn=block.attn1, + hidden_states=norm_hidden_states, + query_rotary_emb=video_rotary_emb, + n_context_tokens=n_context_tokens, + ) + hidden_states = hidden_states + attn_hidden_states * gate_msa + + norm_hidden_states = block.norm2(hidden_states) + attn_hidden_states = block.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + query_rotary_emb=None, + attention_mask=encoder_attention_mask, + ) + hidden_states = hidden_states + attn_hidden_states + + norm_hidden_states = block.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp + hidden_states = hidden_states + block.ff(norm_hidden_states) * gate_mlp + return hidden_states + + +def _streaming_diffusers_self_attention( + *, + attn: nn.Module, + hidden_states: torch.Tensor, + query_rotary_emb: tuple[torch.Tensor, torch.Tensor], + n_context_tokens: int, +) -> torch.Tensor: + """SANA-WM sink/current streaming mask using Diffusers LTX-2 attention.""" + + sequence_length = hidden_states.shape[1] + if n_context_tokens <= 0 or n_context_tokens >= sequence_length: + return attn( + hidden_states=hidden_states, + encoder_hidden_states=None, + query_rotary_emb=query_rotary_emb, + ) + + from diffusers.models.attention_dispatch import dispatch_attention_fn + from diffusers.models.transformers.transformer_ltx2 import ( + apply_interleaved_rotary_emb, + apply_split_rotary_emb, + ) + + # Diffusers 0.38+ always defines `to_gate_logits`, while 0.37 only has it + # on gated variants. The public SANA-WM refiner config is ungated, so a + # missing attribute means the same thing as `None`. + to_gate_logits = getattr(attn, "to_gate_logits", None) + gate_logits = to_gate_logits(hidden_states) if to_gate_logits is not None else None + + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + if attn.rope_type == "interleaved": + query = apply_interleaved_rotary_emb(query, query_rotary_emb) + key = apply_interleaved_rotary_emb(key, query_rotary_emb) + elif attn.rope_type == "split": + query = apply_split_rotary_emb(query, query_rotary_emb) + key = apply_split_rotary_emb(key, query_rotary_emb) + else: + raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}") + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + processor = attn.processor + backend = getattr(processor, "_attention_backend", None) + parallel_config = getattr(processor, "_parallel_config", None) + context_hidden_states = dispatch_attention_fn( + query[:, :n_context_tokens], + key[:, :n_context_tokens], + value[:, :n_context_tokens], + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=backend, + parallel_config=parallel_config, + ) + current_hidden_states = dispatch_attention_fn( + query[:, n_context_tokens:], + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=backend, + parallel_config=parallel_config, + ) + + hidden_states = torch.cat([context_hidden_states, current_hidden_states], dim=1) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + + if gate_logits is not None: + hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) + gates = 2.0 * torch.sigmoid(gate_logits) + hidden_states = hidden_states * gates.unsqueeze(-1) + hidden_states = hidden_states.flatten(2, 3) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class SanaWMLTX2RefinerStage(PipelineStage): + """Run the SANA-WM stage-2 LTX-2 refiner before VAE decode. + + Modules are injected by `SanaWMTwoStagePipeline`: + * `transformer` (`SanaWMLTX2VideoRefiner`) -- video-only LTX-2 forward + * `connectors` (`LTX2TextConnectors`) + * `text_encoder` (`Gemma3ForConditionalGeneration`) + * `tokenizer` (HF `AutoTokenizer` for the refiner Gemma-3) + """ + + def __init__( + self, + *, + transformer: nn.Module, + connectors: nn.Module, + text_encoder: nn.Module, + tokenizer: Any, + dtype: torch.dtype, + text_max_sequence_length: int = _REFINER_TEXT_MAX_LENGTH, + ) -> None: + super().__init__() + self.transformer = transformer + self.connectors = connectors + self.text_encoder = text_encoder + self.tokenizer = tokenizer + self.dtype = dtype + self.text_max_sequence_length = int(text_max_sequence_length) + + @property + def role_affinity(self) -> RoleType: + return RoleType.DENOISER + + @property + def parallelism_type(self) -> StageParallelismType: + if getattr(self.server_args, "enable_cfg_parallel", False): + return StageParallelismType.MAIN_RANK_ONLY + return StageParallelismType.REPLICATED + + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + if sana_wm_skip_refiner_enabled(): + return [] + if ( + getattr(server_args, "enable_cfg_parallel", False) + and not _is_current_cfg_main_rank() + ): + return [] + + # Declare every component this stage forwards through so + # ComponentResidencyManager moves them onto GPU before the stage runs. + # Without this, `dit_cpu_offload=True` keeps refiner sub-modules on CPU + # and the first matmul fails with "mat2 is on cpu" vs cuda inputs. + # The tokenizer stays on CPU (no nn.Module weights to ferry). + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name=stage_name, + component_name="text_encoder_2", + target_dtype=self.dtype, + ), + ComponentUse( + stage_name=stage_name, + component_name="connectors", + target_dtype=self.dtype, + ), + ComponentUse( + stage_name=stage_name, + component_name="transformer_2", + target_dtype=self.dtype, + memory_intensive=True, + ), + ] + + @staticmethod + def _prompts_for_batch(batch: Req, batch_size: int) -> list[str]: + prompt = batch.extra.get("refiner_prompt") if batch.extra else None + if prompt is None: + prompt = batch.prompt + if isinstance(prompt, str): + return [prompt] * batch_size + if isinstance(prompt, list) and all(isinstance(p, str) for p in prompt): + if len(prompt) == batch_size: + return prompt + if len(prompt) == 1: + return prompt * batch_size + raise ValueError( + "SANA-WM refiner requires a string prompt or one prompt per batch item." + ) + + @torch.inference_mode() + def _encode_prompt( + self, + prompt: str, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + tokenizer = self.tokenizer + if getattr(tokenizer, "padding_side", "right") != "left": + tokenizer.padding_side = "left" + if tokenizer.pad_token is None and tokenizer.eos_token is not None: + tokenizer.pad_token = tokenizer.eos_token + + text_inputs = tokenizer( + [prompt.strip()], + padding="max_length", + max_length=self.text_max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + input_ids = text_inputs.input_ids.to(device) + attention_mask = text_inputs.attention_mask.to(device) + + # Diffusers-backed official path loads HF Gemma3ForConditionalGeneration. + # NVlabs encodes through `.model`; the fallback SGLang-native encoder is + # still callable directly, so keep both surfaces. + with self.use_declared_component( + component_name="text_encoder_2", module=self.text_encoder + ): + text_backbone = getattr(self.text_encoder, "model", self.text_encoder) + outputs = text_backbone( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + per_layer_hidden = getattr(outputs, "hidden_states", None) + if per_layer_hidden is None: + raise RuntimeError( + "SANA-WM refiner text encoder must return per-layer hidden_states." + ) + stacked = torch.stack(per_layer_hidden, dim=-1) # (B, L, D, n_layers) + seq_lengths = attention_mask.sum(dim=-1) + log_sana_wm_tensor_stats("refiner.text_hidden_states_stacked", stacked) + prompt_embeds = _pack_text_embeds( + stacked, + seq_lengths, + padding_side=tokenizer.padding_side, + ).to(dtype=self.dtype) + log_sana_wm_tensor_stats("refiner.prompt_embeds_packed", prompt_embeds) + + with self.use_declared_component( + component_name="connectors", module=self.connectors + ): + video_text_embedding, _, video_attention_mask = self.connectors( + prompt_embeds, attention_mask + ) + log_sana_wm_tensor_stats("refiner.video_text_embedding", video_text_embedding) + log_sana_wm_tensor_stats("refiner.video_attention_mask", video_attention_mask) + return ( + video_text_embedding.to(device=device, dtype=self.dtype), + video_attention_mask.to(device=device), + ) + + def _predict_current_x0( + self, + *, + sink: torch.Tensor, + noisy_current: torch.Tensor, + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + sigma: torch.Tensor, + fps: float, + n_context_tokens: int, + step_idx: int, + ) -> torch.Tensor: + full_latent = torch.cat([sink, noisy_current], dim=2) + batch_size, _, num_frames, height, width = full_latent.shape + patch_size = int(_refiner_config_value(self.transformer, "patch_size")) + patch_size_t = int(_refiner_config_value(self.transformer, "patch_size_t")) + latent_tokens = pack_latents(full_latent, patch_size, patch_size_t) + + raw_timestep = torch.zeros( + batch_size, + latent_tokens.shape[1], + 1, + dtype=torch.float32, + device=latent_tokens.device, + ) + raw_timestep[:, n_context_tokens:, 0] = sigma.float() + + with self.use_declared_component( + component_name="transformer_2", module=self.transformer + ): + if _uses_diffusers_ltx2_refiner(self.transformer): + model_timestep = raw_timestep.squeeze(-1) * float( + _refiner_config_value(self.transformer, "timestep_scale_multiplier") + ) + velocity_tokens = _forward_diffusers_video_only( + self.transformer, + hidden_states=latent_tokens.to(self.dtype), + encoder_hidden_states=prompt_embeds, + timestep=model_timestep, + encoder_attention_mask=prompt_attention_mask, + num_frames=num_frames, + height=height, + width=width, + fps=fps, + n_context_tokens=n_context_tokens, + ) + else: + additive_mask = _as_additive_attention_mask( + prompt_attention_mask, self.dtype + ) + with set_forward_context( + current_timestep=step_idx, + attn_metadata=None, + ): + velocity_tokens = self.transformer( + hidden_states=latent_tokens.to(self.dtype), + encoder_hidden_states=prompt_embeds, + timestep=raw_timestep.squeeze(-1), + encoder_attention_mask=additive_mask, + num_frames=num_frames, + height=height, + width=width, + fps=fps, + n_context_tokens=n_context_tokens, + ) + + denoised = latent_tokens.float() - velocity_tokens.float() * raw_timestep + return denoised[:, n_context_tokens:, :].to(self.dtype) + + @torch.inference_mode() + def _refine_one( + self, + latent: torch.Tensor, + prompt: str, + *, + fps: float, + seed: int, + sink_size: int = 1, + ) -> torch.Tensor: + device = get_local_torch_device() + z = latent.to(device=device, dtype=self.dtype) + if z.shape[2] <= sink_size: + raise ValueError( + f"Stage-1 latent has {z.shape[2]} frames but sink_size={sink_size}." + ) + self.log_info( + "SANA-WM refiner start: latent=%s, fps=%.3f, seed=%d, " + "sink_size=%d, sigmas=%s, diagnostics=%s", + tuple(z.shape), + fps, + seed, + sink_size, + STAGE_2_DISTILLED_SIGMA_VALUES, + "on" if sana_wm_diagnostics_enabled() else "off", + ) + log_sana_wm_tensor_stats("refiner.input_latent", z) + + prompt_embeds, prompt_attention_mask = self._encode_prompt(prompt, device) + + sigmas = torch.tensor( + STAGE_2_DISTILLED_SIGMA_VALUES, dtype=torch.float32, device=device + ) + start_sigma = float(sigmas[0]) + sink = z[:, :, :sink_size].contiguous() + current = z[:, :, sink_size:].contiguous() + log_sana_wm_tensor_stats("refiner.sink_latent", sink) + log_sana_wm_tensor_stats("refiner.current_latent_clean", current) + gen = torch.Generator(device=device).manual_seed(int(seed)) + eps = torch.randn(current.shape, generator=gen, device=device, dtype=self.dtype) + noisy = (1.0 - start_sigma) * current + start_sigma * eps + log_sana_wm_tensor_stats("refiner.current_latent_noisy_initial", noisy) + + patch_size = int(_refiner_config_value(self.transformer, "patch_size")) + patch_size_t = int(_refiner_config_value(self.transformer, "patch_size_t")) + + sink_tokens = pack_latents(sink, patch_size, patch_size_t) + n_context_tokens = sink_tokens.shape[1] + + for step_idx in range(len(sigmas) - 1): + sigma = sigmas[step_idx] + denoised = self._predict_current_x0( + sink=sink, + noisy_current=noisy, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + sigma=sigma, + fps=fps, + n_context_tokens=n_context_tokens, + step_idx=step_idx, + ) + noisy_tokens = pack_latents(noisy, patch_size, patch_size_t) + velocity_tokens = (noisy_tokens.float() - denoised.float()) / sigma.float() + next_tokens = ( + noisy_tokens.float() + + velocity_tokens * (sigmas[step_idx + 1] - sigma).float() + ) + noisy = unpack_latents( + next_tokens.to(self.dtype), + num_frames=noisy.shape[2], + height=noisy.shape[3], + width=noisy.shape[4], + patch_size=patch_size, + patch_size_t=patch_size_t, + ) + velocity_5d = unpack_latents( + velocity_tokens, + num_frames=noisy.shape[2], + height=noisy.shape[3], + width=noisy.shape[4], + patch_size=patch_size, + patch_size_t=patch_size_t, + ) + log_sana_wm_tensor_stats( + f"refiner.step_{step_idx}.velocity_current", + velocity_5d.to(self.dtype), + ) + log_sana_wm_tensor_stats(f"refiner.step_{step_idx}.current_latent", noisy) + + refined = torch.cat([sink, noisy], dim=2) + log_sana_wm_tensor_stats("refiner.output_latent", refined) + return refined + + @torch.inference_mode() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if batch.latents is None: + raise ValueError("SANA-WM refiner requires batch.latents from stage 1.") + if batch.latents.ndim != 5: + raise ValueError( + "SANA-WM refiner expects 5D latents shaped (B, C, T, H, W), " + f"got {tuple(batch.latents.shape)}." + ) + + if sana_wm_skip_refiner_enabled(batch): + if batch.extra is None: + batch.extra = {} + batch.extra["sana_wm_refiner_applied"] = False + self.log_info( + "SANA-WM LTX-2 refiner skipped by SGLANG_SANA_WM_SKIP_REFINER." + ) + return batch + + batch_size = int(batch.latents.shape[0]) + prompts = self._prompts_for_batch(batch, batch_size) + fps = float(getattr(batch, "fps", 16) or 16) + + seeds: list[int] + if batch.seeds is not None and len(batch.seeds) == batch_size: + seeds = [int(s) for s in batch.seeds] + elif batch.seeds is not None and len(batch.seeds) == 1: + seeds = [int(batch.seeds[0])] * batch_size + else: + seeds = [int(getattr(batch, "seed", 0) or 0)] * batch_size + + refined: list[torch.Tensor] = [] + for idx, (prompt, seed) in enumerate(zip(prompts, seeds, strict=True)): + refined.append( + self._refine_one( + batch.latents[idx : idx + 1], + prompt, + fps=fps, + seed=seed, + ) + ) + batch.latents = torch.cat(refined, dim=0).to( + device=batch.latents.device, dtype=batch.latents.dtype + ) + if batch.extra is None: + batch.extra = {} + batch.extra["sana_wm_refiner_applied"] = True + self.log_info("SANA-WM LTX-2 refiner applied to stage-1 latents.") + return batch + + +class SanaWMRefinerDecodingStage(SanaWMDecodingStage): + """Decode refined latents and drop the clean sink anchor frame.""" + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs): + self._drop_refiner_sink = bool( + (getattr(batch, "extra", None) or {}).get("sana_wm_refiner_applied", True) + ) + try: + return super().forward(batch, server_args) + finally: + self._drop_refiner_sink = True + + @torch.no_grad() + def decode( + self, + latents: torch.Tensor, + server_args: ServerArgs, + *, + vae_dtype: torch.dtype, + ) -> torch.Tensor: + frames = super().decode(latents, server_args, vae_dtype=vae_dtype) + log_sana_wm_tensor_stats("refiner.decode.frames_with_sink", frames) + if frames.ndim != 5: + raise ValueError( + "SANA-WM refiner decoding expects decoded video shaped " + f"(B, C, T, H, W), got {tuple(frames.shape)}." + ) + if frames.shape[2] <= 1: + raise ValueError( + "SANA-WM refiner decoding expected a sink frame plus refined " + f"frames, got temporal length {frames.shape[2]}." + ) + if not getattr(self, "_drop_refiner_sink", True): + log_sana_wm_tensor_stats("refiner.decode.frames_output", frames) + return frames + # Match NVlabs `inference_sana_wm.py`: decode with the clean sink anchor, + # then drop the first frame from the returned video. + frames = frames[:, :, 1:].contiguous() + log_sana_wm_tensor_stats("refiner.decode.frames_output", frames) + return frames diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/self_forcing.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/self_forcing.py new file mode 100644 index 000000000..f7b7099fa --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/self_forcing.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SANA-WM self-forcing chunk SAMPLER utilities (not a diffusers scheduler). + +Owns the autoregressive chunk grid, per-block KV-cache accumulation / +eviction, and the explicit per-chunk sigma list for streaming (self-forcing) +SANA-WM generation. The actual stepping scheduler is the shared +``FlowMatchEulerDiscreteScheduler`` (``per_token_timesteps`` path). + +Why not ``SelfForcingFlowMatchScheduler`` (the LingBot causal-DMD scheduler): + + * SANA-WM's distilled sigma grid is an explicit NON-uniform list + ((1000, 960, 889, 727, 0)/1000); its ``set_timesteps`` only expresses + ``linspace(sigma_max -> sigma_min)`` + shift. + * SANA-WM pins the condition frame at timestep 0 INSIDE chunk 0 (per-frame + timesteps within one step); its ``step`` applies a single per-sample + sigma. LingBot instead warms the cond frames into the KV cache, so a + scalar sigma suffices there. + +The model-calling denoise loop stays in the stage. + + * segmentation FRONT-LOADS the remainder into chunk 0; + * the KV accumulator concatenates softmax cache slots on **dim=1** (our + ``(B, N, H, D)`` softmax cache layout); + * GDN/STATE blocks copy-forward the previous chunk's recurrent state. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from sglang.multimodal_gen.runtime.models.dits.sana_wm_components import ( + _NUM_STREAM_CACHE_SLOTS, + _SLOT_CAM_K, + _SLOT_CAM_V, + _SLOT_FFN_TCONV, + _SLOT_K, + _SLOT_SHORTCONV, + _SLOT_TYPE_FLAG, + _SLOT_V, +) + + +@dataclass +class SanaWMSelfForcingSamplerConfig: + """Streaming self-forcing knobs (defaults mirror ``SanaWMPipelineConfig``).""" + + num_frame_per_block: int = 3 + num_cached_blocks: int = 2 + sink_token: bool = True + denoising_step_list: tuple = (1000, 960, 889, 727, 0) + streaming_cfg_scale: float = 1.0 + + @classmethod + def from_pipeline_config(cls, pcfg) -> "SanaWMSelfForcingSamplerConfig": + """Read the streaming knobs off a pipeline config (note the ``or 1.0`` cfg-scale guard).""" + return cls( + num_frame_per_block=int(getattr(pcfg, "num_frame_per_block", 3)), + num_cached_blocks=int(getattr(pcfg, "num_cached_blocks", 2)), + sink_token=bool(getattr(pcfg, "sink_token", True)), + denoising_step_list=tuple( + getattr(pcfg, "denoising_step_list", (1000, 960, 889, 727, 0)) + ), + streaming_cfg_scale=float(getattr(pcfg, "streaming_cfg_scale", 1.0) or 1.0), + ) + + +class SanaWMSelfForcingSampler: + """Self-forcing chunk scheduler: segmentation + per-block KV-cache carry. + + Stateless; methods are static (every input is passed explicitly). A + ``config`` may be attached for callers that prefer the instance form. + """ + + def __init__(self, config: SanaWMSelfForcingSamplerConfig | None = None): + self.config = config or SanaWMSelfForcingSamplerConfig() + + # ----------------------------------------------------------------- # + # Chunk schedule + KV cache accumulation + # ----------------------------------------------------------------- # + @staticmethod + def create_autoregressive_segments( + total_frames: int, num_frame_per_block: int + ) -> list[int]: + base = int(num_frame_per_block) + remained = total_frames % base + num_chunks = total_frames // base + chunk_indices = [0] + for idx in range(num_chunks): + cur = chunk_indices[-1] + base + (remained if idx == 0 else 0) + chunk_indices.append(cur) + return chunk_indices + + @staticmethod + def accumulate_kv_cache( + kv_cache: list, + chunk_idx: int, + chunk_indices: list[int], + num_cached_blocks: int, + sink_token: bool, + num_blocks: int, + ) -> tuple[list, int]: + """Build chunk ``chunk_idx``'s read-only KV prefix from prior chunks. + + GDN/STATE blocks (type flag > 0.5) copy-forward the PREVIOUS chunk's + recurrent state; softmax/CONCAT blocks concatenate the rolling-window + + sink K/V along **dim=1** (token axis of our (B,N,H,D) softmax cache).""" + if chunk_idx == 0: + return kv_cache[0], 0 + + cur = kv_cache[chunk_idx] + start_chunk = ( + max(chunk_idx - num_cached_blocks, 0) if num_cached_blocks > 0 else 0 + ) + valid = list(range(start_chunk, chunk_idx)) + sink_num = 0 + if sink_token and num_cached_blocks > 0: + sink_start = max(chunk_idx - num_cached_blocks + 1, 0) + if sink_start > 0: + valid = [0] + list(range(sink_start, chunk_idx)) + sink_num = chunk_indices[1] - chunk_indices[0] + + for block_id in range(num_blocks): + prev_last = kv_cache[chunk_idx - 1][block_id] + type_flag = prev_last[_SLOT_TYPE_FLAG] + if type_flag is not None and float(type_flag.item()) > 0.5: + # STATE (GDN) block: carry the previous chunk's recurrent state. + cur[block_id] = [ + prev_last[_SLOT_K], + prev_last[_SLOT_V], + prev_last[_SLOT_CAM_K], + prev_last[_SLOT_CAM_V], + prev_last[_SLOT_SHORTCONV], + None, + prev_last[_SLOT_TYPE_FLAG], + None, + None, + prev_last[_SLOT_FFN_TCONV], + ] + continue + + # CONCAT (softmax) block: concat cached K/V over the valid window. + acc: list[torch.Tensor | None] = [None] * _NUM_STREAM_CACHE_SLOTS + for idx in valid: + prev = kv_cache[idx][block_id] + if prev[_SLOT_K] is None: + continue + for slot in (_SLOT_K, _SLOT_V, _SLOT_CAM_K, _SLOT_CAM_V): + if prev[slot] is None: + continue + acc[slot] = ( + prev[slot].clone() + if acc[slot] is None + else torch.cat( + [acc[slot], prev[slot]], dim=1 + ) # (B,N,H,D) token axis + ) + cur[block_id] = [ + acc[_SLOT_K], + acc[_SLOT_V], + acc[_SLOT_CAM_K], + acc[_SLOT_CAM_V], + prev_last[_SLOT_SHORTCONV], + None, + prev_last[_SLOT_TYPE_FLAG], + None, + None, + prev_last[_SLOT_FFN_TCONV], + ] + + SanaWMSelfForcingSampler.evict_stale_kv_cache( + kv_cache, chunk_idx, valid, num_cached_blocks, num_blocks + ) + return cur, sink_num + + @staticmethod + def evict_stale_kv_cache( + kv_cache: list, + chunk_idx: int, + valid: list[int], + num_cached_blocks: int, + num_blocks: int, + ) -> None: + if num_cached_blocks <= 0: + return + keep = set(valid) + keep.add(chunk_idx) + for stale in range(chunk_idx): + if stale in keep: + continue + kv_cache[stale] = [ + [None] * _NUM_STREAM_CACHE_SLOTS for _ in range(num_blocks) + ] + + # ----------------------------------------------------------------- # + # Per-chunk flow-Euler sigma schedule + # ----------------------------------------------------------------- # + @staticmethod + def build_per_chunk_sigmas(denoising_step_list) -> list[float]: + """Explicit flow-Euler sigmas for one chunk's short self-forcing schedule. + + The ``denoising_step_list`` (e.g. (1000, 960, 889, 727, 0)) must end with + 0; sigmas are the non-terminal steps divided by 1000.""" + schedule = list(denoising_step_list) + if len(schedule) < 2 or schedule[-1] != 0: + raise ValueError(f"denoising_step_list must end with 0, got {schedule}") + return [float(t) / 1000.0 for t in schedule[:-1]] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py new file mode 100644 index 000000000..085795dad --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py @@ -0,0 +1,786 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SANA-WM streaming (self-forcing) denoising stage — S1c. + +Generates video chunk-by-chunk via the DiT's chunk-causal ``forward_long``, a +rolling per-block KV cache carrying recurrent GDN state + softmax K/V window. + +Gotchas vs the dense ``SanaWMDenoisingStage`` (one-shot bidirectional): + * fresh ``FlowMatchEulerDiscreteScheduler(shift=1.0)`` with explicit sigmas + (NOT the request scheduler, whose shift would warp the sigmas); + * KV accumulator concats softmax cache slots on **dim=1** (our + ``(B, N, H, D)`` softmax cache layout), not the reference's dim=2; + * the condition frame (frame 0) is re-pinned every step in chunk 0. +""" + +from __future__ import annotations + +# --- debug parity harness (gated by env; no-op in production) --- +import os as _os +import time +from contextlib import nullcontext +from types import SimpleNamespace + +import torch + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + ComponentUse, +) +from sglang.multimodal_gen.runtime.models.dits.sana_wm import ( + _NUM_STREAM_CACHE_SLOTS, +) +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import ( + CausalDMDDenoisingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.self_forcing import ( + SanaWMSelfForcingSampler, + SanaWMSelfForcingSamplerConfig, +) +from sglang.multimodal_gen.runtime.realtime.causal_state import RealtimeCausalDiTState +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +from . import parity_probe +from .base import ( + _align_sana_wm_cfg_text_conditions, + _cat_optional_tensors, + _first_tensor, + _to_device_dtype, + log_sana_wm_tensor_stats, +) + +_SANAWM_INJECT_DIR = _os.environ.get(parity_probe.ENV_INJECT) + + +def self_forcing_denoise_chunk( + *, + transformer, + scheduler, + sigmas, + get_chunk, + set_chunk, + cond_mask: torch.Tensor, + embeds: torch.Tensor, + mask, + camera_conditions, + chunk_plucker, + chunk_kv, + start_f: int, + end_f: int, + frame_index, + do_cfg: bool, + cfg_scale: float, + target_dtype: torch.dtype, + device, + forward_ctx=None, + on_noise_pred=None, +) -> list: + """Self-forcing denoise of ONE chunk + the clean t=0 KV pass; returns the chunk's updated per-block KV cache. + + SINGLE parity-locked implementation shared by the offline and realtime paths + (it used to be copy-pasted — the same drift failure mode as the four parity bugs). + + ``get_chunk``/``set_chunk`` are caller closures so each path keeps its EXACT + tensor lifecycle (kernel-visible layouts, and therefore bitwise behavior, + unchanged); ``set_chunk`` owns the dtype cast and cond-frame re-pinning. + ``forward_ctx`` (optional) maps ``int timestep -> context manager`` (batch + wraps forward_long in set_forward_context; realtime does not). + """ + ctx = forward_ctx if forward_ctx is not None else (lambda _ts: nullcontext()) + + scheduler.set_timesteps(sigmas=sigmas, device=device) + for t in scheduler.timesteps: + chunk_lat = get_chunk() + B, C = chunk_lat.shape[0], chunk_lat.shape[1] + lat_in = torch.cat([chunk_lat, chunk_lat], dim=0) if do_cfg else chunk_lat + ts_tensor = (1.0 - cond_mask) * t.to(device=device, dtype=torch.float32).view( + 1, 1, 1, 1, 1 + ) + ts_in = torch.cat([ts_tensor, ts_tensor], dim=0) if do_cfg else ts_tensor + model_ts = ts_in[:, :1, :, 0, 0] # (B|2B, chunk_frames) + + with ctx(int(t.item()) if t.ndim == 0 else 0): + noise_pred, _ = transformer.forward_long( + hidden_states=lat_in.to(target_dtype), + encoder_hidden_states=embeds, + timestep=model_ts, + encoder_attention_mask=mask, + camera_conditions=camera_conditions, + chunk_plucker=chunk_plucker, + kv_cache=chunk_kv, + save_kv_cache=False, + start_f=start_f, + end_f=end_f, + frame_index=frame_index, + ) + if do_cfg: + noise_uncond, noise_text = noise_pred.chunk(2) + noise_pred = noise_uncond + cfg_scale * (noise_text - noise_uncond) + + if on_noise_pred is not None: + on_noise_pred(int(t.item()), noise_pred) + + denoised = scheduler.step( + -noise_pred.reshape(B, C, -1).transpose(1, 2), + t, + chunk_lat.reshape(B, C, -1).transpose(1, 2), + per_token_timesteps=ts_tensor.reshape(B, C, -1)[:, 0], + return_dict=False, + )[0] + set_chunk(denoised.transpose(1, 2).reshape(chunk_lat.shape)) + + # Clean pass (t=0) to write this chunk's KV for the next chunk. + chunk_lat = get_chunk() + lat_in = torch.cat([chunk_lat, chunk_lat], dim=0) if do_cfg else chunk_lat + ts_zero = torch.zeros( + lat_in.shape[0], 1, chunk_lat.shape[2], device=device, dtype=torch.float32 + ) + with ctx(0): + _, updated_cache = transformer.forward_long( + hidden_states=lat_in.to(target_dtype), + encoder_hidden_states=embeds, + timestep=ts_zero, + encoder_attention_mask=mask, + camera_conditions=camera_conditions, + chunk_plucker=chunk_plucker, + kv_cache=chunk_kv, + save_kv_cache=True, + start_f=start_f, + end_f=end_f, + frame_index=frame_index, + ) + return updated_cache + + +class SanaWMStreamCacheState(RealtimeCausalDiTState): + """Per-session streaming DiT state, framework-pattern (cf. LingBot). + + The Wan-shaped ``kv_cache`` field stays None — SANA-WM's cache is the + heterogeneous per-block 10-slot list (GDN recurrent matrix states + softmax + concat windows + conv tails), carried in the fields below.""" + + def __init__(self): + super().__init__() + # kv[chunk][block][slot] — grown per chunk, stale chunks evicted. + self.stream_kv_cache: list = [] + self.chunk_indices: list[int] = [0] + # Growing stage-1 latent buffer (cond frame + denoised chunks). + self.latents: torch.Tensor | None = None + # Per-session flow-Euler scheduler (fresh shift=1.0; persists across ticks). + self.scheduler: FlowMatchEulerDiscreteScheduler | None = None + + def dispose(self) -> None: + super().dispose() + self.stream_kv_cache = [] + self.chunk_indices = [0] + self.latents = None + self.scheduler = None + + +class SanaWMStreamingDenoisingStage(CausalDMDDenoisingStage): + """Autoregressive self-forcing streaming denoise — SANA-WM's causal-DMD variant. + + Same family as ``LingBotWorldCausalDMDDenoisingStage`` (offline whole-clip vs + realtime per-chunk). Differences from the Wan-style base are the cache contract + (heterogeneous 10-slot list instead of preallocated ``CausalSelfAttentionKVCache``), + the front-loaded chunk grid (chunk 0 carries the remainder), and in-chunk + condition-frame pinning instead of KV warm-up — hence the cache-management overrides. + """ + + def __init__( + self, transformer, scheduler=None, *, keep_resident: bool = False + ) -> None: + # Skip CausalDMDDenoisingStage.__init__: it reads Wan-specific arch + # fields (num_frames_per_block / sliding_window_num_frames / sink_size) + # that SANA-WM defines per-run via the pipeline config instead. + DenoisingStage.__init__(self, transformer, scheduler) + self.num_transformer_blocks = len(transformer.blocks) + # Realtime pipelines keep the DiT device-resident for the session's + # lifetime (the per-tick offload round-trip would dominate latency); + # the offline pipeline keeps the default offload behavior. + self._keep_resident = bool(keep_resident) + + def component_uses(self, server_args: ServerArgs, stage_name: str | None = None): + if not self._keep_resident: + return super().component_uses(server_args, stage_name) + + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name, + "transformer", + target_dtype=PRECISION_TO_TYPE[ + server_args.pipeline_config.dit_precision + ], + memory_intensive=True, + keep_ready_after_warmup=True, + ), + ] + + # Chunk schedule + KV cache: delegate to SanaWMSelfForcingSampler. + @staticmethod + def _autoregressive_segments( + total_frames: int, num_frame_per_block: int + ) -> list[int]: + return SanaWMSelfForcingSampler.create_autoregressive_segments( + total_frames, num_frame_per_block + ) + + @staticmethod + def _accumulate_kv_cache( + kv_cache: list, + chunk_idx: int, + chunk_indices: list[int], + num_cached_blocks: int, + sink_token: bool, + num_blocks: int, + ) -> tuple[list, int]: + return SanaWMSelfForcingSampler.accumulate_kv_cache( + kv_cache, + chunk_idx, + chunk_indices, + num_cached_blocks, + sink_token, + num_blocks, + ) + + @staticmethod + def _evict_stale_kv_cache( + kv_cache: list, + chunk_idx: int, + valid: list[int], + num_cached_blocks: int, + num_blocks: int, + ) -> None: + SanaWMSelfForcingSampler.evict_stale_kv_cache( + kv_cache, chunk_idx, valid, num_cached_blocks, num_blocks + ) + + # Realtime per-chunk path (sessions): per-session state in SanaWMStreamCacheState. + @torch.no_grad() + def _forward_realtime_chunk(self, batch: Req, server_args: ServerArgs) -> Req: + if batch.latents is None or batch.latents.ndim != 5: + raise ValueError( + "SANA-WM realtime denoising expects this tick's pre-noised chunk " + "latents (B, C, n, H, W) from the latent-preparation stage." + ) + pcfg = server_args.pipeline_config + device = get_local_torch_device() + target_dtype = PRECISION_TO_TYPE.get( + getattr(pcfg, "dit_precision", "bf16"), torch.bfloat16 + ) + state = batch.session.get_or_create_state(SanaWMStreamCacheState) + if batch.block_idx == 0 and state.latents is not None: + state.dispose() # session restart on chunk 0 (mirrors the base stage) + + sc = self._resolve_stream_conditioning( + batch, server_args, device=device, target_dtype=target_dtype + ) + sampler_cfg = sc.sampler_cfg + incoming = batch.latents.to(device=device, dtype=target_dtype).clone() + plan = list(batch.extra.get("sana_wm_chunk_plan") or [incoming.shape[2]]) + if sum(plan) != incoming.shape[2]: + raise ValueError( + f"chunk plan {plan} does not cover the incoming {incoming.shape[2]} frames" + ) + if state.scheduler is None: + state.scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0) + + # Device-only move, NO dtype cast: Module.to(dtype=...) would cast the + # DiT's complex RoPE buffers to real, discarding the imaginary part + # (parity root cause #1). No use_declared_component round-trip either — + # the DiT stays device-resident for the session's lifetime. + transformer = self.transformer.to(device=device).eval() + num_blocks = len(transformer.blocks) + _dump_dir = parity_probe.probe_dir(parity_probe.ENV_RT_DUMP) + if _dump_dir and state.chunk_idx == 0: # parity harness + parity_probe.dump_tensor(_dump_dir, "cond_embeds", sc.embeds) + parity_probe.dump_tensor(_dump_dir, "cond_mask", sc.mask) + parity_probe.dump_obj( + _dump_dir, + "dit_fingerprint", + parity_probe.weights_fingerprint(transformer), + ) + + offset = 0 + for n in plan: + chunk_lat = incoming[:, :, offset : offset + n] + offset += n + chunk_idx = state.chunk_idx + # chunk 0's incoming includes the conditioning frame at index 0. + start_f = state.chunk_indices[-1] if chunk_idx > 0 else 0 + end_f = (start_f + n) if chunk_idx > 0 else n + state.chunk_indices.append(end_f) + state.stream_kv_cache.append( + [[None] * _NUM_STREAM_CACHE_SLOTS for _ in range(num_blocks)] + ) + chunk_kv, sink_num = self._accumulate_kv_cache( + state.stream_kv_cache, + chunk_idx, + state.chunk_indices, + sampler_cfg.num_cached_blocks, + sampler_cfg.sink_token, + num_blocks, + ) + # Evict entries outside the accumulate window (sink + last + # num_cached_blocks) — unbounded sessions leak concat K/V otherwise. + if chunk_idx > 0 and sampler_cfg.num_cached_blocks > 0: + start_chunk = max(chunk_idx - sampler_cfg.num_cached_blocks, 0) + valid = list(range(start_chunk, chunk_idx)) + if sampler_cfg.sink_token: + sink_start = max(chunk_idx - sampler_cfg.num_cached_blocks + 1, 0) + if sink_start > 0: + valid = [0] + list(range(sink_start, chunk_idx)) + self._evict_stale_kv_cache( + state.stream_kv_cache, + chunk_idx, + valid, + sampler_cfg.num_cached_blocks, + num_blocks, + ) + if _dump_dir: # parity harness + parity_probe.dump_obj( + _dump_dir, + f"kv_probe_{chunk_idx:03d}", + parity_probe.kv_cache_checksums(chunk_kv, sink_num), + ) + parity_probe.dump_tensor( + _dump_dir, f"cond_camera_{chunk_idx:03d}", sc.camera + ) + parity_probe.dump_tensor( + _dump_dir, f"cond_plucker_{chunk_idx:03d}", sc.plucker + ) + frame_index = ( + torch.arange(start_f, end_f, device=device, dtype=torch.long) + if sink_num > 0 + else None + ) + + B, C = chunk_lat.shape[0], chunk_lat.shape[1] + cond_mask = torch.zeros( + B, + C, + end_f - start_f, + *chunk_lat.shape[3:], + device=device, + dtype=torch.float32, + ) + cond_local = [] + if chunk_idx == 0: + cond_mask[:, :, 0] = 1.0 + cond_local = [0] + init_chunk = chunk_lat.clone() + _local = {"lat": chunk_lat} + + def _get_chunk(_local=_local): + return _local["lat"] + + def _set_chunk( + denoised, _local=_local, init_chunk=init_chunk, cond_local=cond_local + ): + lat = denoised.to(target_dtype) + for loc in cond_local: + lat[:, :, loc] = init_chunk[:, :, loc] + _local["lat"] = lat + + def _on_noise_pred(ts_int, noise_pred, chunk_idx=chunk_idx): + if _dump_dir and chunk_idx == 0: # parity harness + parity_probe.dump_tensor( + _dump_dir, f"noise_pred_c0_t{ts_int}", noise_pred + ) + + # Same forward context the mega-stage held across the tick (the + # per-chunk denoise core relies on an ambient context here; the + # offline path supplies per-step contexts instead). + with set_forward_context( + current_timestep=batch.block_idx, + attn_metadata=None, + forward_batch=batch, + ): + state.stream_kv_cache[chunk_idx] = self_forcing_denoise_chunk( + transformer=transformer, + scheduler=state.scheduler, + sigmas=sc.explicit_sigmas, + get_chunk=_get_chunk, + set_chunk=_set_chunk, + cond_mask=cond_mask, + embeds=sc.embeds, + mask=sc.mask, + camera_conditions=sc.camera, + chunk_plucker=sc.plucker, + chunk_kv=chunk_kv, + start_f=start_f, + end_f=end_f, + frame_index=frame_index, + do_cfg=sc.do_cfg, + cfg_scale=sc.cfg_scale, + target_dtype=target_dtype, + device=device, + on_noise_pred=_on_noise_pred, + ) + chunk_lat = _local["lat"] + state.latents = ( + chunk_lat + if state.latents is None + else torch.cat([state.latents, chunk_lat], dim=2) + ) + parity_probe.dump_tensor( # parity harness + _dump_dir, f"stage1_{chunk_idx:03d}_{start_f}_{end_f}", chunk_lat + ) + state.chunk_idx += 1 + state.current_chunk_start_frame = end_f + + # Downstream chain stages (refiner/decode) consume the growing buffer. + batch.latents = state.latents + return batch + + def _resolve_stream_conditioning( + self, + batch: Req, + server_args: ServerArgs, + *, + device: torch.device, + target_dtype: torch.dtype, + iload=None, + ): + """Resolve sampler config + text/camera conditioning (shared by the offline loop and realtime path).""" + pcfg = server_args.pipeline_config + sampler_cfg = SanaWMSelfForcingSamplerConfig.from_pipeline_config(pcfg) + explicit_sigmas = SanaWMSelfForcingSampler.build_per_chunk_sigmas( + sampler_cfg.denoising_step_list + ) + + # Streaming uses its OWN cfg scale (official StreamingGenerationConfig.cfg_scale=1.0 + # => no CFG on the distilled 4-step model). The general guidance_scale (e.g. 4.5) + # is for the dense path; using it here ran CFG=4.5 vs the reference's none. + cfg_scale = sampler_cfg.streaming_cfg_scale + do_cfg = bool(batch.do_classifier_free_guidance) and cfg_scale > 1.0 + if server_args.enable_cfg_parallel and do_cfg: + raise NotImplementedError( + "SANA-WM streaming does not support CFG parallel; run replicated." + ) + + # --- text conditioning --- + pos_embeds = _to_device_dtype( + _first_tensor(pcfg.get_pos_prompt_embeds(batch)), + device=device, + dtype=target_dtype, + ) + pos_mask = _to_device_dtype( + _first_tensor(batch.prompt_attention_mask), device=device + ) + if pos_embeds is None: + raise ValueError("SANA-WM streaming requires positive prompt embeds.") + neg_embeds = neg_mask = None + if do_cfg: + neg_embeds = _to_device_dtype( + _first_tensor(pcfg.get_neg_prompt_embeds(batch)), + device=device, + dtype=target_dtype, + ) + neg_mask = _to_device_dtype( + _first_tensor(batch.negative_attention_mask), device=device + ) + if neg_embeds is None: + raise ValueError( + "SANA-WM streaming CFG requires negative prompt embeds." + ) + pos_embeds, neg_embeds, pos_mask, neg_mask = ( + _align_sana_wm_cfg_text_conditions( + pos_embeds, neg_embeds, pos_mask, neg_mask + ) + ) + embeds_in = torch.cat([neg_embeds, pos_embeds], dim=0) if do_cfg else pos_embeds + mask_in = _cat_optional_tensors(neg_mask, pos_mask) if do_cfg else pos_mask + if _SANAWM_INJECT_DIR and iload is not None and not do_cfg: + _cond = iload("cond").to(device=device, dtype=target_dtype) + while _cond.dim() > embeds_in.dim(): + _cond = _cond.squeeze(1) + embeds_in = _cond + mask_in = iload("cond_mask").to(device=device) + + # --- camera / plücker (FULL length; forward_long windows internally) --- + extra = batch.extra or {} + camera_conditions = _to_device_dtype( + extra.get("camera_conditions"), device=device, dtype=target_dtype + ) + chunk_plucker = _to_device_dtype( + extra.get("chunk_plucker"), device=device, dtype=target_dtype + ) + cam_in = ( + torch.cat([camera_conditions, camera_conditions], dim=0) + if do_cfg and camera_conditions is not None + else camera_conditions + ) + plk_in = ( + torch.cat([chunk_plucker, chunk_plucker], dim=0) + if do_cfg and chunk_plucker is not None + else chunk_plucker + ) + if _SANAWM_INJECT_DIR and iload is not None: + cam_in = iload("raymap").to(device=device, dtype=target_dtype) + plk_in = iload("chunk_plucker").to(device=device, dtype=target_dtype) + if do_cfg: + cam_in = torch.cat([cam_in, cam_in], dim=0) + plk_in = torch.cat([plk_in, plk_in], dim=0) + + return SimpleNamespace( + sampler_cfg=sampler_cfg, + explicit_sigmas=explicit_sigmas, + cfg_scale=cfg_scale, + do_cfg=do_cfg, + embeds=embeds_in, + mask=mask_in, + camera=cam_in, + plucker=plk_in, + ) + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + # LingBot-style dispatch: realtime sessions denoise ONE chunk per call + # with per-session state; otherwise run the whole clip offline. + if batch.session is not None: + return self._forward_realtime_chunk(batch, server_args) + return self._forward_offline(batch, server_args) + + @torch.no_grad() + def _forward_offline(self, batch: Req, server_args: ServerArgs) -> Req: + if batch.latents is None or batch.latents.ndim != 5: + raise ValueError( + "SANA-WM streaming denoising expects 5D latents (B, C, T, H, W)." + ) + + pcfg = server_args.pipeline_config + device = get_local_torch_device() + target_dtype = PRECISION_TO_TYPE.get( + getattr(pcfg, "dit_precision", "bf16"), torch.bfloat16 + ) + + # .clone() detaches from the loader's InferenceMode tensor so the + # per-chunk in-place latent updates below are allowed. + latents = batch.latents.to(device=device, dtype=target_dtype).clone() + init_latents = latents.clone() + B, C, total_frames, H, W = latents.shape + + def _iload(_name): + return torch.load(f"{_SANAWM_INJECT_DIR}/{_name}.pt", map_location=device) + + _dump_dir = parity_probe.probe_dir(parity_probe.ENV_FORK_DUMP) + + def _fdump(_name, _t): + parity_probe.dump_tensor(_dump_dir, _name, _t) + + if ( + _SANAWM_INJECT_DIR + ): # parity harness: run the OFFICIAL's exact stage-1 inputs + latents = ( + _iload("z_full_initial").to(device=device, dtype=target_dtype).clone() + ) + init_latents = latents.clone() + B, C, total_frames, H, W = latents.shape + + _fdump( + "init_noise", init_latents + ) # parity harness: seeded pre-noise (cond @ frame 0) + + sc = self._resolve_stream_conditioning( + batch, server_args, device=device, target_dtype=target_dtype, iload=_iload + ) + sampler_cfg = sc.sampler_cfg + num_frame_per_block = sampler_cfg.num_frame_per_block + num_cached_blocks = sampler_cfg.num_cached_blocks + sink_token = sampler_cfg.sink_token + explicit_sigmas = sc.explicit_sigmas + cfg_scale = sc.cfg_scale + do_cfg = sc.do_cfg + embeds_in, mask_in, cam_in, plk_in = sc.embeds, sc.mask, sc.camera, sc.plucker + + # parity harness: full-length conditioning fed to forward_long (windowed + # internally per chunk via [start_f:end_f]). + _fdump("cond_embeds", embeds_in) + _fdump("cond_mask", mask_in) + _fdump("cond_camera", cam_in) + _fdump("cond_plucker", plk_in) + + scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0) + + chunk_indices = self._autoregressive_segments(total_frames, num_frame_per_block) + num_chunks = len(chunk_indices) - 1 + if num_chunks < 1: + raise ValueError( + f"streaming needs >= {num_frame_per_block} latent frames, got {total_frames}." + ) + + start_time = time.perf_counter() + with self.use_declared_component( + component_name="transformer", module=self.transformer + ) as transformer: + assert transformer is not None + self.transformer = transformer + num_blocks = len(transformer.blocks) + if _dump_dir: # parity harness: weights fingerprint + parity_probe.dump_obj( + _dump_dir, + "dit_fingerprint", + parity_probe.weights_fingerprint(transformer), + ) + kv_cache = [ + [[None] * _NUM_STREAM_CACHE_SLOTS for _ in range(num_blocks)] + for _ in range(num_chunks) + ] + + self.log_info( + "SANA-WM streaming denoise: latent=%s, chunks=%d (block=%d frames), " + "steps/chunk=%d, cfg=%s", + tuple(latents.shape), + num_chunks, + num_frame_per_block, + len(explicit_sigmas), + do_cfg, + ) + + for chunk_idx in self.progress_bar(range(num_chunks)): + chunk_kv, sink_num = self._accumulate_kv_cache( + kv_cache, + chunk_idx, + chunk_indices, + num_cached_blocks, + sink_token, + num_blocks, + ) + if _dump_dir: # parity harness: accumulated-KV checksums + parity_probe.dump_obj( + _dump_dir, + f"kv_probe_{chunk_idx:03d}", + parity_probe.kv_cache_checksums(chunk_kv, sink_num), + ) + start_f = chunk_indices[chunk_idx] + end_f = chunk_indices[chunk_idx + 1] + chunk_frames = end_f - start_f + frame_index = ( + torch.arange(start_f, end_f, device=device, dtype=torch.long) + if sink_num > 0 + else None + ) + + # Condition mask: frame 0 only (chunk 0). Re-pinned each step. + cond_mask = torch.zeros( + B, C, chunk_frames, H, W, device=device, dtype=torch.float32 + ) + cond_local = [] + if start_f == 0: + cond_mask[:, :, 0] = 1.0 + cond_local = [0] + + # Buffer-view closures: the shared loop reads/writes views of + # the full-latent buffer (kernel-visible layouts unchanged). + def _get_chunk(start_f=start_f, end_f=end_f): + return latents[:, :, start_f:end_f] + + def _set_chunk( + denoised, start_f=start_f, end_f=end_f, cond_local=cond_local + ): + latents[:, :, start_f:end_f] = denoised.to(latents.dtype) + for loc in cond_local: + latents[:, :, start_f + loc] = init_latents[:, :, start_f + loc] + + def _forward_ctx(ts_int): + return set_forward_context( + current_timestep=ts_int, + attn_metadata=None, + forward_batch=batch, + ) + + def _on_noise_pred(ts_int, noise_pred, chunk_idx=chunk_idx): + if chunk_idx == 0: # parity harness: per-step model output + _fdump(f"noise_pred_c0_t{ts_int}", noise_pred) + + kv_cache[chunk_idx] = self_forcing_denoise_chunk( + transformer=transformer, + scheduler=scheduler, + sigmas=explicit_sigmas, + get_chunk=_get_chunk, + set_chunk=_set_chunk, + cond_mask=cond_mask, + embeds=embeds_in, + mask=mask_in, + camera_conditions=cam_in, + chunk_plucker=plk_in, + chunk_kv=chunk_kv, + start_f=start_f, + end_f=end_f, + frame_index=frame_index, + do_cfg=do_cfg, + cfg_scale=cfg_scale, + target_dtype=target_dtype, + device=device, + forward_ctx=_forward_ctx, + on_noise_pred=_on_noise_pred, + ) + _fdump( + f"stage1_{chunk_idx:03d}_{start_f}_{end_f}", + latents[:, :, start_f:end_f], + ) + + log_sana_wm_tensor_stats("stream.output_latents", latents) + self.log_info( + "SANA-WM streaming denoise finished in %.4f s; first_frame_max_delta=%.6g", + time.perf_counter() - start_time, + float((latents[:, :, :1] - init_latents[:, :, :1]).abs().max().item()), + ) + batch.latents = pcfg.post_denoising_loop(latents, batch) + return batch + + +class SanaWMStreamingDecodingStage(DecodingStage): + """Streaming causal-VAE decode over the SAME autoregressive grid the denoise stage used. + + Carries a per-conv decoder cache across chunks so the causal LTX-2 VAE produces + seam-free frames (the `decode_per_frame_with_cache` equivalent at chunk granularity). + Subclasses DecodingStage directly (NOT SanaWMDecodingStage, whose long-video config + re-enables the stateless tiled decode). + """ + + @torch.no_grad() + def decode( + self, + latents: torch.Tensor, + server_args: ServerArgs, + *, + vae_dtype: torch.dtype, + ) -> torch.Tensor: + if not hasattr(self.vae, "decode_chunk"): + raise ValueError( + "SANA-WM streaming decode requires AutoencoderKLCausalLTX2Video " + "(decode_chunk). Point --component_paths.vae at the ltx2_causal_vae " + "weights when streaming." + ) + device = get_local_torch_device() + latents = latents.to(device) + pcfg = server_args.pipeline_config + num_frame_per_block = int(getattr(pcfg, "num_frame_per_block", 3)) + total_frames = latents.shape[2] + segments = SanaWMStreamingDenoisingStage._autoregressive_segments( + total_frames, num_frame_per_block + ) + conv_cache = self.vae.reset_decoder_cache() + chunks = [] + for i in range(len(segments) - 1): + s, e = segments[i], segments[i + 1] + z = self.scale_and_shift(latents[:, :, s:e].to(vae_dtype), server_args) + # No autocast: match the official decode (VAE already runs in vae_dtype, + # z is cast once above) for consistent rounding across chunk boundaries. + pixel = self.vae.decode_chunk(z, conv_cache) + pixel = (pixel / 2 + 0.5).clamp(0, 1) + chunks.append(pixel.float().cpu()) + + frames = torch.cat(chunks, dim=2) + log_sana_wm_tensor_stats("stream.decode.frames", frames) + return frames diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming_refiner.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming_refiner.py new file mode 100644 index 000000000..6325f0dd7 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming_refiner.py @@ -0,0 +1,726 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SANA-WM chunked streaming LTX-2 refiner — S2b. + +Runs the refiner block-by-block carrying a refiner KV cache: a ``sink`` prefix +(captured pre-RoPE, re-RoPE'd per block at shifted positions) plus a sliding +``history`` of refined-block K/V (post-RoPE), injected into each ``attn1`` as a +KV prefix. Ports the reference NVlabs ``RefinerChunkRunner`` +(minimal-sanawm/refiner.py), operating on mg's already-loaded diffusers LTX-2 +refiner transformer. +""" + +from __future__ import annotations + +import math +import time +from typing import Any + +import torch +from torch import nn + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.models.dits.sana_wm_refiner_transformer import ( + pack_latents, + unpack_latents, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.server_args import ServerArgs + +from . import parity_probe +from .refiner import ( + STAGE_2_DISTILLED_SIGMA_VALUES, + SanaWMLTX2RefinerStage, + _as_additive_attention_mask, + _unwrap_diffusers_ltx2_refiner, + log_sana_wm_tensor_stats, + sana_wm_skip_refiner_enabled, +) + + +# --------------------------------------------------------------------------- # +# Per-attn1 KV-prefix / capture hooks (port of refiner.py:753-788) +# --------------------------------------------------------------------------- # +def set_kv_prefix_on_blocks(transformer: nn.Module, kv_prefix_per_layer) -> None: + if kv_prefix_per_layer is None: + clear_kv_prefix_on_blocks(transformer) + return + for block, prefix in zip(transformer.transformer_blocks, kv_prefix_per_layer): + block.attn1._tf_kv_prefix = prefix + + +def clear_kv_prefix_on_blocks(transformer: nn.Module) -> None: + for block in transformer.transformer_blocks: + block.attn1._tf_kv_prefix = None + + +def set_capture_flag_on_blocks( + transformer: nn.Module, mode: str, *, enable: bool +) -> None: + if mode == "pre_rope": + attr, clear_attr = "_kv_cache_capture", "_cached_kv_pre" + elif mode == "post_rope": + attr, clear_attr = "_tf_capture_kv", "_cached_kv_post" + else: + raise ValueError(f"unsupported capture mode: {mode}") + for block in transformer.transformer_blocks: + setattr(block.attn1, attr, bool(enable)) + if enable and hasattr(block.attn1, clear_attr): + setattr(block.attn1, clear_attr, None) + + +def collect_captured_kv_from_blocks(transformer: nn.Module, mode: str): + attr = "_cached_kv_pre" if mode == "pre_rope" else "_cached_kv_post" + out = [] + for block in transformer.transformer_blocks: + cached = getattr(block.attn1, attr, None) + if cached is None: + raise RuntimeError(f"missing captured KV on {attr}") + out.append(cached) + setattr(block.attn1, attr, None) + return out + + +# --------------------------------------------------------------------------- # +# Absolute-position RoPE (port of refiner.py:721-750) +# --------------------------------------------------------------------------- # +def build_rotary_emb_for_absolute_positions( + *, transformer, batch_size, frame_positions, height, width, device, fps +): + rope = transformer.rope + patch_size_t = int(rope.patch_size_t) + patch_size = int(rope.patch_size) + f_positions = torch.tensor(frame_positions, dtype=torch.float32, device=device) + if patch_size_t > 1: + f_positions = f_positions[::patch_size_t] + grid_h = torch.arange(0, height, patch_size, dtype=torch.float32, device=device) + grid_w = torch.arange(0, width, patch_size, dtype=torch.float32, device=device) + grid = torch.meshgrid(f_positions, grid_h, grid_w, indexing="ij") + grid = torch.stack(grid, dim=0) + patch_delta = torch.tensor( + (patch_size_t, patch_size, patch_size), dtype=grid.dtype, device=device + ) + patch_ends = grid + patch_delta.view(3, 1, 1, 1) + latent_coords = ( + torch.stack([grid, patch_ends], dim=-1) + .flatten(1, 3) + .unsqueeze(0) + .repeat(batch_size, 1, 1, 1) + ) + scale = torch.tensor(rope.scale_factors, device=device) + broadcast_shape = [1] * latent_coords.ndim + broadcast_shape[1] = -1 + pixel_coords = latent_coords * scale.view(*broadcast_shape) + pixel_coords[:, 0, ...] = ( + pixel_coords[:, 0, ...] + rope.causal_offset - rope.scale_factors[0] + ).clamp(min=0) + pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / float(fps) + return rope(pixel_coords, device=device) + + +# --------------------------------------------------------------------------- # +# Self-attention with KV-prefix injection + capture (port of refiner.py:464-576) +# --------------------------------------------------------------------------- # +def streaming_self_attention( + *, attn, hidden_states, query_rotary_emb, n_context_tokens +): + sequence_length = hidden_states.shape[1] + has_streaming_hooks = ( + getattr(attn, "_kv_cache_capture", False) + or getattr(attn, "_tf_capture_kv", False) + or getattr(attn, "_tf_kv_prefix", None) is not None + ) + if n_context_tokens >= sequence_length and not has_streaming_hooks: + return attn( + hidden_states=hidden_states, + encoder_hidden_states=None, + query_rotary_emb=query_rotary_emb, + ) + + from diffusers.models.attention_dispatch import dispatch_attention_fn + from diffusers.models.transformers.transformer_ltx2 import ( + apply_interleaved_rotary_emb, + apply_split_rotary_emb, + ) + + # diffusers 0.38+ always defines `to_gate_logits`; 0.37 only on gated variants + # (the SANA-WM refiner is ungated) -> getattr so 0.37 works too. + _to_gate_logits = getattr(attn, "to_gate_logits", None) + gate_logits = ( + _to_gate_logits(hidden_states) if _to_gate_logits is not None else None + ) + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + query = attn.norm_q(query) + key = attn.norm_k(key) + if getattr(attn, "_kv_cache_capture", False): + attn._cached_kv_pre = (key.detach().clone(), value.detach().clone()) + + if attn.rope_type == "interleaved": + query = apply_interleaved_rotary_emb(query, query_rotary_emb) + key = apply_interleaved_rotary_emb(key, query_rotary_emb) + elif attn.rope_type == "split": + query = apply_split_rotary_emb(query, query_rotary_emb) + key = apply_split_rotary_emb(key, query_rotary_emb) + else: + raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}") + if getattr(attn, "_tf_capture_kv", False): + attn._cached_kv_post = (key.detach().clone(), value.detach().clone()) + + tf_prefix = getattr(attn, "_tf_kv_prefix", None) + if isinstance(tf_prefix, dict) and tf_prefix.get("mode") == "rf_shifted_sink": + prefix_k_parts = [] + prefix_v_parts = [] + sink_k_pre = tf_prefix.get("sink_k_pre") + sink_v = tf_prefix.get("sink_v") + if sink_k_pre is not None and sink_v is not None and sink_k_pre.shape[1] > 0: + sink_pe = tf_prefix.get("sink_pe") + if sink_pe is None: + raise RuntimeError("rf_shifted_sink prefix requires sink_pe") + if attn.rope_type == "interleaved": + sink_k = apply_interleaved_rotary_emb(sink_k_pre.to(key.dtype), sink_pe) + else: + sink_k = apply_split_rotary_emb(sink_k_pre.to(key.dtype), sink_pe) + prefix_k_parts.append(sink_k) + prefix_v_parts.append(sink_v.to(value.dtype)) + history_k = tf_prefix.get("history_k") + history_v = tf_prefix.get("history_v") + if history_k is not None and history_v is not None and history_k.shape[1] > 0: + prefix_k_parts.append(history_k.to(key.dtype)) + prefix_v_parts.append(history_v.to(value.dtype)) + if prefix_k_parts: + key = torch.cat([*prefix_k_parts, key], dim=1) + value = torch.cat([*prefix_v_parts, value], dim=1) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + processor = attn.processor + backend = getattr(processor, "_attention_backend", None) + parallel_config = getattr(processor, "_parallel_config", None) + if n_context_tokens <= 0 or n_context_tokens >= query.shape[1]: + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=backend, + parallel_config=parallel_config, + ) + else: + context = dispatch_attention_fn( + query[:, :n_context_tokens], + key[:, :n_context_tokens], + value[:, :n_context_tokens], + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=backend, + parallel_config=parallel_config, + ) + current = dispatch_attention_fn( + query[:, n_context_tokens:], + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=backend, + parallel_config=parallel_config, + ) + hidden_states = torch.cat([context, current], dim=1) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + if gate_logits is not None: + hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) + hidden_states = hidden_states * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = attn.to_out[0](hidden_states) + return attn.to_out[1](hidden_states) + + +def forward_video_block( + *, + block, + hidden_states, + encoder_hidden_states, + temb, + video_rotary_emb, + encoder_attention_mask, + n_context_tokens, +): + batch = hidden_states.size(0) + norm_hidden_states = block.norm1(hidden_states) + num_ada_params = block.scale_shift_table.shape[0] + ada_values = block.scale_shift_table[None, None].to(temb.device) + temb.reshape( + batch, temb.size(1), num_ada_params, -1 + ) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind( + dim=2 + ) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + attn_hidden_states = streaming_self_attention( + attn=block.attn1, + hidden_states=norm_hidden_states, + query_rotary_emb=video_rotary_emb, + n_context_tokens=n_context_tokens, + ) + hidden_states = hidden_states + attn_hidden_states * gate_msa + norm_hidden_states = block.norm2(hidden_states) + attn_hidden_states = block.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + query_rotary_emb=None, + attention_mask=encoder_attention_mask, + ) + hidden_states = hidden_states + attn_hidden_states + norm_hidden_states = block.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp + return hidden_states + block.ff(norm_hidden_states) * gate_mlp + + +class _RefinerCore: + """Adapts the unwrapped diffusers refiner transformer to the + DiffusersLTX2Refiner interface RefinerChunkRunner expects.""" + + def __init__( + self, transformer: nn.Module, device: torch.device, dtype: torch.dtype + ): + self.transformer = transformer + self.device = device + self.dtype = dtype + + def _forward_video_only_with_rope( + self, + *, + hidden_states, + encoder_hidden_states, + timestep, + encoder_attention_mask, + video_rotary_emb, + n_context_tokens, + ): + transformer = self.transformer + batch = hidden_states.size(0) + if encoder_attention_mask is not None: + encoder_attention_mask = _as_additive_attention_mask( + encoder_attention_mask, hidden_states.dtype + ) + hidden_states = transformer.proj_in(hidden_states) + temb, embedded_timestep = transformer.time_embed( + timestep.flatten(), + batch_size=batch, + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(batch, -1, temb.size(-1)) + embedded_timestep = embedded_timestep.view( + batch, -1, embedded_timestep.size(-1) + ) + encoder_hidden_states = transformer.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view( + batch, -1, hidden_states.size(-1) + ) + for block in transformer.transformer_blocks: + hidden_states = forward_video_block( + block=block, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + video_rotary_emb=video_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + n_context_tokens=n_context_tokens, + ) + scale_shift = ( + transformer.scale_shift_table[None, None] + embedded_timestep[:, :, None] + ) + shift, scale = scale_shift[:, :, 0], scale_shift[:, :, 1] + hidden_states = transformer.norm_out(hidden_states) + hidden_states = hidden_states * (1 + scale) + shift + return transformer.proj_out(hidden_states) + + def _predict_x0_active_block( + self, + *, + active, + active_positions, + sigma_cur, + prompt_embeds, + prompt_attention_mask, + fps, + kv_prefix_per_layer, + ): + ps = int(self.transformer.config.patch_size) + pst = int(self.transformer.config.patch_size_t) + latent_tokens = pack_latents(active, ps, pst) + batch, seq_len, _ = latent_tokens.shape + timestep = torch.full( + (batch, seq_len), + float(sigma_cur) * float(self.transformer.config.timestep_scale_multiplier), + dtype=torch.float32, + device=self.device, + ) + video_rotary_emb = build_rotary_emb_for_absolute_positions( + transformer=self.transformer, + batch_size=batch, + frame_positions=active_positions, + height=int(active.shape[3]), + width=int(active.shape[4]), + device=self.device, + fps=float(fps), + ) + set_kv_prefix_on_blocks(self.transformer, kv_prefix_per_layer) + try: + velocity = self._forward_video_only_with_rope( + hidden_states=latent_tokens, + encoder_hidden_states=prompt_embeds, + timestep=timestep, + encoder_attention_mask=prompt_attention_mask, + video_rotary_emb=video_rotary_emb, + n_context_tokens=0, + ) + finally: + clear_kv_prefix_on_blocks(self.transformer) + raw_sigma = torch.full( + (batch, seq_len, 1), + float(sigma_cur), + dtype=torch.float32, + device=self.device, + ) + denoised = latent_tokens.float() - velocity.float() * raw_sigma + return unpack_latents( + denoised.to(self.dtype), + num_frames=int(active.shape[2]), + height=int(active.shape[3]), + width=int(active.shape[4]), + patch_size=ps, + patch_size_t=pst, + ) + + def _capture_block_kv( + self, + *, + clean_block, + frame_positions, + prompt_embeds, + prompt_attention_mask, + fps, + capture_mode, + kv_prefix_per_layer, + ): + ps = int(self.transformer.config.patch_size) + pst = int(self.transformer.config.patch_size_t) + latent_tokens = pack_latents(clean_block, ps, pst) + batch, seq_len, _ = latent_tokens.shape + timestep = torch.zeros(batch, seq_len, dtype=torch.float32, device=self.device) + video_rotary_emb = build_rotary_emb_for_absolute_positions( + transformer=self.transformer, + batch_size=batch, + frame_positions=frame_positions, + height=int(clean_block.shape[3]), + width=int(clean_block.shape[4]), + device=self.device, + fps=float(fps), + ) + set_kv_prefix_on_blocks(self.transformer, kv_prefix_per_layer) + set_capture_flag_on_blocks(self.transformer, capture_mode, enable=True) + try: + _ = self._forward_video_only_with_rope( + hidden_states=latent_tokens, + encoder_hidden_states=prompt_embeds, + timestep=timestep, + encoder_attention_mask=prompt_attention_mask, + video_rotary_emb=video_rotary_emb, + n_context_tokens=0, + ) + finally: + set_capture_flag_on_blocks(self.transformer, capture_mode, enable=False) + clear_kv_prefix_on_blocks(self.transformer) + return collect_captured_kv_from_blocks(self.transformer, capture_mode) + + +class RefinerChunkRunner: + """Port of the reference RefinerChunkRunner (refiner.py:579-718).""" + + def __init__( + self, + refiner: _RefinerCore, + *, + prompt_embeds, + prompt_attention_mask, + fps, + sigmas, + source_sink_frames, + block_size, + kv_max_frames, + seed, + spatial_shape, + ): + self.refiner = refiner + self.prompt_embeds = prompt_embeds + self.prompt_attention_mask = prompt_attention_mask + self.fps = float(fps) + self.sigmas = sigmas + self.source_sink_frames = int(source_sink_frames) + self.block_size = int(block_size) + self.kv_max_frames = int(kv_max_frames) + self.max_history_frames = self.kv_max_frames - self.source_sink_frames + self.generator = torch.Generator(device=refiner.device).manual_seed(int(seed)) + self.device = refiner.device + self.dtype = refiner.dtype + self.height, self.width = int(spatial_shape[0]), int(spatial_shape[1]) + transformer = refiner.transformer + self.tokens_per_frame = ( + int(self.height // transformer.config.patch_size) + * int(self.width // transformer.config.patch_size) + * int(transformer.config.patch_size_t) + ) + self.sink_kv_pre = None + self.history_kv_post = [None] * len(transformer.transformer_blocks) + self.history_frames = 0 + + @torch.inference_mode() + def refine_block( + self, *, block_idx, clean_block, block_start, block_end, sink_seed_frames=None + ): + # parity harness (env-gated, no-op in prod): per-block input/config/output + # checksums; both the batch stage and the realtime stage run through here. + _probe_dir = parity_probe.probe_dir( + parity_probe.ENV_RT_DUMP, parity_probe.ENV_FORK_DUMP + ) + _probe = None + if _probe_dir: + _ck = parity_probe.checksum + _probe = { + "block_idx": int(block_idx), + "block_start": int(block_start), + "block_end": int(block_end), + "clean_block": _ck(clean_block), + "sink_seed_frames": _ck(sink_seed_frames), + "prompt_embeds": _ck(self.prompt_embeds), + "prompt_attention_mask": _ck(self.prompt_attention_mask), + "fps": self.fps, + "sigmas": [float(s) for s in self.sigmas], + "source_sink_frames": self.source_sink_frames, + "block_size": self.block_size, + "kv_max_frames": self.kv_max_frames, + "history_frames": self.history_frames, + "generator_state": float( + self.generator.get_state().double().sum().item() + ), + } + del block_idx + refiner = self.refiner + if block_start < self.source_sink_frames: + raise ValueError("refiner block overlaps source sink") + if self.sink_kv_pre is None: + if self.source_sink_frames == 0: + self.sink_kv_pre = [(None, None) for _ in self.history_kv_post] + elif sink_seed_frames is None: + raise ValueError("first refine_block call requires sink_seed_frames") + else: + self.sink_kv_pre = refiner._capture_block_kv( + clean_block=sink_seed_frames.contiguous(), + frame_positions=list(range(self.source_sink_frames)), + prompt_embeds=self.prompt_embeds, + prompt_attention_mask=self.prompt_attention_mask, + fps=self.fps, + capture_mode="pre_rope", + kv_prefix_per_layer=None, + ) + batch = int(clean_block.shape[0]) + sink_rope_offset = block_start - self.history_frames - self.source_sink_frames + sink_pe = None + if self.source_sink_frames > 0: + sink_pe = build_rotary_emb_for_absolute_positions( + transformer=refiner.transformer, + batch_size=batch, + frame_positions=list( + range(sink_rope_offset, sink_rope_offset + self.source_sink_frames) + ), + height=self.height, + width=self.width, + device=self.device, + fps=self.fps, + ) + kv_prefix_per_layer = [] + for layer_idx, sink_kv in enumerate(self.sink_kv_pre): + history = self.history_kv_post[layer_idx] + kv_prefix_per_layer.append( + { + "mode": "rf_shifted_sink", + "sink_k_pre": sink_kv[0], + "sink_v": sink_kv[1], + "sink_pe": sink_pe, + "history_k": history[0] if history is not None else None, + "history_v": history[1] if history is not None else None, + } + ) + sigma0 = float(self.sigmas[0].item()) + eps = torch.randn( + clean_block.shape, + generator=self.generator, + device=self.device, + dtype=self.dtype, + ) + x_t = ((1.0 - sigma0) * clean_block.float() + sigma0 * eps.float()).to( + self.dtype + ) + active_positions = list(range(int(block_start), int(block_end))) + for level in range(int(self.sigmas.numel()) - 1): + sigma_cur = float(self.sigmas[level].item()) + sigma_next = float(self.sigmas[level + 1].item()) + pred_x0 = refiner._predict_x0_active_block( + active=x_t, + active_positions=active_positions, + sigma_cur=sigma_cur, + prompt_embeds=self.prompt_embeds, + prompt_attention_mask=self.prompt_attention_mask, + fps=self.fps, + kv_prefix_per_layer=kv_prefix_per_layer, + ) + if sigma_cur <= 1.0e-6: + x_t = pred_x0.to(self.dtype) + else: + ratio = sigma_next / sigma_cur + x_t = (ratio * x_t.float() + (1.0 - ratio) * pred_x0.float()).to( + self.dtype + ) + block_kv_post = refiner._capture_block_kv( + clean_block=x_t, + frame_positions=active_positions, + prompt_embeds=self.prompt_embeds, + prompt_attention_mask=self.prompt_attention_mask, + fps=self.fps, + capture_mode="post_rope", + kv_prefix_per_layer=kv_prefix_per_layer, + ) + for layer_idx, new_kv in enumerate(block_kv_post): + old = self.history_kv_post[layer_idx] + self.history_kv_post[layer_idx] = ( + new_kv + if old is None + else ( + torch.cat([old[0], new_kv[0]], dim=1), + torch.cat([old[1], new_kv[1]], dim=1), + ) + ) + self.history_frames += int(block_end - block_start) + if ( + self.max_history_frames > 0 + and self.history_frames > self.max_history_frames + ): + keep_tokens = self.max_history_frames * self.tokens_per_frame + for layer_idx, old in enumerate(self.history_kv_post): + if old is not None: + self.history_kv_post[layer_idx] = ( + old[0][:, -keep_tokens:], + old[1][:, -keep_tokens:], + ) + self.history_frames = self.max_history_frames + if _probe is not None: + _probe["refined"] = parity_probe.checksum(x_t) + parity_probe.dump_obj( + _probe_dir, f"refiner_probe_{_probe['block_idx']:03d}", _probe + ) + return x_t + + +class SanaWMStreamingRefinerStage(SanaWMLTX2RefinerStage): + """Chunked streaming LTX-2 refiner — refines stage-1 latents block-by-block + carrying a sink/history KV cache. Inherits loading/encode/residency from + SanaWMLTX2RefinerStage; only the refine loop changes.""" + + def __init__( + self, + *, + block_size: int = 3, + kv_max_frames: int = 11, + sink_size: int = 1, + seed: int = 42, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.block_size = int(block_size) + self.kv_max_frames = int(kv_max_frames) + self.sink_size = int(sink_size) + self.seed = int(seed) + + @torch.inference_mode() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if sana_wm_skip_refiner_enabled(): + return batch + if batch.latents is None or batch.latents.ndim != 5: + raise ValueError("SANA-WM streaming refiner expects 5D stage-1 latents.") + device = get_local_torch_device() + latents = batch.latents.to(device=device, dtype=self.dtype).clone() + B, C, T, H, W = latents.shape + prompt = self._prompts_for_batch(batch, B)[0] + fps = float((batch.extra or {}).get("fps", getattr(batch, "fps", 16.0)) or 16.0) + + prompt_embeds, prompt_mask = self._encode_prompt(prompt, device) + sigmas = torch.tensor( + STAGE_2_DISTILLED_SIGMA_VALUES, dtype=torch.float32, device=device + ) + + n_active = T - self.sink_size + if n_active <= 0: + self.log_info( + "SANA-WM streaming refiner: no active frames (T=%d <= sink=%d); skipping.", + T, + self.sink_size, + ) + return batch + n_blocks = math.ceil(n_active / self.block_size) + self.log_info( + "SANA-WM streaming refiner: latent=%s, sink=%d, block=%d, blocks=%d, kv_max=%d, seed=%d", + tuple(latents.shape), + self.sink_size, + self.block_size, + n_blocks, + self.kv_max_frames, + self.seed, + ) + + t0 = time.perf_counter() + with self.use_declared_component( + component_name="transformer_2", module=self.transformer + ) as transformer_mod: + self.transformer = transformer_mod + unwrapped = _unwrap_diffusers_ltx2_refiner(self.transformer) + core = _RefinerCore(unwrapped, device, self.dtype) + runner = RefinerChunkRunner( + core, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_mask, + fps=fps, + sigmas=sigmas, + source_sink_frames=self.sink_size, + block_size=self.block_size, + kv_max_frames=self.kv_max_frames, + seed=self.seed, + spatial_shape=(H, W), + ) + for i in range(n_blocks): + start_f = self.sink_size + i * self.block_size + end_f = min(start_f + self.block_size, T) + sink_seed = latents[:, :, : self.sink_size] if i == 0 else None + refined = runner.refine_block( + block_idx=i, + clean_block=latents[:, :, start_f:end_f].contiguous(), + block_start=start_f, + block_end=end_f, + sink_seed_frames=sink_seed, + ) + latents[:, :, start_f:end_f] = refined.to(latents.dtype) + + log_sana_wm_tensor_stats("stream.refiner.latents", latents) + self.log_info( + "SANA-WM streaming refiner applied (%d blocks) in %.4f s.", + n_blocks, + time.perf_counter() - t0, + ) + batch.latents = latents + if batch.extra is None: + batch.extra = {} + batch.extra["sana_wm_refiner_applied"] = True + return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_diffusion.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_diffusion.py new file mode 100644 index 000000000..df38eee2b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_diffusion.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + ComponentUse, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( + PipelineStage, + StageParallelismType, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( + scale_and_shift_latents, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + + +@dataclass(frozen=True) +class RealtimeStageComponent: + component_name: str + precision_attr: str | None = None + memory_intensive: bool = False + keep_ready_after_warmup: bool = True + + +class RealtimeDiffusionStage(PipelineStage): + """common contract for stateful realtime diffusion stages + + model-specific subclasses keep conditioning, latent scheduling, and model + forward semantics + """ + + component_specs: tuple[RealtimeStageComponent, ...] = ( + RealtimeStageComponent("transformer", "dit_precision", memory_intensive=True), + RealtimeStageComponent("vae", "vae_precision"), + ) + + def __init__( + self, + *, + transformer: torch.nn.Module | None = None, + vae: torch.nn.Module | None = None, + model_path: str | None = None, + default_height: int | None = None, + default_width: int | None = None, + ) -> None: + super().__init__() + self.transformer = transformer + self.vae = vae + self.model_path = model_path + self.default_height = default_height + self.default_width = default_width + + @property + def role_affinity(self): + return RoleType.MONOLITHIC + + @property + def parallelism_type(self) -> StageParallelismType: + return StageParallelismType.REPLICATED + + def require_session(self, batch: Req, *, context: str | None = None): + if batch.session is None: + label = context or self.__class__.__name__ + raise ValueError(f"{label} requires a realtime session") + return batch.session + + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + stage_name = self._component_stage_name(stage_name) + uses: list[ComponentUse] = [] + for spec in self.component_specs: + target_dtype = None + if spec.precision_attr is not None: + precision = getattr(server_args.pipeline_config, spec.precision_attr) + target_dtype = PRECISION_TO_TYPE[precision] + uses.append( + ComponentUse( + stage_name, + spec.component_name, + target_dtype=target_dtype, + memory_intensive=spec.memory_intensive, + keep_ready_after_warmup=spec.keep_ready_after_warmup, + ) + ) + return uses + + def target_pixel_size(self, batch: Req) -> tuple[int, int]: + height = batch.height if batch.height is not None else self.default_height + width = batch.width if batch.width is not None else self.default_width + if height is None or width is None: + raise ValueError( + f"{self.__class__.__name__} needs batch.height/batch.width or defaults" + ) + return int(height), int(width) + + def _empty_output(self, batch: Req) -> OutputBatch: + target_h, target_w = self.target_pixel_size(batch) + output = torch.empty( + (1, 3, 0, target_h, target_w), + dtype=torch.float32, + device=get_local_torch_device(), + ) + return OutputBatch(output=output, metrics=batch.metrics) + + def scale_and_shift_latents( + self, + latents: torch.Tensor, + server_args: ServerArgs, + ) -> torch.Tensor: + return scale_and_shift_latents(latents, server_args, self.vae) + + def ensure_causal_vae_conv_cache(self, state) -> dict: + if state.conv_cache is None: + state.conv_cache = self.vae.reset_decoder_cache() + return state.conv_cache diff --git a/python/sglang/multimodal_gen/runtime/realtime/__init__.py b/python/sglang/multimodal_gen/runtime/realtime/__init__.py index 2dd5abad1..cfdcb609f 100644 --- a/python/sglang/multimodal_gen/runtime/realtime/__init__.py +++ b/python/sglang/multimodal_gen/runtime/realtime/__init__.py @@ -2,6 +2,9 @@ """session-scoped realtime state, control events, and runtime-only helpers""" +from sglang.multimodal_gen.runtime.realtime.camera_controls import ( + RealtimeCameraControlState, +) from sglang.multimodal_gen.runtime.realtime.causal_state import RealtimeCausalDiTState from sglang.multimodal_gen.runtime.realtime.condition_events import ( ConditionEvent, @@ -25,6 +28,7 @@ __all__ = [ "ControlSignal", "ControlStateSamplingQueue", "ControlStateTransition", + "RealtimeCameraControlState", "RealtimeCausalDiTState", "RealtimeSession", "RealtimeSessionCache", diff --git a/python/sglang/multimodal_gen/runtime/realtime/camera_controls.py b/python/sglang/multimodal_gen/runtime/realtime/camera_controls.py new file mode 100644 index 000000000..3f70990de --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/realtime/camera_controls.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import deque +from collections.abc import Callable +from typing import Any + +from sglang.multimodal_gen.runtime.realtime.condition_events import ( + ControlSignal, + ControlStateSamplingQueue, + ControlStateTransition, +) + +CameraActionNormalizer = Callable[[list[Any]], list[str]] +CameraActionValidator = Callable[[Any], list[list[str]]] + + +def _identity_actions(actions: list[Any]) -> list[str]: + return list(actions) + + +class RealtimeCameraControlState: + """Session-local camera-control buffer shared by realtime model adapters. + + Camera controls arrive in two shapes: + + 1. Script mode: ``list[list[str]]`` where each item is one output-frame's + held actions. The script is consumed once from a FIFO and padded with + neutral ``[]`` frames after it runs out. + 2. State mode: timestamped transitions such as "W is currently held". + ``ControlStateSamplingQueue`` expands that continuous state into the next + chunk and can pulse a short key press for a minimum number of frames. + + The two modes are intentionally exclusive. A new script clears state mode, + and new state transitions clear script mode, so adapters never merge two + camera timelines accidentally. ``sample_camera_actions`` returns ``None`` + only when no control should be sent; otherwise it returns exactly + ``chunk_size`` frames, with ``[]`` meaning neutral/no-op for that frame. + """ + + def __init__( + self, + *, + min_pulse_items: int = 1, + script_maxlen: int = 512, + max_transitions: int = 512, + normalize_state_actions: CameraActionNormalizer = _identity_actions, + ) -> None: + self.camera_state = ControlStateSamplingQueue( + default_item=[], + min_pulse_items=min_pulse_items, + max_transitions=max_transitions, + ) + self.camera_script_queue: deque[ControlSignal] = deque(maxlen=script_maxlen) + self.latest_sampled_event_id: int | None = None + self._normalize_state_actions = normalize_state_actions + + def clear(self) -> None: + """Reset all camera controls owned by this realtime session.""" + self.camera_state.clear() + self.camera_script_queue.clear() + self.latest_sampled_event_id = None + + def receive_camera_script( + self, + camera_actions: list[list[str]], + *, + event_id: int | None = None, + ) -> None: + """Replace active controls with a finite per-frame script.""" + self.camera_script_queue.clear() + self.camera_state.clear() + for actions in camera_actions: + self.camera_script_queue.append( + ControlSignal( + kind="camera_actions", + payload=list(actions), + seq_id=event_id, + ) + ) + + def receive_camera_state_transitions( + self, + transitions: list[ControlStateTransition], + ) -> None: + """Replace the script with continuous state transitions.""" + self.camera_script_queue.clear() + self.camera_state.push_many(transitions) + + def receive_camera_actions( + self, + camera_actions: list[list[str]], + *, + event_id: int | None = None, + ) -> None: + self.receive_camera_script(camera_actions, event_id=event_id) + + def receive_camera_state( + self, + actions: list[str], + *, + event_id: int | None = None, + timestamp_ms: int | None = None, + ) -> None: + self.receive_camera_state_transitions( + [ + self._camera_state_transition( + actions, + event_id=event_id, + timestamp_ms=timestamp_ms, + ) + ] + ) + + def receive_camera_event_payload( + self, + payload: Any, + *, + event_id: int | None, + validate_camera_actions: CameraActionValidator, + ) -> str: + """Parse an external camera event and install it as script or state.""" + if isinstance(payload, dict) and payload.get("mode") == "state": + transitions = self._camera_transitions_from_event_payload( + payload, + event_id=event_id, + ) + self.receive_camera_state_transitions(transitions) + return f"kind=camera_actions, mode=state, transitions={len(transitions)}" + + camera_actions = validate_camera_actions(payload) + self.receive_camera_script(camera_actions, event_id=event_id) + return f"kind=camera_actions, mode=script, frames={len(camera_actions)}" + + def sample_camera_actions(self, chunk_size: int) -> list[list[str]] | None: + """Return the next chunk-sized camera action window. + + Script mode has priority because it represents an explicit finite + timeline. State mode is sampled only when no script is pending. + """ + if self.camera_script_queue: + return self._sample_camera_script(chunk_size) + action_list = self.camera_state.sample_chunk(chunk_size) + if action_list is None: + return None + self.latest_sampled_event_id = self.camera_state.latest_sampled_seq_id() + return [list(actions) for actions in action_list] + + def _sample_camera_script(self, chunk_size: int) -> list[list[str]]: + chunk: list[list[str]] = [] + latest_event_id = self.latest_sampled_event_id + while self.camera_script_queue and len(chunk) < chunk_size: + signal = self.camera_script_queue.popleft() + chunk.append(list(signal.payload)) + latest_event_id = signal.seq_id + while len(chunk) < chunk_size: + chunk.append([]) + self.latest_sampled_event_id = latest_event_id + return chunk + + def _camera_state_transition( + self, + actions: list[Any], + *, + event_id: int | None, + timestamp_ms: int | None, + ) -> ControlStateTransition: + return ControlStateTransition( + payload=self._normalize_state_actions(actions), + seq_id=event_id, + timestamp_ms=timestamp_ms, + ) + + def _camera_transitions_from_event_payload( + self, + payload: dict[str, Any], + *, + event_id: int | None, + ) -> list[ControlStateTransition]: + transitions = payload.get("transitions") + if not isinstance(transitions, list): + raise ValueError("camera_actions state payload requires transitions") + result = [] + for transition in transitions: + if not isinstance(transition, dict): + raise ValueError("camera_actions transition must be a map") + actions = transition.get("actions") + if not isinstance(actions, list): + raise ValueError("camera_actions transition actions must be a list") + timestamp_ms = transition.get("client_ts_ms") + if timestamp_ms is not None: + timestamp_ms = int(timestamp_ms) + result.append( + self._camera_state_transition( + actions, + event_id=event_id, + timestamp_ms=timestamp_ms, + ) + ) + return result diff --git a/python/sglang/multimodal_gen/runtime/realtime/causal_state.py b/python/sglang/multimodal_gen/runtime/realtime/causal_state.py index d7de5eb16..b54a13f2f 100644 --- a/python/sglang/multimodal_gen/runtime/realtime/causal_state.py +++ b/python/sglang/multimodal_gen/runtime/realtime/causal_state.py @@ -20,3 +20,16 @@ class RealtimeCausalDiTState(BaseRealtimeState): self.runtime_cache.clear() self.current_chunk_start_frame = 0 self.chunk_idx = 0 + + +class RealtimeCausalDecodeState(BaseRealtimeState): + """persist causal VAE decode cache and output frontier across chunks""" + + def __init__(self): + super().__init__() + self.conv_cache: dict | None = None + self.next_dec_idx: int = 0 + + def dispose(self) -> None: + self.conv_cache = None + self.next_dec_idx = 0 diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index bfacbb9fd..071f1c314 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -778,9 +778,11 @@ class ServerArgs(DisaggServerArgsMixin): # SamplingParams use classifier-free guidance (negative_prompt is not None), # because non-CFG models (e.g. FLUX) crash when CFG parallel splits ranks. if cfg_unspecified: + deployment_config = self.pipeline_config.get_model_deployment_config() cfg_group_size = self.dp_size * self.tp_size * 2 if ( self.performance_mode != "manual" + and deployment_config.auto_enable_cfg_parallel and self.num_gpus >= 2 and self.num_gpus % cfg_group_size == 0 and sp_unspecified diff --git a/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py index 5df6c958e..93e7e2292 100644 --- a/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py @@ -490,8 +490,10 @@ class ServerArgsAutoTuner: def _enable_cfg_parallel_if_supported(self) -> None: args = self.server_args + deployment_config = self._deployment_config() if ( - args.enable_cfg_parallel is None + deployment_config.auto_enable_cfg_parallel + and args.enable_cfg_parallel is None and not self._has_explicit_parallel_policy() and args._model_default_uses_cfg() ): diff --git a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py index ae8d6bf62..fc36551fc 100644 --- a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py +++ b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py @@ -29,6 +29,14 @@ BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = { "overlay_repo_id": "MickJ/LTX-2.3-overlay", "overlay_revision": "e0cc94f279ec16bb87c230134d40319f6ce40c5e", }, + "Efficient-Large-Model/SANA-WM_bidirectional": { + "overlay_repo_id": "sjmshsh/SANA-WM_bidirectional-overlay", + "overlay_revision": "e611beacbcc0cf33c676306ae0eb89f149e044ad", + }, + "Efficient-Large-Model/SANA-WM_streaming": { + "overlay_repo_id": "AgainstEntropy/SANA-WM_streaming-overlay", + "overlay_revision": "62c6840871ecc3559189047513ba0670e1bf62e7", + }, } @@ -150,9 +158,17 @@ def resolve_model_overlay_target( if os.path.exists(model_name_or_path): # Local source dirs do not have a repo id, so match them by basename. base_name = os.path.basename(os.path.normpath(model_name_or_path)) + normalized_path = ( + os.path.normpath(model_name_or_path).lower().replace(os.sep, "/") + ) for source_model_id, spec in registry.items(): if base_name == source_model_id.rsplit("/", 1)[-1]: return source_model_id, spec + cache_repo_fragment = ( + f"models--{source_model_id.lower().replace('/', '--')}" + ) + if cache_repo_fragment in normalized_path: + return source_model_id, spec return None @@ -577,7 +593,6 @@ def maybe_load_overlay_model_index( # A local overlay repo already contains the model_index we need. if load_overlay_manifest_if_present(model_name_or_path) is not None: return load_model_index_from_dir(model_name_or_path) - return None overlay_target = resolve_model_overlay_target(model_name_or_path) if overlay_target is not None: diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index e1ac03378..bd6b2c4b4 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -26,6 +26,7 @@ from sglang.multimodal_gen.test.server.testcase_configs import ( MULTI_FRAME_I2I_sampling_params, MULTI_IMAGE_TI2I_sampling_params, MULTI_IMAGE_TI2I_UPLOAD_sampling_params, + SANA_WM_TI2V_CI_sampling_params, T2I_sampling_params, T2V_sampling_params, _make_modelopt_ci_case, @@ -44,6 +45,7 @@ from sglang.multimodal_gen.test.test_utils import ( DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_LAYERED_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST, + DEFAULT_SANA_WM_STREAMING_MODEL_NAME_FOR_TEST, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, DEFAULT_WAN_2_1_I2V_14B_480P_MODEL_NAME_FOR_TEST, DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST, @@ -376,6 +378,17 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [ model_path="FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", ), ), + DiffusionTestCase( + "sana_wm_ti2v", + DiffusionServerArgs( + model_path=DEFAULT_SANA_WM_STREAMING_MODEL_NAME_FOR_TEST, + ), + SANA_WM_TI2V_CI_sampling_params, + run_perf_check=False, + run_component_accuracy_check=False, + run_models_api_check=False, + run_t2v_input_reference_check=False, + ), # flaky # === Helios T2V === # DiffusionTestCase( diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 28f9756b1..02107bdc0 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -545,6 +545,15 @@ TI2V_sampling_params = DiffusionSamplingParams( direct_url_test=True, ) +SANA_WM_TI2V_CI_sampling_params = DiffusionSamplingParams( + prompt=TI2V_sampling_params.prompt, + image_path=TI2V_sampling_params.image_path, + direct_url_test=True, + output_size="384x640", + num_frames=17, + extras={"num_inference_steps": 12, "seed": 0, "guidance_scale": 4.5}, +) + TURBOWAN_I2V_sampling_params = DiffusionSamplingParams( prompt="The man in the picture slowly turns his head, his expression enigmatic and otherworldly. The camera performs a slow, cinematic dolly out, focusing on his face. Moody lighting, neon signs glowing in the background, shallow depth of field.", image_path="https://is1-ssl.mzstatic.com/image/thumb/Music114/v4/5f/fa/56/5ffa56c2-ea1f-7a17-6bad-192ff9b6476d/825646124206.jpg/600x600bb.jpg", diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index cb53038b6..92c1e72f2 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -33,7 +33,7 @@ if TYPE_CHECKING: logger = init_logger(__name__) -SGL_TEST_FILES_CI_DATA_REVISION = "af6e712a2c49ab5fcd81dde58e2f54c78e77683b" +SGL_TEST_FILES_CI_DATA_REVISION = "caa56302ccf2d289e4488ed06d952edf5d2314cf" SGL_TEST_FILES_CONSISTENCY_GT_ROOT = ( "https://raw.githubusercontent.com/" f"sgl-project/ci-data/{SGL_TEST_FILES_CI_DATA_REVISION}/" @@ -154,6 +154,12 @@ DEFAULT_WAN_2_2_I2V_A14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-I2V-A14B-Diffusers # MOVA video generation models DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST = "OpenMOSS-Team/MOVA-360p" +# SANA-WM world model (TI2V with optional camera conditioning) +DEFAULT_SANA_WM_MODEL_NAME_FOR_TEST = "Efficient-Large-Model/SANA-WM_bidirectional" +DEFAULT_SANA_WM_STREAMING_MODEL_NAME_FOR_TEST = ( + "Efficient-Large-Model/SANA-WM_streaming" +) + def print_value_formatted(description: str, value: int | float | str): """Helper function to print a metric value formatted.""" diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_output_transport.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_output_transport.py index 0a36c2009..a8b2e2808 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_output_transport.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_output_transport.py @@ -562,6 +562,7 @@ def test_raw_rgb_realtime_output_adapter_offloads_preview_encoding(monkeypatch): async def run(): ws = _WebSocket() adapter = RawRGBRealtimeOutputAdapter() + frame_count = realtime_output_adapter.ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE + 1 batch = SimpleNamespace( block_idx=0, request_id="req-webp-offload", @@ -575,8 +576,8 @@ def test_raw_rgb_realtime_output_adapter_offloads_preview_encoding(monkeypatch): result = OutputBatch( raw_frame_batches=[ [ - bytes([255, 0, 0, 0, 255, 0]), - bytes([0, 0, 255, 255, 255, 0]), + bytes([idx % 256, 0, 0, 0, 255, idx % 256]) + for idx in range(frame_count) ] ], raw_frame_content_type=RAW_RGB_CONTENT_TYPE, @@ -595,18 +596,31 @@ def test_raw_rgb_realtime_output_adapter_offloads_preview_encoding(monkeypatch): payloads = asyncio.run(run()) assert [call[0] for call in calls] == [ - realtime_output_adapter._encode_rgb_frame_to_webp, - realtime_output_adapter._encode_rgb_frame_to_webp, - ] - [(header, payload)] = _unpack_frame_batch_messages(payloads) - assert header["content_type"] == WEBP_FRAME_CONTENT_TYPE - assert header["encoding"] == "webp" - assert header["num_frames"] == 2 - assert header["frame_batch_index"] == 0 - assert header["num_frame_batches"] == 1 - assert header["is_final_frame_batch"] is True - assert len(header["payload_lengths"]) == 2 - assert payload.startswith(b"RIFF") + realtime_output_adapter._encode_rgb_frame_to_webp + ] * (realtime_output_adapter.ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE + 1) + (first_header, first_payload), (second_header, second_payload) = ( + _unpack_frame_batch_messages(payloads) + ) + assert first_header["content_type"] == WEBP_FRAME_CONTENT_TYPE + assert first_header["encoding"] == "webp" + assert first_header["num_frames"] == ( + realtime_output_adapter.ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE + ) + assert first_header["frame_batch_index"] == 0 + assert first_header["num_frame_batches"] == 2 + assert first_header["is_final_frame_batch"] is False + assert len(first_header["payload_lengths"]) == ( + realtime_output_adapter.ENCODED_PREVIEW_FRAMES_PER_WS_MESSAGE + ) + assert second_header["content_type"] == WEBP_FRAME_CONTENT_TYPE + assert second_header["encoding"] == "webp" + assert second_header["num_frames"] == 1 + assert second_header["frame_batch_index"] == 1 + assert second_header["num_frame_batches"] == 2 + assert second_header["is_final_frame_batch"] is True + assert len(second_header["payload_lengths"]) == 1 + assert first_payload.startswith(b"RIFF") + assert second_payload.startswith(b"RIFF") def test_raw_rgb_realtime_output_adapter_can_send_jpeg_preview_frames(): diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py index db75c3d46..2ed9507a2 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py @@ -22,6 +22,9 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.realtime import ( from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.adapters import ( lingbot_world_realtime_adapter as lingbot_realtime, ) +from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.adapters import ( + sana_wm_realtime_adapter as sana_wm_realtime, +) from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import ( GenerateSession, ) @@ -35,10 +38,16 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import ( LingBotWorldCausalDMDDenoisingStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_diffusion import ( + RealtimeDiffusionStage, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_input_validation import ( RealtimeInputValidationStage, RealtimeInputValidationState, ) +from sglang.multimodal_gen.runtime.realtime.causal_state import ( + RealtimeCausalDecodeState, +) from sglang.multimodal_gen.runtime.realtime.condition_events import ( ControlStateTransition, ) @@ -65,6 +74,51 @@ class _State(BaseRealtimeState): self.disposed = True +def test_realtime_diffusion_stage_declares_long_lived_components(): + stage = RealtimeDiffusionStage() + server_args = SimpleNamespace( + pipeline_config=SimpleNamespace(dit_precision="bf16", vae_precision="fp32") + ) + + uses = stage.component_uses(server_args, stage_name="realtime") + + assert [use.component_name for use in uses] == ["transformer", "vae"] + assert [use.stage_name for use in uses] == ["realtime", "realtime"] + assert uses[0].target_dtype == torch.bfloat16 + assert uses[0].memory_intensive + assert uses[0].keep_ready_after_warmup + assert uses[1].target_dtype == torch.float32 + assert not uses[1].memory_intensive + assert uses[1].keep_ready_after_warmup + + +def test_realtime_diffusion_stage_requires_session(): + stage = RealtimeDiffusionStage(default_height=480, default_width=832) + req = _Req(session=None) + + try: + stage.require_session(req, context="test realtime") + except ValueError as exc: + assert "test realtime requires a realtime session" in str(exc) + else: + raise AssertionError("expected missing realtime session to fail") + + session = object() + req.session = session + assert stage.require_session(req) is session + + +def test_realtime_causal_decode_state_dispose_resets_frontier(): + state = RealtimeCausalDecodeState() + state.conv_cache = {"cache": object()} + state.next_dec_idx = 7 + + state.dispose() + + assert state.conv_cache is None + assert state.next_dec_idx == 0 + + def test_realtime_session_cache_reuses_and_releases_state(): cache = RealtimeSessionCache(max_sessions=1) first = _Req(realtime_session_id="session-a", block_idx=0, session=None) @@ -148,6 +202,53 @@ def test_lingbot_realtime_camera_state_compacts_multiple_pending_updates(): assert state.latest_sampled_event_id == 9 +def test_sana_wm_realtime_camera_state_uses_sana_normalizer(): + state = sana_wm_realtime.SanaWMRealtimeAdapterState() + result = state.receive_camera_event_payload( + { + "mode": "state", + "transitions": [ + {"actions": ["W"], "client_ts_ms": 100}, + {"actions": [], "client_ts_ms": 120}, + ], + }, + event_id=11, + ) + + assert result == "kind=camera_actions, mode=state, transitions=2" + assert state.sample_camera_actions(10) == [["w"]] * 8 + [[], []] + assert state.latest_sampled_event_id == 11 + + +def test_sana_wm_realtime_adapter_preserves_requested_size(): + async def fake_save_image_to_path(image, target_path): + return target_path + + old_save_image_to_path = sana_wm_realtime.save_image_to_path + old_get_global_server_args = sana_wm_realtime.get_global_server_args + sana_wm_realtime.save_image_to_path = fake_save_image_to_path + sana_wm_realtime.get_global_server_args = lambda: SimpleNamespace( + input_save_path=None + ) + try: + adapter = sana_wm_realtime.SanaWMRealtimeAdapter() + session = GenerateSession() + session.set_adapter(adapter) + request = RealtimeVideoGenerationsRequest( + type="init", + prompt="walk forward", + first_frame=b"fake-image", + size="832x480", + ) + + asyncio.run(adapter.on_init(session, request)) + + assert request.size == "832x480" + finally: + sana_wm_realtime.save_image_to_path = old_save_image_to_path + sana_wm_realtime.get_global_server_args = old_get_global_server_args + + def test_lingbot_realtime_adapter_ingests_generic_events(): adapter = lingbot_realtime.LingBotWorldRealtimeAdapter() session = GenerateSession() diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py new file mode 100644 index 000000000..75e594b37 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py @@ -0,0 +1,1477 @@ +import argparse +import os +import sys +import tempfile +import types +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.multimodal_gen.configs.models.dits.sana_wm import SanaWMConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType +from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig +from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams +from sglang.multimodal_gen.registry import _get_config_info, get_pipeline_config_classes +from sglang.multimodal_gen.runtime import server_args as server_args_module +from sglang.multimodal_gen.runtime.models.dits.sana_wm import ( + _downscale_to_reference_rms, + _gdn_chunk_scan_forward, + _gdn_scan_forward, + _RMSNorm, + _sana_wm_chunk_index_from_chunk_size, + _sana_wm_chunked_attention, + _sana_wm_normalize_chunk_index, + _single_path_delta_chunk_scan_forward, + _single_path_delta_scan_forward, + compute_chunk_plucker, +) +from sglang.multimodal_gen.runtime.pipelines.sana_wm_pipeline import ( + SanaWMPipeline, + SanaWMTwoStagePipeline, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( + StageParallelismType, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm import ( + SanaWMBeforeDenoisingStage, + SanaWMDenoisingStage, + SanaWMTextEncodingStage, + _align_sana_wm_cfg_text_conditions, + configure_sana_wm_ltx2_vae_for_long_video, + parse_sana_wm_action_string, + sana_wm_action_to_camera_to_world, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.refiner import ( + OfficialDiffusersLTX2RefinerModule, + OfficialGemma3TextEncoderModule, + SanaWMLTX2RefinerStage, + SanaWMRefinerDecodingStage, + _refiner_config_value, + _streaming_diffusers_self_attention, + _uses_diffusers_ltx2_refiner, + sana_wm_skip_refiner_enabled, +) +from sglang.multimodal_gen.runtime.server_args import set_global_server_args +from sglang.multimodal_gen.runtime.utils.model_overlay import ( + resolve_model_overlay_target, +) +from sglang.multimodal_gen.test.test_utils import DEFAULT_SANA_WM_MODEL_NAME_FOR_TEST + +_SANA_WM_REFINER_STAGE_MODULE = ( + "sglang.multimodal_gen.runtime.pipelines_core.stages." + "model_specific_stages.sana_wm.refiner" +) + + +class _GlobalStageArgsMixin: + def setUp(self) -> None: + super().setUp() + self._prev_global_server_args = server_args_module._global_server_args + set_global_server_args( + SimpleNamespace( + comfyui_mode=False, + enable_cfg_parallel=False, + enable_torch_compile=False, + attention_backend=None, + ) + ) + + def tearDown(self) -> None: + set_global_server_args(self._prev_global_server_args) + super().tearDown() + + +class TestSanaWMPipelineConfig(unittest.TestCase): + def setUp(self) -> None: + self.config = SanaWMPipelineConfig() + + def test_task_type_is_ti2v(self) -> None: + self.assertEqual(self.config.task_type, ModelTaskType.TI2V) + + def test_adjust_num_frames_rounds_to_temporal_stride(self) -> None: + self.assertEqual(self.config.adjust_num_frames(50), 49) + self.assertEqual(self.config.adjust_num_frames(49), 49) + + def test_prepare_latent_shape_returns_5d(self) -> None: + batch = SimpleNamespace(height=704, width=1280) + shape = self.config.prepare_latent_shape(batch, batch_size=1, num_frames=49) + self.assertEqual(shape, (1, 128, 7, 22, 40)) + + def test_prepare_latent_shape_requires_spatial_stride_alignment(self) -> None: + batch = SimpleNamespace(height=705, width=1280) + with self.assertRaisesRegex(ValueError, "divisible"): + self.config.prepare_latent_shape(batch, batch_size=1, num_frames=49) + + def test_prepare_latent_shape_uses_axis_specific_spatial_strides(self) -> None: + self.config.vae_stride = (8, 16, 32) + batch = SimpleNamespace(height=704, width=1280) + shape = self.config.prepare_latent_shape(batch, batch_size=1, num_frames=49) + self.assertEqual(shape, (1, 128, 7, 44, 40)) + + def test_prepare_pos_cond_kwargs_passes_camera_from_batch_extra(self) -> None: + # SanaWMBeforeDenoisingStage packs c2w + intrinsics into the upstream + # latent-frame (B, T_lat, 20) raymap plus a (B, 48, T_lat, H, W) + # ``chunk_plucker`` tensor. The pipeline config just forwards those + # tensors verbatim. + camera_conditions = torch.zeros(1, 7, 20) + chunk_plucker = torch.zeros(1, 48, 7, 22, 40) + batch = SimpleNamespace( + prompt_attention_mask=torch.ones(1, 16), + extra={ + "camera_conditions": camera_conditions, + "chunk_plucker": chunk_plucker, + }, + ) + kwargs = self.config.prepare_pos_cond_kwargs( + batch=batch, + device=torch.device("cpu"), + rotary_emb=None, + dtype=torch.bfloat16, + ) + self.assertIs(kwargs["camera_conditions"], camera_conditions) + self.assertIs(kwargs["chunk_plucker"], chunk_plucker) + + def test_get_model_deployment_config_enables_dit_layerwise_offload(self) -> None: + deployment = self.config.get_model_deployment_config() + self.assertTrue(deployment.auto_dit_layerwise_offload) + self.assertEqual(deployment.fsdp_auto_min_available_memory_gb, 60) + + def test_text_encoder_padding_matches_cfg_concat_contract(self) -> None: + self.assertEqual( + self.config.text_encoder_extra_args[0]["padding"], "max_length" + ) + self.assertTrue(self.config.text_encoder_extra_args[0]["return_attention_mask"]) + self.assertTrue(self.config.chi_prompt) + + def test_inference_flow_shift_matches_official_reference(self) -> None: + self.assertEqual(self.config.flow_shift, 9.95) + self.assertEqual(self.config.inference_flow_shift, 9.8) + + def test_dit_arch_text_norm_defaults_match_upstream(self) -> None: + arch = SanaWMConfig().arch_config + self.assertTrue(arch.y_norm) + self.assertEqual(arch.y_norm_scale_factor, 0.01) + self.assertEqual(arch.y_norm_eps, 1e-5) + self.assertEqual(arch.timestep_norm_scale_factor, 1.0) + + def test_prepare_neg_cond_kwargs_keeps_camera_for_cfg(self) -> None: + camera_conditions = torch.zeros(1, 7, 20) + chunk_plucker = torch.zeros(1, 48, 7, 22, 40) + batch = SimpleNamespace( + negative_attention_mask=torch.ones(1, 16), + extra={ + "camera_conditions": camera_conditions, + "chunk_plucker": chunk_plucker, + }, + ) + kwargs = self.config.prepare_neg_cond_kwargs( + batch=batch, + device=torch.device("cpu"), + rotary_emb=None, + dtype=torch.bfloat16, + ) + self.assertIn("encoder_attention_mask", kwargs) + self.assertIs(kwargs["camera_conditions"], camera_conditions) + self.assertIs(kwargs["chunk_plucker"], chunk_plucker) + + def test_decode_scale_and_shift_uses_ltx2_latent_stats(self) -> None: + vae = SimpleNamespace( + config=SimpleNamespace(scaling_factor=2.0), + latents_mean=torch.tensor([1.0, 2.0]), + latents_std=torch.tensor([2.0, 4.0]), + ) + scale, shift = self.config.get_decode_scale_and_shift( + torch.device("cpu"), torch.float32, vae + ) + + self.assertTrue( + torch.equal(scale, torch.tensor([1.0, 0.5]).view(1, 2, 1, 1, 1)) + ) + self.assertTrue( + torch.equal(shift, torch.tensor([1.0, 2.0]).view(1, 2, 1, 1, 1)) + ) + + def test_sana_wm_ltx2_vae_tiling_defaults_match_upstream(self) -> None: + self.assertTrue(self.config.vae_tiling) + self.assertTrue(self.config.vae_framewise_encoding) + self.assertTrue(self.config.vae_framewise_decoding) + self.assertEqual(self.config.vae_tile_sample_min_num_frames, 96) + self.assertEqual(self.config.vae_tile_sample_stride_num_frames, 64) + self.assertEqual(self.config.vae_config.tile_sample_min_num_frames, 96) + self.assertEqual(self.config.vae_config.tile_sample_stride_num_frames, 64) + self.assertEqual(self.config.vae_config.blend_num_frames, 32) + + def test_configure_sana_wm_ltx2_vae_enables_framewise_decode(self) -> None: + class FakeLTX2VAE: + def __init__(self): + self.use_tiling = False + self.use_framewise_encoding = False + self.use_framewise_decoding = False + self.tile_sample_min_num_frames = 16 + self.tile_sample_stride_num_frames = 8 + self.enable_tiling_kwargs = None + + def enable_tiling(self, **kwargs): + self.use_tiling = True + self.enable_tiling_kwargs = kwargs + self.tile_sample_min_num_frames = kwargs["tile_sample_min_num_frames"] + self.tile_sample_stride_num_frames = kwargs[ + "tile_sample_stride_num_frames" + ] + + vae = FakeLTX2VAE() + configure_sana_wm_ltx2_vae_for_long_video(vae, self.config) + + self.assertTrue(vae.use_tiling) + self.assertTrue(vae.use_framewise_encoding) + self.assertTrue(vae.use_framewise_decoding) + self.assertEqual(vae.tile_sample_min_num_frames, 96) + self.assertEqual(vae.tile_sample_stride_num_frames, 64) + self.assertEqual( + vae.enable_tiling_kwargs, + { + "tile_sample_min_num_frames": 96, + "tile_sample_stride_num_frames": 64, + }, + ) + + def test_configure_sana_wm_ltx2_vae_honors_nested_vae_config(self) -> None: + vae = SimpleNamespace( + use_tiling=False, + use_framewise_encoding=False, + use_framewise_decoding=False, + tile_sample_min_num_frames=16, + tile_sample_stride_num_frames=8, + ) + + def enable_tiling(**kwargs): + vae.use_tiling = True + vae.enable_tiling_kwargs = kwargs + + vae.enable_tiling = enable_tiling + self.config.vae_config.tile_sample_min_num_frames = 128 + self.config.vae_config.tile_sample_stride_num_frames = 80 + + configure_sana_wm_ltx2_vae_for_long_video(vae, self.config) + + self.assertEqual(vae.tile_sample_min_num_frames, 128) + self.assertEqual(vae.tile_sample_stride_num_frames, 80) + self.assertEqual( + vae.enable_tiling_kwargs, + { + "tile_sample_min_num_frames": 128, + "tile_sample_stride_num_frames": 80, + }, + ) + + def test_cfg_text_conditions_pad_positive_and_negative_to_same_length(self) -> None: + pos = torch.ones(1, 173, 4) + neg = torch.ones(1, 1, 4) * 2 + pos_mask = torch.ones(1, 173, dtype=torch.long) + neg_mask = torch.ones(1, 1, dtype=torch.long) + + pos, neg, pos_mask, neg_mask = _align_sana_wm_cfg_text_conditions( + pos, neg, pos_mask, neg_mask + ) + + self.assertEqual(pos.shape, (1, 173, 4)) + self.assertEqual(neg.shape, (1, 173, 4)) + self.assertEqual(pos_mask.shape, (1, 173)) + self.assertEqual(neg_mask.shape, (1, 173)) + self.assertTrue(torch.equal(neg[:, :1], torch.ones(1, 1, 4) * 2)) + self.assertTrue(torch.equal(neg[:, 1:], torch.zeros(1, 172, 4))) + self.assertTrue( + torch.equal(neg_mask[:, :1], torch.ones(1, 1, dtype=torch.long)) + ) + self.assertTrue( + torch.equal(neg_mask[:, 1:], torch.zeros(1, 172, dtype=torch.long)) + ) + + +class TestSanaWMSamplingParams(unittest.TestCase): + def test_defaults_match_video_ti2v_contract(self) -> None: + params = SanaWMSamplingParams() + self.assertEqual(params.height, 704) + self.assertEqual(params.width, 1280) + self.assertEqual(params.num_frames, 49) + self.assertEqual(params.num_inference_steps, 20) + self.assertEqual(params.guidance_scale, 4.5) + self.assertEqual(params.negative_prompt, "") + + def test_build_request_extra_omits_camera_by_default(self) -> None: + params = SanaWMSamplingParams() + extra = params.build_request_extra() + self.assertNotIn("camera_to_world", extra) + self.assertNotIn("intrinsics", extra) + + def test_build_request_extra_includes_camera_when_set(self) -> None: + cam = torch.eye(4).unsqueeze(0).expand(49, 4, 4) + intr = torch.eye(3).unsqueeze(0).expand(49, 3, 3) + params = SanaWMSamplingParams(camera_to_world=cam, intrinsics=intr) + extra = params.build_request_extra() + self.assertIs(extra["camera_to_world"], cam) + self.assertIs(extra["intrinsics"], intr) + + def test_build_request_extra_includes_action_when_set(self) -> None: + params = SanaWMSamplingParams( + action="w-8,jw-8", + translation_speed=0.06, + rotation_speed_deg=2.0, + pitch_limit_deg=70.0, + ) + + extra = params.build_request_extra() + + self.assertEqual(extra["action"], "w-8,jw-8") + self.assertEqual(extra["translation_speed"], 0.06) + self.assertEqual(extra["rotation_speed_deg"], 2.0) + self.assertEqual(extra["pitch_limit_deg"], 70.0) + + def test_build_request_extra_rejects_action_with_direct_camera(self) -> None: + cam = torch.eye(4).unsqueeze(0).expand(49, 4, 4) + params = SanaWMSamplingParams(action="w-8", camera_to_world=cam) + + with self.assertRaisesRegex(ValueError, "either action or camera_to_world"): + params.build_request_extra() + + def test_cli_args_accept_upstream_action_aliases(self) -> None: + parser = argparse.ArgumentParser() + SanaWMSamplingParams.add_cli_args(parser) + + args = parser.parse_args( + [ + "--action", + "w-8,jw-8", + "--translation_speed", + "0.06", + "--rotation-speed-deg", + "2.0", + ] + ) + cli_args = SanaWMSamplingParams.get_cli_args(args) + + self.assertEqual(cli_args["action"], "w-8,jw-8") + self.assertEqual(cli_args["translation_speed"], 0.06) + self.assertEqual(cli_args["rotation_speed_deg"], 2.0) + + +class TestSanaWMRegistry(unittest.TestCase): + def setUp(self) -> None: + _get_config_info.cache_clear() + + def test_model_path_resolves_sana_wm_pipeline_config(self) -> None: + info = _get_config_info(DEFAULT_SANA_WM_MODEL_NAME_FOR_TEST) + self.assertIsNotNone(info) + self.assertIs(info.pipeline_config_cls, SanaWMPipelineConfig) + self.assertIs(info.sampling_param_cls, SanaWMSamplingParams) + + def test_two_stage_pipeline_registers_sana_wm_config_classes(self) -> None: + classes = get_pipeline_config_classes("SanaWMTwoStagePipeline") + self.assertIsNotNone(classes) + pipeline_config_cls, sampling_param_cls = classes + self.assertIs(pipeline_config_cls, SanaWMPipelineConfig) + self.assertIs(sampling_param_cls, SanaWMSamplingParams) + + def test_overlay_resolver_matches_hf_cache_snapshot_paths(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + snapshot_dir = f"{tmp_dir}/hub/models--Lightricks--LTX-2.3/snapshots/abc123" + os.makedirs(snapshot_dir) + target = resolve_model_overlay_target(snapshot_dir) + self.assertIsNotNone(target) + source_model_id, _ = target + self.assertEqual(source_model_id, "Lightricks/LTX-2.3") + + +class TestSanaWMTwoStagePipeline(unittest.TestCase): + def test_resolve_refiner_paths_defaults_to_model_refiner_dir(self) -> None: + pipeline = object.__new__(SanaWMTwoStagePipeline) + pipeline.model_path = "/models/sana-wm" + server_args = SimpleNamespace(component_paths={}) + refiner_root, refiner_gemma_root = pipeline._resolve_refiner_paths(server_args) + self.assertEqual(refiner_root, "/models/sana-wm/refiner") + self.assertEqual(refiner_gemma_root, "/models/sana-wm/refiner/text_encoder") + + def test_resolve_refiner_paths_accepts_component_overrides(self) -> None: + pipeline = object.__new__(SanaWMTwoStagePipeline) + pipeline.model_path = "/models/sana-wm" + server_args = SimpleNamespace( + component_paths={ + "refiner": "/custom/refiner", + "refiner_text_encoder": "/custom/refiner/text_encoder", + } + ) + refiner_root, refiner_gemma_root = pipeline._resolve_refiner_paths(server_args) + self.assertEqual(refiner_root, "/custom/refiner") + self.assertEqual(refiner_gemma_root, "/custom/refiner/text_encoder") + + def test_refiner_modules_are_loaded_from_official_subtrees(self) -> None: + self.assertEqual( + SanaWMTwoStagePipeline._REFINER_SUB_MODULES, + ( + ("transformer_2", "refiner/transformer"), + ("connectors", "refiner/connectors"), + ("text_encoder_2", "refiner/text_encoder"), + ("tokenizer_2", "refiner/text_encoder"), + ), + ) + + +class TestSanaWMPipeline(unittest.TestCase): + def test_validate_parallelism_rejects_tensor_parallelism(self) -> None: + with self.assertRaisesRegex(ValueError, "tensor parallelism"): + SanaWMPipeline._validate_parallelism_args( + SimpleNamespace(tp_size=2, sp_degree=1) + ) + + def test_validate_parallelism_rejects_sequence_parallelism(self) -> None: + with self.assertRaisesRegex(ValueError, "sequence parallelism"): + SanaWMPipeline._validate_parallelism_args( + SimpleNamespace(tp_size=1, sp_degree=2) + ) + + +class TestSanaWMBeforeDenoisingStage(_GlobalStageArgsMixin, unittest.TestCase): + def test_action_string_rolls_out_camera_to_world(self) -> None: + self.assertEqual( + parse_sana_wm_action_string(" w-2，d-1, none-1 "), + [["w"], ["w"], ["d"], []], + ) + + poses = sana_wm_action_to_camera_to_world( + "w-2,d-1", + translation_speed=0.05, + rotation_speed_deg=1.2, + ) + + self.assertEqual(poses.shape, (4, 4, 4)) + self.assertAlmostEqual(float(poses[2, 2, 3]), 0.1, places=5) + self.assertAlmostEqual(float(poses[3, 0, 3]), 0.05, places=5) + + def test_action_rollout_uses_upstream_strafe_yaw_coupling(self) -> None: + # Strafe-only action yaws the camera (forward axis tilts off +Z), matching + # the upstream SANA-WM action convention. + strafed = sana_wm_action_to_camera_to_world( + "d-10", translation_speed=0.04, rotation_speed_deg=1.2 + ) + self.assertGreater(abs(float(strafed[-1, 0, 2])), 1e-4) + + def test_action_string_rejects_unknown_keys_and_bad_duration(self) -> None: + with self.assertRaisesRegex(ValueError, "unknown action keys"): + parse_sana_wm_action_string("x-1") + with self.assertRaisesRegex(ValueError, "invalid duration"): + parse_sana_wm_action_string("w-0") + + def test_vae_encode_image_extracts_latent_dist_and_normalizes_ltx2(self) -> None: + class DummyLatentDist: + def mode(self): + return torch.tensor([[[[[3.0]]], [[[10.0]]]]]) + + class DummyVAE: + dtype = torch.float32 + device = torch.device("cpu") + config = SimpleNamespace(scaling_factor=2.0) + latents_mean = torch.tensor([1.0, 2.0]) + latents_std = torch.tensor([2.0, 4.0]) + + def encode(self, image): + return SimpleNamespace(latent_dist=DummyLatentDist()) + + stage = SanaWMBeforeDenoisingStage( + vae=DummyVAE(), + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + + encoded = stage._vae_encode_image( + torch.ones(1, 3, 1, 1), + dtype=torch.float32, + device=torch.device("cpu"), + ) + + expected = torch.tensor([[[[[2.0]]], [[[4.0]]]]]) + self.assertTrue(torch.equal(encoded, expected)) + + def test_first_frame_vae_use_is_declared_for_component_residency(self) -> None: + pipeline_config = SanaWMPipelineConfig() + stage = SanaWMBeforeDenoisingStage( + vae=object(), + transformer=None, + scheduler=None, + pipeline_config=pipeline_config, + ) + + uses = stage.component_uses( + SimpleNamespace(pipeline_config=pipeline_config), + stage_name="sana_wm_before_denoising", + ) + + self.assertEqual([use.component_name for use in uses], ["vae"]) + self.assertEqual(uses[0].target_dtype, torch.bfloat16) + + def test_prepare_noise_latents_accepts_per_sample_generators(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + + def make_generators(): + return [ + torch.Generator(device="cpu").manual_seed(11), + torch.Generator(device="cpu").manual_seed(23), + ] + + latents = stage._prepare_noise_latents( + (2, 1, 1, 2, 2), + dtype=torch.float32, + device=torch.device("cpu"), + generator=make_generators(), + ) + latents_again = stage._prepare_noise_latents( + (2, 1, 1, 2, 2), + dtype=torch.float32, + device=torch.device("cpu"), + generator=make_generators(), + ) + + self.assertEqual(latents.shape, (2, 1, 1, 2, 2)) + self.assertTrue(torch.equal(latents, latents_again)) + self.assertFalse(torch.equal(latents[0], latents[1])) + + def test_prepare_noise_latents_rejects_mismatched_generator_count(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + + with self.assertRaisesRegex(ValueError, "length must match"): + stage._prepare_noise_latents( + (2, 1, 1, 1, 1), + dtype=torch.float32, + device=torch.device("cpu"), + generator=[ + torch.Generator(device="cpu").manual_seed(1), + torch.Generator(device="cpu").manual_seed(2), + torch.Generator(device="cpu").manual_seed(3), + ], + ) + + def test_generator_from_seed_accepts_per_sample_seed_list(self) -> None: + generators = SanaWMBeforeDenoisingStage._generator_from_seed( + [11, 23], + batch_size=2, + device=torch.device("cpu"), + ) + + self.assertEqual(len(generators), 2) + + with self.assertRaisesRegex(ValueError, "seed list length"): + SanaWMBeforeDenoisingStage._generator_from_seed( + [1, 2, 3], + batch_size=2, + device=torch.device("cpu"), + ) + + def test_condition_image_batch_helper_accepts_single_or_batch_list(self) -> None: + self.assertEqual( + SanaWMBeforeDenoisingStage._condition_images_for_batch(["img"], 2), + ["img"], + ) + self.assertEqual( + SanaWMBeforeDenoisingStage._condition_images_for_batch(["a", "b"], 2), + ["a", "b"], + ) + + with self.assertRaisesRegex(ValueError, "one image or one image per batch"): + SanaWMBeforeDenoisingStage._condition_images_for_batch(["a", "b", "c"], 2) + + def test_first_frame_preprocess_records_official_crop_geometry(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + latents = torch.zeros(1, 128, 7, 22, 40) + batch = SimpleNamespace(extra={}) + + with patch.object( + stage, + "_vae_encode_image", + return_value=torch.ones(1, 128, 1, 22, 40), + ): + out = stage._splice_first_frame( + latents, + torch.zeros(3, 100, 100), + dtype=torch.float32, + device=torch.device("cpu"), + batch=batch, + ) + + info = batch.extra["sana_wm_condition_image_preprocess"] + self.assertEqual(info["source_size"], (100, 100)) + self.assertEqual(info["resized_size"], (1280, 1280)) + self.assertEqual(info["crop_offset"], (0, 288)) + self.assertEqual(info["target_size"], (1280, 704)) + self.assertTrue(torch.equal(out[:, :, :1], torch.ones(1, 128, 1, 22, 40))) + + def test_first_frame_splice_supports_batched_condition_images(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + latents = torch.zeros(2, 128, 7, 22, 40) + batch = SimpleNamespace(extra={}) + + with patch.object( + stage, + "_vae_encode_image", + side_effect=[ + torch.ones(1, 128, 1, 22, 40), + torch.full((1, 128, 1, 22, 40), 2.0), + ], + ): + out = stage._splice_first_frame( + latents, + [torch.zeros(3, 100, 100), torch.zeros(3, 200, 100)], + dtype=torch.float32, + device=torch.device("cpu"), + batch=batch, + ) + + info = batch.extra["sana_wm_condition_image_preprocess"] + self.assertEqual(len(info), 2) + self.assertTrue(torch.equal(out[0, :, :1], torch.ones(128, 1, 22, 40))) + self.assertTrue(torch.equal(out[1, :, :1], torch.full((128, 1, 22, 40), 2.0))) + + def test_default_static_camera_builds_latent_raymap_and_chunk_plucker(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + batch = SimpleNamespace(extra={}, height=384, width=640) + + camera_conditions, chunk_plucker, source = stage._build_camera_conditioning( + batch, + batch_size=1, + num_frames=17, + latent_shape=(1, 128, 3, 12, 20), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + self.assertEqual(source, "default_static") + self.assertEqual(camera_conditions.shape, (1, 3, 20)) + self.assertEqual(chunk_plucker.shape, (1, 48, 3, 12, 20)) + self.assertTrue( + torch.equal(camera_conditions[0, 0, :16], torch.eye(4).reshape(-1)) + ) + self.assertTrue( + torch.allclose( + camera_conditions[0, 0, 16:], + torch.tensor([16.0, 16.0, 10.0, 6.0]), + ) + ) + + def test_request_intrinsics_are_transformed_for_condition_image_crop(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + batch = SimpleNamespace( + extra={ + "diffusers_kwargs": { + "intrinsics": [50.0, 50.0, 50.0, 50.0], + }, + "sana_wm_condition_image_preprocess": { + "source_size": (100, 100), + "resized_size": (1280, 1280), + "crop_offset": (0, 288), + "target_size": (1280, 704), + }, + }, + height=704, + width=1280, + ) + + camera_conditions, chunk_plucker, source = stage._build_camera_conditioning( + batch, + batch_size=1, + num_frames=49, + latent_shape=(1, 128, 7, 22, 40), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + self.assertEqual(source, "default_static_request_intrinsics") + self.assertEqual(chunk_plucker.shape, (1, 48, 7, 22, 40)) + self.assertTrue( + torch.allclose( + camera_conditions[0, 0, 16:], + torch.tensor([20.0, 20.0, 20.0, 11.0]), + ) + ) + + def test_explicit_camera_request_detects_diffusers_path_kwargs(self) -> None: + batch = SimpleNamespace( + extra={"diffusers_kwargs": {"camera_to_world_path": "/tmp/cam.npy"}} + ) + + self.assertTrue(SanaWMBeforeDenoisingStage._has_explicit_camera_request(batch)) + + def test_explicit_camera_request_detects_action_kwargs(self) -> None: + batch = SimpleNamespace(extra={"diffusers_kwargs": {"action": "w-8"}}) + + self.assertTrue(SanaWMBeforeDenoisingStage._has_explicit_camera_request(batch)) + + def test_action_conditioning_uses_existing_camera_path(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + batch = SimpleNamespace( + extra={ + "diffusers_kwargs": { + "action": "w-8", + "intrinsics": [50.0, 50.0, 32.0, 32.0], + } + }, + height=64, + width=64, + ) + + camera_conditions, chunk_plucker, source = stage._build_camera_conditioning( + batch, + batch_size=1, + num_frames=9, + latent_shape=(1, 128, 2, 2, 2), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + self.assertEqual(source, "action") + self.assertEqual(camera_conditions.shape, (1, 2, 20)) + self.assertEqual(chunk_plucker.shape, (1, 48, 2, 2, 2)) + # Z translation after 8 forward frames at the default (streaming) + # translation speed. Reference the constant so the assertion cannot go + # stale again if the default moves (it was 0.05 once, 0.04 now). + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.base import ( + _SANA_WM_DEFAULT_TRANSLATION_SPEED, + ) + + self.assertAlmostEqual( + float(camera_conditions[0, 1, 11]), + 8 * _SANA_WM_DEFAULT_TRANSLATION_SPEED, + places=5, + ) + + def test_camera_conditioning_accepts_unbatched_chunk_plucker(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + batch = SimpleNamespace( + extra={"chunk_plucker": torch.zeros(48, 3, 12, 20)}, + height=384, + width=640, + ) + + camera_conditions, chunk_plucker, source = stage._build_camera_conditioning( + batch, + batch_size=1, + num_frames=17, + latent_shape=(1, 128, 3, 12, 20), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + self.assertEqual(source, "default_static") + self.assertEqual(camera_conditions.shape, (1, 3, 20)) + self.assertEqual(chunk_plucker.shape, (1, 48, 3, 12, 20)) + + def test_camera_conditioning_accepts_diffusers_prebuilt_raymap(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + camera_conditions = torch.zeros(1, 17, 20) + camera_conditions[..., :16] = torch.eye(4).reshape(1, 1, 16) + camera_conditions[..., 16:] = torch.tensor([16.0, 16.0, 10.0, 6.0]) + batch = SimpleNamespace( + extra={"diffusers_kwargs": {"camera_conditions": camera_conditions}}, + height=384, + width=640, + ) + + camera_conditions, chunk_plucker, source = stage._build_camera_conditioning( + batch, + batch_size=1, + num_frames=17, + latent_shape=(1, 128, 3, 12, 20), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + self.assertEqual(source, "prebuilt_original_frames") + self.assertEqual(camera_conditions.shape, (1, 3, 20)) + self.assertEqual(chunk_plucker.shape, (1, 48, 3, 12, 20)) + + def test_camera_conditioning_rejects_bad_prepacked_shapes(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + + with self.assertRaisesRegex(ValueError, "camera_conditions must have shape"): + stage._build_camera_conditioning( + SimpleNamespace( + extra={"camera_conditions": torch.zeros(1, 1, 3, 20)}, + height=384, + width=640, + ), + batch_size=1, + num_frames=17, + latent_shape=(1, 128, 3, 12, 20), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + with self.assertRaisesRegex(ValueError, "batch dimension"): + stage._build_camera_conditioning( + SimpleNamespace( + extra={"camera_conditions": torch.zeros(3, 3, 20)}, + height=384, + width=640, + ), + batch_size=2, + num_frames=17, + latent_shape=(2, 128, 3, 12, 20), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + with self.assertRaisesRegex(ValueError, "require chunk_plucker"): + stage._build_camera_conditioning( + SimpleNamespace( + extra={"camera_conditions": torch.zeros(1, 3, 20)}, + height=384, + width=640, + ), + batch_size=1, + num_frames=17, + latent_shape=(1, 128, 3, 12, 20), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + with self.assertRaisesRegex(ValueError, "batch dimension"): + stage._build_camera_conditioning( + SimpleNamespace( + extra={"chunk_plucker": torch.zeros(3, 48, 3, 12, 20)}, + height=384, + width=640, + ), + batch_size=2, + num_frames=17, + latent_shape=(2, 128, 3, 12, 20), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + with self.assertRaisesRegex(ValueError, "chunk_plucker shape mismatch"): + stage._build_camera_conditioning( + SimpleNamespace( + extra={"chunk_plucker": torch.zeros(48, 2, 12, 20)}, + height=384, + width=640, + ), + batch_size=1, + num_frames=17, + latent_shape=(1, 128, 3, 12, 20), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + def test_action_and_camera_path_are_mutually_exclusive(self) -> None: + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=None, + pipeline_config=SanaWMPipelineConfig(), + ) + batch = SimpleNamespace( + extra={ + "diffusers_kwargs": { + "action": "w-8", + "camera_to_world": torch.eye(4).unsqueeze(0), + } + }, + height=64, + width=64, + ) + + with self.assertRaisesRegex(ValueError, "mutually exclusive"): + stage._build_camera_conditioning( + batch, + batch_size=1, + num_frames=9, + latent_shape=(1, 128, 2, 2, 2), + device=torch.device("cpu"), + dtype=torch.float32, + ) + + def test_action_num_frames_uses_longest_batched_action(self) -> None: + batch = SimpleNamespace(extra={"diffusers_kwargs": {"action": ["w-8", "w-16"]}}) + + self.assertEqual( + SanaWMBeforeDenoisingStage._action_num_frames_for_request(batch), + 17, + ) + + def test_intrinsics_3x3_sequence_can_exceed_requested_frame_count(self) -> None: + intrinsics = torch.eye(3).unsqueeze(0).repeat(321, 1, 1) + + coerced = SanaWMBeforeDenoisingStage._coerce_intrinsics_vec4( + intrinsics, + batch_size=1, + num_frames=49, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + self.assertEqual(coerced.shape, (1, 49, 4)) + self.assertTrue(torch.equal(coerced[0, 0], torch.tensor([1.0, 1.0, 0.0, 0.0]))) + + def test_latent_frame_camera_conditions_samples_stride_indices(self) -> None: + original = torch.arange(1 * 17 * 20, dtype=torch.float32).reshape(1, 17, 20) + + sampled = SanaWMBeforeDenoisingStage._latent_frame_camera_conditions( + original, + num_frames=17, + latent_frames=3, + vae_temporal_stride=8, + ) + + self.assertTrue(torch.equal(sampled[:, 0], original[:, 0])) + self.assertTrue(torch.equal(sampled[:, 1], original[:, 8])) + self.assertTrue(torch.equal(sampled[:, 2], original[:, 16])) + + def test_chunk_plucker_accepts_ltx_frame_count(self) -> None: + camera_conditions = torch.zeros(1, 17, 20) + camera_conditions[..., :16] = torch.eye(4).reshape(1, 1, 16) + camera_conditions[..., 16:] = torch.tensor([16.0, 16.0, 10.0, 6.0]) + + chunk_plucker = compute_chunk_plucker( + camera_conditions, + HW=(3, 12, 20), + vae_temporal_stride=8, + patch_size=(1, 1, 1), + ) + + self.assertEqual(chunk_plucker.shape, (1, 48, 3, 12, 20)) + + def test_post_ucpe_rms_stabilization_clamps_inflated_tensors(self) -> None: + ref = torch.ones(1, 2, 4, 3) + transformed = ref * 8.0 + + stabilized = _downscale_to_reference_rms(ref, transformed) + + ref_rms = ref.square().mean(dim=2, keepdim=True).sqrt() + stabilized_rms = stabilized.square().mean(dim=2, keepdim=True).sqrt() + self.assertTrue(torch.all(stabilized_rms <= ref_rms + 1e-5)) + + def test_rmsnorm_scale_factor_initializes_weight(self) -> None: + norm = _RMSNorm(4, scale_factor=0.01) + self.assertTrue(torch.allclose(norm.weight, torch.full((4,), 0.01))) + + def test_prepare_timesteps_uses_inference_flow_shift(self) -> None: + class FakeScheduler: + def __init__(self): + self.shift = None + self.timesteps = None + self.sigmas = None + + def set_timesteps(self, num_inference_steps, device, shift=None): + self.shift = shift + self.timesteps = torch.arange(num_inference_steps, device=device) + self.sigmas = torch.tensor([float(shift), 0.0], device=device) + + scheduler = FakeScheduler() + stage = SanaWMBeforeDenoisingStage( + vae=None, + transformer=None, + scheduler=scheduler, + pipeline_config=SanaWMPipelineConfig(), + ) + batch = SimpleNamespace(num_inference_steps=3, scheduler=None) + + stage._prepare_timesteps(batch, SimpleNamespace(), torch.device("cpu")) + + self.assertEqual(scheduler.shift, 9.8) + self.assertTrue(torch.equal(batch.timesteps, torch.arange(3))) + + +class TestSanaWMTextEncodingStage(unittest.TestCase): + def test_official_prompt_window_keeps_bos_and_tail(self) -> None: + tensor = torch.arange(5).reshape(1, 5, 1) + + selected = SanaWMTextEncodingStage._select_official_prompt_window( + tensor, max_length=3 + ) + + self.assertTrue(torch.equal(selected.squeeze(-1), torch.tensor([[0, 3, 4]]))) + + +class TestSanaWMDenoisingStage(unittest.TestCase): + def test_parallelism_type_follows_cfg_parallel_flag(self) -> None: + stage = object.__new__(SanaWMDenoisingStage) + + stage.server_args = SimpleNamespace(enable_cfg_parallel=False) + self.assertEqual(stage.parallelism_type, StageParallelismType.REPLICATED) + + stage.server_args = SimpleNamespace(enable_cfg_parallel=True) + self.assertEqual(stage.parallelism_type, StageParallelismType.CFG_PARALLEL) + + def test_cfg_parallel_formula_matches_serial_cfg(self) -> None: + pos = torch.tensor([2.0]) + neg = torch.tensor([-1.0]) + guidance_scale = 4.5 + serial = neg + guidance_scale * (pos - neg) + + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.base.cfg_model_parallel_all_reduce", + side_effect=lambda partial: partial + (1.0 - guidance_scale) * neg, + ): + combined_from_pos_rank = SanaWMDenoisingStage._combine_cfg_parallel_noise( + pos, guidance_scale, cfg_rank=0 + ) + + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.base.cfg_model_parallel_all_reduce", + side_effect=lambda partial: partial + guidance_scale * pos, + ): + combined_from_neg_rank = SanaWMDenoisingStage._combine_cfg_parallel_noise( + neg, guidance_scale, cfg_rank=1 + ) + + self.assertTrue(torch.allclose(combined_from_pos_rank, serial)) + self.assertTrue(torch.allclose(combined_from_neg_rank, serial)) + + +class TestSanaWMNativeDiTChunking(unittest.TestCase): + def test_softmax_chunking_is_disabled_by_default_for_upstream_parity(self) -> None: + arch = SanaWMConfig().arch_config + + self.assertFalse(arch.use_chunked_softmax_attention) + + def test_main_gdn_backend_defaults_to_auto(self) -> None: + arch = SanaWMConfig().arch_config + + self.assertEqual(arch.gdn_backend, "auto") + + def test_first_chunk_plus_one_chunk_indices_match_upstream(self) -> None: + self.assertEqual( + _sana_wm_chunk_index_from_chunk_size( + 21, 4, strategy="first_chunk_plus_one" + ), + [0, 5, 9, 13, 17], + ) + + def test_normalize_chunk_index_adds_start_and_final_boundary(self) -> None: + self.assertEqual(_sana_wm_normalize_chunk_index([3], 5), [0, 3, 5]) + + def test_chunked_attention_matches_prefix_attention_per_chunk(self) -> None: + torch.manual_seed(0) + q = torch.randn(1, 5, 1, 4) + k = torch.randn(1, 5, 1, 4) + v = torch.randn(1, 5, 1, 4) + scale = 0.5 + + out = _sana_wm_chunked_attention( + q, + k, + v, + HW=(5, 1, 1), + chunk_size=2, + chunk_split_strategy="first_chunk_plus_one", + chunk_index=None, + softmax_scale=scale, + ) + + first = torch.nn.functional.scaled_dot_product_attention( + q[:, :3].transpose(1, 2), + k[:, :3].transpose(1, 2), + v[:, :3].transpose(1, 2), + dropout_p=0.0, + scale=scale, + ).transpose(1, 2) + second = torch.nn.functional.scaled_dot_product_attention( + q[:, 3:].transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + dropout_p=0.0, + scale=scale, + ).transpose(1, 2) + expected = torch.cat([first, second], dim=1) + + self.assertTrue(torch.allclose(out, expected)) + + def test_chunked_attention_returns_none_when_bidirectional(self) -> None: + q = torch.randn(1, 3, 1, 4) + + out = _sana_wm_chunked_attention( + q, + q, + q, + HW=(3, 1, 1), + chunk_size=10, + chunk_split_strategy="first_chunk_plus_one", + chunk_index=None, + softmax_scale=0.5, + ) + + self.assertIsNone(out) + + def test_gdn_chunk_scan_matches_recurrent_scan(self) -> None: + torch.manual_seed(0) + B, H, D, T, S = 1, 2, 4, 5, 3 + N = T * S + q = torch.randn(B, H, D, N, dtype=torch.float64) + k = torch.randn(B, H, D, N, dtype=torch.float64) + v = torch.randn(B, H, D, N, dtype=torch.float64) + q_rot = torch.randn(B, H, D, N, dtype=torch.float64) + k_rot = torch.randn(B, H, D, N, dtype=torch.float64) + beta = torch.sigmoid(torch.randn(B, H, T, S, dtype=torch.float64)) + decay = torch.sigmoid(torch.randn(B, H, T, dtype=torch.float64)) + + recurrent = _gdn_scan_forward( + q, + k, + v, + q_rot, + k_rot, + beta, + decay, + return_components=True, + ) + chunked = _gdn_chunk_scan_forward( + q, + k, + v, + q_rot, + k_rot, + beta, + decay, + chunk_size=2, + return_components=True, + ) + + self.assertTrue(torch.allclose(chunked[0], recurrent[0], atol=1e-9)) + self.assertTrue(torch.allclose(chunked[1], recurrent[1], atol=1e-9)) + + def test_camera_chunk_scan_matches_recurrent_scan(self) -> None: + torch.manual_seed(0) + B, H, D, T, S = 1, 2, 4, 5, 3 + N = T * S + q_rot = torch.randn(B, H, D, N, dtype=torch.float64) + k_rot = torch.randn(B, H, D, N, dtype=torch.float64) + v = torch.randn(B, H, D, N, dtype=torch.float64) + beta = torch.sigmoid(torch.randn(B, H, T, S, dtype=torch.float64)) + decay = torch.sigmoid(torch.randn(B, H, T, dtype=torch.float64)) + + recurrent = _single_path_delta_scan_forward(q_rot, k_rot, v, beta, decay) + chunked = _single_path_delta_chunk_scan_forward( + q_rot, + k_rot, + v, + beta, + decay, + chunk_size=2, + ) + + self.assertTrue(torch.allclose(chunked, recurrent, atol=1e-9)) + + +class TestSanaWMRefinerStage(_GlobalStageArgsMixin, unittest.TestCase): + def test_diffusers_refiner_detection_uses_official_class_name(self) -> None: + class LTX2VideoTransformer3DModel(torch.nn.Module): + pass + + self.assertTrue(_uses_diffusers_ltx2_refiner(LTX2VideoTransformer3DModel())) + self.assertTrue( + _uses_diffusers_ltx2_refiner( + OfficialDiffusersLTX2RefinerModule(LTX2VideoTransformer3DModel()) + ) + ) + + def test_refiner_config_value_prefers_diffusers_config(self) -> None: + module = SimpleNamespace( + patch_size=999, + config=SimpleNamespace(patch_size=1), + ) + + self.assertEqual(_refiner_config_value(module, "patch_size"), 1) + + def test_official_refiner_wrappers_expose_layerwise_blocks(self) -> None: + self.assertEqual( + OfficialDiffusersLTX2RefinerModule.layer_names, + ["module.transformer_blocks"], + ) + self.assertIn( + "module.model.language_model.layers", + OfficialGemma3TextEncoderModule.layer_names, + ) + + def test_refiner_component_uses_follow_execution_order(self) -> None: + stage = SanaWMLTX2RefinerStage( + transformer=torch.nn.Identity(), + connectors=torch.nn.Identity(), + text_encoder=torch.nn.Identity(), + tokenizer=SimpleNamespace(pad_token="", eos_token=""), + dtype=torch.bfloat16, + ) + + names = [ + use.component_name + for use in stage.component_uses( + SimpleNamespace(), stage_name="sana_wm_refiner" + ) + ] + + self.assertEqual(names, ["text_encoder_2", "connectors", "transformer_2"]) + + def test_refiner_parallelism_runs_on_main_rank_when_cfg_parallel(self) -> None: + stage = SanaWMLTX2RefinerStage( + transformer=torch.nn.Identity(), + connectors=torch.nn.Identity(), + text_encoder=torch.nn.Identity(), + tokenizer=SimpleNamespace(pad_token="", eos_token=""), + dtype=torch.bfloat16, + ) + + stage.server_args = SimpleNamespace(enable_cfg_parallel=False) + self.assertEqual(stage.parallelism_type, StageParallelismType.REPLICATED) + + stage.server_args = SimpleNamespace(enable_cfg_parallel=True) + self.assertEqual(stage.parallelism_type, StageParallelismType.MAIN_RANK_ONLY) + + def test_refiner_component_uses_skip_non_main_cfg_rank(self) -> None: + stage = SanaWMLTX2RefinerStage( + transformer=torch.nn.Identity(), + connectors=torch.nn.Identity(), + text_encoder=torch.nn.Identity(), + tokenizer=SimpleNamespace(pad_token="", eos_token=""), + dtype=torch.bfloat16, + ) + + with ( + patch( + f"{_SANA_WM_REFINER_STAGE_MODULE}.torch.distributed.is_available", + return_value=True, + ), + patch( + f"{_SANA_WM_REFINER_STAGE_MODULE}.torch.distributed.is_initialized", + return_value=True, + ), + patch( + f"{_SANA_WM_REFINER_STAGE_MODULE}.get_classifier_free_guidance_rank", + return_value=1, + ), + ): + uses = stage.component_uses( + SimpleNamespace(enable_cfg_parallel=True), + stage_name="sana_wm_refiner", + ) + + self.assertEqual(uses, []) + + def test_streaming_diffusers_attention_accepts_ungated_ltx2_attention(self) -> None: + """Diffusers 0.37 LTX2Attention omits `to_gate_logits` for ungated configs.""" + + class IdentityLinear(torch.nn.Module): + def forward(self, x): + return x + + class FakeProcessor: + _attention_backend = None + _parallel_config = None + + class UngatedAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.to_q = IdentityLinear() + self.to_k = IdentityLinear() + self.to_v = IdentityLinear() + self.norm_q = IdentityLinear() + self.norm_k = IdentityLinear() + self.to_out = torch.nn.ModuleList( + [IdentityLinear(), torch.nn.Identity()] + ) + self.heads = 2 + self.rope_type = "split" + self.processor = FakeProcessor() + + def dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=None, + parallel_config=None, + ): + return query + + attention_dispatch = types.ModuleType("diffusers.models.attention_dispatch") + attention_dispatch.dispatch_attention_fn = dispatch_attention_fn + transformer_ltx2 = types.ModuleType( + "diffusers.models.transformers.transformer_ltx2" + ) + transformer_ltx2.apply_interleaved_rotary_emb = lambda x, _rope: x + transformer_ltx2.apply_split_rotary_emb = lambda x, _rope: x + + with patch.dict( + sys.modules, + { + "diffusers": types.ModuleType("diffusers"), + "diffusers.models": types.ModuleType("diffusers.models"), + "diffusers.models.attention_dispatch": attention_dispatch, + "diffusers.models.transformers": types.ModuleType( + "diffusers.models.transformers" + ), + "diffusers.models.transformers.transformer_ltx2": transformer_ltx2, + }, + ): + hidden_states = torch.randn(1, 3, 4) + rotary = (torch.empty(0), torch.empty(0)) + out = _streaming_diffusers_self_attention( + attn=UngatedAttention(), + hidden_states=hidden_states, + query_rotary_emb=rotary, + n_context_tokens=1, + ) + + self.assertEqual(out.shape, hidden_states.shape) + + def test_skip_refiner_flag_accepts_request_extra(self) -> None: + batch = SimpleNamespace(extra={"diffusers_kwargs": {"skip_refiner": True}}) + self.assertTrue(sana_wm_skip_refiner_enabled(batch)) + + def test_prompt_resolution_broadcasts_single_prompt(self) -> None: + batch = Req(prompt="drive forward") + prompts = SanaWMLTX2RefinerStage._prompts_for_batch(batch, batch_size=2) + self.assertEqual(prompts, ["drive forward", "drive forward"]) + + def test_prompt_resolution_accepts_batch_prompt_list(self) -> None: + batch = Req(prompt=["left", "right"]) + prompts = SanaWMLTX2RefinerStage._prompts_for_batch(batch, batch_size=2) + self.assertEqual(prompts, ["left", "right"]) + + def test_refiner_prompt_encoding_uses_hf_gemma_backbone(self) -> None: + class DummyTokenizer: + padding_side = "right" + pad_token = None + eos_token = "" + + def __call__(self, *args, **kwargs): + return SimpleNamespace( + input_ids=torch.tensor([[1, 2, 0, 0]]), + attention_mask=torch.tensor([[1, 1, 0, 0]]), + ) + + class DummyBackbone: + def __init__(self): + self.called = False + + def __call__(self, **kwargs): + self.called = True + hidden0 = torch.arange(8, dtype=torch.float32).reshape(1, 4, 2) + hidden1 = hidden0 + 10 + return SimpleNamespace(hidden_states=(hidden0, hidden1)) + + class DummyTextEncoder: + def __init__(self): + self.called = False + self.model = DummyBackbone() + + def __call__(self, **kwargs): + self.called = True + raise AssertionError( + "Gemma3ForConditionalGeneration.model was not used" + ) + + class DummyConnectors: + def __call__(self, prompt_embeds, attention_mask): + return prompt_embeds, None, attention_mask + + stage = SanaWMLTX2RefinerStage( + transformer=torch.nn.Identity(), + connectors=DummyConnectors(), + text_encoder=DummyTextEncoder(), + tokenizer=DummyTokenizer(), + dtype=torch.float32, + text_max_sequence_length=4, + ) + + prompt_embeds, attention_mask = stage._encode_prompt( + "drive forward", torch.device("cpu") + ) + + self.assertTrue(stage.text_encoder.model.called) + self.assertFalse(stage.text_encoder.called) + self.assertEqual(prompt_embeds.shape, (1, 4, 4)) + self.assertTrue(torch.equal(attention_mask, torch.tensor([[1, 1, 0, 0]]))) + + def test_refiner_decoding_drops_clean_sink_frame_after_decode(self) -> None: + stage = SanaWMRefinerDecodingStage(vae=None) + decoded = torch.arange(1 * 3 * 4 * 2 * 2).reshape(1, 3, 4, 2, 2) + + with patch.object(DecodingStage, "decode", return_value=decoded): + frames = stage.decode( + torch.empty(1, 128, 4, 2, 2), + SimpleNamespace(pipeline_config=SanaWMPipelineConfig()), + vae_dtype=torch.bfloat16, + ) + + self.assertTrue(torch.equal(frames, decoded[:, :, 1:])) + + def test_refiner_decoding_keeps_sink_frame_when_refiner_skipped(self) -> None: + stage = SanaWMRefinerDecodingStage(vae=None) + decoded = torch.arange(1 * 3 * 4 * 2 * 2).reshape(1, 3, 4, 2, 2) + stage._drop_refiner_sink = False + + with patch.object(DecodingStage, "decode", return_value=decoded): + frames = stage.decode( + torch.empty(1, 128, 4, 2, 2), + SimpleNamespace(pipeline_config=SanaWMPipelineConfig()), + vae_dtype=torch.bfloat16, + ) + + self.assertTrue(torch.equal(frames, decoded)) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py new file mode 100644 index 000000000..cbed7d8bc --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Realtime chain stages — latent-prep plan arithmetic + noise discipline. + +The chunk PLAN must reproduce the engine-era tick behavior exactly: +front-loaded segments within a fixed horizon (chunk 0 carries the remainder +and the conditioning frame), uniform chunks past the horizon (seamless +continuation), and the noise buffer sliced bitwise within the horizon. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from PIL import Image + +from sglang.multimodal_gen.runtime import server_args as _sa_mod +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_chain import ( + SanaWMCameraCondStage, + SanaWMNoiseState, + SanaWMRealtimeLatentPrepStage, + SanaWMSessionInputsState, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_stage import ( + SanaWMRealtimeStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import ( + SanaWMStreamCacheState, +) +from sglang.multimodal_gen.runtime.realtime.session import RealtimeSession +from sglang.multimodal_gen.runtime.server_args import set_global_server_args + +MC = 8 + + +class _TestRealtimeStage(SanaWMRealtimeStage): + def forward(self, batch, server_args): + raise NotImplementedError + + +@pytest.fixture +def _global_args(): + prev = _sa_mod._global_server_args + set_global_server_args( + SimpleNamespace( + comfyui_mode=False, + enable_cfg_parallel=False, + enable_torch_compile=False, + attention_backend=None, + ) + ) + try: + yield + finally: + set_global_server_args(prev) + + +def _prep_stage(): + return SanaWMRealtimeLatentPrepStage( + use_refiner=True, transformer=None, vae=None, model_path="" + ) + + +def _realtime_stage(): + return object.__new__(_TestRealtimeStage) + + +def _camera_stage(): + return object.__new__(SanaWMCameraCondStage) + + +def _batch(session, block_idx, image_latent): + return SimpleNamespace( + session=session, + block_idx=block_idx, + image_latent=image_latent, + seed=7, + generator=None, + extra={}, + latents=None, + ) + + +def _server_args(): + return SimpleNamespace(pipeline_config=SimpleNamespace(dit_precision="fp32")) + + +def test_realtime_intrinsics_default_to_centered_heuristic(): + stage = _realtime_stage() + state = SimpleNamespace( + intrinsics_raw=None, + intrinsics_image=Image.new("RGB", (832, 480)), + ) + + intrinsics = stage._prepare_intrinsics( + SimpleNamespace(condition_inputs={}), + state, + num_frames=3, + device=torch.device("cpu"), + ) + + assert intrinsics.shape == (3, 4) + assert intrinsics[0].tolist() == pytest.approx([665.6, 665.6, 416.0, 240.0]) + assert intrinsics[2].tolist() == pytest.approx([665.6, 665.6, 416.0, 240.0]) + + +def test_realtime_first_frame_uses_requested_size(): + stage = _realtime_stage() + batch = SimpleNamespace( + condition_image=Image.new("RGB", (640, 360)), + image_path=None, + height=480, + width=832, + ) + + cropped, original, src_size, resized_size, crop_offset = stage._prepare_image(batch) + + assert cropped.size == (832, 480) + assert original.size == (640, 360) + assert src_size == (640, 360) + assert resized_size == (853, 480) + assert crop_offset == (10, 0) + + +def test_realtime_camera_conditioning_uses_requested_size(): + stage = _camera_stage() + inputs = SanaWMSessionInputsState() + inputs.src_size = (640, 360) + inputs.resized_size = (853, 480) + inputs.crop_offset = (10, 0) + inputs.target_height = 480 + inputs.target_width = 832 + inputs.intrinsics_image = Image.new("RGB", (640, 360)) + inputs.open_ended = True + batch = SimpleNamespace( + condition_inputs={}, + extra={}, + height=480, + width=832, + num_frames=17, + ) + + camera, plucker = stage._build_camera_windows( + batch, + inputs, + target_latent=3, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + assert camera.shape == (1, 3, 20) + assert plucker.shape == (1, 48, 3, 15, 26) + + +def test_latent_prep_plan_and_noise_discipline(_global_args): + stage = _prep_stage() + session = RealtimeSession() + fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32) + + inputs = session.get_or_create_state(SanaWMSessionInputsState) + inputs.latent_t = 5 # fixed horizon: segments [0, 3, 5] for nfpb=2 + inputs.num_frame_per_block = 2 + inputs.sink_size = 1 + cache = session.get_or_create_state(SanaWMStreamCacheState) + + # Tick 0: front-loaded chunk 0 (cond + remainder) + full-horizon buffer. + batch = stage.forward(_batch(session, 0, fl), _server_args()) + noise = session.get_or_create_state(SanaWMNoiseState) + assert noise.segments == [0, 3, 5] + assert noise.noise_buffer is not None and noise.noise_buffer.shape[2] == 5 + assert batch.extra["sana_wm_chunk_plan"] == [3] + assert batch.latents.shape[2] == 3 + assert torch.equal(batch.latents[:, :, :1].cpu(), fl) # cond frame in front + assert torch.equal( # noise sliced from the buffer, bitwise + batch.latents[:, :, 1:].cpu(), noise.noise_buffer[:, :, 1:3].cpu() + ) + + # Tick 1: stage-1 advanced to frame 3 (simulated) -> next segment [3, 5). + cache.chunk_indices = [0, 3] + cache.chunk_idx = 1 + batch = stage.forward(_batch(session, 1, fl), _server_args()) + assert batch.extra["sana_wm_chunk_plan"] == [2] + assert torch.equal(batch.latents.cpu(), noise.noise_buffer[:, :, 3:5].cpu()) + + # Tick 2: horizon exhausted -> seamless continuation, uniform chunk from + # the seeded fallback generator (no reset, no rollover). + cache.chunk_indices = [0, 3, 5] + cache.chunk_idx = 2 + batch = stage.forward(_batch(session, 2, fl), _server_args()) + assert batch.extra["sana_wm_chunk_plan"] == [2] + assert batch.latents.shape[2] == 2 + assert torch.isfinite(batch.latents).all() + + +def test_latent_prep_open_ended_uniform_chunk0(_global_args): + stage = _prep_stage() + session = RealtimeSession() + fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32) + + inputs = session.get_or_create_state(SanaWMSessionInputsState) + inputs.latent_t = None # open-ended + inputs.num_frame_per_block = 3 + inputs.sink_size = 1 + session.get_or_create_state(SanaWMStreamCacheState) + + batch = stage.forward(_batch(session, 0, fl), _server_args()) + noise = session.get_or_create_state(SanaWMNoiseState) + assert noise.segments is None and noise.noise_buffer is None + # Uniform grid chunk 0 = cond frame + nfpb new frames. + assert batch.extra["sana_wm_chunk_plan"] == [4] + assert batch.latents.shape[2] == 4 + assert torch.equal(batch.latents[:, :, :1].cpu(), fl) diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_cached.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_cached.py new file mode 100644 index 000000000..48f14bab5 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_cached.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +"""S1a-2a tests — chunk-causal cached scan combiners for streaming `forward_long`. + +`_gdn_scan_cached` / `_single_path_delta_scan_cached` are the per-chunk combiners +the streaming path calls: a FORWARD (inclusive) pass that carries recurrent state +across chunks plus a BACKWARD (exclusive) pass recomputed intra-chunk. Two +invariants pin them down: + +1. **Reduce-to-bidirectional** (#26153-only): a single chunk with no carried + state must equal the validated dense `_gdn_scan_bidirectional` / + `_single_path_delta_scan_bidirectional`. +2. **Port-fidelity** (vs the reference `minimal-sanawm`): output AND carried state + must match `recurrent_gdn_cached` / `recurrent_delta_cached` chunk-for-chunk, + including a 2-chunk state-carry run. + +FP64, CPU. +""" + +from __future__ import annotations + +import pathlib +import sys + +import pytest +import torch + +from sglang.multimodal_gen.runtime.models.dits.sana_wm import ( + BidirectionalGDNUCPESinglePathLiteLA, + GLUMBConvTemp, + _gdn_scan_bidirectional, + _gdn_scan_cached, + _ShortConvolution, + _single_path_delta_scan_bidirectional, + _single_path_delta_scan_cached, + _temporal_short_conv_cached, +) + +B, H, D, T, S = 1, 2, 4, 6, 3 +N = T * S +HW = (T, S, 1) # H_sp * W_sp == S + + +def _rand(*shape): + return torch.randn(*shape, dtype=torch.float64) + + +def _inputs(): + torch.manual_seed(0) + q, k, v = _rand(B, H, D, N), _rand(B, H, D, N), _rand(B, H, D, N) + q_rot, k_rot = _rand(B, H, D, N), _rand(B, H, D, N) + beta = torch.rand(B, H, T, S, dtype=torch.float64) + decay = torch.rand(B, H, T, dtype=torch.float64) * 0.5 + 0.4 + return q, k, v, q_rot, k_rot, beta, decay + + +def _frame_slice(x, f0, f1): + # x is (B, H, D, N=T*S), frame-major (T outer, S inner) + return x[..., f0 * S : f1 * S].contiguous() + + +# --------------------------------------------------------------------------- # +# 1. Reduce-to-bidirectional (uses only #26153 functions) +# --------------------------------------------------------------------------- # + + +def test_gdn_cached_single_chunk_reduces_to_bidirectional(): + q, k, v, q_rot, k_rot, beta, decay = _inputs() + ref = _gdn_scan_bidirectional(q, k, v, q_rot, k_rot, beta, decay, HW) + out, (kv, z) = _gdn_scan_cached(q, k, v, q_rot, k_rot, beta, decay) + torch.testing.assert_close(out, ref, atol=1e-9, rtol=0) + assert kv.shape == (B, H, D, D) and z.shape == (B, H, D, 1) + + +def test_cam_cached_single_chunk_reduces_to_bidirectional(): + _, _, v, q_rot, k_rot, beta, decay = _inputs() + ref = _single_path_delta_scan_bidirectional(q_rot, k_rot, v, beta, decay, HW) + out, kv = _single_path_delta_scan_cached(q_rot, k_rot, v, beta, decay) + torch.testing.assert_close(out, ref, atol=1e-9, rtol=0) + assert kv.shape == (B, H, D, D) + + +# --------------------------------------------------------------------------- # +# 2. Port-fidelity vs reference (gold standard; skip if reference absent) +# --------------------------------------------------------------------------- # + +_REF_DIR = pathlib.Path("/sgl-workspace/sglang/myUtils/sana/minimal-sanawm") + + +def _load_reference(): + if not (_REF_DIR / "components.py").exists(): + pytest.skip(f"reference impl not found at {_REF_DIR}") + if str(_REF_DIR) not in sys.path: + sys.path.insert(0, str(_REF_DIR)) + try: + import components as ref # type: ignore + except Exception as exc: # pragma: no cover - environment dependent + pytest.skip(f"reference components.py not importable: {exc}") + return ref + + +def test_gdn_cached_matches_reference_single_and_chunked(): + ref = _load_reference() + q, k, v, q_rot, k_rot, beta, decay = _inputs() + + # Single chunk, no carried state. + mine, (kv_m, z_m) = _gdn_scan_cached(q, k, v, q_rot, k_rot, beta, decay) + r_out, kv_r, z_r, _ = ref.recurrent_gdn_cached( + q, k, v, q_rot, k_rot, beta, decay, eps=1e-6, kv_state=None, z_state=None + ) + torch.testing.assert_close(mine, r_out, atol=1e-9, rtol=0) + torch.testing.assert_close(kv_m, kv_r, atol=1e-9, rtol=0) + torch.testing.assert_close(z_m, z_r, atol=1e-9, rtol=0) + + # Two chunks carrying state — mine vs reference, chunk-for-chunk. + split = 2 + + def chunk(x, f0, f1): + return _frame_slice(x, f0, f1) + + args0 = ( + chunk(q, 0, split), + chunk(k, 0, split), + chunk(v, 0, split), + chunk(q_rot, 0, split), + chunk(k_rot, 0, split), + beta[:, :, :split], + decay[:, :, :split], + ) + args1 = ( + chunk(q, split, T), + chunk(k, split, T), + chunk(v, split, T), + chunk(q_rot, split, T), + chunk(k_rot, split, T), + beta[:, :, split:], + decay[:, :, split:], + ) + + m0, (mkv, mz) = _gdn_scan_cached(*args0) + m1, _ = _gdn_scan_cached(*args1, init_state_kv=mkv, init_state_z=mz) + + r0, rkv, rz, _ = ref.recurrent_gdn_cached( + *args0, eps=1e-6, kv_state=None, z_state=None + ) + r1, _, _, _ = ref.recurrent_gdn_cached(*args1, eps=1e-6, kv_state=rkv, z_state=rz) + + torch.testing.assert_close(m0, r0, atol=1e-9, rtol=0) + torch.testing.assert_close(m1, r1, atol=1e-9, rtol=0) + + +def test_cam_cached_matches_reference_single_and_chunked(): + ref = _load_reference() + _, _, v, q_rot, k_rot, beta, decay = _inputs() + + mine, kv_m = _single_path_delta_scan_cached(q_rot, k_rot, v, beta, decay) + r_out, kv_r = ref.recurrent_delta_cached(q_rot, k_rot, v, beta, decay, state=None) + torch.testing.assert_close(mine, r_out, atol=1e-9, rtol=0) + torch.testing.assert_close(kv_m, kv_r, atol=1e-9, rtol=0) + + split = 3 + a0 = ( + _frame_slice(q_rot, 0, split), + _frame_slice(k_rot, 0, split), + _frame_slice(v, 0, split), + beta[:, :, :split], + decay[:, :, :split], + ) + a1 = ( + _frame_slice(q_rot, split, T), + _frame_slice(k_rot, split, T), + _frame_slice(v, split, T), + beta[:, :, split:], + decay[:, :, split:], + ) + + m0, mkv = _single_path_delta_scan_cached(*a0) + m1, _ = _single_path_delta_scan_cached(*a1, init_state_kv=mkv) + r0, rkv = ref.recurrent_delta_cached(*a0, state=None) + r1, _ = ref.recurrent_delta_cached(*a1, state=rkv) + + torch.testing.assert_close(m0, r0, atol=1e-9, rtol=0) + torch.testing.assert_close(m1, r1, atol=1e-9, rtol=0) + + +# --------------------------------------------------------------------------- # +# 3. Cached short conv on K (cache slot 4) +# --------------------------------------------------------------------------- # + +C_CONV, KERN = 5, 4 + + +def _conv(): + conv = _ShortConvolution(C_CONV, KERN).double() + with torch.no_grad(): # randomize off the identity init for a non-trivial filter + conv.weight.copy_(torch.randn(C_CONV, 1, KERN, dtype=torch.float64)) + return conv + + +def _nslice(x, f0, f1): + # x is (B, N=T*S, C), frame-major + return x[:, f0 * S : f1 * S, :].contiguous() + + +def test_short_conv_cached_single_chunk_reduces_to_bidirectional(): + torch.manual_seed(1) + x = torch.randn(B, N, C_CONV, dtype=torch.float64) + conv = _conv() + ref = BidirectionalGDNUCPESinglePathLiteLA._temporal_short_conv( + x, conv, (T, S, 1), bidirectional=True + ) + out, prefix = _temporal_short_conv_cached(x, conv, (T, S, 1)) + torch.testing.assert_close(out, ref, atol=1e-9, rtol=0) + assert prefix.shape == (B * S, KERN - 1, C_CONV) + + +def test_short_conv_cached_forward_continuity_across_chunks(): + # The forward (causal) direction must be chunk-invariant via the prefix carry. + torch.manual_seed(2) + x = torch.randn(B, N, C_CONV, dtype=torch.float64) + conv = _conv() + whole, _ = _temporal_short_conv_cached(x, conv, (T, S, 1), bidirectional=False) + + split = 2 + o0, p0 = _temporal_short_conv_cached( + _nslice(x, 0, split), conv, (split, S, 1), bidirectional=False + ) + o1, _ = _temporal_short_conv_cached( + _nslice(x, split, T), conv, (T - split, S, 1), prefix=p0, bidirectional=False + ) + chunked = torch.cat([o0, o1], dim=1) # reassemble frame-major + torch.testing.assert_close(chunked, whole, atol=1e-9, rtol=0) + + +# --------------------------------------------------------------------------- # +# 4. Cached FFN temporal tail (GLUMBConvTemp, cache slot 9) +# --------------------------------------------------------------------------- # + +C_FFN, HID = 6, 8 +HFFN, WFFN = S, 1 # S spatial tokens per frame + + +def _ffn(): + m = GLUMBConvTemp(C_FFN, HID, t_kernel_size=3).double().eval() + with torch.no_grad(): # zero-init t_conv -> randomize for a non-trivial temporal filter + m.t_conv.weight.copy_(torch.randn_like(m.t_conv.weight)) + return m + + +def test_ffn_cached_no_prefix_reduces_to_dense(): + torch.manual_seed(3) + x = torch.randn(B, N, C_FFN, dtype=torch.float64) + m = _ffn() + dense = m(x, (T, HFFN, WFFN)) + cached, tail = m(x, (T, HFFN, WFFN), save_ffn_tail=True) + torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0) + assert tail.shape[2] == m.t_conv.kernel_size[0] // 2 + + +def test_ffn_cached_final_chunk_matches_whole(): + # The prefix supplies real left context, so the LAST chunk (right edge = end + # of sequence, same as the whole pass) matches the whole-sequence output. + torch.manual_seed(4) + x = torch.randn(B, N, C_FFN, dtype=torch.float64) + m = _ffn() + whole, _ = m(x, (T, HFFN, WFFN), save_ffn_tail=True) + + split = 2 + _, tail0 = m(_nslice(x, 0, split), (split, HFFN, WFFN), save_ffn_tail=True) + o1, _ = m( + _nslice(x, split, T), + (T - split, HFFN, WFFN), + ffn_tail=tail0, + save_ffn_tail=True, + ) + torch.testing.assert_close(whole[:, split * S :, :], o1, atol=1e-9, rtol=0) diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_forward_long.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_forward_long.py new file mode 100644 index 000000000..78c04201a --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_forward_long.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: Apache-2.0 +"""S1a-2c tests — streaming `forward_long` assembly on the #26153 SANA-WM DiT. + +Built up in stages, each pinned at fp64 / CPU / atol=1e-9: + Stage 0 — RoPE windowing: `WanRotaryPosEmbed` must produce freqs at GLOBAL + frame positions for a chunk `[start, end)`, equal to the corresponding slice + of the full table. This is the prerequisite that keeps a chunk's queries + aligned with the carried K and the softmax concat-window. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from sglang.multimodal_gen.configs.models.dits.sana_wm import ( + SanaWMArchConfig, + SanaWMConfig, +) +from sglang.multimodal_gen.runtime import server_args as _sa_mod +from sglang.multimodal_gen.runtime.models.dits.sana_wm import ( + _CACHE_TYPE_STATE, + _SLOT_CAM_K, + _SLOT_FFN_TCONV, + _SLOT_K, + _SLOT_TYPE_FLAG, + _SLOT_V, + BidirectionalGDNUCPESinglePathLiteLA, + SanaWMBlock, + SanaWMTransformer3DModel, + WanRotaryPosEmbed, + _build_ucpe_apply_fns, + _slice_rope_to_current_chunk, + process_camera_conditions_ucpe, +) +from sglang.multimodal_gen.runtime.server_args import set_global_server_args + +HEAD_DIM = 112 +H, W = 2, 3 +S = H * W +T = 6 + + +def _rope(): + return WanRotaryPosEmbed(attention_head_dim=HEAD_DIM, patch_size=(1, 1, 1)) + + +def test_rope_full_equals_range_from_zero(): + rope = _rope() + dense = rope((T, H, W), torch.device("cpu")) + ranged = rope(((0, T), H, W), torch.device("cpu")) + torch.testing.assert_close(ranged, dense, atol=1e-9, rtol=0) + + +def test_rope_window_equals_slice_of_full(): + rope = _rope() + start, end = 2, 5 + full = rope((end, H, W), torch.device("cpu")) # frames 0..end-1 + window = rope(((start, end), H, W), torch.device("cpu")) # frames start..end-1 + # full is frame-major (1, 1, end*S, D/2); slice frames [start:end]. + torch.testing.assert_close( + window, full[:, :, start * S : end * S, :], atol=1e-9, rtol=0 + ) + + +def test_rope_frame_index_overrides_to_global_positions(): + rope = _rope() + start, end = 2, 5 + fidx = torch.arange(start, end) + by_index = rope((end - start, H, W), torch.device("cpu"), frame_index=fidx) + by_range = rope(((start, end), H, W), torch.device("cpu")) + torch.testing.assert_close(by_index, by_range, atol=1e-9, rtol=0) + + +def test_slice_rope_to_current_chunk_is_noop_when_sized(): + rope = _rope() + freqs = rope(((1, 4), H, W), torch.device("cpu")) # 3 frames -> 3*S tokens + assert _slice_rope_to_current_chunk(freqs, 3 * S) is freqs + assert _slice_rope_to_current_chunk(None, 3 * S) is None + # wider table trimmed to the trailing chunk + wide = rope(((0, 4), H, W), torch.device("cpu")) # 4 frames + trimmed = _slice_rope_to_current_chunk(wide, 3 * S) + torch.testing.assert_close(trimmed, wide[:, :, -3 * S :, :], atol=1e-9, rtol=0) + + +# --------------------------------------------------------------------------- # +# Stage 1 — attention-level forward_long: cached branch methods +# --------------------------------------------------------------------------- # + +AB, AHEADS, ADIM = ( + 1, + 2, + 16, +) # head_dim/2 divisible by 4 (UCPE homog. 4-vectors), like 112 +AC = AHEADS * ADIM # in_dim == heads*head_dim +AHW = (T, H, W) +AN = T * S + + +def _arope(): + # RoPE table sized to the attention's head_dim (NOT the Stage-0 HEAD_DIM). + return WanRotaryPosEmbed(attention_head_dim=ADIM, patch_size=(1, 1, 1)) + + +def _attn(softmax_main=False): + m = ( + BidirectionalGDNUCPESinglePathLiteLA( + in_dim=AC, + heads=AHEADS, + head_dim=ADIM, + update_rule="torch_recurrent", # match the cached recurrent scan exactly + cam_update_rule="torch_recurrent", + softmax_main=softmax_main, + use_chunked_softmax_attention=softmax_main, + ) + .double() + .eval() + ) + return m + + +def _softmax_attn(): + # Build a GDN module (no LocalAttention -> no server-args/backend needed), + # then flip to softmax mode: softmax blocks have no short conv, and the + # cached softmax path only reads softmax_attn.softmax_scale (a scalar). + m = _attn(softmax_main=False) + m.softmax_main = True + m.conv_k = None + m.conv_k_cam = None + m.softmax_attn = SimpleNamespace(softmax_scale=ADIM**-0.5) + return m + + +def _x(): + torch.manual_seed(5) + return torch.randn(AB, AN, AC, dtype=torch.float64) + + +def _prope(start=0, end=T): + # Build UCPE apply fns co-windowed with freqs for frames [start, end). + torch.manual_seed(7) + cam = torch.randn(AB, T, 20, dtype=torch.float64)[:, start:end] + raymats = process_camera_conditions_ucpe( + cam, HW=(end - start, H, W), patch_size=(1, 1, 1) + ) + raymats_flat = raymats.reshape(AB, -1, 4, 4) + freqs = _arope()(((start, end), H, W), torch.device("cpu")) + return _build_ucpe_apply_fns(ADIM, raymats_flat, freqs), freqs + + +def _empty_cache(): + return [None] * 10 + + +def test_main_gdn_cached_single_chunk_reduces_to_dense(): + attn = _attn(softmax_main=False) + x = _x() + rope_emb = _arope()((T, H, W), torch.device("cpu")) + dense, _, _ = attn._main_branch_gdn(x, AHW, rope_emb) + cache = _empty_cache() + cached, _, _ = attn._main_branch_gdn_cached(x, AHW, rope_emb, cache, True) + torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0) + assert cache[_SLOT_K].shape == (AB, AHEADS, ADIM, ADIM) + assert cache[_SLOT_V].shape == (AB, AHEADS, ADIM, 1) + assert float(cache[_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_STATE + + +def test_cam_gdn_cached_single_chunk_reduces_to_dense(): + attn = _attn(softmax_main=False) + x = _x() + (apply_q, apply_kv, apply_o), _ = _prope() + _, beta, decay = attn._main_branch_gdn(x, AHW, None) + dense = attn._cam_branch(x, AHW, apply_q, apply_kv, apply_o, beta, decay) + cache = _empty_cache() + cached = attn._cam_branch_cached( + x, AHW, apply_q, apply_kv, apply_o, beta, decay, cache, True + ) + torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0) + assert cache[_SLOT_CAM_K].shape == (AB, AHEADS, ADIM, ADIM) + + +def test_softmax_main_cached_last_chunk_matches_whole(): + attn = _softmax_attn() + x = _x() + rope_full = _arope()(((0, T), H, W), torch.device("cpu")) + whole, _, _ = attn._main_branch_softmax_cached( + x, AHW, rope_full, _empty_cache(), True + ) + + split = 2 + rope0 = _arope()(((0, split), H, W), torch.device("cpu")) + rope1 = _arope()(((split, T), H, W), torch.device("cpu")) + cache = _empty_cache() + _o0, _, _ = attn._main_branch_softmax_cached( + x[:, : split * S, :], (split, H, W), rope0, cache, True + ) + o1, _, _ = attn._main_branch_softmax_cached( + x[:, split * S :, :], (T - split, H, W), rope1, cache, True + ) + # Last chunk's queries attend to [chunk0_K || chunk1_K] == full K. + torch.testing.assert_close(o1, whole[:, split * S :, :], atol=1e-9, rtol=0) + + +def test_forward_long_gdn_reduces_to_dense_no_camera(): + attn = _attn(softmax_main=False) + x = _x() + rope_emb = _arope()((T, H, W), torch.device("cpu")) + dense = attn(x, AHW, rope_emb, None) + cache = _empty_cache() + cached, ret = attn.forward_long( + x, AHW, rope_emb, None, kv_cache=cache, save_kv_cache=True + ) + torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0) + assert ret is cache and cache[_SLOT_K] is not None + + +def test_forward_long_gdn_reduces_to_dense_with_camera(): + attn = _attn(softmax_main=False) + x = _x() + (apply_q, apply_kv, apply_o), freqs = _prope() + prope_fns = (apply_q, apply_kv, apply_o) + dense = attn(x, AHW, freqs, prope_fns) + cache = _empty_cache() + cached, _ = attn.forward_long( + x, AHW, freqs, prope_fns, kv_cache=cache, save_kv_cache=True + ) + torch.testing.assert_close(cached, dense, atol=1e-9, rtol=0) + assert cache[_SLOT_CAM_K] is not None + + +# --------------------------------------------------------------------------- # +# Stage 2 — block-level forward_long (cross-attn stubbed; needs server args +# only to construct the cross-attn's LocalAttention) +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def _global_args(): + prev = _sa_mod._global_server_args + set_global_server_args( + SimpleNamespace( + comfyui_mode=False, + enable_cfg_parallel=False, + enable_torch_compile=False, + attention_backend=None, + ) + ) + try: + yield + finally: + set_global_server_args(prev) + + +class _ZeroCross(torch.nn.Module): + def forward(self, x, y, mask=None): + return torch.zeros_like(x) + + +def _block(): + b = ( + SanaWMBlock( + hidden_size=AC, + num_heads=AHEADS, + head_dim=ADIM, + mlp_ratio=2.0, + t_kernel_size=3, + qk_norm=True, + cross_norm=True, + conv_kernel_size=4, + k_conv_only=True, + softmax_main=False, + use_chunk_plucker_post_attn=False, + update_rule="torch_recurrent", + cam_update_rule="torch_recurrent", + ) + .double() + .eval() + ) + # Stub cross-attn (unchanged/uncached in forward_long) to avoid the CUDA + # attention backend; both forward and forward_long add the same zero. + b.cross_attn = _ZeroCross() + return b + + +def test_block_forward_long_reduces_to_dense(_global_args): + block = _block() + x = _x() + y = torch.randn(AB, 4, AC, dtype=torch.float64) + t6 = torch.randn(AB, 1, T, 6 * AC, dtype=torch.float64) + (apply_q, apply_kv, apply_o), freqs = _prope() + prope_fns = (apply_q, apply_kv, apply_o) + + dense = block(x, y, t6, AHW, freqs, prope_fns, None, None) + cache = _empty_cache() + out, ret = block.forward_long( + x, y, t6, AHW, freqs, prope_fns, None, None, kv_cache=cache, save_kv_cache=True + ) + torch.testing.assert_close(out, dense, atol=1e-9, rtol=0) + assert ret is cache + assert cache[_SLOT_K] is not None # main GDN state + assert cache[_SLOT_CAM_K] is not None # cam state + assert cache[_SLOT_FFN_TCONV] is not None # FFN temporal tail + + +# --------------------------------------------------------------------------- # +# Stage 3 — model-level forward_long (tiny CPU model; depth 2 => all-GDN blocks, +# no softmax/LocalAttention in the main path; cross-attn stubbed) +# --------------------------------------------------------------------------- # + +MB, MC, MT, MH, MW = 1, 8, 4, 2, 2 + + +def _tiny_model(): + arch = SanaWMArchConfig( + in_channels=MC, + out_channels=MC, + num_layers=2, # blocks 0,1 -> (i+1)%4 != 0 -> no softmax blocks + num_attention_heads=2, + attention_head_dim=16, + linear_head_dim=16, + num_cross_attention_heads=2, + cross_attention_head_dim=16, + cross_attention_dim=32, + caption_channels=32, + model_max_length=8, + softmax_every_n=4, + update_rule="torch_recurrent", + cam_update_rule="torch_recurrent", + chunk_size=None, + ) + m = SanaWMTransformer3DModel(SanaWMConfig(arch_config=arch)).double().eval() + for b in m.blocks: + b.cross_attn = _ZeroCross() + return m + + +def _model_inputs(): + torch.manual_seed(0) + return dict( + hidden_states=torch.randn(MB, MC, MT, MH, MW, dtype=torch.float64), + encoder_hidden_states=torch.randn(MB, 4, 32, dtype=torch.float64), + timestep=torch.randint( + 1, 1000, (MB, 1, MT) + ).double(), # framewise for both paths + camera_conditions=torch.randn(MB, MT, 20, dtype=torch.float64), + chunk_plucker=torch.randn(MB, 48, MT, MH, MW, dtype=torch.float64), + ) + + +def test_model_forward_long_single_chunk_reduces_to_dense(_global_args): + m = _tiny_model() + inp = _model_inputs() + with torch.no_grad(): + dense = m(**inp) + out, cache = m.forward_long( + **inp, kv_cache=None, save_kv_cache=True, start_f=0, end_f=MT + ) + torch.testing.assert_close(out, dense, atol=1e-9, rtol=0) + assert len(cache) == len(m.blocks) + assert cache[0][_SLOT_K] is not None + assert cache[0][_SLOT_CAM_K] is not None + assert cache[0][_SLOT_FFN_TCONV] is not None + + +def test_model_forward_long_two_chunks_runs_and_windows(_global_args): + m = _tiny_model() + inp = _model_inputs() + split = 2 + with torch.no_grad(): + o0, c0 = m.forward_long( + hidden_states=inp["hidden_states"][:, :, :split], + encoder_hidden_states=inp["encoder_hidden_states"], + timestep=inp["timestep"][:, :, :split], + camera_conditions=inp["camera_conditions"], # full; sliced internally + chunk_plucker=inp["chunk_plucker"], + kv_cache=None, + save_kv_cache=True, + start_f=0, + end_f=split, + ) + o1, c1 = m.forward_long( + hidden_states=inp["hidden_states"][:, :, split:], + encoder_hidden_states=inp["encoder_hidden_states"], + timestep=inp["timestep"][:, :, split:], + camera_conditions=inp["camera_conditions"], + chunk_plucker=inp["chunk_plucker"], + kv_cache=c0, + save_kv_cache=True, + start_f=split, + end_f=MT, + ) + assert o0.shape == (MB, MC, split, MH, MW) + assert o1.shape == (MB, MC, MT - split, MH, MW) + assert torch.isfinite(o0).all() and torch.isfinite(o1).all() + assert c1 is not None and len(c1) == len(m.blocks) diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_realtime_path.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_realtime_path.py new file mode 100644 index 000000000..4f42793b5 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_realtime_path.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Realtime (per-chunk) path of SanaWMStreamingDenoisingStage. + +Successor of the retired SanaWMChunkGenerator tests: a tiny CPU-safe model is +ticked through stage.forward with a session attached (one chunk per call, +chunk noise supplied directly — the latent-prep stage is exercised separately). +The stage must carry the per-session 10-slot KV cache + the growing latent +across ticks, evict stale chunk entries, stay deterministic for identical +inputs, and reset on block_idx == 0. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from sglang.multimodal_gen.configs.models.dits.sana_wm import ( + SanaWMArchConfig, + SanaWMConfig, +) +from sglang.multimodal_gen.runtime import server_args as _sa_mod +from sglang.multimodal_gen.runtime.models.dits.sana_wm import SanaWMTransformer3DModel +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import ( + SanaWMStreamCacheState, + SanaWMStreamingDenoisingStage, +) +from sglang.multimodal_gen.runtime.realtime.session import RealtimeSession +from sglang.multimodal_gen.runtime.server_args import set_global_server_args + +MC = 8 + + +class _ZeroCross(torch.nn.Module): + def forward(self, x, y, mask=None): + return torch.zeros_like(x) + + +@pytest.fixture +def _global_args(): + prev = _sa_mod._global_server_args + set_global_server_args( + SimpleNamespace( + comfyui_mode=False, + enable_cfg_parallel=False, + enable_torch_compile=False, + attention_backend=None, + # DenoisingStage.__init__ reads this for its CFG-parallel plumbing. + pipeline_config=SimpleNamespace( + dit_config=SimpleNamespace(hidden_size=32, num_attention_heads=2) + ), + ) + ) + try: + yield + finally: + set_global_server_args(prev) + + +def _tiny_model(): + arch = SanaWMArchConfig( + in_channels=MC, + out_channels=MC, + num_layers=2, # GDN-only -> CPU-safe main path + num_attention_heads=2, + attention_head_dim=16, + linear_head_dim=16, + num_cross_attention_heads=2, + cross_attention_head_dim=16, + cross_attention_dim=32, + caption_channels=32, + model_max_length=8, + softmax_every_n=4, + update_rule="torch_recurrent", + cam_update_rule="torch_recurrent", + chunk_size=None, + ) + m = SanaWMTransformer3DModel(SanaWMConfig(arch_config=arch)).float().eval() + for b in m.blocks: + b.cross_attn = _ZeroCross() + return m + + +def _stage_and_args(nfpb: int = 2): + stage = SanaWMStreamingDenoisingStage(transformer=_tiny_model()) + prompt = torch.zeros(1, 4, 32, dtype=torch.float32) + pcfg = SimpleNamespace( + dit_precision="fp32", + num_frame_per_block=nfpb, + num_cached_blocks=2, + sink_token=True, + denoising_step_list=(1000, 700, 0), + streaming_cfg_scale=1.0, + get_pos_prompt_embeds=lambda batch: [prompt], + get_neg_prompt_embeds=lambda batch: [], + ) + server_args = SimpleNamespace(pipeline_config=pcfg, enable_cfg_parallel=False) + return stage, server_args + + +def _tick(session, block_idx: int, chunk_lat: torch.Tensor, plan: list[int]): + return SimpleNamespace( + session=session, + block_idx=block_idx, + latents=chunk_lat, + extra={"sana_wm_chunk_plan": plan}, + prompt_attention_mask=None, + negative_attention_mask=None, + do_classifier_free_guidance=False, + ) + + +def _noise(*shape, seed: int): + g = torch.Generator().manual_seed(seed) + return torch.randn(*shape, dtype=torch.float32, generator=g) + + +def test_realtime_path_multi_tick_carries_state(_global_args): + stage, server_args = _stage_and_args(nfpb=2) + session = RealtimeSession() + fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32) + + # Tick 0 (chunk 0): conditioning frame + 2 new frames. + chunk0 = torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=1)], dim=2) + out = stage.forward(_tick(session, 0, chunk0, [3]), server_args) + state = session.get_or_create_state(SanaWMStreamCacheState) + assert out.latents.shape[2] == 3 + assert state.chunk_idx == 1 and state.chunk_indices == [0, 3] + assert state.stream_kv_cache[0][0][0] is not None # GDN state stored + # The condition frame is held fixed. + assert torch.allclose(state.latents[:, :, 0].cpu(), fl[:, :, 0]) + + # Tick 1: 2 more frames, KV carried. + out = stage.forward( + _tick(session, 1, _noise(1, MC, 2, 2, 2, seed=2), [2]), server_args + ) + assert out.latents.shape[2] == 5 + assert state.chunk_idx == 2 and state.chunk_indices == [0, 3, 5] + assert torch.isfinite(state.latents).all() + + # Boundary-style tick: a TWO-chunk plan in one call. + out = stage.forward( + _tick(session, 2, _noise(1, MC, 4, 2, 2, seed=3), [2, 2]), server_args + ) + assert out.latents.shape[2] == 9 + assert state.chunk_idx == 4 and state.chunk_indices == [0, 3, 5, 7, 9] + + +def test_realtime_path_evicts_stale_kv(_global_args): + stage, server_args = _stage_and_args(nfpb=2) + session = RealtimeSession() + fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32) + stage.forward( + _tick(session, 0, torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=1)], 2), [3]), + server_args, + ) + for i in range(1, 5): + stage.forward( + _tick(session, i, _noise(1, MC, 2, 2, 2, seed=10 + i), [2]), server_args + ) + state = session.get_or_create_state(SanaWMStreamCacheState) + assert state.chunk_indices == [0, 3, 5, 7, 9, 11] + + def _has_any(entry): + return any(slot is not None for block in entry for slot in block) + + kept = [i for i, e in enumerate(state.stream_kv_cache) if _has_any(e)] + # Sink chunk + the last num_cached_blocks chunks (accumulate's read window). + assert kept == [0, 3, 4] + + +def test_realtime_path_is_deterministic(_global_args): + def _run(): + stage, server_args = _stage_and_args(nfpb=2) + torch.manual_seed(0) # tiny-model init inside _stage_and_args uses global RNG + session = RealtimeSession() + fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32) + stage.forward( + _tick(session, 0, torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=5)], 2), [3]), + server_args, + ) + stage.forward( + _tick(session, 1, _noise(1, MC, 2, 2, 2, seed=6), [2]), server_args + ) + return session.get_or_create_state(SanaWMStreamCacheState).latents.cpu() + + torch.manual_seed(1234) + a = _run() + torch.manual_seed(1234) + b = _run() + assert torch.equal(a, b) + + +def test_realtime_path_resets_on_block_zero(_global_args): + stage, server_args = _stage_and_args(nfpb=2) + session = RealtimeSession() + fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32) + chunk0 = torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=1)], 2) + stage.forward(_tick(session, 0, chunk0, [3]), server_args) + stage.forward(_tick(session, 1, _noise(1, MC, 2, 2, 2, seed=2), [2]), server_args) + state = session.get_or_create_state(SanaWMStreamCacheState) + assert state.chunk_idx == 2 + + # block_idx == 0 restarts the session in place. + stage.forward(_tick(session, 0, chunk0, [3]), server_args) + assert state.chunk_idx == 1 and state.chunk_indices == [0, 3] + assert state.latents.shape[2] == 3 diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_stage.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_stage.py new file mode 100644 index 000000000..c2471df0a --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_stage.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +"""S1c tests — streaming denoising stage: KV accumulator + chunk-loop composition. + +The accumulator is the load-bearing new logic: GDN/STATE blocks copy-forward the +previous chunk's recurrent state; softmax/CONCAT blocks concatenate the rolling ++ sink K/V along **dim=1** (our (B,N,H,D) softmax cache layout — the reference's +dim=2 would concat the head axis and silently corrupt every softmax block). +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from sglang.multimodal_gen.configs.models.dits.sana_wm import ( + SanaWMArchConfig, + SanaWMConfig, +) +from sglang.multimodal_gen.runtime import server_args as _sa_mod +from sglang.multimodal_gen.runtime.models.dits.sana_wm import ( + _CACHE_TYPE_CONCAT, + _CACHE_TYPE_STATE, + _NUM_STREAM_CACHE_SLOTS, + _SLOT_CAM_K, + _SLOT_CAM_V, + _SLOT_FFN_TCONV, + _SLOT_K, + _SLOT_SHORTCONV, + _SLOT_TYPE_FLAG, + _SLOT_V, + SanaWMTransformer3DModel, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import ( + SanaWMStreamingDenoisingStage as Stage, +) +from sglang.multimodal_gen.runtime.server_args import set_global_server_args + +B, Hh, D = 1, 2, 4 + + +# --------------------------------------------------------------------------- # +# 1. Chunk schedule +# --------------------------------------------------------------------------- # + + +def test_autoregressive_segments_first_chunk_absorbs_remainder(): + assert Stage._autoregressive_segments(13, 3) == [0, 4, 7, 10, 13] + seg = Stage._autoregressive_segments(121, 3) + assert seg[0] == 0 and seg[1] == 4 and seg[-1] == 121 + assert len(seg) - 1 == 40 # num chunks + assert all(seg[i + 1] - seg[i] == 3 for i in range(1, len(seg) - 1)) + + +# --------------------------------------------------------------------------- # +# 2. KV accumulator +# --------------------------------------------------------------------------- # + + +def _state_block(seed): + g = torch.Generator().manual_seed(seed) + blk = [None] * _NUM_STREAM_CACHE_SLOTS + blk[_SLOT_K] = torch.randn(B, Hh, D, D, generator=g) # recurrent state_kv + blk[_SLOT_V] = torch.randn(B, Hh, D, 1, generator=g) # state_z + blk[_SLOT_CAM_K] = torch.randn(B, Hh, D, D, generator=g) # cam state + blk[_SLOT_SHORTCONV] = torch.randn(B * 6, 3, Hh * D, generator=g) + blk[_SLOT_TYPE_FLAG] = torch.tensor([_CACHE_TYPE_STATE]) + blk[_SLOT_FFN_TCONV] = torch.randn(B, Hh * D, 1, 6, generator=g) + return blk + + +def _concat_block(n_tok, seed): + g = torch.Generator().manual_seed(seed) + blk = [None] * _NUM_STREAM_CACHE_SLOTS + for slot in (_SLOT_K, _SLOT_V, _SLOT_CAM_K, _SLOT_CAM_V): + blk[slot] = torch.randn(B, n_tok, Hh, D, generator=g) # (B, N, H, D) + blk[_SLOT_TYPE_FLAG] = torch.tensor([_CACHE_TYPE_CONCAT]) + blk[_SLOT_FFN_TCONV] = torch.randn(B, Hh * D, 1, 6, generator=g) + return blk + + +def test_accumulate_state_block_copies_previous_chunk(): + # 3 chunks, 1 STATE block. chunk2 must copy-forward chunk1's state. + kv = [[_state_block(c)] for c in range(3)] + cur, sink_num = Stage._accumulate_kv_cache( + kv, + chunk_idx=2, + chunk_indices=[0, 4, 7, 10], + num_cached_blocks=2, + sink_token=True, + num_blocks=1, + ) + assert cur[0][_SLOT_K] is kv[1][0][_SLOT_K] # carried from chunk 1 + assert cur[0][_SLOT_V] is kv[1][0][_SLOT_V] + assert cur[0][_SLOT_CAM_K] is kv[1][0][_SLOT_CAM_K] + assert float(cur[0][_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_STATE + + +def test_accumulate_concat_block_concats_on_token_axis_dim1(): + # chunk0 K has 4 tokens, chunk1 K has 3 -> chunk2 prefix = 4+3=7 on dim=1. + n0, n1 = 4, 3 + kv = [ + [_concat_block(n0, 0)], + [_concat_block(n1, 1)], + [[None] * _NUM_STREAM_CACHE_SLOTS], + ] + cur, _ = Stage._accumulate_kv_cache( + kv, + chunk_idx=2, + chunk_indices=[0, n0, n0 + n1, n0 + n1 + 3], + num_cached_blocks=5, + sink_token=False, + num_blocks=1, + ) + acc_k = cur[0][_SLOT_K] + assert acc_k.shape == (B, n0 + n1, Hh, D) # dim=1 grew; head axis (dim 2) intact + # exact content: chunk0 then chunk1 along dim=1 + torch.testing.assert_close(acc_k[:, :n0], kv[0][0][_SLOT_K], atol=0, rtol=0) + torch.testing.assert_close(acc_k[:, n0:], kv[1][0][_SLOT_K], atol=0, rtol=0) + assert float(cur[0][_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_CONCAT + + +def test_accumulate_sink_includes_chunk_zero(): + # num_cached_blocks=2 at chunk 3 -> sink_start=2>0 -> valid = [0, 2] + n = 3 + kv = [[_concat_block(n, c)] for c in range(3)] + [ + [[None] * _NUM_STREAM_CACHE_SLOTS] + ] + cur, sink_num = Stage._accumulate_kv_cache( + kv, + chunk_idx=3, + chunk_indices=[0, n, 2 * n, 3 * n, 4 * n], + num_cached_blocks=2, + sink_token=True, + num_blocks=1, + ) + # valid = [0, 2] -> 2 chunks of n tokens each + assert cur[0][_SLOT_K].shape == (B, 2 * n, Hh, D) + assert sink_num == n # chunk_indices[1] - chunk_indices[0] + + +def test_accumulate_chunk_zero_is_empty(): + kv = [[[None] * _NUM_STREAM_CACHE_SLOTS]] + cur, sink_num = Stage._accumulate_kv_cache( + kv, 0, [0, 3], num_cached_blocks=2, sink_token=True, num_blocks=1 + ) + assert sink_num == 0 and cur[0][_SLOT_K] is None + + +# --------------------------------------------------------------------------- # +# 3. Model-level chunk-loop composition (depth >= 4 -> a softmax/CONCAT block) +# --------------------------------------------------------------------------- # + +MC, MT, MHt, MWt = 8, 9, 2, 2 # 9 latent frames, block=3 -> chunks [0,3),[3,6),[6,9) + + +class _ZeroCross(torch.nn.Module): + def forward(self, x, y, mask=None): + return torch.zeros_like(x) + + +@pytest.fixture +def _global_args(): + prev = _sa_mod._global_server_args + set_global_server_args( + SimpleNamespace( + comfyui_mode=False, + enable_cfg_parallel=False, + enable_torch_compile=False, + attention_backend=None, + ) + ) + try: + yield + finally: + set_global_server_args(prev) + + +def _depth4_model(): + arch = SanaWMArchConfig( + in_channels=MC, + out_channels=MC, + num_layers=4, # block 3 -> softmax/CONCAT + num_attention_heads=2, + attention_head_dim=16, + linear_head_dim=16, + num_cross_attention_heads=2, + cross_attention_head_dim=16, + cross_attention_dim=32, + caption_channels=32, + model_max_length=8, + softmax_every_n=4, + update_rule="torch_recurrent", + cam_update_rule="torch_recurrent", + chunk_size=None, + ) + m = SanaWMTransformer3DModel(SanaWMConfig(arch_config=arch)).double().eval() + for b in m.blocks: + b.cross_attn = _ZeroCross() + return m + + +def test_streaming_loop_runs_and_accumulates_concat_block(_global_args): + """Drive forward_long chunk-by-chunk (the stage's core loop) on a depth-4 + model: accumulate -> denoise(save=False) -> clean(save=True). Verify finite + output, threaded GDN state, and that the softmax block (idx 3) accumulates a + growing K window across chunks (the CONCAT path the depth-2 fixture can't).""" + m = _depth4_model() + assert [b.softmax_main for b in m.blocks] == [False, False, False, True] + torch.manual_seed(0) + latents = torch.randn(B, MC, MT, MHt, MWt, dtype=torch.float64) + y = torch.randn(B, 4, 32, dtype=torch.float64) + cam = torch.randn(B, MT, 20, dtype=torch.float64) + plk = torch.randn(B, 48, MT, MHt, MWt, dtype=torch.float64) + + seg = Stage._autoregressive_segments(MT, 3) # [0,3,6,9] + num_chunks = len(seg) - 1 + kv = [ + [[None] * _NUM_STREAM_CACHE_SLOTS for _ in range(4)] for _ in range(num_chunks) + ] + concat_k_lens = [] + + for ci in range(num_chunks): + chunk_kv, sink_num = Stage._accumulate_kv_cache( + kv, ci, seg, num_cached_blocks=2, sink_token=True, num_blocks=4 + ) + s, e = seg[ci], seg[ci + 1] + # softmax block 3 prefix K length entering this chunk + pk = chunk_kv[3][_SLOT_K] + concat_k_lens.append(0 if pk is None else pk.shape[1]) + fidx = torch.arange(s, e) if sink_num > 0 else None + lat = latents[:, :, s:e] + ts = torch.full((B, e - s), 500.0, dtype=torch.float64) + # denoise step (save=False reads the accumulated prefix) + out, _ = m.forward_long( + hidden_states=lat, + encoder_hidden_states=y, + timestep=ts, + camera_conditions=cam, + chunk_plucker=plk, + kv_cache=chunk_kv, + save_kv_cache=False, + start_f=s, + end_f=e, + frame_index=fidx, + ) + assert out.shape == (B, MC, e - s, MHt, MWt) + assert torch.isfinite(out).all() + # clean pass writes this chunk's KV + ts0 = torch.zeros(B, 1, e - s, dtype=torch.float64) + _, updated = m.forward_long( + hidden_states=lat, + encoder_hidden_states=y, + timestep=ts0, + camera_conditions=cam, + chunk_plucker=plk, + kv_cache=chunk_kv, + save_kv_cache=True, + start_f=s, + end_f=e, + frame_index=fidx, + ) + kv[ci] = updated + # block 3 (softmax) stores current-chunk K as (B, N_tok, H, D) where + # N_tok = frames * H * W; blocks 0-2 (GDN) store recurrent state (B,heads,hd,hd). + n_tok = (e - s) * MHt * MWt + assert kv[ci][3][_SLOT_K].shape == (B, n_tok, 2, 16) # heads=2, head_dim=16 + assert kv[ci][0][_SLOT_K].shape == (B, 2, 16, 16) # GDN state_kv + assert float(kv[ci][3][_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_CONCAT + assert float(kv[ci][0][_SLOT_TYPE_FLAG].item()) == _CACHE_TYPE_STATE + + # softmax prefix grows (in TOKENS = frames*H*W): chunk0 sees 0, chunk1 sees + # chunk0's 3*4=12, chunk2 sees the rolling+sink window. + S_tok = MHt * MWt + assert concat_k_lens[0] == 0 + assert concat_k_lens[1] == 3 * S_tok + assert concat_k_lens[2] >= 3 * S_tok diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_vae.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_vae.py new file mode 100644 index 000000000..ef641bae6 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_vae.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""S2 tests — streaming causal-VAE decode: conv_cache + upsampler trim. + +The streaming decode threads a per-conv `conv_cache` dict so a chunked causal +decode (carrying the cache across chunks) is bit-comparable to a monolithic +causal decode of the whole clip. The two logic-bearing pieces are the conv leaf +(`LTX2VideoCausalConv3d`, prepend prev tail / store last k-1 frames) and the +upsampler temporal trim (`LTXVideoUpsampler3d`, drop the anchor frame ONLY at the +true clip start). FP64, CPU. +""" + +from __future__ import annotations + +import pytest +import torch + +from sglang.multimodal_gen.runtime.models.vaes.ltx_2_vae import ( + LTX2VideoCausalConv3d, + LTX2VideoResnetBlock3d, + LTXVideoUpsampler3d, +) + +torch.manual_seed(0) + + +def _chunks(t, splits): + out, i = [], 0 + for s in splits: + out.append(t[:, :, i : i + s]) + i += s + return out + + +@pytest.mark.parametrize("splits", [[6], [3, 3], [1, 2, 3], [2, 4]]) +def test_causal_conv_chunked_equals_monolithic(splits): + conv = ( + LTX2VideoCausalConv3d(in_channels=4, out_channels=5, kernel_size=3) + .double() + .eval() + ) + x = torch.randn(1, 4, 6, 3, 3, dtype=torch.float64) + with torch.no_grad(): + whole = conv(x, causal=True) + cache = {} + parts = [ + conv(c, causal=True, conv_cache=cache, cache_key="c") + for c in _chunks(x, splits) + ] + chunked = torch.cat(parts, dim=2) + assert chunked.shape == whole.shape + torch.testing.assert_close(chunked, whole, atol=1e-9, rtol=0) + + +@pytest.mark.parametrize("splits", [[6], [3, 3], [2, 4], [1, 2, 3]]) +def test_upsampler_chunked_equals_monolithic(splits): + # residual + temporal upsample (stride[0]=2) -> exercises the trim gating. + up = ( + LTXVideoUpsampler3d( + in_channels=8, stride=(2, 2, 2), residual=True, upscale_factor=1 + ) + .double() + .eval() + ) + x = torch.randn(1, 8, 6, 4, 4, dtype=torch.float64) + with torch.no_grad(): + whole = up(x, causal=True) + cache = {} + parts = [ + up(c, causal=True, conv_cache=cache, cache_key="u") + for c in _chunks(x, splits) + ] + chunked = torch.cat(parts, dim=2) + # monolithic trims the anchor frame once; chunked must reproduce that exactly. + assert chunked.shape == whole.shape + torch.testing.assert_close(chunked, whole, atol=1e-9, rtol=0) + + +@pytest.mark.parametrize("splits", [[6], [3, 3], [2, 4]]) +def test_resnet_chunked_equals_monolithic(splits): + blk = LTX2VideoResnetBlock3d(in_channels=4, out_channels=4).double().eval() + x = torch.randn(1, 4, 6, 3, 3, dtype=torch.float64) + with torch.no_grad(): + whole = blk(x, causal=True) + cache = {} + parts = [ + blk(c, causal=True, conv_cache=cache, cache_key="r") + for c in _chunks(x, splits) + ] + chunked = torch.cat(parts, dim=2) + torch.testing.assert_close(chunked, whole, atol=1e-9, rtol=0) + + +def test_conv_cache_none_is_unchanged_dense_behavior(): + # Backward-compat: with no conv_cache, output == the original causal forward. + conv = ( + LTX2VideoCausalConv3d(in_channels=4, out_channels=4, kernel_size=3) + .double() + .eval() + ) + x = torch.randn(1, 4, 5, 2, 2, dtype=torch.float64) + with torch.no_grad(): + a = conv(x, causal=True) + b = conv(x, causal=True, conv_cache=None, cache_key=None) + torch.testing.assert_close(a, b, atol=0, rtol=0) diff --git a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py b/python/sglang/multimodal_gen/test/unit/test_sampling_params.py index 386ab3672..8ca545fdb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py +++ b/python/sglang/multimodal_gen/test/unit/test_sampling_params.py @@ -279,6 +279,32 @@ class TestSamplingParamsCliArgs(unittest.TestCase): self.assertIn("width", explicit_fields) self.assertIn("height", explicit_fields) + def test_cli_path_preserves_diffusers_kwargs_in_request_extra(self): + server_args = MagicMock() + server_args.backend = "sglang" + server_args.model_id = None + server_args.pipeline_config = MagicMock() + diffusers_kwargs = {"camera_to_world_path": "/tmp/camera.npy"} + + with patch.object( + SamplingParams, + "from_pretrained", + side_effect=lambda *args, **kwargs: Flux2SamplingParams(), + ): + params = SamplingParams.from_user_sampling_params_args( + "dummy-model", + server_args=server_args, + prompt="p", + image_path="/tmp/in.png", + diffusers_kwargs=diffusers_kwargs, + ) + + self.assertEqual(params.diffusers_kwargs, diffusers_kwargs) + self.assertEqual( + params.build_request_extra()["diffusers_kwargs"], + diffusers_kwargs, + ) + def test_dataclasses_replace_preserves_explicit_fields(self): """`dataclasses.replace` drops `_explicit_fields`; DiffGenerator must restore it.""" import dataclasses diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 396ffb794..e89f990a2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -20,6 +20,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConf from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImagePipelineConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import ( + SanaWMPipelineConfig, + SanaWMRealtimeConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.wan import ( FastWan2_2_TI2V_5B_Config, TurboWanT2V480PConfig, @@ -665,6 +669,7 @@ class TestOffloadDefaults(unittest.TestCase): mova_deployment = MOVAPipelineConfig().get_model_deployment_config() zimage_deployment = ZImagePipelineConfig().get_model_deployment_config() ltx_deployment = LTX2PipelineConfig().get_model_deployment_config() + sana_wm_deployment = SanaWMPipelineConfig().get_model_deployment_config() self.assertIsNone(qwen_deployment.fsdp_auto_min_available_memory_gb) self.assertFalse(qwen_deployment.auto_dit_layerwise_offload) @@ -686,6 +691,35 @@ class TestOffloadDefaults(unittest.TestCase): ltx_deployment.auto_disable_component_offload_components, ("dit",) ) + self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60) + self.assertTrue(sana_wm_deployment.auto_dit_layerwise_offload) + + def test_auto_multi_gpu_sana_wm_prefers_fsdp_and_cfg_parallel(self): + args = self._from_dict_with_pipeline_config( + SanaWMPipelineConfig(), + kwargs={ + "model_path": "Efficient-Large-Model/SANA-WM_bidirectional", + "num_gpus": 2, + "performance_mode": "auto", + }, + ) + + self.assertTrue(args.use_fsdp_inference) + self.assertTrue(args.enable_cfg_parallel) + + def test_auto_multi_gpu_sana_wm_realtime_disables_cfg_parallel(self): + args = self._from_dict_with_pipeline_config( + SanaWMRealtimeConfig(), + kwargs={ + "model_path": "Efficient-Large-Model/SANA-WM_streaming", + "num_gpus": 2, + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.use_fsdp_inference) + self.assertFalse(args.enable_cfg_parallel) + def test_manual_mode_preserves_unset_performance_args(self): args = self._from_dict_with_pipeline_config( QwenImagePipelineConfig(), @@ -1346,6 +1380,10 @@ class TestModelIdResolution(unittest.TestCase): info = _get_config_info(path) self.assertIsNotNone(info) + def test_sana_wm_model_path_resolves_registry(self): + info = _get_config_info("Efficient-Large-Model/SANA-WM_bidirectional") + self.assertIs(info.pipeline_config_cls, SanaWMPipelineConfig) + def test_model_id_unknown_falls_back_without_crash(self): # unrecognized model_id: should warn and fall back to path-based detection # with an unresolvable path, expect RuntimeError from the detector step