diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 31f3ac35d..58bca4809 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -39,6 +39,7 @@ th {
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `SanaWmPipeline` | SANA-WM | `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-WM.md) |
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md
index 5844d1114..34b642b9e 100644
--- a/docs/user_guide/diffusion_features.md
+++ b/docs/user_guide/diffusion_features.md
@@ -168,6 +168,14 @@ The following tables show which models support each feature:
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
+| **SANA-WM** | ❌ | ❌ | ❌5 | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+
+> Notes:
+> 5. SANA-WM cannot support sequence parallelism: its bidirectional gated delta
+> recurrence carries state across frames, so a rank cannot denoise a slice of
+> the token sequence in isolation. Doing so would need a distributed scan or
+> an all-gather before every GDN block. The remaining ❌ columns are simply
+> unvalidated on this model, not known-broken.
> **Step execution note:** Helios supports single-request step execution only;
> use `max_num_seqs=1`.
@@ -175,7 +183,7 @@ The following tables show which models support each feature:
**Frame Interpolation Support**
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
-- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
+- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni, SANA-WM
### AudioGen
diff --git a/recipes/NVIDIA/SANA-WM.md b/recipes/NVIDIA/SANA-WM.md
new file mode 100644
index 000000000..2926cc46f
--- /dev/null
+++ b/recipes/NVIDIA/SANA-WM.md
@@ -0,0 +1,193 @@
+# SANA-WM
+
+> Camera-controllable first-frame image-to-video world model.
+
+## Summary
+
+- Vendor: Efficient-Large-Model / NVlabs SANA
+- Model: `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` (standard diffusers layout, converted offline from the NVlabs release; Stage-1 transformer + VAE only)
+- Task: First-frame image-to-video generation with camera control
+- Mode: Online serving with the OpenAI-compatible video API
+- Model weights: about 13 GB for the Stage-1 transformer (10 GB) and VAE (2.3 GB)
+- Local disk: reserve about 40 GB for the Hugging Face cache and runtime artifacts
+- Recommended GPU: 24 GB or larger CUDA GPU
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want to serve SANA-WM through `/v1/videos` or
+`/v1/videos/sync`. The model takes a text prompt, a first-frame image, and
+either an action DSL string or explicit camera poses. vLLM-Omni serves the
+SANA-WM Stage-1 DiT, decoded through the SANA VAE. The optional LTX-2 refiner
+stage is not supported by this integration; it is a planned follow-up.
+
+## References
+
+- Upstream model card:
+- Video API: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA RTX PRO 6000 Blackwell 96GB
+
+#### Capacity
+
+- Model storage: the Stage-1 transformer is about 10 GB and the VAE about
+ 2.3 GB. The Gemma text encoder is a separate 4.9 GB repo. The engine prefetches
+ the whole model repo at startup (`allow_patterns=["*"]`), which is why the
+ Stage-1 weights live in their own repo — the two-stage one carries an
+ additional 84 GB `refiner/` that this path never loads.
+- Text encoder: the pipeline tries `google/gemma-2-2b-it` first, then falls back
+ to the ungated mirror `Efficient-Large-Model/gemma-2-2b-it`. The first repo is
+ gated, so without an accepted licence and `hf auth login` you will see one
+ failed load before the fallback succeeds — that is expected, not an error to
+ chase. Set `VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER` to pin a specific repo or a
+ local path and skip both.
+- Disk sizing: provision about 40 GB of local disk or Hugging Face cache volume
+ so the model, temporary downloads, and generated artifacts fit without cache
+ eviction.
+- GPU sizing: the default 1280x704, 161-frame, 60-step serving profile peaks at
+ 22.6 GB of device memory and takes about 133 s to generate on one RTX PRO
+ 6000 Blackwell. The peak lands in the VAE decode, which is why the pipeline
+ forces VAE tiling on regardless of `vae_use_tiling` — without it the same
+ request costs about 9 GB more, and a 321-frame one OOMs outright. On smaller
+ GPUs, lower `width`, `height`, or `num_frames` before serving production
+ requests.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA driver with CUDA runtime supported by your PyTorch
+ build
+- Recommended operator library: Triton, installed through the vLLM/vLLM-Omni
+ Python environment
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+The repo ships the standard Diffusers layout (`model_index.json` +
+`transformer/`, `vae/`), and its `model_index.json` names `SanaWmPipeline`, so
+the pipeline class resolves on its own.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+vllm serve BBBBruce/SANA-WM_bidirectional-stage1-diffusers \
+ --omni \
+ --host 0.0.0.0 \
+ --port 8091
+```
+
+No deploy config: single-stage diffusion models are deliberately absent from
+`OMNI_PIPELINES` (`vllm_omni/config/pipeline_registry.py`), so stage resolution
+falls back to the default stage config and a YAML's stage settings — including
+`default_sampling_params` — would not be applied. The production generation
+settings therefore live in the model (`num_inference_steps=60`,
+`guidance_scale=5.0`); a request that omits a field gets them. The examples
+below still pass every field explicitly so the numbers are visible.
+
+If you point this at the older two-stage repo
+(`BBBBruce/SANA-WM_bidirectional-diffusers`), startup fails with `Model class
+SanaWmTwoStagesPipeline not found in diffusion model registry`, because that
+repo's `model_index.json` names a class this build does not register. Add
+`--model-class-name SanaWmPipeline` to override it.
+
+#### Verification
+
+Use a short smoke request first:
+
+```bash
+curl -sS -X POST http://localhost:8091/v1/videos/sync \
+ -H "Accept: video/mp4" \
+ -F "prompt=A slow forward camera move through a quiet city street." \
+ -F "negative_prompt=blurry, low quality, distorted, watermark" \
+ -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+ -F "width=1280" \
+ -F "height=704" \
+ -F "num_frames=9" \
+ -F "fps=16" \
+ -F "num_inference_steps=2" \
+ -F "guidance_scale=5.0" \
+ -F "seed=42" \
+ --form-string 'extra_params={"sana_wm":{"action":"w-8","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
+ -o sana_wm_smoke.mp4
+```
+
+For a production-length request, note that the action durations must sum to
+`num_frames - 1` — the rollout includes the identity start pose — and a
+mismatch is rejected rather than padded or truncated:
+
+```bash
+curl -sS -X POST http://localhost:8091/v1/videos/sync \
+ -H "Accept: video/mp4" \
+ -F "prompt=A slow forward camera move through a quiet city street." \
+ -F "negative_prompt=blurry, low quality, distorted, watermark" \
+ -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+ -F "width=1280" \
+ -F "height=704" \
+ -F "num_frames=161" \
+ -F "fps=16" \
+ -F "num_inference_steps=60" \
+ -F "guidance_scale=5.0" \
+ -F "seed=42" \
+ --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
+ -o sana_wm_output.mp4
+```
+
+Use `POST /v1/videos` instead when you want job storage and polling rather than
+inline MP4 bytes. It accepts the same form fields as `/v1/videos/sync`.
+
+```bash
+create_response=$(curl -sS -X POST http://localhost:8091/v1/videos \
+ -H "Accept: application/json" \
+ -F "prompt=A slow forward camera move through a quiet city street." \
+ -F "negative_prompt=blurry, low quality, distorted, watermark" \
+ -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+ -F "width=1280" \
+ -F "height=704" \
+ -F "num_frames=161" \
+ -F "fps=16" \
+ -F "num_inference_steps=60" \
+ -F "guidance_scale=5.0" \
+ -F "seed=42" \
+ --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}')
+
+video_id=$(echo "$create_response" | jq -r '.id')
+curl -sS "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o sana_wm_output.mp4
+```
+
+#### Notes
+
+- Sequence parallelism is not supported. The bidirectional gated delta
+ recurrence carries state across frames, so a rank cannot denoise a slice
+ of the token sequence in isolation; supporting it needs a distributed scan
+ or an all-gather before the GDN blocks.
+
+- `input_reference` is required for the first frame. Use `image_reference` only
+ when you need a JSON-safe image URL or data URL instead of a multipart upload.
+- `sana_wm` must provide exactly one of `action` or `camera`.
+- Action strings use comma-separated `-` segments. Supported
+ keys are `w`, `a`, `s`, `d` for translation and `i`, `j`, `k`, `l` for
+ pitch/yaw rotation. The durations must sum to `num_frames - 1`.
+- Explicit camera control (alternative to `action`): pass
+ `"camera": {"poses": [...]}` where `poses` is a list of `num_frames`
+ camera-to-world 4x4 matrices (row-major, OpenCV `+X right, +Y down, +Z forward`
+ convention), e.g.
+ `extra_params={"sana_wm":{"camera":{"poses":[[[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]], ...]},"intrinsics":{...}}}`.
+ Most callers should prefer `action`; explicit poses exist for callers that
+ already have a per-frame trajectory.
+- Explicit `intrinsics` are recommended and take the mapping form
+ `{"fx":640,"fy":640,"cx":640,"cy":352}` (for 1280x704). This `{fx,fy,cx,cy}`
+ mapping is the only accepted intrinsics form; omit `intrinsics` to derive them
+ from the output resolution. All four values must be finite, and `fx`/`fy`
+ must be positive — the ray map divides by them.
+- The video API returns decoded MP4 bytes and has no `output_type` field, so
+ raw Stage-1 latents are reachable only from the offline API
+ (`OmniDiffusionSamplingParams(output_type="latent")`); see
+ [`tests/e2e/offline_inference/test_sana_wm.py`](../../tests/e2e/offline_inference/test_sana_wm.py).
+ Putting `output_type` in `extra_params` does not work: it lands in
+ `sampling_params.extra_args`, while the pipeline reads the top-level field.
diff --git a/recipes/README.md b/recipes/README.md
index 32c68e8c1..8b93124d0 100644
--- a/recipes/README.md
+++ b/recipes/README.md
@@ -70,6 +70,7 @@ recipes/
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
+| [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
diff --git a/tests/diffusion/models/sana_wm/__init__.py b/tests/diffusion/models/sana_wm/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py b/tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py
new file mode 100644
index 000000000..cea8aeb76
--- /dev/null
+++ b/tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's camera-branch reduce-scatter.
+
+``out_proj_cam`` is row-parallel, so every rank produces a partial sum over the
+full hidden width, while the GDN stream it is added to is TP-local. Instead of
+all-reducing and then discarding all but one slice, the branch reduce-scatters
+straight to the local shard. These tests pin the shard each rank ends up with,
+including the bias slice, on both code paths:
+
+* the process-group path, which drives ``torch.distributed.reduce_scatter_tensor``
+ directly because vLLM-Omni's diffusion ``GroupCoordinator`` does not implement
+ ``reduce_scatter``;
+* the all-reduce-and-slice fallback taken when the coordinator exposes no
+ ``device_group``.
+
+CPU-only: the collectives and the rank/world-size lookups are mocked, so a
+single process can simulate an arbitrary TP degree.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmSelfAttention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+HIDDEN = 24
+
+
+def _host(bias: torch.Tensor | None) -> SimpleNamespace:
+ """A stand-in ``self`` exposing only what the method reads."""
+ return SimpleNamespace(out_proj_cam=SimpleNamespace(bias=bias))
+
+
+def _fake_reduce_scatter_tensor(rank: int, tp_size: int):
+ """Model ``reduce_scatter_tensor``: split dim 0 and keep chunk ``rank``.
+
+ The test drives a single process, so the other ranks' partials are already
+ folded into the input; scattering is all this has to reproduce.
+ """
+
+ def _impl(output: torch.Tensor, input_: torch.Tensor, group=None) -> None:
+ output.copy_(input_.chunk(tp_size, dim=0)[rank])
+
+ return _impl
+
+
+def _tp_group(with_device_group: bool):
+ return SimpleNamespace(device_group=object() if with_device_group else None)
+
+
+@pytest.mark.parametrize("tp_size", [2, 4])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_process_group_path_matches_all_reduce_then_slice(tp_size: int, rank: int) -> None:
+ """The scattered shard equals the slice the all-reduce path would take."""
+ local_width = HIDDEN // tp_size
+ torch.manual_seed(0)
+ contrib = torch.randn(2, 3, HIDDEN)
+ bias = torch.randn(HIDDEN)
+ reference = torch.zeros(2, 3, local_width)
+
+ with (
+ patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+ patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+ patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(True)),
+ patch("torch.distributed.reduce_scatter_tensor", side_effect=_fake_reduce_scatter_tensor(rank, tp_size)),
+ ):
+ out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+ start = rank * local_width
+ expected = (contrib + bias)[..., start : start + local_width]
+ assert out.shape == (2, 3, local_width)
+ torch.testing.assert_close(out, expected)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_fallback_path_matches_the_process_group_path(rank: int) -> None:
+ """Without a device_group the method all-reduces and slices, same result."""
+ tp_size = 2
+ local_width = HIDDEN // tp_size
+ torch.manual_seed(1)
+ contrib = torch.randn(2, 3, HIDDEN)
+ bias = torch.randn(HIDDEN)
+ reference = torch.zeros(2, 3, local_width)
+
+ with (
+ patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+ patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+ patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(False)),
+ patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t),
+ ):
+ out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+ start = rank * local_width
+ torch.testing.assert_close(out, (contrib + bias)[..., start : start + local_width])
+
+
+def test_tp1_adds_bias_without_a_collective() -> None:
+ contrib = torch.randn(2, 3, HIDDEN)
+ bias = torch.randn(HIDDEN)
+
+ with (
+ patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1),
+ patch("torch.distributed.reduce_scatter_tensor") as collective,
+ ):
+ out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, contrib)
+
+ assert collective.call_count == 0
+ torch.testing.assert_close(out, contrib + bias)
+
+
+def test_rejects_a_width_that_does_not_shard() -> None:
+ """A local width that is not 1/tp of the contribution is a wiring bug."""
+ contrib = torch.randn(1, 2, HIDDEN)
+ reference = torch.zeros(1, 2, HIDDEN // 3) # tp=2 would need HIDDEN // 2
+
+ with (
+ patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+ patch("torch.distributed.reduce_scatter_tensor"),
+ pytest.raises(ValueError, match="camera contribution width mismatch"),
+ ):
+ SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(None), contrib, reference)
diff --git a/tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py b/tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py
new file mode 100644
index 000000000..b01ef02ce
--- /dev/null
+++ b/tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py
@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's TensorParallelRMSNorm q/k fusion.
+
+`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
+and issues a single tensor-parallel all-reduce instead of two. These tests lock
+in that it is numerically identical to applying `norm_q`/`norm_k` separately,
+and that it really issues exactly one (fused) all-reduce.
+
+CPU-only: the collective is mocked, so a single process simulates one rank
+(all-reduce is identity). This exercises the cat/split logic a real multi-rank
+run depends on.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
+ TensorParallelRMSNorm,
+ fused_qk_rms_norm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+
+def _make_norms(dim: int, tp_size: int) -> tuple[TensorParallelRMSNorm, TensorParallelRMSNorm]:
+ """Two norms with distinct random weights so a swapped q/k slice is caught."""
+ torch.manual_seed(0)
+ norm_q = TensorParallelRMSNorm(dim, tp_size=tp_size)
+ norm_k = TensorParallelRMSNorm(dim, tp_size=tp_size)
+ norm_q.weight.data.normal_()
+ norm_k.weight.data.normal_()
+ return norm_q, norm_k
+
+
+def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+ """Single-rank all-reduce: the reduction over one rank is the identity."""
+ return tensor
+
+
+def test_fused_matches_separate() -> None:
+ dim = 64
+ norm_q, norm_k = _make_norms(dim, tp_size=2)
+ q = torch.randn(2, 5, dim)
+ k = torch.randn(2, 5, dim)
+
+ with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce):
+ ref_q, ref_k = norm_q(q), norm_k(k)
+ fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+ torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+ torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_issues_one_all_reduce() -> None:
+ dim = 32
+ norm_q, norm_k = _make_norms(dim, tp_size=4)
+ q = torch.randn(1, 3, dim)
+ k = torch.randn(1, 3, dim)
+
+ with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+ fused_qk_rms_norm(norm_q, norm_k, q, k)
+ assert collective.call_count == 1
+ # Both sums travel in one tensor: last dim is 2, not 1.
+ assert collective.call_args.args[0].shape[-1] == 2
+
+ collective.reset_mock()
+ norm_q(q), norm_k(k)
+ assert collective.call_count == 2
+
+
+@pytest.mark.parametrize(
+ "norms",
+ [
+ (RMSNorm(16), RMSNorm(16)),
+ (nn.Identity(), nn.Identity()),
+ ],
+ ids=["plain_rms_norm", "identity"],
+)
+def test_falls_back_for_non_tp_norms(norms: tuple[nn.Module, nn.Module]) -> None:
+ """qk_norm=False and TP=1 both hand in norms the fused path cannot use."""
+ norm_q, norm_k = norms
+ q = torch.randn(2, 4, 16)
+ k = torch.randn(2, 4, 16)
+
+ with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+ fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+ assert collective.call_count == 0
+
+ torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+ torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)
+
+
+def test_tp1_skips_the_collective() -> None:
+ norm_q, norm_k = _make_norms(16, tp_size=1)
+ q = torch.randn(2, 4, 16)
+ k = torch.randn(2, 4, 16)
+
+ with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+ fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+ assert collective.call_count == 0
+
+ torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+ torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)
diff --git a/tests/diffusion/models/sana_wm/test_request.py b/tests/diffusion/models/sana_wm/test_request.py
new file mode 100644
index 000000000..7e40eedaf
--- /dev/null
+++ b/tests/diffusion/models/sana_wm/test_request.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for SANA-WM request normalization."""
+
+import pytest
+
+from vllm_omni.diffusion.models.sana_wm.request import (
+ SANA_WM_DEFAULT_HEIGHT,
+ SANA_WM_DEFAULT_NUM_FRAMES,
+ SANA_WM_DEFAULT_WIDTH,
+ normalize_sana_wm_payload,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _prompt(**sana_wm):
+ """A valid request, with the action rollout kept in step with num_frames."""
+ block = {"num_frames": SANA_WM_DEFAULT_NUM_FRAMES, **sana_wm}
+ block.setdefault("action", f"w-{block['num_frames'] - 1}")
+ return {
+ "prompt": "a quiet city street",
+ "multi_modal_data": {"image": object()},
+ "sana_wm": block,
+ }
+
+
+def _identity_poses(count):
+ return [[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]] for _ in range(count)]
+
+
+def _camera_prompt(**camera):
+ prompt = _prompt(num_frames=9)
+ prompt["sana_wm"].pop("action")
+ prompt["sana_wm"]["camera"] = {"poses": _identity_poses(9), **camera}
+ return prompt
+
+
+def test_defaults_are_latent_aligned():
+ payload = normalize_sana_wm_payload(_prompt())["additional_information"]["sana_wm"]
+ assert payload["num_frames"] == SANA_WM_DEFAULT_NUM_FRAMES
+ assert payload["height"] == SANA_WM_DEFAULT_HEIGHT
+ assert payload["width"] == SANA_WM_DEFAULT_WIDTH
+
+
+@pytest.mark.parametrize("num_frames", [9, 33, 161])
+def test_accepts_aligned_geometry(num_frames):
+ payload = normalize_sana_wm_payload(_prompt(num_frames=num_frames))["additional_information"]["sana_wm"]
+ assert payload["num_frames"] == num_frames
+
+
+@pytest.mark.parametrize("num_frames", [8, 10, 160])
+def test_rejects_unaligned_num_frames(num_frames):
+ # The VAE compresses 8x temporally; an unaligned count would be floored.
+ with pytest.raises(ValueError, match="num_frames"):
+ normalize_sana_wm_payload(_prompt(num_frames=num_frames))
+
+
+@pytest.mark.parametrize(("height", "width"), [(700, 1280), (704, 1270), (703, 1279)])
+def test_rejects_unaligned_resolution(height, width):
+ with pytest.raises(ValueError, match="divisible by 32"):
+ normalize_sana_wm_payload(_prompt(height=height, width=width))
+
+
+def test_normalization_is_idempotent():
+ once = normalize_sana_wm_payload(_prompt(num_frames=33))
+ twice = normalize_sana_wm_payload(once)
+ assert once["additional_information"]["sana_wm"] == twice["additional_information"]["sana_wm"]
+
+
+def test_rejects_action_shorter_than_num_frames():
+ # Used to shrink the camera tensor silently, then fail deep in the
+ # transformer on a camera/latent frame mismatch.
+ with pytest.raises(ValueError, match="rolls out 33 frames but num_frames is 161"):
+ normalize_sana_wm_payload(_prompt(action="w-32"))
+
+
+def test_rejects_action_longer_than_num_frames():
+ with pytest.raises(ValueError, match="rolls out 161 frames but num_frames is 9"):
+ normalize_sana_wm_payload(_prompt(num_frames=9, action="w-160"))
+
+
+def test_rejects_oversized_action_without_expanding_it():
+ # A 12-byte request body: expanding this rollout before checking the frame
+ # count allocated ~16 GB. Rejection has to be arithmetic, so this test also
+ # fails (by hanging/OOM) if the duration is ever expanded again.
+ with pytest.raises(ValueError, match="rolls out 1000000001 frames"):
+ normalize_sana_wm_payload(_prompt(action="w-1000000000"))
+
+
+@pytest.mark.parametrize("action", ["", " ", "w-0", "w-abc", "w", "x-4"])
+def test_rejects_malformed_action(action):
+ with pytest.raises(ValueError):
+ normalize_sana_wm_payload(_prompt(action=action))
+
+
+@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf")])
+def test_rejects_unusable_focal_length(value):
+ # compute_raymap divides by fx/fy; a zero or non-finite focal length turns
+ # the whole Plucker map into NaN and the response into a NaN video.
+ intrinsics = {"fx": value, "fy": 640, "cx": 640, "cy": 352}
+ with pytest.raises(ValueError, match="fx"):
+ normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_principal_point(value):
+ intrinsics = {"fx": 640, "fy": 640, "cx": value, "cy": 352}
+ with pytest.raises(ValueError, match="must be finite"):
+ normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_speeds(value):
+ # `parsed <= 0` is False for NaN, so the positivity check alone let it pass.
+ with pytest.raises(ValueError, match="must be finite"):
+ normalize_sana_wm_payload(_prompt(translation_speed=value))
+
+
+def test_rejects_non_finite_pose_entry():
+ prompt = _camera_prompt()
+ prompt["sana_wm"]["camera"]["poses"][3][2][1] = float("nan")
+ with pytest.raises(ValueError, match="must be finite"):
+ normalize_sana_wm_payload(prompt)
+
+
+def test_accepts_default_camera_metadata():
+ payload = normalize_sana_wm_payload(_camera_prompt())["additional_information"]["sana_wm"]
+ assert payload["camera"]["poses"]
+ # The unread fields are not carried into the canonical payload.
+ assert set(payload["camera"]) == {"poses"}
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [("format", "w2c_4x4"), ("coordinate_system", "opengl")],
+)
+def test_rejects_unsupported_camera_convention(field, value):
+ # Nothing downstream reads these, so accepting one would mean interpreting
+ # the poses under the wrong convention.
+ with pytest.raises(ValueError, match=field.split("_")[0]):
+ normalize_sana_wm_payload(_camera_prompt(**{field: value}))
diff --git a/tests/e2e/offline_inference/test_sana_wm.py b/tests/e2e/offline_inference/test_sana_wm.py
new file mode 100644
index 000000000..8b1c460c2
--- /dev/null
+++ b/tests/e2e/offline_inference/test_sana_wm.py
@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import numpy as np
+import pytest
+
+pytestmark = [
+ pytest.mark.advanced_model,
+ pytest.mark.diffusion,
+ pytest.mark.gpu,
+]
+
+
+def _make_sana_wm_e2e_image() -> Any:
+ from PIL import Image
+
+ from vllm_omni.diffusion.models.sana_wm import SANA_WM_OUTPUT_HEIGHT, SANA_WM_OUTPUT_WIDTH
+
+ return Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
+
+
+def _sana_wm_e2e_tensor_parallel_size() -> int:
+ return int(
+ os.environ.get(
+ "SANA_WM_E2E_TENSOR_PARALLEL_SIZE",
+ os.environ.get("SANA_WM_E2E_TP", "1"),
+ )
+ )
+
+
+def _coerce_video_array(video: Any) -> np.ndarray:
+ import torch
+
+ if isinstance(video, list):
+ assert video, "SANA-WM e2e produced an empty video list."
+ video = video[0]
+ if isinstance(video, torch.Tensor):
+ video = video.detach().cpu().float().numpy()
+ video_array = np.asarray(video)
+ if video_array.ndim == 5:
+ assert video_array.shape[0] == 1
+ video_array = video_array[0]
+ assert video_array.ndim == 4
+ return video_array
+
+
+def _assert_sana_wm_e2e_shape(
+ video: np.ndarray,
+ *,
+ output_type: str,
+ num_frames: int,
+) -> None:
+ if output_type == "latent":
+ # Latent output is channel-first: (C, T, H, W) after removing the
+ # batch dimension. Stage-1 emits LTX-2 latents with 128 channels.
+ assert video.shape[0] == 128
+ assert 0 < video.shape[1] <= num_frames
+ return
+ assert 0 < video.shape[0] <= num_frames
+
+
+def _run_sana_wm_e2e(*, output_type: str) -> np.ndarray:
+ import torch
+
+ from vllm_omni.diffusion.models.sana_wm import (
+ SANA_WM_MODEL_ID,
+ SANA_WM_OUTPUT_HEIGHT,
+ SANA_WM_OUTPUT_WIDTH,
+ )
+ from vllm_omni.entrypoints.omni import Omni
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+ from vllm_omni.outputs import OmniRequestOutput
+
+ if not torch.cuda.is_available():
+ pytest.skip("Sana-WM e2e requires CUDA.")
+
+ model = os.environ.get("SANA_WM_E2E_MODEL", SANA_WM_MODEL_ID)
+ model_class_name = os.environ.get("SANA_WM_E2E_MODEL_CLASS", "SanaWmPipeline")
+ num_frames = int(os.environ.get("SANA_WM_E2E_NUM_FRAMES", "9"))
+ stage1_steps = int(
+ os.environ.get(
+ "SANA_WM_E2E_STAGE1_STEPS",
+ os.environ.get("SANA_WM_E2E_NUM_INFERENCE_STEPS", "1"),
+ )
+ )
+ action = os.environ.get("SANA_WM_E2E_ACTION", f"w-{max(num_frames - 1, 1)}")
+ tensor_parallel_size = _sana_wm_e2e_tensor_parallel_size()
+ image = _make_sana_wm_e2e_image()
+ omni = Omni(
+ model=model,
+ model_class_name=model_class_name,
+ enforce_eager=True,
+ tensor_parallel_size=tensor_parallel_size,
+ )
+ extra_args: dict[str, Any] = {}
+ # The native Stage-1 path is latent-token capped; raise it for full-size
+ # e2e runs (704x1280 exceeds the default 4096-token cap).
+ native_max_tokens = os.environ.get("SANA_WM_E2E_NATIVE_MAX_TOKENS", "").strip()
+ if native_max_tokens:
+ extra_args["sana_wm_native_max_tokens"] = int(native_max_tokens)
+ # Return the requested output type so the shape assertion matches the
+ # pipeline output instead of the raw Stage-1 latent default.
+
+ output = omni.generate(
+ {
+ "prompt": "A slow forward camera move through a quiet city street.",
+ "multi_modal_data": {"image": image},
+ "sana_wm": {
+ "action": action,
+ "num_frames": num_frames,
+ "translation_speed": 0.055,
+ "rotation_speed_deg": 1.2,
+ # Avoid the optional Pi3X dependency by passing deterministic
+ # camera intrinsics directly.
+ "intrinsics": {
+ "fx": SANA_WM_OUTPUT_WIDTH / 2,
+ "fy": SANA_WM_OUTPUT_WIDTH / 2,
+ "cx": SANA_WM_OUTPUT_WIDTH / 2,
+ "cy": SANA_WM_OUTPUT_HEIGHT / 2,
+ },
+ },
+ },
+ OmniDiffusionSamplingParams(
+ height=SANA_WM_OUTPUT_HEIGHT,
+ width=SANA_WM_OUTPUT_WIDTH,
+ num_frames=num_frames,
+ seed=0,
+ fps=16,
+ num_inference_steps=stage1_steps,
+ guidance_scale=1.0,
+ guidance_scale_provided=True,
+ output_type=output_type,
+ extra_args=extra_args,
+ ),
+ )
+
+ request_output = output[0] if isinstance(output, list) else output
+ assert isinstance(request_output, OmniRequestOutput)
+ assert request_output.error is None
+ assert request_output.images
+
+ frames = request_output.images[0]
+ video = _coerce_video_array(frames)
+ _assert_sana_wm_e2e_shape(video, output_type=output_type, num_frames=num_frames)
+ return video
+
+
+def test_sana_wm_native_generates_video() -> None:
+ output_type = os.environ.get("SANA_WM_E2E_OUTPUT_TYPE", "np")
+ video = _run_sana_wm_e2e(output_type=output_type)
+ if output_type == "latent":
+ assert video.ndim == 4
+ else:
+ assert video.shape[-1] in (3, 4)
diff --git a/tests/e2e/online_serving/test_sana_wm.py b/tests/e2e/online_serving/test_sana_wm.py
new file mode 100644
index 000000000..883e4c367
--- /dev/null
+++ b/tests/e2e/online_serving/test_sana_wm.py
@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Online serving smoke for SANA-WM image-to-video (first-frame I2V via ``/v1/videos``).
+
+Boots the server straight from the model id — the Stage-1 repo's
+``model_index.json`` names ``SanaWmPipeline``, so the class resolves on its own
+— submits one async ``/v1/videos`` job with a first-frame reference image, and
+asserts the job completes and returns video bytes.
+
+From ``tests/``::
+
+ pytest -s -v e2e/online_serving/test_sana_wm.py -m "advanced_model and diffusion" --run-level=advanced_model
+"""
+
+import base64
+import json
+import os
+from io import BytesIO
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from vllm_omni.diffusion.models.sana_wm import (
+ SANA_WM_MODEL_ID,
+ SANA_WM_OUTPUT_HEIGHT,
+ SANA_WM_OUTPUT_WIDTH,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = SANA_WM_MODEL_ID
+PROMPT = "A slow forward camera move through a quiet city street."
+NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark"
+
+# Smoke knobs: a handful of frames / steps so the job finishes quickly.
+SMOKE_NUM_FRAMES = 9
+SMOKE_NUM_INFERENCE_STEPS = 2
+# First-frame I2V needs a camera trajectory; "w-" is the forward-move action DSL.
+# Passed via the ``sana_wm`` form field (a JSON string parsed server-side); explicit
+# intrinsics avoid the optional Pi3X camera-calibration dependency.
+SANA_WM_PARAMS = {
+ "action": f"w-{SMOKE_NUM_FRAMES - 1}",
+ "translation_speed": 0.055,
+ "rotation_speed_deg": 1.2,
+ "intrinsics": {
+ "fx": SANA_WM_OUTPUT_WIDTH / 2,
+ "fy": SANA_WM_OUTPUT_WIDTH / 2,
+ "cx": SANA_WM_OUTPUT_WIDTH / 2,
+ "cy": SANA_WM_OUTPUT_HEIGHT / 2,
+ },
+}
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _first_frame_data_url() -> str:
+ """A solid-color first frame encoded as a PNG ``data:`` URL for I2V."""
+ from PIL import Image
+
+ image = Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
+ buffer = BytesIO()
+ image.save(buffer, format="PNG")
+ encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+ return f"data:image/png;base64,{encoded}"
+
+
+def _get_diffusion_feature_cases(model: str):
+ """Single default server row."""
+ return [
+ pytest.param(
+ OmniServerParams(model=model),
+ id="default",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_image_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+ """Default SANA-WM I2V smoke: async ``/v1/videos`` job completes and returns video bytes."""
+ request_config = {
+ "model": omni_server.model,
+ "image_reference": _first_frame_data_url(),
+ "form_data": {
+ "prompt": PROMPT,
+ "negative_prompt": NEGATIVE_PROMPT,
+ # The camera payload rides inside extra_params; there is no
+ # top-level sana_wm form field on the video endpoints, so sending
+ # one is silently dropped and the request fails in the
+ # preprocessor for missing camera control.
+ "extra_params": json.dumps({"sana_wm": SANA_WM_PARAMS}),
+ "height": SANA_WM_OUTPUT_HEIGHT,
+ "width": SANA_WM_OUTPUT_WIDTH,
+ "num_frames": SMOKE_NUM_FRAMES,
+ "fps": 16,
+ "num_inference_steps": SMOKE_NUM_INFERENCE_STEPS,
+ "guidance_scale": 5.0,
+ "seed": 42,
+ },
+ }
+ openai_client.send_video_diffusion_request(request_config)
diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py
index b9b7df4df..a83edfcb9 100644
--- a/tests/entrypoints/openai_api/test_video_server.py
+++ b/tests/entrypoints/openai_api/test_video_server.py
@@ -2168,6 +2168,31 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
assert captured.quality == "high"
+def test_sync_sana_wm_extra_params_payload_passes_to_engine_prompt(test_client, mocker: MockerFixture):
+ _mock_encode_video_bytes(mocker)
+ response = test_client.post(
+ "/v1/videos/sync",
+ data={
+ "prompt": "drive forward",
+ "extra_params": json.dumps(
+ {
+ "sana_wm": {"action": "d-4", "rotation_speed_deg": 1.5},
+ "sana_wm_native_max_tokens": 30000,
+ }
+ ),
+ },
+ )
+
+ assert response.status_code == 200
+ engine = test_client.app.state.openai_serving_video._engine_client
+ captured_params = engine.captured_sampling_params_list[0]
+ # The camera block reaches the model through extra_args; the Sana preprocess
+ # hook is what lifts it onto the prompt.
+ assert captured_params.extra_args["sana_wm"]["action"] == "d-4"
+ assert captured_params.extra_args["sana_wm"]["rotation_speed_deg"] == 1.5
+ assert captured_params.extra_args["sana_wm_native_max_tokens"] == 30000
+
+
def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture):
"""Frame interpolation parameters should be forwarded on the sync path."""
encode_mock = _mock_encode_video_bytes(mocker)
diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py
index 7faf175f8..737eba753 100644
--- a/tests/model_tests/diffusion/test_alignment.py
+++ b/tests/model_tests/diffusion/test_alignment.py
@@ -80,6 +80,7 @@ EXCLUDED_MODELS = [
"Pi0Pipeline",
"SoulXSingerPipeline",
"SoulXSingerSVCPipeline",
+ "SanaWmPipeline",
]
diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py
index ecc53cdc6..e11a7541e 100644
--- a/vllm_omni/diffusion/model_metadata.py
+++ b/vllm_omni/diffusion/model_metadata.py
@@ -56,6 +56,10 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
"WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True),
"WanVACEPipeline": DiffusionModelMetadata(attention_mask_free=True),
"WanS2VPipeline": DiffusionModelMetadata(attention_mask_free=True),
+ "SanaWmPipeline": DiffusionModelMetadata(
+ supports_multimodal_inputs=True,
+ max_multimodal_image_inputs=1,
+ ),
}
diff --git a/vllm_omni/diffusion/models/sana_wm/__init__.py b/vllm_omni/diffusion/models/sana_wm/__init__.py
new file mode 100644
index 000000000..f6f776194
--- /dev/null
+++ b/vllm_omni/diffusion/models/sana_wm/__init__.py
@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SANA-WM diffusion model integration."""
+
+from vllm_omni.diffusion.models.sana_wm.config import SanaWmConfig
+from vllm_omni.diffusion.models.sana_wm.pipeline_sana_wm import (
+ SANA_WM_MODEL_ID,
+ SANA_WM_OUTPUT_HEIGHT,
+ SANA_WM_OUTPUT_WIDTH,
+ SanaWmPipeline,
+ get_sana_wm_pre_process_func,
+)
+from vllm_omni.diffusion.models.sana_wm.request import normalize_sana_wm_payload
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmTransformer3DModel
+
+__all__ = [
+ "SANA_WM_MODEL_ID",
+ "SANA_WM_OUTPUT_HEIGHT",
+ "SANA_WM_OUTPUT_WIDTH",
+ "SanaWmConfig",
+ "SanaWmPipeline",
+ "SanaWmTransformer3DModel",
+ "get_sana_wm_pre_process_func",
+ "normalize_sana_wm_payload",
+]
diff --git a/vllm_omni/diffusion/models/sana_wm/camera_control.py b/vllm_omni/diffusion/models/sana_wm/camera_control.py
new file mode 100644
index 000000000..914fbc0d3
--- /dev/null
+++ b/vllm_omni/diffusion/models/sana_wm/camera_control.py
@@ -0,0 +1,356 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Camera-control helpers for SANA-WM.
+
+This module mirrors the public NVlabs/Sana SANA-WM camera preparation path:
+WASD/IJKL action rollout or explicit camera-to-world poses are converted into
+relative poses, per-latent-frame ray metadata, and per-VAE-chunk Plucker maps.
+Model-local projection layers consume these tensors in the Stage-1 DiT.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+
+SANA_WM_DEFAULT_TRANSLATION_SPEED = 0.05
+SANA_WM_DEFAULT_ROTATION_SPEED_DEG = 1.2
+SANA_WM_DEFAULT_PITCH_LIMIT_DEG = 85.0
+SANA_WM_ALLOWED_ACTION_KEYS = frozenset("wasdijkl")
+SANA_WM_DEFAULT_VAE_STRIDE = (8, 32, 32)
+
+
+@dataclass(frozen=True)
+class SanaWmCameraCondition:
+ poses: Any | None = None
+ intrinsics: Any | None = None
+ action: str | None = None
+ num_frames: int | None = None
+ height: int = 704
+ width: int = 1280
+ 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
+
+
+def _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 _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 _parse_action_segments(action: str) -> list[tuple[list[str], int]]:
+ """Parse an action string into ``(keys, duration)`` pairs without expanding it.
+
+ A single segment names an unbounded rollout (``"w-1000000000"`` is valid
+ syntax), so expanding here would allocate one list entry per frame before
+ anything checks the frame count. Callers total the durations and validate
+ against ``num_frames`` first; only :func:`_parse_action_string` expands.
+ """
+ cleaned = "".join(action.replace("，", ",").split())
+ if not cleaned:
+ raise ValueError("Sana-WM action string is empty.")
+
+ segments: list[tuple[list[str], int]] = []
+ for segment in cleaned.split(","):
+ if not segment or "-" not in segment:
+ raise ValueError(f"Invalid Sana-WM action segment {segment!r}: expected '-'.")
+ keys_part, duration = segment.rsplit("-", 1)
+ if not duration.isdigit() or int(duration) <= 0:
+ raise ValueError(f"Sana-WM action segment {segment!r} has a non-positive duration {duration!r}.")
+ keys_lower = keys_part.lower()
+ if keys_lower == "none":
+ keys: list[str] = []
+ else:
+ bad = sorted({key for key in keys_lower if key not in SANA_WM_ALLOWED_ACTION_KEYS})
+ if bad:
+ allowed = "".join(sorted(SANA_WM_ALLOWED_ACTION_KEYS))
+ raise ValueError(f"Sana-WM action segment {segment!r} contains unknown keys {bad}; allowed: {allowed}.")
+ keys = sorted(set(keys_lower))
+ segments.append((keys, int(duration)))
+ return segments
+
+
+def action_rollout_num_frames(action: str) -> int:
+ """Frames ``action`` rolls out: the identity start pose plus one per step.
+
+ Totals the durations arithmetically, so an oversized request is rejected
+ before anything allocates per-frame state.
+ """
+ return sum(duration for _, duration in _parse_action_segments(action)) + 1
+
+
+def _parse_action_string(action: str) -> list[list[str]]:
+ per_frame: list[list[str]] = []
+ for keys, duration in _parse_action_segments(action):
+ per_frame.extend([keys] * duration)
+ return per_frame
+
+
+def action_string_to_c2w(
+ 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,
+) -> np.ndarray:
+ """Roll out an OpenCV-convention ``(N + 1, 4, 4)`` c2w trajectory."""
+
+ per_frame = _parse_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)
+ current_pitch = 0.0
+ poses = [current.copy()]
+
+ 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)
+ new_pitch = current_pitch + pitch_delta
+ if -pitch_limit_rad <= new_pitch <= pitch_limit_rad:
+ current_pitch = new_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)
+ new_rotation = _rot_y(yaw_delta) @ rotation @ _rot_x(pitch_delta)
+
+ forward = new_rotation[:, 2].copy()
+ forward[1] = 0.0
+ right = new_rotation[:, 0].copy()
+ right[1] = 0.0
+ forward_norm = float(np.linalg.norm(forward))
+ right_norm = float(np.linalg.norm(right))
+ if forward_norm > 0:
+ forward /= forward_norm + 1e-6
+ if right_norm > 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] = new_rotation
+ current[:3, 3] = translation + move
+ poses.append(current.copy())
+
+ return np.stack(poses, axis=0).astype(np.float32)
+
+
+def get_pose_inverse(transform: torch.Tensor) -> torch.Tensor:
+ """Invert homogeneous rigid transforms with shape ``(..., 4, 4)``."""
+
+ rotation = transform[..., :3, :3]
+ translation = transform[..., :3, 3]
+ rotation_inv = rotation.transpose(-1, -2)
+ translation_inv = -torch.matmul(rotation_inv, translation.unsqueeze(-1)).squeeze(-1)
+ output = torch.eye(4, dtype=transform.dtype, device=transform.device).repeat(transform.shape[:-2] + (1, 1))
+ output[..., :3, :3] = rotation_inv
+ output[..., :3, 3] = translation_inv
+ return output
+
+
+def compute_raymap(
+ intrinsics: torch.Tensor,
+ poses: torch.Tensor,
+ height: int,
+ width: int,
+ *,
+ use_plucker: bool = True,
+) -> torch.Tensor:
+ """Compute SANA-WM geometry ray maps.
+
+ Args:
+ intrinsics: ``(T, 4)`` tensor in ``[fx, fy, cx, cy]`` order.
+ poses: ``(T, 4, 4)`` camera-to-world matrices.
+ height: raymap height in latent pixels.
+ width: raymap width in latent pixels.
+ use_plucker: return ``[direction, moment]`` if true, otherwise
+ ``[origin, direction]``.
+ """
+
+ if intrinsics.ndim != 2 or intrinsics.shape[-1] != 4:
+ raise ValueError(f"Sana-WM intrinsics must have shape (T, 4), got {tuple(intrinsics.shape)}.")
+ if poses.ndim != 3 or poses.shape[-2:] != (4, 4):
+ raise ValueError(f"Sana-WM poses must have shape (T, 4, 4), got {tuple(poses.shape)}.")
+ if intrinsics.shape[0] != poses.shape[0]:
+ raise ValueError("Sana-WM intrinsics and poses must have the same temporal length.")
+
+ frames = intrinsics.shape[0]
+ device = intrinsics.device
+ dtype = intrinsics.dtype
+ y_grid, x_grid = torch.meshgrid(
+ torch.arange(height, device=device, dtype=dtype),
+ torch.arange(width, device=device, dtype=dtype),
+ indexing="ij",
+ )
+ x_grid = x_grid[None].expand(frames, -1, -1)
+ y_grid = y_grid[None].expand(frames, -1, -1)
+
+ fx = intrinsics[:, 0].view(frames, 1, 1)
+ fy = intrinsics[:, 1].view(frames, 1, 1)
+ cx = intrinsics[:, 2].view(frames, 1, 1)
+ cy = intrinsics[:, 3].view(frames, 1, 1)
+ dirs_cam = torch.stack([(x_grid - cx) / fx, (y_grid - cy) / fy, torch.ones_like(x_grid)], dim=-1)
+
+ rotation = poses[:, :3, :3]
+ translation = poses[:, :3, 3]
+ dirs_world = torch.einsum("tij,thwj->thwi", rotation, dirs_cam)
+ dirs_world = dirs_world / torch.norm(dirs_world, dim=-1, keepdim=True).clamp_min(1e-12)
+ origins = translation.view(frames, 1, 1, 3).expand_as(dirs_world)
+
+ if use_plucker:
+ moments = torch.cross(origins, dirs_world, dim=-1)
+ return torch.cat([dirs_world, moments], dim=-1)
+ return torch.cat([origins, dirs_world], dim=-1)
+
+
+def intrinsics_to_vec4_array(intrinsics: Any, *, num_frames: int, height: int, width: int) -> np.ndarray:
+ """Normalize the intrinsics payload to an ``(F, 4)`` ``[fx, fy, cx, cy]`` array.
+
+ Accepts either ``None`` (derive from the output resolution) or the
+ ``{fx, fy, cx, cy}`` mapping — the only form the request contract exposes.
+ """
+
+ if intrinsics is None:
+ focal = float(max(height, width))
+ vec = np.array([focal, focal, width / 2.0, height / 2.0], dtype=np.float32)
+ return np.broadcast_to(vec, (num_frames, 4)).copy()
+
+ if isinstance(intrinsics, dict):
+ vec = np.array(
+ [intrinsics["fx"], intrinsics["fy"], intrinsics["cx"], intrinsics["cy"]],
+ dtype=np.float32,
+ )
+ return np.broadcast_to(vec, (num_frames, 4)).copy()
+
+ raise ValueError("Sana-WM intrinsics must be a {fx, fy, cx, cy} mapping.")
+
+
+def _pack_camera_conditions(
+ poses: torch.Tensor,
+ intrinsics_latent: torch.Tensor,
+ *,
+ num_frames: int,
+ latent_height: int,
+ latent_width: int,
+ vae_time_stride: int,
+) -> dict[str, torch.Tensor]:
+ # This yields exactly the latent frame count the VAE produces:
+ # ``len(arange(0, n, s)) == ceil(n / s) == (n - 1) // s + 1``.
+ time_indices = torch.arange(0, num_frames, vae_time_stride)
+
+ raymap = torch.cat([poses[time_indices].reshape(len(time_indices), -1), intrinsics_latent[time_indices]], dim=-1)
+
+ chunks: list[torch.Tensor] = []
+ spatial_raymap_frames: list[torch.Tensor] = []
+ chunk_starts = time_indices - (vae_time_stride - 1)
+ for start, t_idx in zip(chunk_starts, time_indices):
+ start_idx = max(0, int(start))
+ end_idx = start_idx + vae_time_stride
+ chunk_poses = poses[start_idx:end_idx]
+ chunk_intrinsics = intrinsics_latent[start_idx:end_idx]
+ if chunk_poses.shape[0] < vae_time_stride:
+ pad = vae_time_stride - chunk_poses.shape[0]
+ chunk_poses = torch.cat([chunk_poses, chunk_poses[-1:].repeat(pad, 1, 1)], dim=0)
+ chunk_intrinsics = torch.cat([chunk_intrinsics, chunk_intrinsics[-1:].repeat(pad, 1)], dim=0)
+ plucker = compute_raymap(chunk_intrinsics, chunk_poses, latent_height, latent_width, use_plucker=True)
+ chunks.append(plucker.permute(0, 3, 1, 2).reshape(-1, latent_height, latent_width))
+
+ # Per-frame spatial ray-direction map for raymap_embedder: take the
+ # representative frame in each chunk (the anchor at t_idx) and use only
+ # the direction channels (last 3 of 6 in the origin+direction layout).
+ t = int(t_idx)
+ od = compute_raymap(
+ intrinsics_latent[t : t + 1],
+ poses[t : t + 1],
+ latent_height,
+ latent_width,
+ use_plucker=False,
+ ) # [1, H, W, 6]
+ spatial_raymap_frames.append(od[0, :, :, 3:]) # [H, W, 3] — direction only
+
+ chunk_plucker = torch.stack(chunks).permute(1, 0, 2, 3)
+ # spatial_raymap: [3, F_latent, H, W] — consumed by raymap_embedder
+ spatial_raymap = torch.stack(spatial_raymap_frames).permute(3, 0, 1, 2)
+ return {"raymap": raymap, "chunk_plucker": chunk_plucker, "spatial_raymap": spatial_raymap}
+
+
+def build_plucker_condition(
+ condition: SanaWmCameraCondition,
+ *,
+ vae_stride: tuple[int, int, int] | list[int] = SANA_WM_DEFAULT_VAE_STRIDE,
+) -> dict[str, torch.Tensor]:
+ """Build native Stage-1 camera tensors from normalized request metadata."""
+
+ if bool(condition.action) == (condition.poses is not None):
+ raise ValueError("Sana-WM camera condition requires exactly one of action or poses.")
+
+ if condition.action:
+ poses_c2w = action_string_to_c2w(
+ condition.action,
+ translation_speed=condition.translation_speed,
+ rotation_speed_deg=condition.rotation_speed_deg,
+ pitch_limit_deg=condition.pitch_limit_deg,
+ )
+ else:
+ poses_c2w = np.asarray(condition.poses, dtype=np.float32)
+
+ if poses_c2w.ndim != 3 or poses_c2w.shape[-2:] != (4, 4):
+ raise ValueError(f"Sana-WM c2w poses must have shape (F, 4, 4), got {poses_c2w.shape}.")
+
+ num_frames = condition.num_frames or int(poses_c2w.shape[0])
+ if poses_c2w.shape[0] < num_frames:
+ if condition.action:
+ num_frames = int(poses_c2w.shape[0])
+ else:
+ raise ValueError(f"Sana-WM c2w poses length {poses_c2w.shape[0]} is shorter than num_frames {num_frames}.")
+ poses_c2w = poses_c2w[:num_frames]
+
+ intrinsics = intrinsics_to_vec4_array(
+ condition.intrinsics,
+ num_frames=num_frames,
+ height=condition.height,
+ width=condition.width,
+ )
+
+ vae_time_stride = int(vae_stride[0])
+ vae_spatial_stride = int(vae_stride[-1])
+ latent_height = condition.height // vae_spatial_stride
+ latent_width = condition.width // vae_spatial_stride
+
+ poses = torch.from_numpy(poses_c2w).float()
+ first_inv = get_pose_inverse(poses[0:1]).squeeze(0)
+ poses_rel = torch.matmul(first_inv, poses[1:])
+ poses = torch.cat([torch.eye(4, dtype=poses.dtype).unsqueeze(0), poses_rel], dim=0)
+
+ intrinsics_latent = torch.from_numpy(intrinsics).float()
+ intrinsics_latent[:, [0, 2]] *= latent_width / float(condition.width)
+ intrinsics_latent[:, [1, 3]] *= latent_height / float(condition.height)
+
+ return _pack_camera_conditions(
+ poses,
+ intrinsics_latent,
+ num_frames=num_frames,
+ latent_height=latent_height,
+ latent_width=latent_width,
+ vae_time_stride=vae_time_stride,
+ )
diff --git a/vllm_omni/diffusion/models/sana_wm/config.py b/vllm_omni/diffusion/models/sana_wm/config.py
new file mode 100644
index 000000000..bf4eb428c
--- /dev/null
+++ b/vllm_omni/diffusion/models/sana_wm/config.py
@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Config for the SANA-WM Stage-1 transformer.
+
+Loaded from the standard diffusers ``transformer/config.json`` (a flat dict
+keyed by the dataclass field names). The bespoke ``config.yaml`` release format
+is handled once, offline, at checkpoint-conversion time — the runtime only ever
+sees the converted, diffusers-standard config.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from typing import Any
+
+
+def _as_tuple3(value: Any, default: tuple[int, int, int]) -> tuple[int, int, int]:
+ if value is None:
+ return default
+ if isinstance(value, int):
+ return (value, value, value)
+ if isinstance(value, (list, tuple)):
+ if len(value) == 2:
+ return (1, int(value[0]), int(value[1]))
+ if len(value) == 3:
+ return (int(value[0]), int(value[1]), int(value[2]))
+ return default
+
+
+# LTX-2 VAE compression ratios (SANA-WM ships ``AutoencoderKLLTX2Video``).
+# They live here rather than in the pipeline so the request-normalisation module
+# can validate latent geometry without importing the pipeline.
+SANA_WM_VAE_SPATIAL_COMPRESSION = 32
+SANA_WM_VAE_TEMPORAL_COMPRESSION = 8
+
+
+@dataclass(frozen=True)
+class SanaWmConfig:
+ """SANA-WM Stage-1 architecture and runtime config.
+
+ Defaults mirror the first public HF release. ``from_json`` (reading
+ ``transformer/config.json``) is the normal construction path.
+ """
+
+ architecture_name: str | None = "SanaMSVideoCamCtrl_1600M_P1_D20"
+ num_blocks: int = 20
+ hidden_size: int = 2240
+ mlp_ratio: float = 3.0
+ attn_type: str = "BidirectionalGDNTriton"
+ softmax_every_n: int = 4
+ linear_head_dim: int = 112
+ conv_kernel_size: int = 4
+ t_kernel_size: int = 3
+ k_conv_only: bool = True
+ ffn_type: str = "GLUMBConvTemp"
+ pos_embed_type: str = "wan_rope"
+ patch_size: tuple[int, int, int] = (1, 1, 1)
+ qk_norm: bool = True
+ cross_norm: bool = True
+ mixed_precision: str = "bf16"
+ fp32_attention: bool = True
+ image_size: int = 720
+ cam_attn_compress: int = 1
+ use_chunk_plucker_post_attn: bool = True
+ chunk_plucker_channels: int = 48
+ chunk_plucker_post_attn_blocks: int = 20
+ inference_flow_shift: float = 9.8
+ scheduler_type: str = "flow_dpm-solver"
+ chi_prompt: list[str] = field(default_factory=list)
+ model_max_length: int = 300
+
+ @classmethod
+ def from_dict(cls, data: Mapping[str, Any]) -> SanaWmConfig:
+ """Build from a flat mapping of field name -> value.
+
+ Unknown keys (e.g. the diffusers ``_class_name`` marker or the
+ ``latent_channels`` / ``prompt_channels`` constructor kwargs) are
+ ignored. ``patch_size`` is coerced from a JSON list back to a tuple.
+ """
+ field_names = {f.name for f in fields(cls)}
+ known = {k: v for k, v in data.items() if k in field_names}
+ if "patch_size" in known:
+ known["patch_size"] = _as_tuple3(known["patch_size"], cls.patch_size)
+ return cls(**known)
+
+ @classmethod
+ def from_json(cls, path: str | Path) -> SanaWmConfig:
+ with open(path, encoding="utf-8") as f:
+ data = json.load(f)
+ return cls.from_dict(data if isinstance(data, Mapping) else {})
diff --git a/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py b/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py
new file mode 100644
index 000000000..9c2843040
--- /dev/null
+++ b/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py
@@ -0,0 +1,942 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sana-WM pipeline integration.
+
+This module wires the registry-visible surface, release-layout validation, and
+the Stage-1 sampling loop. The DiT, Gated-DeltaNet and camera stack are
+vLLM-Omni-native layers (``sana_wm_transformer.py`` / ``ucpe.py``); nothing here
+calls into NVlabs code at runtime.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, ClassVar
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
+from vllm_omni.diffusion.models.ltx2.ltx2_latents import (
+ denormalize_latents,
+ normalize_latents,
+ pack_latents,
+ resolve_video_latent_shape,
+ unpack_latents,
+)
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.models.sana_wm.camera_control import (
+ SanaWmCameraCondition,
+ build_plucker_condition,
+)
+from vllm_omni.diffusion.models.sana_wm.config import (
+ SANA_WM_VAE_SPATIAL_COMPRESSION,
+ SANA_WM_VAE_TEMPORAL_COMPRESSION,
+ SanaWmConfig,
+)
+from vllm_omni.diffusion.models.sana_wm.request import normalize_sana_wm_payload
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
+ SANA_WM_STAGE1_PROMPT_CHANNELS,
+ SanaWmTransformer3DModel,
+)
+from vllm_omni.diffusion.models.schedulers import FlowMatchEulerDiscreteScheduler
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+ DiffusionPipelineProfilerMixin,
+)
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+logger = init_logger(__name__)
+
+# The loader expects the standard diffusers layout (transformer/config.json +
+# transformer/diffusion_pytorch_model.safetensors + model_index.json), so this
+# repo is an offline conversion of the NVlabs release into that layout. It holds
+# Stage-1 only — transformer/ + vae/, with model_index.json naming
+# SanaWmPipeline — because startup prefetches the whole model repo
+# (allow_patterns=["*"]); a repo that also carried the refiner would pull tens of
+# GB this path never loads.
+SANA_WM_MODEL_ID = "BBBBruce/SANA-WM_bidirectional-stage1-diffusers"
+
+SANA_WM_STAGE1_DIT_FILE = "transformer/diffusion_pytorch_model.safetensors"
+SANA_WM_CONFIG_FILE = "transformer/config.json"
+SANA_WM_VAE_CONFIG_FILE = "vae/config.json"
+SANA_WM_VAE_WEIGHT_FILE = "vae/diffusion_pytorch_model.safetensors"
+SANA_WM_STAGE1_TEXT_ENCODER_ID = "google/gemma-2-2b-it"
+SANA_WM_STAGE1_TEXT_ENCODER_FALLBACK_ID = "Efficient-Large-Model/gemma-2-2b-it"
+SANA_WM_STAGE1_TEXT_ENCODER_ENV = "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER"
+SANA_WM_OUTPUT_HEIGHT = 704
+SANA_WM_OUTPUT_WIDTH = 1280
+SANA_WM_NUM_TRAIN_TIMESTEPS = 1000
+# Maximum supported clip length (20s @16fps) — the native envelope, not a
+# request default. Request-side defaults live in
+# ``models/sana_wm/request.py``.
+SANA_WM_NATIVE_NUM_FRAMES = 321
+# Fail-fast latent-token cap for the native Stage-1 path, sized to the native
+# 704x1280 output at the maximum clip length so it only rejects genuinely
+# oversized requests. Overridable per-request via ``sana_wm_native_max_tokens``.
+SANA_WM_NATIVE_MAX_TOKENS = (
+ ((SANA_WM_NATIVE_NUM_FRAMES - 1) // SANA_WM_VAE_TEMPORAL_COMPRESSION + 1)
+ * (SANA_WM_OUTPUT_HEIGHT // SANA_WM_VAE_SPATIAL_COMPRESSION)
+ * (SANA_WM_OUTPUT_WIDTH // SANA_WM_VAE_SPATIAL_COMPRESSION)
+)
+# The NVlabs production inference settings for this release. They belong in the
+# model rather than in a deploy YAML: single-stage diffusion models resolve
+# through ``create_default_diffusion``, which never reads a YAML's
+# ``default_sampling_params``, so a caller that omits either field lands here.
+SANA_WM_DEFAULT_NUM_INFERENCE_STEPS = 60
+SANA_WM_DEFAULT_GUIDANCE_SCALE = 5.0
+
+SANA_WM_STAGE1_PATTERNS = (
+ SANA_WM_CONFIG_FILE,
+ SANA_WM_STAGE1_DIT_FILE,
+ SANA_WM_VAE_CONFIG_FILE,
+ SANA_WM_VAE_WEIGHT_FILE,
+)
+SANA_WM_STAGE1_DIT_SUBFOLDER = "transformer"
+SANA_WM_STAGE1_DIT_BASENAME = Path(SANA_WM_STAGE1_DIT_FILE).name
+
+
+@dataclass(frozen=True)
+class SanaWmLocalPaths:
+ """Resolved local file paths for a SANA-WM snapshot."""
+
+ root: Path
+ config: Path
+ stage1_dit: Path
+ vae_config: Path
+ vae_weights: Path
+
+
+@dataclass(frozen=True)
+class SanaWmNativeParams:
+ """Resolved generation settings for one Stage-1 request.
+
+ Built by ``_native_params`` from the request payload and sampling params,
+ and consumed by ``_run_native_backend`` — this is the production path.
+ """
+
+ height: int
+ width: int
+ num_frames: int
+ num_inference_steps: int
+ seed: int
+ # Defaults to the NVlabs production scale, so a caller that never mentions
+ # guidance still gets the two-branch CFG forward the model was tuned for.
+ # An explicit ``guidance_scale <= 1.0`` still selects the single-branch path.
+ cfg_scale: float = SANA_WM_DEFAULT_GUIDANCE_SCALE
+
+
+def build_sana_wm_download_patterns() -> tuple[str, ...]:
+ """Return the minimal HF allow-patterns needed for SANA-WM."""
+
+ return tuple(SANA_WM_STAGE1_PATTERNS)
+
+
+def resolve_sana_wm_local_paths(snapshot_dir: str | Path) -> SanaWmLocalPaths:
+ root = Path(snapshot_dir)
+ return SanaWmLocalPaths(
+ root=root,
+ config=root / SANA_WM_CONFIG_FILE,
+ stage1_dit=root / SANA_WM_STAGE1_DIT_FILE,
+ vae_config=root / SANA_WM_VAE_CONFIG_FILE,
+ vae_weights=root / SANA_WM_VAE_WEIGHT_FILE,
+ )
+
+
+def validate_sana_wm_local_paths(paths: SanaWmLocalPaths) -> None:
+ required = [
+ paths.config,
+ paths.stage1_dit,
+ paths.vae_config,
+ paths.vae_weights,
+ ]
+ missing = [path for path in required if not path.exists()]
+ if missing:
+ joined = ", ".join(str(path) for path in missing)
+ raise FileNotFoundError(f"Missing SANA-WM checkpoint files: {joined}")
+
+
+def resolve_or_download_sana_wm_checkpoint(
+ model: str = SANA_WM_MODEL_ID,
+ *,
+ revision: str | None = None,
+ cache_dir: str | None = None,
+) -> SanaWmLocalPaths:
+ """Resolve a local SANA-WM tree or download the required HF files."""
+
+ root = Path(model)
+ if root.is_dir():
+ snapshot_dir = root
+ else:
+ snapshot_dir = Path(
+ download_weights_from_hf_specific(
+ model,
+ cache_dir,
+ list(build_sana_wm_download_patterns()),
+ revision=revision,
+ require_all=True,
+ )
+ )
+ paths = resolve_sana_wm_local_paths(snapshot_dir)
+ validate_sana_wm_local_paths(paths)
+ return paths
+
+
+def build_sana_wm_output_envelope(
+ *,
+ output: Any,
+ output_type: str,
+ metadata: dict[str, Any],
+) -> dict[str, Any]:
+ """Wrap a Stage-1 result in the canonical output envelope.
+
+ ``normalize_diffusion_postprocess_output`` splits ``{"payload": ...,
+ "metadata": ...}`` into the API-facing payload and the metadata groups, so
+ the model-specific diagnostics ride along under a ``sana_wm`` group instead
+ of the removed ``DiffusionOutput.custom_output`` field.
+ """
+ payload_key = "latents" if output_type == "latent" else "video"
+ return {"payload": {payload_key: output}, "metadata": {"sana_wm": metadata}}
+
+
+def get_sana_wm_pre_process_func(od_config: OmniDiffusionConfig):
+ del od_config
+
+ def pre_process_func(request: Any) -> Any:
+ from vllm_omni.inputs.data import OmniTextPrompt
+
+ sampling_params = getattr(request, "sampling_params", None)
+ prompt = request.prompt
+ prompt_mapping = OmniTextPrompt(prompt=prompt) if isinstance(prompt, str) else prompt
+ prompt_mapping = dict(prompt_mapping)
+
+ # The camera block travels as extra_params["sana_wm"], which serving
+ # merges into extra_args. Offline callers may also put it straight on
+ # the prompt.
+ if sampling_params is not None:
+ extra_args = getattr(sampling_params, "extra_args", None)
+ if isinstance(extra_args, Mapping):
+ block = extra_args.get("sana_wm")
+ if isinstance(block, Mapping):
+ merged = dict(block)
+ merged.update(prompt_mapping.get("sana_wm") or {})
+ prompt_mapping["sana_wm"] = merged
+ # Geometry always comes from the standard sampling fields.
+ for field in ("num_frames", "height", "width"):
+ value = getattr(sampling_params, field, None)
+ if value is not None and not (field == "num_frames" and value == 1):
+ prompt_mapping[field] = value
+
+ request.prompt = normalize_sana_wm_payload(prompt_mapping)
+
+ # Preprocess is the only place the request is normalised, so publish the
+ # canonical geometry back onto sampling_params for the rest of the stack.
+ if sampling_params is not None:
+ payload = request.prompt["additional_information"]["sana_wm"]
+ sampling_params.num_frames = payload["num_frames"]
+ sampling_params.height = payload["height"]
+ sampling_params.width = payload["width"]
+ return request
+
+ return pre_process_func
+
+
+class SanaWmPipeline(
+ nn.Module,
+ CFGParallelMixin,
+ SupportImageInput,
+ SupportsComponentDiscovery,
+ ProgressBarMixin,
+ DiffusionPipelineProfilerMixin,
+):
+ """Stage-1 SANA-WM image-to-video pipeline."""
+
+ support_image_input: ClassVar[bool] = True
+ color_format: ClassVar[str] = "RGB"
+ # Skip the startup warmup: the generic dummy request carries a blank image
+ # with no camera poses and no action string, which SANA-WM's request
+ # validation always rejects. The first real request does the warming.
+ dummy_run_num_frames: ClassVar[int] = 0
+ _dit_modules: ClassVar[list[str]] = ["transformer"]
+ _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+ _vae_modules: ClassVar[list[str]] = ["vae"]
+ _resident_modules: ClassVar[list[str]] = []
+
+ def __init__(self, *, od_config: OmniDiffusionConfig | None = None, prefix: str = "") -> None:
+ super().__init__()
+ self.od_config = od_config
+ self.prefix = prefix
+
+ self.sana_wm_config = SanaWmConfig()
+ self.quant_config = getattr(od_config, "quantization_config", None) if od_config is not None else None
+
+ self.tokenizer: Any | None = None
+ self.text_encoder: nn.Module | None = None
+ self.vae: nn.Module | None = None
+ self.release_paths: SanaWmLocalPaths | None = None
+ self.weights_sources = []
+ if od_config is not None and od_config.model is not None:
+ # Resolve the checkpoint once at construction (the standard
+ # startup-time flow, like every other diffusion pipeline) so the
+ # transformer below is built with the transformer/config.json shape
+ # rather than defaults.
+ self.resolve_checkpoint()
+ self.weights_sources = [
+ DiffusersPipelineLoader.ComponentSource(
+ model_or_path=od_config.model,
+ subfolder=SANA_WM_STAGE1_DIT_SUBFOLDER,
+ revision=od_config.revision,
+ # The subfolder file carries module-local keys (standard
+ # diffusers layout); prefix them into the pipeline
+ # namespace for the loader's strict-coverage check.
+ prefix="transformer.",
+ fall_back_to_pt=False,
+ allow_patterns_overrides=[SANA_WM_STAGE1_DIT_BASENAME],
+ )
+ ]
+ # Built eagerly: the loader constructs the pipeline under the target
+ # device / default-dtype context, and the weight loader then streams
+ # checkpoint tensors directly into these modules.
+ self.transformer = SanaWmTransformer3DModel(
+ config=self.sana_wm_config,
+ quant_config=self.quant_config,
+ prefix=f"{prefix}.transformer" if prefix else "transformer",
+ )
+ self.device = get_local_device()
+ if od_config is not None and od_config.model is not None:
+ self._build_aux_components()
+
+ def resolve_checkpoint(self) -> SanaWmLocalPaths:
+ if self.release_paths is not None:
+ return self.release_paths
+ if self.od_config is None or self.od_config.model is None:
+ raise ValueError("Sana-WM checkpoint resolution requires od_config.model.")
+ self.release_paths = resolve_or_download_sana_wm_checkpoint(
+ self.od_config.model,
+ revision=self.od_config.revision,
+ )
+ self.sana_wm_config = SanaWmConfig.from_json(self.release_paths.config)
+ return self.release_paths
+
+ @staticmethod
+ def _extra_args(sampling_params: Any | None) -> dict[str, Any]:
+ extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
+ return dict(extra_args or {})
+
+ def _native_params(self, payload: dict[str, Any], sampling_params: Any | None) -> SanaWmNativeParams:
+ extra_args = self._extra_args(sampling_params)
+ height = int(getattr(sampling_params, "height", None) or payload["height"])
+ width = int(getattr(sampling_params, "width", None) or payload["width"])
+ num_frames = int(getattr(sampling_params, "num_frames", None) or payload["num_frames"])
+ steps = int(
+ getattr(sampling_params, "num_inference_steps", None)
+ or extra_args.get("num_inference_steps")
+ or SANA_WM_DEFAULT_NUM_INFERENCE_STEPS
+ )
+ seed = int(getattr(sampling_params, "seed", None) or extra_args.get("seed", 0))
+ if min(height, width, num_frames, steps) <= 0:
+ raise ValueError("Sana-WM native height, width, num_frames, and steps must be positive.")
+ cfg_scale = SANA_WM_DEFAULT_GUIDANCE_SCALE
+ if sampling_params is not None and getattr(sampling_params, "guidance_scale_provided", False):
+ cfg_scale = float(getattr(sampling_params, "guidance_scale", 1.0) or 1.0)
+ return SanaWmNativeParams(
+ height=height,
+ width=width,
+ num_frames=num_frames,
+ num_inference_steps=steps,
+ seed=seed,
+ cfg_scale=cfg_scale,
+ )
+
+ def _native_dtype(self, device: torch.device) -> torch.dtype:
+ dtype = getattr(self.od_config, "dtype", None) if self.od_config is not None else None
+ if isinstance(dtype, torch.dtype):
+ return dtype
+ return torch.float32 if device.type == "cpu" else torch.bfloat16
+
+ def _runtime_device_dtype(self) -> tuple[torch.device, torch.dtype]:
+ return self.device, self._native_dtype(self.device)
+
+ def _build_aux_components(self) -> None:
+ """Construct the text encoder and VAE at startup, on CPU.
+
+ ``ModuleDiscovery`` skips attributes that are ``None``, so components
+ created lazily on the first request are invisible to sequential
+ offload, HSDP auxiliary placement, and the shared device/memory
+ lifecycle. Building them here and deferring placement to the common
+ logic mirrors LTX-2's ``initialize_pipeline_components``.
+ """
+ dtype = self._native_dtype(self.device)
+ self._load_stage1_text_encoder(dtype=dtype)
+ self._load_vae(dtype=dtype)
+ self._place_aux_components()
+
+ def _place_aux_components(self) -> None:
+ """Move the discovered auxiliary components onto the local device.
+
+ Skipped when offload or HSDP owns placement — those paths expect the
+ components to still be on CPU when they take over.
+ """
+ parallel_config = getattr(self.od_config, "parallel_config", None)
+ if (
+ getattr(self.od_config, "enable_cpu_offload", False)
+ or getattr(self.od_config, "enable_layerwise_offload", False)
+ or getattr(parallel_config, "use_hsdp", False)
+ ):
+ return
+ modules = ModuleDiscovery.discover(self)
+ for module in (*modules.encoders, *modules.vaes, *modules.resident_modules):
+ module.to(self.device)
+
+ @staticmethod
+ def _component_local_files_only(model_id: Any) -> bool:
+ """Decide download policy per component, from its own model id.
+
+ The Stage-1 release directory is local once ``resolve_checkpoint`` has
+ run, but the Gemma text encoder lives in a separate Hub repo. Deriving
+ locality from the release paths forced ``local_files_only=True`` onto
+ Gemma as well, which forbids the download on a box that has not cached
+ it yet.
+ """
+ return Path(str(model_id)).expanduser().exists()
+
+ def _load_stage1_text_encoder(self, *, dtype: torch.dtype) -> None:
+ from transformers import AutoModel, AutoTokenizer
+
+ model_ids = [os.environ.get(SANA_WM_STAGE1_TEXT_ENCODER_ENV, "").strip()]
+ model_ids.extend([SANA_WM_STAGE1_TEXT_ENCODER_ID, SANA_WM_STAGE1_TEXT_ENCODER_FALLBACK_ID])
+ errors: list[str] = []
+ for model_id in dict.fromkeys(model_id for model_id in model_ids if model_id):
+ try:
+ self.tokenizer = AutoTokenizer.from_pretrained(
+ model_id,
+ local_files_only=self._component_local_files_only(model_id),
+ )
+ text_encoder_model_id = model_id
+ break
+ except OSError as exc:
+ errors.append(f"{model_id}: {exc}")
+ else:
+ raise OSError("Could not load SANA-WM Stage-1 text encoder. Tried: " + "; ".join(errors))
+ if getattr(self.tokenizer, "pad_token", None) is None:
+ self.tokenizer.pad_token = self.tokenizer.eos_token
+ self.tokenizer.padding_side = "right"
+ with torch.device("cpu"):
+ # AutoModel, not AutoModelForCausalLM: only the final hidden state is
+ # used, so the LM head would compute a [seq, 256k] logits tensor
+ # that is thrown away on every prompt encode.
+ self.text_encoder = AutoModel.from_pretrained(
+ text_encoder_model_id,
+ torch_dtype=dtype,
+ local_files_only=self._component_local_files_only(text_encoder_model_id),
+ )
+
+ def _ensure_stage1_text_encoder(self, *, device: torch.device, dtype: torch.dtype) -> None:
+ if self.text_encoder is None:
+ self._load_stage1_text_encoder(dtype=dtype)
+ self.text_encoder.to(device=device, dtype=dtype)
+
+ def _load_vae(self, *, dtype: torch.dtype) -> None:
+ if self.release_paths is None and self.od_config is not None and self.od_config.model is not None:
+ self.resolve_checkpoint()
+ if self.release_paths is None:
+ raise ValueError("Sana-WM VAE loading requires a resolved checkpoint.")
+
+ from diffusers import AutoencoderKLLTX2Video
+
+ with torch.device("cpu"):
+ self.vae = AutoencoderKLLTX2Video.from_pretrained(
+ str(self.release_paths.root),
+ subfolder="vae",
+ torch_dtype=dtype,
+ local_files_only=self._component_local_files_only(self.release_paths.root),
+ )
+ self.vae.use_framewise_encoding = True
+ self.vae.use_framewise_decoding = True
+ # NVlabs reads these from its YAML VAE config with 64/96 defaults; the
+ # diffusers VAE config.json only ships scaling_factor, so the defaults
+ # are load-bearing (the keys are genuinely absent here).
+ self.vae.tile_sample_stride_num_frames = int(getattr(self.vae.config, "tile_sample_stride_num_frames", 64))
+ self.vae.tile_sample_min_num_frames = int(getattr(self.vae.config, "tile_sample_min_num_frames", 96))
+
+ def _ensure_vae(self, *, device: torch.device, dtype: torch.dtype) -> None:
+ if self.vae is None:
+ self._load_vae(dtype=dtype)
+ # Tiling is re-asserted here, not in _load_vae: registry.py assigns
+ # ``model.vae.use_tiling = od_config.vae_use_tiling`` after the pipeline
+ # is constructed, and that flag defaults to False. Matching NVlabs'
+ # inference_video_scripts/inference_sana_wm.py, the SANA VAE has to
+ # tile — a 161-frame decode costs about 9 GiB more without it, and a
+ # 321-frame request (41 latent frames) decodes as one 3D volume and
+ # OOMs. Overriding is deliberate, but say so rather than dropping the
+ # caller's setting silently.
+ if not self.vae.use_tiling:
+ logger.warning("Sana-WM forces VAE tiling on; ignoring vae_use_tiling=False (the VAE OOMs without it).")
+ self.vae.use_tiling = True
+ self.vae.to(device=device, dtype=dtype)
+
+ @staticmethod
+ def _preprocess_first_frame(
+ image: Any,
+ *,
+ height: int,
+ width: int,
+ device: torch.device,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ """Return the first-frame image as ``[1, 3, 1, H, W]`` in ``[-1, 1]``.
+
+ ``VideoProcessor.preprocess`` covers the PIL / ndarray / tensor cases.
+ ``resample="bicubic"`` keeps the filter PIL's ``Image.resize`` defaulted
+ to; the processor's own default is lanczos, which would move the
+ conditioning latent.
+ """
+ import numpy as np
+ from diffusers.video_processor import VideoProcessor
+
+ # The processor normalises to [-1, 1] assuming float inputs in [0, 1].
+ if isinstance(image, np.ndarray) and image.dtype == np.uint8:
+ image = image.astype(np.float32) / 255.0
+
+ processor = VideoProcessor(vae_scale_factor=SANA_WM_VAE_SPATIAL_COMPRESSION, resample="bicubic")
+ frame = processor.preprocess(image, height=height, width=width)
+ return frame.unsqueeze(2).to(device=device, dtype=dtype)
+
+ def _vae_normalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+ """LTX-2 VAE per-channel latent normalisation matching NVlabs.
+
+ ``z_norm = (z_raw - latents_mean) * scaling_factor / latents_std``
+
+ The LTX-2 VAE (``AutoencoderKLLTX2Video`` in diffusers) ships with
+ per-channel ``latents_mean`` and ``latents_std`` tensors that must be
+ applied alongside ``scaling_factor`` for the round-trip to be
+ identity. Applying `* scaling_factor` alone gives decoded videos a
+ systematic G/B colour shift (pred RGB means
+ ~(110, 220, 208) vs ref ~(94, 114, 138)).
+
+ Returns the normalised latent in the same dtype as the input.
+ """
+ if self.vae is None:
+ raise RuntimeError("Sana-WM VAE did not initialize.")
+ return normalize_latents(
+ latent,
+ self.vae.latents_mean,
+ self.vae.latents_std,
+ float(self.vae.config.scaling_factor),
+ )
+
+ def _vae_denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+ """Inverse of :meth:`_vae_normalize_latent` — match NVlabs decode."""
+ if self.vae is None:
+ raise RuntimeError("Sana-WM VAE did not initialize.")
+ return denormalize_latents(
+ latent,
+ self.vae.latents_mean,
+ self.vae.latents_std,
+ float(self.vae.config.scaling_factor),
+ )
+
+ def _vae_encode_first_frame(
+ self,
+ image: Any,
+ *,
+ height: int,
+ width: int,
+ latent_height: int,
+ latent_width: int,
+ device: torch.device,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ """Encode the first-frame image to latent space via the LTX-2 VAE.
+
+ Returns shape ``[1, 128, 1, latent_height, latent_width]``.
+ """
+ self._ensure_vae(device=device, dtype=dtype)
+ if self.vae is None:
+ raise RuntimeError("Sana-WM VAE did not initialize.")
+ frame = self._preprocess_first_frame(image, height=height, width=width, device=device, dtype=dtype)
+ with torch.inference_mode():
+ encoded = self.vae.encode(frame)
+ first_latent = encoded.latent_dist.mean
+ first_latent = self._vae_normalize_latent(first_latent.to(dtype=dtype))
+ # The encoder's own rounding can land a pixel off the latent grid the
+ # noise tensor was built on, and the two are concatenated below.
+ if first_latent.shape[-2:] != (latent_height, latent_width):
+ first_latent = F.interpolate(
+ first_latent.squeeze(2),
+ size=(latent_height, latent_width),
+ mode="bilinear",
+ align_corners=False,
+ ).unsqueeze(2)
+ return first_latent # [1, 128, 1, lh, lw]
+
+ def _stage1_prompt_text(self, prompt: dict[str, Any]) -> str:
+ user_prompt = str(prompt.get("prompt") or "")
+ chi_prompt = "\n".join(part for part in self.sana_wm_config.chi_prompt if part)
+ if chi_prompt:
+ return chi_prompt + user_prompt
+ return user_prompt
+
+ def _native_prompt_embeds(
+ self,
+ prompt: dict[str, Any],
+ *,
+ device: torch.device,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ prompt_text = self._stage1_prompt_text(prompt)
+ self._ensure_stage1_text_encoder(device=device, dtype=dtype)
+ tokenizer = self.tokenizer
+ if tokenizer is None or self.text_encoder is None:
+ raise RuntimeError("Sana-WM Stage-1 text encoder did not initialize.")
+ chi_prompt = "\n".join(part for part in self.sana_wm_config.chi_prompt if part)
+ model_max_length = int(self.sana_wm_config.model_max_length)
+ max_length_all = model_max_length
+ if chi_prompt:
+ # NVlabs encodes the chi-prompt-prefixed text with extra room for
+ # the chi tokens, then keeps BOS plus the final model window.
+ max_length_all = len(tokenizer.encode(chi_prompt)) + model_max_length - 2
+ encoded = tokenizer(
+ [prompt_text],
+ padding="max_length",
+ max_length=max_length_all,
+ truncation=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ ).to(device)
+ # last_hidden_state is the same tensor the CausalLM wrapper exposed as
+ # hidden_states[-1] (both are post-final-norm), without retaining every
+ # intermediate layer.
+ hidden_states = self.text_encoder(**encoded).last_hidden_state
+ attention_mask = encoded.attention_mask
+ if chi_prompt:
+ select_index = [0] + list(range(-model_max_length + 1, 0))
+ hidden_states = hidden_states[:, select_index]
+ attention_mask = attention_mask[:, select_index]
+ if hidden_states.shape[-1] != SANA_WM_STAGE1_PROMPT_CHANNELS:
+ raise ValueError(
+ "Sana-WM Stage-1 Gemma hidden size mismatch: expected "
+ f"{SANA_WM_STAGE1_PROMPT_CHANNELS}, got {hidden_states.shape[-1]}."
+ )
+ # Return RAW Gemma hidden states: the transformer's ``attention_y_norm``
+ # RMSNorm does the normalisation, so normalising here double-normalises.
+ prompt_attention_mask = attention_mask.to(device=device, dtype=torch.float32)
+ # An all-ones mask carries no information. Keep it ``None`` so
+ # cross-attention takes the mask-free attention path instead of handing
+ # every block, on every step, a mask the backend must reshape to no
+ # effect. Evaluated once per prompt encode, so the device sync is paid
+ # twice per request (cond + uncond), not once per denoising step.
+ if bool(prompt_attention_mask.all()):
+ prompt_attention_mask = None
+ self._last_prompt_attention_mask = prompt_attention_mask
+ return hidden_states.to(device=device, dtype=dtype)
+
+ def predict_noise(self, **kwargs: Any) -> torch.Tensor:
+ """Single transformer forward.
+
+ Overrides the mixin default, which assumes a tuple-returning transformer
+ and takes ``result[0]``; SANA-WM's returns the noise prediction directly.
+ """
+ return self.transformer(**kwargs)
+
+ def _decode_native_latents(
+ self,
+ latents: torch.Tensor,
+ *,
+ output_type: str,
+ device: torch.device,
+ dtype: torch.dtype,
+ ) -> Any:
+ if output_type == "latent":
+ return latents
+ self._ensure_vae(device=device, dtype=dtype)
+ if self.vae is None:
+ raise RuntimeError("Sana-WM VAE did not initialize.")
+ # NVlabs always passes ``temb=None`` to ``vae.decode`` for this VAE
+ # (``builder.py::vae_decode -> LTX2VAE_diffusers``); match that.
+ denorm = self._vae_denormalize_latent(latents.to(getattr(self.vae, "dtype", dtype)))
+ video = self.vae.decode(denorm, temb=None, return_dict=False)[0]
+ # Post-process to the requested output_type (np/pil/pt). Let failures
+ # surface: silently returning the raw decoder tensor would hand the
+ # caller a wrong-format/range video with no error.
+ from diffusers.video_processor import VideoProcessor
+
+ processor = VideoProcessor(vae_scale_factor=getattr(self.vae, "spatial_compression_ratio", 32))
+ return processor.postprocess_video(video, output_type=output_type)
+
+ def _run_native_backend(
+ self,
+ *,
+ prompt: dict[str, Any],
+ payload: dict[str, Any],
+ sampling_params: Any | None,
+ ) -> DiffusionOutput:
+ params = self._native_params(payload, sampling_params)
+ latent_frames, latent_height, latent_width = resolve_video_latent_shape(
+ params.height,
+ params.width,
+ params.num_frames,
+ vae_spatial_compression_ratio=SANA_WM_VAE_SPATIAL_COMPRESSION,
+ vae_temporal_compression_ratio=SANA_WM_VAE_TEMPORAL_COMPRESSION,
+ )
+ token_count = latent_frames * latent_height * latent_width
+ extra_args = self._extra_args(sampling_params)
+ max_tokens = int(extra_args.get("sana_wm_native_max_tokens", SANA_WM_NATIVE_MAX_TOKENS))
+ if token_count > max_tokens:
+ raise ValueError(
+ "Sana-WM native latent token count exceeds the configured cap. "
+ f"Requested latent tokens={token_count}, max={max_tokens}. Request a smaller "
+ "`height`/`width`/`num_frames`, or raise the cap via "
+ "`sana_wm_native_max_tokens`."
+ )
+
+ device, dtype = self._runtime_device_dtype()
+ generator = torch.Generator(device=device)
+ generator.manual_seed(params.seed)
+ noise = torch.randn(
+ (1, 128, latent_frames, latent_height, latent_width),
+ device=device,
+ dtype=dtype,
+ generator=generator,
+ )
+
+ # Shared flow-matching Euler scheduler. NVlabs ``LTXFlowEuler`` drives
+ # diffusers' ``FlowMatchEulerDiscreteScheduler``; this is vLLM-Omni's
+ # vendored copy of it, including the ``per_token_timesteps`` step branch.
+ # The schedule applies the shift twice (once via the shifted ``sigma_min``
+ # persisted by ``__init__``, once in ``set_timesteps``), which is what
+ # reproduces the NVlabs timestep table (e.g. shift=9.8, N=3 ->
+ # [1000, 909.0, 87.7]).
+ scheduler = FlowMatchEulerDiscreteScheduler(
+ num_train_timesteps=SANA_WM_NUM_TRAIN_TIMESTEPS,
+ shift=self.sana_wm_config.inference_flow_shift,
+ )
+ scheduler.set_timesteps(params.num_inference_steps, device=device)
+ timesteps = scheduler.timesteps
+
+ # ``normalize_sana_wm_payload`` guarantees the key is present but says
+ # nothing about its type. Falling back to pure noise on an unusable
+ # image would silently drop the conditioning and produce a
+ # plausible-looking but wrong video, so this raises instead.
+ import numpy as _np
+
+ first_frame_image = (prompt.get("multi_modal_data") or {}).get("image")
+ if not (hasattr(first_frame_image, "convert") or isinstance(first_frame_image, (_np.ndarray, torch.Tensor))):
+ raise TypeError(
+ "Sana-WM first-frame image must be a PIL Image, numpy ndarray, or "
+ f"torch.Tensor; got {type(first_frame_image).__name__}."
+ )
+ # Per-frame timestep contract matching NVlabs ``LTXFlowEuler.sample``:
+ # frame 0 is the clean VAE-encoded conditioning latent and its timestep
+ # is held at 0 while the sampling loop drives the other frames from
+ # noise.
+ first_latent = self._vae_encode_first_frame(
+ first_frame_image,
+ height=params.height,
+ width=params.width,
+ latent_height=latent_height,
+ latent_width=latent_width,
+ device=device,
+ dtype=dtype,
+ )
+ # Place the CLEAN VAE-encoded first frame at frame 0 and rely on
+ # per-frame timesteps + the `condition_mask` torch.where restore (below)
+ # to keep it invariant across denoising steps.
+ latents = torch.cat([first_latent, noise[:, :, 1:]], dim=2)
+
+ prompt_embeds = self._native_prompt_embeds(
+ prompt,
+ device=device,
+ dtype=dtype,
+ )
+ prompt_attention_mask = self._last_prompt_attention_mask
+
+ # Under CFG, encode the negative prompt up front; each denoise step
+ # then runs a two-branch transformer forward. The negative prompt comes
+ # from the request (``prompt["negative_prompt"]``, set by
+ # ``serving_video``); absent, it defaults to the empty string matching
+ # NVlabs ``GenerationParams.negative_prompt = ""``.
+ do_cfg = params.cfg_scale > 1.0
+ if do_cfg:
+ negative_prompt_obj = {"prompt": str(prompt.get("negative_prompt") or "")}
+ negative_prompt_embeds = self._native_prompt_embeds(
+ negative_prompt_obj,
+ device=device,
+ dtype=dtype,
+ )
+ negative_prompt_attention_mask = self._last_prompt_attention_mask
+ # Restore the cond mask so later readers of
+ # `_last_prompt_attention_mask` see the positive branch.
+ self._last_prompt_attention_mask = prompt_attention_mask
+ else:
+ negative_prompt_embeds = None
+ negative_prompt_attention_mask = None
+
+ camera = payload.get("camera") or {}
+ condition = SanaWmCameraCondition(
+ poses=camera.get("poses") if isinstance(camera, dict) else None,
+ intrinsics=payload.get("intrinsics"),
+ action=payload.get("action"),
+ num_frames=params.num_frames,
+ height=params.height,
+ width=params.width,
+ translation_speed=float(payload.get("translation_speed", 0.05)),
+ rotation_speed_deg=float(payload.get("rotation_speed_deg", 1.2)),
+ )
+ camera_tensors = build_plucker_condition(condition)
+ plucker = camera_tensors["chunk_plucker"].to(device=device, dtype=dtype)
+ raymap = camera_tensors["raymap"].to(device=device, dtype=dtype)
+ spatial_raymap = camera_tensors.get("spatial_raymap")
+ if spatial_raymap is not None:
+ spatial_raymap = spatial_raymap.to(device=device, dtype=dtype)
+
+ # Build a per-frame condition mask matching NVlabs ``LTXFlowEuler``:
+ # frame 0 is the conditioning frame at sigma=0 (preserved after
+ # each ``scheduler.step``), all other frames carry the current
+ # sampling sigma. We omit the legacy `add_noise_to_image_conditioning_latents`
+ # motion-continuity term because the public SANA-WM config sets
+ # ``condition_frame_info={0: 0.0}`` (image_cond_noise_scale=0).
+ cam_batch = latents.shape[0]
+ cam_frames = latents.shape[2]
+ condition_mask = torch.zeros(cam_batch, 1, cam_frames, 1, 1, device=latents.device, dtype=latents.dtype)
+ condition_mask[:, :, 0] = 1.0
+
+ for timestep in timesteps:
+ # (B, 1, F) per-frame timestep, frame 0 forced to 0. Stays fp32:
+ # casting through the bf16 latent dtype would quantise it before
+ # ``SanaWmTimestepEmbedder``'s sinusoidal embed.
+ model_timestep = timestep.float().expand(cam_batch, 1, cam_frames).clone()
+ model_timestep[:, :, 0] = 0.0
+ positive_kwargs = {
+ "hidden_states": latents,
+ "timestep": model_timestep,
+ "encoder_hidden_states": prompt_embeds,
+ "encoder_attention_mask": prompt_attention_mask,
+ "plucker": plucker,
+ "raymap": raymap,
+ "spatial_raymap": spatial_raymap,
+ }
+ negative_kwargs = (
+ {
+ **positive_kwargs,
+ "encoder_hidden_states": negative_prompt_embeds,
+ "encoder_attention_mask": negative_prompt_attention_mask,
+ }
+ if do_cfg
+ else None
+ )
+ # ``cfg_normalize=False`` keeps the plain NVlabs combine,
+ # ``uncond + scale * (cond - uncond)``. Under cfg_parallel_size > 1
+ # the mixin runs one branch per rank and all-gathers instead of
+ # evaluating both branches serially.
+ noise_pred = self.predict_noise_maybe_with_cfg(
+ do_cfg,
+ params.cfg_scale,
+ positive_kwargs,
+ negative_kwargs,
+ cfg_normalize=False,
+ )
+
+ # The per-frame contract pairs with a per-token flow-matching Euler
+ # step. Both step branches end at ``prev = sample + dt *
+ # model_output``, but the scheduler defines dt as
+ # ``sigma - sigma_next`` per-token against ``sigma_next - sigma``
+ # standard, so the per-token call takes the sign-flipped prediction
+ # (matching NVlabs ``LTXFlowEuler.sample``); passing the standard
+ # sign here would integrate towards more noise.
+ #
+ # Broadcast the (B, 1, F) timestep to (B, 1, F, H, W) and flatten
+ # F*H*W; frame-0 conditioning tokens are already 0. The flatten
+ # order (F, H, W) matches ``pack_latents``.
+ pt_t = (
+ model_timestep.unsqueeze(-1)
+ .unsqueeze(-1)
+ .expand(cam_batch, 1, cam_frames, latents.shape[3], latents.shape[4])
+ .reshape(cam_batch, -1)
+ )
+ # The per-token branch consumes/returns the packed (B, N, C) layout
+ # and keeps its result in fp32; cast back so the latent dtype stays
+ # stable across steps.
+ stepped_packed = scheduler.step(
+ -pack_latents(noise_pred),
+ timestep,
+ pack_latents(latents),
+ per_token_timesteps=pt_t,
+ return_dict=False,
+ )[0]
+ stepped = unpack_latents(
+ stepped_packed.to(latents.dtype),
+ cam_frames,
+ latents.shape[3],
+ latents.shape[4],
+ )
+
+ # Match NVlabs LTXFlowEuler exactly. This is stricter than a plain
+ # condition-frame restore: with bf16 latents, the `t=1000`
+ # comparison rounds `1 - 1e-6` to `1`, so the first full-noise step
+ # is discarded for generated tokens as well.
+ tokens_to_denoise_mask = timestep / float(SANA_WM_NUM_TRAIN_TIMESTEPS) - 1e-6 < (1.0 - condition_mask)
+ latents = torch.where(tokens_to_denoise_mask, stepped, latents)
+
+ output_type = getattr(sampling_params, "output_type", None) or "np"
+ output = self._decode_native_latents(latents, output_type=output_type, device=device, dtype=dtype)
+ return DiffusionOutput(
+ output=build_sana_wm_output_envelope(
+ output=output,
+ output_type=output_type,
+ metadata={
+ "backend": "native_gdn",
+ "output_space": output_type,
+ "chi_prompt_applied": bool(self.sana_wm_config.chi_prompt),
+ "num_frames": params.num_frames,
+ "height": params.height,
+ "width": params.width,
+ "latent_tokens": token_count,
+ "sampling_steps": params.num_inference_steps,
+ },
+ ),
+ )
+
+ def forward(self, req: DiffusionRequestBatch, *args: Any, **kwargs: Any) -> DiffusionOutput:
+ del args, kwargs
+ if len(req.prompts) != 1:
+ raise ValueError("Sana-WM native backend currently supports exactly one prompt per request.")
+
+ prompt = req.prompts[0]
+ if isinstance(prompt, str):
+ raise ValueError("Sana-WM requires a mapping prompt with first-frame image and camera/action metadata.")
+ # get_sana_wm_pre_process_func is the single canonical entry point: it
+ # normalises the request and republishes the geometry onto
+ # sampling_params. normalize_sana_wm_payload is idempotent, so this only
+ # covers offline callers that bypass the preprocess hook.
+ prompt = normalize_sana_wm_payload(prompt)
+ payload = prompt["additional_information"]["sana_wm"]
+
+ # The checkpoint is resolved once in __init__; this is a cached no-op
+ # kept as a guard for pipelines constructed without od_config.model.
+ if self.od_config is not None and self.od_config.model is not None:
+ self.resolve_checkpoint()
+
+ return self._run_native_backend(
+ prompt=prompt,
+ payload=payload,
+ sampling_params=req.sampling_params,
+ )
+
+ def load_weights(self, weights: Iterable[tuple[str, Any]]) -> set[str]:
+ # Keys arrive pipeline-prefixed (``transformer.*``, per the component
+ # source). AutoWeightsLoader streams them to the submodule and returns
+ # the prefixed names the loader's strict-coverage check compares.
+ loader = AutoWeightsLoader(self)
+ return loader.load_weights(weights)
diff --git a/vllm_omni/diffusion/models/sana_wm/request.py b/vllm_omni/diffusion/models/sana_wm/request.py
new file mode 100644
index 000000000..010675045
--- /dev/null
+++ b/vllm_omni/diffusion/models/sana_wm/request.py
@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Request normalization for Sana-WM image-to-video.
+
+The released SANA-WM checkpoint is first-frame image-to-video. This module keeps
+request-side normalization lightweight: it validates the image/camera/action
+contract and stores a canonical payload under
+``additional_information["sana_wm"]``. Model-internal raymap / Plucker
+projection lives with the transformer.
+
+This lives beside the model rather than under ``stage_input_processors``
+because nothing loads it as one: the pipeline's ``pre_process_func`` is
+registered from ``pipeline_sana_wm``, the same as every other diffusion model,
+and no deploy config wires a ``custom_process_input_func`` here. Keeping it in
+the stage package forced the pipeline to import *upwards* — the only such
+import in ``diffusion/models/`` — which in turn made the VAE compression
+constants impossible to share and so duplicated.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from vllm_omni.diffusion.models.sana_wm.camera_control import action_rollout_num_frames
+from vllm_omni.diffusion.models.sana_wm.config import (
+ SANA_WM_VAE_SPATIAL_COMPRESSION,
+ SANA_WM_VAE_TEMPORAL_COMPRESSION,
+)
+
+SANA_WM_CANONICAL_KEY = "sana_wm"
+# Request fallbacks used whenever the field is absent, mirroring the deploy YAML
+# (161 frames) and the model's native resolution. Both are genuinely reachable:
+# a served request that omits num_frames never gets one injected, and callers
+# may ask for a different resolution.
+SANA_WM_DEFAULT_NUM_FRAMES = 161
+SANA_WM_DEFAULT_HEIGHT = 704
+SANA_WM_DEFAULT_WIDTH = 1280
+SANA_WM_DEFAULT_CAMERA_FORMAT = "c2w_4x4"
+SANA_WM_DEFAULT_COORDINATE_SYSTEM = "official"
+
+
+def _unwrap_single(value: Any) -> Any:
+ if isinstance(value, list):
+ return value[0] if value else None
+ return value
+
+
+def _as_dict(value: Any, *, name: str) -> dict[str, Any]:
+ if value is None:
+ return {}
+ if not isinstance(value, Mapping):
+ raise ValueError(f"Sana-WM {name} must be a mapping, got {type(value).__name__}.")
+ return dict(value)
+
+
+def _is_sequence(value: Any) -> bool:
+ return isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray))
+
+
+def _as_finite_float(value: Any, *, name: str) -> float:
+ """Parse a float and reject NaN/inf.
+
+ Non-finite geometry does not fail loudly downstream: it propagates through
+ the raymap into the Plucker maps and comes back as an all-NaN video with a
+ 200 response.
+ """
+ try:
+ parsed = float(value)
+ except (TypeError, ValueError) as exc:
+ raise ValueError(f"Sana-WM {name} must be a number, got {value!r}.") from exc
+ if not math.isfinite(parsed):
+ raise ValueError(f"Sana-WM {name} must be finite, got {parsed!r}.")
+ return parsed
+
+
+def _validate_numeric_matrix4x4(matrix: Any, *, name: str) -> None:
+ if not _is_sequence(matrix) or len(matrix) != 4:
+ raise ValueError(f"Sana-WM {name} must be a 4x4 matrix.")
+ for row_idx, row in enumerate(matrix):
+ if not _is_sequence(row) or len(row) != 4:
+ raise ValueError(f"Sana-WM {name}[{row_idx}] must contain 4 values.")
+ for col_idx, value in enumerate(row):
+ _as_finite_float(value, name=f"{name}[{row_idx}][{col_idx}]")
+
+
+def _validate_camera_poses(poses: Any, *, num_frames: int | None = None) -> None:
+ if not _is_sequence(poses) or len(poses) == 0:
+ raise ValueError("Sana-WM camera poses must be a non-empty sequence of 4x4 matrices.")
+ if num_frames is not None and len(poses) != num_frames:
+ raise ValueError(f"Sana-WM camera poses length {len(poses)} must equal num_frames {num_frames}.")
+ for idx, pose in enumerate(poses):
+ _validate_numeric_matrix4x4(pose, name=f"camera.poses[{idx}]")
+
+
+def _validate_intrinsics(intrinsics: Any) -> None:
+ # Only the ``{fx, fy, cx, cy}`` mapping is accepted (the one form the recipe
+ # and examples use). Camera intrinsics are optional; when omitted the model
+ # derives them from the output resolution.
+ if intrinsics is None:
+ return
+ if not isinstance(intrinsics, Mapping):
+ raise ValueError("Sana-WM intrinsics must be a {fx, fy, cx, cy} mapping.")
+ required = {"fx", "fy", "cx", "cy"}
+ missing = sorted(required - set(intrinsics))
+ if missing:
+ raise ValueError(f"Sana-WM intrinsics mapping is missing keys: {missing}.")
+ for key in required:
+ value = _as_finite_float(intrinsics[key], name=f"intrinsics[{key!r}]")
+ # ``compute_raymap`` divides by fx/fy. cx/cy stay unconstrained beyond
+ # finiteness — a principal point outside the frame is unusual, not wrong.
+ if key in ("fx", "fy") and value <= 0:
+ raise ValueError(f"Sana-WM intrinsics[{key!r}] must be positive, got {value}.")
+
+
+def _as_positive_int(value: Any, *, name: str) -> int:
+ try:
+ parsed = int(value)
+ except (TypeError, ValueError) as exc:
+ raise ValueError(f"Sana-WM {name} must be an integer, got {value!r}.") from exc
+ if parsed <= 0:
+ raise ValueError(f"Sana-WM {name} must be positive, got {parsed}.")
+ return parsed
+
+
+def _as_positive_float(value: Any, *, name: str) -> float:
+ # Finiteness first: every comparison against NaN is False, so a bare
+ # ``parsed <= 0`` would accept it.
+ parsed = _as_finite_float(value, name=name)
+ if parsed <= 0:
+ raise ValueError(f"Sana-WM {name} must be positive, got {parsed}.")
+ return parsed
+
+
+def _extract_image(mm_data: Mapping[str, Any]) -> Any:
+ return _unwrap_single(mm_data.get("image"))
+
+
+def _first_present(*values: Any) -> Any:
+ for value in values:
+ if value is not None:
+ return value
+ return None
+
+
+def _validate_latent_geometry(*, num_frames: int, height: int, width: int) -> None:
+ """Reject geometry the VAE would silently floor.
+
+ The SANA-WM VAE compresses 32x spatially and 8x temporally, so a request
+ that is not aligned loses pixels/frames instead of failing.
+ """
+ if height % SANA_WM_VAE_SPATIAL_COMPRESSION or width % SANA_WM_VAE_SPATIAL_COMPRESSION:
+ raise ValueError(
+ f"Sana-WM height and width must be divisible by {SANA_WM_VAE_SPATIAL_COMPRESSION}, "
+ f"got {height} and {width}."
+ )
+ if (num_frames - 1) % SANA_WM_VAE_TEMPORAL_COMPRESSION:
+ raise ValueError(
+ f"Sana-WM num_frames must satisfy (num_frames - 1) % {SANA_WM_VAE_TEMPORAL_COMPRESSION} == 0, "
+ f"got {num_frames}."
+ )
+
+
+def normalize_sana_wm_payload(prompt: Mapping[str, Any]) -> dict[str, Any]:
+ """Return a prompt copy with canonical Sana-WM request metadata.
+
+ The ``sana_wm`` block is read from the top-level ``sana_wm`` key, falling
+ back to ``additional_information["sana_wm"]`` so the function is idempotent
+ when called again on its own output. The first-frame image is read from
+ ``multi_modal_data["image"]``.
+ """
+
+ if not isinstance(prompt, Mapping):
+ raise ValueError(f"Sana-WM prompt must be a mapping, got {type(prompt).__name__}.")
+
+ result = dict(prompt)
+ mm_data = _as_dict(result.get("multi_modal_data"), name="multi_modal_data")
+ image = _extract_image(mm_data)
+ if image is None:
+ raise ValueError("Sana-WM requires a first-frame image in `multi_modal_data['image']`.")
+
+ additional = _as_dict(result.get("additional_information"), name="additional_information")
+ raw = _first_present(
+ result.get(SANA_WM_CANONICAL_KEY),
+ additional.get(SANA_WM_CANONICAL_KEY),
+ {},
+ )
+ raw = _as_dict(raw, name=SANA_WM_CANONICAL_KEY)
+
+ num_frames = _as_positive_int(
+ _first_present(raw.get("num_frames"), result.get("num_frames"), SANA_WM_DEFAULT_NUM_FRAMES),
+ name="num_frames",
+ )
+ height = _as_positive_int(
+ _first_present(raw.get("height"), result.get("height"), SANA_WM_DEFAULT_HEIGHT),
+ name="height",
+ )
+ width = _as_positive_int(_first_present(raw.get("width"), result.get("width"), SANA_WM_DEFAULT_WIDTH), name="width")
+ _validate_latent_geometry(num_frames=num_frames, height=height, width=width)
+
+ camera = raw.get("camera")
+ camera = _as_dict(camera, name="camera") if camera is not None else None
+ action = raw.get("action")
+ if action is not None:
+ if not isinstance(action, str) or not action.strip():
+ raise ValueError("Sana-WM action must be a non-empty string when provided.")
+ # Rejected here rather than in the model: a mismatch used to shorten the
+ # camera tensor silently and then fail deep in the transformer on a
+ # camera/latent frame mismatch, and an unbounded duration used to be
+ # expanded one entry per frame before any of that.
+ rollout_frames = action_rollout_num_frames(action)
+ if rollout_frames != num_frames:
+ raise ValueError(
+ f"Sana-WM action rolls out {rollout_frames} frames but num_frames is {num_frames}. "
+ "Action segment durations must sum to num_frames - 1."
+ )
+ if camera is not None and action is not None:
+ raise ValueError("Sana-WM accepts exactly one of `camera` or `action`, not both.")
+ if camera is None and action is None:
+ raise ValueError("Sana-WM requires either explicit camera poses or an action DSL string.")
+
+ canonical_camera = None
+ if camera is not None:
+ poses = camera.get("poses")
+ if poses is None:
+ raise ValueError("Sana-WM camera payload requires `poses`.")
+ _validate_camera_poses(poses, num_frames=num_frames)
+ # Only the defaults are accepted: nothing downstream reads these two
+ # fields, so honouring another value would mean interpreting w2c poses
+ # as c2w and silently generating the wrong trajectory.
+ camera_format = camera.get("format", SANA_WM_DEFAULT_CAMERA_FORMAT)
+ if camera_format != SANA_WM_DEFAULT_CAMERA_FORMAT:
+ raise ValueError(
+ f"Sana-WM only accepts camera format {SANA_WM_DEFAULT_CAMERA_FORMAT!r}; "
+ f"convert the poses before sending, got {camera_format!r}."
+ )
+ coordinate_system = camera.get("coordinate_system", SANA_WM_DEFAULT_COORDINATE_SYSTEM)
+ if coordinate_system != SANA_WM_DEFAULT_COORDINATE_SYSTEM:
+ raise ValueError(
+ f"Sana-WM only accepts coordinate system {SANA_WM_DEFAULT_COORDINATE_SYSTEM!r}, "
+ f"got {coordinate_system!r}."
+ )
+ canonical_camera = {"poses": poses}
+
+ intrinsics = raw.get("intrinsics")
+ _validate_intrinsics(intrinsics)
+ translation_speed = raw.get("translation_speed")
+ rotation_speed_deg = raw.get("rotation_speed_deg")
+
+ canonical = {
+ "num_frames": num_frames,
+ "height": height,
+ "width": width,
+ "camera": canonical_camera,
+ "action": action,
+ "intrinsics": intrinsics,
+ }
+ if translation_speed is not None:
+ canonical["translation_speed"] = _as_positive_float(translation_speed, name="translation_speed")
+ if rotation_speed_deg is not None:
+ canonical["rotation_speed_deg"] = _as_positive_float(rotation_speed_deg, name="rotation_speed_deg")
+
+ mm_data["image"] = image
+ result["multi_modal_data"] = mm_data
+ additional[SANA_WM_CANONICAL_KEY] = canonical
+ result["additional_information"] = additional
+ result.pop(SANA_WM_CANONICAL_KEY, None)
+ return result
diff --git a/vllm_omni/diffusion/models/sana_wm/sana_wm_transformer.py b/vllm_omni/diffusion/models/sana_wm/sana_wm_transformer.py
new file mode 100644
index 000000000..07537ec09
--- /dev/null
+++ b/vllm_omni/diffusion/models/sana_wm/sana_wm_transformer.py
@@ -0,0 +1,2065 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SANA-WM Stage-1 transformer.
+
+Native vLLM-Omni port of the NVlabs SANA-WM DiT. Modules are built eagerly at
+construction (under the loader's target-device context) and checkpoint tensors
+are streamed in via ``load_weights``. The Bidirectional Gated DeltaNet
+recurrence is implemented in pure PyTorch in this module.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable
+from typing import Any, ClassVar
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.embeddings import Timesteps
+from torch import nn
+from vllm.distributed import (
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+ tensor_model_parallel_all_reduce,
+)
+from vllm.distributed.parallel_state import get_tp_group
+from vllm.model_executor.layers.linear import (
+ ColumnParallelLinear,
+ MergedColumnParallelLinear,
+ QKVParallelLinear,
+ RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.sana_wm.config import SanaWmConfig
+from vllm_omni.diffusion.models.sana_wm.ucpe import (
+ SanaWmCamGeometry,
+ cam_prep_func,
+ prepare_cam_geometry,
+)
+
+SANA_WM_STAGE1_LATENT_CHANNELS = 128
+SANA_WM_STAGE1_PROMPT_CHANNELS = 2304
+SANA_WM_STAGE1_TIMESTEP_CHANNELS = 256
+
+
+def _shard_param_across_tp(param: torch.Tensor | None, dim: int = 0) -> None:
+ """Attach vLLM's TP shard loader to a plain (non-parallel-layer) parameter.
+
+ vLLM parallel layers narrow full checkpoint tensors to the local shard via
+ their own ``weight_loader``. SANA-WM also has plain parameters that become
+ TP-local - the GDN vectors (``A_log``, ``dt_bias``), the depthwise temporal
+ convs and the q/k norms - so they need the same behavior at load time.
+ Attached only under TP > 1: at TP = 1 the checkpoint tensor already matches
+ the parameter and ``sharded_weight_loader`` would require an initialized TP
+ group that standalone/unit-test builds do not have.
+ """
+ if param is None or get_tensor_model_parallel_world_size() <= 1:
+ return
+ param.weight_loader = sharded_weight_loader(dim)
+
+
+def _is_sana_wm_transformer_block(name: str, module: Any) -> bool:
+ del module
+ parts = name.split(".")
+ return len(parts) == 2 and parts[0] == "blocks" and parts[1].isdigit()
+
+
+def _to_3tuple(value: int | tuple[int, int] | tuple[int, int, int]) -> tuple[int, int, int]:
+ if isinstance(value, int):
+ return (value, value, value)
+ if len(value) == 2:
+ return (1, int(value[0]), int(value[1]))
+ return (int(value[0]), int(value[1]), int(value[2]))
+
+
+# ---------------------------------------------------------------------------
+# Bidirectional Gated DeltaNet recurrence (pure PyTorch)
+#
+# Ported from the Apache-2.0 NVlabs/Sana reference. SANA-WM's bidirectional
+# video-latent recurrence is a different contract than vLLM's autoregressive
+# Qwen3-Next GDN cache path, so it stays model-local.
+# ---------------------------------------------------------------------------
+
+
+def _validate_gdn_inputs(
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ beta: torch.Tensor,
+ decay: torch.Tensor,
+ spatial_tokens: int,
+) -> tuple[int, int, int, int, int]:
+ if query.ndim != 4 or key.ndim != 4 or value.ndim != 4:
+ raise ValueError("Sana-WM GDN query/key/value must be shaped [B, H, D, N].")
+ if query.shape != key.shape or query.shape != value.shape:
+ raise ValueError(
+ "Sana-WM GDN query/key/value shapes must match, got "
+ f"{tuple(query.shape)}, {tuple(key.shape)}, {tuple(value.shape)}."
+ )
+ if beta.ndim != 4:
+ raise ValueError("Sana-WM GDN beta must be shaped [B, H, T, S].")
+ if decay.ndim != 3:
+ raise ValueError("Sana-WM GDN decay must be shaped [B, H, T].")
+ if spatial_tokens <= 0:
+ raise ValueError("Sana-WM GDN spatial_tokens must be positive.")
+
+ batch_size, num_heads, head_dim, token_count = query.shape
+ if token_count % spatial_tokens != 0:
+ raise ValueError(f"Sana-WM GDN token count {token_count} is not divisible by spatial_tokens={spatial_tokens}.")
+ frames = token_count // spatial_tokens
+ if beta.shape != (batch_size, num_heads, frames, spatial_tokens):
+ raise ValueError(
+ "Sana-WM GDN beta shape mismatch: expected "
+ f"{(batch_size, num_heads, frames, spatial_tokens)}, got {tuple(beta.shape)}."
+ )
+ if decay.shape != (batch_size, num_heads, frames):
+ raise ValueError(
+ f"Sana-WM GDN decay shape mismatch: expected {(batch_size, num_heads, frames)}, got {tuple(decay.shape)}."
+ )
+ return batch_size, num_heads, head_dim, token_count, frames
+
+
+def _delta_scan(
+ query_rot: torch.Tensor,
+ key_rot: torch.Tensor,
+ value: torch.Tensor,
+ beta: torch.Tensor,
+ decay: torch.Tensor,
+ *,
+ spatial_tokens: int,
+ query: torch.Tensor | None = None,
+ key: torch.Tensor | None = None,
+ skip_z: bool = False,
+ flip_output: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """One-directional gated delta-rule recurrence over frames.
+
+ The KV stream consumes only the rotary-embedded Q/K; the Z (denominator)
+ stream consumes only the non-rotary Q/K. ``skip_z=True`` drops the Z stream
+ entirely and returns ``(numerator, None)`` — the numerator-only variant the
+ camera branch needs (NVlabs ``torch_recurrent_cam_single_path_delta_rule``),
+ so ``query``/``key`` may be omitted there.
+
+ ``flip_output=True`` emits the frames in reverse of the order they were
+ computed, which is what a caller running on flipped inputs wants; it costs a
+ list reversal instead of a flip over the assembled result.
+ """
+ if not skip_z and (query is None or key is None):
+ raise ValueError("Sana-WM delta scan needs non-rotary query/key unless skip_z is set.")
+
+ batch_size, num_heads, head_dim, token_count = query_rot.shape
+ frames = beta.shape[2]
+ if token_count != frames * spatial_tokens:
+ raise ValueError(f"Sana-WM delta scan token_count={token_count} != frames*S={frames * spatial_tokens}.")
+
+ def to_frames(tensor: torch.Tensor) -> torch.Tensor:
+ return tensor.view(batch_size, num_heads, head_dim, frames, spatial_tokens).permute(0, 1, 3, 2, 4)
+
+ query_rot_f = to_frames(query_rot)
+ key_rot_f = to_frames(key_rot)
+ value_f = to_frames(value)
+ state_kv = torch.zeros(batch_size, num_heads, head_dim, head_dim, device=query_rot.device, dtype=query_rot.dtype)
+ numerators: list[torch.Tensor] = []
+
+ if skip_z:
+ query_f = key_f = None
+ state_z = None
+ denominators = None
+ else:
+ query_f = to_frames(query)
+ key_f = to_frames(key)
+ state_z = torch.zeros(batch_size, num_heads, head_dim, 1, device=query_rot.device, dtype=query_rot.dtype)
+ denominators = []
+
+ for frame_idx in range(frames):
+ query_rot_t = query_rot_f[:, :, frame_idx]
+ key_rot_t = key_rot_f[:, :, frame_idx]
+ value_t = value_f[:, :, frame_idx]
+ beta_t = beta[:, :, frame_idx].unsqueeze(2)
+ decay_t = decay[:, :, frame_idx].view(batch_size, num_heads, 1, 1)
+
+ state_kv = state_kv * decay_t
+ value_pred = torch.matmul(state_kv, key_rot_t)
+ delta_value = (value_t - value_pred) * beta_t
+ state_kv = state_kv + torch.matmul(delta_value, key_rot_t.transpose(-1, -2))
+ numerators.append(torch.matmul(state_kv, query_rot_t))
+
+ if skip_z:
+ continue
+ query_t = query_f[:, :, frame_idx]
+ key_t = key_f[:, :, frame_idx]
+ state_z = state_z * decay_t
+ z_pred = torch.matmul(state_z.transpose(-1, -2), key_t)
+ delta_z = (1.0 - z_pred) * beta_t
+ state_z = state_z + torch.matmul(key_t, delta_z.transpose(-1, -2))
+ denominators.append(torch.matmul(state_z.transpose(-1, -2), query_t))
+
+ def restore(tensors: list[torch.Tensor], dim: int) -> torch.Tensor:
+ stacked = torch.stack(tensors[::-1] if flip_output else tensors, dim=2)
+ return stacked.permute(0, 1, 3, 2, 4).reshape(batch_size, num_heads, dim, token_count)
+
+ return restore(numerators, head_dim), (None if skip_z else restore(denominators, 1))
+
+
+def _bidirectional_delta_scan(
+ query_rot: torch.Tensor,
+ key_rot: torch.Tensor,
+ value: torch.Tensor,
+ beta: torch.Tensor,
+ decay: torch.Tensor,
+ *,
+ spatial_tokens: int,
+ query: torch.Tensor | None = None,
+ key: torch.Tensor | None = None,
+ skip_z: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Forward + backward delta scan, summed.
+
+ The backward direction shifts K/V/beta by one frame (zero pad) and the decay
+ by one frame (neutral 1.0 pad), matching the NVlabs ``flip_and_shift``
+ convention. The backward scan also emits its frames in forward order, so the
+ two directions add directly instead of flipping the assembled result back.
+ """
+ frames = beta.shape[2]
+
+ def reverse(tensor: torch.Tensor, *, shift_value: float | None = None) -> torch.Tensor:
+ return _reverse_frames(tensor, frames=frames, spatial_tokens=spatial_tokens, shift_value=shift_value)
+
+ num_fwd, den_fwd = _delta_scan(
+ query_rot,
+ key_rot,
+ value,
+ beta,
+ decay,
+ spatial_tokens=spatial_tokens,
+ query=query,
+ key=key,
+ skip_z=skip_z,
+ )
+
+ bwd_kwargs: dict[str, torch.Tensor] = {}
+ if not skip_z:
+ bwd_kwargs["query"] = reverse(query)
+ bwd_kwargs["key"] = reverse(key, shift_value=0.0)
+
+ num_bwd, den_bwd = _delta_scan(
+ reverse(query_rot),
+ reverse(key_rot, shift_value=0.0),
+ reverse(value, shift_value=0.0),
+ _flip_and_shift(beta, dim=2, shift_value=0.0),
+ _flip_and_shift(decay, dim=2, shift_value=1.0),
+ spatial_tokens=spatial_tokens,
+ skip_z=skip_z,
+ flip_output=True,
+ **bwd_kwargs,
+ )
+
+ return num_fwd + num_bwd, (None if skip_z else den_fwd + den_bwd)
+
+
+def _reverse_frames(
+ tensor: torch.Tensor,
+ *,
+ frames: int,
+ spatial_tokens: int,
+ shift_value: float | None = None,
+) -> torch.Tensor:
+ """Reverse frame order on a ``[..., T*S]`` tensor, optionally shifting by one frame.
+
+ The frame axis is the outer half of the flat token axis, so unfolding it in
+ place keeps the tensor's own layout: the flip writes a contiguous result and
+ the reshape back is a view. Routing this through a ``[..., T, D, S]``
+ permutation instead would cost a second full materialization.
+ """
+ lead = tensor.shape[:-1]
+ reversed_ = torch.flip(tensor.reshape(*lead, frames, spatial_tokens), dims=[-2])
+ if shift_value is not None:
+ padding = torch.full((*lead, 1, spatial_tokens), shift_value, device=tensor.device, dtype=tensor.dtype)
+ reversed_ = torch.cat([padding, reversed_.narrow(-2, 0, frames - 1)], dim=-2)
+ return reversed_.reshape(*lead, frames * spatial_tokens)
+
+
+def _flip_and_shift(tensor: torch.Tensor, *, dim: int, shift_value: float) -> torch.Tensor:
+ flipped = torch.flip(tensor, dims=[dim])
+ shifted = flipped.narrow(dim, 0, tensor.shape[dim] - 1)
+ pad_shape = list(tensor.shape)
+ pad_shape[dim] = 1
+ padding = torch.full(pad_shape, shift_value, device=tensor.device, dtype=tensor.dtype)
+ return torch.cat([padding, shifted], dim=dim)
+
+
+def reference_bidirectional_gated_delta_net(
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ *,
+ beta: torch.Tensor,
+ decay: torch.Tensor,
+ spatial_tokens: int,
+ query_rot: torch.Tensor | None = None,
+ key_rot: torch.Tensor | None = None,
+ eps: float = 1e-8,
+) -> torch.Tensor:
+ """Run the SANA-WM bidirectional gated delta recurrence in PyTorch.
+
+ Inputs follow the layout used by the official Stage-1 operator:
+ ``query/key/value/query_rot/key_rot`` are ``[B, H, D, T*S]``, ``beta`` is
+ ``[B, H, T, S]``, and ``decay`` is ``[B, H, T]``.
+ """
+
+ _validate_gdn_inputs(query, key, value, beta, decay, spatial_tokens)
+ if query_rot is None:
+ query_rot = query
+ if key_rot is None:
+ key_rot = key
+ if query_rot.shape != query.shape or key_rot.shape != key.shape:
+ raise ValueError("Sana-WM GDN rotary query/key shapes must match query/key.")
+
+ dtype_orig = query.dtype
+ numerator, denominator = _bidirectional_delta_scan(
+ query_rot.float(),
+ key_rot.float(),
+ value.float(),
+ beta.float(),
+ decay.float(),
+ spatial_tokens=spatial_tokens,
+ query=query.float(),
+ key=key.float(),
+ )
+ return (numerator / (denominator + eps)).to(dtype_orig)
+
+
+class TensorParallelRMSNorm(nn.Module):
+ """RMSNorm whose statistics span a tensor-parallel sharded last dimension.
+
+ Models that normalize q/k across *all* heads (``qk_norm="rms_norm_across_heads"``)
+ cannot use a plain RMSNorm once the q/k projections are column-parallel: each
+ rank would then compute the RMS over its local head shard only. This
+ all-reduces the squared sum so the denominator matches the global width,
+ while the affine weight stays local.
+
+ Model-local on purpose: LTX-2 and DreamZero carry their own near-identical
+ copies, and converging all three belongs in its own change rather than in
+ the SANA-WM port.
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ eps: float = 1e-6,
+ elementwise_affine: bool = True,
+ tp_size: int = 1,
+ ) -> None:
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.tp_size = max(int(tp_size), 1)
+ self.global_hidden_size = hidden_size * self.tp_size
+ self.eps = eps
+ if elementwise_affine:
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ else:
+ self.register_parameter("weight", None)
+
+ def _local_sum_sq(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ x_float = x.float()
+ return x_float, x_float.pow(2).sum(dim=-1, keepdim=True)
+
+ def _scale(self, x_float: torch.Tensor, global_sum_sq: torch.Tensor, input_dtype: torch.dtype) -> torch.Tensor:
+ out = x_float * torch.rsqrt(global_sum_sq / self.global_hidden_size + self.eps)
+ if self.weight is not None:
+ out = out * self.weight.float()
+ return out.to(dtype=input_dtype)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x_float, sum_sq = self._local_sum_sq(x)
+ if self.tp_size > 1:
+ sum_sq = tensor_model_parallel_all_reduce(sum_sq)
+ return self._scale(x_float, sum_sq, x.dtype)
+
+
+def fused_qk_rms_norm(
+ norm_q: nn.Module,
+ norm_k: nn.Module,
+ q: torch.Tensor,
+ k: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Apply q/k :class:`TensorParallelRMSNorm` with a SINGLE fused all-reduce.
+
+ Self-attention normalizes q and k every step, and under TP each norm issues
+ its own tiny all-reduce of the per-token sum-of-squares. These collectives
+ are latency-bound, so packing both sums into one tensor halves the count
+ (2 collectives -> 1) for free.
+
+ Numerically identical to ``norm_q(q), norm_k(k)``: all-reduce is
+ elementwise, so packing along the last dim reduces each slice independently
+ with the same fp32 accumulation. Falls back to independent application when
+ either norm is not a :class:`TensorParallelRMSNorm` (e.g. ``nn.Identity``
+ when qk_norm is off, or plain ``RMSNorm`` at TP=1).
+ """
+ if not (isinstance(norm_q, TensorParallelRMSNorm) and isinstance(norm_k, TensorParallelRMSNorm)):
+ return norm_q(q), norm_k(k)
+ if norm_q.tp_size <= 1:
+ return norm_q(q), norm_k(k)
+
+ q_float, q_sum_sq = norm_q._local_sum_sq(q)
+ k_float, k_sum_sq = norm_k._local_sum_sq(k)
+ packed = tensor_model_parallel_all_reduce(torch.cat([q_sum_sq, k_sum_sq], dim=-1))
+ q_sum_sq, k_sum_sq = packed[..., 0:1], packed[..., 1:2]
+ return norm_q._scale(q_float, q_sum_sq, q.dtype), norm_k._scale(k_float, k_sum_sq, k.dtype)
+
+
+def _make_sharded_qk_rms_norm(hidden_size: int, eps: float = 1e-6) -> nn.Module:
+ """Build the q/k norm, TP-aware.
+
+ SANA-WM normalizes q/k across all heads, so once the projections are
+ column-parallel the RMS denominator has to span the TP ranks.
+ """
+ tp_size = get_tensor_model_parallel_world_size()
+ if tp_size > 1:
+ return TensorParallelRMSNorm(hidden_size, eps=eps, tp_size=tp_size)
+ return RMSNorm(hidden_size, eps=eps)
+
+
+class SanaWmTextProjection(nn.Module):
+ def __init__(
+ self,
+ prompt_channels: int,
+ hidden_size: int,
+ *,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.fc1 = ColumnParallelLinear(
+ prompt_channels,
+ hidden_size,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.fc1" if prefix else "fc1",
+ )
+ self.fc2 = RowParallelLinear(
+ hidden_size,
+ hidden_size,
+ bias=True,
+ input_is_parallel=True,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.fc2" if prefix else "fc2",
+ )
+ self.act = nn.SiLU()
+
+ def forward(self, prompt_embeds: torch.Tensor) -> torch.Tensor:
+ return self.fc2(self.act(self.fc1(prompt_embeds)))
+
+
+class SanaWmTextEmbedder(nn.Module):
+ def __init__(
+ self,
+ prompt_channels: int,
+ hidden_size: int,
+ max_length: int,
+ *,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.y_embedding = nn.Parameter(torch.zeros(max_length, prompt_channels))
+ self.y_proj = SanaWmTextProjection(
+ prompt_channels,
+ hidden_size,
+ quant_config=quant_config,
+ prefix=f"{prefix}.y_proj" if prefix else "y_proj",
+ )
+
+ def forward(
+ self,
+ prompt_embeds: torch.Tensor | None,
+ *,
+ batch_size: int,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ if prompt_embeds is None:
+ prompt_embeds = self.y_embedding.unsqueeze(0).expand(batch_size, -1, -1)
+ return self.y_proj(prompt_embeds.to(dtype=dtype))
+
+
+class SanaWmTimestepEmbedder(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_size: int,
+ *,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ linear_1 = ColumnParallelLinear(
+ in_features,
+ hidden_size,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.mlp.0" if prefix else "mlp.0",
+ )
+ linear_2 = RowParallelLinear(
+ hidden_size,
+ hidden_size,
+ bias=True,
+ input_is_parallel=True,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.mlp.2" if prefix else "mlp.2",
+ )
+ self.act = nn.SiLU()
+ self.mlp = nn.ModuleList([linear_1, self.act, linear_2])
+ # diffusers' Timesteps reproduces this model's sinusoid exactly:
+ # downscale_freq_shift=0 makes the exponent denominator half_dim, and
+ # flip_sin_to_cos=True emits [cos, sin] in that order. The MLP stays
+ # model-local because its checkpoint keys are mlp.0/mlp.2 and it is
+ # tensor-parallel.
+ self.timesteps_proj = Timesteps(num_channels=in_features, flip_sin_to_cos=True, downscale_freq_shift=0)
+
+ def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+ emb = self.timesteps_proj(timestep.flatten())
+ hidden_states = self.mlp[0](emb.to(dtype=self.mlp[0].weight.dtype))
+ hidden_states = self.act(hidden_states)
+ return self.mlp[2](hidden_states)
+
+
+class SanaWmPatchEmbedMS3D(nn.Module):
+ """Official-style 3D patch embedder used by SANA-WM Stage-1."""
+
+ def __init__(
+ self,
+ patch_size: tuple[int, int, int],
+ in_channels: int,
+ hidden_size: int,
+ *,
+ kernel_size: tuple[int, int, int] | None = None,
+ bias: bool = True,
+ ) -> None:
+ super().__init__()
+ self.patch_size = _to_3tuple(patch_size)
+ self.kernel_size = _to_3tuple(kernel_size or patch_size)
+ self.proj = nn.Conv3d(
+ in_channels,
+ hidden_size,
+ kernel_size=self.kernel_size,
+ stride=self.patch_size,
+ bias=bias,
+ )
+ self.norm = nn.Identity()
+
+ def project_with_shape(self, latents: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int, int]]:
+ hidden_states = self.norm(self.proj(latents))
+ _, _, frames, height, width = hidden_states.shape
+ return hidden_states.flatten(2).transpose(1, 2), (frames, height, width)
+
+ def forward(self, latents: torch.Tensor) -> torch.Tensor:
+ return self.project_with_shape(latents)[0]
+
+
+class SanaWmWanRotaryPosEmbed(nn.Module):
+ """Wan-style 3D RoPE table used by the official SANA-WM GDN blocks."""
+
+ def __init__(
+ self,
+ attention_head_dim: int,
+ *,
+ max_seq_len: int = 1024,
+ theta: float = 10000.0,
+ ) -> None:
+ super().__init__()
+ self.theta = theta
+ h_dim = w_dim = 2 * (attention_head_dim // 6)
+ t_dim = attention_head_dim - h_dim - w_dim
+ self._split_sizes = (
+ t_dim // 2,
+ h_dim // 2,
+ w_dim // 2,
+ )
+ self.freqs = self._build_freqs(max_seq_len)
+
+ def _build_1d_freq(self, dim: int, positions: torch.Tensor) -> torch.Tensor:
+ if dim <= 0:
+ return torch.empty(positions.shape[0], 0, dtype=torch.complex128)
+ freqs = 1.0 / (
+ self.theta ** (torch.arange(0, dim, 2, dtype=torch.float64, device=positions.device)[: dim // 2] / dim)
+ )
+ phase = torch.outer(positions.to(torch.float64), freqs)
+ return torch.polar(torch.ones_like(phase), phase)
+
+ def _build_freqs(self, max_seq_len: int) -> torch.Tensor:
+ positions = torch.arange(max_seq_len, dtype=torch.float64)
+ dims = (
+ self._split_sizes[0] * 2,
+ self._split_sizes[1] * 2,
+ self._split_sizes[2] * 2,
+ )
+ return torch.cat([self._build_1d_freq(dim, positions) for dim in dims], dim=1)
+
+ def forward(self, spatial_shape: tuple[int, int, int], device: torch.device) -> torch.Tensor:
+ frames, height, width = spatial_shape
+ if max(spatial_shape) > self.freqs.shape[0]:
+ self.freqs = self._build_freqs(max(spatial_shape)).to(self.freqs.device)
+ freqs = self.freqs.to(device=device)
+ freqs_f, freqs_h, freqs_w = freqs.split(self._split_sizes, dim=1)
+ f_dim, h_dim, w_dim = self._split_sizes
+ parts = [
+ freqs_f[:frames].view(frames, 1, 1, f_dim).expand(frames, height, width, f_dim),
+ freqs_h[:height].view(1, height, 1, h_dim).expand(frames, height, width, h_dim),
+ freqs_w[:width].view(1, 1, width, w_dim).expand(frames, height, width, w_dim),
+ ]
+ return torch.cat(parts, dim=-1).reshape(1, 1, frames * height * width, -1)
+
+
+class SanaWmSelfAttention(nn.Module):
+ def __init__(
+ self,
+ config: SanaWmConfig,
+ *,
+ use_gdn: bool = True,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ hidden_size = config.hidden_size
+ self.total_num_heads = max(hidden_size // max(config.linear_head_dim, 1), 1)
+ self.head_dim = hidden_size // self.total_num_heads
+ self.num_heads = self.total_num_heads
+ self.num_kv_heads = self.total_num_heads
+ self.eps = 1e-8
+ self.use_gdn = use_gdn and "GDN" in config.attn_type
+ self.patch_size = _to_3tuple(config.patch_size)
+ # Total camera branch width from the checkpoint. TP ColumnParallel
+ # layers expose a per-rank local slice at runtime; keep both sizes so
+ # layer construction uses the global contract while local conv/norm
+ # tensors match q/k/v outputs.
+ self.total_cam_dim = hidden_size // max(config.cam_attn_compress, 1)
+ self.cam_dim = self.total_cam_dim
+ self.qkv = QKVParallelLinear(
+ hidden_size=hidden_size,
+ head_size=self.head_dim,
+ total_num_heads=self.total_num_heads,
+ bias=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.qkv" if prefix else "qkv",
+ )
+ self.num_heads = self.qkv.num_heads
+ self.num_kv_heads = self.qkv.num_kv_heads
+ self.proj = RowParallelLinear(
+ self.total_num_heads * self.head_dim,
+ hidden_size,
+ bias=True,
+ input_is_parallel=True,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.proj" if prefix else "proj",
+ )
+ self.beta_proj = ColumnParallelLinear(
+ hidden_size,
+ self.total_num_heads,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.beta_proj" if prefix else "beta_proj",
+ )
+ self.gate_proj = ColumnParallelLinear(
+ hidden_size,
+ self.total_num_heads,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.gate_proj" if prefix else "gate_proj",
+ )
+ self.output_gate = ColumnParallelLinear(
+ hidden_size,
+ self.total_num_heads * self.head_dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.output_gate" if prefix else "output_gate",
+ )
+ self.q_proj_cam = ColumnParallelLinear(
+ hidden_size,
+ self.total_cam_dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.q_proj_cam" if prefix else "q_proj_cam",
+ )
+ self.k_proj_cam = ColumnParallelLinear(
+ hidden_size,
+ self.total_cam_dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.k_proj_cam" if prefix else "k_proj_cam",
+ )
+ self.v_proj_cam = ColumnParallelLinear(
+ hidden_size,
+ self.total_cam_dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.v_proj_cam" if prefix else "v_proj_cam",
+ )
+ self.out_proj_cam = RowParallelLinear(
+ self.total_cam_dim,
+ hidden_size,
+ bias=True,
+ input_is_parallel=True,
+ # The camera contribution is summed into the TP-local GDN stream, so
+ # the cross-rank sum and the slice down to the local width collapse
+ # into a single reduce-scatter (see _reduce_scatter_cam_contrib).
+ # RowParallelLinear must therefore not all-reduce, and its bias is
+ # applied after the scatter rather than inside the matmul.
+ reduce_results=False,
+ skip_bias_add=True,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.out_proj_cam" if prefix else "out_proj_cam",
+ )
+ local_inner_dim = self.num_heads * self.head_dim
+ norm_cls = _make_sharded_qk_rms_norm if config.qk_norm else (lambda *_args, **_kwargs: nn.Identity())
+ self.q_norm = norm_cls(local_inner_dim)
+ self.k_norm = norm_cls(local_inner_dim)
+
+ # These names mirror the official GDN checkpoint.
+ self.A_log = nn.Parameter(torch.zeros(self.num_heads))
+ self.dt_bias = nn.Parameter(torch.zeros(self.num_heads))
+ self.register_buffer("recall_gate", torch.zeros(1))
+ self.conv_k = (
+ nn.Conv1d(
+ local_inner_dim,
+ local_inner_dim,
+ kernel_size=config.conv_kernel_size,
+ groups=local_inner_dim,
+ bias=False,
+ )
+ if self.use_gdn and config.conv_kernel_size > 0
+ else None
+ )
+ self.conv_q = None
+ self.conv_v = None
+
+ self.cam_dim = int(getattr(self.q_proj_cam, "output_size_per_partition", self.total_cam_dim))
+
+ cam_compress = max(config.cam_attn_compress, 1)
+ self.cam_heads = max(self.num_heads // cam_compress, 1)
+ if self.cam_dim % self.cam_heads != 0:
+ raise ValueError(f"Sana-WM local cam_dim={self.cam_dim} must be divisible by cam_heads={self.cam_heads}.")
+ self.cam_head_dim = self.cam_dim // self.cam_heads
+ # Under TP, cam_head_dim must equal main head_dim so the WAN RoPE
+ # (built for main head_dim) can be correctly sliced for the cam branch.
+ # A mismatch indicates the cam ColumnParallelLinear layers were not
+ # partitioned consistently with the main QKVParallelLinear.
+ if self.cam_head_dim != self.head_dim:
+ raise ValueError(
+ f"Sana-WM cam_head_dim={self.cam_head_dim} != head_dim={self.head_dim}. "
+ f"Under TP the cam layers must be partitioned by the same TP degree as the "
+ f"main QKV (cam_dim={self.cam_dim}, cam_heads={self.cam_heads}, "
+ f"total_cam_dim={self.total_cam_dim})."
+ )
+ self.q_norm_cam = norm_cls(self.cam_dim)
+ self.k_norm_cam = norm_cls(self.cam_dim)
+ self.conv_k_cam = (
+ nn.Conv1d(
+ self.cam_dim,
+ self.cam_dim,
+ kernel_size=config.conv_kernel_size,
+ groups=self.cam_dim,
+ bias=False,
+ )
+ if self.use_gdn and config.conv_kernel_size > 0
+ else None
+ )
+ self.conv_q_cam = None
+ self.conv_v_cam = None
+ # Softmax hybrid blocks (every ``softmax_every_n``-th) and the camera
+ # branch both run plain non-causal self-attention, so both go through
+ # the shared Attention for platform backend selection and the fp32
+ # fallback. The camera branch gets its own module because
+ # ``cam_attn_compress`` may give it fewer heads.
+ self.softmax_attn = Attention(
+ num_heads=self.num_heads,
+ head_size=self.head_dim,
+ num_kv_heads=self.num_kv_heads,
+ softmax_scale=self.head_dim**-0.5,
+ causal=False,
+ role="self",
+ qkv_layout="BSND",
+ skip_sequence_parallel=True,
+ prefix=prefix,
+ )
+ self.softmax_attn_cam = Attention(
+ num_heads=self.cam_heads,
+ head_size=self.cam_head_dim,
+ num_kv_heads=self.cam_heads,
+ softmax_scale=self.cam_head_dim**-0.5,
+ causal=False,
+ role="self",
+ qkv_layout="BSND",
+ skip_sequence_parallel=True,
+ prefix=prefix,
+ )
+ self._init_short_convs()
+ self._mark_tp_sharded_params()
+
+ def _mark_tp_sharded_params(self) -> None:
+ """Declare which plain parameters are TP-local (see ``_shard_param_across_tp``)."""
+ # Per-head GDN vectors: one entry per local head.
+ _shard_param_across_tp(self.A_log)
+ _shard_param_across_tp(self.dt_bias)
+ # Depthwise temporal convs and q/k norms: one entry per local channel.
+ for module in (self.conv_q, self.conv_k, self.conv_v, self.conv_q_cam, self.conv_k_cam, self.conv_v_cam):
+ if module is not None:
+ _shard_param_across_tp(module.weight)
+ for norm in (self.q_norm, self.k_norm, self.q_norm_cam, self.k_norm_cam):
+ _shard_param_across_tp(getattr(norm, "weight", None))
+
+ @staticmethod
+ def _init_short_conv(conv: nn.Conv1d | None) -> None:
+ if conv is None:
+ return
+ with torch.no_grad():
+ conv.weight.zero_()
+ conv.weight[:, 0, -1] = 1.0
+
+ def _init_short_convs(self) -> None:
+ for conv in (self.conv_q, self.conv_k, self.conv_v, self.conv_q_cam, self.conv_k_cam, self.conv_v_cam):
+ self._init_short_conv(conv)
+
+ @staticmethod
+ def _apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ rotated = torch.view_as_complex(hidden_states.permute(0, 1, 3, 2).to(torch.float64).unflatten(3, (-1, 2)))
+ output = torch.view_as_real(rotated * freqs).flatten(3, 4).permute(0, 1, 3, 2)
+ return output.type_as(hidden_states)
+
+ @staticmethod
+ def _reshape_to_temporal(
+ hidden_states: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ ) -> tuple[torch.Tensor, int, int, int]:
+ batch_size, token_count, hidden_size = hidden_states.shape
+ frames, height, width = spatial_shape
+ spatial_tokens = height * width
+ if token_count != frames * spatial_tokens:
+ raise ValueError(f"Sana-WM temporal conv expects N=T*H*W, got N={token_count}, THW={spatial_shape}.")
+ hidden_states = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size)
+ hidden_states = hidden_states.permute(0, 2, 1, 3).reshape(batch_size * spatial_tokens, frames, hidden_size)
+ return hidden_states, batch_size, spatial_tokens, frames
+
+ @staticmethod
+ def _reshape_from_temporal(
+ hidden_states: torch.Tensor,
+ batch_size: int,
+ spatial_tokens: int,
+ frames: int,
+ ) -> torch.Tensor:
+ hidden_size = hidden_states.shape[-1]
+ return (
+ hidden_states.reshape(batch_size, spatial_tokens, frames, hidden_size)
+ .permute(0, 2, 1, 3)
+ .reshape(batch_size, frames * spatial_tokens, hidden_size)
+ )
+
+ @staticmethod
+ def _causal_conv_1d(hidden_states: torch.Tensor, conv: nn.Conv1d) -> torch.Tensor:
+ dtype = hidden_states.dtype
+ conv_input = hidden_states.transpose(1, 2).to(conv.weight.dtype)
+ conv_input = F.pad(conv_input, (conv.kernel_size[0] - 1, 0))
+ output = conv(conv_input).transpose(1, 2)
+ return output.to(dtype)
+
+ def _bidirectional_temporal_short_conv(
+ self,
+ hidden_states: torch.Tensor,
+ conv: nn.Conv1d,
+ spatial_shape: tuple[int, int, int],
+ ) -> torch.Tensor:
+ hidden_states, batch_size, spatial_tokens, frames = self._reshape_to_temporal(hidden_states, spatial_shape)
+ forward_states = self._causal_conv_1d(hidden_states, conv)
+ backward_states = self._causal_conv_1d(hidden_states.flip(1), conv).flip(1)
+ center_weight = conv.weight[:, 0, -1].to(hidden_states.dtype)
+ center_states = hidden_states * center_weight.view(1, 1, -1)
+ output = forward_states + backward_states - center_states
+ return self._reshape_from_temporal(output, batch_size, spatial_tokens, frames)
+
+ def _compute_frame_gates(
+ self,
+ hidden_states: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ batch_size, token_count, hidden_size = hidden_states.shape
+ frames, height, width = spatial_shape
+ spatial_tokens = height * width
+ if token_count != frames * spatial_tokens:
+ raise ValueError(f"Sana-WM GDN token layout mismatch: N={token_count}, expected {frames * spatial_tokens}.")
+ beta = torch.sigmoid(self.beta_proj(hidden_states))
+ beta = beta.reshape(batch_size, frames, spatial_tokens, self.num_heads).permute(0, 3, 1, 2)
+ frame_states = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size).mean(dim=2)
+ gate = self.gate_proj(frame_states).float()
+ decay = torch.exp(
+ -self.A_log.float().exp().view(1, 1, -1) * F.softplus(gate + self.dt_bias.float().view(1, 1, -1))
+ )
+ return beta, decay.transpose(1, 2)
+
+ def _forward_gdn_raw(
+ self,
+ hidden_states: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ rotary_emb: torch.Tensor | None,
+ *,
+ precomputed_gates: tuple[torch.Tensor, torch.Tensor] | None = None,
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+ """Run main-branch bidirectional GDN and return the (B, N, q_size)
+ raw output BEFORE ``output_gate`` and ``proj`` are applied.
+
+ Returns ``(raw_output, (beta, decay))`` so the caller can share the
+ beta/decay gates with the camera branch (matches NVlabs' shared
+ ``precomputed_gates`` plumbing).
+ """
+ batch_size, token_count, _hidden_size = hidden_states.shape
+ frames, height, width = spatial_shape
+ spatial_tokens = height * width
+ if token_count != frames * spatial_tokens:
+ raise ValueError(f"Sana-WM GDN expects N=T*H*W, got N={token_count}, THW={spatial_shape}.")
+
+ qkv = self.qkv(hidden_states)
+ q_size = self.num_heads * self.head_dim
+ kv_size = self.num_kv_heads * self.head_dim
+ query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
+ if self.conv_k is not None:
+ key = self._bidirectional_temporal_short_conv(key, self.conv_k, spatial_shape)
+ if precomputed_gates is None:
+ beta, decay = self._compute_frame_gates(hidden_states, spatial_shape)
+ else:
+ beta, decay = precomputed_gates
+
+ query, key = fused_qk_rms_norm(self.q_norm, self.k_norm, query, key)
+ query = query.reshape(batch_size, token_count, self.num_heads, self.head_dim)
+ key = key.reshape(batch_size, token_count, self.num_heads, self.head_dim)
+ value = value.reshape(batch_size, token_count, self.num_heads, self.head_dim)
+
+ query = F.relu(query).permute(0, 2, 3, 1)
+ key = F.relu(key).permute(0, 2, 3, 1)
+ value = value.permute(0, 2, 3, 1)
+ key = key * ((self.head_dim**-0.5) * (spatial_tokens**-0.5))
+ if rotary_emb is not None:
+ query_rot = self._apply_rotary_emb(query, rotary_emb)
+ key_rot = self._apply_rotary_emb(key, rotary_emb)
+ else:
+ query_rot = query
+ key_rot = key
+
+ output = reference_bidirectional_gated_delta_net(
+ query,
+ key,
+ value,
+ beta=beta,
+ decay=decay,
+ spatial_tokens=spatial_tokens,
+ query_rot=query_rot,
+ key_rot=key_rot,
+ eps=self.eps,
+ )
+ output = output.permute(0, 3, 1, 2).reshape(batch_size, token_count, q_size)
+ return output, (beta, decay)
+
+ def _apply_output_gate_and_proj(
+ self,
+ combined: torch.Tensor,
+ hidden_states: torch.Tensor,
+ ) -> torch.Tensor:
+ """Apply SiLU output gate (driven by hidden_states) + final projection.
+
+ The gate reads ``hidden_states``, not ``combined``, and NVlabs applies
+ it once to the main+camera sum rather than per branch.
+ """
+ gate = F.silu(self.output_gate(hidden_states).float())
+ gated = combined * gate
+ return self.proj(gated.to(self.proj.weight.dtype))
+
+ def _reduce_scatter_cam_contrib(self, contrib: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
+ """Reduce the partial camera projection straight down to the local shard.
+
+ ``out_proj_cam`` is row-parallel over the camera features, so each rank
+ produces a partial sum spanning the full hidden width. The GDN stream it
+ is added to is TP-local (``output_gate`` and ``proj`` both expect the
+ local width), so an all-reduce would hand every rank the full result
+ only for all but one slice of it to be discarded. Reduce-scatter does
+ both steps in one collective and moves 1/tp of the bytes.
+
+ ``skip_bias_add`` keeps the bias out of the matmul so it can be applied
+ to the scattered shard; that slice is the only remaining use of the TP
+ rank here, and it is a 1-D view rather than a copy of the hidden tensor.
+
+ ``tensor_model_parallel_reduce_scatter`` is not usable here: it calls
+ ``get_tp_group().reduce_scatter()``, and vLLM-Omni's diffusion
+ ``GroupCoordinator`` deliberately does not implement that method (see
+ the note in ``diffusion/distributed/parallel_state.py``). Drive
+ ``torch.distributed`` on the group's process group instead, and fall
+ back to all-reduce plus a slice if the group is not exposed.
+ """
+ bias = self.out_proj_cam.bias
+ tp_size = get_tensor_model_parallel_world_size()
+ if tp_size == 1:
+ return contrib if bias is None else contrib + bias
+
+ local_width = reference.shape[-1]
+ if local_width <= 0 or contrib.shape[-1] != local_width * tp_size:
+ raise ValueError(
+ "Sana-WM TP camera contribution width mismatch: "
+ f"contrib={contrib.shape[-1]} expected={local_width * tp_size} (local={local_width}, tp={tp_size})."
+ )
+
+ device_group = getattr(get_tp_group(), "device_group", None)
+ if device_group is None:
+ out = tensor_model_parallel_all_reduce(contrib.contiguous())
+ start = get_tensor_model_parallel_rank() * local_width
+ out = out[..., start : start + local_width].contiguous()
+ else:
+ # reduce_scatter_tensor splits along dim 0, so move the rank axis to
+ # the front first. Rank r then receives the summed columns
+ # [r * local_width, (r + 1) * local_width), matching what the
+ # all-reduce-and-slice fallback produces.
+ batch, tokens, _ = contrib.shape
+ source = (
+ contrib.reshape(batch, tokens, tp_size, local_width)
+ .permute(2, 0, 1, 3)
+ .contiguous()
+ .reshape(tp_size * batch, tokens, local_width)
+ )
+ out = torch.empty((batch, tokens, local_width), device=contrib.device, dtype=contrib.dtype)
+ torch.distributed.reduce_scatter_tensor(out, source, group=device_group)
+
+ if bias is not None:
+ start = get_tensor_model_parallel_rank() * local_width
+ out = out + bias[start : start + local_width]
+ return out
+
+ def _forward_gdn(
+ self,
+ hidden_states: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ rotary_emb: torch.Tensor | None,
+ ) -> torch.Tensor:
+ """Main-branch only GDN forward with output_gate + proj applied."""
+ raw, _ = self._forward_gdn_raw(hidden_states, spatial_shape, rotary_emb)
+ return self._apply_output_gate_and_proj(raw, hidden_states)
+
+ @staticmethod
+ def _downscale_to_reference_rms(
+ ref: torch.Tensor,
+ transformed: torch.Tensor,
+ eps: float = 1e-6,
+ ) -> torch.Tensor:
+ """Downscale ``transformed`` so its per-token channel RMS does not
+ exceed ``ref``'s. Mirrors NVlabs ``_downscale_to_reference_rms``.
+
+ Inputs are ``(B, H, N, D)`` (channel dim last). The RMS-scale is
+ clamped at 1.0 — the function only ever shrinks, never amplifies.
+ """
+ ref_rms = ref.square().mean(dim=-1, keepdim=True).add(eps).sqrt()
+ tr_rms = transformed.square().mean(dim=-1, keepdim=True).add(eps).sqrt()
+ scale = (ref_rms / tr_rms.clamp_min(eps)).clamp(max=1.0)
+ return transformed * scale
+
+ @staticmethod
+ def _ucpe_rotary_freqs(rotary_emb: torch.Tensor | None) -> torch.Tensor | None:
+ # SanaWmRope emits (1, 1, N, D//2) complex; the cam path wants (N, D//2).
+ if rotary_emb is None:
+ return None
+ return rotary_emb.squeeze(0).squeeze(0)
+
+ def _forward_softmax_raw(
+ self,
+ hidden_states: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ rotary_emb: torch.Tensor | None,
+ ) -> torch.Tensor:
+ """Softmax self-attention raw output before output_gate/proj.
+
+ Mirrors NVlabs ``_forward_softmax_attn(..., apply_output_gate=False)``
+ for every-N-th hybrid blocks. GDN-specific gates are intentionally not
+ used here.
+ """
+ batch_size, token_count, hidden_size = hidden_states.shape
+ frames, height, width = spatial_shape
+ if token_count != frames * height * width:
+ raise ValueError(f"Sana-WM softmax attention expects N=T*H*W, got N={token_count}, THW={spatial_shape}.")
+
+ qkv = self.qkv(hidden_states)
+ q_size = self.num_heads * self.head_dim
+ kv_size = self.num_kv_heads * self.head_dim
+ query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
+
+ query, key = fused_qk_rms_norm(self.q_norm, self.k_norm, query, key)
+ query = query.reshape(batch_size, token_count, self.num_heads, self.head_dim)
+ key = key.reshape(batch_size, token_count, self.num_kv_heads, self.head_dim)
+ value = value.reshape(batch_size, token_count, self.num_kv_heads, self.head_dim)
+
+ if rotary_emb is not None:
+ query = self._apply_rotary_emb(query.permute(0, 2, 3, 1), rotary_emb).permute(0, 3, 1, 2)
+ key = self._apply_rotary_emb(key.permute(0, 2, 3, 1), rotary_emb).permute(0, 3, 1, 2)
+
+ # NVlabs runs the attention in bf16 even under ``fp32_attention=True``,
+ # since FlashAttention only supports bf16/fp16 and fp32 falls back to
+ # the math backend. Keep Q/K/V at the module dtype so an upstream step
+ # that promoted to fp32 does not silently change the kernel. See
+ # ``sana_gdn_camctrl_blocks.py::_forward_softmax_attn_sdpa``.
+ dtype_orig = hidden_states.dtype
+ attn = self.softmax_attn(query.to(dtype_orig), key.to(dtype_orig), value.to(dtype_orig))
+ return attn.reshape(batch_size, token_count, q_size)
+
+ def _forward_softmax_cam_branch(
+ self,
+ hidden_states: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ cam_geometry: SanaWmCamGeometry,
+ rotary_emb: torch.Tensor | None,
+ ) -> torch.Tensor:
+ """UCPE camera branch for softmax hybrid blocks.
+
+ Matches NVlabs ``_forward_cam_branch_softmax``: camera Q/K/V are
+ projected from the same hidden stream, transformed by UCPE, run through
+ SDPA, inverse-transformed, then returned as raw ``cam_dim`` features.
+ """
+ batch_size, token_count, _ = hidden_states.shape
+ frames, height, width = spatial_shape
+ if token_count != frames * height * width:
+ raise ValueError(f"Sana-WM softmax cam branch expects N=T*H*W, got N={token_count}, THW={spatial_shape}.")
+
+ q_cam = self.q_proj_cam(hidden_states)
+ k_cam = self.k_proj_cam(hidden_states)
+ v_cam = self.v_proj_cam(hidden_states)
+
+ if self.conv_q_cam is not None:
+ q_cam = self._bidirectional_temporal_short_conv(q_cam, self.conv_q_cam, spatial_shape)
+ if self.conv_k_cam is not None:
+ k_cam = self._bidirectional_temporal_short_conv(k_cam, self.conv_k_cam, spatial_shape)
+ if self.conv_v_cam is not None:
+ v_cam = self._bidirectional_temporal_short_conv(v_cam, self.conv_v_cam, spatial_shape)
+
+ q_cam, k_cam = fused_qk_rms_norm(self.q_norm_cam, self.k_norm_cam, q_cam, k_cam)
+ q_cam = q_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim)
+ k_cam = k_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim)
+ v_cam = v_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim)
+
+ q_cam = q_cam.transpose(1, 2).contiguous() # (B, H_cam, N, D)
+ k_cam = k_cam.transpose(1, 2).contiguous()
+ v_cam = v_cam.transpose(1, 2).contiguous()
+
+ q_cam_trans = cam_geometry.apply_q(q_cam)
+ kv_cam_trans = cam_geometry.apply_kv(torch.cat([k_cam, v_cam], dim=1))
+ k_cam_trans, v_cam_trans = torch.chunk(kv_cam_trans, chunks=2, dim=1)
+
+ q_cam_trans = self._downscale_to_reference_rms(q_cam, q_cam_trans)
+ k_cam_trans = self._downscale_to_reference_rms(k_cam, k_cam_trans)
+ v_cam_trans = self._downscale_to_reference_rms(v_cam, v_cam_trans)
+
+ # The shared Attention takes BSND and selects a backend that supports
+ # ``cam_head_dim`` directly, so odd head dims need no padding up to an
+ # SDPA-friendly size. Q/K/V stay at the module dtype for the same
+ # reason as the main softmax path above.
+ dtype_orig = hidden_states.dtype
+ out = self.softmax_attn_cam(
+ q_cam_trans.transpose(1, 2).to(dtype_orig),
+ k_cam_trans.transpose(1, 2).to(dtype_orig),
+ v_cam_trans.transpose(1, 2).to(dtype_orig),
+ )
+ out = cam_geometry.apply_output(out.transpose(1, 2))
+ return out.transpose(1, 2).reshape(batch_size, token_count, self.cam_dim)
+
+ def _forward_cam_branch(
+ self,
+ hidden_states: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ cam_geometry: SanaWmCamGeometry,
+ rotary_emb: torch.Tensor | None,
+ *,
+ precomputed_gates: tuple[torch.Tensor, torch.Tensor],
+ ) -> torch.Tensor:
+ """UCPE camera branch — matches the NVlabs production
+ ``BidirectionalGDNUCPESinglePathLiteLABothTriton`` variant declared
+ in the SANA-WM 1600M release config.
+
+ Two things differ from the main GDN branch and are easy to get wrong:
+ Q/K/V all come from ``hidden_states`` (UCPE is not cross-attention —
+ camera information enters through the per-pixel projection matrices),
+ and the recurrence is single-path, i.e. numerator only, with no Z
+ denominator and no final divide.
+
+ Returns ``(B, N, cam_dim)`` raw — caller applies ``out_proj_cam`` and
+ the shared ``output_gate`` + ``proj``.
+ """
+ batch_size, token_count, _ = hidden_states.shape
+ frames, height, width = spatial_shape
+ spatial_tokens = height * width
+ if token_count != frames * spatial_tokens:
+ raise ValueError(f"Sana-WM cam branch expects N=T*H*W, got N={token_count}, THW={spatial_shape}.")
+
+ q_cam = self.q_proj_cam(hidden_states)
+ k_cam = self.k_proj_cam(hidden_states)
+ v_cam = self.v_proj_cam(hidden_states)
+
+ if self.conv_k_cam is not None:
+ k_cam = self._bidirectional_temporal_short_conv(k_cam, self.conv_k_cam, spatial_shape)
+ if self.conv_q_cam is not None:
+ q_cam = self._bidirectional_temporal_short_conv(q_cam, self.conv_q_cam, spatial_shape)
+ if self.conv_v_cam is not None:
+ v_cam = self._bidirectional_temporal_short_conv(v_cam, self.conv_v_cam, spatial_shape)
+
+ # Normalise through the model's own norm modules: under TP these are
+ # TensorParallelRMSNorm, which all-reduces the squared sum so the
+ # denominator spans the global channel width. Doing the RMS inside
+ # cam_prep_func would normalise over this rank's head shard only. q and
+ # k share a shape here, so the two reductions fuse into one.
+ q_cam, k_cam = fused_qk_rms_norm(self.q_norm_cam, self.k_norm_cam, q_cam, k_cam)
+ q_normed = q_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim)
+ k_normed = k_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim)
+
+ # cam_prep_func returns NVlabs' ``(B, H, D, N)`` scan layout directly,
+ # so nothing needs permuting before the recurrence below.
+ v_raw = v_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim).contiguous()
+ k_scale = (self.cam_head_dim**-0.5) * (spatial_tokens**-0.5)
+ q_rot_bhdn, k_rot_bhdn, v_bhdn, inflation_sq = cam_prep_func(
+ q_normed.contiguous(),
+ k_normed.contiguous(),
+ v_raw,
+ proj_q=cam_geometry.proj_q,
+ proj_kv=cam_geometry.proj_kv,
+ rope_cos=cam_geometry.rope_cos,
+ rope_sin=cam_geometry.rope_sin,
+ k_scale=k_scale,
+ )
+
+ # Discount beta by the raw K inflation instead of shrinking Q/K/V first:
+ # the released NVlabs ``BothTriton`` cam path skips the Python
+ # ``PostUCPERenorm`` step and feeds the scan the raw UCPE-transformed
+ # tensors.
+ beta, decay = precomputed_gates
+ inflation_per_token = inflation_sq.reshape(batch_size, self.cam_heads, frames, spatial_tokens)
+ frame_inflation_sq = inflation_per_token.mean(dim=-1) # (B, H_cam, T)
+ # β/decay come from the main branch, so the cam branch can only reuse
+ # them head-for-head. cam_attn_compress > 1 would give the cam branch
+ # fewer heads and needs a grouping rule this port does not implement;
+ # __init__ already rejects it via the cam_head_dim check.
+ if beta.shape[1] != self.cam_heads:
+ raise ValueError(
+ "Sana-WM cam branch requires cam_heads == main heads (cam_attn_compress=1),"
+ f" got cam_heads={self.cam_heads} main_heads={beta.shape[1]}."
+ )
+ beta_cam = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0)
+ decay_cam = decay
+
+ cam_out_bhdn, _ = _bidirectional_delta_scan(
+ q_rot_bhdn.float(),
+ k_rot_bhdn.float(),
+ v_bhdn.float(),
+ beta_cam.float(),
+ decay_cam.float(),
+ spatial_tokens=spatial_tokens,
+ skip_z=True,
+ )
+ cam_out_bhdn = cam_out_bhdn.to(q_rot_bhdn.dtype)
+
+ # ``apply_output`` expects (B, H, N, D); the scan emits (B, H, D, N).
+ cam_out_bhnd = cam_out_bhdn.transpose(-1, -2).contiguous()
+ cam_out_bhnd = cam_geometry.apply_output(cam_out_bhnd)
+ cam_out_bhdn = cam_out_bhnd.transpose(-1, -2).contiguous()
+
+ # (B, H_cam, D_cam, N) → (B, N, cam_dim)
+ return cam_out_bhdn.permute(0, 3, 1, 2).reshape(batch_size, token_count, self.cam_dim)
+
+ # The GDN/UCPE branch runs a sequential recurrence whose trip count is the
+ # latent frame count, i.e. data-dependent at trace time. Run the whole
+ # attention eagerly (graph break) rather than letting torch.compile unroll
+ # it per shape; the rest of the block (MLP / norms / cross-attention) still
+ # compiles.
+ @torch.compiler.disable
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ rotary_emb: torch.Tensor | None = None,
+ cam_geometry: SanaWmCamGeometry | None = None,
+ ) -> torch.Tensor:
+ """Forward with dual-branch GDN+UCPE when ``cam_geometry`` is
+ provided, otherwise main-branch-only GDN or softmax fallback.
+
+ ``cam_geometry`` holds the per-pixel 4x4 transforms built once per
+ transformer forward from the raw ``(B, T_latent, 20)`` raymap. The cam
+ branch consumes ``hidden_states`` as its sole Q/K/V source and applies
+ those transforms to the Q/K/V channels via UCPE.
+ """
+ if self.use_gdn:
+ if cam_geometry is None:
+ # No camera conditioning on this request — main branch only.
+ return self._forward_gdn(hidden_states, spatial_shape, rotary_emb)
+ # β/decay are computed here rather than inside each branch: NVlabs
+ # drives both recurrences from one set of gates, and the camera
+ # branch then discounts beta by its own K inflation.
+ beta, decay = self._compute_frame_gates(hidden_states, spatial_shape)
+ main_raw, _ = self._forward_gdn_raw(
+ hidden_states,
+ spatial_shape,
+ rotary_emb,
+ precomputed_gates=(beta, decay),
+ )
+ cam_raw = self._forward_cam_branch(
+ hidden_states,
+ spatial_shape,
+ cam_geometry,
+ rotary_emb,
+ precomputed_gates=(beta, decay),
+ )
+ cam_contrib = self.out_proj_cam(cam_raw)
+ cam_contrib = self._reduce_scatter_cam_contrib(cam_contrib, main_raw)
+ combined = main_raw + cam_contrib.to(main_raw.dtype)
+ attn_out = self._apply_output_gate_and_proj(combined, hidden_states)
+ return attn_out
+
+ main_raw = self._forward_softmax_raw(hidden_states, spatial_shape, rotary_emb)
+ if cam_geometry is not None:
+ cam_raw = self._forward_softmax_cam_branch(
+ hidden_states,
+ spatial_shape,
+ cam_geometry,
+ rotary_emb,
+ )
+ cam_contrib = self.out_proj_cam(cam_raw)
+ cam_contrib = self._reduce_scatter_cam_contrib(cam_contrib, main_raw)
+ main_raw = main_raw + cam_contrib.to(main_raw.dtype)
+ return self._apply_output_gate_and_proj(main_raw, hidden_states)
+
+
+class SanaWmCrossAttention(nn.Module):
+ def __init__(
+ self,
+ config: SanaWmConfig,
+ *,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ hidden_size = config.hidden_size
+ self.total_num_heads = max(hidden_size // max(config.linear_head_dim, 1), 1)
+ self.head_dim = hidden_size // self.total_num_heads
+ self.num_heads = self.total_num_heads
+ inner = self.total_num_heads * self.head_dim
+ self.q_linear = ColumnParallelLinear(
+ hidden_size,
+ inner,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.q_linear" if prefix else "q_linear",
+ )
+ # The checkpoint stores K and V fused as one (2 * inner, hidden)
+ # tensor. A plain ColumnParallelLinear would shard that by the
+ # concatenated output dim, handing all of K to one rank and all of V
+ # to another; MergedColumnParallelLinear knows the block structure
+ # and shards K and V by heads independently, so `chunk(2)` below
+ # yields this rank's K and V directly - no gather, no manual slice.
+ self.kv_linear = MergedColumnParallelLinear(
+ hidden_size,
+ [inner, inner],
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.kv_linear" if prefix else "kv_linear",
+ )
+ self.num_heads = self.q_linear.output_size_per_partition // self.head_dim
+ self.proj = RowParallelLinear(
+ inner,
+ hidden_size,
+ bias=True,
+ input_is_parallel=True,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.proj" if prefix else "proj",
+ )
+ local_inner = self.num_heads * self.head_dim
+ norm_cls = _make_sharded_qk_rms_norm if config.cross_norm else (lambda *_args, **_kwargs: nn.Identity())
+ self.q_norm = norm_cls(local_inner)
+ self.k_norm = norm_cls(local_inner)
+ # q is column-parallel and k is sliced to the local width before k_norm,
+ # so both norms hold one weight entry per local channel.
+ _shard_param_across_tp(getattr(self.q_norm, "weight", None))
+ _shard_param_across_tp(getattr(self.k_norm, "weight", None))
+ self.softmax_attn = Attention(
+ num_heads=self.num_heads,
+ head_size=self.head_dim,
+ num_kv_heads=self.num_heads,
+ softmax_scale=1.0 / (self.head_dim**0.5),
+ causal=False,
+ role="cross",
+ qkv_layout="BSND",
+ skip_sequence_parallel=True,
+ disable_kv_quant=True,
+ prefix=prefix,
+ )
+
+ def _reshape_to_seq_heads(self, tensor: torch.Tensor) -> torch.Tensor:
+ batch, seq_len, hidden_size = tensor.shape
+ return tensor.reshape(batch, seq_len, self.num_heads, hidden_size // self.num_heads)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ query = self.q_norm(self.q_linear(hidden_states))
+ # MergedColumnParallelLinear emits [K_local; V_local], so the chunk is
+ # already this rank's head shard.
+ key, value = self.kv_linear(encoder_hidden_states).chunk(2, dim=-1)
+ key = self.k_norm(key)
+ if attention_mask is not None and attention_mask.ndim != 2:
+ raise ValueError(
+ f"Sana-WM cross-attention mask must be shaped [B, text_len], got {tuple(attention_mask.shape)}."
+ )
+ query = self._reshape_to_seq_heads(query)
+ key = self._reshape_to_seq_heads(key)
+ value = self._reshape_to_seq_heads(value)
+ # The prompt padding mask goes to the shared Attention through
+ # ``AttentionMetadata``: backends take the [B, text_len] form directly
+ # and convert/reshape it themselves (``_maybe_reshape_attn_mask``), so
+ # cross-attention keeps backend selection. ``None`` keeps the mask-free
+ # fast path - callers normalise an all-ones mask to ``None`` rather than
+ # paying for a no-op mask.
+ attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
+ attn = self.softmax_attn(query, key, value, attn_metadata)
+ return self.proj(attn.flatten(2, 3))
+
+
+class _ConvWrapper(nn.Module):
+ def __init__(self, conv: nn.Module) -> None:
+ super().__init__()
+ self.conv = conv
+
+
+class SanaWmMbConvFfn(nn.Module):
+ def __init__(self, config: SanaWmConfig) -> None:
+ super().__init__()
+ hidden_size = config.hidden_size
+ expanded = int(hidden_size * config.mlp_ratio) * 2
+ t_padding = config.t_kernel_size // 2
+ self.glu_act = nn.SiLU()
+ self.inverted_conv = _ConvWrapper(nn.Conv2d(hidden_size, expanded, kernel_size=1))
+ self.depth_conv = _ConvWrapper(nn.Conv2d(expanded, expanded, kernel_size=3, padding=1, groups=expanded))
+ self.point_conv = _ConvWrapper(nn.Conv2d(expanded // 2, hidden_size, kernel_size=1, bias=False))
+ self.t_conv = nn.Conv2d(
+ hidden_size,
+ hidden_size,
+ kernel_size=(config.t_kernel_size, 1),
+ padding=(t_padding, 0),
+ bias=False,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, spatial_shape: tuple[int, int, int]) -> torch.Tensor:
+ """Match NVlabs ``GLUMBConvTemp.forward`` exactly:
+
+ 1. Reshape ``(B, N=F*H*W, C) → (B*F, H, W, C) → (B*F, C, H, W)`` so the
+ spatial MBConv runs PER FRAME with proper 2D (H, W) geometry — the
+ 3×3 ``depth_conv`` must not span an (F, H*W) plane.
+ 2. Apply expand → depthwise 3×3 → GLU → contract on per-frame
+ (B*F, C, H, W).
+ 3. Temporal aggregation: reshape back to ``(B, C, F, H*W)``
+ and add ``t_conv`` (kernel ``(3, 1)``) along the F axis.
+ 4. Reshape to ``(B, N, C)``.
+ """
+ batch, _, hidden_size = hidden_states.shape
+ frames, height, width = spatial_shape
+ spatial_tokens = height * width
+ # (B, F*H*W, C) → (B*F, H, W, C) → (B*F, C, H, W)
+ x = (
+ hidden_states.reshape(batch * frames, spatial_tokens, hidden_size)
+ .reshape(batch * frames, height, width, hidden_size)
+ .permute(0, 3, 1, 2)
+ )
+ # Keep this NHWC-derived NCHW view non-contiguous like NVlabs: on bf16
+ # CUDA convs the stride selects a different kernel, and a contiguous
+ # copy breaks late-step parity.
+ # NVlabs ``ConvLayer`` is `Conv → norm → act` with
+ # `act=("silu", "silu", None)`, so inverted_conv activates and the other
+ # two do not. ``_ConvWrapper`` stores only the raw Conv2d (to match the
+ # `inverted_conv.conv.weight` checkpoint key), hence the explicit SiLU.
+ x_inv = self.glu_act(self.inverted_conv.conv(x))
+ x_dep = self.depth_conv.conv(x_inv)
+ value, gate = x_dep.chunk(2, dim=1)
+ x_glu = value * self.glu_act(gate)
+ x_pt = self.point_conv.conv(x_glu) # (B*F, hidden, H, W)
+ x = x_pt
+ # Temporal aggregation: (B*F, C, H, W) → (B, F, C, H*W) → (B, C, F, H*W)
+ x_temporal = x.reshape(batch, frames, hidden_size, spatial_tokens).permute(0, 2, 1, 3)
+ t_conv_out = self.t_conv(x_temporal)
+ x_temporal = x_temporal + t_conv_out
+ # → (B, F, H*W, C) → (B, N, C)
+ final = x_temporal.permute(0, 2, 3, 1).reshape(batch, frames * spatial_tokens, hidden_size).contiguous()
+ return final
+
+
+class SanaWmBlock(nn.Module):
+ def __init__(
+ self,
+ config: SanaWmConfig,
+ *,
+ block_idx: int = 0,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ hidden_size = config.hidden_size
+ use_gdn = config.softmax_every_n <= 0 or (block_idx + 1) % config.softmax_every_n != 0
+ use_plucker_proj = config.use_chunk_plucker_post_attn and (
+ config.chunk_plucker_post_attn_blocks < 0 or block_idx < config.chunk_plucker_post_attn_blocks
+ )
+ self.norm1 = AdaLayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.attn = SanaWmSelfAttention(
+ config,
+ use_gdn=use_gdn,
+ quant_config=quant_config,
+ prefix=f"{prefix}.attn" if prefix else "attn",
+ )
+ self.cross_attn = SanaWmCrossAttention(
+ config,
+ quant_config=quant_config,
+ prefix=f"{prefix}.cross_attn" if prefix else "cross_attn",
+ )
+ self.norm2 = AdaLayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.mlp = SanaWmMbConvFfn(config)
+ if use_plucker_proj:
+ self.plucker_proj: nn.Module | None = ColumnParallelLinear(
+ hidden_size,
+ hidden_size,
+ bias=True,
+ gather_output=True,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.plucker_proj" if prefix else "plucker_proj",
+ )
+ else:
+ self.plucker_proj = None
+ self.scale_shift_table = nn.Parameter(torch.zeros(6, hidden_size))
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ timestep_modulation: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ rotary_emb: torch.Tensor | None = None,
+ camera_hidden_states: torch.Tensor | None = None,
+ cam_geometry: SanaWmCamGeometry | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ # Per-frame timestep contract (see SanaWmTransformer3DModel.forward).
+ if timestep_modulation.ndim > 2:
+ return self._forward_frame_aware(
+ hidden_states,
+ encoder_hidden_states,
+ timestep_modulation,
+ spatial_shape,
+ rotary_emb,
+ camera_hidden_states,
+ cam_geometry,
+ encoder_attention_mask,
+ )
+ batch_size = hidden_states.shape[0]
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+ self.scale_shift_table[None] + timestep_modulation.reshape(batch_size, 6, -1)
+ ).chunk(6, dim=1)
+ attn_input = self.norm1(hidden_states, scale_msa, shift_msa)
+ attn_output = self.attn(attn_input, spatial_shape, rotary_emb, cam_geometry)
+ if camera_hidden_states is not None and self.plucker_proj is not None:
+ attn_output = attn_output + self.plucker_proj(camera_hidden_states)
+ hidden_states = hidden_states + gate_msa * attn_output
+ hidden_states = hidden_states + self.cross_attn(
+ hidden_states,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ )
+ mlp_input = self.norm2(hidden_states, scale_mlp, shift_mlp)
+ hidden_states = hidden_states + gate_mlp * self.mlp(mlp_input, spatial_shape)
+ return hidden_states
+
+ def _forward_frame_aware(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ timestep_modulation: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ rotary_emb: torch.Tensor | None,
+ camera_hidden_states: torch.Tensor | None,
+ cam_geometry: SanaWmCamGeometry | None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """Per-frame timestep modulation path matching NVlabs
+ ``SanaVideoMSCamCtrlBlock.forward_frame_aware``.
+
+ ``timestep_modulation`` is ``(B, 1, F, 6*D)``. We split into the
+ six per-frame ``(B, F, 1, D)`` chunks via ``scale_shift_table``
+ and apply ``shift/scale/gate`` per frame, broadcasting over the
+ spatial tokens within each frame.
+ """
+ batch_size, token_count, hidden_size = hidden_states.shape
+ frames, height, width = spatial_shape
+ spatial_tokens = height * width
+ if token_count != frames * spatial_tokens:
+ raise ValueError(f"Sana-WM frame-aware block expects N=T*H*W, got N={token_count}, THW={spatial_shape}.")
+ if timestep_modulation.shape[2] != frames:
+ raise ValueError(
+ "Sana-WM frame-aware block: timestep frame axis must match spatial frames, "
+ f"got modulation frames={timestep_modulation.shape[2]} vs spatial frames={frames}."
+ )
+
+ # (B, 1, F, 6*D) -> (B, F, 6, D), then add the (1, 1, 6, D) table.
+ t_per_frame = timestep_modulation.reshape(batch_size, frames, 6, -1)
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+ self.scale_shift_table[None, None, :, :] + t_per_frame
+ ).chunk(6, dim=-2) # each: (B, F, 1, D)
+
+ # Apply per-frame modulation by reshaping x to (B, F, S, D), broadcasting
+ # scale/shift over the spatial axis, then flattening back. The reshape
+ # comes first because the norm only reduces over the last axis, so it is
+ # unaffected, and the modulation then broadcasts as written.
+ x_4d = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size)
+ x_msa_in = self.norm1(x_4d, scale_msa, shift_msa).reshape(batch_size, token_count, hidden_size)
+ attn_output = self.attn(x_msa_in, spatial_shape, rotary_emb, cam_geometry)
+ if camera_hidden_states is not None and self.plucker_proj is not None:
+ attn_output = attn_output + self.plucker_proj(camera_hidden_states)
+ attn_output_4d = attn_output.reshape(batch_size, frames, spatial_tokens, hidden_size)
+ hidden_states = hidden_states + (gate_msa * attn_output_4d).reshape(batch_size, token_count, hidden_size)
+
+ hidden_states = hidden_states + self.cross_attn(
+ hidden_states,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ )
+
+ x_4d = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size)
+ x_mlp_in = self.norm2(x_4d, scale_mlp, shift_mlp).reshape(batch_size, token_count, hidden_size)
+ mlp_out = self.mlp(x_mlp_in, spatial_shape)
+ mlp_out_4d = mlp_out.reshape(batch_size, frames, spatial_tokens, hidden_size)
+ hidden_states = hidden_states + (gate_mlp * mlp_out_4d).reshape(batch_size, token_count, hidden_size)
+ return hidden_states
+
+
+class SanaWmFinalLayer(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ patch_size: tuple[int, int, int],
+ out_channels: int,
+ *,
+ quant_config: Any = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.norm_final = AdaLayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.scale_shift_table = nn.Parameter(torch.zeros(2, hidden_size))
+ self.out_channels = out_channels
+ out_features = math.prod(patch_size) * out_channels
+ self.linear: nn.Module = ColumnParallelLinear(
+ hidden_size,
+ out_features,
+ bias=True,
+ gather_output=True,
+ return_bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.linear" if prefix else "linear",
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ timestep_embed: torch.Tensor,
+ spatial_shape: tuple[int, int, int] | None = None,
+ ) -> torch.Tensor:
+ # Per-frame timestep contract (see SanaWmTransformer3DModel.forward).
+ if timestep_embed.ndim > 2:
+ return self._forward_frame_aware(hidden_states, timestep_embed, spatial_shape)
+ shift, scale = (self.scale_shift_table[None] + timestep_embed[:, None]).chunk(2, dim=1)
+ return self.linear(self.norm_final(hidden_states, scale, shift))
+
+ def _forward_frame_aware(
+ self,
+ hidden_states: torch.Tensor,
+ timestep_embed: torch.Tensor,
+ spatial_shape: tuple[int, int, int] | None,
+ ) -> torch.Tensor:
+ """Per-frame final-layer modulation matching NVlabs
+ ``T2IFinalLayer.forward_frame_aware``.
+
+ ``timestep_embed`` is ``(B, 1, F, D)``. We transpose to
+ ``(B, F, 1, D)`` so it adds correctly to the ``(1, 1, 2, D)``
+ ``scale_shift_table`` and produces per-frame ``(B, F, 1, D)``
+ ``shift`` / ``scale`` that broadcast over the spatial tokens
+ within each frame.
+ """
+ batch_size, token_count, hidden_size = hidden_states.shape
+ frames = timestep_embed.shape[2]
+ if spatial_shape is not None:
+ spatial_tokens = spatial_shape[1] * spatial_shape[2]
+ else:
+ spatial_tokens = token_count // frames
+ if frames * spatial_tokens != token_count:
+ raise ValueError(
+ "Sana-WM frame-aware final layer: token count mismatch "
+ f"(N={token_count}, F={frames}, S={spatial_tokens})."
+ )
+ # (B, 1, F, D) -> (B, F, 1, D); add (1, 1, 2, D); chunk into shift/scale: each (B, F, 1, D).
+ t_per_frame = timestep_embed.transpose(1, 2)
+ shift, scale = (self.scale_shift_table[None, None, :, :] + t_per_frame).chunk(2, dim=-2)
+ x_4d = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size)
+ x_mod = self.norm_final(x_4d, scale, shift).reshape(batch_size, token_count, hidden_size)
+ return self.linear(x_mod)
+
+
+class SanaWmTransformer3DModel(nn.Module):
+ """SANA-WM Stage-1 DiT: bidirectional Gated DeltaNet blocks with a softmax
+ attention block every ``softmax_every_n``."""
+
+ _repeated_blocks: ClassVar[list[str]] = ["SanaWmBlock"]
+ _layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["blocks"]
+ _hsdp_shard_conditions: ClassVar[list[Any]] = [_is_sana_wm_transformer_block]
+
+ def __init__(
+ self,
+ config: SanaWmConfig | None = None,
+ *,
+ device: torch.device | str | None = None,
+ dtype: torch.dtype | None = None,
+ latent_channels: int = SANA_WM_STAGE1_LATENT_CHANNELS,
+ prompt_channels: int = SANA_WM_STAGE1_PROMPT_CHANNELS,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.config = config or SanaWmConfig()
+ self.quant_config = quant_config
+ self.prefix = prefix
+ self._latent_channels = latent_channels
+ self._prompt_channels = prompt_channels
+ self.patch_size = _to_3tuple(self.config.patch_size)
+ self.x_embedder = SanaWmPatchEmbedMS3D(self.patch_size, self._latent_channels, self.config.hidden_size)
+ self.y_embedder = SanaWmTextEmbedder(
+ self._prompt_channels,
+ self.config.hidden_size,
+ self.config.model_max_length,
+ quant_config=self.quant_config,
+ prefix=f"{self.prefix}.y_embedder" if self.prefix else "y_embedder",
+ )
+ self.t_embedder = SanaWmTimestepEmbedder(
+ SANA_WM_STAGE1_TIMESTEP_CHANNELS,
+ self.config.hidden_size,
+ quant_config=self.quant_config,
+ prefix=f"{self.prefix}.t_embedder" if self.prefix else "t_embedder",
+ )
+ # t_block: SiLU → Linear(hidden, 6*hidden). Weight key is t_block.1.{weight,bias}
+ # to match the Stage-1 checkpoint layout (nn.Sequential index 1).
+ _t_block_linear: nn.Module = ColumnParallelLinear(
+ self.config.hidden_size,
+ 6 * self.config.hidden_size,
+ bias=True,
+ gather_output=True,
+ return_bias=False,
+ quant_config=self.quant_config,
+ prefix=f"{self.prefix}.t_block.1" if self.prefix else "t_block.1",
+ )
+ self.t_block = nn.Sequential(nn.SiLU(), _t_block_linear)
+ self.plucker_embedder = SanaWmPatchEmbedMS3D(
+ self.patch_size,
+ self.config.chunk_plucker_channels,
+ self.config.hidden_size,
+ )
+ self.raymap_embedder = SanaWmPatchEmbedMS3D(self.patch_size, 3, self.config.hidden_size)
+ self.blocks = nn.ModuleList(
+ [
+ SanaWmBlock(
+ self.config,
+ block_idx=i,
+ quant_config=self.quant_config,
+ prefix=f"{self.prefix}.blocks.{i}" if self.prefix else f"blocks.{i}",
+ )
+ for i in range(self.config.num_blocks)
+ ]
+ )
+ self.final_layer = SanaWmFinalLayer(
+ self.config.hidden_size,
+ self.patch_size,
+ self._latent_channels,
+ quant_config=self.quant_config,
+ prefix=f"{self.prefix}.final_layer" if self.prefix else "final_layer",
+ )
+ self.pos_embed = nn.Parameter(torch.zeros(1, 484, self.config.hidden_size))
+ self.rope = SanaWmWanRotaryPosEmbed(self.config.linear_head_dim)
+ self.attention_y_norm = RMSNorm(self.config.hidden_size)
+ if device is not None or dtype is not None:
+ self.to(device=device, dtype=dtype)
+
+ def _positional_embedding(self, token_count: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+ pos_embed = self.pos_embed.to(device=device, dtype=dtype)
+ if pos_embed.shape[1] == token_count:
+ return pos_embed
+ pos_embed = pos_embed.transpose(1, 2)
+ pos_embed = F.interpolate(pos_embed, size=token_count, mode="linear", align_corners=False)
+ return pos_embed.transpose(1, 2)
+
+ @staticmethod
+ def _match_tokens(hidden_states: torch.Tensor, expected_tokens: int) -> torch.Tensor:
+ if hidden_states.shape[1] == expected_tokens:
+ return hidden_states
+ if hidden_states.shape[1] > expected_tokens:
+ return hidden_states[:, :expected_tokens]
+ pad = hidden_states[:, -1:].expand(-1, expected_tokens - hidden_states.shape[1], -1)
+ return torch.cat([hidden_states, pad], dim=1)
+
+ def _camera_hidden_states_from_conditions(
+ self,
+ *,
+ plucker: torch.Tensor | None,
+ spatial_raymap: torch.Tensor | None,
+ spatial_shape: tuple[int, int, int],
+ batch_size: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ ) -> torch.Tensor | None:
+ if plucker is None:
+ return None
+ if plucker.ndim == 4:
+ plucker = plucker.unsqueeze(0)
+ plucker = plucker.to(device=device, dtype=dtype)
+ if plucker.shape[0] == 1 and batch_size > 1:
+ plucker = plucker.expand(batch_size, -1, -1, -1, -1)
+ camera_hidden_states = self.plucker_embedder(plucker)
+ expected_tokens = spatial_shape[0] * spatial_shape[1] * spatial_shape[2]
+ camera_hidden_states = self._match_tokens(camera_hidden_states, expected_tokens)
+
+ # Fuse per-pixel ray-direction map via raymap_embedder only on the
+ # non-plucker path. NVlabs skips the absmap/raymap embedder whenever
+ # chunk-plucker post-attention injection is enabled.
+ if spatial_raymap is not None and not self.config.use_chunk_plucker_post_attn:
+ if spatial_raymap.ndim == 4: # [C, F, H, W] → [1, C, F, H, W]
+ spatial_raymap = spatial_raymap.unsqueeze(0)
+ spatial_raymap = spatial_raymap.to(device=device, dtype=dtype)
+ if spatial_raymap.shape[0] == 1 and batch_size > 1:
+ spatial_raymap = spatial_raymap.expand(batch_size, -1, -1, -1, -1)
+ ray_hidden = self._match_tokens(self.raymap_embedder(spatial_raymap), expected_tokens)
+ camera_hidden_states = camera_hidden_states + ray_hidden
+
+ return camera_hidden_states
+
+ def _unpatchify(self, hidden_states: torch.Tensor, spatial_shape: tuple[int, int, int]) -> torch.Tensor:
+ batch_size = hidden_states.shape[0]
+ frames, height, width = spatial_shape
+ patch_frames, patch_height, patch_width = self.patch_size
+ hidden_states = hidden_states.reshape(
+ batch_size,
+ frames,
+ height,
+ width,
+ patch_frames,
+ patch_height,
+ patch_width,
+ self._latent_channels,
+ )
+ hidden_states = torch.einsum("nfhwopqc->ncfohpwq", hidden_states)
+ return hidden_states.reshape(
+ batch_size,
+ self._latent_channels,
+ frames * patch_frames,
+ height * patch_height,
+ width * patch_width,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ timestep: torch.Tensor | float | int,
+ *,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ camera_hidden_states: torch.Tensor | None = None,
+ plucker: torch.Tensor | None = None,
+ raymap: torch.Tensor | None = None,
+ spatial_raymap: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ if hidden_states.ndim != 5:
+ raise ValueError("Sana-WM transformer expects latent input shaped [B, C, F, H, W].")
+ batch_size = hidden_states.shape[0]
+ latent_shape = hidden_states.shape[2:]
+ hidden_states, spatial_shape = self.x_embedder.project_with_shape(hidden_states)
+ rotary_emb = None
+ if self.config.pos_embed_type == "wan_rope":
+ rotary_emb = self.rope(spatial_shape, hidden_states.device)
+ else:
+ hidden_states = hidden_states + self._positional_embedding(
+ hidden_states.shape[1], hidden_states.dtype, hidden_states.device
+ )
+
+ encoder_hidden_states = self.y_embedder(
+ encoder_hidden_states.to(device=hidden_states.device) if encoder_hidden_states is not None else None,
+ batch_size=batch_size,
+ dtype=hidden_states.dtype,
+ )
+ encoder_hidden_states = self.attention_y_norm(encoder_hidden_states)
+ # No mask means "attend to every prompt token"; keep it ``None`` so
+ # cross-attention takes the mask-free attention path rather than
+ # building an all-ones mask the backend would have to reshape per block.
+ if encoder_attention_mask is not None:
+ encoder_attention_mask = encoder_attention_mask.to(device=hidden_states.device)
+ if encoder_attention_mask.shape != encoder_hidden_states.shape[:2]:
+ raise ValueError(
+ "Sana-WM encoder_attention_mask must match text tokens "
+ f"{tuple(encoder_hidden_states.shape[:2])}, got {tuple(encoder_attention_mask.shape)}."
+ )
+
+ if not torch.is_tensor(timestep):
+ timestep = torch.tensor([timestep], device=hidden_states.device, dtype=torch.float32)
+ timestep = timestep.to(device=hidden_states.device)
+ if timestep.ndim == 0:
+ timestep = timestep.expand(batch_size)
+ elif timestep.ndim == 1 and timestep.shape[0] == 1 and batch_size > 1:
+ timestep = timestep.expand(batch_size)
+ # Per-frame timestep contract: a rank > 1 ``timestep`` (``(B, 1, F)``
+ # from the LTX flow-matching sampler) embeds on the flattened token axis
+ # and unflattens back, giving ``(B, 1, F, D)`` / ``(B, 1, F, 6*D)``.
+ # SanaWmBlock and SanaWmFinalLayer dispatch on ``ndim > 2`` to the
+ # frame-aware paths that broadcast modulation over spatial tokens.
+ timestep_shape = tuple(timestep.shape)
+ time_embed = self.t_embedder(timestep) # (numel, D)
+ _t_silu = self.t_block[0](time_embed)
+ timestep_modulation = self.t_block[1](_t_silu).to(hidden_states.dtype)
+ if len(timestep_shape) > 1:
+ time_embed = time_embed.unflatten(0, timestep_shape)
+ timestep_modulation = timestep_modulation.unflatten(0, timestep_shape)
+
+ if camera_hidden_states is None:
+ camera_hidden_states = self._camera_hidden_states_from_conditions(
+ plucker=plucker,
+ spatial_raymap=spatial_raymap,
+ spatial_shape=spatial_shape,
+ batch_size=batch_size,
+ dtype=hidden_states.dtype,
+ device=hidden_states.device,
+ )
+
+ # Build (B, T_latent, 20) camera_conditions from the raw raymap, then
+ # derive the UCPE geometry ONCE for the whole forward. The ray grid,
+ # its SE(3) inverse and the camera RoPE tables are identical in every
+ # block, so building them per block cost ~2400 rebuilds per request
+ # (20 blocks x 60 steps x 2 CFG branches). When the caller passes
+ # pre-embedded camera_hidden_states without raymap, the cam branch
+ # becomes a no-op (main-branch only).
+ cam_geometry: SanaWmCamGeometry | None = None
+ if raymap is not None:
+ # ``_pack_camera_conditions`` emits an unbatched (T_latent, 20).
+ camera_conditions = raymap.unsqueeze(0)
+ if batch_size > 1:
+ camera_conditions = camera_conditions.expand(batch_size, -1, -1)
+ camera_conditions = camera_conditions.to(device=hidden_states.device, dtype=hidden_states.dtype)
+ # cam_head_dim comes from the attention module rather than being
+ # re-derived here: __init__ already asserts it equals the main
+ # head_dim and is consistent across blocks.
+ attn = self.blocks[0].attn
+ cam_geometry = prepare_cam_geometry(
+ camera_conditions=camera_conditions,
+ spatial_shape=spatial_shape,
+ patch_size=self.patch_size,
+ head_dim=attn.cam_head_dim,
+ rotary_emb=attn._ucpe_rotary_freqs(rotary_emb),
+ )
+
+ for block in self.blocks:
+ hidden_states = block(
+ hidden_states,
+ encoder_hidden_states,
+ timestep_modulation,
+ spatial_shape,
+ rotary_emb,
+ camera_hidden_states,
+ cam_geometry,
+ encoder_attention_mask,
+ )
+
+ hidden_states = self.final_layer(hidden_states, time_embed.to(hidden_states.dtype), spatial_shape)
+ hidden_states = self._unpatchify(hidden_states, spatial_shape)
+ if hidden_states.shape[2:] != latent_shape:
+ hidden_states = F.interpolate(hidden_states, size=latent_shape, mode="trilinear", align_corners=False)
+ return hidden_states
+
+ def load_weights(self, weights: Iterable[tuple[str, Any]]) -> set[str]:
+ """Stream checkpoint tensors into the eagerly-built modules.
+
+ Follows the wan2_2 idiom: one ``named_parameters`` lookup table plus the
+ per-tensor ``weight_loader``, which narrows full checkpoint tensors to
+ the TP-local shard at copy time for both vLLM parallel layers and the
+ plain parameters marked by ``_shard_param_across_tp``. The source tensor
+ is dropped each iteration — no copy of the checkpoint is retained.
+ """
+ params_dict = dict(self.named_parameters())
+ buffers_dict = dict(self.named_buffers())
+ loaded: set[str] = set()
+ unmapped: list[str] = []
+ duplicates: list[str] = []
+
+ for source_name, tensor in weights:
+ # Weights ship in the standard diffusers ``transformer/`` layout, so
+ # the checkpoint key is already the module-local parameter name.
+ if source_name in loaded:
+ duplicates.append(source_name)
+ continue
+ if not isinstance(tensor, torch.Tensor):
+ raise TypeError(f"Sana-WM weight {source_name!r} must be a torch.Tensor, got {type(tensor).__name__}.")
+ target = params_dict.get(source_name)
+ if target is None:
+ target = buffers_dict.get(source_name)
+ if target is None:
+ unmapped.append(source_name)
+ continue
+ weight_loader = getattr(target, "weight_loader", None)
+ if callable(weight_loader):
+ # vLLM parallel layers and the plain params marked by
+ # `_shard_param_across_tp` narrow to the local shard here.
+ weight_loader(target, tensor)
+ else:
+ if tuple(target.shape) != tuple(tensor.shape):
+ raise ValueError(
+ f"Sana-WM weight shape mismatch for {source_name}: "
+ f"expected {tuple(target.shape)}, got {tuple(tensor.shape)}."
+ )
+ with torch.no_grad():
+ target.copy_(tensor.to(device=target.device, dtype=target.dtype))
+ loaded.add(source_name)
+
+ if unmapped or duplicates:
+ details = []
+ if unmapped:
+ details.append(f"unmapped={unmapped[:10]}")
+ if duplicates:
+ details.append(f"duplicates={duplicates[:10]}")
+ raise ValueError("Invalid SANA-WM Stage-1 checkpoint keys: " + "; ".join(details))
+ return loaded
diff --git a/vllm_omni/diffusion/models/sana_wm/ucpe.py b/vllm_omni/diffusion/models/sana_wm/ucpe.py
new file mode 100644
index 000000000..88a3a115d
--- /dev/null
+++ b/vllm_omni/diffusion/models/sana_wm/ucpe.py
@@ -0,0 +1,536 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SANA-WM UCPE (Unified Camera Pose Embedding) per-block attention transforms.
+
+Ported from NVlabs/Sana ``sana_camctrl_blocks.py`` for the inference-only
+SANA-WM bidirectional 1600M release. Scope is intentionally narrow:
+
+* Pinhole camera only (xi=0); the UCM ``xi`` parameter is fixed.
+* Inference path only — no training-time camera-branch dropout.
+* Online computation only — the precomputed ``cam_pos_embeds`` shortcut
+ used by NVlabs' fused kernels is omitted because vLLM-Omni builds the
+ matrices once per transformer forward and shares them across all blocks.
+* ``apply_vo=True`` only — the SANA-WM checkpoint always uses the
+ inverse-output transform.
+
+Public surface:
+ ``prepare_cam_geometry(...)`` -> :class:`SanaWmCamGeometry`, built once per
+ transformer forward and threaded through every block. It carries both the
+ ``(apply_q, apply_kv, apply_output)`` closures used by the softmax camera
+ branch and the projection/RoPE tensors used by ``cam_prep_func(...)``, the
+ NVlabs fused camera-prep contract on the GDN camera branch.
+
+Each closure transforms a tensor of shape ``(B, num_heads, N, D)`` where
+``N = T * H * W`` (matching the GDN token layout) and ``D == head_dim``,
+returning a tensor of the same shape.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# Pinhole ray-grid construction
+# ---------------------------------------------------------------------------
+
+
+def _make_pixel_grid(
+ height: int,
+ width: int,
+ device: torch.device,
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ """Return ``(H, W, 3)`` grid ``(x, y, 1)`` of pixel coordinates."""
+ xs = torch.arange(width, device=device, dtype=dtype)
+ ys = torch.arange(height, device=device, dtype=dtype)
+ y_grid, x_grid = torch.meshgrid(ys, xs, indexing="ij")
+ return torch.stack((x_grid, y_grid, torch.ones_like(x_grid)), dim=-1)
+
+
+def _unproject_pinhole(
+ fx: torch.Tensor,
+ fy: torch.Tensor,
+ cx: torch.Tensor,
+ cy: torch.Tensor,
+ height: int,
+ width: int,
+) -> torch.Tensor:
+ """Unproject a pixel grid into camera-frame direction vectors.
+
+ Args:
+ fx, fy, cx, cy: shape ``(B, T)`` intrinsics in token-grid units — the
+ caller divides all four by patch size before calling.
+ height, width: token-grid size.
+
+ Returns:
+ ``(B, T, H, W, 3)`` direction tensor (NOT normalised).
+ """
+ B, T = fx.shape
+ device = fx.device
+ dtype = fx.dtype
+ grid = _make_pixel_grid(height, width, device, dtype) # (H, W, 3)
+ u = grid[..., 0] # (H, W)
+ v = grid[..., 1]
+ fx_e = fx.view(B, T, 1, 1)
+ fy_e = fy.view(B, T, 1, 1)
+ cx_e = cx.view(B, T, 1, 1)
+ cy_e = cy.view(B, T, 1, 1)
+ x = (u - cx_e) / fx_e
+ y = (v - cy_e) / fy_e
+ z = torch.ones_like(x)
+ # Match NVlabs UCM normalisation at xi=0: alpha = 1, gamma = 1/(1+r^2),
+ # X = gamma*x, Y = gamma*y, Z = gamma. The resulting vector has unit
+ # norm. Computed compactly via direct stack + normalise.
+ d_cam = torch.stack((x, y, z), dim=-1) # (B, T, H, W, 3)
+ d_cam = d_cam / d_cam.norm(dim=-1, keepdim=True).clamp_min(1e-12)
+ return d_cam
+
+
+def _world_to_ray_mats(
+ d_cam: torch.Tensor, # (B, T, H, W, 3) unit rays in camera frame
+ c2w: torch.Tensor, # (B, T, 4, 4) camera-to-world SE(3)
+) -> torch.Tensor:
+ """Build per-pixel ``ray <- world`` SE(3) transforms ``(B, T, H, W, 4, 4)``.
+
+ The ray frame is anchored at the camera origin with the +Z axis along
+ the per-pixel viewing direction and a +Y axis defined by the camera-Y
+ cross product. This matches NVlabs ``world_to_ray_mats``.
+ """
+ B, T, H, W, _ = d_cam.shape
+ device = d_cam.device
+ dtype = d_cam.dtype
+ rot_cam = c2w[..., :3, :3]
+ trans_cam = c2w[..., :3, 3]
+ d_world = torch.einsum("btij,bthwj->bthwi", rot_cam, d_cam)
+ cam_y = rot_cam[..., :, 1].unsqueeze(2).unsqueeze(3).expand(B, T, H, W, 3)
+ z_ray = F.normalize(d_world, dim=-1, eps=1e-6)
+ 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)
+ rot_local_to_world = torch.stack((x_ray, y_ray, z_ray), dim=-1) # (B,T,H,W,3,3)
+ rot_world_to_local = rot_local_to_world.transpose(-1, -2)
+ t_world = trans_cam.unsqueeze(2).unsqueeze(3).expand(B, T, H, W, 3)
+ t_w2l = -torch.einsum("bthwij,bthwj->bthwi", rot_world_to_local, t_world)
+ raymats = torch.zeros(B, T, H, W, 4, 4, device=device, dtype=dtype)
+ raymats[..., :3, :3] = rot_world_to_local
+ raymats[..., :3, 3] = t_w2l
+ raymats[..., 3, 3] = 1.0
+ # Guard NaNs from degenerate pixels (e.g. nan direction at infinity).
+ nan_mask = torch.isnan(d_world).any(-1)
+ if nan_mask.any():
+ eye = torch.eye(4, device=device, dtype=dtype)
+ raymats[nan_mask] = eye
+ return raymats
+
+
+def _process_camera_conditions(
+ camera_conditions: torch.Tensor, # (B, T, 20)
+ hw: tuple[int, int, int],
+ patch_size: tuple[int, int, int] = (1, 1, 1),
+) -> torch.Tensor:
+ """Convert ``(B, T, 20)`` SANA-WM camera conditions to ``(B, T, H, W, 4, 4)``
+ ``ray <- world`` transforms.
+
+ Layout of the last axis is ``[c2w_flat (16), fx, fy, cx, cy]``.
+ Intrinsics follow the NVlabs UCPE contract: the packed camera payload is
+ expressed on the pre-patch latent image grid, while ``HW`` is the
+ transformer token grid. ``patch_size`` maps those units to the token grid
+ before unprojection. The released SANA-WM P1 config uses ``(1, 1, 1)``.
+
+ Args:
+ camera_conditions: ``(B, T, 20)`` from ``_pack_camera_conditions``.
+ hw: ``(T, H, W)`` token-grid shape. Equal to the latent grid only
+ because the released config patches by ``(1, 1, 1)``.
+ """
+ B, T_cond, last = camera_conditions.shape
+ if last != 20:
+ raise ValueError(f"SANA-WM camera_conditions last dim must be 20, got {last}.")
+ T_latent, H_latent, W_latent = hw
+ if T_cond != T_latent:
+ raise ValueError(f"SANA-WM camera_conditions frames {T_cond} != latent frames {T_latent}.")
+ c2w = camera_conditions[..., :16].reshape(B, T_latent, 4, 4)
+ patch_t, patch_h, patch_w = patch_size
+ del patch_t
+ fx = camera_conditions[..., 16] / float(patch_w)
+ fy = camera_conditions[..., 17] / float(patch_h)
+ cx = camera_conditions[..., 18] / float(patch_w)
+ cy = camera_conditions[..., 19] / float(patch_h)
+ d_cam = _unproject_pinhole(fx, fy, cx, cy, H_latent, W_latent)
+ raymats = _world_to_ray_mats(d_cam, c2w)
+ return raymats
+
+
+def _invert_se3(transforms: torch.Tensor) -> torch.Tensor:
+ """Closed-form inverse of a stack of 4x4 SE(3) matrices."""
+ if transforms.shape[-2:] != (4, 4):
+ raise ValueError(f"SE3 inverse expects (..., 4, 4), got {tuple(transforms.shape)}.")
+ rot_inv = transforms[..., :3, :3].transpose(-1, -2)
+ out = torch.zeros_like(transforms)
+ out[..., :3, :3] = rot_inv
+ out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv, transforms[..., :3, 3])
+ out[..., 3, 3] = 1.0
+ return out
+
+
+# ---------------------------------------------------------------------------
+# Per-token transform primitives
+# ---------------------------------------------------------------------------
+
+
+def _apply_ray_projmat(
+ feats: torch.Tensor, # (B, num_heads, N, D)
+ matrix: torch.Tensor, # (B, N, 4, 4)
+) -> torch.Tensor:
+ """Apply a per-token 4x4 projection to features grouped by 4 channels.
+
+ ``D`` must be a multiple of 4. Out-of-place; preserves ``feats.shape``.
+ """
+ B, num_heads, N, D = feats.shape
+ if D % 4 != 0:
+ raise ValueError(f"UCPE projmat expects D divisible by 4, got D={D}.")
+ feats_grouped = feats.reshape(B, num_heads, N, D // 4, 4)
+ # matrix: (B, N, 4, 4); feats_grouped: (B, num_heads, N, K, 4)
+ out = torch.einsum("bnij,bhnkj->bhnki", matrix, feats_grouped)
+ return out.reshape(B, num_heads, N, D)
+
+
+def _apply_complex_rope(
+ hidden_states: torch.Tensor,
+ freqs: torch.Tensor,
+ *,
+ inverse: bool = False,
+) -> torch.Tensor:
+ """Apply complex RoPE: treat last-dim pairs as ``(real, imag)`` and multiply
+ by ``freqs``. Computation runs in fp64 to match the NVlabs reference.
+
+ Args:
+ hidden_states: ``(..., D)`` with ``D % 2 == 0``.
+ freqs: complex tensor broadcastable to ``(..., D // 2)``.
+ inverse: when ``True`` apply ``freqs.conj()`` (used for ``apply_o``).
+ """
+ x_real = hidden_states.to(torch.float64).contiguous()
+ x_complex = torch.view_as_complex(x_real.unflatten(-1, (-1, 2)))
+ if inverse:
+ freqs = freqs.conj()
+ x_out = torch.view_as_real(x_complex * freqs).flatten(-2, -1)
+ return x_out.to(hidden_states.dtype)
+
+
+def _apply_block_diagonal(
+ feats: torch.Tensor,
+ func_size_pairs: list[tuple[Callable[[torch.Tensor], torch.Tensor], int]],
+) -> torch.Tensor:
+ """Split ``feats`` along the last axis by ``block_sizes`` and apply one
+ callable per block, then concatenate. Preserves ``feats.shape``.
+ """
+ funcs, block_sizes = zip(*func_size_pairs)
+ if feats.shape[-1] != sum(block_sizes):
+ raise ValueError(f"UCPE block-diagonal block sum {sum(block_sizes)} != feats D {feats.shape[-1]}.")
+ blocks = torch.split(feats, list(block_sizes), dim=-1)
+ return torch.cat([func(block) for func, block in zip(funcs, blocks)], dim=-1)
+
+
+def _slice_rope_for_cam(
+ rotary_emb: torch.Tensor | None,
+ rope_dim: int,
+) -> torch.Tensor | None:
+ """Re-slice WAN-style RoPE frequencies for a smaller ``rope_dim`` using the
+ same ``(T, H, W)`` split as the main branch. Mirrors NVlabs slicing logic.
+
+ The T/H/W band layout is read off the input shape
+ (``rotary_emb.shape[-1] * 2``) rather than a separately supplied head_dim,
+ so the split stays correct without the caller passing it. ``cam_head_dim``
+ still has to equal the main ``head_dim`` — ``SanaWmSelfAttention.__init__``
+ rejects anything else.
+ """
+ if rotary_emb is None:
+ return None
+ # Derive WAN RoPE T/H/W band offsets from the actual input tensor size,
+ # not from a head_dim argument. This is correct regardless of whether
+ # cam_head_dim == main_head_dim.
+ in_head_dim = rotary_emb.shape[-1] * 2
+ orig_t = in_head_dim // 2 - 2 * (in_head_dim // 6)
+ orig_h = in_head_dim // 6
+ new_t = rope_dim // 2 - 2 * (rope_dim // 6)
+ new_h = rope_dim // 6
+ new_w = rope_dim // 6
+ t_part = rotary_emb[..., :new_t]
+ h_part = rotary_emb[..., orig_t : orig_t + new_h]
+ w_part = rotary_emb[..., orig_t + orig_h : orig_t + orig_h + new_w]
+ return torch.cat([t_part, h_part, w_part], dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# Apply-fn preparation
+# ---------------------------------------------------------------------------
+
+
+def _identity(x: torch.Tensor) -> torch.Tensor:
+ return x
+
+
+def _prepare_ray_apply_fns(
+ head_dim: int,
+ proj: torch.Tensor, # (B, N, 4, 4) ray <- world
+ proj_t: torch.Tensor, # (B, N, 4, 4) ray <- world transpose
+ proj_inv: torch.Tensor, # (B, N, 4, 4) world <- ray
+ rotary_emb: torch.Tensor | None = None,
+) -> tuple[Callable, Callable, Callable]:
+ """Build ``(apply_q, apply_kv, apply_o)`` block-diagonal callables.
+
+ Each callable transforms a ``(B, num_heads, N, head_dim)`` tensor:
+ the first ``head_dim // 2`` channels go through the per-token 4x4
+ projection, the remaining ``head_dim // 2`` channels go through complex
+ RoPE (or identity when ``rotary_emb`` is None).
+ """
+ if rotary_emb is not None:
+ rope_fwd = partial(_apply_complex_rope, freqs=rotary_emb, inverse=False)
+ rope_inv = partial(_apply_complex_rope, freqs=rotary_emb, inverse=True)
+ else:
+ rope_fwd = _identity
+ rope_inv = _identity
+
+ half = head_dim // 2
+
+ transforms_q = [
+ (partial(_apply_ray_projmat, matrix=proj_t), half),
+ (rope_fwd, half),
+ ]
+ transforms_kv = [
+ (partial(_apply_ray_projmat, matrix=proj_inv), half),
+ (rope_fwd, half),
+ ]
+ transforms_o = [
+ (partial(_apply_ray_projmat, matrix=proj), half),
+ (rope_inv, half),
+ ]
+
+ apply_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q)
+ apply_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv)
+ apply_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o)
+ return apply_q, apply_kv, apply_o
+
+
+@dataclass(frozen=True)
+class SanaWmCamGeometry:
+ """UCPE geometry shared by every block of one transformer forward.
+
+ The ray grid, its SE(3) inverse and the camera RoPE tables depend only on
+ the camera payload, the latent grid, the patch size and the head dim — all
+ fixed for the whole forward — yet each block's camera branch used to
+ rebuild them. At 20 blocks x 60 steps x 2 CFG branches that is ~2400
+ rebuilds of an 18480-token ray grid and its 4x4 inverses per request.
+
+ ``apply_q``/``apply_kv``/``apply_output`` serve the softmax camera branch;
+ ``proj_q``/``proj_kv``/``rope_cos``/``rope_sin`` serve the fused
+ ``cam_prep_func`` contract on the GDN branch. Both are derived from the
+ same projection matrices, so they are built together.
+ """
+
+ proj_q: torch.Tensor
+ proj_kv: torch.Tensor
+ rope_cos: torch.Tensor
+ rope_sin: torch.Tensor
+ apply_q: Callable[[torch.Tensor], torch.Tensor]
+ apply_kv: Callable[[torch.Tensor], torch.Tensor]
+ apply_output: Callable[[torch.Tensor], torch.Tensor]
+
+
+def _prepare_ucpe_rope_tables(
+ rotary_emb_cam: torch.Tensor,
+ token_count: int,
+ half_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Convert complex RoPE freqs to NVlabs' interleaved cos/sin tables.
+
+ The output convention is ``out[d] = x[d] * cos[d] + x[d ^ 1] * sin[d]``.
+ For every complex pair, ``sin[2i] = -imag(freq[i])`` and
+ ``sin[2i + 1] = imag(freq[i])``.
+ """
+ freqs = rotary_emb_cam
+ while freqs.ndim > 2 and freqs.shape[0] == 1:
+ freqs = freqs.squeeze(0)
+ if freqs.ndim != 2:
+ raise ValueError(f"Sana-WM cam RoPE freqs must be 2D after squeeze, got {tuple(freqs.shape)}.")
+ if freqs.shape[0] != token_count:
+ raise ValueError(f"Sana-WM cam RoPE token count {freqs.shape[0]} != {token_count}.")
+ if freqs.shape[1] * 2 != half_dim:
+ raise ValueError(f"Sana-WM cam RoPE dim {freqs.shape[1] * 2} != half_dim {half_dim}.")
+ cos_half = freqs.real.float()
+ sin_half = freqs.imag.float()
+ rope_cos = cos_half.repeat_interleave(2, dim=-1).contiguous()
+ rope_sin = torch.stack((-sin_half, sin_half), dim=-1).reshape(token_count, half_dim).contiguous()
+ return rope_cos, rope_sin
+
+
+def prepare_cam_geometry(
+ *,
+ camera_conditions: torch.Tensor,
+ spatial_shape: tuple[int, int, int],
+ patch_size: tuple[int, int, int],
+ head_dim: int,
+ rotary_emb: torch.Tensor | None,
+) -> SanaWmCamGeometry:
+ """Build the per-forward UCPE geometry: ray projections, RoPE tables, apply fns.
+
+ Call once per transformer forward and share the result across every block —
+ see :class:`SanaWmCamGeometry` for why.
+ """
+ if head_dim % 4 != 0:
+ raise ValueError(f"UCPE head_dim must be divisible by 4, got {head_dim}.")
+ if head_dim % 2 != 0 or (head_dim // 2) % 4 != 0:
+ raise ValueError(f"Sana-WM cam prep head_dim={head_dim} must satisfy D % 2 == 0 and (D/2) % 4 == 0.")
+
+ batch = camera_conditions.shape[0]
+ frames, height, width = spatial_shape
+ token_count = frames * height * width
+ half_dim = head_dim // 2
+
+ raymats = _process_camera_conditions(camera_conditions, spatial_shape, patch_size=patch_size)
+ proj = raymats.reshape(batch, token_count, 4, 4).contiguous()
+ proj_q = proj.transpose(-1, -2).contiguous()
+ proj_kv = _invert_se3(proj).contiguous()
+
+ rotary_emb_cam = _slice_rope_for_cam(rotary_emb, half_dim)
+ if rotary_emb_cam is None:
+ rope_cos = torch.ones(token_count, half_dim, device=camera_conditions.device, dtype=torch.float32)
+ rope_sin = torch.zeros(token_count, half_dim, device=camera_conditions.device, dtype=torch.float32)
+ else:
+ rope_cos, rope_sin = _prepare_ucpe_rope_tables(rotary_emb_cam, token_count, half_dim)
+
+ apply_q, apply_kv, apply_output = _prepare_ray_apply_fns(head_dim, proj, proj_q, proj_kv, rotary_emb=rotary_emb_cam)
+ return SanaWmCamGeometry(
+ proj_q=proj_q,
+ proj_kv=proj_kv,
+ rope_cos=rope_cos,
+ rope_sin=rope_sin,
+ apply_q=apply_q,
+ apply_kv=apply_kv,
+ apply_output=apply_output,
+ )
+
+
+def _validate_cam_prep_inputs(
+ q_raw: torch.Tensor,
+ k_raw: torch.Tensor,
+ v_raw: torch.Tensor,
+ proj_q: torch.Tensor,
+ proj_kv: torch.Tensor,
+ rope_cos: torch.Tensor,
+ rope_sin: torch.Tensor,
+) -> int:
+ """Validate the fused cam-prep contract and return ``head_dim // 2``."""
+ if q_raw.shape != k_raw.shape or q_raw.shape != v_raw.shape:
+ raise ValueError(f"Sana-WM cam prep q/k/v shapes must match, got {q_raw.shape}, {k_raw.shape}, {v_raw.shape}.")
+ if q_raw.ndim != 4:
+ raise ValueError(f"Sana-WM cam prep q/k/v must be (B,N,H,D), got {tuple(q_raw.shape)}.")
+ batch, token_count, _num_heads, head_dim = q_raw.shape
+ if head_dim % 2 != 0 or (head_dim // 2) % 4 != 0:
+ raise ValueError(f"Sana-WM cam prep head_dim={head_dim} must satisfy D % 2 == 0 and (D/2) % 4 == 0.")
+ half_dim = head_dim // 2
+ if proj_q.shape != (batch, token_count, 4, 4):
+ raise ValueError(f"Sana-WM cam prep proj_q shape {tuple(proj_q.shape)} != {(batch, token_count, 4, 4)}.")
+ if proj_kv.shape != (batch, token_count, 4, 4):
+ raise ValueError(f"Sana-WM cam prep proj_kv shape {tuple(proj_kv.shape)} != {(batch, token_count, 4, 4)}.")
+ if rope_cos.shape != (token_count, half_dim) or rope_sin.shape != (token_count, half_dim):
+ raise ValueError(
+ "Sana-WM cam prep rope table shapes must be "
+ f"{(token_count, half_dim)}, got {tuple(rope_cos.shape)} and {tuple(rope_sin.shape)}."
+ )
+ return half_dim
+
+
+def _apply_ray_projection(feats: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
+ batch, token_count, num_heads, half_dim = feats.shape
+ grouped = feats.reshape(batch, token_count, num_heads, half_dim // 4, 4)
+ out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), grouped)
+ return out.reshape(batch, token_count, num_heads, half_dim)
+
+
+def _apply_interleaved_rope(feats: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor) -> torch.Tensor:
+ half_dim = feats.shape[-1]
+ pair_index = torch.arange(half_dim, device=feats.device) ^ 1
+ paired = feats.index_select(-1, pair_index)
+ return feats * rope_cos.reshape(1, rope_cos.shape[0], 1, half_dim) + paired * rope_sin.reshape(
+ 1, rope_sin.shape[0], 1, half_dim
+ )
+
+
+def cam_prep_func(
+ q_normed: torch.Tensor,
+ k_normed: torch.Tensor,
+ v_raw: torch.Tensor,
+ *,
+ proj_q: torch.Tensor,
+ proj_kv: torch.Tensor,
+ rope_cos: torch.Tensor,
+ rope_sin: torch.Tensor,
+ k_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Native equivalent of NVlabs ``cam_prep_func``.
+
+ The caller is responsible for the q/k RMSNorm, which must run through the
+ model's norm modules so a tensor-parallel shard still normalises over the
+ global channel width. This function only applies ReLU, the K scale, the
+ UCPE ray projection and RoPE.
+
+ Args:
+ q_normed: ``(B, N, H, D)`` camera Q, already RMSNorm'd.
+ k_normed: ``(B, N, H, D)`` camera K, already RMSNorm'd. K must also
+ include the temporal short convolution.
+ v_raw: ``(B, N, H, D)`` raw camera V.
+ proj_q: ``(B, N, 4, 4)`` Q ray projection matrices.
+ proj_kv: ``(B, N, 4, 4)`` K/V ray projection matrices.
+ rope_cos: ``(N, D//2)`` interleaved-pair RoPE cosine table.
+ rope_sin: ``(N, D//2)`` interleaved-pair RoPE sine table.
+ k_scale: ``D^-0.5 * spatial_tokens^-0.5``.
+
+ Returns:
+ ``q_trans``, ``k_trans``, ``v_trans`` in ``(B, H, D, N)`` layout, and
+ ``inflation_sq`` in ``(B, H, N)``.
+ """
+ half_dim = _validate_cam_prep_inputs(
+ q_normed,
+ k_normed,
+ v_raw,
+ proj_q,
+ proj_kv,
+ rope_cos,
+ rope_sin,
+ )
+
+ q_norm = torch.relu(q_normed.float())
+ k_norm = torch.relu(k_normed.float()) * float(k_scale)
+ value = v_raw.float()
+
+ q_half = _apply_ray_projection(q_norm[..., :half_dim], proj_q)
+ k_half = _apply_ray_projection(k_norm[..., :half_dim], proj_kv)
+ v_half = _apply_ray_projection(value[..., :half_dim], proj_kv)
+
+ q_rope = _apply_interleaved_rope(q_norm[..., half_dim:], rope_cos.float(), rope_sin.float())
+ k_rope = _apply_interleaved_rope(k_norm[..., half_dim:], rope_cos.float(), rope_sin.float())
+ v_rope = _apply_interleaved_rope(value[..., half_dim:], rope_cos.float(), rope_sin.float())
+
+ k_pre_sq = k_norm.square().sum(dim=-1).permute(0, 2, 1).contiguous()
+ k_post = torch.cat((k_half, k_rope), dim=-1)
+ k_post_sq = k_post.square().sum(dim=-1).permute(0, 2, 1).contiguous()
+ inflation_sq = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12)
+
+ def to_bhdn(tensor: torch.Tensor) -> torch.Tensor:
+ return tensor.permute(0, 2, 3, 1).contiguous().to(v_raw.dtype)
+
+ q_out = to_bhdn(torch.cat((q_half, q_rope), dim=-1))
+ k_out = to_bhdn(k_post)
+ v_out = to_bhdn(torch.cat((v_half, v_rope), dim=-1))
+ return q_out, k_out, v_out, inflation_sq
+
+
+__all__ = [
+ "SanaWmCamGeometry",
+ "cam_prep_func",
+ "prepare_cam_geometry",
+]
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index 47063bc4a..b248e6690 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -171,6 +171,11 @@ _DIFFUSION_MODELS = {
"pipeline_ming_imagegen",
"MingImagePipeline",
),
+ "SanaWmPipeline": (
+ "sana_wm",
+ "pipeline_sana_wm",
+ "SanaWmPipeline",
+ ),
"InternVLAA1Pipeline": (
"internvla_a1",
"pipeline_internvla_a1",
@@ -633,6 +638,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
"LingBotVideoPipeline": "get_lingbot_video_pre_process_func",
"HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func",
"MagiHumanPipeline": "get_magi_human_pre_process_func",
+ "SanaWmPipeline": "get_sana_wm_pre_process_func",
"Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
"Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",
"SoulXSingerPipeline": "get_soulxsinger_pre_process_func",