From 50b8a150f9df5ec17ff3a1a99ab113eb57d786d8 Mon Sep 17 00:00:00 2001 From: leo <92263211+liuyao0322@users.noreply.github.com> Date: Tue, 25 Aug 2026 09:46:02 +0800 Subject: [PATCH] [Model] Add native SANA-Video 2B T2V and I2V support (#5508) Signed-off-by: leo <92263211+napleon-liu@users.noreply.github.com> Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com> Signed-off-by: liuyao0322 Co-authored-by: leo <92263211+napleon-liu@users.noreply.github.com> --- docs/models/supported_models.md | 2 + docs/user_guide/diffusion_features.md | 5 +- .../offline_inference/image_to_video.md | 46 +- .../offline_inference/text_to_video.md | 27 + .../image_to_video/README.md | 46 +- .../image_to_video/image_to_video.py | 29 +- .../text_to_video/text_to_video.md | 36 + .../text_to_video/text_to_video.py | 20 + .../online_serving/image_to_video/README.md | 14 + .../image_to_video/run_curl_sana_video.sh | 67 ++ .../image_to_video/run_server_sana_video.sh | 13 + .../run_server_sana_video_diffusers.sh | 18 + .../online_serving/text_to_video/README.md | 5 + .../text_to_video/run_curl_sana_video.sh | 16 + .../text_to_video/run_server_sana_video.sh | 11 + .../run_server_sana_video_diffusers.sh | 12 + recipes/NVIDIA/SANA-Video-2B.md | 245 ++++ recipes/README.md | 2 + .../sana_video/test_pipeline_sana_video.py | 992 +++++++++++++++ .../sana_video/test_transformer_sana_video.py | 602 ++++++++++ tests/e2e/accuracy/sana_video/README.md | 63 + .../accuracy/sana_video/generate_goldens.py | 380 ++++++ .../sana_video/test_sana_video_golden.py | 328 +++++ .../test_sana_video_i2v_alignment.py | 170 +++ .../test_sana_video_pipeline_alignment.py | 432 +++++++ .../test_sana_video_expansion.py | 126 ++ .../test_text_to_video_prompts.py | 3 + .../diffusion/diff_model_builders.py | 77 +- tests/model_tests/diffusion/model_settings.py | 6 + tests/model_tests/diffusion/test_alignment.py | 6 + vllm_omni/diffusion/io_support.py | 6 +- .../pipeline_diffusers_adapter.py | 36 +- .../diffusers_adapter/pipeline_utils.py | 84 +- .../diffusion/models/sana_video/__init__.py | 32 + .../models/sana_video/pipeline_output.py | 19 + .../models/sana_video/pipeline_sana_video.py | 1055 ++++++++++++++++ .../sana_video/pipeline_sana_video_i2v.py | 382 ++++++ .../sana_video/transformer_sana_video.py | 1066 +++++++++++++++++ vllm_omni/diffusion/registry.py | 13 + vllm_omni/diffusion/request.py | 21 +- vllm_omni/model_extras/registry.py | 7 + vllm_omni/model_extras/sana_video.py | 10 + 42 files changed, 6512 insertions(+), 18 deletions(-) create mode 100755 examples/online_serving/image_to_video/run_curl_sana_video.sh create mode 100755 examples/online_serving/image_to_video/run_server_sana_video.sh create mode 100755 examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh create mode 100755 examples/online_serving/text_to_video/run_curl_sana_video.sh create mode 100755 examples/online_serving/text_to_video/run_server_sana_video.sh create mode 100755 examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh create mode 100644 recipes/NVIDIA/SANA-Video-2B.md create mode 100644 tests/diffusion/models/sana_video/test_pipeline_sana_video.py create mode 100644 tests/diffusion/models/sana_video/test_transformer_sana_video.py create mode 100644 tests/e2e/accuracy/sana_video/README.md create mode 100644 tests/e2e/accuracy/sana_video/generate_goldens.py create mode 100644 tests/e2e/accuracy/sana_video/test_sana_video_golden.py create mode 100644 tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py create mode 100644 tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py create mode 100644 tests/e2e/online_serving/test_sana_video_expansion.py create mode 100644 vllm_omni/diffusion/models/sana_video/__init__.py create mode 100644 vllm_omni/diffusion/models/sana_video/pipeline_output.py create mode 100644 vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py create mode 100644 vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py create mode 100644 vllm_omni/diffusion/models/sana_video/transformer_sana_video.py create mode 100644 vllm_omni/model_extras/sana_video.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f1d18ea48..b9724e8eb 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -52,6 +52,8 @@ th { | `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) | | `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) | | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — | +| `SanaVideoPipeline` | SANA-Video-2B T2V (native 480p/720p; Diffusers adapter validated at both resolutions) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) | +| `SanaImageToVideoPipeline` | SANA-Video-2B I2V (native and Diffusers adapter validated at 480p/720p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) | | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 538c3ee71..73862febe 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -162,6 +162,7 @@ The following tables show which models support each feature: | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ | +| **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | @@ -182,8 +183,8 @@ 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, SANA-WM +- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines; SANA-Video-2B native text-to-video and image-to-video pipelines +- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni, SANA-WM; SANA-Video-2B Diffusers-adapter text-to-video and image-to-video pipelines ### AudioGen diff --git a/docs/user_guide/examples/offline_inference/image_to_video.md b/docs/user_guide/examples/offline_inference/image_to_video.md index db7a31ce5..ff57063b8 100644 --- a/docs/user_guide/examples/offline_inference/image_to_video.md +++ b/docs/user_guide/examples/offline_inference/image_to_video.md @@ -4,8 +4,8 @@ Source `. + ## Key Arguments ### Common diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index 050448571..02f233642 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -1,7 +1,7 @@ # Image-To-Video This shared example generates videos from images with VACE, Wan2.2, LTX-2, -HunyuanVideo-1.5, Cosmos3, and other compatible pipelines. +HunyuanVideo-1.5, SANA-Video, Cosmos3, and other compatible pipelines. - `image_to_video.py`: command-line script for single video generation with advanced options. @@ -27,6 +27,8 @@ This folder provides a unified CLI script for image-to-video generation using vL | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card | | `Lightricks/LTX-2` | 512 x 768 | 121 | 40 | video 3.0 / audio 7.0 | Memory use depends on frame count and tensor parallelism | +| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480 x 832 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; Wan VAE | +| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704 x 1280 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; LTX-2 Video VAE | !!! info Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading. @@ -227,6 +229,48 @@ python image_to_video.py \ See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints, pipeline selection, T2V, defaults, and advanced options. +### SANA-Video-2B + +The checkpoint metadata selects the upstream T2V pipeline, so pass +`SanaImageToVideoPipeline` explicitly for I2V: + +```bash +python image_to_video.py \ + --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --model-class-name SanaImageToVideoPipeline \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms sway in the breeze as petals drift past the camera." \ + --negative-prompt "blurry, low quality, temporal artifacts" \ + --height 480 \ + --width 832 \ + --num-frames 81 \ + --num-inference-steps 50 \ + --guidance-scale 6.0 \ + --extra-body '{"motion_score": 30}' \ + --fps 16 \ + --seed 42 \ + --output sana_video_i2v_480p.mp4 +``` + +Use the 720p checkpoint with `--height 704 --width 1280`. In this example, +81 frames at 16 FPS means the standard approximately five-second SANA-Video +checkpoint request. It does not enable the separate LongSANA/LongLive +minute-scale autoregressive workflow. + +The native `SanaImageToVideoPipeline` supports both checkpoint variants. The +480p path uses vLLM-Omni's `DistributedAutoencoderKLWan` wrapper and the 720p +path uses `DistributedAutoencoderKLLTX2Video`; both wrap their corresponding +Diffusers VAE implementations. The native denoising loop intentionally keeps +Diffusers' checkpoint-compatible `DPMSolverMultistepScheduler`. + +The Diffusers adapter has completed end-to-end I2V serving validation at both +480p and 720p. For online adapter serving, select the matching checkpoint and +use `WIDTH=1280 HEIGHT=704` for the 720p request. Adapter T2V is also validated +at both resolutions; see the recipe for the exact backend matrix. + +See the [SANA-Video recipe](../../../recipes/NVIDIA/SANA-Video-2B.md) for +online serving, native/adapter validation boundaries, and hardware notes. + ## Advanced Features ### CFG Parallel diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 5635ce16c..1a7e40b12 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -3,13 +3,14 @@ """ Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3, -HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE. +HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE. Supports: - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder - Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B) - LTX2 image-to-video pipeline - HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning +- SANA-Video 2B: first-frame latent conditioning at 480p or 720p - Wan2.1 VACE: first/last-frame, inpainting, and reference conditioning Usage: @@ -38,6 +39,13 @@ Usage: --image input.jpg --prompt "A cat playing with yarn" \ --flow-shift 5.0 --guidance-scale 6.0 + # SANA-Video 2B I2V (480p) + python image_to_video.py --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --model-class-name SanaImageToVideoPipeline \ + --image input.jpg --prompt "A cat turns toward the camera." \ + --height 480 --width 832 --num-frames 81 --num-inference-steps 50 \ + --guidance-scale 6.0 + # Cosmos3 I2V (image conditioning) python image_to_video.py --model nvidia/Cosmos3-Nano \ --image input.jpg --prompt "The scene comes to life with smooth, natural motion." \ @@ -107,13 +115,16 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( "Generate a video from one or more images " - "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE)." + "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)." ) ) parser.add_argument( "--model", default="Wan-AI/Wan2.2-I2V-A14B-Diffusers", - help="Diffusers I2V model ID or local path (Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE).", + help=( + "Diffusers I2V model ID or local path " + "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)." + ), ) parser.add_argument( "--model-class-name", @@ -377,6 +388,7 @@ def main(): is_ltx2_distilled = "distilled" in model_class_name_lower or "distilled" in model_name is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name + is_sana = model_class_name == "SanaImageToVideoPipeline" or "sana-video" in model_name is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower()) is_cosmos_edge = is_cosmos and ("edge" in model_name or "edge" in model_class_name_lower) @@ -434,6 +446,17 @@ def main(): 512 * 768, 32, ) + elif is_sana: + is_720p = "720p" in model_name + d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = ( + 16, + 6.0, + 81, + 50, + 5.0, + (704 * 1280) if is_720p else (480 * 832), + 32 if is_720p else 16, + ) else: # Wan2.2 / HunyuanVideo-1.5 d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16 diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index 339c7dac9..0cb75a8d4 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -14,6 +14,8 @@ A unified script for text-to-video generation. Supports multiple models with mod | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) | | `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — | +| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE wrapper | +| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | BF16 DiT + LTX-2 Video VAE wrapper | ## Local CLI Usage @@ -151,6 +153,40 @@ python text_to_video.py \ --output quick_test.mp4 ``` +### SANA-Video-2B + +Select the native T2V pipeline explicitly: + +```bash +python text_to_video.py \ + --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --model-class-name SanaVideoPipeline \ + --prompt "A cinematic tracking shot of a sailboat crossing the ocean at sunset." \ + --height 480 \ + --width 832 \ + --num-frames 81 \ + --num-inference-steps 50 \ + --guidance-scale 6.0 \ + --extra-body '{"motion_score": 30}' \ + --fps 16 \ + --output sana_video_480p.mp4 +``` + +For the 720p checkpoint, switch the model to +`Efficient-Large-Model/SANA-Video_2B_720p_diffusers`, use +`--height 704 --width 1280`, and write to a different output path. Both native +checkpoint variants are supported. The Diffusers adapter T2V compatibility +path is also validated at both resolutions through online serving. + +The native 480p path uses vLLM-Omni's `DistributedAutoencoderKLWan` wrapper; +the native 720p path uses `DistributedAutoencoderKLLTX2Video`. Both wrappers +retain their underlying Diffusers autoencoders. The pipeline intentionally +uses Diffusers' checkpoint-compatible `DPMSolverMultistepScheduler`. + +For image-to-video, use the shared +[`image_to_video.py`](../image_to_video/README.md#sana-video-2b) example with +`--model-class-name SanaImageToVideoPipeline` and `--image `. + ### Cosmos3 ```bash diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 1bc2d5b11..ce46fba49 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -114,12 +114,32 @@ _MODEL_PRESETS = { "fps": 24, "output": "ltx23_output.mp4", }, + "sana_480p": { + "height": 480, + "width": 832, + "num_frames": 81, + "num_inference_steps": 50, + "guidance_scale": 6.0, + "fps": 16, + "output": "sana_video_480p.mp4", + }, + "sana_720p": { + "height": 704, + "width": 1280, + "num_frames": 81, + "num_inference_steps": 50, + "guidance_scale": 6.0, + "fps": 16, + "output": "sana_video_720p.mp4", + }, } def _detect_preset(model: str, model_class_name: str | None = None) -> dict: model_lower = model.lower() class_lower = (model_class_name or "").lower() + if "sana-video" in model_lower or "sana_video" in model_lower or "sanavideo" in class_lower: + return _MODEL_PRESETS["sana_720p" if "720p" in model_lower else "sana_480p"] if "lingbot" in model_lower or "lingbotvideo" in class_lower: return _MODEL_PRESETS["lingbot"] if "ltx" in class_lower or "ltx" in model_lower: diff --git a/examples/online_serving/image_to_video/README.md b/examples/online_serving/image_to_video/README.md index e6d9b070f..7fac613df 100644 --- a/examples/online_serving/image_to_video/README.md +++ b/examples/online_serving/image_to_video/README.md @@ -1,5 +1,19 @@ # Image-To-Video +This example demonstrates how to deploy image-to-video models for online +video generation using vLLM-Omni. + +## Supported Models + +| Model | Model ID | +|-------|----------| +| Wan2.2 I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | +| SANA-Video 2B (480p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | +| SANA-Video 2B (720p) | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | + +For SANA-Video 2B image-to-video online serving, see the +[SANA-Video 2B recipe](../../../recipes/NVIDIA/SANA-Video-2B.md). + This example demonstrates online image-to-video generation with vLLM-Omni. The startup script defaults to Wan2.2, but `MODEL` can select any supported image-to-video model. The existing curl helper remains a concrete Wan2.2 diff --git a/examples/online_serving/image_to_video/run_curl_sana_video.sh b/examples/online_serving/image_to_video/run_curl_sana_video.sh new file mode 100755 index 000000000..35806335f --- /dev/null +++ b/examples/online_serving/image_to_video/run_curl_sana_video.sh @@ -0,0 +1,67 @@ +#!/bin/bash +# Submit and download a standard SANA-Video-2B image-to-video request. + +set -euo pipefail + +INPUT_IMAGE="${INPUT_IMAGE:-}" +BASE_URL="${BASE_URL:-http://localhost:8099}" +OUTPUT_PATH="${OUTPUT_PATH:-sana_video_i2v.mp4}" +WIDTH="${WIDTH:-832}" +HEIGHT="${HEIGHT:-480}" +POLL_INTERVAL="${POLL_INTERVAL:-2}" + +if [ -z "$INPUT_IMAGE" ] || [ ! -f "$INPUT_IMAGE" ]; then + echo "Set INPUT_IMAGE to an existing image file." + exit 1 +fi + +create_response="$( + curl -sS -X POST "${BASE_URL}/v1/videos" \ + -H "Accept: application/json" \ + -F "prompt=A cat turns toward the camera with smooth, natural motion." \ + -F "negative_prompt=blurry, low quality, temporal artifacts" \ + -F "input_reference=@${INPUT_IMAGE}" \ + -F "width=${WIDTH}" \ + -F "height=${HEIGHT}" \ + -F "num_frames=81" \ + -F "fps=16" \ + -F "num_inference_steps=50" \ + -F "guidance_scale=6.0" \ + -F "seed=42" \ + -F 'extra_params={"motion_score":30}' +)" + +video_id="$(echo "${create_response}" | jq -r '.id')" +if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then + echo "Failed to create SANA-Video job:" + echo "${create_response}" | jq . + exit 1 +fi + +while true; do + status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")" + status="$(echo "${status_response}" | jq -r '.status')" + + case "${status}" in + queued|in_progress) + echo "Video job ${video_id} status: ${status}" + sleep "${POLL_INTERVAL}" + ;; + completed) + break + ;; + failed) + echo "SANA-Video generation failed:" + echo "${status_response}" | jq . + exit 1 + ;; + *) + echo "Unexpected status response:" + echo "${status_response}" | jq . + exit 1 + ;; + esac +done + +curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}" +echo "Saved SANA-Video to ${OUTPUT_PATH}" diff --git a/examples/online_serving/image_to_video/run_server_sana_video.sh b/examples/online_serving/image_to_video/run_server_sana_video.sh new file mode 100755 index 000000000..257cb2b00 --- /dev/null +++ b/examples/online_serving/image_to_video/run_server_sana_video.sh @@ -0,0 +1,13 @@ +#!/bin/bash +# SANA-Video-2B native image-to-video serving. + +set -euo pipefail + +MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}" +PORT="${PORT:-8099}" + +vllm serve "$MODEL" \ + --omni \ + --model-class-name SanaImageToVideoPipeline \ + --dtype bfloat16 \ + --port "$PORT" diff --git a/examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh b/examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh new file mode 100755 index 000000000..bf0573529 --- /dev/null +++ b/examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# SANA-Video-2B image-to-video serving through the Diffusers adapter. +# Both 480p and 720p checkpoints are validated. For 720p: +# MODEL=Efficient-Large-Model/SANA-Video_2B_720p_diffusers bash run_server_sana_video_diffusers.sh +# INPUT_IMAGE=input.jpg WIDTH=1280 HEIGHT=704 bash run_curl_sana_video.sh + +set -euo pipefail + +MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}" +PORT="${PORT:-8099}" + +vllm serve "$MODEL" \ + --omni \ + --model-class-name SanaImageToVideoPipeline \ + --diffusion-load-format diffusers \ + --diffusion-attention-backend TORCH_SDPA \ + --dtype bfloat16 \ + --port "$PORT" diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md index 51512fe0c..cf303ccb1 100644 --- a/examples/online_serving/text_to_video/README.md +++ b/examples/online_serving/text_to_video/README.md @@ -11,6 +11,11 @@ This example demonstrates how to deploy text-to-video models for online video ge | Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | | LTX-2 | `Lightricks/LTX-2` | | Helios (Base / Mid / Distilled) | `BestWishYsh/Helios-Base`, `Helios-Mid`, `Helios-Distilled` | +| SANA-Video 2B (480p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | +| SANA-Video 2B (720p) | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | + +For SANA-Video 2B text-to-video online serving, see the +[SANA-Video 2B recipe](../../../recipes/NVIDIA/SANA-Video-2B.md). ## Wan2.2 T2V diff --git a/examples/online_serving/text_to_video/run_curl_sana_video.sh b/examples/online_serving/text_to_video/run_curl_sana_video.sh new file mode 100755 index 000000000..490a743e5 --- /dev/null +++ b/examples/online_serving/text_to_video/run_curl_sana_video.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +BASE_URL="${BASE_URL:-http://localhost:8091}" +OUTPUT_PATH="${OUTPUT_PATH:-sana_video.mp4}" + +curl -sS -X POST "$BASE_URL/v1/videos/sync" \ + -F "prompt=A cinematic tracking shot of a sailboat crossing the ocean at sunset" \ + -F "height=480" \ + -F "width=832" \ + -F "num_frames=81" \ + -F "fps=16" \ + -F "num_inference_steps=50" \ + -F "guidance_scale=6.0" \ + -F "seed=42" \ + -F 'extra_params={"motion_score":30}' \ + --output "$OUTPUT_PATH" diff --git a/examples/online_serving/text_to_video/run_server_sana_video.sh b/examples/online_serving/text_to_video/run_server_sana_video.sh new file mode 100755 index 000000000..446a75239 --- /dev/null +++ b/examples/online_serving/text_to_video/run_server_sana_video.sh @@ -0,0 +1,11 @@ +#!/bin/bash +# SANA-Video-2B text-to-video serving. + +MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}" +PORT="${PORT:-8091}" + +vllm serve "$MODEL" \ + --omni \ + --model-class-name SanaVideoPipeline \ + --dtype bfloat16 \ + --port "$PORT" diff --git a/examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh b/examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh new file mode 100755 index 000000000..adfb125c0 --- /dev/null +++ b/examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# SANA-Video through the black-box Diffusers backend adapter. + +MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}" +PORT="${PORT:-8091}" + +vllm serve "$MODEL" \ + --omni \ + --diffusion-load-format diffusers \ + --diffusion-attention-backend TORCH_SDPA \ + --dtype bfloat16 \ + --port "$PORT" diff --git a/recipes/NVIDIA/SANA-Video-2B.md b/recipes/NVIDIA/SANA-Video-2B.md new file mode 100644 index 000000000..c85d8ca40 --- /dev/null +++ b/recipes/NVIDIA/SANA-Video-2B.md @@ -0,0 +1,245 @@ +# SANA-Video 2B + +> Native text-to-video and image-to-video generation at 480p and 720p, +> plus a validated Diffusers-adapter compatibility path + +## Summary + +- Vendor: NVIDIA +- Model: `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, + `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` +- Task: Text-to-video and image-to-video +- Mode: Offline inference and OpenAI-compatible online serving +- Maintainer: Community + +## When to use this recipe + +Use the native `SanaVideoPipeline` for T2V and +`SanaImageToVideoPipeline` for I2V. Both native pipelines support the 480p +and 720p checkpoints. Use `--diffusion-load-format diffusers` when you need +the black-box Diffusers compatibility baseline; adapter T2V and I2V are +validated at both resolutions. + +The native pipeline loads the 480p checkpoint through +`DistributedAutoencoderKLWan` and the 720p checkpoint through +`DistributedAutoencoderKLLTX2Video`. These are vLLM-Omni distributed wrappers +around the corresponding Diffusers autoencoders, not independent VAE +implementations. The denoising loop also intentionally loads Diffusers' +`DPMSolverMultistepScheduler` from the checkpoint to preserve its scheduler +configuration and numerical behavior. + +## References + +- Upstream project: +- Model cards: + [480p](https://huggingface.co/Efficient-Large-Model/SANA-Video_2B_480p_diffusers), + [720p](https://huggingface.co/Efficient-Large-Model/SANA-Video_2B_720p_diffusers) +- Diffusers documentation: +- Online serving guides: + [Text-to-Video](../../docs/user_guide/examples/online_serving/text_to_video.md), + [Image-to-Video](../../docs/user_guide/examples/online_serving/image_to_video.md) +- Offline T2V example: + [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py) +- Offline I2V example: + [`examples/offline_inference/image_to_video/image_to_video.py`](../../examples/offline_inference/image_to_video/image_to_video.py) +- Support discussion: [vLLM-Omni issue #5432](https://github.com/vllm-project/vllm-omni/issues/5432) + +## Hardware Support + +## GPU + +### 1x RTX 5090 32GB + +#### Environment + +- OS: Ubuntu 22.04.5 LTS +- Python: 3.12.3 +- Driver / runtime: NVIDIA driver 580.95.05; PyTorch 2.11.0+cu130 +- Diffusers: 0.38.0 +- vLLM version: 0.26.0 +- vLLM-Omni version or commit: PR #5508, commit `22037901` + +#### Command + +##### Native text-to-video inference + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ + --model Efficient-Large-Model/SANA-Video_2B_720p_diffusers \ + --model-class-name SanaVideoPipeline \ + --prompt "A cat walking on the grass, facing the camera." \ + --negative-prompt "blurry, low quality, temporal artifacts" \ + --height 704 --width 1280 --num-frames 81 \ + --num-inference-steps 50 --guidance-scale 6 \ + --extra-body '{"motion_score": 30}' \ + --fps 16 --seed 42 --output sana_video_720p.mp4 +``` + +For 480p, select `SANA-Video_2B_480p_diffusers` and use +`--height 480 --width 832`. + +##### Native image-to-video inference + +SANA checkpoints declare `SanaVideoPipeline` in `model_index.json`, so I2V +must be selected explicitly with `--model-class-name +SanaImageToVideoPipeline`. + +```bash +python examples/offline_inference/image_to_video/image_to_video.py \ + --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --model-class-name SanaImageToVideoPipeline \ + --image input.jpg \ + --prompt "A cat turns toward the camera with smooth, natural motion." \ + --negative-prompt "blurry, low quality, temporal artifacts" \ + --height 480 --width 832 --num-frames 81 \ + --num-inference-steps 50 --guidance-scale 6 \ + --extra-body '{"motion_score": 30}' \ + --fps 16 --seed 42 --output sana_video_i2v_480p.mp4 +``` + +The same pipeline class supports the 720p checkpoint through vLLM-Omni's +distributed LTX-2 VAE wrapper; use `--height 704 --width 1280`. + +For online I2V serving: + +```bash +MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + bash examples/online_serving/image_to_video/run_server_sana_video.sh + +INPUT_IMAGE=input.jpg OUTPUT_PATH=sana_video_i2v.mp4 \ + bash examples/online_serving/image_to_video/run_curl_sana_video.sh +``` + +##### Native online serving + +```bash +MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + bash examples/online_serving/text_to_video/run_server_sana_video.sh + +bash examples/online_serving/text_to_video/run_curl_sana_video.sh +``` + +To run the black-box compatibility backend for T2V, replace the server script +with `run_server_sana_video_diffusers.sh`. The same `/v1/videos` request +works; `num_frames` is adapted to Diffusers' `frames` argument. The script +selects `TORCH_SDPA` because SANA-Video uses an attention mask that the +AITER-backed Diffusers attention path does not accept. + +##### Diffusers-adapter image-to-video serving + +The validated I2V adapter commands are: + +```bash +# 480p +MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + bash examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh + +INPUT_IMAGE=input.jpg OUTPUT_PATH=sana_video_i2v_adapter.mp4 \ + bash examples/online_serving/image_to_video/run_curl_sana_video.sh + +# 720p +MODEL=Efficient-Large-Model/SANA-Video_2B_720p_diffusers \ + bash examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh + +INPUT_IMAGE=input.jpg WIDTH=1280 HEIGHT=704 \ + OUTPUT_PATH=sana_video_i2v_adapter_720p.mp4 \ + bash examples/online_serving/image_to_video/run_curl_sana_video.sh +``` + +#### Verification + +Check the encoded 720p output metadata after running a generation command: + +```bash +ffprobe -v error -select_streams v:0 \ + -show_entries stream=width,height,r_frame_rate,nb_frames \ + -of default=noprint_wrappers=1 sana_video_720p.mp4 +``` + +The standard 720p request above should report: + +```text +width=1280 +height=704 +r_frame_rate=16/1 +nb_frames=81 +``` + +For 480p, expect `width=832` and `height=480`. + +The automated serving matrix covers both checkpoint variants: + +| Backend | 480p T2V | 720p T2V | 480p I2V | 720p I2V | +| --- | --- | --- | --- | --- | +| Native vLLM-Omni | Validated | Validated | Validated | Validated | +| Diffusers adapter | Validated | Validated | Validated | Validated | + +Use the native `SanaVideoPipeline` and `SanaImageToVideoPipeline` for the +primary SANA execution paths. The Diffusers adapter is retained as a +validated compatibility/reference backend. + +#### Notes + +- Key flags: select I2V explicitly with `--model-class-name + SanaImageToVideoPipeline`; SANA checkpoint `model_index.json` files declare + the T2V class. Pass `motion_score` through `--extra-body`. The supplied + Diffusers-adapter server scripts select `TORCH_SDPA` because the AITER-backed + path does not accept SANA's attention mask. +- Output profile: 81 frames at 16 FPS is the standard checkpoint profile and + produces approximately five seconds of video. Minute-scale generation + requires the separate LongSANA/LongLive block-autoregressive workflow. +- Backend boundary: the native pipelines and Transformer are owned by + vLLM-Omni. The 480p Wan VAE and 720p LTX-2 VAE run through vLLM-Omni + distributed wrappers derived from the corresponding Diffusers VAE classes. + The denoising loop intentionally retains the checkpoint-compatible + Diffusers `DPMSolverMultistepScheduler`. +- Known limitations: + - Sequence/tensor/CFG parallelism, Cache-DiT, TeaCache, and step execution + are not validated for the native pipeline. + - The Diffusers backend is a compatibility path and does not provide native + vLLM-Omni parallelism or continuous batching. + - Native describes pipeline and Transformer ownership, not a zero-Diffusers + dependency guarantee. + +### 1x NVIDIA L20X + +#### Environment + +- Driver: 570.133.20 +- PyTorch: 2.13.0+cu130 +- Diffusers: 0.38.0 +- vLLM: 0.27.0 +- Nsight Systems: 2026.1 +- vLLM-Omni commit: `54fcc6d` + +#### Same-workload T2V comparison + +Both backends used +`Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, the same T2V prompt +and seeds, 832×480 output, 81 frames, and four inference steps. One warmup +request was excluded, followed by three measured requests. + +| Metric | Native vLLM-Omni | Diffusers adapter | Native delta | +| --- | ---: | ---: | ---: | +| E2E latency | 7.1600 ± 0.0043 s | 6.9761 ± 0.0517 s | +2.636% | +| Throughput | 0.1397 req/s | 0.1433 req/s | -2.568% | +| Peak reserved VRAM | 23,936 MiB | 19,060 MiB | +25.582% | +| Initialization | 11.669 s | 20.236 s | -42.336% | + +The peak-memory result is not a Transformer-only comparison. For the 480p +checkpoint, the native pipeline deliberately loads and decodes with the Wan +VAE in FP32. This matches the upstream SANA-Video reference and the recorded +accuracy-golden setup for this integration. The adapter launch uses +`--dtype bfloat16` as the Diffusers pipeline-wide dtype, so its Wan VAE also +runs in BF16. + +Peak VRAM is the CUDA allocator's maximum reserved memory over the complete +request. It includes VAE decode activations and allocator-reserved blocks, not +only live Transformer tensors. The different VAE precision policies are a +material code-level difference and a likely contributor to the higher native +high-water mark; the available full-request measurements do not isolate how +much of the delta they account for. The 25.582% result should therefore not be +interpreted as an isolated native Transformer memory regression. Exact +attribution requires both paths to use the same VAE dtype. Changing the native +480p VAE to BF16 would instead change the recorded golden/reference setup and +requires a new accuracy study. diff --git a/recipes/README.md b/recipes/README.md index 1bf35745a..af26807ef 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -57,6 +57,8 @@ recipes/ | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU | | [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified | +| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom | +| [`NVIDIA/SANA-Video-2B.md`](./NVIDIA/SANA-Video-2B.md) | Native and Diffusers-adapter 480p/720p T2V+I2V | 1x RTX 5090 32GB | | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) | diff --git a/tests/diffusion/models/sana_video/test_pipeline_sana_video.py b/tests/diffusion/models/sana_video/test_pipeline_sana_video.py new file mode 100644 index 000000000..f6f469aff --- /dev/null +++ b/tests/diffusion/models/sana_video/test_pipeline_sana_video.py @@ -0,0 +1,992 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from contextlib import nullcontext +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image +from torch import nn + +from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _DummyComponent: + pass + + +def _make_request_batch(prompt, request_id="sana-video-test", **sampling_overrides): + sampling = OmniDiffusionSamplingParams(**sampling_overrides) + request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id=request_id) + return DiffusionRequestBatch([request]) + + +def test_sana_video_pipeline_import_and_registry(): + from vllm_omni.diffusion.models.sana_video import ( + SanaImageToVideoPipeline, + SanaVideoPipeline, + SanaVideoTransformer3DModel, + get_sana_video_i2v_post_process_func, + get_sana_video_i2v_pre_process_func, + get_sana_video_post_process_func, + ) + from vllm_omni.diffusion.registry import ( + _DIFFUSION_MODELS, + _DIFFUSION_POST_PROCESS_FUNCS, + _DIFFUSION_PRE_PROCESS_FUNCS, + ) + + assert SanaImageToVideoPipeline is not None + assert SanaVideoPipeline is not None + assert SanaVideoTransformer3DModel is not None + assert get_sana_video_i2v_post_process_func is not None + assert get_sana_video_i2v_pre_process_func is not None + assert get_sana_video_post_process_func is not None + assert _DIFFUSION_MODELS["SanaVideoPipeline"] == ( + "sana_video", + "pipeline_sana_video", + "SanaVideoPipeline", + ) + assert _DIFFUSION_POST_PROCESS_FUNCS["SanaVideoPipeline"] == "get_sana_video_post_process_func" + assert _DIFFUSION_MODELS["SanaImageToVideoPipeline"] == ( + "sana_video", + "pipeline_sana_video_i2v", + "SanaImageToVideoPipeline", + ) + assert _DIFFUSION_POST_PROCESS_FUNCS["SanaImageToVideoPipeline"] == "get_sana_video_i2v_post_process_func" + assert _DIFFUSION_PRE_PROCESS_FUNCS["SanaImageToVideoPipeline"] == "get_sana_video_i2v_pre_process_func" + + +def test_component_discovery_declarations(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + assert SanaVideoPipeline._dit_modules == ["transformer"] + assert SanaVideoPipeline._encoder_modules == ["text_encoder"] + assert SanaVideoPipeline._vae_modules == ["vae"] + assert SanaVideoPipeline.supports_step_execution is False + + +def test_sana_video_loads_concrete_gemma_tokenizer(monkeypatch): + from vllm_omni.diffusion.models.sana_video import pipeline_sana_video + + captured: dict[str, object] = {} + expected = object() + + def fake_from_pretrained_with_prefetch(loader, model, **kwargs): + captured.update(loader=loader, model=model, kwargs=kwargs) + return expected + + monkeypatch.setattr( + pipeline_sana_video, + "from_pretrained_with_prefetch", + fake_from_pretrained_with_prefetch, + ) + + result = pipeline_sana_video._load_sana_tokenizer( + "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + ["tokenizer", "text_encoder"], + local_files_only=False, + ) + + assert result is expected + assert captured["loader"] == pipeline_sana_video.GemmaTokenizer.from_pretrained + assert captured["model"] == "Efficient-Large-Model/SANA-Video_2B_480p_diffusers" + assert captured["kwargs"] == { + "subfolder": "tokenizer", + "prefetch_list": ["tokenizer", "text_encoder"], + "local_files_only": False, + } + + +def test_sana_video_declares_extra_body_params(): + from vllm_omni.model_extras import get_extra_body_params + + assert get_extra_body_params("SanaVideoPipeline") == { + "clean_caption", + "motion_score", + "use_resolution_binning", + } + assert get_extra_body_params("SanaImageToVideoPipeline") == { + "clean_caption", + "motion_score", + "use_resolution_binning", + } + + +def test_sana_video_i2v_preprocesses_image_and_preserves_aspect_ratio(): + from vllm_omni.diffusion.models.sana_video import get_sana_video_i2v_pre_process_func + + request = OmniDiffusionRequest( + prompt={"prompt": "a robot walks", "multi_modal_data": {"image": Image.new("RGB", (640, 360))}}, + sampling_params=OmniDiffusionSamplingParams(), + request_id="sana-video-i2v-preprocess", + ) + processed = get_sana_video_i2v_pre_process_func(SimpleNamespace())(request) + + assert processed.sampling_params.height == 448 + assert processed.sampling_params.width == 832 + assert processed.prompt["multi_modal_data"]["image"].size == (832, 448) + + +def test_sana_video_i2v_remote_720p_uses_loaded_transformer_config(): + from vllm_omni.diffusion.models.sana_video import get_sana_video_i2v_pre_process_func + + od_config = SimpleNamespace( + model="Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + tf_model_config={"sample_size": 22}, + ) + request = OmniDiffusionRequest( + prompt={"prompt": "a robot walks", "multi_modal_data": {"image": Image.new("RGB", (1280, 704))}}, + sampling_params=OmniDiffusionSamplingParams(), + request_id="sana-video-i2v-remote-720p", + ) + + processed = get_sana_video_i2v_pre_process_func(od_config)(request) + + assert processed.sampling_params.height == 704 + assert processed.sampling_params.width == 1280 + assert processed.prompt["multi_modal_data"]["image"].size == (1280, 704) + + +@pytest.mark.parametrize( + "frames", + [ + [[Image.new("RGB", (16, 16)) for _ in range(2)]], + np.zeros((1, 2, 16, 16, 3), dtype=np.float32), + [[np.zeros((16, 16, 3), dtype=np.float32) for _ in range(2)]], + ], +) +@pytest.mark.parametrize( + "factory_name", + ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"], +) +def test_sana_video_diffusers_preserves_postprocessed_frames(frames, factory_name): + from vllm_omni.diffusion.models import sana_video + + post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="diffusers")) + + result = post_process(frames) + + assert result["payload"]["video"] is frames + + +@pytest.mark.parametrize( + "factory_name", + ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"], +) +def test_sana_video_diffusers_still_postprocesses_native_tensor(factory_name): + from vllm_omni.diffusion.models import sana_video + + decoded_video = torch.zeros(1, 3, 2, 16, 16) + post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="diffusers")) + + result = post_process(decoded_video) + + assert isinstance(result["payload"]["video"], np.ndarray) + assert result["payload"]["video"].shape == (1, 2, 16, 16, 3) + + +@pytest.mark.parametrize( + "factory_name", + ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"], +) +def test_sana_video_postprocess_preserves_requested_latents(factory_name): + from vllm_omni.diffusion.models import sana_video + + latents = torch.randn(1, 4, 2, 2, 2) + post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="native")) + + result = post_process(latents, sampling_params=SimpleNamespace(output_type="latent")) + + assert result is latents + + +def test_sana_video_720p_model_id_fallback(monkeypatch): + from vllm_omni.diffusion.models.sana_video import pipeline_sana_video + + monkeypatch.setattr(pipeline_sana_video, "get_hf_file_to_dict", lambda *_args, **_kwargs: None) + od_config = SimpleNamespace( + model="Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + tf_model_config={}, + ) + + assert pipeline_sana_video.resolve_sana_video_sample_size(od_config) == 22 + + +@pytest.mark.parametrize( + ("requested_output_type", "expected_internal_output_type"), + [ + (None, "raw"), + ("np", "raw"), + ("latent", "latent"), + ], +) +def test_sana_video_i2v_forward_maps_image_request(requested_output_type, expected_internal_output_type): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate_i2v(**kwargs): + calls.append(kwargs) + return torch.zeros(kwargs["num_videos_per_prompt"], 3, 9, 192, 320) + + pipeline._generate_i2v = fake_generate_i2v + image = Image.new("RGB", (320, 192)) + custom_timesteps = torch.tensor([900, 100]) + req = _make_request_batch( + {"prompt": "a robot walks", "negative_prompt": "blurry", "multi_modal_data": {"image": image}}, + height=192, + width=320, + num_frames=9, + num_inference_steps=2, + timesteps=custom_timesteps, + guidance_scale=4.5, + num_outputs_per_prompt=2, + eta=0.25, + seed=42, + output_type=requested_output_type, + extra_args={"clean_caption": True, "motion_score": 30, "use_resolution_binning": False}, + ) + + output = pipeline.forward(req) + + assert output.output.shape == (2, 3, 9, 192, 320) + assert calls[0]["image"] is image + assert calls[0]["prompt"] == "a robot walks motion score: 30." + assert calls[0]["negative_prompt"] == "blurry" + assert calls[0]["frames"] == 9 + assert calls[0]["timesteps"] is custom_timesteps + assert calls[0]["sigmas"] is None + assert calls[0]["num_videos_per_prompt"] == 2 + assert calls[0]["eta"] == 0.25 + assert calls[0]["clean_caption"] is True + assert calls[0]["generator"].initial_seed() == 42 + assert calls[0]["output_type"] == expected_internal_output_type + + +def test_sana_video_i2v_forward_maps_custom_sigmas(): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate_i2v(**kwargs): + calls.append(kwargs) + return torch.zeros(1, 3, 9, 192, 320) + + pipeline._generate_i2v = fake_generate_i2v + sigmas = [1.0, 0.5] + output = pipeline.forward( + _make_request_batch( + { + "prompt": "a robot walks", + "multi_modal_data": {"image": Image.new("RGB", (320, 192))}, + }, + height=192, + width=320, + num_frames=9, + sigmas=sigmas, + generator_device="cpu", + extra_args={"use_resolution_binning": False}, + ) + ) + + assert output.output.shape == (1, 3, 9, 192, 320) + assert calls[0]["timesteps"] is None + assert calls[0]["sigmas"] is sigmas + + +def test_sana_video_i2v_omitted_num_frames_uses_model_default(): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate_i2v(**kwargs): + calls.append(kwargs) + return torch.zeros(1, 3, 81, 2, 2) + + pipeline._generate_i2v = fake_generate_i2v + req = _make_request_batch( + { + "prompt": "a robot walks", + "multi_modal_data": {"image": Image.new("RGB", (320, 192))}, + }, + height=192, + width=320, + num_inference_steps=1, + generator_device="cpu", + extra_args={"use_resolution_binning": False}, + ) + + assert req.sampling_params.num_frames == 1 + output = pipeline.forward(req) + + assert output.output.shape == (1, 3, 81, 2, 2) + assert calls[0]["frames"] == 81 + + +def test_sana_video_i2v_clean_caption_request_is_explicitly_rejected(): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + pipeline.text_encoder = SimpleNamespace(dtype=torch.float32) + pipeline.tokenizer = None + pipeline.check_inputs = lambda **_kwargs: None + + req = _make_request_batch( + { + "prompt": "a robot walks", + "multi_modal_data": {"image": Image.new("RGB", (320, 192))}, + }, + height=192, + width=320, + num_frames=9, + generator_device="cpu", + extra_args={"clean_caption": True, "use_resolution_binning": False}, + ) + + with pytest.raises(ValueError, match="does not support `clean_caption=True`"): + pipeline.forward(req) + + +@pytest.mark.parametrize( + ("custom_timesteps", "custom_sigmas"), + [ + (torch.tensor([4.0, 2.0]), None), + (None, [1.0, 0.5]), + ], +) +def test_sana_video_i2v_consumes_batch_schedule_and_step_kwargs(custom_timesteps, custom_sigmas): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + calls = SimpleNamespace( + encode_prompt=[], + prepare_latents=[], + set_timesteps=[], + step=[], + transformer=[], + ) + + class StubTransformer: + dtype = torch.float32 + config = SimpleNamespace( + sample_size=30, + in_channels=4, + out_channels=4, + patch_size=(1, 1, 1), + ) + + def __call__( + self, + hidden_states, + *, + encoder_hidden_states, + encoder_attention_mask, + timestep, + return_dict, + ): + calls.transformer.append( + ( + hidden_states.shape, + encoder_hidden_states.shape, + encoder_attention_mask.shape, + timestep.shape, + return_dict, + ) + ) + return (torch.zeros_like(hidden_states),) + + class StubScheduler: + order = 1 + + def set_timesteps( + self, + num_inference_steps=None, + *, + device=None, + timesteps=None, + sigmas=None, + ): + calls.set_timesteps.append( + { + "num_inference_steps": num_inference_steps, + "timesteps": timesteps, + "sigmas": sigmas, + } + ) + schedule_length = len(timesteps if timesteps is not None else sigmas) + self.timesteps = torch.arange(schedule_length, 0, -1, dtype=torch.float32, device=device) + + def step( + self, + _noise_pred, + _timestep, + current_latents, + *, + eta, + generator, + return_dict, + ): + calls.step.append({"eta": eta, "generator": generator, "return_dict": return_dict}) + return (current_latents + 1,) + + class VaeMustNotBeUsed: + @property + def dtype(self): + pytest.fail("latent output must return before accessing the VAE") + + def decode(self, *_args, **_kwargs): + pytest.fail("latent output must not call VAE decode") + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = StubTransformer() + pipeline.scheduler = StubScheduler() + pipeline.vae = VaeMustNotBeUsed() + pipeline.check_inputs = lambda **_kwargs: None + pipeline.video_processor = SimpleNamespace( + preprocess=lambda _image, height, width: torch.zeros(1, 3, height, width), + ) + pipeline.progress_bar = lambda **_kwargs: nullcontext(SimpleNamespace(update=lambda: None)) + + def capture_encode_prompt(*_args, **kwargs): + calls.encode_prompt.append(kwargs) + batch_size = kwargs["num_videos_per_prompt"] + return ( + torch.zeros(batch_size, 1, 1), + torch.ones(batch_size, 1, dtype=torch.bool), + torch.zeros(batch_size, 1, 1), + torch.ones(batch_size, 1, dtype=torch.bool), + ) + + pipeline.encode_prompt = capture_encode_prompt + initial_latents = torch.zeros(2, 4, 2, 2, 2) + + def capture_prepare_latents(*_args, **kwargs): + calls.prepare_latents.append(kwargs) + return initial_latents + + pipeline._prepare_i2v_latents = capture_prepare_latents + generator = torch.Generator(device="cpu").manual_seed(42) + + output = pipeline._generate_i2v( + image=Image.new("RGB", (2, 2)), + prompt="test", + negative_prompt="", + height=2, + width=2, + frames=2, + num_inference_steps=50, + timesteps=custom_timesteps, + sigmas=custom_sigmas, + guidance_scale=4.5, + num_videos_per_prompt=2, + eta=0.25, + generator=generator, + latents=None, + clean_caption=True, + use_resolution_binning=False, + max_sequence_length=1, + output_type="latent", + ) + + assert calls.encode_prompt[0]["num_videos_per_prompt"] == 2 + assert calls.encode_prompt[0]["clean_caption"] is True + assert calls.prepare_latents[0]["batch_size"] == 2 + assert len(calls.set_timesteps) == 1 + assert calls.set_timesteps[0]["num_inference_steps"] is None + assert calls.set_timesteps[0]["timesteps"] is custom_timesteps + assert calls.set_timesteps[0]["sigmas"] is custom_sigmas + assert len(calls.step) == 2 + assert all(call == {"eta": 0.25, "generator": generator, "return_dict": False} for call in calls.step) + assert all( + call + == ( + torch.Size([4, 4, 2, 2, 2]), + torch.Size([4, 1, 1]), + torch.Size([4, 1]), + torch.Size([4, 1, 2, 2, 2]), + False, + ) + for call in calls.transformer + ) + expected = initial_latents.clone() + expected[:, :, 1:] += 2 + torch.testing.assert_close(output, expected) + + +def test_sana_video_i2v_uses_diffusers_complex_instruction_default(): + import inspect + + from diffusers import SanaImageToVideoPipeline as DiffusersSanaImageToVideoPipeline + + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + reference_default = ( + inspect.signature(DiffusersSanaImageToVideoPipeline.__call__).parameters["complex_human_instruction"].default + ) + native_default = ( + inspect.signature(SanaImageToVideoPipeline._generate_i2v).parameters["complex_human_instruction"].default + ) + assert native_default == reference_default + + class StopAfterPromptEncodingError(Exception): + pass + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.check_inputs = lambda **_kwargs: None + captured = {} + + def capture_encode_prompt(*_args, **kwargs): + captured.update(kwargs) + raise StopAfterPromptEncodingError + + pipeline.encode_prompt = capture_encode_prompt + with pytest.raises(StopAfterPromptEncodingError): + pipeline._generate_i2v( + image=Image.new("RGB", (320, 192)), + prompt="a robot walks", + negative_prompt="blurry", + height=192, + width=320, + frames=9, + num_inference_steps=2, + guidance_scale=4.5, + generator=torch.Generator(device="cpu").manual_seed(42), + latents=None, + use_resolution_binning=False, + max_sequence_length=300, + ) + + assert captured["complex_human_instruction"] == reference_default + + +def test_sana_video_t2v_complex_instruction_has_no_mutable_default(): + import inspect + + from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( + SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION, + SanaVideoPipeline, + ) + + parameter = inspect.signature(SanaVideoPipeline._generate).parameters["complex_human_instruction"] + assert parameter.default is None + assert isinstance(SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION, tuple) + + +def test_diffusers_adapter_t2v_omitted_num_frames_uses_sana_default(): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils + + adapter = object.__new__(DiffusersAdapterPipeline) + adapter._accept_call_kwargs = {"prompt", "negative_prompt", "frames", "generator"} + adapter._pipeline_utils = SanaVideoPipelineUtils() + adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None) + + req = _make_request_batch( + {"prompt": "a robot walks", "negative_prompt": "blurry"}, + seed=42, + generator_device="cpu", + ) + assert req.sampling_params.num_frames == 1 + + kwargs = adapter._build_call_kwargs(req) + + assert kwargs["frames"] == 81 + assert "num_frames" not in kwargs + assert kwargs["generator"].initial_seed() == 42 + + +def test_diffusers_adapter_i2v_omitted_num_frames_uses_sana_default(): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils + + adapter = object.__new__(DiffusersAdapterPipeline) + adapter._accept_call_kwargs = {"prompt", "negative_prompt", "image", "frames", "generator"} + adapter._pipeline_utils = SanaVideoPipelineUtils() + adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None) + image = Image.new("RGB", (320, 192)) + + req = _make_request_batch( + {"prompt": "a robot walks", "negative_prompt": "blurry", "multi_modal_data": {"image": image}}, + seed=42, + generator_device="cpu", + ) + assert req.sampling_params.num_frames == 1 + + kwargs = adapter._build_call_kwargs(req) + + assert kwargs["image"] is image + assert kwargs["prompt"] == "a robot walks" + assert kwargs["negative_prompt"] == "blurry" + assert kwargs["frames"] == 81 + assert kwargs["generator"].initial_seed() == 42 + + +@pytest.mark.parametrize( + ("sampling_overrides", "expected_frames"), + [ + ({}, 49), + ({"num_frames": 9}, 9), + ], +) +def test_diffusers_adapter_sana_resolves_frame_precedence(sampling_overrides, expected_frames): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils + + adapter = object.__new__(DiffusersAdapterPipeline) + adapter._accept_call_kwargs = {"prompt", "frames", "generator"} + adapter._pipeline_utils = SanaVideoPipelineUtils() + adapter.od_config = SimpleNamespace(diffusers_call_kwargs={"frames": 49}, output_type=None) + + kwargs = adapter._build_call_kwargs( + _make_request_batch( + "a robot walks", + seed=42, + generator_device="cpu", + **sampling_overrides, + ) + ) + + assert kwargs["frames"] == expected_frames + + +@pytest.mark.parametrize("pipeline_class_name", ["SanaVideoPipeline", "SanaImageToVideoPipeline"]) +def test_diffusers_adapter_selects_sana_pipeline_utils(pipeline_class_name): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import ( + SanaVideoPipelineUtils, + get_pipeline_utils, + ) + + assert isinstance(get_pipeline_utils(pipeline_class_name), SanaVideoPipelineUtils) + + +def test_diffusers_adapter_resolves_requested_sana_i2v_pipeline(): + from diffusers import SanaImageToVideoPipeline, SanaVideoPipeline + + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import ( + SanaVideoPipelineUtils, + get_pipeline_utils_for_config, + resolve_diffusers_pipeline_class, + ) + + od_config = SimpleNamespace( + model_class_name="SanaImageToVideoPipeline", + diffusers_pipeline_cls=SanaVideoPipeline, + ) + + assert isinstance(get_pipeline_utils_for_config(od_config), SanaVideoPipelineUtils) + assert resolve_diffusers_pipeline_class(od_config) is SanaImageToVideoPipeline + + +def test_diffusers_adapter_pipeline_override_is_sana_specific(): + from diffusers import DiffusionPipeline + + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import resolve_diffusers_pipeline_class + + class OtherPipeline(DiffusionPipeline): + pass + + od_config = SimpleNamespace( + model_class_name="DiffusersAdapterPipeline", + diffusers_pipeline_cls=OtherPipeline, + ) + + assert resolve_diffusers_pipeline_class(od_config) is OtherPipeline + + +def test_diffusers_adapter_loads_resolved_pipeline_class(monkeypatch): + from vllm_omni.diffusion.models.diffusers_adapter import pipeline_diffusers_adapter + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + + loaded = [] + expected_pipeline = object() + + class ResolvedPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + loaded.append((model_id, kwargs)) + return expected_pipeline + + monkeypatch.setattr( + pipeline_diffusers_adapter, + "resolve_diffusers_pipeline_class", + lambda _od_config: ResolvedPipeline, + ) + adapter = object.__new__(DiffusersAdapterPipeline) + adapter.od_config = SimpleNamespace(diffusers_pipeline_cls=object) + + result = adapter._load_pipeline_from_pretrained("sana-checkpoint", {"torch_dtype": torch.bfloat16}) + + assert result is expected_pipeline + assert loaded == [("sana-checkpoint", {"torch_dtype": torch.bfloat16})] + + +def test_diffusers_adapter_sana_i2v_override_reports_image_capability(): + from diffusers import SanaVideoPipeline + + from vllm_omni.diffusion.io_support import supports_multimodal_input + + od_config = SimpleNamespace( + diffusion_load_format="diffusers", + model_class_name="SanaImageToVideoPipeline", + diffusers_pipeline_cls=SanaVideoPipeline, + ) + + assert supports_multimodal_input(od_config) == (True, False) + + +def test_diffusers_adapter_sana_i2v_warmup_supplies_image(): + from diffusers import SanaVideoPipeline + + from vllm_omni.diffusion.diffusion_engine import DiffusionEngine + + engine = object.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace( + diffusion_load_format="diffusers", + model_class_name="SanaImageToVideoPipeline", + diffusers_pipeline_cls=SanaVideoPipeline, + ) + captured_requests = [] + engine.pre_process_func = lambda request: request + + def capture_request(request): + captured_requests.append(request) + return SimpleNamespace(error=None) + + engine.add_req_and_wait_for_response = capture_request + + engine._dummy_run() + + assert len(captured_requests) == 1 + prompt = captured_requests[0].prompt + assert isinstance(prompt, dict) + assert isinstance(prompt["multi_modal_data"]["image"], Image.Image) + + +def test_diffusers_adapter_disables_resolution_binning_for_warmup(): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils + + adapter = object.__new__(DiffusersAdapterPipeline) + adapter._accept_call_kwargs = {"prompt", "frames", "use_resolution_binning", "generator"} + adapter._pipeline_utils = SanaVideoPipelineUtils() + adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None) + + kwargs = adapter._build_call_kwargs( + _make_request_batch( + "dummy run", + request_id=DUMMY_DIFFUSION_REQUEST_ID, + generator_device="cpu", + ) + ) + + assert kwargs["use_resolution_binning"] is False + assert kwargs["frames"] == 1 + + +def test_pipeline_is_torch_module_and_supports_eval(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + nn.Module.__init__(pipeline) + + assert isinstance(pipeline, nn.Module) + assert pipeline.eval() is pipeline + assert pipeline.training is False + + +@pytest.mark.parametrize( + ("vae_scale", "patch_size", "valid_size", "invalid_size", "alignment"), + [ + (8, (1, 2, 2), 624, 632, 16), + (32, (1, 1, 1), 704, 712, 32), + ], +) +def test_check_inputs_uses_variant_spatial_alignment(vae_scale, patch_size, valid_size, invalid_size, alignment): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.vae_scale_factor_spatial = vae_scale + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(patch_size=patch_size)) + + pipeline.check_inputs(prompt="test", height=valid_size, width=valid_size) + with pytest.raises(ValueError, match=f"divisible by {alignment}"): + pipeline.check_inputs(prompt="test", height=invalid_size, width=valid_size) + + +@pytest.mark.parametrize( + ("requested_output_type", "expected_internal_output_type"), + [ + (None, "raw"), + ("np", "raw"), + ("latent", "latent"), + ], +) +def test_forward_maps_omni_request_to_sana_generation_args(requested_output_type, expected_internal_output_type): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return SimpleNamespace(frames=torch.zeros(1, 3, 9, 192, 320)) + + pipeline._generate = fake_generate + req = _make_request_batch( + {"prompt": "a robot walks", "negative_prompt": "blurry"}, + height=192, + width=320, + num_frames=9, + num_inference_steps=2, + guidance_scale=4.5, + seed=42, + output_type=requested_output_type, + extra_args={"motion_score": 30, "use_resolution_binning": False}, + ) + + output = pipeline.forward(req) + + assert output.output.shape == (1, 3, 9, 192, 320) + assert calls[0]["prompt"] == "a robot walks motion score: 30." + assert calls[0]["negative_prompt"] == "blurry" + assert calls[0]["height"] == 192 + assert calls[0]["width"] == 320 + assert calls[0]["frames"] == 9 + assert calls[0]["num_inference_steps"] == 2 + assert calls[0]["guidance_scale"] == 4.5 + assert calls[0]["use_resolution_binning"] is False + assert calls[0]["generator"].initial_seed() == 42 + assert calls[0]["output_type"] == expected_internal_output_type + + +def test_sana_video_t2v_omitted_num_frames_uses_model_default(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return SimpleNamespace(frames=torch.zeros(1, 3, 81, 2, 2)) + + pipeline._generate = fake_generate + req = _make_request_batch( + "a robot walks", + height=192, + width=320, + num_inference_steps=1, + generator_device="cpu", + ) + + assert req.sampling_params.num_frames == 1 + output = pipeline.forward(req) + + assert output.output.shape == (1, 3, 81, 2, 2) + assert calls[0]["frames"] == 81 + + +def test_sana_video_t2v_dummy_run_keeps_single_frame(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return SimpleNamespace(frames=torch.zeros(1, 3, 1, 192, 320)) + + pipeline._generate = fake_generate + + output = pipeline.forward( + _make_request_batch( + "dummy run", + request_id=DUMMY_DIFFUSION_REQUEST_ID, + height=192, + width=320, + num_frames=1, + num_inference_steps=1, + generator_device="cpu", + ) + ) + + assert output.output.shape == (1, 3, 1, 192, 320) + assert calls[0]["frames"] == 1 + + +@pytest.mark.parametrize( + ("sample_size", "expected_height", "expected_width"), + [ + (30, 480, 832), + (22, 704, 1280), + ], +) +def test_t2v_uses_variant_default_resolution(sample_size, expected_height, expected_width): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=sample_size)) + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return SimpleNamespace(frames=torch.zeros(1, 3, 9, expected_height, expected_width)) + + pipeline._generate = fake_generate + output = pipeline.forward( + _make_request_batch( + "a robot walks", + num_frames=9, + num_inference_steps=1, + guidance_scale=1.0, + seed=42, + generator_device="cpu", + ) + ) + + assert output.output.shape == (1, 3, 9, expected_height, expected_width) + assert calls[0]["height"] == expected_height + assert calls[0]["width"] == expected_width + + +def test_forward_requires_exactly_one_nonempty_prompt(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + + with pytest.raises(ValueError, match="Prompt is required"): + pipeline.forward(_make_request_batch("")) + + first = OmniDiffusionRequest( + prompt="first", + sampling_params=OmniDiffusionSamplingParams(), + request_id="first", + ) + second = OmniDiffusionRequest( + prompt="second", + sampling_params=OmniDiffusionSamplingParams(), + request_id="second", + ) + with pytest.raises(ValueError, match="exactly one prompt"): + pipeline.forward(DiffusionRequestBatch([first, second])) diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video.py b/tests/diffusion/models/sana_video/test_transformer_sana_video.py new file mode 100644 index 000000000..804692d7a --- /dev/null +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video.py @@ -0,0 +1,602 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +_TINY_CONFIG = { + "in_channels": 4, + "out_channels": 4, + "num_attention_heads": 2, + "attention_head_dim": 12, + "num_layers": 1, + "num_cross_attention_heads": 2, + "cross_attention_head_dim": 12, + "cross_attention_dim": 24, + "caption_channels": 8, + "mlp_ratio": 2.0, + "patch_size": (1, 2, 2), + "sample_size": 4, + "rope_max_seq_len": 32, +} + +_GOLDEN_PREFIX = torch.tensor( + [ + -0.1159307063, + 0.2003862262, + -0.3754127920, + -0.2403523028, + -1.2376221418, + -0.6513092518, + -0.3358457983, + -0.9894289970, + 0.7858567238, + 0.5089095831, + -0.5362522602, + -0.3732060790, + ] +) + + +@pytest.mark.parametrize( + ("elementwise_affine", "bias", "expected_state_keys"), + [ + (False, False, set()), + (False, True, set()), + (True, False, {"weight"}), + (True, True, {"weight", "bias"}), + ], +) +@pytest.mark.parametrize( + ("input_dtype", "parameter_dtype"), + [ + (torch.float16, torch.float16), + (torch.bfloat16, torch.bfloat16), + (torch.float32, torch.float32), + (torch.float64, torch.float64), + (torch.float16, torch.float32), + (torch.bfloat16, torch.float32), + (torch.float32, torch.float16), + (torch.float32, torch.bfloat16), + ], +) +def test_sana_rms_norm_matches_diffusers( + elementwise_affine, + bias, + expected_state_keys, + input_dtype, + parameter_dtype, +): + from diffusers.models.normalization import RMSNorm as DiffusersRMSNorm + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaRMSNorm + + reference = DiffusersRMSNorm(8, eps=1e-5, elementwise_affine=elementwise_affine, bias=bias).to( + dtype=parameter_dtype + ) + actual = SanaRMSNorm(8, eps=1e-5, elementwise_affine=elementwise_affine, bias=bias).to(dtype=parameter_dtype) + + assert actual.eps == reference.eps + assert actual.elementwise_affine == reference.elementwise_affine + assert actual.dim == reference.dim + assert set(actual.state_dict()) == expected_state_keys + assert set(actual.state_dict()) == set(reference.state_dict()) + + if elementwise_affine: + weight = torch.linspace(0.5, 1.5, 8, dtype=parameter_dtype) + actual.weight.data.copy_(weight) + reference.weight.data.copy_(weight) + if bias: + norm_bias = torch.linspace(-0.25, 0.25, 8, dtype=parameter_dtype) + actual.bias.data.copy_(norm_bias) + reference.bias.data.copy_(norm_bias) + else: + assert actual.weight is None + assert actual.bias is None + + hidden_states = torch.linspace(-2.0, 2.0, 48, dtype=input_dtype).reshape(2, 3, 8) + expected = reference(hidden_states) + result = actual(hidden_states) + + assert result.dtype == expected.dtype + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("num_channels", "flip_sin_to_cos", "downscale_freq_shift", "scale"), + [ + (256, True, 0, 1), + (7, False, 1, 0.5), + (12, True, 0.5, 2), + ], +) +@pytest.mark.parametrize("timestep_dtype", [torch.int64, torch.float32, torch.float64]) +def test_sana_timesteps_matches_diffusers( + num_channels, + flip_sin_to_cos, + downscale_freq_shift, + scale, + timestep_dtype, +): + from diffusers.models.embeddings import Timesteps + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaTimesteps + + reference = Timesteps( + num_channels=num_channels, + flip_sin_to_cos=flip_sin_to_cos, + downscale_freq_shift=downscale_freq_shift, + scale=scale, + ) + actual = SanaTimesteps( + num_channels=num_channels, + flip_sin_to_cos=flip_sin_to_cos, + downscale_freq_shift=downscale_freq_shift, + scale=scale, + ) + timesteps = torch.tensor([0, 1, 500, 999], dtype=timestep_dtype) + + assert actual.state_dict() == reference.state_dict() == {} + assert actual.num_channels == reference.num_channels + assert actual.flip_sin_to_cos == reference.flip_sin_to_cos + assert actual.downscale_freq_shift == reference.downscale_freq_shift + assert actual.scale == reference.scale + torch.testing.assert_close(actual(timesteps), reference(timesteps), rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("kwargs", "sample_shape", "condition_shape"), + [ + ({}, (2, 256), None), + ( + { + "out_dim": 12, + "post_act_fn": "silu", + "cond_proj_dim": 6, + "sample_proj_bias": False, + }, + (2, 256), + (2, 6), + ), + ], +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_sana_timestep_embedding_matches_diffusers( + kwargs, + sample_shape, + condition_shape, + dtype, +): + from diffusers.models.embeddings import TimestepEmbedding + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaTimestepEmbedding, + ) + + torch.manual_seed(17) + reference = TimestepEmbedding(256, 24, **kwargs).to(dtype=dtype) + actual = SanaTimestepEmbedding(256, 24, **kwargs).to(dtype=dtype) + assert set(actual.state_dict()) == set(reference.state_dict()) + actual.load_state_dict(reference.state_dict()) + + sample = torch.randn(sample_shape, dtype=dtype) + condition = torch.randn(condition_shape, dtype=dtype) if condition_shape else None + expected = reference(sample, condition) + result = actual(sample, condition) + + assert result.shape == expected.shape + assert result.dtype == expected.dtype + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +def test_sana_timestep_embedding_rejects_unsupported_activation(): + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaTimestepEmbedding, + ) + + with pytest.raises(ValueError, match="act_fn='silu'"): + SanaTimestepEmbedding(8, 16, act_fn="gelu") + with pytest.raises(ValueError, match="post_act_fn"): + SanaTimestepEmbedding(8, 16, post_act_fn="gelu") + + module = SanaTimestepEmbedding(8, 16) + with pytest.raises(ValueError, match="cond_proj_dim"): + module(torch.randn(1, 8), torch.randn(1, 4)) + + +@pytest.mark.parametrize("use_additional_conditions", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_sana_ada_layer_norm_single_matches_diffusers( + use_additional_conditions, + dtype, +): + from diffusers.models.normalization import AdaLayerNormSingle + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaAdaLayerNormSingle, + ) + + torch.manual_seed(19) + reference = AdaLayerNormSingle( + 24, + use_additional_conditions=use_additional_conditions, + ).to(dtype=dtype) + actual = SanaAdaLayerNormSingle( + 24, + use_additional_conditions=use_additional_conditions, + ).to(dtype=dtype) + assert set(actual.state_dict()) == set(reference.state_dict()) + actual.load_state_dict(reference.state_dict()) + + timestep = torch.tensor([10, 500]) + added_cond_kwargs = None + if use_additional_conditions: + added_cond_kwargs = { + "resolution": torch.tensor([[480, 832], [704, 1280]]), + "aspect_ratio": torch.tensor([[832 / 480], [1280 / 704]]), + } + + expected = reference( + timestep, + added_cond_kwargs=added_cond_kwargs, + batch_size=2, + hidden_dtype=dtype, + ) + result = actual( + timestep, + added_cond_kwargs=added_cond_kwargs, + batch_size=2, + hidden_dtype=dtype, + ) + + assert len(result) == len(expected) == 2 + for actual_tensor, expected_tensor in zip(result, expected): + assert actual_tensor.shape == expected_tensor.shape + assert actual_tensor.dtype == expected_tensor.dtype + torch.testing.assert_close(actual_tensor, expected_tensor, rtol=0, atol=0) + + +@pytest.mark.parametrize("act_fn", ["gelu_tanh", "silu", "silu_fp32"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_sana_pixart_text_projection_matches_diffusers(act_fn, dtype): + from diffusers.models.embeddings import PixArtAlphaTextProjection + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaPixArtAlphaTextProjection, + ) + + torch.manual_seed(23) + reference = PixArtAlphaTextProjection( + in_features=8, + hidden_size=24, + out_features=12, + act_fn=act_fn, + ).to(dtype=dtype) + actual = SanaPixArtAlphaTextProjection( + in_features=8, + hidden_size=24, + out_features=12, + act_fn=act_fn, + ).to(dtype=dtype) + assert set(actual.state_dict()) == set(reference.state_dict()) + actual.load_state_dict(reference.state_dict()) + + caption = torch.randn(2, 5, 8, dtype=dtype) + expected = reference(caption) + result = actual(caption) + + assert result.shape == expected.shape + assert result.dtype == expected.dtype + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("freqs_dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize(("dim", "max_seq_len", "theta"), [(4, 8, 10000.0), (12, 32, 256.0)]) +def test_native_rope_matches_diffusers(dim, max_seq_len, theta, freqs_dtype): + from diffusers.models.embeddings import get_1d_rotary_pos_embed + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + _get_1d_rotary_pos_embed, + ) + + expected_cos, expected_sin = get_1d_rotary_pos_embed( + dim, + max_seq_len, + theta, + use_real=True, + repeat_interleave_real=True, + freqs_dtype=freqs_dtype, + ) + actual_cos, actual_sin = _get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype) + + torch.testing.assert_close(actual_cos, expected_cos, rtol=0, atol=0) + torch.testing.assert_close(actual_sin, expected_sin, rtol=0, atol=0) + + +def test_tiny_transformer_matches_diffusers_and_frozen_output(): + from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer + from diffusers.configuration_utils import ConfigMixin + from diffusers.models.modeling_utils import ModelMixin + + from vllm_omni.diffusion.attention.layer import Attention + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaLinearAttention, + SanaRMSNorm, + SanaVideoTransformerConfig, + SanaVideoTransformerOutput, + ) + + torch.manual_seed(7) + reference = DiffusersTransformer(**_TINY_CONFIG).eval() + model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval() + assert not isinstance(model, ModelMixin) + assert not isinstance(model, ConfigMixin) + assert isinstance(model.config, SanaVideoTransformerConfig) + assert set(model.state_dict()) == set(reference.state_dict()) + model.load_state_dict(reference.state_dict()) + + block = model.transformer_blocks[0] + assert isinstance(block.attn1, SanaLinearAttention) + assert isinstance(block.attn2.attn, Attention) + assert block.attn2.attn.role == "cross" + assert isinstance(block.attn1.norm_q, SanaRMSNorm) + assert isinstance(block.attn1.norm_k, SanaRMSNorm) + assert isinstance(block.attn2.norm_q, SanaRMSNorm) + assert isinstance(block.attn2.norm_k, SanaRMSNorm) + assert isinstance(model.caption_norm, SanaRMSNorm) + + torch.manual_seed(11) + hidden_states = torch.randn(1, 4, 3, 4, 4) + encoder_hidden_states = torch.randn(1, 5, 8) + encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 0]]) + timestep = torch.tensor([500.0]) + + with torch.no_grad(): + expected = reference( + hidden_states, + encoder_hidden_states, + timestep, + encoder_attention_mask=encoder_attention_mask, + ).sample + actual_output = model( + hidden_states, + encoder_hidden_states, + timestep, + encoder_attention_mask=encoder_attention_mask, + ) + actual = actual_output.sample + + assert isinstance(actual_output, SanaVideoTransformerOutput) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(actual.flatten()[: len(_GOLDEN_PREFIX)], _GOLDEN_PREFIX, rtol=1e-5, atol=1e-5) + + +def test_encoder_attention_mask_uses_bool_sdpa_contract(monkeypatch): + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + torch.manual_seed(37) + model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval() + captured = {} + omni_attention = model.transformer_blocks[0].attn2.attn + + def fake_sdpa_forward(query, key, value, attn_metadata): + captured["query_shape"] = query.shape + captured["key_shape"] = key.shape + captured["mask"] = attn_metadata.attn_mask.detach().clone() + return torch.zeros_like(query) + + def unexpected_backend_forward(*args, **kwargs): + raise AssertionError("masked SANA cross-attention must use SDPA") + + monkeypatch.setattr(omni_attention.sdpa_fallback, "forward", fake_sdpa_forward) + monkeypatch.setattr(omni_attention, "forward", unexpected_backend_forward) + + hidden_states = torch.randn(2, 4, 3, 4, 4) + encoder_hidden_states = torch.randn(2, 5, 8) + input_mask = torch.tensor( + [ + [1, 1, 1, 0, 0], + [1, 1, 1, 1, 0], + ], + dtype=torch.int64, + ) + + with torch.no_grad(): + output = model( + hidden_states, + encoder_hidden_states, + torch.tensor([100.0, 500.0]), + encoder_attention_mask=input_mask, + ).sample + + assert output.shape == hidden_states.shape + assert captured["query_shape"][:2] == (2, 12) + assert captured["key_shape"][:2] == (2, 5) + assert captured["query_shape"][1] != captured["key_shape"][1] + assert captured["mask"].shape == (2, 5) + assert captured["mask"].dtype == torch.bool + assert torch.equal(captured["mask"], input_mask.bool()) + + +def test_encoder_attention_mask_rejects_additive_bias(): + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval() + additive_mask = torch.tensor([[0.0, 0.0, 0.0, -10000.0, -10000.0]]) + + with pytest.raises(TypeError, match="not a floating-point additive attention bias"): + model( + torch.randn(1, 4, 3, 4, 4), + torch.randn(1, 5, 8), + torch.tensor([500.0]), + encoder_attention_mask=additive_mask, + ) + + +def test_linear_attention_requires_rotary_embeddings(): + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaLinearAttention + + attention = SanaLinearAttention( + dim=24, + num_heads=2, + head_dim=12, + dropout=0.0, + bias=True, + qk_norm="rms_norm_across_heads", + ) + + with pytest.raises(ValueError, match="requires rotary_emb"): + attention(torch.randn(1, 4, 24)) + + +def test_linear_attention_bfloat16_matches_diffusers_mixed_precision(): + from diffusers.models.attention_processor import Attention as DiffusersAttention + from diffusers.models.transformers.transformer_sana_video import SanaLinearAttnProcessor3_0 + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaLinearAttention + + torch.manual_seed(17) + reference = DiffusersAttention( + query_dim=24, + heads=2, + dim_head=12, + kv_heads=2, + qk_norm="rms_norm_across_heads", + dropout=0.0, + bias=True, + cross_attention_dim=None, + processor=SanaLinearAttnProcessor3_0(), + ).to(dtype=torch.bfloat16) + actual = SanaLinearAttention( + dim=24, + num_heads=2, + head_dim=12, + dropout=0.0, + bias=True, + qk_norm="rms_norm_across_heads", + ).to(dtype=torch.bfloat16) + actual.load_state_dict(reference.state_dict()) + + hidden_states = torch.randn(1, 5, 24, dtype=torch.bfloat16) + rotary_emb = ( + torch.randn(1, 5, 1, 12, dtype=torch.bfloat16), + torch.randn(1, 5, 1, 12, dtype=torch.bfloat16), + ) + + with torch.no_grad(): + expected = reference(hidden_states, rotary_emb=rotary_emb) + result = actual(hidden_states, rotary_emb=rotary_emb) + + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +def test_transformer_block_without_cross_attention_keeps_feed_forward_norm(): + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformerBlock + + block = SanaVideoTransformerBlock( + dim=24, + num_attention_heads=2, + attention_head_dim=12, + num_cross_attention_heads=2, + cross_attention_head_dim=12, + cross_attention_dim=None, + mlp_ratio=2.0, + ) + + assert block.attn2 is None + assert isinstance(block.norm2, torch.nn.LayerNorm) + + +def test_guidance_transformer_matches_diffusers_and_tuple_output(): + from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer + + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaTimestepEmbedding, + SanaTimesteps, + ) + + config = _TINY_CONFIG | {"guidance_embeds": True} + torch.manual_seed(29) + reference = DiffusersTransformer(**config).eval() + model = SanaVideoTransformer3DModel(**config).eval() + assert set(model.state_dict()) == set(reference.state_dict()) + model.load_state_dict(reference.state_dict()) + assert isinstance(model.time_embed.time_proj, SanaTimesteps) + assert isinstance(model.time_embed.timestep_embedder, SanaTimestepEmbedding) + assert isinstance(model.time_embed.guidance_condition_proj, SanaTimesteps) + assert isinstance(model.time_embed.guidance_embedder, SanaTimestepEmbedding) + + torch.manual_seed(31) + hidden_states = torch.randn(1, 4, 3, 4, 4) + encoder_hidden_states = torch.randn(1, 5, 8) + encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 0]]) + timestep = torch.tensor([500.0]) + guidance = torch.tensor([6.0]) + + with torch.no_grad(): + expected = reference( + hidden_states, + encoder_hidden_states, + timestep, + guidance=guidance, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + actual = model( + hidden_states, + encoder_hidden_states, + timestep, + guidance=guidance, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + + assert isinstance(actual, tuple) + assert len(actual) == len(expected) == 1 + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + + +def test_native_transformer_config_filters_diffusers_metadata(): + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformerConfig + + config = SanaVideoTransformerConfig.from_dict( + _TINY_CONFIG | {"_class_name": "SanaVideoTransformer3DModel", "_diffusers_version": "0.38.0"} + ) + assert config.patch_size == (1, 2, 2) + + with pytest.raises(ValueError, match="unsupported_field"): + SanaVideoTransformerConfig.from_dict(_TINY_CONFIG | {"unsupported_field": True}) + + +def test_dual_vae_variant_resolution(): + from diffusers.pipelines.pipeline_utils import DiffusionPipeline + + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + DistributedAutoencoderKLWan, + ) + from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin + from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( + SanaVideoPipeline, + _resolve_vae_class_and_dtype, + ) + + assert not issubclass(SanaVideoPipeline, DiffusionPipeline) + assert issubclass(SanaVideoPipeline, ProgressBarMixin) + + vae_class, dtype = _resolve_vae_class_and_dtype("AutoencoderKLWan", torch.bfloat16) + assert vae_class is DistributedAutoencoderKLWan + assert dtype is torch.float32 + + vae_class, dtype = _resolve_vae_class_and_dtype("AutoencoderKLLTX2Video", torch.bfloat16) + assert vae_class is DistributedAutoencoderKLLTX2Video + assert dtype is torch.bfloat16 + + with pytest.raises(ValueError, match="Unsupported SANA-Video VAE"): + _resolve_vae_class_and_dtype("UnknownVAE", torch.bfloat16) diff --git a/tests/e2e/accuracy/sana_video/README.md b/tests/e2e/accuracy/sana_video/README.md new file mode 100644 index 000000000..cab38b7a2 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/README.md @@ -0,0 +1,63 @@ +# SANA-Video golden outputs + +The v2 golden set covers both text-to-video and image-to-video for each +checkpoint variant. Generate references with the Diffusers version pinned by +this repository and the immutable Hub revisions embedded in +`generate_goldens.py`: + +```bash +python tests/e2e/accuracy/sana_video/generate_goldens.py \ + --variant 480p --task t2v +python tests/e2e/accuracy/sana_video/generate_goldens.py \ + --variant 480p --task i2v --input-image /path/to/fixed-input.png +python tests/e2e/accuracy/sana_video/generate_goldens.py \ + --variant 720p --task t2v +python tests/e2e/accuracy/sana_video/generate_goldens.py \ + --variant 720p --task i2v --input-image /path/to/fixed-input.png +``` + +Use the same reviewed input image for both I2V variants. The generator +normalizes it to an RGB PNG and includes that exact file and its SHA256 in the +case manifest, so an I2V reference is self-contained. + +Each case is written under: + +```text +/// +``` + +It contains: + +- `transformer_case.safetensors`: deterministic transformer inputs and output. +- `pipeline.mp4`: the encoded reference used by online similarity tests. +- `pipeline_reference.safetensors`: final denoised latents and decoded, + pre-MP4 frames quantized to the exact uint8 video input. +- `input.png`: the normalized conditioning image for I2V only. +- `metadata.json`: model, generator, scheduler, input and sampling provenance. +- `manifest.json`: the size and SHA256 of every artifact. + +## Publication requirements + +A canonical golden must be generated from the pinned Hugging Face model ID and +immutable commit. The generator also requires a clean vLLM-Omni worktree and +records the repository revision plus the generator script hash. + +`--model` is available for local development. It hashes every file in the local +checkpoint directory and records optional `--local-source-model` and +`--local-source-revision` claims, but marks the result `publishable: false`. +The golden test rejects such a manifest. A local directory hash proves which +bytes were used; it does not prove that those bytes match a claimed upstream +commit. + +After reviewing the metadata, input image and generated videos, upload each +publishable case directory without renaming files to: + +```text +s3://vllm-public-assets/omni-assets/sana-video/v2/// +``` + +Set `SANA_VIDEO_GOLDEN_BASE_URL` to the corresponding `v2` HTTP base URL when +running the tests. Uploading is intentionally separate from generation so test +runs cannot overwrite a frozen reference. The storage location must prevent +in-place mutation; publish a new version prefix instead of replacing any v2 +file when a model, dependency, input or sampling parameter changes. diff --git a/tests/e2e/accuracy/sana_video/generate_goldens.py b/tests/e2e/accuracy/sana_video/generate_goldens.py new file mode 100644 index 000000000..8ab25dd18 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/generate_goldens.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate frozen SANA-Video transformer and pipeline reference artifacts.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +import time +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from diffusers import SanaImageToVideoPipeline, SanaVideoPipeline +from diffusers.utils import export_to_video +from PIL import Image +from safetensors.torch import save_file + +PROMPT = "A cat walking on the grass, facing the camera. motion score: 30." +NEGATIVE_PROMPT = ( + "A chaotic sequence with misshapen, deformed limbs in heavy motion blur, sudden disappearance, jump cuts, " + "jerky movements, rapid shot changes, frames out of sync, inconsistent character shapes, temporal artifacts, " + "jitter, and ghosting effects, creating a disorienting visual experience." +) +SCHEMA_VERSION = 2 +SEED = 42 +NUM_FRAMES = 81 +FPS = 16 +NUM_INFERENCE_STEPS = 50 +GUIDANCE_SCALE = 6.0 +VARIANTS = { + "480p": { + "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + "revision": "fed3bce411c58a0f688a31afe8f52e61acc2b15f", + "height": 480, + "width": 832, + "vae_dtype": torch.float32, + "transformer_case_shape": (1, 16, 3, 8, 8), + }, + "720p": { + "model": "Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + "revision": "8bda5e623d0f48cd6da3b387b10ca35d15cf1c4e", + "height": 704, + "width": 1280, + "vae_dtype": torch.bfloat16, + "transformer_case_shape": (1, 128, 2, 4, 4), + }, +} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _json_sha256(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _directory_fingerprint(root: Path) -> dict[str, Any]: + """Fingerprint every local checkpoint file without recording host paths.""" + if not root.is_dir(): + raise ValueError(f"Local checkpoint must be a directory: {root}") + files = [] + digest = hashlib.sha256() + for path in sorted(candidate for candidate in root.rglob("*") if candidate.is_file()): + relative = path.relative_to(root).as_posix() + file_sha256 = _sha256(path) + size = path.stat().st_size + entry = f"{relative}\0{size}\0{file_sha256}\n".encode() + digest.update(entry) + files.append({"path": relative, "size": size, "sha256": file_sha256}) + if not files: + raise ValueError(f"Local checkpoint contains no files: {root}") + return { + "algorithm": "sha256-tree-v1", + "sha256": digest.hexdigest(), + "file_count": len(files), + "total_size": sum(entry["size"] for entry in files), + } + + +def _repository_provenance() -> dict[str, Any]: + repo_root = Path(__file__).resolve().parents[4] + revision = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repo_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + dirty = bool( + subprocess.run( + ["git", "status", "--porcelain", "--untracked-files=no"], + cwd=repo_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + ) + return { + "repository": "vllm-project/vllm-omni", + "revision": revision, + "dirty": dirty, + "generator_sha256": _sha256(Path(__file__)), + } + + +def _cpu(tensor: torch.Tensor) -> torch.Tensor: + return tensor.detach().cpu().contiguous() + + +def _normalize_input_image(source: Path, destination: Path) -> Image.Image: + if not source.is_file(): + raise ValueError(f"I2V input image does not exist: {source}") + with Image.open(source) as image: + normalized = image.convert("RGB") + normalized.load() + normalized.save(destination, format="PNG", optimize=False) + return normalized + + +def _frames_to_uint8(frames: np.ndarray) -> torch.Tensor: + tensor = torch.from_numpy(np.asarray(frames)) + if tensor.is_floating_point(): + # Match diffusers.utils.export_to_video, which truncates rather than + # rounds when converting float frames to the encoder's uint8 input. + tensor = tensor.clamp(0, 1) * 255 + return tensor.to(torch.uint8).contiguous() + + +def _model_provenance( + config: dict[str, Any], + model: str | None, + local_source_model: str | None, + local_source_revision: str | None, +) -> tuple[dict[str, Any], bool]: + if model is None: + return ( + { + "kind": "huggingface_hub", + "model_id": config["model"], + "revision": config["revision"], + }, + True, + ) + + local_path = Path(model).expanduser().resolve() + if bool(local_source_model) != bool(local_source_revision): + raise ValueError("--local-source-model and --local-source-revision must be specified together.") + provenance = { + "kind": "local_checkpoint", + "checkpoint_fingerprint": _directory_fingerprint(local_path), + "claimed_source_model": local_source_model, + "claimed_source_revision": local_source_revision, + "verification": "unverified-local-copy", + } + # A directory hash proves which local bytes were used, but not that those + # bytes match a claimed Hub commit. Such output is useful for development + # comparisons but must never be uploaded as the canonical frozen reference. + return provenance, False + + +def generate( + variant: str, + task: str, + model: str | None, + input_image: Path | None, + output_root: Path, + local_source_model: str | None = None, + local_source_revision: str | None = None, +) -> Path: + config = dict(VARIANTS[variant]) + model_provenance, publishable = _model_provenance( + config, + model, + local_source_model, + local_source_revision, + ) + repository_provenance = _repository_provenance() + if publishable and repository_provenance["dirty"]: + raise RuntimeError("Refusing to generate publishable golden assets from a dirty vllm-omni worktree.") + if model is not None: + config["model"] = str(Path(model).expanduser().resolve()) + config["revision"] = None + + output_dir = output_root / variant / task + output_dir.mkdir(parents=True, exist_ok=True) + normalized_image: Image.Image | None = None + input_path: Path | None = None + if task == "i2v": + if input_image is None: + raise ValueError("--input-image is required for task=i2v.") + input_path = output_dir / "input.png" + normalized_image = _normalize_input_image(input_image, input_path) + elif input_image is not None: + raise ValueError("--input-image is only valid for task=i2v.") + + pipeline_class = SanaImageToVideoPipeline if task == "i2v" else SanaVideoPipeline + pipe = pipeline_class.from_pretrained( + config["model"], + revision=config["revision"], + torch_dtype=torch.bfloat16, + ) + pipe.text_encoder.to(torch.bfloat16) + pipe.vae.to(config["vae_dtype"]) + pipe.to("cuda") + + torch.manual_seed(1234) + hidden_states = torch.randn(config["transformer_case_shape"], dtype=torch.bfloat16, device="cuda") + encoder_hidden_states = torch.randn(1, 8, 2304, dtype=torch.bfloat16, device="cuda") + encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 0]], device="cuda") + timestep = torch.tensor([500.0], device="cuda") + with torch.inference_mode(): + transformer_output = pipe.transformer( + hidden_states, + encoder_hidden_states, + timestep, + encoder_attention_mask=encoder_attention_mask, + ).sample + + transformer_path = output_dir / "transformer_case.safetensors" + save_file( + { + "hidden_states": _cpu(hidden_states), + "encoder_hidden_states": _cpu(encoder_hidden_states), + "encoder_attention_mask": _cpu(encoder_attention_mask), + "timestep": _cpu(timestep), + "output": _cpu(transformer_output), + }, + transformer_path, + ) + + final_latents: dict[str, torch.Tensor] = {} + + def capture_final_latents(_pipe, step, _timestep, callback_kwargs): + if step == NUM_INFERENCE_STEPS - 1: + final_latents["latents"] = _cpu(callback_kwargs["latents"]) + return callback_kwargs + + torch.accelerator.reset_peak_memory_stats() + generator = torch.Generator(device="cuda").manual_seed(SEED) + generation_kwargs = { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": config["height"], + "width": config["width"], + "frames": NUM_FRAMES, + "guidance_scale": GUIDANCE_SCALE, + "num_inference_steps": NUM_INFERENCE_STEPS, + "generator": generator, + "output_type": "np", + "callback_on_step_end": capture_final_latents, + "callback_on_step_end_tensor_inputs": ["latents"], + } + if normalized_image is not None: + generation_kwargs["image"] = normalized_image + + torch.accelerator.synchronize() + started = time.perf_counter() + with torch.inference_mode(): + frames = pipe(**generation_kwargs).frames[0] + torch.accelerator.synchronize() + generation_seconds = time.perf_counter() - started + if "latents" not in final_latents: + raise RuntimeError("Diffusers did not expose the final denoised latents to the generation callback.") + + video_path = output_dir / "pipeline.mp4" + export_to_video(frames, video_path, fps=FPS) + reference_path = output_dir / "pipeline_reference.safetensors" + save_file( + { + "final_latents": final_latents["latents"], + "decoded_frames_uint8": _frames_to_uint8(frames), + }, + reference_path, + ) + + scheduler_config = dict(pipe.scheduler.config) + input_metadata = None + if input_path is not None: + input_metadata = { + "filename": input_path.name, + "sha256": _sha256(input_path), + "size": input_path.stat().st_size, + "format": "png-rgb", + } + metadata = { + "schema_version": SCHEMA_VERSION, + "variant": variant, + "task": task, + "publishable": publishable, + "reference_implementation": { + "library": "diffusers", + "version": __import__("diffusers").__version__, + "pipeline_class": pipeline_class.__name__, + "torch_version": torch.__version__, + }, + "model_provenance": model_provenance, + "generator_provenance": repository_provenance, + "scheduler": { + "class": type(pipe.scheduler).__name__, + "config_sha256": _json_sha256(scheduler_config), + "config": scheduler_config, + }, + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "seed": SEED, + "generator_device": "cuda", + "dtype": "bfloat16", + "vae_dtype": str(config["vae_dtype"]).removeprefix("torch."), + "height": config["height"], + "width": config["width"], + "num_frames": NUM_FRAMES, + "fps": FPS, + "num_inference_steps": NUM_INFERENCE_STEPS, + "guidance_scale": GUIDANCE_SCALE, + "input_image": input_metadata, + "generation_seconds": generation_seconds, + "peak_reserved_gib": torch.accelerator.max_memory_reserved() / 1024**3, + } + metadata_path = output_dir / "metadata.json" + metadata_path.write_text(json.dumps(metadata, indent=2, default=str) + "\n") + + artifact_paths = [transformer_path, video_path, reference_path, metadata_path] + if input_path is not None: + artifact_paths.append(input_path) + manifest = { + "schema_version": SCHEMA_VERSION, + "variant": variant, + "task": task, + "publishable": publishable, + "model_provenance": model_provenance, + "metadata_sha256": _sha256(metadata_path), + "files": {path.name: {"sha256": _sha256(path), "size": path.stat().st_size} for path in artifact_paths}, + } + manifest_path = output_dir / "manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + return manifest_path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--variant", choices=VARIANTS, required=True) + parser.add_argument("--task", choices=("t2v", "i2v"), required=True) + parser.add_argument( + "--model", + help=( + "Optional local checkpoint. Local output is fingerprinted and marked development-only; " + "it cannot be published as the canonical golden." + ), + ) + parser.add_argument("--input-image", type=Path, help="Fixed local RGB input for task=i2v.") + parser.add_argument("--local-source-model", help="Claimed upstream model ID for local checkpoint audit records.") + parser.add_argument("--local-source-revision", help="Claimed upstream commit for local checkpoint audit records.") + parser.add_argument("--output-root", type=Path, default=Path("/tmp/sana-video-goldens-v2")) + args = parser.parse_args() + print( + generate( + args.variant, + args.task, + args.model, + args.input_image, + args.output_root, + args.local_source_model, + args.local_source_revision, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_golden.py b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py new file mode 100644 index 000000000..c0c468315 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import base64 +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import pytest +import requests +import torch +from safetensors.torch import load_file +from torch import nn + +from tests.e2e.accuracy.helpers import assert_video_metadata, assert_video_similarity_metrics, probe_video +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +GOLDEN_BASE_URL = os.getenv("SANA_VIDEO_GOLDEN_BASE_URL") +SCHEMA_VERSION = 2 +PROMPT = "A cat walking on the grass, facing the camera. motion score: 30." +NEGATIVE_PROMPT = ( + "A chaotic sequence with misshapen, deformed limbs in heavy motion blur, sudden disappearance, jump cuts, " + "jerky movements, rapid shot changes, frames out of sync, inconsistent character shapes, temporal artifacts, " + "jitter, and ghosting effects, creating a disorienting visual experience." +) +VARIANTS = { + "480p": ("Efficient-Large-Model/SANA-Video_2B_480p_diffusers", 480, 832), + "720p": ("Efficient-Large-Model/SANA-Video_2B_720p_diffusers", 704, 1280), +} +REVISIONS = { + "480p": "fed3bce411c58a0f688a31afe8f52e61acc2b15f", + "720p": "8bda5e623d0f48cd6da3b387b10ca35d15cf1c4e", +} +SSIM_THRESHOLDS = { + ("480p", "t2v"): 0.91, + ("480p", "i2v"): 0.93, + ("720p", "t2v"): 0.93, + ("720p", "i2v"): 0.93, +} + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.diffusion, + pytest.mark.skipif( + not GOLDEN_BASE_URL, + reason="Set SANA_VIDEO_GOLDEN_BASE_URL after publishing the frozen v2 S3 assets.", + ), +] + + +class _TransformerCheckpoint(nn.Module): + def __init__(self, transformer: nn.Module, model: str, revision: str) -> None: + super().__init__() + from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader + + self.transformer = transformer + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model, + subfolder="transformer", + revision=revision, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _json_sha256(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _assert_hex_digest(value: Any, *, length: int = 64) -> None: + assert isinstance(value, str) and len(value) == length + assert all(character in "0123456789abcdef" for character in value) + + +def _validate_manifest(payload: Any, variant: str, task: str) -> dict[str, Any]: + assert isinstance(payload, dict) + assert payload["schema_version"] == SCHEMA_VERSION + assert payload["variant"] == variant + assert payload["task"] == task + # Local checkpoint generation is deliberately marked non-publishable. Do + # not permit such output to become the canonical reference by accident. + assert payload["publishable"] is True + assert payload["model_provenance"] == { + "kind": "huggingface_hub", + "model_id": VARIANTS[variant][0], + "revision": REVISIONS[variant], + } + _assert_hex_digest(payload["metadata_sha256"]) + + expected_files = { + "metadata.json", + "pipeline.mp4", + "pipeline_reference.safetensors", + "transformer_case.safetensors", + } + if task == "i2v": + expected_files.add("input.png") + files = payload["files"] + assert isinstance(files, dict) + assert set(files) == expected_files + for filename, expected in files.items(): + # Prevent a compromised manifest from writing outside output_dir. + assert filename == Path(filename).name + assert isinstance(expected, dict) and set(expected) == {"sha256", "size"} + _assert_hex_digest(expected["sha256"]) + assert isinstance(expected["size"], int) and expected["size"] > 0 + return payload + + +def _validate_metadata(payload: Any, manifest: dict[str, Any], variant: str, task: str) -> None: + assert isinstance(payload, dict) + _, height, width = VARIANTS[variant] + assert payload["schema_version"] == SCHEMA_VERSION + assert payload["variant"] == variant + assert payload["task"] == task + assert payload["publishable"] is True + assert payload["model_provenance"] == manifest["model_provenance"] + assert payload["prompt"] == PROMPT + assert payload["negative_prompt"] == NEGATIVE_PROMPT + assert payload["seed"] == 42 + assert payload["generator_device"] == "cuda" + assert payload["dtype"] == "bfloat16" + assert payload["height"] == height + assert payload["width"] == width + assert payload["num_frames"] == 81 + assert payload["fps"] == 16 + assert payload["num_inference_steps"] == 50 + assert payload["guidance_scale"] == 6.0 + + implementation = payload["reference_implementation"] + assert implementation["library"] == "diffusers" + assert implementation["pipeline_class"] == ("SanaImageToVideoPipeline" if task == "i2v" else "SanaVideoPipeline") + assert implementation["version"] + assert implementation["torch_version"] + + generator = payload["generator_provenance"] + assert generator["repository"] == "vllm-project/vllm-omni" + _assert_hex_digest(generator["revision"], length=40) + _assert_hex_digest(generator["generator_sha256"]) + assert generator["dirty"] is False + + scheduler = payload["scheduler"] + assert scheduler["class"] == "DPMSolverMultistepScheduler" + _assert_hex_digest(scheduler["config_sha256"]) + assert scheduler["config_sha256"] == _json_sha256(scheduler["config"]) + + input_image = payload["input_image"] + if task == "i2v": + assert input_image["filename"] == "input.png" + _assert_hex_digest(input_image["sha256"]) + assert input_image["size"] == manifest["files"]["input.png"]["size"] + assert input_image["sha256"] == manifest["files"]["input.png"]["sha256"] + assert input_image["format"] == "png-rgb" + else: + assert input_image is None + + +def _download_case(variant: str, task: str, output_dir: Path, filenames: set[str]) -> dict[str, Any]: + output_dir.mkdir(parents=True, exist_ok=True) + base_url = f"{GOLDEN_BASE_URL.rstrip('/')}/{variant}/{task}" + response = requests.get(f"{base_url}/manifest.json", timeout=60) + response.raise_for_status() + payload = _validate_manifest(response.json(), variant, task) + filenames = filenames | {"metadata.json"} + if task == "i2v": + filenames.add("input.png") + assert filenames <= payload["files"].keys() + for filename in filenames: + expected = payload["files"][filename] + path = output_dir / filename + response = requests.get(f"{base_url}/{filename}", timeout=300) + response.raise_for_status() + path.write_bytes(response.content) + assert path.stat().st_size == expected["size"] + assert _sha256(path) == expected["sha256"] + assert _sha256(output_dir / "metadata.json") == payload["metadata_sha256"] + metadata = json.loads((output_dir / "metadata.json").read_text()) + _validate_metadata(metadata, payload, variant, task) + return payload + + +@pytest.mark.benchmark +@pytest.mark.parametrize("variant", VARIANTS) +@pytest.mark.parametrize( + "_hardware", + [pytest.param(None, marks=hardware_marks(res={"cuda": "H100"}))], +) +def test_sana_video_transformer_matches_frozen_golden( + variant: str, + _hardware, + accuracy_artifact_root: Path, +) -> None: + del _hardware + from vllm.config import LoadConfig + from vllm.transformers_utils.config import get_hf_file_to_dict + + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + model, _, _ = VARIANTS[variant] + revision = REVISIONS[variant] + output_dir = accuracy_artifact_root / "sana_video" / variant / "t2v" + _download_case(variant, "t2v", output_dir, {"transformer_case.safetensors"}) + case = load_file(output_dir / "transformer_case.safetensors") + transformer_config = get_hf_file_to_dict("transformer/config.json", model, revision=revision) + assert transformer_config is not None + transformer = SanaVideoTransformer3DModel.from_config(transformer_config).to( + device="cuda", + dtype=torch.bfloat16, + ) + checkpoint = _TransformerCheckpoint(transformer, model, revision) + od_config = OmniDiffusionConfig( + model=model, + model_class_name="SanaVideoPipeline", + dtype=torch.bfloat16, + revision=revision, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config) + transformer.load_weights( + (name.removeprefix("transformer."), tensor) + for name, tensor in loader.get_all_weights(checkpoint) + if name.startswith("transformer.") + ) + transformer.eval() + with torch.inference_mode(): + actual = transformer( + case["hidden_states"].to("cuda"), + case["encoder_hidden_states"].to("cuda"), + case["timestep"].to("cuda"), + encoder_attention_mask=case["encoder_attention_mask"].to("cuda"), + ).sample.float() + expected = case["output"].to("cuda").float() + error = actual - expected + max_abs = error.abs().max() + relative_l2 = torch.linalg.vector_norm(error) / torch.linalg.vector_norm(expected) + cosine = torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0) + # CUDNN_ATTN BF16 reduction order is not bitwise deterministic on + # Blackwell. Bound rare outliers separately while keeping the aggregate + # relative-L2 and cosine checks strict. + max_abs_threshold = 0.2 if variant == "480p" else 0.1 + relative_l2_threshold = 0.017 if variant == "480p" else 0.015 + assert max_abs.item() <= max_abs_threshold + assert relative_l2.item() <= relative_l2_threshold + assert cosine.item() >= 0.999 + + +SERVER_CASES = [ + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--model-class-name", + pipeline_class, + ], + ), + id=f"{variant}-{task}", + marks=hardware_marks(res={"cuda": "H100"}), + ) + for variant, (model, _, _) in VARIANTS.items() + for task, pipeline_class in ( + ("t2v", "SanaVideoPipeline"), + ("i2v", "SanaImageToVideoPipeline"), + ) +] + + +@pytest.mark.benchmark +@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) +def test_sana_video_pipeline_matches_frozen_golden( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + accuracy_artifact_root: Path, +) -> None: + variant = "720p" if "720p" in omni_server.model else "480p" + task = "i2v" if "SanaImageToVideoPipeline" in omni_server.serve_args else "t2v" + _, height, width = VARIANTS[variant] + output_dir = accuracy_artifact_root / "sana_video" / variant / task + _download_case(variant, task, output_dir, {"pipeline.mp4"}) + + request_config = { + "model": omni_server.model, + "form_data": { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": height, + "width": width, + "num_frames": 81, + "fps": 16, + "num_inference_steps": 50, + "guidance_scale": 6.0, + "seed": 42, + }, + } + if task == "i2v": + encoded = base64.b64encode((output_dir / "input.png").read_bytes()).decode() + request_config["image_reference"] = f"data:image/png;base64,{encoded}" + result = openai_client.send_video_diffusion_request(request_config)[0] + actual_path = output_dir / "actual.mp4" + actual_path.write_bytes(result.videos[0]) + golden_path = output_dir / "pipeline.mp4" + + metadata = probe_video(actual_path) + assert_video_metadata(metadata, width=width, height=height, fps=16, frame_count=81) + assert_video_similarity_metrics( + label=f"sana_video_{variant}_{task}", + online_path=actual_path, + offline_path=golden_path, + # The native prompt cross-attention uses vLLM-Omni Attention rather + # than Diffusers SDPA. Frozen stage checks prove exact inputs and + # scheduler state, while the BF16 backend reduction difference grows + # across CFG denoising. Calibrate only the affected 480p T2V case; + # its frozen result is 0.919861 SSIM / 28.7442 dB PSNR. + ssim_threshold=SSIM_THRESHOLDS[(variant, task)], + psnr_threshold=28.0, + ) diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py b/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py new file mode 100644 index 000000000..8dd35eca2 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from tests.helpers.mark import hardware_marks + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.diffusion, + pytest.mark.benchmark, +] + +VARIANTS = { + "480p": { + "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + "vae_class": "AutoencoderKLWan", + "native_vae_class": "DistributedAutoencoderKLWan", + "dtype": torch.float32, + "latent_channels": 16, + "spatial_scale": 8, + "temporal_scale": 4, + "frames": 5, + "max_abs": 1e-6, + "relative_l2": 1e-6, + }, + "720p": { + "model": "Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + "vae_class": "AutoencoderKLLTX2Video", + "native_vae_class": "DistributedAutoencoderKLLTX2Video", + "dtype": torch.bfloat16, + "latent_channels": 128, + "spatial_scale": 32, + "temporal_scale": 8, + "frames": 9, + "max_abs": 0.04, + "relative_l2": 0.004, + }, +} + + +@pytest.mark.parametrize("variant", VARIANTS) +@pytest.mark.parametrize( + "_hardware", + [pytest.param(None, marks=hardware_marks(res={"cuda": "H100"}))], +) +def test_sana_video_i2v_first_frame_latents_match_diffusers(variant: str, _hardware) -> None: + del _hardware + from diffusers import ( + AutoencoderKLLTX2Video, + AutoencoderKLWan, + ) + from diffusers import ( + SanaImageToVideoPipeline as DiffusersSanaImageToVideoPipeline, + ) + + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + DistributedAutoencoderKLWan, + ) + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + config = VARIANTS[variant] + vae_classes = { + "AutoencoderKLLTX2Video": AutoencoderKLLTX2Video, + "AutoencoderKLWan": AutoencoderKLWan, + } + native_vae_classes = { + "DistributedAutoencoderKLLTX2Video": DistributedAutoencoderKLLTX2Video, + "DistributedAutoencoderKLWan": DistributedAutoencoderKLWan, + } + reference_vae = ( + vae_classes[config["vae_class"]] + .from_pretrained( + config["model"], + subfolder="vae", + torch_dtype=config["dtype"], + ) + .to("cuda") + ) + native_vae = ( + native_vae_classes[config["native_vae_class"]] + .from_config(dict(reference_vae.config)) + .to(device="cuda", dtype=config["dtype"]) + ) + native_vae.load_state_dict(reference_vae.state_dict()) + reference_vae.eval() + native_vae.eval() + + height, width = 192, 320 + num_latent_frames = (config["frames"] - 1) // config["temporal_scale"] + 1 + image = torch.linspace( + -1, + 1, + 3 * height * width, + dtype=torch.float32, + device="cuda", + ).reshape(1, 3, height, width) + noise = torch.randn( + ( + 1, + config["latent_channels"], + num_latent_frames, + height // config["spatial_scale"], + width // config["spatial_scale"], + ), + generator=torch.Generator(device="cuda").manual_seed(7), + device="cuda", + dtype=torch.float32, + ) + reference_pipeline = SimpleNamespace( + vae=reference_vae, + vae_scale_factor_temporal=config["temporal_scale"], + vae_scale_factor_spatial=config["spatial_scale"], + ) + native_pipeline = SimpleNamespace( + vae=native_vae, + vae_scale_factor_temporal=config["temporal_scale"], + vae_scale_factor_spatial=config["spatial_scale"], + device=torch.device("cuda"), + ) + + with torch.inference_mode(): + expected = DiffusersSanaImageToVideoPipeline.prepare_latents( + reference_pipeline, + image, + 1, + config["latent_channels"], + height, + width, + config["frames"], + torch.float32, + torch.device("cuda"), + None, + noise.clone(), + ) + actual = SanaImageToVideoPipeline._prepare_i2v_latents( + native_pipeline, + image, + 1, + config["latent_channels"], + height, + width, + config["frames"], + torch.float32, + None, + noise.clone(), + ) + + # Conditioning replaces only the first latent frame. + torch.testing.assert_close(actual[:, :, 1:], noise[:, :, 1:], rtol=0, atol=0) + error = (actual[:, :, :1] - expected[:, :, :1]).float() + expected_first_frame = expected[:, :, :1].float() + max_abs = error.abs().max() + relative_l2 = torch.linalg.vector_norm(error) / torch.linalg.vector_norm(expected_first_frame) + cosine = torch.nn.functional.cosine_similarity( + actual[:, :, :1].flatten(), + expected[:, :, :1].flatten(), + dim=0, + ) + assert max_abs.item() <= config["max_abs"] + assert relative_l2.item() <= config["relative_l2"] + assert cosine.item() >= 0.9999 diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py new file mode 100644 index 000000000..a88aec7e2 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py @@ -0,0 +1,432 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import gc +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_marks + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.diffusion, + pytest.mark.benchmark, +] + +PROMPT = "A small red panda walks through a bamboo forest." +NEGATIVE_PROMPT = "blur, distortion, abrupt cuts" +SEED = 42 +GUIDANCE_SCALE = 6.0 +NUM_INFERENCE_STEPS = 2 +HEIGHT = 192 +WIDTH = 320 +NUM_FRAMES = 9 + +VARIANTS = { + "480p": { + "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + "revision": "fed3bce411c58a0f688a31afe8f52e61acc2b15f", + "model_env": "SANA_VIDEO_480P_MODEL", + "max_abs": 0.2, + "relative_l2": 0.02, + "cosine": 0.999, + "cfg_max_abs": 0.5, + "cfg_relative_l2": 0.05, + "cfg_cosine": 0.999, + "step_max_abs": 0.075, + "step_relative_l2": 0.012, + "step_cosine": 0.9999, + }, + "720p": { + "model": "Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + "revision": "8bda5e623d0f48cd6da3b387b10ca35d15cf1c4e", + "model_env": "SANA_VIDEO_720P_MODEL", + "max_abs": 0.1, + "relative_l2": 0.015, + "cosine": 0.999, + "cfg_max_abs": 0.6, + "cfg_relative_l2": 0.065, + "cfg_cosine": 0.998, + "step_max_abs": 0.08, + "step_relative_l2": 0.014, + "step_cosine": 0.9999, + }, +} + + +@dataclass(frozen=True) +class ErrorMetrics: + shape: tuple[int, ...] + actual_dtype: torch.dtype + expected_dtype: torch.dtype + max_abs: float + mean_abs: float + relative_l2: float + cosine: float + + def __str__(self) -> str: + return ( + f"shape={self.shape}, actual_dtype={self.actual_dtype}, " + f"expected_dtype={self.expected_dtype}, max_abs={self.max_abs:.8g}, " + f"mean_abs={self.mean_abs:.8g}, relative_l2={self.relative_l2:.8g}, " + f"cosine={self.cosine:.8g}" + ) + + +def _metrics(actual: torch.Tensor, expected: torch.Tensor) -> ErrorMetrics: + assert actual.shape == expected.shape + actual_float = actual.detach().float() + expected_float = expected.detach().float() + error = actual_float - expected_float + expected_norm = torch.linalg.vector_norm(expected_float) + relative_l2 = torch.linalg.vector_norm(error) / expected_norm.clamp_min(torch.finfo(torch.float32).tiny) + if actual_float.numel() == 1: + cosine = torch.tensor( + 1.0 if bool(torch.equal(actual_float, expected_float)) else 0.0, + device=actual.device, + ) + else: + cosine = F.cosine_similarity(actual_float.flatten(), expected_float.flatten(), dim=0) + return ErrorMetrics( + shape=tuple(actual.shape), + actual_dtype=actual.dtype, + expected_dtype=expected.dtype, + max_abs=error.abs().max().item(), + mean_abs=error.abs().mean().item(), + relative_l2=relative_l2.item(), + cosine=cosine.item(), + ) + + +def _assert_stage( + stage: str, + actual: torch.Tensor, + expected: torch.Tensor, + *, + max_abs: float, + relative_l2: float, + cosine: float, +) -> ErrorMetrics: + result = _metrics(actual, expected) + print(f"\nSANA-Video alignment [{stage}]: {result}") + assert result.actual_dtype == result.expected_dtype, f"{stage}: {result}" + assert result.max_abs <= max_abs, f"{stage}: {result}" + assert result.relative_l2 <= relative_l2, f"{stage}: {result}" + assert result.cosine >= cosine, f"{stage}: {result}" + return result + + +def _model_source(variant: str) -> tuple[str, str | None]: + config = VARIANTS[variant] + model = os.getenv(config["model_env"], config["model"]) + revision = None if Path(model).exists() else config["revision"] + return model, revision + + +def _native_method_proxy(reference_pipeline: Any): + """Build the minimum native-pipeline surface needed by stage helpers. + + Sharing the tokenizer and text encoder makes this a comparison of pipeline + orchestration rather than a second copy of the same multi-billion-parameter + encoder. Transformer execution still uses separate Diffusers and native + modules below. + """ + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + proxy = object.__new__(SanaVideoPipeline) + torch.nn.Module.__init__(proxy) + proxy.tokenizer = reference_pipeline.tokenizer + proxy.text_encoder = reference_pipeline.text_encoder + proxy.device = torch.device("cuda") + proxy.vae_scale_factor_temporal = reference_pipeline.vae_scale_factor_temporal + proxy.vae_scale_factor_spatial = reference_pipeline.vae_scale_factor_spatial + return proxy + + +def _cfg_prediction( + prediction: torch.Tensor, + *, + guidance_scale: float, + latent_channels: int, + out_channels: int, +) -> torch.Tensor: + unconditional, conditional = prediction.float().chunk(2) + prediction = unconditional + guidance_scale * (conditional - unconditional) + if out_channels // 2 == latent_channels: + prediction = prediction.chunk(2, dim=1)[0] + return prediction + + +@pytest.mark.parametrize("variant", VARIANTS) +@pytest.mark.parametrize( + "_hardware", + [pytest.param(None, marks=hardware_marks(res={"cuda": "H100"}))], +) +def test_sana_video_quick_pipeline_stages_match_diffusers( + variant: str, + _hardware, +) -> None: + """Compare the first denoising step without replacing native attention. + + This deliberately stops after the first scheduler step. It keeps the test + useful as a stage-localizing PR/nightly check while the frozen 81-frame, + 50-step golden test covers the final latent and decoded-video boundary. + """ + del _hardware + from diffusers import SanaVideoPipeline as DiffusersSanaVideoPipeline + from diffusers.pipelines.sana_video.pipeline_sana_video import ( + retrieve_timesteps as diffusers_retrieve_timesteps, + ) + from diffusers.schedulers import DPMSolverMultistepScheduler + + from vllm_omni.diffusion.models.sana_video import ( + SanaVideoPipeline, + SanaVideoTransformer3DModel, + ) + from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( + retrieve_timesteps as native_retrieve_timesteps, + ) + + model, revision = _model_source(variant) + config = VARIANTS[variant] + load_kwargs: dict[str, Any] = { + "torch_dtype": torch.bfloat16, + # No stage in this test decodes latents. Avoid loading a multi-GB VAE + # that would add memory and runtime without increasing coverage. + "vae": None, + } + if revision is not None: + load_kwargs["revision"] = revision + + reference = DiffusersSanaVideoPipeline.from_pretrained(model, **load_kwargs).to("cuda") + assert type(reference) is DiffusersSanaVideoPipeline + assert type(reference.transformer).__module__.startswith("diffusers.") + native_proxy = _native_method_proxy(reference) + assert type(native_proxy) is SanaVideoPipeline + assert type(native_proxy).__module__.startswith("vllm_omni.") + try: + with torch.inference_mode(): + reference_prompt = reference.encode_prompt( + PROMPT, + True, + negative_prompt=NEGATIVE_PROMPT, + device=torch.device("cuda"), + clean_caption=False, + max_sequence_length=300, + complex_human_instruction=None, + ) + native_prompt = SanaVideoPipeline.encode_prompt( + native_proxy, + PROMPT, + True, + negative_prompt=NEGATIVE_PROMPT, + device=torch.device("cuda"), + clean_caption=False, + max_sequence_length=300, + complex_human_instruction=None, + ) + prompt_stage_names = ( + "prompt_embeddings", + "prompt_attention_mask", + "negative_prompt_embeddings", + "negative_prompt_attention_mask", + ) + for stage, actual, expected in zip(prompt_stage_names, native_prompt, reference_prompt, strict=True): + _assert_stage( + stage, + actual, + expected, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + + reference_prompt_embeds, reference_prompt_mask, reference_negative_embeds, reference_negative_mask = ( + reference_prompt + ) + encoder_hidden_states = torch.cat([reference_negative_embeds, reference_prompt_embeds]) + encoder_attention_mask = torch.cat([reference_negative_mask, reference_prompt_mask]) + + latent_channels = reference.transformer.config.in_channels + reference_latents = reference.prepare_latents( + 1, + latent_channels, + HEIGHT, + WIDTH, + NUM_FRAMES, + torch.float32, + torch.device("cuda"), + torch.Generator(device="cuda").manual_seed(SEED), + ) + native_latents = SanaVideoPipeline.prepare_latents( + native_proxy, + 1, + latent_channels, + HEIGHT, + WIDTH, + NUM_FRAMES, + torch.float32, + torch.device("cuda"), + torch.Generator(device="cuda").manual_seed(SEED), + ) + _assert_stage( + "initial_noise_latents", + native_latents, + reference_latents, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + + reference_scheduler = DPMSolverMultistepScheduler.from_config(reference.scheduler.config) + native_scheduler = DPMSolverMultistepScheduler.from_config(reference.scheduler.config) + isolated_scheduler = DPMSolverMultistepScheduler.from_config(reference.scheduler.config) + reference_timesteps, reference_num_steps = diffusers_retrieve_timesteps( + reference_scheduler, + NUM_INFERENCE_STEPS, + torch.device("cuda"), + ) + native_timesteps, native_num_steps = native_retrieve_timesteps( + native_scheduler, + NUM_INFERENCE_STEPS, + torch.device("cuda"), + ) + isolated_retrieve_timesteps, isolated_num_steps = native_retrieve_timesteps( + isolated_scheduler, + NUM_INFERENCE_STEPS, + torch.device("cuda"), + ) + assert reference_num_steps == native_num_steps == isolated_num_steps + _assert_stage( + "timesteps", + native_timesteps, + reference_timesteps, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + _assert_stage( + "sigmas", + native_scheduler.sigmas, + reference_scheduler.sigmas, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + torch.testing.assert_close(isolated_retrieve_timesteps, reference_timesteps, rtol=0, atol=0) + + native_transformer = SanaVideoTransformer3DModel.from_config(dict(reference.transformer.config)).to( + device="cuda", + dtype=torch.bfloat16, + ) + incompatible = native_transformer.load_state_dict(reference.transformer.state_dict(), strict=True) + assert not incompatible.missing_keys + assert not incompatible.unexpected_keys + assert type(native_transformer) is SanaVideoTransformer3DModel + assert type(native_transformer).__module__.startswith("vllm_omni.") + assert type(native_transformer) is not type(reference.transformer) + reference_parameter = next(reference.transformer.parameters()) + native_parameter = next(native_transformer.parameters()) + assert native_parameter.data_ptr() != reference_parameter.data_ptr() + torch.testing.assert_close(native_parameter, reference_parameter, rtol=0, atol=0) + native_transformer.eval() + reference.transformer.eval() + + latent_model_input = torch.cat([reference_latents, reference_latents]) + timestep = reference_timesteps[0].expand(latent_model_input.shape[0]) + with torch.inference_mode(): + reference_raw_prediction = reference.transformer( + latent_model_input.to(dtype=torch.bfloat16), + encoder_hidden_states=encoder_hidden_states.to(dtype=torch.bfloat16), + encoder_attention_mask=encoder_attention_mask, + timestep=timestep, + return_dict=False, + )[0] + native_raw_prediction = native_transformer( + latent_model_input.to(dtype=torch.bfloat16), + encoder_hidden_states=encoder_hidden_states.to(dtype=torch.bfloat16), + encoder_attention_mask=encoder_attention_mask, + timestep=timestep, + return_dict=False, + )[0] + + _assert_stage( + "first_transformer_raw_prediction", + native_raw_prediction, + reference_raw_prediction, + max_abs=config["max_abs"], + relative_l2=config["relative_l2"], + cosine=config["cosine"], + ) + reference_cfg_prediction = _cfg_prediction( + reference_raw_prediction, + guidance_scale=GUIDANCE_SCALE, + latent_channels=latent_channels, + out_channels=reference.transformer.config.out_channels, + ) + native_cfg_prediction = _cfg_prediction( + native_raw_prediction, + guidance_scale=GUIDANCE_SCALE, + latent_channels=latent_channels, + out_channels=native_transformer.config.out_channels, + ) + _assert_stage( + "first_cfg_prediction", + native_cfg_prediction, + reference_cfg_prediction, + max_abs=config["cfg_max_abs"], + relative_l2=config["cfg_relative_l2"], + cosine=config["cfg_cosine"], + ) + + reference_first_step = reference_scheduler.step( + reference_cfg_prediction, + reference_timesteps[0], + reference_latents, + return_dict=False, + )[0] + # Isolate scheduler orchestration from native-attention error by + # feeding the Diffusers prediction through an independently configured + # scheduler initialized by the native helper. + isolated_first_step = isolated_scheduler.step( + reference_cfg_prediction, + isolated_retrieve_timesteps[0], + reference_latents, + return_dict=False, + )[0] + _assert_stage( + "first_scheduler_step_isolated", + isolated_first_step, + reference_first_step, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + + native_first_step = native_scheduler.step( + native_cfg_prediction, + native_timesteps[0], + native_latents, + return_dict=False, + )[0] + _assert_stage( + "first_scheduler_step_end_to_end", + native_first_step, + reference_first_step, + max_abs=config["step_max_abs"], + relative_l2=config["step_relative_l2"], + cosine=config["step_cosine"], + ) + finally: + del native_proxy + del reference + gc.collect() + torch.accelerator.empty_cache() diff --git a/tests/e2e/online_serving/test_sana_video_expansion.py b/tests/e2e/online_serving/test_sana_video_expansion.py new file mode 100644 index 000000000..31e0f8552 --- /dev/null +++ b/tests/e2e/online_serving/test_sana_video_expansion.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""SANA-Video 2B online smoke tests for the native and Diffusers backends.""" + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL_480P = os.environ.get( + "SANA_VIDEO_480P_MODEL", + "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", +) +MODEL_720P = os.environ.get( + "SANA_VIDEO_720P_MODEL", + "Efficient-Large-Model/SANA-Video_2B_720p_diffusers", +) +PROMPT = "A cat walking on grass toward the camera. motion score: 30." +NEGATIVE_PROMPT = "blurry, low quality, temporal artifacts" + +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _t2v_backend_cases(): + cases = [] + for variant, model in (("480p", MODEL_480P), ("720p", MODEL_720P)): + cases.extend( + [ + pytest.param( + OmniServerParams(model=model, server_args=["--model-class-name", "SanaVideoPipeline"]), + id=f"native-{variant}", + marks=SINGLE_CARD_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--diffusion-load-format", + "diffusers", + "--diffusion-attention-backend", + "TORCH_SDPA", + ], + ), + id=f"diffusers-adapter-{variant}", + marks=SINGLE_CARD_MARKS, + ), + ] + ) + return cases + + +@pytest.mark.parametrize("omni_server", _t2v_backend_cases(), indirect=True) +def test_sana_video_t2v_backends(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + request_config = { + "model": omni_server.model, + "form_data": { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": 192, + "width": 320, + "num_frames": 9, + "fps": 8, + "num_inference_steps": 2, + "guidance_scale": 4.0, + "seed": 42, + }, + } + openai_client.send_video_diffusion_request(request_config) + + +def _i2v_cases(): + cases = [] + for variant, model in (("480p", MODEL_480P), ("720p", MODEL_720P)): + cases.extend( + [ + pytest.param( + OmniServerParams(model=model, server_args=["--model-class-name", "SanaImageToVideoPipeline"]), + id=f"native-{variant}", + marks=SINGLE_CARD_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--model-class-name", + "SanaImageToVideoPipeline", + "--diffusion-load-format", + "diffusers", + "--diffusion-attention-backend", + "TORCH_SDPA", + ], + ), + id=f"diffusers-adapter-{variant}", + marks=SINGLE_CARD_MARKS, + ), + ] + ) + return cases + + +@pytest.mark.parametrize("omni_server", _i2v_cases(), indirect=True) +def test_sana_video_i2v_variants(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + request_config = { + "model": omni_server.model, + "form_data": { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": 192, + "width": 320, + "num_frames": 9, + "fps": 8, + "num_inference_steps": 2, + "guidance_scale": 4.0, + "seed": 42, + }, + "image_reference": (f"data:image/jpeg;base64,{generate_synthetic_image(320, 192, seed=42)['base64']}"), + } + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/examples/offline_inference/test_text_to_video_prompts.py b/tests/examples/offline_inference/test_text_to_video_prompts.py index 342638ccc..c588b7a38 100644 --- a/tests/examples/offline_inference/test_text_to_video_prompts.py +++ b/tests/examples/offline_inference/test_text_to_video_prompts.py @@ -46,6 +46,9 @@ def test_text_to_video_builds_canonical_prompt( @pytest.mark.parametrize( ("model", "model_class_name", "preset_name"), [ + ("Efficient-Large-Model/SANA-Video_2B_480p_diffusers", None, "sana_480p"), + ("Efficient-Large-Model/SANA-Video_2B_720p_diffusers", None, "sana_720p"), + ("/models/custom-checkpoint", "SanaVideoPipeline", "sana_480p"), ("robbyant/lingbot-video-dense-1.3b", None, "lingbot"), ("/models/custom-checkpoint", "LingBotVideoPipeline", "lingbot"), ("Lightricks/LTX-2", "LTX2DistilledPipeline", "ltx2_distilled"), diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index 51881ddf3..394d7a3e1 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -1,7 +1,16 @@ from functools import partial from pathlib import Path -from tests.helpers.tiny_model import build_tiny_from_configs +import torch +from diffusers import ( + AutoencoderKLWan, + DPMSolverMultistepScheduler, + SanaVideoPipeline, + SanaVideoTransformer3DModel, +) +from transformers import Gemma2Config, Gemma2Model, GemmaTokenizerFast + +from tests.helpers.tiny_model import _get_tiny_model_path, build_tiny_from_configs TINY_CONFIGS_DIR = Path(__file__).parent / "tiny_configs" @@ -40,6 +49,72 @@ def tiny_ltx2_builder() -> str: return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline") +def tiny_sana_video_builder() -> str: + """Build a tiny 480p SANA-Video model without downloading model weights. + + The tokenizer is the only component loaded from the upstream repository. + All modules with weights are initialized locally from intentionally small + configs, and the scheduler is constructed from its weight-free config. + """ + model_id = "Efficient-Large-Model/SANA-Video_2B_480p_diffusers" + model_dir = _get_tiny_model_path("SanaVideoPipeline") + + tokenizer = GemmaTokenizerFast.from_pretrained(model_id, subfolder="tokenizer") + scheduler = DPMSolverMultistepScheduler( + algorithm_type="dpmsolver++", + flow_shift=8.0, + prediction_type="flow_prediction", + use_flow_sigmas=True, + ) + text_encoder = Gemma2Model( + Gemma2Config( + vocab_size=256000, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=512, + sliding_window=256, + layer_types=["sliding_attention"], + ) + ) + transformer = SanaVideoTransformer3DModel( + in_channels=16, + out_channels=16, + num_attention_heads=2, + attention_head_dim=16, + num_layers=1, + num_cross_attention_heads=2, + cross_attention_head_dim=16, + cross_attention_dim=32, + caption_channels=32, + mlp_ratio=2.0, + sample_size=30, + patch_size=(1, 2, 2), + ) + vae = AutoencoderKLWan( + base_dim=8, + decoder_base_dim=8, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + latents_mean=[0.0] * 16, + latents_std=[1.0] * 16, + ) + pipeline = SanaVideoPipeline( + tokenizer=tokenizer, + text_encoder=text_encoder, + vae=vae, + transformer=transformer, + scheduler=scheduler, + ) + pipeline.to(dtype=torch.bfloat16).save_pretrained(model_dir) + return model_dir + + def _shrink_flux_clip_text_encoder(config: dict) -> dict: config["num_hidden_layers"] = 2 return config diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index 399a3d765..e671522cd 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -49,6 +49,12 @@ DIFFUSION_TEST_SETTINGS = { [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD], ], ), + "SanaVideoPipeline": DiffusionModelTestOpts( + model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + builder=diff_model_builders.tiny_sana_video_builder, + supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO], + check_multi_output=False, + ), "QwenImagePipeline": DiffusionModelTestOpts( model="Qwen/Qwen-Image", builder=diff_model_builders.tiny_qwen_image_builder, diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index f3aa11525..6c35b231a 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -63,6 +63,12 @@ EXCLUDED_MODELS = [ "HunyuanVideo15Pipeline", "HunyuanVideo15ImageToVideoPipeline", "LingBotVideoPipeline", + # The common harness does not yet define an image-to-video task or supply + # an input image. SANA I2V is covered by + # tests/diffusion/models/sana_video/test_pipeline_sana_video.py and + # tests/e2e/online_serving/test_sana_video_expansion.py. SANA T2V is exercised through + # the common tiny-model setting in model_settings.py. + "SanaImageToVideoPipeline", # Requires camera/action assets and persistent AR state; covered by its # dedicated contract tests and real-checkpoint E2E smoke instead. "LingBotWorldCausalDMDPipeline", diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py index 368577a0a..8fecc50df 100644 --- a/vllm_omni/diffusion/io_support.py +++ b/vllm_omni/diffusion/io_support.py @@ -7,11 +7,15 @@ import inspect from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata +from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import resolve_diffusers_pipeline_class from vllm_omni.diffusion.registry import DiffusionModelRegistry def supports_multimodal_input(od_config: OmniDiffusionConfig) -> tuple[bool, bool]: - if od_config.diffusion_load_format == "diffusers" and (pipe_cls := od_config.diffusers_pipeline_cls) is not None: + if ( + od_config.diffusion_load_format == "diffusers" + and (pipe_cls := resolve_diffusers_pipeline_class(od_config)) is not None + ): signature = inspect.signature(pipe_cls.__call__) support_image_input = "image" in signature.parameters support_audio_input = ( diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index a4a9bc0dd..006af60e1 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -23,7 +23,11 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline from torch import nn from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import BasePipelineUtils, get_pipeline_utils +from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import ( + BasePipelineUtils, + get_pipeline_utils_for_config, + resolve_diffusers_pipeline_class, +) from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import ( apply_diffusers_quantization_config, convert_diffusers_quantization_config, @@ -117,9 +121,7 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): } convert_diffusers_quantization_config(load_kwargs) - pipeline_class = self.od_config.diffusers_pipeline_cls - pipeline_class_name = pipeline_class.__name__ if pipeline_class is not None else None - self._pipeline_utils = get_pipeline_utils(pipeline_class_name) + self._pipeline_utils = get_pipeline_utils_for_config(self.od_config) self._pipeline_utils.update_load_kwargs(self.od_config, load_kwargs) component_names = ( self._load_diffusers_component_names(model_id, load_kwargs) @@ -129,7 +131,7 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): apply_diffusers_quantization_config(self.od_config, load_kwargs, component_names) logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}") - self._pipeline = DiffusionPipeline.from_pretrained(model_id, **load_kwargs) + self._pipeline = self._load_pipeline_from_pretrained(model_id, load_kwargs) self._pipeline_utils.apply_post_load_updates(self._pipeline, self.od_config) self._pipeline.to(self.device) @@ -162,6 +164,23 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): # Attention backend self._set_attention_backend() + def _load_pipeline_from_pretrained( + self, + model_id: str, + load_kwargs: dict[str, Any], + ) -> DiffusionPipeline: + pipeline_class = resolve_diffusers_pipeline_class(self.od_config) + configured_pipeline_class = self.od_config.diffusers_pipeline_cls + # Preserve DiffusionPipeline's normal dynamic class resolution unless + # a model-specific adapter hook explicitly selected another class. + loader_class = ( + pipeline_class + if pipeline_class is not None and pipeline_class is not configured_pipeline_class + else DiffusionPipeline + ) + logger.info("Loading Diffusers pipeline class %s", loader_class.__name__) + return loader_class.from_pretrained(model_id, **load_kwargs) + # ------------------------------------------------------------------ # Step-wise execution — explicitly rejected # ------------------------------------------------------------------ @@ -376,6 +395,13 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): if self._accept_call_kwargs is None or key in self._accept_call_kwargs: kwargs[key] = value + self._pipeline_utils.update_call_kwargs( + req, + sampling, + self._accept_call_kwargs, + kwargs, + ) + # Special format fields in sampling params if output_type := sampling.output_type or self.od_config.output_type: kwargs["output_type"] = output_type diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py index 4c810bb68..dc5ea2038 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py @@ -1,17 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from typing import Any from diffusers.pipelines.pipeline_utils import DiffusionPipeline from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.request import resolve_video_num_frames +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams class BasePipelineUtils: """No-op hooks for pipeline-specific diffusers adapter behavior.""" + def resolve_pipeline_class( + self, + od_config: OmniDiffusionConfig, + configured_pipeline_class: type[DiffusionPipeline] | None, + ) -> type[DiffusionPipeline] | None: + """Resolve the concrete Diffusers class used to load a checkpoint.""" + return configured_pipeline_class + def update_load_kwargs(self, od_config: OmniDiffusionConfig, load_kwargs: dict[str, Any]) -> None: pass @@ -21,6 +31,59 @@ class BasePipelineUtils: def validate_runtime_sampling_params(self, sampling: OmniDiffusionSamplingParams) -> None: pass + def update_call_kwargs( + self, + req: DiffusionRequestBatch, + sampling: OmniDiffusionSamplingParams, + accepted_call_kwargs: set[str] | None, + call_kwargs: dict[str, Any], + ) -> None: + pass + + +class SanaVideoPipelineUtils(BasePipelineUtils): + def resolve_pipeline_class( + self, + od_config: OmniDiffusionConfig, + configured_pipeline_class: type[DiffusionPipeline] | None, + ) -> type[DiffusionPipeline] | None: + # Both released SANA-Video checkpoints declare SanaVideoPipeline in + # model_index.json. Honour an explicit Omni I2V class selection by + # loading the same components into Diffusers' I2V pipeline instead. + if od_config.model_class_name == "SanaImageToVideoPipeline": + from diffusers import SanaImageToVideoPipeline + + return SanaImageToVideoPipeline + return configured_pipeline_class + + def update_call_kwargs( + self, + req: DiffusionRequestBatch, + sampling: OmniDiffusionSamplingParams, + accepted_call_kwargs: set[str] | None, + call_kwargs: dict[str, Any], + ) -> None: + # Diffusers SANA-Video calls the shared request's `num_frames` + # parameter `frames`. + if ( + accepted_call_kwargs is not None + and "frames" in accepted_call_kwargs + and "num_frames" not in accepted_call_kwargs + ): + configured_default = call_kwargs.get("frames", 81) + if configured_default is None: + configured_default = 81 + call_kwargs["frames"] = resolve_video_num_frames( + sampling.num_frames, + default_num_frames=configured_default, + is_dummy_run=req.is_dummy_run(), + ) + + # SANA's startup dummy dimensions can map to a resolution bucket that + # conflicts with the upstream pipeline's divisibility validation. + if req.is_dummy_run() and accepted_call_kwargs is not None and "use_resolution_binning" in accepted_call_kwargs: + call_kwargs["use_resolution_binning"] = False + class WanPipelineUtils(BasePipelineUtils): def update_load_kwargs(self, od_config: OmniDiffusionConfig, load_kwargs: dict[str, Any]) -> None: @@ -49,6 +112,8 @@ class WanPipelineUtils(BasePipelineUtils): PIPELINE_UTILS_REGISTRY: dict[str, type[BasePipelineUtils]] = { + "SanaVideoPipeline": SanaVideoPipelineUtils, + "SanaImageToVideoPipeline": SanaVideoPipelineUtils, "WanPipeline": WanPipelineUtils, "WanImageToVideoPipeline": WanPipelineUtils, "WanVACEPipeline": WanPipelineUtils, @@ -62,3 +127,20 @@ def get_pipeline_utils(pipeline_class_name: str | None) -> BasePipelineUtils: return BasePipelineUtils() pipeline_utils_cls = PIPELINE_UTILS_REGISTRY.get(pipeline_class_name, BasePipelineUtils) return pipeline_utils_cls() + + +def get_pipeline_utils_for_config(od_config: OmniDiffusionConfig) -> BasePipelineUtils: + """Return model-specific adapter hooks for the requested or configured class.""" + requested_class_name = od_config.model_class_name + if requested_class_name in PIPELINE_UTILS_REGISTRY: + return get_pipeline_utils(requested_class_name) + + configured_pipeline_class = od_config.diffusers_pipeline_cls + configured_class_name = configured_pipeline_class.__name__ if configured_pipeline_class is not None else None + return get_pipeline_utils(configured_class_name) + + +def resolve_diffusers_pipeline_class(od_config: OmniDiffusionConfig) -> type[DiffusionPipeline] | None: + """Resolve the concrete Diffusers pipeline used for loading and capability checks.""" + pipeline_utils = get_pipeline_utils_for_config(od_config) + return pipeline_utils.resolve_pipeline_class(od_config, od_config.diffusers_pipeline_cls) diff --git a/vllm_omni/diffusion/models/sana_video/__init__.py b/vllm_omni/diffusion/models/sana_video/__init__.py new file mode 100644 index 000000000..cebefc4f8 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/__init__.py @@ -0,0 +1,32 @@ +# Copyright 2025 SANA-Video Authors and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .pipeline_output import SanaVideoPipelineOutput +from .pipeline_sana_video import SanaVideoPipeline, get_sana_video_post_process_func +from .pipeline_sana_video_i2v import ( + SanaImageToVideoPipeline, + get_sana_video_i2v_post_process_func, + get_sana_video_i2v_pre_process_func, +) +from .transformer_sana_video import SanaVideoTransformer3DModel + +__all__ = [ + "SanaImageToVideoPipeline", + "SanaVideoPipeline", + "SanaVideoPipelineOutput", + "SanaVideoTransformer3DModel", + "get_sana_video_i2v_post_process_func", + "get_sana_video_i2v_pre_process_func", + "get_sana_video_post_process_func", +] diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_output.py b/vllm_omni/diffusion/models/sana_video/pipeline_output.py new file mode 100644 index 000000000..6494b112c --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/pipeline_output.py @@ -0,0 +1,19 @@ +from dataclasses import dataclass + +import torch +from diffusers.utils import BaseOutput + + +@dataclass +class SanaVideoPipelineOutput(BaseOutput): + r""" + Output class for Sana-Video pipelines. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape + `(batch_size, num_frames, channels, height, width)`. + """ + + frames: torch.Tensor diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py new file mode 100644 index 000000000..a1fa38d43 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py @@ -0,0 +1,1055 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# Copyright 2025 SANA-Video Authors and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ruff: noqa: E501 + +import inspect +import os +import warnings +from collections.abc import Iterable + +import numpy as np +import PIL.Image +import torch +from diffusers.schedulers import DPMSolverMultistepScheduler +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from torch import nn +from transformers import AutoModel, Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast +from vllm.transformers_utils.config import get_hf_file_to_dict + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, +) +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.utils import _load_json +from vllm_omni.diffusion.postprocess import interpolate_video_tensor +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import resolve_video_num_frames +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +from .pipeline_output import SanaVideoPipelineOutput +from .transformer_sana_video import SanaVideoTransformer3DModel + + +def _resolve_vae_class_and_dtype( + vae_class_name: str, + pipeline_dtype: torch.dtype, +) -> tuple[type[DistributedAutoencoderKLLTX2Video | DistributedAutoencoderKLWan], torch.dtype]: + vae_classes = { + "AutoencoderKLLTX2Video": DistributedAutoencoderKLLTX2Video, + "AutoencoderKLWan": DistributedAutoencoderKLWan, + } + if vae_class_name not in vae_classes: + raise ValueError(f"Unsupported SANA-Video VAE class: {vae_class_name!r}") + vae_dtype = torch.float32 if vae_class_name == "AutoencoderKLWan" else pipeline_dtype + return vae_classes[vae_class_name], vae_dtype + + +def _load_sana_tokenizer( + model: str, + component_subfolders: list[str], + local_files_only: bool, +) -> GemmaTokenizer: + # Transformers v5 AutoTokenizer first tries to resolve an AutoConfig from + # the requested subfolder. Released SANA-Video repositories correctly + # contain tokenizer metadata under tokenizer/, but their model config lives + # under text_encoder/, so that extra tokenizer/config.json lookup fails for + # Hub model IDs. Both released checkpoints use GemmaTokenizer; loading the + # concrete class also matches Diffusers and avoids the unrelated config + # lookup. + return from_pretrained_with_prefetch( + GemmaTokenizer.from_pretrained, + model, + subfolder="tokenizer", + prefetch_list=component_subfolders, + local_files_only=local_files_only, + ) + + +ASPECT_RATIO_480_BIN = { + "0.5": [448.0, 896.0], + "0.57": [480.0, 832.0], + "0.68": [528.0, 768.0], + "0.78": [560.0, 720.0], + "1.0": [624.0, 624.0], + "1.13": [672.0, 592.0], + "1.29": [720.0, 560.0], + "1.46": [768.0, 528.0], + "1.67": [816.0, 496.0], + "1.75": [832.0, 480.0], + "2.0": [896.0, 448.0], +} + + +ASPECT_RATIO_720_BIN = { + "0.5": [672.0, 1344.0], + "0.57": [704.0, 1280.0], + "0.68": [800.0, 1152.0], + "0.78": [832.0, 1088.0], + "1.0": [960.0, 960.0], + "1.13": [1024.0, 896.0], + "1.29": [1088.0, 832.0], + "1.46": [1152.0, 800.0], + "1.67": [1248.0, 736.0], + "1.75": [1280.0, 704.0], + "2.0": [1344.0, 672.0], +} + + +SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION = ( + "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for video generation. Evaluate the level of detail in the user prompt:", + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, motion, and temporal relationships to create vivid and dynamic scenes.", + "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", + "Here are examples of how to transform or refine prompts:", + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat slowly settling into a curled position, peacefully falling asleep on a warm sunny windowsill, with gentle sunlight filtering through surrounding pots of blooming red flowers.", + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps gradually lighting up, a diverse crowd of people in colorful clothing walking past, and a double-decker bus smoothly passing by towering glass skyscrapers.", + "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", + "User Prompt: ", +) + + +def get_sana_video_default_resolution(sample_size: int) -> tuple[int, int]: + if sample_size == 30: + return 480, 832 + if sample_size == 22: + return 704, 1280 + raise ValueError(f"Unsupported SANA-Video sample size: {sample_size}") + + +def resolve_sana_video_sample_size(od_config: OmniDiffusionConfig) -> int: + tf_model_config = getattr(od_config, "tf_model_config", None) + sample_size = tf_model_config.get("sample_size", None) if tf_model_config is not None else None + model = getattr(od_config, "model", None) + if sample_size is None and model: + try: + transformer_config = get_hf_file_to_dict("transformer/config.json", model) or {} + sample_size = transformer_config.get("sample_size") + except (OSError, ValueError): + pass + if sample_size is None: + sample_size = 22 if model and "720p" in str(model).lower() else 30 + return int(sample_size) + + +def get_sana_video_post_process_func(od_config: OmniDiffusionConfig): + video_processor = VideoProcessor() + is_diffusers_adapter = od_config.diffusion_load_format == "diffusers" + + def post_process_func(video: torch.Tensor, output_type: str = "np", sampling_params=None): + if output_type == "latent" or ( + sampling_params is not None and getattr(sampling_params, "output_type", None) == "latent" + ): + return video + + first_frame = video + while isinstance(first_frame, (list, tuple)) and first_frame: + first_frame = first_frame[0] + is_postprocessed_frame_sequence = isinstance(video, np.ndarray) or isinstance( + first_frame, + (PIL.Image.Image, np.ndarray), + ) + if is_diffusers_adapter and is_postprocessed_frame_sequence: + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + raise ValueError( + "Frame interpolation is not supported for SANA-Video with the Diffusers adapter because " + "Diffusers returns already postprocessed frames." + ) + return {"payload": {"video": video}, "metadata": {}} + + video_metadata = {} + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + video, multiplier = interpolate_video_tensor( + video, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + video_metadata["video_fps_multiplier"] = multiplier + return { + "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)}, + "metadata": {"video": video_metadata} if video_metadata else {}, + } + + return post_process_func + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +) -> tuple[torch.Tensor, int]: + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + scheduler_device = torch.device("cpu") + # vLLM initializes models under a CUDA default-device context. Diffusers' + # schedulers are not nn.Modules, so their tensors can be created on CUDA and + # are not moved by pipeline.to(). DPM solvers also perform NumPy operations + # while constructing their schedule. Move scheduler state and temporarily + # override the global default device so every intermediate remains on CPU. + for name, value in vars(scheduler).items(): + if isinstance(value, torch.Tensor) and value.device.type != "cpu": + setattr(scheduler, name, value.cpu()) + previous_default_device = torch.get_default_device() + try: + torch.set_default_device(scheduler_device) + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=scheduler_device, **kwargs) + num_inference_steps = len(scheduler.timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=scheduler_device, **kwargs) + num_inference_steps = len(scheduler.timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=scheduler_device, **kwargs) + finally: + torch.set_default_device(previous_default_device) + scheduler.timesteps = scheduler.timesteps.to(device) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class SanaVideoPipeline( + nn.Module, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, +): + r""" + vLLM-Omni pipeline for text-to-video generation using + [Sana](https://huggingface.co/papers/2509.24695). + + Args: + tokenizer ([`GemmaTokenizer`] or [`GemmaTokenizerFast`]): + The tokenizer used to tokenize the prompt. + text_encoder ([`Gemma2PreTrainedModel`]): + Text encoder model to encode the input prompts. + vae ([`DistributedAutoencoderKLWan`] or [`DistributedAutoencoderKLLTX2Video`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + transformer ([`SanaVideoTransformer3DModel`]): + Conditional Transformer to denoise the input latents. + scheduler ([`DPMSolverMultistepScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded video latents. + """ + + _dit_modules = ["transformer"] + _encoder_modules = ["text_encoder"] + _vae_modules = ["vae"] + supports_step_execution = False + + def __init__( + self, + tokenizer: GemmaTokenizer | GemmaTokenizerFast | None = None, + text_encoder: Gemma2PreTrainedModel | None = None, + vae: DistributedAutoencoderKLLTX2Video | DistributedAutoencoderKLWan | None = None, + transformer: SanaVideoTransformer3DModel | None = None, + scheduler: DPMSolverMultistepScheduler | None = None, + *, + od_config: OmniDiffusionConfig | None = None, + prefix: str = "", + ): + super().__init__() + + self.od_config = od_config + self.device = get_local_device() + self.weights_sources: list[DiffusersPipelineLoader.ComponentSource] = [] + if od_config is not None: + tokenizer, text_encoder, vae, transformer, scheduler = self._load_components(od_config, prefix) + + if tokenizer is None or text_encoder is None or vae is None or transformer is None or scheduler is None: + raise ValueError("SanaVideoPipeline requires tokenizer, text_encoder, vae, transformer, and scheduler.") + + self.tokenizer = tokenizer + self.text_encoder = text_encoder + self.vae = vae + self.transformer = transformer + self.scheduler = scheduler + if od_config is None: + first_parameter = next(self.transformer.parameters(), None) + if first_parameter is not None: + self.device = first_parameter.device + + if isinstance(self.vae, DistributedAutoencoderKLLTX2Video): + self.vae_scale_factor_temporal = self.vae.config.temporal_compression_ratio + self.vae_scale_factor_spatial = self.vae.config.spatial_compression_ratio + elif isinstance(self.vae, DistributedAutoencoderKLWan): + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial + else: + raise TypeError(f"Unsupported SANA-Video VAE type: {type(self.vae).__name__}") + + self.vae_scale_factor = self.vae_scale_factor_spatial + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + if od_config is not None: + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, + profiler_targets=["forward"], + ) + + def _load_components( + self, + od_config: OmniDiffusionConfig, + prefix: str, + ) -> tuple[ + GemmaTokenizer | GemmaTokenizerFast, + Gemma2PreTrainedModel, + DistributedAutoencoderKLLTX2Video | DistributedAutoencoderKLWan, + SanaVideoTransformer3DModel, + DPMSolverMultistepScheduler, + ]: + model = od_config.model + local_files_only = os.path.exists(model) + dtype = getattr(od_config, "dtype", torch.bfloat16) + device = self.device + # Transformer weights are streamed by DiffusersPipelineLoader below. + # Prefetch only components loaded through from_pretrained here. + component_subfolders = ["tokenizer", "text_encoder", "vae", "scheduler"] + prefetch_subfolders(model, component_subfolders, local_files_only=local_files_only) + + tokenizer = _load_sana_tokenizer( + model, + component_subfolders, + local_files_only=local_files_only, + ) + text_encoder = from_pretrained_with_prefetch( + AutoModel.from_pretrained, + model, + subfolder="text_encoder", + prefetch_list=component_subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ).to(device) + + model_index = _load_json(model, "model_index.json", local_files_only) + vae_class_name = model_index.get("vae", [None, "AutoencoderKLWan"])[1] + # The released 480p checkpoint uses Wan VAE weights that are decoded + # in FP32 in the upstream reference recipe. The 720p LTX2 VAE is + # validated in BF16. + vae_class, vae_dtype = _resolve_vae_class_and_dtype(vae_class_name, dtype) + vae = from_pretrained_with_prefetch( + vae_class.from_pretrained, + model, + subfolder="vae", + prefetch_list=component_subfolders, + local_files_only=local_files_only, + torch_dtype=vae_dtype, + ).to(device) + + transformer_config = _load_json(model, "transformer/config.json", local_files_only) + transformer = SanaVideoTransformer3DModel.from_config(transformer_config).to(dtype=dtype, device=device) + scheduler = DPMSolverMultistepScheduler.from_pretrained( + model, + subfolder="scheduler", + local_files_only=local_files_only, + ) + self.weights_sources[:] = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model, + subfolder="transformer", + revision=None, + prefix=f"{prefix}transformer." if prefix else "transformer.", + fall_back_to_pt=True, + ) + ] + return tokenizer, text_encoder, vae, transformer, scheduler + + # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds + def _get_gemma_prompt_embeds( + self, + prompt: str | list[str], + device: torch.device, + dtype: torch.dtype, + clean_caption: bool = False, + max_sequence_length: int = 300, + complex_human_instruction: list[str] | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + device: (`torch.device`, *optional*): + torch device to place the resulting embeddings on + clean_caption (`bool`, defaults to `False`): + If `True`, the function will preprocess and clean the provided caption before encoding. + max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. + complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): + If `complex_human_instruction` is not empty, the function will use the complex Human instruction for + the prompt. + """ + prompt = [prompt] if isinstance(prompt, str) else prompt + + if getattr(self, "tokenizer", None) is not None: + self.tokenizer.padding_side = "right" + + prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) + + # prepare complex human instruction + if not complex_human_instruction: + max_length_all = max_sequence_length + else: + chi_prompt = "\n".join(complex_human_instruction) + prompt = [chi_prompt + p for p in prompt] + num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt)) + max_length_all = num_chi_prompt_tokens + max_sequence_length - 2 + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_length_all, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + + prompt_attention_mask = text_inputs.attention_mask + prompt_attention_mask = prompt_attention_mask.to(device) + + prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) + prompt_embeds = prompt_embeds[0].to(dtype=dtype, device=device) + + return prompt_embeds, prompt_attention_mask + + def encode_prompt( + self, + prompt: str | list[str], + do_classifier_free_guidance: bool = True, + negative_prompt: str = "", + num_videos_per_prompt: int = 1, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + clean_caption: bool = False, + max_sequence_length: int = 300, + complex_human_instruction: list[str] | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `list[str]`, *optional*): + The prompt not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` + instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For + PixArt-Alpha, this should be "". + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + whether to use classifier free guidance or not + num_videos_per_prompt (`int`, *optional*, defaults to 1): + number of videos that should be generated per prompt + device: (`torch.device`, *optional*): + torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. For Sana, it's should be the embeddings of the "" string. + clean_caption (`bool`, defaults to `False`): + If `True`, the function will preprocess and clean the provided caption before encoding. + max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. + complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): + If `complex_human_instruction` is not empty, the function will use the complex Human instruction for + the prompt. + """ + + if device is None: + device = self.device + + if self.text_encoder is not None: + dtype = self.text_encoder.dtype + else: + dtype = None + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if getattr(self, "tokenizer", None) is not None: + self.tokenizer.padding_side = "right" + + # See Section 3.1. of the paper. + max_length = max_sequence_length + select_index = [0] + list(range(-max_length + 1, 0)) + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + device=device, + dtype=dtype, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=complex_human_instruction, + ) + + prompt_embeds = prompt_embeds[:, select_index] + prompt_attention_mask = prompt_attention_mask[:, select_index] + + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) + prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = [negative_prompt] * batch_size if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=negative_prompt, + device=device, + dtype=dtype, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=False, + ) + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device) + + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1) + negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_videos_per_prompt, 1) + else: + negative_prompt_embeds = None + negative_prompt_attention_mask = None + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + patch_size = getattr(self.transformer.config, "patch_size", (1, 1, 1)) + spatial_alignment = self.vae_scale_factor_spatial * max(patch_size[-2:]) + if height % spatial_alignment != 0 or width % spatial_alignment != 0: + raise ValueError( + f"`height` and `width` have to be divisible by {spatial_alignment} but are {height} and {width}." + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}." + ) + + def _text_preprocessing(self, text, clean_caption=False): + if clean_caption: + raise ValueError( + "SANA-Video native pipeline does not support `clean_caption=True`; " + "pre-clean the prompt or use the Diffusers adapter." + ) + if not isinstance(text, (tuple, list)): + text = [text] + return [str(item).lower().strip() for item in text] + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int = 16, + height: int = 480, + width: int = 832, + num_frames: int = 81, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + transformer_weights = ( + (name.removeprefix("transformer."), tensor) for name, tensor in weights if name.startswith("transformer.") + ) + return {f"transformer.{name}" for name in self.transformer.load_weights(transformer_weights)} + + def __call__(self, req: DiffusionRequestBatch) -> DiffusionOutput: + return self.forward(req) + + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + if len(req.prompts) != 1: + raise ValueError("SANA-Video currently supports exactly one prompt per request.") + prompt_obj = req.prompts[0] + prompt = prompt_obj if isinstance(prompt_obj, str) else (prompt_obj.get("prompt") or "") + negative_prompt = "" if isinstance(prompt_obj, str) else (prompt_obj.get("negative_prompt") or "") + if not prompt: + raise ValueError("Prompt is required for SANA-Video generation.") + + sampling = req.sampling_params + default_height, default_width = get_sana_video_default_resolution(self.transformer.config.sample_size) + height = sampling.height or default_height + width = sampling.width or default_width + num_frames = resolve_video_num_frames( + sampling.num_frames, + default_num_frames=81, + is_dummy_run=req.is_dummy_run(), + ) + num_steps = sampling.num_inference_steps if sampling.num_inference_steps is not None else 50 + guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 6.0 + generator = sampling.generator + if generator is None and sampling.seed is not None: + generator = torch.Generator(device=get_local_device()).manual_seed(sampling.seed) + + extra_args = sampling.extra_args or {} + motion_score = extra_args.get("motion_score") + if motion_score is not None: + prompt = f"{prompt} motion score: {int(motion_score)}." + + result = self._generate( + prompt=prompt, + negative_prompt=negative_prompt, + num_inference_steps=num_steps, + timesteps=sampling.timesteps, + sigmas=sampling.sigmas, + guidance_scale=guidance_scale, + num_videos_per_prompt=sampling.num_outputs_per_prompt or 1, + height=height, + width=width, + frames=num_frames, + eta=sampling.eta, + generator=generator, + latents=sampling.latents, + output_type="latent" if sampling.output_type == "latent" else "raw", + clean_caption=bool(extra_args.get("clean_caption", False)), + use_resolution_binning=bool(extra_args.get("use_resolution_binning", True)), + max_sequence_length=sampling.max_sequence_length or 300, + ) + return DiffusionOutput(output=result.frames) + + @torch.no_grad() + def _generate( + self, + prompt: str | list[str] = None, + negative_prompt: str = "", + num_inference_steps: int = 50, + timesteps: list[int] = None, + sigmas: list[float] = None, + guidance_scale: float = 6.0, + num_videos_per_prompt: int | None = 1, + height: int = 480, + width: int = 832, + frames: int = 81, + eta: float = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + clean_caption: bool = False, + use_resolution_binning: bool = True, + max_sequence_length: int = 300, + complex_human_instruction: list[str] | None = None, + ) -> SanaVideoPipelineOutput | tuple: + """ + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. + instead. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the video generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality video at the + expense of slower inference. + timesteps (`list[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + sigmas (`list[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 4.5): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate videos that are closely linked to + the text `prompt`, usually at the expense of lower video quality. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + height (`int`, *optional*, defaults to 480): + The height in pixels of the generated video. + width (`int`, *optional*, defaults to 832): + The width in pixels of the generated video. + frames (`int`, *optional*, defaults to 81): + The number of frames in the generated video. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only + applies to [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + prompt_attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask for text embeddings. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + negative_prompt_attention_mask (`torch.Tensor`, *optional*): + Pre-generated attention mask for negative text embeddings. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated video. Choose between mp4 or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`SanaVideoPipelineOutput`] instead of a plain tuple. + clean_caption (`bool`, *optional*, defaults to `True`): + Advanced caption cleaning is not supported by the native pipeline. Passing `True` raises an error; + pre-clean the prompt or use the Diffusers adapter. + use_resolution_binning (`bool` defaults to `True`): + If set to `True`, the requested height and width are first mapped to the closest resolutions using + `ASPECT_RATIO_480_BIN` or `ASPECT_RATIO_720_BIN`. After the produced latents are decoded into videos, + they are resized back to the requested resolution. Useful for generating non-square videos. + max_sequence_length (`int` defaults to `300`): + Maximum sequence length to use with the `prompt`. + complex_human_instruction (`list[str]`, *optional*): + Instructions for complex human attention: + https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. + + Examples: + + Returns: + [`~pipelines.sana_video.pipeline_output.SanaVideoPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.sana_video.pipeline_output.SanaVideoPipelineOutput`] is + returned, otherwise a `tuple` is returned where the first element is a list with the generated videos + """ + if complex_human_instruction is None: + complex_human_instruction = list(SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION) + + # 1. Check inputs. Raise error if not correct + if use_resolution_binning: + if self.transformer.config.sample_size == 30: + aspect_ratio_bin = ASPECT_RATIO_480_BIN + elif self.transformer.config.sample_size == 22: + aspect_ratio_bin = ASPECT_RATIO_720_BIN + else: + raise ValueError("Invalid sample size") + orig_height, orig_width = height, width + height, width = self.video_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) + + self.check_inputs( + prompt=prompt, + height=height, + width=width, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + + self._guidance_scale = guidance_scale + self._interrupt = False + + # 2. Default height and width to transformer + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self.device + + # 3. Encode input prompt + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt, + self.do_classifier_free_guidance, + negative_prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + device=device, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=complex_human_instruction, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + # 4. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps, sigmas + ) + + # 5. Prepare latents. + latent_channels = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + latent_channels, + height, + width, + frames, + torch.float32, + device, + generator, + latents, + ) + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + transformer_dtype = self.transformer.dtype + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + + # predict noise model_output + noise_pred = self.transformer( + latent_model_input.to(dtype=transformer_dtype), + encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype), + encoder_attention_mask=prompt_attention_mask, + timestep=timestep, + return_dict=False, + )[0] + noise_pred = noise_pred.float() + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # learned sigma + if self.transformer.config.out_channels // 2 == latent_channels: + noise_pred = noise_pred.chunk(2, dim=1)[0] + + # compute previous image: x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if output_type == "latent": + video = latents + else: + latents = latents.to(self.vae.dtype) + if isinstance(self.vae, DistributedAutoencoderKLLTX2Video): + latents_mean = self.vae.latents_mean + latents_std = self.vae.latents_std + z_dim = self.vae.config.latent_channels + elif isinstance(self.vae, DistributedAutoencoderKLWan): + latents_mean = torch.tensor(self.vae.config.latents_mean) + latents_std = torch.tensor(self.vae.config.latents_std) + z_dim = self.vae.config.z_dim + else: + latents_mean = torch.zeros(latents.shape[1], device=latents.device, dtype=latents.dtype) + latents_std = torch.ones(latents.shape[1], device=latents.device, dtype=latents.dtype) + z_dim = latents.shape[1] + + latents_mean = latents_mean.view(1, z_dim, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = 1.0 / latents_std.view(1, z_dim, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents / latents_std + latents_mean + try: + video = self.vae.decode(latents, return_dict=False)[0] + except torch.OutOfMemoryError as e: + warnings.warn( + f"{e}. \n" + f"Try to use VAE tiling for large images. For example: \n" + f"pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)" + ) + raise + + if use_resolution_binning: + video = self.video_processor.resize_and_crop_tensor(video, orig_width, orig_height) + + if output_type != "raw": + video = self.video_processor.postprocess_video(video, output_type=output_type) + + if not return_dict: + return (video,) + + return SanaVideoPipelineOutput(frames=video) diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py new file mode 100644 index 000000000..8d0fe9919 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import os +import warnings +from math import sqrt + +import PIL.Image +import torch +from diffusers.utils.torch_utils import randn_tensor + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, +) +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( + ASPECT_RATIO_480_BIN, + ASPECT_RATIO_720_BIN, + SanaVideoPipeline, + get_sana_video_default_resolution, + get_sana_video_post_process_func, + resolve_sana_video_sample_size, + retrieve_timesteps, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest, resolve_video_num_frames +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniTextPrompt + +SANA_VIDEO_I2V_COMPLEX_HUMAN_INSTRUCTION = [ + ( + "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for " + "video generation. Evaluate the level of detail in the user prompt:" + ), + ( + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, motion, and " + "temporal relationships to create vivid and dynamic scenes." + ), + ("- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating."), + "Here are examples of how to transform or refine prompts:", + ( + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat slowly settling into a curled position, " + "peacefully falling asleep on a warm sunny windowsill, with gentle sunlight filtering through surrounding " + "pots of blooming red flowers." + ), + ( + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street " + "lamps gradually lighting up, a diverse crowd of people in colorful clothing walking past, and a double-decker " + "bus smoothly passing by towering glass skyscrapers." + ), + ( + "Please generate only the enhanced description for the prompt below and avoid including any additional " + "commentary or evaluations:" + ), + "User Prompt: ", +] + + +def get_sana_video_i2v_post_process_func(od_config: OmniDiffusionConfig): + return get_sana_video_post_process_func(od_config) + + +def get_sana_video_i2v_pre_process_func(od_config: OmniDiffusionConfig): + default_height, default_width = get_sana_video_default_resolution(resolve_sana_video_sample_size(od_config)) + + def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: + prompt = request.prompt + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt, multi_modal_data={}) + + multi_modal_data = prompt.get("multi_modal_data") or {} + raw_image = multi_modal_data.get("image") + if raw_image is None: + raise ValueError('SANA-Video I2V requires `"multi_modal_data": {"image": }`.') + if isinstance(raw_image, str): + if not os.path.isfile(raw_image): + raise ValueError(f"SANA-Video I2V image path does not exist: {raw_image}") + image = PIL.Image.open(raw_image).convert("RGB") + elif isinstance(raw_image, PIL.Image.Image): + image = raw_image.convert("RGB") + else: + raise TypeError(f"Unsupported SANA-Video I2V image type: {type(raw_image).__name__}") + + if request.sampling_params.height is None or request.sampling_params.width is None: + max_area = default_height * default_width + aspect_ratio = image.height / image.width + inferred_height = round(sqrt(max_area * aspect_ratio)) // 32 * 32 + inferred_width = round(sqrt(max_area / aspect_ratio)) // 32 * 32 + if request.sampling_params.height is None: + request.sampling_params.height = inferred_height + if request.sampling_params.width is None: + request.sampling_params.width = inferred_width + + image = image.resize( + (request.sampling_params.width, request.sampling_params.height), + PIL.Image.Resampling.LANCZOS, + ) + multi_modal_data["image"] = image + prompt["multi_modal_data"] = multi_modal_data + request.prompt = prompt + return request + + return pre_process_func + + +def _retrieve_latents(encoder_output: torch.Tensor) -> torch.Tensor: + if hasattr(encoder_output, "latent_dist"): + return encoder_output.latent_dist.mode() + if hasattr(encoder_output, "latents"): + return encoder_output.latents + raise AttributeError("Could not access latents from the SANA-Video VAE encoder output.") + + +class SanaImageToVideoPipeline(SanaVideoPipeline, SupportImageInput): + """Native vLLM-Omni SANA-Video 2B image-to-video pipeline. + + The first latent frame is encoded from the input image and remains fixed; + denoising updates only the subsequent video frames. + """ + + def _prepare_i2v_latents( + self, + image: torch.Tensor, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + num_frames: int, + dtype: torch.dtype, + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None, + ) -> torch.Tensor: + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"Generator list length {len(generator)} does not match effective batch size {batch_size}." + ) + if latents is None: + latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype) + else: + latents = latents.to(device=self.device, dtype=dtype) + + image = image.unsqueeze(2).to(device=self.device, dtype=self.vae.dtype) + image_latents = _retrieve_latents(self.vae.encode(image)) + image_latents = image_latents.repeat(batch_size, 1, 1, 1, 1) + + if isinstance(self.vae, DistributedAutoencoderKLLTX2Video): + latents_mean = self.vae.latents_mean + latents_std = self.vae.latents_std + elif isinstance(self.vae, DistributedAutoencoderKLWan): + latents_mean = torch.tensor(self.vae.config.latents_mean) + latents_std = torch.tensor(self.vae.config.latents_std) + else: + raise TypeError(f"Unsupported SANA-Video I2V VAE type: {type(self.vae).__name__}") + + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(image_latents.device, image_latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(image_latents.device, image_latents.dtype) + image_latents = (image_latents - latents_mean) / latents_std + latents[:, :, :1] = image_latents.to(dtype) + return latents + + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + if len(req.prompts) != 1: + raise ValueError("SANA-Video I2V currently supports exactly one prompt per request.") + prompt_obj = req.prompts[0] + if isinstance(prompt_obj, str): + raise ValueError("SANA-Video I2V requires an image in `multi_modal_data`.") + + prompt = prompt_obj.get("prompt") or "" + negative_prompt = prompt_obj.get("negative_prompt") or "" + image = (prompt_obj.get("multi_modal_data") or {}).get("image") + if not prompt: + raise ValueError("Prompt is required for SANA-Video I2V generation.") + if not isinstance(image, PIL.Image.Image): + raise TypeError("SANA-Video I2V expects a preprocessed PIL image.") + sampling = req.sampling_params + default_height, default_width = get_sana_video_default_resolution(self.transformer.config.sample_size) + height = sampling.height or default_height + width = sampling.width or default_width + frames = resolve_video_num_frames( + sampling.num_frames, + default_num_frames=81, + is_dummy_run=req.is_dummy_run(), + ) + num_steps = sampling.num_inference_steps if sampling.num_inference_steps is not None else 50 + guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 6.0 + generator = sampling.generator + if generator is None and sampling.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(sampling.seed) + + extra_args = sampling.extra_args or {} + motion_score = extra_args.get("motion_score") + if motion_score is not None: + prompt = f"{prompt} motion score: {int(motion_score)}." + + video = self._generate_i2v( + image=image, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + frames=frames, + num_inference_steps=num_steps, + timesteps=sampling.timesteps, + sigmas=sampling.sigmas, + guidance_scale=guidance_scale, + num_videos_per_prompt=sampling.num_outputs_per_prompt or 1, + eta=sampling.eta, + generator=generator, + latents=sampling.latents, + clean_caption=bool(extra_args.get("clean_caption", False)), + use_resolution_binning=bool(extra_args.get("use_resolution_binning", True)), + max_sequence_length=sampling.max_sequence_length or 300, + output_type="latent" if sampling.output_type == "latent" else "raw", + ) + return DiffusionOutput(output=video) + + @torch.no_grad() + def _generate_i2v( + self, + *, + image: PIL.Image.Image, + prompt: str, + negative_prompt: str, + height: int, + width: int, + frames: int, + num_inference_steps: int, + guidance_scale: float, + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None, + use_resolution_binning: bool, + max_sequence_length: int, + timesteps: list[int] | torch.Tensor | None = None, + sigmas: list[float] | None = None, + num_videos_per_prompt: int = 1, + eta: float = 0.0, + clean_caption: bool = False, + output_type: str = "raw", + complex_human_instruction: list[str] = SANA_VIDEO_I2V_COMPLEX_HUMAN_INSTRUCTION, + ) -> torch.Tensor: + orig_height, orig_width = height, width + if use_resolution_binning: + if self.transformer.config.sample_size == 30: + aspect_ratio_bin = ASPECT_RATIO_480_BIN + elif self.transformer.config.sample_size == 22: + aspect_ratio_bin = ASPECT_RATIO_720_BIN + else: + raise ValueError(f"Unsupported SANA-Video sample size: {self.transformer.config.sample_size}") + height, width = self.video_processor.classify_height_width_bin( + height, + width, + ratios=aspect_ratio_bin, + ) + + self.check_inputs( + prompt=prompt, + height=height, + width=width, + negative_prompt=negative_prompt, + ) + self._guidance_scale = guidance_scale + self._interrupt = False + + prompt_embeds, prompt_mask, negative_embeds, negative_mask = self.encode_prompt( + prompt, + self.do_classifier_free_guidance, + negative_prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + device=self.device, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=complex_human_instruction, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_embeds, prompt_embeds]) + prompt_mask = torch.cat([negative_mask, prompt_mask]) + + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps=num_inference_steps, + device=self.device, + timesteps=timesteps, + sigmas=sigmas, + ) + image_tensor = self.video_processor.preprocess(image, height=height, width=width).to( + self.device, + dtype=torch.float32, + ) + latent_channels = self.transformer.config.in_channels + latents = self._prepare_i2v_latents( + image_tensor, + batch_size=num_videos_per_prompt, + num_channels_latents=latent_channels, + height=height, + width=width, + num_frames=frames, + dtype=torch.float32, + generator=generator, + latents=latents, + ) + + patch_t, patch_h, patch_w = self.transformer.config.patch_size + conditioning_mask = latents.new_zeros( + num_videos_per_prompt, + 1, + latents.shape[2] // patch_t, + latents.shape[3] // patch_h, + latents.shape[4] // patch_w, + ) + conditioning_mask[:, :, 0] = 1 + if self.do_classifier_free_guidance: + conditioning_mask = torch.cat([conditioning_mask, conditioning_mask]) + + self._num_timesteps = len(timesteps) + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + transformer_dtype = self.transformer.dtype + with self.progress_bar(total=num_inference_steps) as progress_bar: + for index, timestep_value in enumerate(timesteps): + latent_input = torch.cat([latents, latents]) if self.do_classifier_free_guidance else latents + timestep = timestep_value.expand(conditioning_mask.shape) * (1 - conditioning_mask) + noise_pred = self.transformer( + latent_input.to(transformer_dtype), + encoder_hidden_states=prompt_embeds.to(transformer_dtype), + encoder_attention_mask=prompt_mask, + timestep=timestep, + return_dict=False, + )[0].float() + if self.do_classifier_free_guidance: + uncond, cond = noise_pred.chunk(2) + noise_pred = uncond + guidance_scale * (cond - uncond) + if self.transformer.config.out_channels // 2 == latent_channels: + noise_pred = noise_pred.chunk(2, dim=1)[0] + + denoised = self.scheduler.step( + noise_pred[:, :, 1:], + timestep_value, + latents[:, :, 1:], + **extra_step_kwargs, + return_dict=False, + )[0] + latents = torch.cat([latents[:, :, :1], denoised], dim=2) + if index == len(timesteps) - 1 or (index + 1) % self.scheduler.order == 0: + progress_bar.update() + + if output_type == "latent": + return latents + + latents = latents.to(self.vae.dtype) + if isinstance(self.vae, DistributedAutoencoderKLLTX2Video): + latents_mean = self.vae.latents_mean + latents_std = self.vae.latents_std + latent_dim = self.vae.config.latent_channels + else: + latents_mean = torch.tensor(self.vae.config.latents_mean) + latents_std = torch.tensor(self.vae.config.latents_std) + latent_dim = self.vae.config.z_dim + latents_mean = latents_mean.view(1, latent_dim, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, latent_dim, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std + latents_mean + + try: + video = self.vae.decode(latents, return_dict=False)[0] + except torch.OutOfMemoryError as error: + warnings.warn( + f"{error}. Enable VAE tiling for SANA-Video I2V high-resolution decoding.", + stacklevel=2, + ) + raise + if use_resolution_binning: + video = self.video_processor.resize_and_crop_tensor(video, orig_width, orig_height) + return video diff --git a/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py b/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py new file mode 100644 index 000000000..eb0f5843b --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py @@ -0,0 +1,1066 @@ +# Copyright 2025 The HuggingFace Team and SANA-Video Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import numbers +from collections.abc import Iterable +from dataclasses import dataclass, fields + +import torch +from torch import nn +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention as OmniAttention + + +@dataclass +class SanaVideoTransformerConfig: + in_channels: int = 16 + out_channels: int | None = 16 + num_attention_heads: int = 20 + attention_head_dim: int = 112 + num_layers: int = 20 + num_cross_attention_heads: int | None = 20 + cross_attention_head_dim: int | None = 112 + cross_attention_dim: int | None = 2240 + caption_channels: int = 2304 + mlp_ratio: float = 2.5 + dropout: float = 0.0 + attention_bias: bool = False + sample_size: int = 30 + patch_size: tuple[int, int, int] = (1, 2, 2) + norm_elementwise_affine: bool = False + norm_eps: float = 1e-6 + interpolation_scale: int | None = None + guidance_embeds: bool = False + guidance_embeds_scale: float = 0.1 + qk_norm: str | None = "rms_norm_across_heads" + rope_max_seq_len: int = 1024 + + @classmethod + def from_dict(cls, config: dict) -> "SanaVideoTransformerConfig": + field_names = {field.name for field in fields(cls)} + unknown = {key for key in config if key not in field_names and not key.startswith("_")} + if unknown: + raise ValueError(f"Unsupported SANA-Video transformer config fields: {sorted(unknown)}") + values = {key: value for key, value in config.items() if key in field_names} + if "patch_size" in values: + values["patch_size"] = tuple(values["patch_size"]) + return cls(**values) + + +@dataclass +class SanaVideoTransformerOutput: + sample: torch.Tensor + + +class SanaRMSNorm(nn.Module): + """RMSNorm with the parameter and dtype semantics used by SANA checkpoints.""" + + def __init__( + self, + dim: int | Iterable[int], + eps: float, + elementwise_affine: bool = True, + bias: bool = False, + ) -> None: + super().__init__() + + self.eps = eps + self.elementwise_affine = elementwise_affine + + if isinstance(dim, numbers.Integral): + dim = (dim,) + self.dim = torch.Size(dim) + + self.weight: nn.Parameter | None = None + self.bias: nn.Parameter | None = None + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(self.dim)) + if bias: + self.bias = nn.Parameter(torch.zeros(self.dim)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + + if self.weight is not None: + if self.weight.dtype in (torch.float16, torch.bfloat16): + hidden_states = hidden_states.to(self.weight.dtype) + hidden_states = hidden_states * self.weight + if self.bias is not None: + hidden_states = hidden_states + self.bias + else: + hidden_states = hidden_states.to(input_dtype) + + return hidden_states + + +def _get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +) -> torch.Tensor: + """Create the sinusoidal timestep embeddings used by SANA checkpoints.""" + if timesteps.ndim != 1: + raise ValueError(f"Timesteps must be a 1D tensor, got shape {tuple(timesteps.shape)}") + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange( + start=0, + end=half_dim, + dtype=torch.float32, + device=timesteps.device, + ) + exponent = exponent / (half_dim - downscale_freq_shift) + embeddings = timesteps[:, None].float() * torch.exp(exponent)[None, :] + embeddings = scale * embeddings + embeddings = torch.cat([torch.sin(embeddings), torch.cos(embeddings)], dim=-1) + + if flip_sin_to_cos: + embeddings = torch.cat([embeddings[:, half_dim:], embeddings[:, :half_dim]], dim=-1) + if embedding_dim % 2 == 1: + embeddings = torch.nn.functional.pad(embeddings, (0, 1, 0, 0)) + + return embeddings + + +class SanaTimesteps(nn.Module): + """Parameter-free timestep projection compatible with Diffusers Timesteps.""" + + def __init__( + self, + num_channels: int, + flip_sin_to_cos: bool, + downscale_freq_shift: float, + scale: int = 1, + ) -> None: + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + return _get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + ) + + +class SanaTimestepEmbedding(nn.Module): + """Two-layer timestep MLP with Diffusers-compatible parameter names.""" + + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int | None = None, + post_act_fn: str | None = None, + cond_proj_dim: int | None = None, + sample_proj_bias: bool = True, + ) -> None: + super().__init__() + if act_fn != "silu": + raise ValueError(f"SANA timestep embeddings only support act_fn='silu', got {act_fn!r}") + if post_act_fn not in (None, "silu"): + raise ValueError(f"SANA timestep embeddings only support post_act_fn=None or 'silu', got {post_act_fn!r}") + + self.linear_1 = nn.Linear(in_channels, time_embed_dim, bias=sample_proj_bias) + self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) if cond_proj_dim is not None else None + self.act = nn.SiLU() + self.linear_2 = nn.Linear( + time_embed_dim, + out_dim if out_dim is not None else time_embed_dim, + bias=sample_proj_bias, + ) + self.post_act = nn.SiLU() if post_act_fn == "silu" else None + + def forward( + self, + sample: torch.Tensor, + condition: torch.Tensor | None = None, + ) -> torch.Tensor: + if condition is not None: + if self.cond_proj is None: + raise ValueError("condition was provided but cond_proj_dim was not configured") + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + sample = self.act(sample) + sample = self.linear_2(sample) + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class SanaCombinedTimestepSizeEmbeddings(nn.Module): + """PixArt timestep/size conditioning used by SANA's adaLN-single.""" + + def __init__( + self, + embedding_dim: int, + size_emb_dim: int, + use_additional_conditions: bool = False, + ) -> None: + super().__init__() + self.outdim = size_emb_dim + self.time_proj = SanaTimesteps( + num_channels=256, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + self.timestep_embedder = SanaTimestepEmbedding( + in_channels=256, + time_embed_dim=embedding_dim, + ) + + self.use_additional_conditions = use_additional_conditions + if use_additional_conditions: + self.additional_condition_proj = SanaTimesteps( + num_channels=256, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + self.resolution_embedder = SanaTimestepEmbedding( + in_channels=256, + time_embed_dim=size_emb_dim, + ) + self.aspect_ratio_embedder = SanaTimestepEmbedding( + in_channels=256, + time_embed_dim=size_emb_dim, + ) + + def forward( + self, + timestep: torch.Tensor, + resolution: torch.Tensor | None, + aspect_ratio: torch.Tensor | None, + batch_size: int | None, + hidden_dtype: torch.dtype | None, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) + + if not self.use_additional_conditions: + return timesteps_emb + if resolution is None or aspect_ratio is None or batch_size is None: + raise ValueError( + "resolution, aspect_ratio, and batch_size are required when use_additional_conditions=True" + ) + + resolution_emb = self.additional_condition_proj(resolution.flatten()).to(hidden_dtype) + resolution_emb = self.resolution_embedder(resolution_emb).reshape(batch_size, -1) + aspect_ratio_emb = self.additional_condition_proj(aspect_ratio.flatten()).to(hidden_dtype) + aspect_ratio_emb = self.aspect_ratio_embedder(aspect_ratio_emb).reshape(batch_size, -1) + return timesteps_emb + torch.cat([resolution_emb, aspect_ratio_emb], dim=1) + + +class SanaAdaLayerNormSingle(nn.Module): + """SANA's checkpoint-compatible adaLN-single conditioning module.""" + + def __init__( + self, + embedding_dim: int, + use_additional_conditions: bool = False, + ) -> None: + super().__init__() + self.emb = SanaCombinedTimestepSizeEmbeddings( + embedding_dim, + size_emb_dim=embedding_dim // 3, + use_additional_conditions=use_additional_conditions, + ) + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + + def forward( + self, + timestep: torch.Tensor, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + batch_size: int | None = None, + hidden_dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + added_cond_kwargs = added_cond_kwargs or { + "resolution": None, + "aspect_ratio": None, + } + embedded_timestep = self.emb( + timestep, + **added_cond_kwargs, + batch_size=batch_size, + hidden_dtype=hidden_dtype, + ) + return self.linear(self.silu(embedded_timestep)), embedded_timestep + + +class _SanaFP32SiLU(nn.Module): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.silu(inputs.float(), inplace=False).to(inputs.dtype) + + +class SanaPixArtAlphaTextProjection(nn.Module): + """Caption projection with Diffusers-compatible parameters and activations.""" + + def __init__( + self, + in_features: int, + hidden_size: int, + out_features: int | None = None, + act_fn: str = "gelu_tanh", + ) -> None: + super().__init__() + out_features = hidden_size if out_features is None else out_features + self.linear_1 = nn.Linear(in_features, hidden_size, bias=True) + if act_fn == "gelu_tanh": + self.act_1 = nn.GELU(approximate="tanh") + elif act_fn == "silu": + self.act_1 = nn.SiLU() + elif act_fn == "silu_fp32": + self.act_1 = _SanaFP32SiLU() + else: + raise ValueError(f"Unknown activation function: {act_fn}") + self.linear_2 = nn.Linear(hidden_size, out_features, bias=True) + + def forward(self, caption: torch.Tensor) -> torch.Tensor: + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + return self.linear_2(hidden_states) + + +def _get_1d_rotary_pos_embed( + dim: int, + max_seq_len: int, + theta: float, + freqs_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate the real-valued 1D RoPE representation used by SANA-Video.""" + if dim % 2 != 0: + raise ValueError(f"RoPE dimension must be even, got {dim}") + + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype) / dim)) + positions = torch.arange(max_seq_len, dtype=freqs_dtype) + freqs = torch.outer(positions, freqs) + freqs_cos = freqs.cos().float().repeat_interleave(2, dim=-1) + freqs_sin = freqs.sin().float().repeat_interleave(2, dim=-1) + return freqs_cos, freqs_sin + + +class GLUMBTempConv(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + expand_ratio: float = 4, + norm_type: str | None = None, + residual_connection: bool = True, + ) -> None: + super().__init__() + + hidden_channels = int(expand_ratio * in_channels) + self.norm_type = norm_type + self.residual_connection = residual_connection + + self.nonlinearity = nn.SiLU() + self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0) + self.conv_depth = nn.Conv2d(hidden_channels * 2, hidden_channels * 2, 3, 1, 1, groups=hidden_channels * 2) + self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False) + + self.norm = None + if norm_type == "rms_norm": + self.norm = SanaRMSNorm(out_channels, eps=1e-5, elementwise_affine=True, bias=True) + + self.conv_temp = nn.Conv2d(out_channels, out_channels, kernel_size=(3, 1), stride=1, padding=(1, 0), bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.residual_connection: + residual = hidden_states + batch_size, num_frames, height, width, num_channels = hidden_states.shape + hidden_states = hidden_states.view(batch_size * num_frames, height, width, num_channels).permute(0, 3, 1, 2) + + hidden_states = self.conv_inverted(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.conv_depth(hidden_states) + hidden_states, gate = torch.chunk(hidden_states, 2, dim=1) + hidden_states = hidden_states * self.nonlinearity(gate) + + hidden_states = self.conv_point(hidden_states) + + # Temporal aggregation + hidden_states_temporal = hidden_states.view(batch_size, num_frames, num_channels, height * width).permute( + 0, 2, 1, 3 + ) + hidden_states = hidden_states_temporal + self.conv_temp(hidden_states_temporal) + hidden_states = hidden_states.permute(0, 2, 3, 1).view(batch_size, num_frames, height, width, num_channels) + + if self.norm_type == "rms_norm": + # move channel to the last dimension so we apply RMSnorm across channel dimension + hidden_states = self.norm(hidden_states.movedim(1, -1)).movedim(-1, 1) + + if self.residual_connection: + hidden_states = hidden_states + residual + + return hidden_states + + +class SanaLinearAttention(nn.Module): + """SANA's ReLU-kernel linear self-attention. + + This is intentionally not routed through the unified softmax attention + layer because doing so would change the model's attention algorithm. + """ + + def __init__( + self, + dim: int, + num_heads: int, + head_dim: int, + dropout: float, + bias: bool, + qk_norm: str | None, + ) -> None: + super().__init__() + inner_dim = num_heads * head_dim + self.heads = num_heads + self.to_q = nn.Linear(dim, inner_dim, bias=bias) + self.to_k = nn.Linear(dim, inner_dim, bias=bias) + self.to_v = nn.Linear(dim, inner_dim, bias=bias) + self.norm_q = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None + self.norm_k = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None + self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=True), nn.Dropout(dropout)]) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + if rotary_emb is None: + raise ValueError("SanaLinearAttention requires rotary_emb.") + + original_dtype = hidden_states.dtype + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.unflatten(2, (self.heads, -1)) + key = key.unflatten(2, (self.heads, -1)) + value = value.unflatten(2, (self.heads, -1)) + # B,N,H,C + + query = torch.relu(query) + key = torch.relu(key) + + def apply_rotary_emb( + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ): + x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + out = torch.empty_like(hidden_states) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out.type_as(hidden_states) + + query_rotate = apply_rotary_emb(query, *rotary_emb) + key_rotate = apply_rotary_emb(key, *rotary_emb) + + # B,H,C,N + query = query.permute(0, 2, 3, 1) + key = key.permute(0, 2, 3, 1) + query_rotate = query_rotate.permute(0, 2, 3, 1) + key_rotate = key_rotate.permute(0, 2, 3, 1) + value = value.permute(0, 2, 3, 1) + + query_rotate, key_rotate, value = query_rotate.float(), key_rotate.float(), value.float() + + # Keep the unrotated denominator contraction in the input dtype. This + # intentionally matches Diffusers 0.38.0's SanaLinearAttnProcessor3_0; + # only the rotated numerator path is accumulated in FP32. + z = 1 / (key.sum(dim=-1, keepdim=True).transpose(-2, -1) @ query + 1e-15) + + scores = torch.matmul(value, key_rotate.transpose(-1, -2)) + hidden_states = torch.matmul(scores, query_rotate) + + hidden_states = hidden_states * z + # B,H,C,N + hidden_states = hidden_states.flatten(1, 2).transpose(1, 2) + hidden_states = hidden_states.to(original_dtype) + + hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + + return hidden_states + + +class WanRotaryPosEmbed(nn.Module): + def __init__( + self, + attention_head_dim: int, + patch_size: tuple[int, int, int], + max_seq_len: int, + theta: float = 10000.0, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + + h_dim = w_dim = 2 * (attention_head_dim // 6) + t_dim = attention_head_dim - h_dim - w_dim + + self.t_dim = t_dim + self.h_dim = h_dim + self.w_dim = w_dim + + freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 + + freqs_cos = [] + freqs_sin = [] + + for dim in [t_dim, h_dim, w_dim]: + freq_cos, freq_sin = _get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype) + freqs_cos.append(freq_cos) + freqs_sin.append(freq_sin) + + self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False) + self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w + + split_sizes = [self.t_dim, self.h_dim, self.w_dim] + + freqs_cos = self.freqs_cos.split(split_sizes, dim=1) + freqs_sin = self.freqs_sin.split(split_sizes, dim=1) + + freqs_cos_f = freqs_cos[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_h = freqs_cos[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_sin_f = freqs_sin[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_h = freqs_sin[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_cos = torch.cat([freqs_cos_f, freqs_cos_h, freqs_cos_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + freqs_sin = torch.cat([freqs_sin_f, freqs_sin_h, freqs_sin_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + + return freqs_cos, freqs_sin + + +class SanaModulatedNorm(nn.Module): + def __init__(self, dim: int, elementwise_affine: bool = False, eps: float = 1e-6): + super().__init__() + self.norm = nn.LayerNorm(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, scale_shift_table: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm(hidden_states) + shift, scale = (scale_shift_table[None, None] + temb[:, :, None].to(scale_shift_table.device)).unbind(dim=2) + hidden_states = hidden_states * (1 + scale) + shift + return hidden_states + + +class SanaCombinedTimestepGuidanceEmbeddings(nn.Module): + def __init__(self, embedding_dim): + super().__init__() + self.time_proj = SanaTimesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = SanaTimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.guidance_condition_proj = SanaTimesteps( + num_channels=256, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + self.guidance_embedder = SanaTimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + + def forward(self, timestep: torch.Tensor, guidance: torch.Tensor = None, hidden_dtype: torch.dtype = None): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) + + guidance_proj = self.guidance_condition_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=hidden_dtype)) + conditioning = timesteps_emb + guidance_emb + + return self.linear(self.silu(conditioning)), conditioning + + +class SanaCrossAttention(nn.Module): + """SANA cross-attention backed by vLLM-Omni's unified attention layer.""" + + def __init__( + self, + dim: int, + cross_attention_dim: int, + num_heads: int, + head_dim: int, + dropout: float, + bias: bool, + out_bias: bool, + qk_norm: str | None, + prefix: str = "", + ) -> None: + super().__init__() + inner_dim = num_heads * head_dim + self.heads = num_heads + self.head_dim = head_dim + self.to_q = nn.Linear(dim, inner_dim, bias=bias) + self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias) + self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias) + self.norm_q = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None + self.norm_k = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None + self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=out_bias), nn.Dropout(dropout)]) + self.attn = OmniAttention( + num_heads=num_heads, + head_size=head_dim, + num_kv_heads=num_heads, + softmax_scale=1.0 / (head_dim**0.5), + causal=False, + role="cross", + qkv_layout="BSND", + prefix=prefix, + skip_sequence_parallel=True, + disable_kv_quant=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + query = self.to_q(hidden_states) + key = self.to_k(encoder_hidden_states) + value = self.to_v(encoder_hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.unflatten(2, (self.heads, self.head_dim)) + key = key.unflatten(2, (self.heads, self.head_dim)) + value = value.unflatten(2, (self.heads, self.head_dim)) + + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + if attention_mask is not None: + # The Flash masked-varlen path applies one padding mask to Q/K/V, + # while SANA's text mask describes only cross-attention K/V. + # Use SDPA until Flash supports an independent key-padding mask. + hidden_states = self.attn.sdpa_fallback.forward(query, key, value, attn_metadata) + else: + hidden_states = self.attn(query, key, value, None) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(query) + + hidden_states = self.to_out[0](hidden_states) + return self.to_out[1](hidden_states) + + +class SanaVideoTransformerBlock(nn.Module): + r""" + Transformer block introduced in [Sana-Video](https://huggingface.co/papers/2509.24695). + """ + + def __init__( + self, + dim: int = 2240, + num_attention_heads: int = 20, + attention_head_dim: int = 112, + dropout: float = 0.0, + num_cross_attention_heads: int | None = 20, + cross_attention_head_dim: int | None = 112, + cross_attention_dim: int | None = 2240, + attention_bias: bool = True, + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + attention_out_bias: bool = True, + mlp_ratio: float = 3.0, + qk_norm: str | None = "rms_norm_across_heads", + rope_max_seq_len: int = 1024, + prefix: str = "", + ) -> None: + super().__init__() + + # 1. Self Attention + self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) + self.attn1 = SanaLinearAttention( + dim=dim, + num_heads=num_attention_heads, + head_dim=attention_head_dim, + dropout=dropout, + bias=attention_bias, + qk_norm=qk_norm, + ) + + # 2. Cross Attention + self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + self.attn2 = None + if cross_attention_dim is not None: + self.attn2 = SanaCrossAttention( + dim=dim, + cross_attention_dim=cross_attention_dim, + num_heads=num_cross_attention_heads, + head_dim=cross_attention_head_dim, + dropout=dropout, + bias=True, + out_bias=attention_out_bias, + qk_norm=qk_norm, + prefix=f"{prefix}.attn2" if prefix else "attn2", + ) + + # 3. Feed-forward + self.ff = GLUMBTempConv(dim, dim, mlp_ratio, norm_type=None, residual_connection=False) + + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + timestep: torch.LongTensor | None = None, + frames: int = None, + height: int = None, + width: int = None, + rotary_emb: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size = hidden_states.shape[0] + + # 1. Modulation + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None, None] + timestep.reshape(batch_size, timestep.shape[1], 6, -1) + ).unbind(dim=2) + + # 2. Self Attention + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + norm_hidden_states = norm_hidden_states.to(hidden_states.dtype) + + attn_output = self.attn1(norm_hidden_states, rotary_emb=rotary_emb) + hidden_states = hidden_states + gate_msa * attn_output + + # 3. Cross Attention + if self.attn2 is not None: + attn_output = self.attn2( + hidden_states, + encoder_hidden_states, + encoder_attention_mask, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + norm_hidden_states = norm_hidden_states.unflatten(1, (frames, height, width)) + ff_output = self.ff(norm_hidden_states) + ff_output = ff_output.flatten(1, 3) + hidden_states = hidden_states + gate_mlp * ff_output + + return hidden_states + + +class SanaVideoTransformer3DModel(nn.Module): + r""" + A 3D Transformer model introduced in [Sana-Video](https://huggingface.co/papers/2509.24695) family of models. + + Args: + in_channels (`int`, defaults to `16`): + The number of channels in the input. + out_channels (`int`, *optional*, defaults to `16`): + The number of channels in the output. + num_attention_heads (`int`, defaults to `20`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`, defaults to `112`): + The number of channels in each head. + num_layers (`int`, defaults to `20`): + The number of layers of Transformer blocks to use. + num_cross_attention_heads (`int`, *optional*, defaults to `20`): + The number of heads to use for cross-attention. + cross_attention_head_dim (`int`, *optional*, defaults to `112`): + The number of channels in each head for cross-attention. + cross_attention_dim (`int`, *optional*, defaults to `2240`): + The number of channels in the cross-attention output. + caption_channels (`int`, defaults to `2304`): + The number of channels in the caption embeddings. + mlp_ratio (`float`, defaults to `2.5`): + The expansion ratio to use in the GLUMBConv layer. + dropout (`float`, defaults to `0.0`): + The dropout probability. + attention_bias (`bool`, defaults to `False`): + Whether to use bias in the attention layer. + sample_size (`int`, defaults to `32`): + The base size of the input latent. + patch_size (`int`, defaults to `1`): + The size of the patches to use in the patch embedding layer. + norm_elementwise_affine (`bool`, defaults to `False`): + Whether to use elementwise affinity in the normalization layer. + norm_eps (`float`, defaults to `1e-6`): + The epsilon value for the normalization layer. + qk_norm (`str`, *optional*, defaults to `None`): + The normalization to use for the query and key. + """ + + _no_split_modules = ["SanaVideoTransformerBlock", "SanaModulatedNorm"] + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return AutoWeightsLoader(self).load_weights(weights) + + def __init__( + self, + in_channels: int = 16, + out_channels: int | None = 16, + num_attention_heads: int = 20, + attention_head_dim: int = 112, + num_layers: int = 20, + num_cross_attention_heads: int | None = 20, + cross_attention_head_dim: int | None = 112, + cross_attention_dim: int | None = 2240, + caption_channels: int = 2304, + mlp_ratio: float = 2.5, + dropout: float = 0.0, + attention_bias: bool = False, + sample_size: int = 30, + patch_size: tuple[int, int, int] = (1, 2, 2), + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + interpolation_scale: int | None = None, + guidance_embeds: bool = False, + guidance_embeds_scale: float = 0.1, + qk_norm: str | None = "rms_norm_across_heads", + rope_max_seq_len: int = 1024, + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + self.config = SanaVideoTransformerConfig( + in_channels=in_channels, + out_channels=out_channels, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + num_layers=num_layers, + num_cross_attention_heads=num_cross_attention_heads, + cross_attention_head_dim=cross_attention_head_dim, + cross_attention_dim=cross_attention_dim, + caption_channels=caption_channels, + mlp_ratio=mlp_ratio, + dropout=dropout, + attention_bias=attention_bias, + sample_size=sample_size, + patch_size=tuple(patch_size), + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + interpolation_scale=interpolation_scale, + guidance_embeds=guidance_embeds, + guidance_embeds_scale=guidance_embeds_scale, + qk_norm=qk_norm, + rope_max_seq_len=rope_max_seq_len, + ) + inner_dim = num_attention_heads * attention_head_dim + + # 1. Patch & position embedding + self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) + self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) + + # 2. Additional condition embeddings + if guidance_embeds: + self.time_embed = SanaCombinedTimestepGuidanceEmbeddings(inner_dim) + else: + self.time_embed = SanaAdaLayerNormSingle(inner_dim) + + self.caption_projection = SanaPixArtAlphaTextProjection( + in_features=caption_channels, + hidden_size=inner_dim, + ) + self.caption_norm = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) + + # 3. Transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + SanaVideoTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + num_cross_attention_heads=num_cross_attention_heads, + cross_attention_head_dim=cross_attention_head_dim, + cross_attention_dim=cross_attention_dim, + attention_bias=attention_bias, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + mlp_ratio=mlp_ratio, + qk_norm=qk_norm, + prefix=f"transformer_blocks.{layer_idx}", + ) + for layer_idx in range(num_layers) + ] + ) + + # 4. Output blocks + self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) + self.norm_out = SanaModulatedNorm(inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner_dim, math.prod(patch_size) * out_channels) + + @property + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + @classmethod + def from_config(cls, config: dict | SanaVideoTransformerConfig) -> "SanaVideoTransformer3DModel": + if isinstance(config, dict): + config = SanaVideoTransformerConfig.from_dict(config) + return cls(**{field.name: getattr(config, field.name) for field in fields(config)}) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + guidance: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + controlnet_block_samples: tuple[torch.Tensor] | None = None, + return_dict: bool = True, + ) -> tuple[torch.Tensor, ...] | SanaVideoTransformerOutput: + """ + The [`SanaVideoTransformer3DModel`] forward method. + + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, in_channels, num_frames, height, width)`): + Input `hidden_states`. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, sequence_len, embed_dims)`): + Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. + timestep (`torch.LongTensor`): + Used to indicate denoising step. + guidance (`torch.Tensor`, *optional*): + Guidance scale embedding. + encoder_attention_mask (`torch.Tensor`, *optional*): + Cross-attention mask applied to `encoder_hidden_states`. + attention_mask (`torch.Tensor`, *optional*): + Retained for Diffusers API compatibility. SANA linear self-attention does not consume this mask. + controlnet_block_samples (`tuple` of `torch.Tensor`, *optional*): + A list of tensors that if specified are added to the residuals of transformer blocks. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return `SanaVideoTransformerOutput` instead of a plain tuple. + + Returns: + If `return_dict` is True, a [`SanaVideoTransformerOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + # SANA's linear self-attention does not consume ``attention_mask``; + # retain the public argument for Diffusers-compatible call sites. + if encoder_attention_mask is not None: + if encoder_attention_mask.ndim != 2: + raise ValueError( + "encoder_attention_mask must be a 2D padding mask with " + f"shape (batch_size, key_length); got {tuple(encoder_attention_mask.shape)}" + ) + if encoder_attention_mask.shape != encoder_hidden_states.shape[:2]: + raise ValueError( + "encoder_attention_mask shape must match encoder_hidden_states: " + f"expected {tuple(encoder_hidden_states.shape[:2])}, " + f"got {tuple(encoder_attention_mask.shape)}" + ) + if encoder_attention_mask.is_floating_point() or encoder_attention_mask.is_complex(): + raise TypeError( + "encoder_attention_mask must be a boolean or integer padding mask, " + "not a floating-point additive attention bias" + ) + encoder_attention_mask = encoder_attention_mask.to( + device=encoder_hidden_states.device, + dtype=torch.bool, + ) + + # 1. Input + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.config.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + rotary_emb = self.rope(hidden_states) + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if guidance is not None: + timestep, embedded_timestep = self.time_embed( + timestep.flatten(), guidance=guidance, hidden_dtype=hidden_states.dtype + ) + else: + timestep, embedded_timestep = self.time_embed( + timestep.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype + ) + + timestep = timestep.view(batch_size, -1, timestep.size(-1)) + embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) + + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) + + encoder_hidden_states = self.caption_norm(encoder_hidden_states) + + # 2. Transformer blocks + for index_block, block in enumerate(self.transformer_blocks): + hidden_states = block( + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + post_patch_num_frames, + post_patch_height, + post_patch_width, + rotary_emb, + ) + if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples): + hidden_states = hidden_states + controlnet_block_samples[index_block - 1] + + # 3. Normalization + hidden_states = self.norm_out(hidden_states, embedded_timestep, self.scale_shift_table) + + hidden_states = self.proj_out(hidden_states) + + # 5. Unpatchify + hidden_states = hidden_states.reshape( + batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (output,) + + return SanaVideoTransformerOutput(sample=output) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 921ceca1b..f5d6ffe07 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -291,6 +291,16 @@ _DIFFUSION_MODELS = { "pipeline_lingbot_video", "LingBotVideoPipeline", ), + "SanaVideoPipeline": ( + "sana_video", + "pipeline_sana_video", + "SanaVideoPipeline", + ), + "SanaImageToVideoPipeline": ( + "sana_video", + "pipeline_sana_video_i2v", + "SanaImageToVideoPipeline", + ), "MagiHumanPipeline": ( "magi_human", "pipeline_magi_human", @@ -586,6 +596,8 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func", "HunyuanImage3Pipeline": "get_hunyuan_image3_post_process_func", "LingBotVideoPipeline": "get_lingbot_video_post_process_func", + "SanaVideoPipeline": "get_sana_video_post_process_func", + "SanaImageToVideoPipeline": "get_sana_video_i2v_post_process_func", "MagiHumanPipeline": "get_magi_human_post_process_func", "OmniVoicePipeline": "get_omnivoice_post_process_func", "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", @@ -631,6 +643,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "HeliosPyramidPipeline": "get_helios_pre_process_func", "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_pre_process_func", "LingBotVideoPipeline": "get_lingbot_video_pre_process_func", + "SanaImageToVideoPipeline": "get_sana_video_i2v_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", diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 82ff7ebe2..6173895c3 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -1,6 +1,6 @@ # adapted from sglang and fastvideo # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -17,6 +17,25 @@ DUMMY_DIFFUSION_REQUEST_ID = "dummy_req_id" _DUMMY_DIFFUSION_REQUEST_ID_PREFIX = f"{DUMMY_DIFFUSION_REQUEST_ID}/" +def resolve_video_num_frames( + num_frames: int | None, + *, + default_num_frames: int, + is_dummy_run: bool, +) -> int: + """Resolve the shared image-model frame sentinel for a video pipeline. + + ``OmniDiffusionSamplingParams`` defaults ``num_frames`` to one for image + models, so an omitted video API field reaches model code as ``1``. Video + pipelines with a different default must materialize it at their boundary. + Startup profiling intentionally requests one frame, however, and must stay + lightweight. + """ + if num_frames is None or (num_frames == 1 and not is_dummy_run): + return default_num_frames + return num_frames + + @dataclass class OmniDiffusionRequest: """ diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 0bd57cb02..b8eb6bbe8 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -51,6 +51,7 @@ from vllm_omni.model_extras.ming_flash_omni import ( from vllm_omni.model_extras.ming_flash_omni import ( build_text_to_image_prompt as build_ming_flash_omni_text_to_image_prompt, ) +from vllm_omni.model_extras.sana_video import SANA_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.sensenova_u1 import ( SENSENOVA_U1_EXTRA_BODY_PARAMS, SENSENOVA_U1_EXTRA_OUTPUT_PARAMS, @@ -213,6 +214,12 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "LTX2DistilledTwoStagePipeline", ) }, + "SanaVideoPipeline": { + "extra_body_params": SANA_VIDEO_EXTRA_BODY_PARAMS, + }, + "SanaImageToVideoPipeline": { + "extra_body_params": SANA_VIDEO_EXTRA_BODY_PARAMS, + }, "WanVACEPipeline": { "extra_body_params": VACE_EXTRA_BODY_PARAMS, "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS, diff --git a/vllm_omni/model_extras/sana_video.py b/vllm_omni/model_extras/sana_video.py new file mode 100644 index 000000000..67234de48 --- /dev/null +++ b/vllm_omni/model_extras/sana_video.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +SANA_VIDEO_EXTRA_BODY_PARAMS = frozenset( + { + "clean_caption", + "motion_score", + "use_resolution_binning", + } +)