[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 <yaoliu548926@gmail.com> Co-authored-by: leo <92263211+napleon-liu@users.noreply.github.com>
This commit is contained in:
leoandleo authored and GitHub committed
1 parent
4af0ddc1c2
commit
50b8a150f9
42 files changed
+6512
-18
No files matched your search
@@ -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) |
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -4,8 +4,8 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
|
||||
|
||||
|
||||
This example demonstrates how to generate videos from images using Wan2.2,
|
||||
LTX-2, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
|
||||
offline inference API.
|
||||
LTX-2, HunyuanVideo-1.5, and SANA-Video Image-to-Video models with
|
||||
vLLM-Omni's offline inference API.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
@@ -66,9 +66,49 @@ 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
|
||||
|
||||
SANA checkpoints identify their upstream T2V pipeline in `model_index.json`.
|
||||
Select the native I2V implementation explicitly:
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
For the 720p checkpoint, use
|
||||
`Efficient-Large-Model/SANA-Video_2B_720p_diffusers` with
|
||||
`--height 704 --width 1280`. The native I2V path supports both checkpoint
|
||||
variants: the 480p checkpoint uses the Wan VAE, while the 720p checkpoint uses
|
||||
the LTX-2 Video VAE.
|
||||
|
||||
For SANA-Video, 81 frames at 16 FPS is the standard checkpoint request
|
||||
(approximately five seconds), not minute-scale long-video generation.
|
||||
Minute-scale SANA generation requires the separate LongSANA/LongLive
|
||||
autoregressive workflow, which this pipeline does not implement.
|
||||
|
||||
See the [SANA-Video recipe](../../../../recipes/NVIDIA/SANA-Video-2B.md) for
|
||||
online serving, backend boundaries, and the tested hardware profile.
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or LTX-2).
|
||||
- `--model`: Model ID or local path (for example Wan I2V/TI2V, LTX-2, or
|
||||
SANA-Video).
|
||||
- `--model-class-name`: Explicit pipeline override. SANA-Video I2V requires
|
||||
`SanaImageToVideoPipeline`.
|
||||
- `--image`: Path to input image (required).
|
||||
- `--prompt`: Text description of desired motion/animation.
|
||||
- `--height/--width`: Output resolution (auto-calculated from image if not set).
|
||||
|
||||
@@ -16,6 +16,8 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
|
||||
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
|
||||
| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE |
|
||||
| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | ~23.6 GiB |
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
@@ -122,6 +124,31 @@ python text_to_video.py \
|
||||
--output quick_test.mp4
|
||||
```
|
||||
|
||||
### SANA-Video-2B
|
||||
|
||||
```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. SANA-Video
|
||||
accepts `motion_score`, `clean_caption`, and `use_resolution_binning` through
|
||||
`--extra-body`. For image-to-video, use the shared
|
||||
`examples/offline_inference/image_to_video/image_to_video.py` example with
|
||||
`--model-class-name SanaImageToVideoPipeline` and `--image <path>`.
|
||||
|
||||
## Key Arguments
|
||||
|
||||
### Common
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 <path>`.
|
||||
|
||||
### Cosmos3
|
||||
|
||||
```bash
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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: <https://github.com/NVlabs/Sana>
|
||||
- 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: <https://huggingface.co/docs/diffusers/api/pipelines/sana_video>
|
||||
- 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.
|
||||
@@ -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) |
|
||||
|
||||
@@ -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]))
|
||||
@@ -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)
|
||||
@@ -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
|
||||
<output-root>/<variant>/<task>/
|
||||
```
|
||||
|
||||
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/<variant>/<task>/
|
||||
```
|
||||
|
||||
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.
|
||||
@@ -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()
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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"),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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": <PIL image or file path>}`.')
|
||||
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
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
)
|
||||
Reference in new issue
Block a user