[Model] Add SANA-WM support(Stage-1 only) (#4061)

Signed-off-by: BruceLoveDecimal <liuqihao970610@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: liuqihao <liuqihao970610@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

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` | ✅︎ |   |   |   | — |

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**                  |     ❌     |     ❌      |          ❌<sup>5</sup> |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+
+> 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
 

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: <https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional>
+- 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 `<keys>-<duration>` 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.

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 |

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)

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)

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}))

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)

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-<n>" 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)

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)

tests/model_tests/diffusion/test_alignment.py

@@ -80,6 +80,7 @@ EXCLUDED_MODELS = [
     "Pi0Pipeline",
     "SoulXSingerPipeline",
     "SoulXSingerSVCPipeline",
+    "SanaWmPipeline",
 ]
 
 

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,
+    ),
 }
 
 

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",
+]

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>-<duration>'.")
+        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,
+    )

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 {})

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)

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

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",
+]

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",