diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index ea74847b9..ed112b84b 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -459,6 +459,12 @@ steps: - pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model" mirror_hardwares: h100_2 + - label: ":full_moon: Diffusion X2V · LTX-2.5 Online Test" + timeout_in_minutes: 60 + commands: + - pytest -s -v tests/e2e/online_serving/test_ltx25.py -m "full_model and cuda" --run-level "full_model" + mirror_hardwares: h100_1 + - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test" timeout_in_minutes: 180 commands: @@ -487,6 +493,13 @@ steps: - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model mirror_hardwares: h100_1 + - label: ":full_moon: Diffusion X2V · LTX-2.5 Official Parity Test" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_distilled_two_stage_matches_official -m full_model --run-level full_model + - pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_full_matches_official -m full_model --run-level full_model + mirror_hardwares: h100_1 + - label: ":full_moon: Diffusion X2V · Perf Test" key: nightly-diffusion-x2v-performance timeout_in_minutes: 180 diff --git a/.gitignore b/.gitignore index f11b57388..7502df8a0 100644 --- a/.gitignore +++ b/.gitignore @@ -254,6 +254,7 @@ outputs/ results/ generated/ output_*/ +/tests/e2e/accuracy/artifacts/ # Configuration overrides configs/local.yaml diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index e6a38640f..5097c83fa 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -43,7 +43,12 @@ th { | `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — | | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — | | `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — | -| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — | +| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — | +| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — | +| `LTX2Pipeline` | LTX-2.5 Full/SFT 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) | +| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) | +| `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` | ✅︎ | | | | — | | `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 11c6366b7..5844d1114 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -161,6 +161,7 @@ The following tables show which models support each feature: | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | @@ -174,7 +175,7 @@ The following tables show which models support each feature: **Frame Interpolation Support** - **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines -- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni +- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni ### AudioGen diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 222502e8a..28c98990a 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -68,6 +68,7 @@ from vllm_omni.model_extras import ( from vllm_omni.model_extras import ( get_extra_body_params, get_model_class_name, + should_preserve_reference_image_size, ) from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -454,7 +455,13 @@ def main(): media_inputs: dict[str, Any] = {} if image is not None: - media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS) + preserve_image_size = should_preserve_reference_image_size( + model_class_name, + model=args.model, + ) + media_inputs["image"] = ( + image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS) + ) if last_image is not None: media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS) if mask_image is not None: diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 129ca860a..406aab39e 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -570,6 +570,8 @@ def main(): num_frames=args.num_frames, extra_args=extra_args, ) + if args.frame_rate is not None: + sampling_kwargs["frame_rate"] = args.frame_rate if args.guidance_scale_high is not None: sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high diff --git a/recipes/LTX/LTX-2.5.md b/recipes/LTX/LTX-2.5.md new file mode 100644 index 000000000..4ff5f6dab --- /dev/null +++ b/recipes/LTX/LTX-2.5.md @@ -0,0 +1,175 @@ +# LTX-2.5 + +> Text-to-video and first-frame image-to-video generation with synchronized audio + +## Overview + +vLLM-Omni supports the gated +[`Lightricks/LTX-2.5-Diffusers`](https://huggingface.co/Lightricks/LTX-2.5-Diffusers) +checkpoint for T2V and first-frame I2V. Every output includes synchronized +48 kHz stereo audio. + +The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5) +repository is not directly loadable with `--model`. It supplies the official +upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept +both model licenses and authenticate before first use. + +## Choose a pipeline + +| Pipeline | Mode | Default size | Schedule | +|---|---|---:|---:| +| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps | +| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps | +| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps | +| `LTX2DistilledTwoStagePipeline` | Distilled two-stage | 1920x1088 | 8 + 3 steps | + +Both two-stage pipelines generate at half resolution, apply the official x2 +latent upsampler, and run a three-step refinement stage. Select the class with +`--model-class-name`; no `--task-type` flag is required. Supplying one initial +image selects I2V, while omitting it selects T2V. + +## Prerequisites + +```bash +hf auth login +export MODEL=Lightricks/LTX-2.5-Diffusers +``` + +Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and +`ffprobe` are on `PATH`. I2V requires PyAV backed by an FFmpeg build with the +`libx264` encoder. + +## Hardware + +| GPU | Status | Recommended scope | +|---|---|---| +| NVIDIA B300 | Verified | All four canonical pipelines | +| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines | +| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines | +| NVIDIA H100 80 GB | FP8 recipe | Distilled one-stage at 960x544 | + +The 1920x1088 two-stage examples require about 114 GB of peak GPU memory. +80 GB GPUs do not have enough safety margin for the canonical two-stage +configuration; use a larger GPU or reduce the output configuration. + +### H100 80 GB FP8 + +Use the distilled one-stage pipeline with FP8 and cuDNN attention: + +```bash +vllm serve Lightricks/LTX-2.5-Diffusers \ + --omni \ + --model-class-name LTX2DistilledOneStagePipeline \ + --quantization fp8 \ + --diffusion-attention-backend CUDNN_ATTN \ + --host 0.0.0.0 \ + --port 8000 +``` + +Use the T2V or I2V request below after the server is ready. + +## Offline inference + +Choose values from the pipeline table. For example, the distilled two-stage +path uses: + +```bash +export PIPELINE=LTX2DistilledTwoStagePipeline +export WIDTH=1920 +export HEIGHT=1088 +export STEPS=8 +``` + +### Text-to-video + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ + --model "${MODEL}" \ + --model-class-name "${PIPELINE}" \ + --prompt "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot." \ + --width "${WIDTH}" \ + --height "${HEIGHT}" \ + --num-frames 121 \ + --num-inference-steps "${STEPS}" \ + --frame-rate 24 \ + --fps 24 \ + --seed 42 \ + --output ltx25-t2v.mp4 +``` + +### First-frame image-to-video + +```bash +python examples/offline_inference/image_to_video/image_to_video.py \ + --model "${MODEL}" \ + --model-class-name "${PIPELINE}" \ + --image /absolute/path/to/first-frame.png \ + --prompt "The red fox walks forward while the camera tracks alongside." \ + --width "${WIDTH}" \ + --height "${HEIGHT}" \ + --num-frames 121 \ + --num-inference-steps "${STEPS}" \ + --frame-rate 24 \ + --fps 24 \ + --seed 42 \ + --output ltx25-i2v.mp4 +``` + +LTX-2.5 uses the official CRF-18 first-frame conditioning path by default. + +## Online serving + +Start one server for the selected pipeline: + +```bash +vllm serve "${MODEL}" \ + --omni \ + --model-class-name "${PIPELINE}" \ + --host 0.0.0.0 \ + --port 8000 \ + --stage-init-timeout 900 +``` + +T2V request: + +```bash +curl -sS --fail-with-body \ + -X POST http://127.0.0.1:8000/v1/videos/sync \ + -F 'prompt=A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot.' \ + -o ltx25-online-t2v.mp4 +``` + +For I2V, add exactly one first frame: + +```bash +export FIRST_FRAME=/absolute/path/to/first-frame.png + +curl -sS --fail-with-body \ + -X POST http://127.0.0.1:8000/v1/videos/sync \ + -F "input_reference=@${FIRST_FRAME};type=image/png" \ + -F 'prompt=The red fox walks forward while the camera tracks alongside.' \ + -o ltx25-online-i2v.mp4 +``` + +Restart the server after changing `PIPELINE` because each class selects a +different weight layout and execution topology. + +## Constraints + +- Online `width`, `height`, `num_frames`, `fps`, and `num_inference_steps` + are optional and use the selected pipeline defaults when omitted. `seed` is + optional and random when omitted. +- `num_frames` must be `8k+1` when overridden. +- One-stage width and height must be divisible by 32; two-stage final + dimensions must be divisible by 64. +- Full/SFT pipelines accept negative prompts. Distilled pipelines are + positive-only and reject negative prompts. +- I2V accepts exactly one initial image per prompt. +- The supported `--model` value is `Lightricks/LTX-2.5-Diffusers`. + +## References + +- [LTX-2.5-Diffusers model card](https://huggingface.co/Lightricks/LTX-2.5-Diffusers) +- [Official LTX-2 repository](https://github.com/Lightricks/LTX-2) +- [Generic T2V example](../../examples/offline_inference/text_to_video/text_to_video.py) +- [Generic I2V example](../../examples/offline_inference/image_to_video/image_to_video.py) diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md index 123880d09..ec4e8fb3b 100644 --- a/recipes/LTX/LTX-2.md +++ b/recipes/LTX/LTX-2.md @@ -8,10 +8,12 @@ |---|---|---| | `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` | | `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` | -| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` | +| `LTX2DistilledOneStagePipeline` | LTX-2 merged-distilled one-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` | +| `LTX2DistilledTwoStagePipeline` | LTX-2 merged-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` | | `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` | | `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`
`Lightricks/LTX-2.3` | -| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`
`Lightricks/LTX-2.3` | +| `LTX2DistilledOneStagePipeline` | LTX-2.3 merged-distilled one-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers` | +| `LTX2DistilledTwoStagePipeline` | LTX-2.3 merged-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`
`Lightricks/LTX-2.3` | Repositories in the table are download units. A full pipeline repository contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler, @@ -26,9 +28,12 @@ image selects I2V. Both one-stage repositories declare this class, so the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable. `LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples, -then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully -merged distilled Transformer in both stages. Both entries support T2V and I2V; -select their class explicitly. +then refines with the distilled LoRA. `LTX2DistilledOneStagePipeline` uses a +merged distilled Transformer without upsampling, while +`LTX2DistilledTwoStagePipeline` uses it in both stages. All entries support +T2V and I2V; select their class explicitly. The deprecated +`LTX2DistilledPipeline` name remains an alias for +`LTX2DistilledTwoStagePipeline`. ## API Migration @@ -52,8 +57,8 @@ The consolidation also removes these registry names without aliases: | `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 | | `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` | | `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 | -| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` | -| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` | +| `LTX2TwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` | +| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` | Passing any second positional argument now raises `TypeError`. These changes affect direct Python callers and explicit `--model-class-name` overrides; @@ -112,10 +117,10 @@ Start a two-stage checkpoint with its explicit class: ```bash vllm serve rootonchair/LTX-2-19b-distilled --omni \ - --model-class-name LTX2DistilledPipeline --stage-init-timeout 600 + --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600 # LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \ - --model-class-name LTX2DistilledPipeline --stage-init-timeout 600 + --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600 # Ordinary LTX-2.3 two-stage vllm serve diffusers/LTX-2.3-Diffusers --omni \ --model-class-name LTX2TwoStagePipeline \ diff --git a/recipes/README.md b/recipes/README.md index e010ade67..6a3b23770 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -55,6 +55,7 @@ recipes/ | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) | | [`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 | | [`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 | diff --git a/requirements/common.txt b/requirements/common.txt index e0233bd8d..368cbe8be 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -1,15 +1,15 @@ # Common dependencies for all platforms -# transformers: kept exactly aligned with the upstream vLLM constraint. -# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0". -# The previous extra "< 5.9.0" exclusion was removed when re-aligning with -# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from +# transformers: required unified multimodal model classes are unavailable +# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction +# regression. transformers 5.9 also changed +# _adjust_missing_and_unexpected_keys from # list+list to set|set, breaking models that declare # _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano) # with "TypeError: unsupported operand type(s) for |: 'list' and 'set'". # That breakage is now handled in code instead of by pinning: # MossTTSNanoForGeneration wraps weight loading in # transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py). -transformers >= 5.5.3 +transformers >= 5.10.1, < 5.15 av>=14.0.0 omegaconf>=2.3.0 diffusers==0.38.0 diff --git a/tests/diffusion/attention/test_cudnn_attn.py b/tests/diffusion/attention/test_cudnn_attn.py index de1adf8c7..1ff3baafb 100644 --- a/tests/diffusion/attention/test_cudnn_attn.py +++ b/tests/diffusion/attention/test_cudnn_attn.py @@ -5,6 +5,7 @@ from contextlib import nullcontext import pytest import torch +from torch.nn.attention import SDPBackend from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl @@ -12,6 +13,48 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] +@pytest.mark.parametrize( + ("query_sequence_length", "key_sequence_length", "expected_backend"), + [ + (1, 1, SDPBackend.MATH), + (8, 1, SDPBackend.MATH), + (1, 8, SDPBackend.MATH), + (8, 8, SDPBackend.CUDNN_ATTENTION), + ], +) +def test_cudnn_routes_single_token_attention_to_math( + monkeypatch, query_sequence_length, key_sequence_length, expected_backend +): + selected_backends = [] + + def fake_sdpa_kernel(backends): + selected_backends.append(backends) + return nullcontext() + + monkeypatch.setattr( + "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel", + fake_sdpa_kernel, + ) + monkeypatch.setattr( + torch.nn.functional, + "scaled_dot_product_attention", + lambda query, _key, _value, **_kwargs: query, + ) + impl = CuDNNAttentionImpl( + num_heads=2, + head_size=4, + softmax_scale=0.5, + ) + query = torch.randn(1, query_sequence_length, 2, 4) + key = torch.randn(1, key_sequence_length, 2, 4) + value = torch.randn_like(key) + + output = impl.forward_cuda(query, key, value) + + assert output.shape == query.shape + assert selected_backends == [[expected_backend]] + + def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch): observed = {} diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 2799f1607..0723c1d62 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -36,7 +36,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] SEPARATE_CFG_TRANSFORMERS = [ DreamIdOmniModel, - LTX2VideoTransformer3DModel, HeliosTransformer3DModel, LongCatImageTransformer2DModel, Cosmos3VFMTransformer, @@ -124,6 +123,13 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model): assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True +def test_ltx2_cache_dit_uses_one_forward_per_denoise_step(): + """LTX batches all guidance passes into one Transformer invocation.""" + adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config + + assert adapter_config.has_separate_cfg is False + + @patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter") @patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit") def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter): diff --git a/tests/diffusion/model_loader/test_hub_prefetch.py b/tests/diffusion/model_loader/test_hub_prefetch.py new file mode 100644 index 000000000..39eb2b277 --- /dev/null +++ b/tests/diffusion/model_loader/test_hub_prefetch.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import contextlib + +import huggingface_hub +import pytest + +from vllm_omni.diffusion.model_loader import hub_prefetch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_prefetch_subfolders_propagates_revision(monkeypatch): + calls = [] + + def fake_snapshot_download(**kwargs): + calls.append(kwargs) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext()) + + hub_prefetch.prefetch_subfolders( + "org/model", + ["transformer"], + revision="pinned-revision", + ) + + assert calls == [ + { + "repo_id": "org/model", + "revision": "pinned-revision", + "allow_patterns": ["transformer/*", "transformer/**", "*.json", "*.txt"], + } + ] + + +def test_from_pretrained_retry_preserves_revision(monkeypatch): + factory_calls = [] + prefetch_calls = [] + + def fake_factory(model, **kwargs): + factory_calls.append((model, kwargs)) + if len(factory_calls) == 1: + raise OSError("partial cache") + return "loaded" + + def fake_prefetch(model, subfolders, **kwargs): + prefetch_calls.append((model, tuple(subfolders), kwargs)) + + monkeypatch.setattr(hub_prefetch, "prefetch_subfolders", fake_prefetch) + monkeypatch.setattr(hub_prefetch.time, "sleep", lambda _seconds: None) + + result = hub_prefetch.from_pretrained_with_prefetch( + fake_factory, + "org/model", + subfolder="transformer", + prefetch_list=["transformer", "vae"], + revision="pinned-revision", + max_attempts=2, + ) + + assert result == "loaded" + assert len(factory_calls) == 2 + assert all(call[1]["revision"] == "pinned-revision" for call in factory_calls) + assert prefetch_calls == [ + ( + "org/model", + ("transformer", "vae"), + {"local_files_only": False, "revision": "pinned-revision"}, + ) + ] diff --git a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py new file mode 100644 index 000000000..cbe52e997 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py @@ -0,0 +1,881 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Focused LTX-2.5 pipeline correctness tests.""" + +import json +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.models.ltx2 import ltx2_components, ltx2_latents, ltx2_runtime +from vllm_omni.diffusion.models.ltx2.ltx2_components import ( + LTX25_DISTILLED_COMPONENT_PROFILE, + LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, + LTX25_FULL_COMPONENT_PROFILE, + LTX25_TWO_STAGE_COMPONENT_PROFILE, + detect_ltx_model_version, + resolve_ltx_checkpoint_kind, + resolve_ltx_component_profile, +) +from vllm_omni.diffusion.models.ltx2.ltx2_denoise import ( + LTXPhaseResult, + _first_frame_keyframes_mask, + _official_ltx_sigmas, +) +from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( + LTX_GUIDANCE_EXECUTOR, + LTXGuidanceSpec, + LTXModalityGuidance, + velocity_from_x0, + x0_from_velocity, +) +from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState +from vllm_omni.diffusion.models.ltx2.ltx2_recipes import ( + LTX2_DISTILLED_TWO_STAGE_RECIPE, + LTX2_ONE_STAGE_RECIPE, + LTX23_ONE_STAGE_RECIPE, + LTX25_DISTILLED_ONE_STAGE_RECIPE, + LTX25_DISTILLED_TWO_STAGE_RECIPE, + LTX25_FULL_RECIPE, + LTX25_TWO_STAGE_RECIPE, + LTX_DISTILLED_SIGMAS, + LTX_POSITIVE_ONLY_RECIPE, + resolve_ltx_pipeline_recipe, +) +from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs +from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2DistilledOneStagePipeline, + LTX2Pipeline, +) +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( + LTX2DistilledTwoStagePipeline, + LTX2TwoStagePipeline, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("model_version", ["2", "2.3", "2.5"]) +def test_ltx_timestep_adaln_expands_per_token_like_official_without_sp(model_version): + pipe = object.__new__(LTX2Pipeline) + pipe.model_version = model_version + pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)) + ts = torch.tensor([750.0, 500.0]) + + kwargs = LTXRuntime._denoise_timestep_kwargs( + pipe, + ts, + SimpleNamespace(), + SimpleNamespace(), + video_token_count=4, + audio_token_count=3, + ) + + torch.testing.assert_close(kwargs["timestep"], ts[:, None].expand(-1, 4)) + torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None].expand(-1, 3)) + assert kwargs["sigma"] is ts + assert kwargs["audio_sigma"] is ts + + +def test_ltx_ancestral_step_clears_sp_audio_padding_after_scheduler_update(monkeypatch): + state = LTXAVState( + video=torch.zeros(1, 2, 2), + audio=torch.zeros(1, 4, 2), + ) + forward_ctx = SimpleNamespace( + original_audio_num_frames=3, + sampler="euler_ancestral", + ) + denoise_ctx = SimpleNamespace(latents=None, audio_latents=None) + pipe = SimpleNamespace( + _predict_noise_for_step=lambda *_args: (torch.zeros_like(state.video), torch.zeros_like(state.audio)) + ) + + def fake_step(_pipeline, actual_forward_ctx, _denoise_ctx, *_args): + assert actual_forward_ctx.sampler == "euler_ancestral" + updated_audio = torch.ones_like(state.audio) + updated_audio[:, 3:] = 99.0 + return torch.ones_like(state.video), updated_audio + + monkeypatch.setattr(ltx2_runtime, "step_denoised_latents", fake_step) + + actual = LTXRuntime._denoise_step( + pipe, + 0, + torch.tensor(1.0), + state, + forward_ctx, + denoise_ctx, + ) + + torch.testing.assert_close(actual.audio[:, :3], torch.ones_like(actual.audio[:, :3])) + torch.testing.assert_close(actual.audio[:, 3:], torch.zeros_like(actual.audio[:, 3:])) + + +def test_ltx25_first_causal_latent_frame_is_marked_for_keyframe_embedding(): + reference = torch.zeros(2, 12, 8) + + mask = _first_frame_keyframes_mask(reference, latent_num_frames=3) + + assert mask.shape == (2, 12, 1) + torch.testing.assert_close(mask[:, :4], torch.ones(2, 4, 1)) + torch.testing.assert_close(mask[:, 4:], torch.zeros(2, 8, 1)) + + +def _make_ltx_request_pipe(cls): + pipe = object.__new__(cls) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 99 + pipe.vae_spatial_compression_ratio = 32 + pipe.vae_temporal_compression_ratio = 8 + return pipe + + +def _resolve_request_inputs_for_test( + pipe, + req, + *, + guidance_scale=4.0, + negative_prompt=None, +): + return pipe._resolve_request_inputs( + req, + prompt=None, + negative_prompt=negative_prompt, + height=None, + width=None, + num_frames=None, + frame_rate=None, + num_inference_steps=None, + guidance_scale=guidance_scale, + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=None, + ) + + +def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch): + pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None)) + od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16) + + monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(ltx2_components, "prefetch_subfolders", lambda *_args, **_kwargs: None) + monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", lambda *_args, **_kwargs: object()) + + with pytest.raises(ImportError) as exc_info: + ltx2_components.initialize_pipeline_components(pipe, od_config) + + assert str(exc_info.value) == ( + "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15." + ) + + +def test_ltx_converted_component_loading_propagates_revision(monkeypatch): + revision = "pinned-revision" + calls = {"components": []} + profile = replace( + LTX25_DISTILLED_COMPONENT_PROFILE, + text_encoder_cls=object, + video_vae_cls=object, + vocoder_cls=object, + vocoder_fallback_cls=None, + scheduler_use_dynamic_shifting=False, + ) + pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage") + od_config = SimpleNamespace( + model="org/converted-ltx25", + revision=revision, + dtype=torch.bfloat16, + quantization_config=None, + ) + + monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(ltx2_components, "_install_connector_attention", lambda _connectors, **_kwargs: None) + monkeypatch.setattr(ltx2_components, "_place_aux_components", lambda _pipeline: None) + + def fake_prefetch(model, subfolders, **kwargs): + calls["prefetch"] = (model, tuple(subfolders), kwargs) + + def fake_tokenizer(*args, **kwargs): + calls["tokenizer"] = (args, kwargs) + return SimpleNamespace(model_max_length=1024) + + def fake_component(component_cls, model, subfolder, **kwargs): + calls["components"].append((component_cls, model, subfolder, kwargs)) + if subfolder == "vae": + return SimpleNamespace(spatial_compression_ratio=32, temporal_compression_ratio=8) + if subfolder == "audio_vae": + return SimpleNamespace( + mel_compression_ratio=4, + temporal_compression_ratio=4, + config=SimpleNamespace(sample_rate=16_000, mel_hop_length=160), + ) + return object() + + def fake_transformer_config(model, subfolder, local_files_only, *, revision): + calls["transformer_config"] = (model, subfolder, local_files_only, revision) + return {"component": "transformer"} + + monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch) + monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer) + monkeypatch.setattr(ltx2_components, "_load_component", fake_component) + monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config) + monkeypatch.setattr( + ltx2_components, + "create_transformer_from_config", + lambda *_args, **_kwargs: SimpleNamespace(config=SimpleNamespace(patch_size=1, patch_size_t=1)), + ) + + def fake_scheduler(*args, **kwargs): + calls["scheduler"] = (args, kwargs) + return SimpleNamespace(config={"use_dynamic_shifting": False, "shift_terminal": None}) + + monkeypatch.setattr(ltx2_components.FlowMatchEulerDiscreteScheduler, "from_pretrained", fake_scheduler) + + ltx2_components.initialize_pipeline_components(pipeline, od_config) + + assert pipeline.weights_sources[0].revision == revision + assert calls["prefetch"][2]["revision"] == revision + assert calls["tokenizer"][1]["revision"] == revision + assert {call[2] for call in calls["components"]} == { + "text_encoder", + "connectors", + "vae", + "audio_vae", + "vocoder", + "latent_upsampler", + } + assert all(call[3]["revision"] == revision for call in calls["components"]) + assert calls["transformer_config"] == ( + od_config.model, + profile.transformer_subfolder, + False, + revision, + ) + assert calls["scheduler"][1]["revision"] == revision + + +def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path): + (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"})) + (tmp_path / "transformer").mkdir() + (tmp_path / "transformer" / "config.json").write_text(json.dumps({"ff_bias": False})) + + assert detect_ltx_model_version(str(tmp_path)) == "2.3" + + +def test_ltx25_checkpoint_selects_full_one_stage_profile(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text( + json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]}) + ) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + + pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False)) + + assert pipe.model_version == "2.5" + assert pipe.component_profile is LTX25_FULL_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX25_FULL_RECIPE + assert pipe.preserve_sp_padded_audio_duration + assert pipe.reports_stage_durations + + +def test_ltx25_full_recipe_matches_official_sft_defaults(): + (phase,) = LTX25_FULL_RECIPE.phases + + assert (LTX25_FULL_RECIPE.height, LTX25_FULL_RECIPE.width) == (544, 960) + assert LTX25_FULL_RECIPE.num_frames == 121 + assert LTX25_FULL_RECIPE.frame_rate == 24.0 + assert LTX25_FULL_RECIPE.num_inference_steps == 30 + assert phase.guidance.video == LTXModalityGuidance( + cfg_scale=3.0, + stg_scale=1.0, + modality_scale=3.0, + rescale_scale=0.7, + stg_blocks=(28,), + ) + assert phase.guidance.audio == LTXModalityGuidance( + cfg_scale=7.0, + stg_scale=1.0, + modality_scale=3.0, + rescale_scale=0.7, + stg_blocks=(28,), + ) + assert phase.sigmas is None + assert phase.noise_scale == 1.0 + assert phase.use_official_sigma_schedule + assert LTX25_FULL_COMPONENT_PROFILE.transformer_subfolder == "transformer_full" + assert LTX25_FULL_COMPONENT_PROFILE.scheduler_use_dynamic_shifting + assert LTX25_FULL_COMPONENT_PROFILE.scheduler_shift_terminal == 0.1 + + +def test_ltx25_full_sigma_schedule_uses_official_default_sequence_anchor(): + scheduler = SimpleNamespace( + config={ + "base_image_seq_len": 1024, + "max_image_seq_len": 4096, + "base_shift": 0.95, + "max_shift": 2.05, + "shift_terminal": 0.1, + } + ) + + sigmas = _official_ltx_sigmas( + scheduler, + steps=30, + device=torch.device("cpu"), + ) + + torch.testing.assert_close( + sigmas[[0, 1, -2, -1]], + torch.tensor([1.0, 0.9949570, 0.1, 0.0]), + rtol=1e-5, + atol=1e-6, + ) + + +def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text( + json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]}) + ) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + + pipe = LTX2DistilledTwoStagePipeline( + od_config=SimpleNamespace( + model=str(tmp_path), + enable_diffusion_pipeline_profiler=False, + ) + ) + + assert pipe.component_profile is LTX25_DISTILLED_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE + + +def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "LTX2Pipeline", "model_version": "2.5"})) + for subfolder, marker in (("transformer", "distilled"), ("transformer_full", "full")): + path = tmp_path / subfolder + path.mkdir() + (path / "config.json").write_text(json.dumps({"marker": marker})) + + config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="LTX2Pipeline") + config.enrich_config() + + assert config.tf_model_config.get("marker") == "full" + + +def test_ltx25_four_public_pipeline_semantics_are_disjoint(): + assert resolve_ltx_component_profile("one_stage", "2.5") is LTX25_FULL_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("one_stage", "2.5") is LTX25_FULL_RECIPE + assert resolve_ltx_component_profile("two_stage", "2.5") is LTX25_TWO_STAGE_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("two_stage", "2.5") is LTX25_TWO_STAGE_RECIPE + assert resolve_ltx_component_profile("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_RECIPE + assert resolve_ltx_component_profile("distilled_two_stage", "2.5") is LTX25_DISTILLED_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.5") is LTX25_DISTILLED_TWO_STAGE_RECIPE + + assert LTX2Pipeline.pipeline_kind == "one_stage" + assert LTX2TwoStagePipeline.pipeline_kind == "two_stage" + assert LTX2DistilledOneStagePipeline.pipeline_kind == "distilled_one_stage" + assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage" + + +@pytest.mark.parametrize( + ("pipeline_kind", "expected"), + [ + ("one_stage", "regular"), + ("two_stage", "regular"), + ("distilled_one_stage", "distilled"), + ("distilled_two_stage", "distilled"), + ("dmd2", None), + ], +) +def test_ltx_checkpoint_kind_is_derived_from_pipeline_kind(pipeline_kind, expected): + assert resolve_ltx_checkpoint_kind(pipeline_kind) == expected + + +def test_ltx_checkpoint_kind_rejects_unknown_pipeline(): + with pytest.raises(ValueError, match="Unsupported LTX pipeline kind"): + resolve_ltx_checkpoint_kind("unknown") + + +def test_ltx25_distilled_two_stage_recipe_matches_model_card_resolution(): + stage1, stage2 = LTX25_DISTILLED_TWO_STAGE_RECIPE.phases + + assert (LTX25_DISTILLED_TWO_STAGE_RECIPE.height, LTX25_DISTILLED_TWO_STAGE_RECIPE.width) == (1088, 1920) + assert stage1.spatial_downscale == 2 + assert stage1.sigmas == LTX_DISTILLED_SIGMAS + assert stage1.sampler == "euler_ancestral" + assert stage2.sampler == "euler" + assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0) + assert stage2.input_transform == "spatial_upsample" + + assert LTX25_DISTILLED_TWO_STAGE_RECIPE.negative_prompt == "" + assert not LTX25_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt + + +def test_ltx25_distilled_one_stage_uses_native_low_resolution(): + (phase,) = LTX25_DISTILLED_ONE_STAGE_RECIPE.phases + + assert (LTX25_DISTILLED_ONE_STAGE_RECIPE.height, LTX25_DISTILLED_ONE_STAGE_RECIPE.width) == (544, 960) + assert phase.spatial_downscale == 1 + assert phase.sigmas == LTX_DISTILLED_SIGMAS + assert phase.sampler == "euler_ancestral" + assert not LTX25_DISTILLED_ONE_STAGE_RECIPE.supports_cache_dit + + +def test_ltx25_regular_two_stage_uses_full_transformer_then_distilled_lora(): + stage1, stage2 = LTX25_TWO_STAGE_RECIPE.phases + + assert (LTX25_TWO_STAGE_RECIPE.height, LTX25_TWO_STAGE_RECIPE.width) == (1088, 1920) + assert LTX25_TWO_STAGE_RECIPE.num_inference_steps == 30 + assert stage1.spatial_downscale == 2 + assert stage1.guidance == LTX25_FULL_RECIPE.request_guidance + assert stage2.adapter_slot == "ltx_distilled" + assert LTX25_TWO_STAGE_RECIPE.allow_request_phase_sigmas + assert LTX25_TWO_STAGE_COMPONENT_PROFILE.transformer_subfolder == "transformer_full" + assert LTX25_TWO_STAGE_COMPONENT_PROFILE.distilled_lora_filename is not None + + +@pytest.mark.parametrize( + "recipe", + [ + LTX2_ONE_STAGE_RECIPE, + LTX23_ONE_STAGE_RECIPE, + LTX_POSITIVE_ONLY_RECIPE, + ], +) +def test_ltx_one_stage_recipes_declare_cache_dit_supported(recipe): + assert recipe.supports_cache_dit + + +@pytest.mark.parametrize( + "recipe", + [ + LTX25_FULL_RECIPE, + LTX25_DISTILLED_ONE_STAGE_RECIPE, + LTX2_DISTILLED_TWO_STAGE_RECIPE, + LTX25_DISTILLED_TWO_STAGE_RECIPE, + ], +) +def test_ltx25_and_multistage_recipes_declare_cache_dit_unsupported(recipe): + assert not recipe.supports_cache_dit + + +def test_ltx_positive_only_guidance_preserves_official_x0_roundtrip(): + sample = torch.full((1, 1), 1e8, dtype=torch.float32) + velocity = torch.ones((1, 1), dtype=torch.bfloat16) + sigma = torch.tensor(0.5) + + actual = LTX_GUIDANCE_EXECUTOR._guide_modality( + sample, + {"cond": velocity}, + sigma, + LTXModalityGuidance(), + ) + expected = velocity_from_x0(sample, x0_from_velocity(sample, velocity, sigma), sigma) + + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + +def test_ltx_velocity_from_x0_materializes_official_scalar_sigma(): + item_calls = [] + + class TrackingTensor(torch.Tensor): + @staticmethod + def __new__(cls): + return torch.Tensor._make_subclass(cls, torch.tensor(0.725), False) + + def item(self): + item_calls.append(True) + return super().item() + + sample = torch.tensor([[1.0]], dtype=torch.bfloat16) + x0 = torch.tensor([[0.5]], dtype=torch.bfloat16) + + velocity_from_x0(sample, x0, TrackingTensor()) + + assert item_calls == [True] + + +def test_ltx_ancestral_positive_only_guidance_preserves_raw_velocity(): + sample = torch.full((1, 1), 1e8, dtype=torch.float32) + velocity = torch.ones((1, 1), dtype=torch.bfloat16) + + actual = LTX_GUIDANCE_EXECUTOR._guide_modality( + sample, + {"cond": velocity}, + torch.tensor(0.5), + LTXModalityGuidance(), + preserve_positive_velocity=True, + ) + + assert actual is velocity + + +@pytest.mark.parametrize("model_version", ["2.5", "2.3"]) +def test_ltx_t2v_denoise_state_uses_official_bfloat16_dtype(model_version): + expected_dtype = torch.bfloat16 + captured = {} + + def prepare_latents(**kwargs): + captured["video_dtype"] = kwargs["dtype"] + return torch.empty(1, 1, 1, dtype=kwargs["dtype"]) + + def prepare_audio_latents(*_args, **kwargs): + captured["audio_dtype"] = kwargs["dtype"] + return torch.empty(1, 1, 1, dtype=kwargs["dtype"]), 1, 1 + + pipeline = SimpleNamespace( + model_version=model_version, + transformer=SimpleNamespace(config=SimpleNamespace(in_channels=128)), + prepare_latents=prepare_latents, + prepare_audio_latents=prepare_audio_latents, + audio_sampling_rate=48_000, + audio_hop_length=160, + audio_vae_temporal_compression_ratio=4, + audio_vae_mel_compression_ratio=4, + audio_vae=SimpleNamespace(config=SimpleNamespace(mel_bins=64, latent_channels=8)), + _resolve_audio_latent_length=lambda length, _latents: length, + ) + request = SimpleNamespace( + num_videos_per_prompt=1, + height=64, + width=64, + num_frames=1, + frame_rate=24.0, + generator=None, + latents=None, + audio_latents=None, + audio_latents_normalized=False, + ) + prompt = SimpleNamespace( + batch_size=1, + positive_connector_prompt_embeds=torch.empty(1, dtype=torch.bfloat16), + positive_connector_audio_prompt_embeds=torch.empty(1, dtype=torch.bfloat16), + ) + + LTXRuntime._prepare_video_latents_stage( + pipeline, + request, + prompt, + device=torch.device("cpu"), + noise_scale=1.0, + ) + LTXRuntime._prepare_audio_latents_stage( + pipeline, + request, + prompt, + device=torch.device("cpu"), + noise_scale=1.0, + ) + + assert captured == {"video_dtype": expected_dtype, "audio_dtype": expected_dtype} + + +def test_ltx_seeded_latents_match_official_packed_rng_layout(): + pipeline = SimpleNamespace( + model_version="2.5", + vae_spatial_compression_ratio=2, + vae_temporal_compression_ratio=2, + transformer_spatial_patch_size=1, + transformer_temporal_patch_size=1, + audio_vae_mel_compression_ratio=2, + od_config=SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)), + ) + device = torch.device("cpu") + actual_generator = torch.Generator(device=device).manual_seed(123) + expected_generator = torch.Generator(device=device).manual_seed(123) + + actual_video = ltx2_latents.prepare_video_latents( + pipeline, + batch_size=1, + num_channels_latents=2, + height=4, + width=6, + num_frames=3, + dtype=torch.float32, + device=device, + generator=actual_generator, + ) + actual_audio, original_length, padded_length = ltx2_latents.prepare_audio_latents( + pipeline, + batch_size=1, + num_channels_latents=2, + audio_latent_length=3, + num_mel_bins=4, + dtype=torch.float32, + device=device, + generator=actual_generator, + ) + + expected_video = torch.randn((1, 12, 2), generator=expected_generator) + expected_audio = torch.randn((1, 3, 4), generator=expected_generator) + + torch.testing.assert_close(actual_video, expected_video, rtol=0.0, atol=0.0) + torch.testing.assert_close(actual_audio, expected_audio, rtol=0.0, atol=0.0) + assert (original_length, padded_length) == (3, 3) + + +def test_ltx25_distilled_two_stage_executes_custom_phase_schedules(): + request_inputs = LTXRequestInputs( + prompt="prompt", + negative_prompt="", + height=64, + width=64, + num_frames=1, + frame_rate=24.0, + num_inference_steps=8, + guidance=LTXGuidanceSpec.positive_only(), + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.25, + decode_noise_scale=0.5, + output_type="np", + max_sequence_length=16, + ) + prompt_context = object() + phase_calls = [] + + def resolve_request_inputs(req, **kwargs): + return request_inputs + + def run_phase(req, inputs, *, prompt_context=None, **kwargs): + phase_recipe = kwargs["phase_recipe"] + assert kwargs["image"] is source_image + phase_calls.append((phase_recipe, inputs, prompt_context)) + if len(phase_calls) == 1: + assert prompt_context is None + assert inputs.decode_timestep == 0.25 + assert inputs.decode_noise_scale == 0.5 + context = prompt_context_sentinel + assert phase_recipe.name == "generate_lowres" + assert (inputs.height, inputs.width) == (32, 32) + assert inputs.num_inference_steps == 2 + assert kwargs["sigmas"] == stage_1_sigmas + assert kwargs["noise_scale"] == 1.0 + video = torch.ones(1, 128, 1, 1, 1) + audio = torch.full((1, 8, 1, 2), 2.0) + audio_for_next_phase = torch.full((1, 8, 1, 2), 5.0) + else: + assert prompt_context is prompt_context_sentinel + assert phase_recipe.name == "refine" + assert (inputs.height, inputs.width) == (64, 64) + assert inputs.latents.shape == (1, 128, 1, 2, 2) + torch.testing.assert_close(inputs.audio_latents, torch.full((1, 8, 1, 2), 5.0)) + assert inputs.audio_latents_normalized + assert inputs.guidance_scale == 1.0 + assert inputs.num_inference_steps == 2 + assert kwargs["sigmas"] == stage_2_sigmas + assert kwargs["noise_scale"] == stage_2_sigmas[0] + assert inputs.decode_timestep == 0.0 + assert inputs.decode_noise_scale is None + context = prompt_context + video = torch.full((1, 128, 1, 2, 2), 3.0) + audio = torch.full((1, 8, 1, 2), 4.0) + audio_for_next_phase = None + return LTXPhaseResult( + forward_context=SimpleNamespace(prompt_context=context), + video=video, + audio=audio, + audio_for_next_phase=audio_for_next_phase, + ) + + def decode_phase(phase): + return DiffusionOutput(output=(phase.video, phase.audio)) + + class FakeUpsampler(torch.nn.Module): + dtype = torch.float32 + + def forward(self, latents): + return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1) + + prompt_context_sentinel = prompt_context + stage_1_sigmas = [1.0, 0.5, 0.0] + stage_2_sigmas = [0.8, 0.2, 0.0] + pipeline = object.__new__(LTX2TwoStagePipeline) + torch.nn.Module.__init__(pipeline) + pipeline.pipeline_recipe = LTX25_DISTILLED_TWO_STAGE_RECIPE + pipeline.device = torch.device("cpu") + pipeline.vae_spatial_compression_ratio = 32 + pipeline.vae_temporal_compression_ratio = 8 + pipeline.latent_upsampler = FakeUpsampler() + object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs) + object.__setattr__(pipeline, "run_phase", run_phase) + object.__setattr__(pipeline, "decode_phase", decode_phase) + + source_image = object() + output = pipeline.forward( + SimpleNamespace(sampling_params_list=[], prompts=[]), + image=source_image, + stage_1_sigmas=stage_1_sigmas, + stage_2_sigmas=stage_2_sigmas, + ) + + assert len(phase_calls) == 2 + assert phase_calls[1][2] is prompt_context_sentinel + torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0)) + torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0)) + + +class TestLTXRequestParsing: + def test_request_resolves_custom_sigmas_from_extra_args(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + sigmas = [1.0, 0.5, 0.0] + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=OmniDiffusionSamplingParams(extra_args={"sigmas": sigmas}), + request_id="ltx-custom-sigmas", + ) + ] + ) + + assert _make_ltx_request_pipe(LTX2Pipeline)._resolve_request_sigmas(req, None) == sigmas + + def test_request_resolves_per_stage_sigmas_from_extra_args(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + stage_1_sigmas = [1.0, 0.4, 0.0] + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=OmniDiffusionSamplingParams(extra_args={"stage_1_sigmas": stage_1_sigmas}), + request_id="ltx-custom-phase-sigmas", + ) + ] + ) + + resolved = _make_ltx_request_pipe(LTX2TwoStagePipeline)._resolve_request_phase_sigmas( + req, [1.0, 0.0], [0.8, 0.0] + ) + + assert resolved == (stage_1_sigmas, [0.8, 0.0]) + + @pytest.mark.parametrize("value", [-1, 52, 1.5, True]) + def test_request_rejects_invalid_image_crf(self, value): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=OmniDiffusionSamplingParams(extra_args={"image_crf": value}), + request_id="ltx-invalid-image-crf", + ) + ] + ) + + with pytest.raises(ValueError, match="image_crf"): + _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req) + + +class TestLTXForwardStages: + def test_gemma_tokenization_matches_official_bos_and_left_padding(self): + pipe = _make_ltx_request_pipe(LTX2Pipeline) + + class FakeTokenizer: + bos_token_id = 2 + eos_token_id = 1 + eos_token = "" + pad_token_id = 0 + pad_token = "" + padding_side = "right" + + def __init__(self): + self.calls = [] + + def __call__(self, text, **kwargs): + self.calls.append((text, kwargs)) + ids = [7, 8, 9, self.eos_token_id] if text == "long" else [6, self.eos_token_id] + return SimpleNamespace(input_ids=ids) + + def pad(self, encoded, *, padding, max_length, return_tensors, return_attention_mask): + assert padding == "max_length" + assert return_tensors == "pt" + assert return_attention_mask is True + rows = [] + masks = [] + for ids in encoded["input_ids"]: + pad = max_length - len(ids) + rows.append([self.pad_token_id] * pad + ids) + masks.append([0] * pad + [1] * len(ids)) + return SimpleNamespace(input_ids=torch.tensor(rows), attention_mask=torch.tensor(masks)) + + class FakeTextEncoder(torch.nn.Module): + dtype = torch.float32 + + def __init__(self): + super().__init__() + self.seen_input_ids = None + self.seen_attention_mask = None + + def forward(self, *, input_ids, attention_mask, output_hidden_states): + assert output_hidden_states is True + self.seen_input_ids = input_ids.clone() + self.seen_attention_mask = attention_mask.clone() + return SimpleNamespace(hidden_states=(input_ids.unsqueeze(-1).float(),)) + + tokenizer = FakeTokenizer() + text_encoder = FakeTextEncoder() + object.__setattr__(pipe, "tokenizer", tokenizer) + object.__setattr__(pipe, "text_encoder", text_encoder) + + pipe._get_gemma_prompt_embeds( + [" long ", "short"], + max_sequence_length=4, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + torch.testing.assert_close(text_encoder.seen_input_ids, torch.tensor([[2, 7, 8, 9], [0, 2, 6, 1]])) + torch.testing.assert_close(text_encoder.seen_attention_mask, torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]])) + assert tokenizer.padding_side == "left" + assert [call[0] for call in tokenizer.calls] == ["long", "short"] diff --git a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py index e309dbf47..cbc66e69c 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py +++ b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py @@ -3,6 +3,7 @@ """Unit tests for LTX image-to-video input and conditioning behavior.""" +import sys from types import SimpleNamespace import pytest @@ -23,7 +24,7 @@ def _make_ltx23_request_pipe(cls): class TestLTXImageToVideoForwardStages: - def test_i2v_pil_preprocessing_matches_official_center_crop(self): + def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self): from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images pixels = torch.zeros(4, 8, 3, dtype=torch.uint8) @@ -35,6 +36,136 @@ class TestLTXImageToVideoForwardStages: torch.testing.assert_close(actual, expected) + def test_ltx25_i2v_pil_preprocessing_preserves_aspect_ratio_after_crf(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning + + pixels = torch.zeros(4, 8, 3, dtype=torch.uint8) + pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8) + image = Image.fromarray(pixels.numpy()) + monkeypatch.setattr(conditioning, "_apply_image_conditioning_crf", lambda image_array, _crf: image_array) + + actual = conditioning._preprocess_i2v_pil_images(image, height=4, width=4, crf=18) + expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0 + + torch.testing.assert_close(actual, expected) + + def test_ltx25_i2v_applies_crf18_before_resize(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + captured_crfs = [] + captured_prepare_kwargs = [] + + def fake_preprocess(image, *, height, width, crf, device, dtype): + del image, device + captured_crfs.append(crf) + return torch.zeros(1, 3, height, width, dtype=dtype) + + def fake_prepare_latents(**kwargs): + captured_prepare_kwargs.append(kwargs) + return kwargs["image"], None + + monkeypatch.setattr(conditioning, "_preprocess_i2v_pil_images", fake_preprocess) + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.model_version = "2.5" + object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4))) + object.__setattr__(pipe, "prepare_latents", fake_prepare_latents) + request_inputs = SimpleNamespace( + height=4, + width=4, + num_frames=3, + num_videos_per_prompt=1, + generator=None, + latents=None, + ) + prompt_context = SimpleNamespace( + batch_size=1, + positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16), + ) + + pipe._prepare_video_latents_stage( + request_inputs, + prompt_context, + device=torch.device("cpu"), + noise_scale=0.0, + image=Image.new("RGB", (8, 4)), + ) + + assert captured_crfs == [18] + assert captured_prepare_kwargs[0]["dtype"] is torch.bfloat16 + assert captured_prepare_kwargs[0]["image"].dtype is torch.bfloat16 + + def test_ltx25_i2v_rejects_inputs_that_cannot_receive_crf18(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.model_version = "2.5" + + with pytest.raises(ValueError, match="requires PIL images"): + pipe._prepare_video_latents_stage( + SimpleNamespace(), + SimpleNamespace(), + device=torch.device("cpu"), + noise_scale=0.0, + image=torch.zeros(3, 4, 4), + ) + + def test_ltx25_i2v_tensor_input_can_opt_out_of_crf18(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + captured_prepare_kwargs = [] + + def fake_prepare_latents(**kwargs): + captured_prepare_kwargs.append(kwargs) + return kwargs["image"], None + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.model_version = "2.5" + object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4))) + object.__setattr__(pipe, "prepare_latents", fake_prepare_latents) + request_inputs = SimpleNamespace( + height=4, + width=4, + num_frames=3, + num_videos_per_prompt=1, + generator=None, + latents=None, + image_crf=0, + ) + prompt_context = SimpleNamespace( + batch_size=1, + positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16), + ) + + pipe._prepare_video_latents_stage( + request_inputs, + prompt_context, + device=torch.device("cpu"), + noise_scale=0.0, + image=torch.zeros(3, 4, 4), + ) + + assert captured_prepare_kwargs[0]["image"].shape == (1, 3, 4, 4) + + def test_ltx25_i2v_crf_preserves_single_pixel_dimension(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf + + image = torch.zeros(1, 4, 3, dtype=torch.uint8).numpy() + monkeypatch.setitem(sys.modules, "av", None) + + assert _apply_image_conditioning_crf(image, 18) is image + + def test_ltx25_i2v_reports_missing_pyav(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf + + monkeypatch.setitem(sys.modules, "av", None) + + with pytest.raises(ImportError, match="PyAV with a libx264 encoder"): + _apply_image_conditioning_crf(torch.zeros(4, 4, 3, dtype=torch.uint8).numpy(), 18) + def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -98,7 +229,7 @@ class TestLTXImageToVideoForwardStages: with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"): pipe._resolve_request_image(mixed_req, None, request_inputs) - def test_denoise_timestep_kwargs_masks_video_only(self): + def test_denoise_timestep_kwargs_masks_video_and_keeps_audio_per_token(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline pipe = object.__new__(LTX2Pipeline) @@ -189,13 +320,18 @@ class TestLTXImageToVideoConditioning: torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]])) - def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch): + @pytest.mark.parametrize( + ("model_version", "sampled_shape"), + [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))], + ) + def test_i2v_5d_latents_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape): import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) + pipe.model_version = model_version pipe.vae_spatial_compression_ratio = 1 pipe.vae_temporal_compression_ratio = 1 pipe.transformer_spatial_patch_size = 1 @@ -237,14 +373,20 @@ class TestLTXImageToVideoConditioning: torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]])) - assert sampled_shapes == [(1, 3, 2)] + assert sampled_shapes == [sampled_shape] - def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch): + @pytest.mark.parametrize( + ("model_version", "sampled_shape"), + [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))], + ) + def test_i2v_image_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape): import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning + import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) + pipe.model_version = model_version pipe.vae_spatial_compression_ratio = 1 pipe.vae_temporal_compression_ratio = 1 pipe.transformer_spatial_patch_size = 1 @@ -266,6 +408,7 @@ class TestLTXImageToVideoConditioning: return torch.ones(shape, device=device, dtype=dtype) monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor) + monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor) out, conditioning_mask = pipe.prepare_latents( image=torch.zeros(1, 3, 1, 1), @@ -279,10 +422,54 @@ class TestLTXImageToVideoConditioning: device=torch.device("cpu"), ) - assert sampled_shapes == [(1, 3, 2)] + assert sampled_shapes == [sampled_shape] torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]])) + def test_ltx25_re_noise_uses_official_scalar_noise_scale(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.model_version = "2.5" + pipe.vae_spatial_compression_ratio = 1 + pipe.vae_temporal_compression_ratio = 1 + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + pipe.vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(scaling_factor=1.0), + ) + object.__setattr__( + pipe, + "_encode_i2v_image_latents", + lambda *_args, **_kwargs: torch.zeros(1, 2, 1, 1, 1, dtype=torch.bfloat16), + ) + captured = {} + + def fake_create_noised_state(latents, noise_scale, generator=None): + del generator + captured["noise_scale"] = noise_scale + return latents + + monkeypatch.setattr(ltx2_latents, "create_noised_state", fake_create_noised_state) + pipe.prepare_latents( + image=torch.zeros(1, 3, 1, 1), + batch_size=1, + num_channels_latents=2, + height=1, + width=1, + num_frames=3, + noise_scale=0.909375, + dtype=torch.bfloat16, + device=torch.device("cpu"), + latents=torch.zeros(1, 2, 3, 1, 1, dtype=torch.bfloat16), + ) + + assert captured["noise_scale"] == 0.909375 + def test_ltx23_i2v_video_step_preserves_conditioning_frame(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline @@ -385,6 +572,99 @@ class TestLTXImageToVideoConditioning: torch.testing.assert_close(audio, torch.tensor([[[-2.0]]])) assert pipeline.step_indices == [0, 1] + def test_ltx25_ancestral_adapter_matches_official_seeded_step_and_preserves_i2v_mask(self): + from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter + + device = torch.device("cpu") + sigmas = torch.tensor([1.0, 0.5, 0.0]) + pipeline = SimpleNamespace(device=device, scheduler=SimpleNamespace(sigmas=sigmas)) + mask = torch.tensor([[1.0, 0.0]]) + adapter = LTXVideoAudioStepAdapter( + pipeline, + SimpleNamespace(sigmas=sigmas.clone()), + latent_num_frames=1, + latent_height=1, + latent_width=1, + image_conditioned=True, + sampler="euler_ancestral", + generator=torch.Generator(device=device).manual_seed(42), + conditioning_mask=mask, + ) + video = torch.tensor([[[3.0], [2.0]]]) + audio = torch.tensor([[[4.0]]]) + video_velocity = torch.tensor([[[7.0], [0.25]]]) + audio_velocity = torch.tensor([[[0.5]]]) + + reference_generator = torch.Generator(device=device).manual_seed(10042) + video_noise = torch.randn(video.shape, generator=reference_generator) + audio_noise = torch.randn(audio.shape, generator=reference_generator) + + def official_step(sample, velocity, noise): + sigma = sigmas[0] + sigma_next = sigmas[1] + denoised = sample.float() - velocity.float() * sigma + sigma_down = sigma_next * (sigma_next / sigma) + ratio = sigma_down / sigma + deterministic = ratio * sample.float() + (1.0 - ratio) * denoised + alpha_next = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt() + return (alpha_next / alpha_down) * deterministic + noise * coeff + + expected_video = official_step(video, video_velocity, video_noise) + expected_video[:, :1] = video[:, :1] + expected_audio = official_step(audio, audio_velocity, audio_noise) + + ((actual_video, actual_audio),) = adapter.step( + (video_velocity, audio_velocity), + (sigmas[0], sigmas[0]), + (video, audio), + ) + torch.testing.assert_close(actual_video, expected_video) + torch.testing.assert_close(actual_audio, expected_audio) + + terminal_video_velocity = torch.full_like(actual_video, 0.125) + terminal_audio_velocity = torch.full_like(actual_audio, 0.25) + ((terminal_video, terminal_audio),) = adapter.step( + (terminal_video_velocity, terminal_audio_velocity), + (sigmas[1], sigmas[1]), + (actual_video, actual_audio), + ) + expected_terminal_video = actual_video - terminal_video_velocity * sigmas[1] + expected_terminal_video[:, :1] = actual_video[:, :1] + expected_terminal_audio = actual_audio - terminal_audio_velocity * sigmas[1] + torch.testing.assert_close(terminal_video, expected_terminal_video) + torch.testing.assert_close(terminal_audio, expected_terminal_audio) + + def test_ltx25_ancestral_quantizes_x0_to_model_dtype_before_fp32_step(self): + from vllm_omni.diffusion.models.ltx2.ltx2_denoise import _ancestral_euler_step_from_velocity + + sample = torch.tensor([[[3.140625, -1.234375]]], dtype=torch.bfloat16) + velocity = torch.tensor([[[0.333984375, 2.71875]]], dtype=torch.bfloat16) + sigmas = torch.tensor([0.73, 0.41, 0.0], dtype=torch.float32) + generator = torch.Generator(device="cpu").manual_seed(10042) + reference_generator = torch.Generator(device="cpu").manual_seed(10042) + + sigma, sigma_next = sigmas[:2] + denoised = (sample.float() - velocity.float() * sigma).to(sample.dtype).float() + sigma_down = sigma_next * (sigma_next / sigma) + sigma_down_ratio = sigma_down / sigma + expected = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised + alpha_next = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + coefficient = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt() + noise = torch.randn( + sample.shape, + generator=reference_generator, + dtype=sample.dtype, + device=sample.device, + ) + expected = ((alpha_next / alpha_down) * expected + noise.float() * coefficient).to(sample.dtype) + + actual = _ancestral_euler_step_from_velocity(sample, velocity, sigmas, 0, generator) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self): from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py index cad5c300d..a9d19945a 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_latents.py +++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py @@ -202,3 +202,68 @@ def test_prepare_audio_latents_rejects_incompatible_provided_length(): device=torch.device("cpu"), latents=latents, ) + + +def test_create_noised_state_matches_official_fp32_lerp(): + latents = torch.linspace(-2, 2, 4096, dtype=torch.bfloat16).reshape(1, 32, 128) + expected_generator = torch.Generator().manual_seed(42) + noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype) + expected = torch.lerp(latents.float(), noise.float(), torch.tensor(0.15, dtype=torch.bfloat16).float()).to( + latents.dtype + ) + + actual = latent_ops.create_noised_state( + latents, + torch.tensor(0.15, dtype=torch.bfloat16), + torch.Generator().manual_seed(42), + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_create_conditioned_noised_state_matches_official_two_lerps(): + latents = torch.linspace(-2, 2, 24, dtype=torch.bfloat16).reshape(1, 3, 8) + clean_latents = latents.clone() + clean_latents[:, 0] = 3.0 + denoise_mask = torch.tensor([[[0.0], [1.0], [1.0]]]) + expected_generator = torch.Generator().manual_seed(42) + noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype) + noised = torch.lerp(latents.float(), noise.float(), 0.909375) + expected = torch.lerp(clean_latents.float(), noised, denoise_mask).to(latents.dtype) + + actual = latent_ops.create_conditioned_noised_state( + latents, + clean_latents, + denoise_mask, + 0.909375, + torch.Generator().manual_seed(42), + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_prepare_supplied_video_latents_uses_official_token_major_layout(): + pipeline = _make_pipeline(LTX2Pipeline) + supplied = torch.arange(512, dtype=torch.float32).reshape(1, 32, 16) + actual_video = pipeline.prepare_latents( + batch_size=1, + num_channels_latents=4, + height=64, + width=64, + num_frames=9, + dtype=torch.float32, + device=torch.device("cpu"), + latents=supplied, + ) + + torch.testing.assert_close(actual_video, supplied, rtol=0, atol=0) + assert actual_video.stride()[1:] == (1, actual_video.shape[1]) + + +def test_clear_audio_padding_keeps_padding_outside_sampler_state(): + updated = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [99.0, -99.0]]]) + + actual = latent_ops.clear_audio_padding(updated, 2) + + torch.testing.assert_close(actual[:, :2], updated[:, :2]) + torch.testing.assert_close(actual[:, 2:], torch.zeros_like(actual[:, 2:])) diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py index 95ca6e868..0573dfb37 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py +++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py @@ -28,6 +28,15 @@ def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker): class TestCFGParallelHelpers: """Test LTX-2.3 CFG helper math without loading model weights.""" + def test_single_guidance_pass_preserves_official_token_layout(self): + from vllm_omni.diffusion.models.ltx2.ltx2_guidance import _repeat_batch + + tokens = torch.randn(1, 16, 32).transpose(1, 2).contiguous().transpose(1, 2) + repeated = _repeat_batch(tokens, 1) + + assert repeated is tokens + assert repeated.stride() == tokens.stride() + def test_combine_cfg_noise_matches_x0_space_formula(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index 2f3ae522e..b3da177db 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -20,6 +20,7 @@ from vllm_omni.diffusion.models.ltx2 import ( from vllm_omni.diffusion.models.ltx2.ltx2_components import ( LTX2_COMPONENT_PROFILE, LTX2_DISTILLED_COMPONENT_PROFILE, + LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, LTX2_TWO_STAGE_COMPONENT_PROFILE, LTX23_COMPONENT_PROFILE, LTX23_DISTILLED_COMPONENT_PROFILE, @@ -47,6 +48,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( ) from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState from vllm_omni.diffusion.models.ltx2.ltx2_recipes import ( + LTX2_DISTILLED_ONE_STAGE_RECIPE, LTX2_DISTILLED_TWO_STAGE_RECIPE, LTX2_ONE_STAGE_RECIPE, LTX2_TWO_STAGE_RECIPE, @@ -64,12 +66,14 @@ from vllm_omni.diffusion.models.ltx2.ltx2_request import ( ) from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2DistilledOneStagePipeline, LTX2I2VDMD2Pipeline, LTX2Pipeline, LTX2T2VDMD2Pipeline, ) from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( LTX2DistilledPipeline, + LTX2DistilledTwoStagePipeline, LTX2TwoStagePipeline, ) @@ -125,15 +129,21 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries(): assert issubclass(LTX2Pipeline, LTXRuntime) assert issubclass(LTX2TwoStagePipeline, LTXRuntime) + assert issubclass(LTX2DistilledOneStagePipeline, LTXRuntime) + assert issubclass(LTX2DistilledTwoStagePipeline, LTXRuntime) assert issubclass(LTX2DistilledPipeline, LTXRuntime) + assert LTX2DistilledPipeline is LTX2DistilledTwoStagePipeline assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE + assert LTX2DistilledOneStagePipeline.component_profile is LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE + assert LTX2DistilledOneStagePipeline.pipeline_recipe is LTX2_DISTILLED_ONE_STAGE_RECIPE assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE for pipeline_cls, profile in ( (LTX2Pipeline, LTX2_COMPONENT_PROFILE), (LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE), + (LTX2DistilledOneStagePipeline, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE), (LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE), ): assert pipeline_cls._dit_modules == list(profile.dit_modules) @@ -254,7 +264,7 @@ def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_pa monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) - with pytest.raises(ValueError, match="only one-stage LTX recipes"): + with pytest.raises(ValueError, match="not qualified for this LTX recipe"): LTX2TwoStagePipeline( od_config=SimpleNamespace( model=str(tmp_path), @@ -709,15 +719,48 @@ def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch): assert len(installed) == 2 assert all(processor.__class__.__name__ == "_LTXConnectorAttnProcessor" for processor in installed) - assert all(processor.has_learned_registers for processor in installed) assert all(kernel["role"] == "ltx2.connector" for kernel in kernels) assert all(kernel["role_category"] == "self" for kernel in kernels) assert all(kernel["skip_sequence_parallel"] for kernel in kernels) assert all(kernel["disable_kv_quant"] for kernel in kernels) + assert all(processor.has_learned_registers for processor in installed) + assert not any(processor.preserve_learned_register_mask for processor in installed) -@pytest.mark.parametrize("has_learned_registers", [False, True]) -def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers): +def test_ltx25_connector_attention_preserves_learned_register_mask(monkeypatch): + installed = [] + + class OmniAttention: + def __init__(self, **_kwargs): + pass + + class Attention: + heads = 32 + head_dim = 128 + inner_kv_dim = 4096 + + def set_processor(self, processor): + installed.append(processor) + + monkeypatch.setattr(ltx2_components, "OmniAttention", OmniAttention) + connectors = SimpleNamespace( + video_connector=SimpleNamespace( + learnable_registers=object(), + transformer_blocks=[SimpleNamespace(attn1=Attention())], + ), + audio_connector=SimpleNamespace( + learnable_registers=object(), + transformer_blocks=[SimpleNamespace(attn1=Attention())], + ), + ) + + _install_connector_attention(connectors, preserve_learned_register_mask=True) + + assert all(processor.has_learned_registers for processor in installed) + assert all(processor.preserve_learned_register_mask for processor in installed) + + +def test_ltx_connector_attention_dispatches_through_omni_kernel(): calls = [] class Backend: @@ -732,12 +775,16 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi calls.append((query, key, value, metadata)) return query + class UpcastingNorm(torch.nn.Module): + def forward(self, tensor): + return tensor.float() + attention = SimpleNamespace( heads=2, head_dim=4, inner_kv_dim=8, - norm_q=torch.nn.Identity(), - norm_k=torch.nn.Identity(), + norm_q=UpcastingNorm(), + norm_k=UpcastingNorm(), to_q=torch.nn.Identity(), to_k=torch.nn.Identity(), to_v=torch.nn.Identity(), @@ -746,24 +793,29 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi rope_type="split", omni_attention=OmniAttention(), ) - hidden_states = torch.randn(2, 3, 8) + hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16) additive_mask = torch.zeros(2, 1, 3, 3) + rotary = ( + torch.randn(2, 2, 3, 2, dtype=torch.float32), + torch.randn(2, 2, 3, 2, dtype=torch.float32), + ) - output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)( + output = ltx2_components._LTXConnectorAttnProcessor()( attention, hidden_states, attention_mask=additive_mask, + query_rotary_emb=rotary, ) query, key, value, metadata = calls[0] assert query.shape == key.shape == value.shape == (2, 3, 2, 4) - if has_learned_registers: - assert metadata is None - else: - assert metadata.attn_mask.shape == (2, 3) - assert metadata.attn_mask.dtype == torch.bool - assert metadata.attn_mask.all() - torch.testing.assert_close(output, hidden_states) + assert query.dtype == key.dtype == value.dtype == torch.bfloat16 + assert metadata.attn_mask.shape == (2, 3) + assert metadata.attn_mask.dtype == torch.bool + assert metadata.attn_mask.all() + expected_rotary = tuple(component.to(torch.bfloat16) for component in rotary) + expected = ltx2_components.apply_split_rotary_emb(hidden_states, expected_rotary, head_dim=4) + torch.testing.assert_close(output, expected, rtol=0, atol=0) def test_ltx2_distilled_pipeline_shares_recipe_runtime(): @@ -1707,6 +1759,7 @@ class TestLTXForwardStages: assert seen["request_inputs"].output_type == "latent" assert seen["kwargs"] == { "request_sigmas": request_sigmas, + "request_phase_sigmas": None, "image": None, } @@ -1824,6 +1877,16 @@ class TestRegistryIntegration: "pipeline_ltx2_two_stage", "LTX2TwoStagePipeline", ) + assert _DIFFUSION_MODELS["LTX2DistilledOneStagePipeline"] == ( + "ltx2", + "pipeline_ltx2", + "LTX2DistilledOneStagePipeline", + ) + assert _DIFFUSION_MODELS["LTX2DistilledTwoStagePipeline"] == ( + "ltx2", + "pipeline_ltx2_two_stage", + "LTX2DistilledTwoStagePipeline", + ) assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == ( "ltx2", "pipeline_ltx2_two_stage", @@ -1848,6 +1911,8 @@ class TestRegistryIntegration: expected = [ "LTX2Pipeline", "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", + "LTX2DistilledTwoStagePipeline", "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", @@ -1861,6 +1926,8 @@ class TestRegistryIntegration: [ "LTX2Pipeline", "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", + "LTX2DistilledTwoStagePipeline", "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index f74b2991d..e9bba1f1f 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -7,18 +7,54 @@ import pytest import torch from torch import nn -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig +from vllm_omni.diffusion.models.ltx2 import ltx2_transformer from vllm_omni.diffusion.models.ltx2.ltx2_transformer import ( - LTX2AudioVideoAttnProcessor, + LTX2FeedForward, LTX2VideoTransformer3DModel, - _LTX2ParallelAttention, + LTX2VideoTransformerBlock, _make_rms_norm, + apply_interleaved_rotary_emb, + apply_keyframes_absolute_embedding, + apply_split_rotary_emb, ) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +def test_ltx_interleaved_rope_matches_official_model_dtype_arithmetic(): + torch.manual_seed(0) + x = torch.randn(2, 3, 8, dtype=torch.bfloat16) + cos = torch.randn(2, 3, 8, dtype=x.dtype) + sin = torch.randn(2, 3, 8, dtype=x.dtype) + x_real, x_imag = x.unflatten(2, (4, 2)).unbind(-1) + rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) + expected = x * cos + rotated * sin + + actual = apply_interleaved_rotary_emb(x, (cos, sin)) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_ltx_split_rope_matches_official_model_dtype_arithmetic(): + torch.manual_seed(1) + x = torch.randn(2, 3, 8, dtype=torch.bfloat16) + cos = torch.randn(2, 3, 4, dtype=x.dtype) + sin = torch.randn(2, 3, 4, dtype=x.dtype) + split_x = x.reshape(2, 3, 2, 4) + first_x, second_x = split_x.chunk(2, dim=-2) + cos_u = cos.unsqueeze(-2) + sin_u = sin.unsqueeze(-2) + expected = split_x * cos_u + expected[..., :1, :].addcmul_(-sin_u, second_x) + expected[..., 1:, :].addcmul_(sin_u, first_x) + expected = expected.reshape_as(x) + + actual = apply_split_rotary_emb(x, (cos, sin), head_dim=8) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_ltx_rms_norm_no_affine_matches_official_torch_module(): norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False) @@ -36,11 +72,53 @@ def test_ltx_rms_norm_affine_weight_remains_parameter(): assert "weight" in norm.state_dict() -def test_ltx_transformer_has_separate_cfg_cache_dit_config(): +def test_ltx_attention_restores_projection_dtype_after_upcasting_qk_norms(): + calls = [] + + class UpcastingNorm(nn.Module): + def forward(self, tensor): + return tensor.float() + + class RecordingAttention: + def __call__(self, query, key, value, metadata): + calls.append((query, key, value, metadata)) + return query + + attention = SimpleNamespace( + heads=2, + head_dim=4, + to_qkv=None, + to_q=nn.Identity(), + to_k=nn.Identity(), + to_v=nn.Identity(), + norm_q=UpcastingNorm(), + norm_k=UpcastingNorm(), + to_gate_logits=None, + attn=RecordingAttention(), + to_out=(nn.Identity(), nn.Identity()), + rope_type="split", + ) + hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16) + rotary = ( + torch.ones(2, 2, 3, 2, dtype=torch.float32), + torch.zeros(2, 2, 3, 2, dtype=torch.float32), + ) + + processor = ltx2_transformer.LTX2AudioVideoAttnProcessor() + processor._slice_rope_for_tp = lambda rope, _: rope + output = processor(attention, hidden_states, query_rotary_emb=rotary) + + query, key, value, metadata = calls[0] + assert query.dtype == key.dtype == value.dtype == torch.bfloat16 + assert metadata is None + torch.testing.assert_close(output, hidden_states) + + +def test_ltx_transformer_has_fused_cfg_cache_dit_config(): adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config") assert isinstance(adapter_config, CacheDiTAdapterConfig) - assert adapter_config.has_separate_cfg + assert not adapter_config.has_separate_cfg def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks(): @@ -76,8 +154,9 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type): def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch): - monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True)) - processor = LTX2AudioVideoAttnProcessor() + processor_cls = ltx2_transformer.LTX2AudioVideoAttnProcessor + monkeypatch.setattr(processor_cls, "_is_sp_enabled", staticmethod(lambda: True)) + processor = processor_cls() hidden_states = torch.zeros(1, 3, 4) encoder_hidden_states = torch.zeros(1, 4, 4) key_padding_mask = torch.tensor([[True, True, True, False]]) @@ -113,7 +192,7 @@ def _fake_ltx_parallel_attention(): return forward - layer = object.__new__(_LTX2ParallelAttention) + layer = object.__new__(ltx2_transformer._LTX2ParallelAttention) nn.Module.__init__(layer) layer.attn_backend = Backend layer.attention = SimpleNamespace(forward=run("native")) @@ -123,33 +202,281 @@ def _fake_ltx_parallel_attention(): @pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)]) -def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length): +def test_ltx_masked_attention_uses_sdpa(query_length, key_length): layer, calls = _fake_ltx_parallel_attention() query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16) key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16) key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]]) - metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :]) + metadata = ltx2_transformer.AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :]) layer._run_local_attention(query, key, value, metadata) assert calls == ["sdpa"] -def test_ltx_masked_self_attention_uses_sdpa(): - layer, calls = _fake_ltx_parallel_attention() - query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16) - key_padding_mask = torch.tensor([[True, True, True, False]]) - metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :]) - - layer._run_local_attention(query, key, value, metadata) - - assert calls == ["sdpa"] - - -def test_ltx_unmasked_cross_attention_keeps_native_backend(): +def test_ltx_unmasked_attention_keeps_native_backend(): layer, calls = _fake_ltx_parallel_attention() query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16) layer._run_local_attention(query, key, value, None) assert calls == ["native"] + + +def test_ltx25_feed_forward_bias_flag_controls_both_linears(monkeypatch): + class FakeParallelLinear(nn.Linear): + def __init__(self, dim_in, dim_out, *, bias=True, **kwargs): + super().__init__(dim_in, dim_out, bias=bias) + + monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeParallelLinear) + monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeParallelLinear) + + without_bias = LTX2FeedForward(8, inner_dim=16, bias=False) + with_bias = LTX2FeedForward(8, inner_dim=16, bias=True) + + assert without_bias.net[0].proj.bias is None + assert without_bias.net[2].bias is None + assert with_bias.net[0].proj.bias is not None + assert with_bias.net[2].bias is not None + + +def _build_tiny_ltx_transformer(monkeypatch, **overrides): + class FakeBlock(nn.Module): + def __init__(self, **kwargs): + super().__init__() + self.init_kwargs = kwargs + + def forward(self, hidden_states, audio_hidden_states, *_args, **_kwargs): + return hidden_states, audio_hidden_states + + monkeypatch.setattr(ltx2_transformer, "LTX2VideoTransformerBlock", FakeBlock) + kwargs = { + "in_channels": 4, + "out_channels": 4, + "num_attention_heads": 2, + "attention_head_dim": 4, + "cross_attention_dim": 8, + "vae_scale_factors": (1, 1, 1), + "pos_embed_max_pos": 4, + "base_height": 8, + "base_width": 8, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 4, + "audio_cross_attention_dim": 8, + "audio_scale_factor": 1, + "audio_pos_embed_max_pos": 4, + "audio_sampling_rate": 16, + "audio_hop_length": 4, + "num_layers": 1, + "caption_channels": 8, + "use_prompt_embeddings": False, + "cross_attn_mod": True, + "audio_cross_attn_mod": True, + } + kwargs.update(overrides) + return LTX2VideoTransformer3DModel(**kwargs) + + +def test_ltx25_keyframe_absolute_embedding_matches_official(): + hidden_states = torch.arange(32, dtype=torch.bfloat16).view(1, 4, 8) + keyframes_mask = torch.tensor([[[1.0], [1.0], [0.0], [0.0]]]) + embedding = torch.linspace(-0.5, 0.5, 8, dtype=torch.bfloat16).view(1, 8) + + actual = apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, embedding) + expected = hidden_states + (keyframes_mask > 0).to(torch.bfloat16) * embedding + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert apply_keyframes_absolute_embedding(hidden_states, None, embedding) is hidden_states + + +def test_ltx25_constructor_flags_and_native_weight_loader(monkeypatch): + legacy_model = _build_tiny_ltx_transformer(monkeypatch) + + assert legacy_model.config.ff_bias is True + assert legacy_model.config.audio_ff_bias is True + assert legacy_model.config.use_prompt_adaln_single is True + assert legacy_model.config.use_keyframes_abs_pos_embedding is False + assert hasattr(legacy_model, "prompt_adaln") + assert hasattr(legacy_model, "audio_prompt_adaln") + assert not hasattr(legacy_model, "keyframes_abs_pos_embedding") + assert legacy_model.transformer_blocks[0].init_kwargs["ff_bias"] is True + assert legacy_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True + + ltx25_model = _build_tiny_ltx_transformer( + monkeypatch, + ff_bias=False, + audio_ff_bias=True, + use_prompt_adaln_single=False, + use_keyframes_abs_pos_embedding=True, + ) + + assert ltx25_model.keyframes_abs_pos_embedding.shape == (1, 8) + torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, torch.zeros(1, 8)) + assert not hasattr(ltx25_model, "prompt_adaln") + assert not hasattr(ltx25_model, "audio_prompt_adaln") + assert ltx25_model.transformer_blocks[0].init_kwargs["ff_bias"] is False + assert ltx25_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True + + monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 1) + expected = torch.arange(8, dtype=torch.float32).view(1, 8) + loaded = ltx25_model.load_weights([("keyframes_abs_pos_embedding", expected)]) + + assert loaded == {"keyframes_abs_pos_embedding"} + torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, expected) + + +def test_ltx_cross_attention_modulations_use_official_timesteps(monkeypatch): + class CaptureAda(nn.Module): + def __init__(self): + super().__init__() + self.calls = [] + + def forward(self, timestep, *, batch_size, hidden_dtype): + self.calls.append(timestep.detach().clone()) + embedding = torch.zeros(timestep.numel(), 8, dtype=hidden_dtype) + return embedding, embedding + + class CaptureRope(nn.Module): + def __init__(self): + super().__init__() + self.output_dtypes = [] + + def forward(self, _coords, *, device, out_dtype=None): + self.output_dtypes.append(out_dtype) + empty = torch.empty(0, device=device, dtype=out_dtype) + return empty, empty + + model = _build_tiny_ltx_transformer(monkeypatch) + rope = CaptureRope() + model.rope = model.audio_rope = model.cross_attn_rope = model.cross_attn_audio_rope = rope + model.prompt_modulation = False + video_scale_shift = CaptureAda() + video_gate = CaptureAda() + audio_scale_shift = CaptureAda() + audio_gate = CaptureAda() + model.av_cross_attn_video_scale_shift = video_scale_shift + model.av_cross_attn_video_a2v_gate = video_gate + model.av_cross_attn_audio_scale_shift = audio_scale_shift + model.av_cross_attn_audio_v2a_gate = audio_gate + + model( + hidden_states=torch.zeros(1, 2, 4), + audio_hidden_states=torch.zeros(1, 1, 4), + encoder_hidden_states=torch.zeros(1, 1, 8), + audio_encoder_hidden_states=torch.zeros(1, 1, 8), + timestep=torch.tensor([[0.0, 2.0]]), + audio_timestep=torch.tensor([[5.0]]), + sigma=torch.tensor([2.0]), + audio_sigma=torch.tensor([5.0]), + video_coords=torch.zeros(1, 3, 2, 2), + audio_coords=torch.zeros(1, 1, 1, 2), + ) + + torch.testing.assert_close(video_scale_shift.calls[0], torch.tensor([0.0, 2.0])) + torch.testing.assert_close(video_gate.calls[0], torch.tensor([5.0])) + torch.testing.assert_close(audio_scale_shift.calls[0], torch.tensor([5.0])) + torch.testing.assert_close(audio_gate.calls[0], torch.tensor([2.0])) + assert len(rope.output_dtypes) == 4 + assert all(dtype is torch.float32 for dtype in rope.output_dtypes) + + +def test_ltx25_static_prompt_table_is_used_without_prompt_adaln(): + class CaptureAttention(nn.Module): + def __init__(self): + super().__init__() + self.encoder_hidden_states = None + + def forward(self, hidden_states, encoder_hidden_states=None, **kwargs): + if encoder_hidden_states is not None: + self.encoder_hidden_states = encoder_hidden_states.detach().clone() + return torch.zeros_like(hidden_states) + + class ZeroOutput(nn.Module): + def forward(self, hidden_states): + return torch.zeros_like(hidden_states) + + block = object.__new__(LTX2VideoTransformerBlock) + nn.Module.__init__(block) + block.video_cross_attn_adaln = True + block.audio_cross_attn_adaln = True + block.cross_attn_adaln = True + block.scale_shift_table = nn.Parameter(torch.zeros(9, 2)) + block.audio_scale_shift_table = nn.Parameter(torch.zeros(9, 2)) + block.prompt_scale_shift_table = nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]])) + block.audio_prompt_scale_shift_table = nn.Parameter(torch.tensor([[0.5, 1.0], [1.5, 2.0]])) + block.norm1 = nn.Identity() + block.audio_norm1 = nn.Identity() + block.norm2 = nn.Identity() + block.audio_norm2 = nn.Identity() + block.norm3 = nn.Identity() + block.audio_norm3 = nn.Identity() + block.attn1 = CaptureAttention() + block.audio_attn1 = CaptureAttention() + block.attn2 = CaptureAttention() + block.audio_attn2 = CaptureAttention() + block.ff = ZeroOutput() + block.audio_ff = ZeroOutput() + + hidden_states = torch.ones(1, 1, 2) + audio_hidden_states = torch.ones(1, 1, 2) + encoder_hidden_states = torch.ones(1, 1, 2) + audio_encoder_hidden_states = torch.ones(1, 1, 2) + block( + hidden_states, + audio_hidden_states, + encoder_hidden_states, + audio_encoder_hidden_states, + temb=torch.zeros(1, 1, 18), + temb_audio=torch.zeros(1, 1, 18), + temb_ca_scale_shift=torch.empty(0), + temb_ca_audio_scale_shift=torch.empty(0), + temb_ca_gate=torch.empty(0), + temb_ca_audio_gate=torch.empty(0), + temb_prompt=None, + temb_prompt_audio=None, + use_a2v_cross_attention=False, + use_v2a_cross_attention=False, + ) + + torch.testing.assert_close(block.attn2.encoder_hidden_states, torch.tensor([[[5.0, 7.0]]])) + torch.testing.assert_close(block.audio_attn2.encoder_hidden_states, torch.tensor([[[3.0, 4.0]]])) + + +def test_ltx_gated_attention_uses_tp_sharded_projection(monkeypatch): + column_calls = [] + + class FakeColumnParallelLinear(nn.Linear): + def __init__(self, input_size, output_size, *, bias=True, **kwargs): + column_calls.append((output_size, kwargs)) + super().__init__(input_size, output_size, bias=bias) + + class FakeRowParallelLinear(nn.Linear): + def __init__(self, input_size, output_size, *, bias=True, **_kwargs): + super().__init__(input_size, output_size, bias=bias) + + class FakeParallelAttention(nn.Module): + def __init__(self, **_kwargs): + super().__init__() + + monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumnParallelLinear) + monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRowParallelLinear) + monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeParallelAttention) + monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2) + + attention = ltx2_transformer.LTX2Attention( + query_dim=8, + heads=4, + kv_heads=4, + dim_head=2, + qk_norm="rms_norm", + apply_gated_attention=True, + pack_qkv=False, + ) + + gate_output_size, gate_kwargs = column_calls[-1] + assert gate_output_size == 4 + assert gate_kwargs["gather_output"] is False + assert gate_kwargs["return_bias"] is False + assert attention.query_num_heads == 2 diff --git a/tests/e2e/accuracy/ltx/run_ltx25_reference.py b/tests/e2e/accuracy/ltx/run_ltx25_reference.py new file mode 100644 index 000000000..cbc8917a8 --- /dev/null +++ b/tests/e2e/accuracy/ltx/run_ltx25_reference.py @@ -0,0 +1,591 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate raw LTX video and audio outputs with the official or Omni runtime.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--backend", choices=("official", "omni"), required=True) + parser.add_argument("--request", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--model") + parser.add_argument("--model-class-name") + parser.add_argument("--official-root", type=Path) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--gemma-root", type=Path) + parser.add_argument("--transformer-path", type=Path) + parser.add_argument("--text-encoder-path", type=Path) + parser.add_argument("--video-vae-path", type=Path) + parser.add_argument("--audio-vae-path", type=Path) + parser.add_argument("--spatial-upsampler-path", type=Path) + parser.add_argument("--distilled-lora-path", type=Path) + parser.add_argument("--connector-model", type=Path) + parser.add_argument( + "--official-pipeline", + choices=("distilled_two_stage", "full_one_stage", "full_two_stage"), + default="distilled_two_stage", + ) + parser.add_argument("--enable-layerwise-offload", action="store_true") + return parser.parse_args() + + +def _save_outputs( + output_dir: Path, + *, + video: torch.Tensor, + audio: torch.Tensor, + audio_sample_rate: int, + fps: float, + metadata: dict[str, Any], +) -> None: + output_dir.mkdir(parents=True, exist_ok=True) + video = video.detach().float().cpu().clamp(0.0, 1.0) + audio = audio.detach().float().cpu() + np.save(output_dir / "video.npy", video.numpy()) + np.save(output_dir / "audio.npy", audio.numpy()) + + from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes + + frames_u8 = video.mul(255.0).round().to(torch.uint8).numpy() + audio_array = audio.numpy() + while audio_array.ndim > 2 and audio_array.shape[0] == 1: + audio_array = audio_array[0] + (output_dir / "output.mp4").write_bytes( + mux_video_audio_bytes( + frames_u8, + audio_array, + fps=float(fps), + audio_sample_rate=int(audio_sample_rate), + video_codec_options={"preset": "ultrafast"}, + ) + ) + + frame_indices = sorted({0, video.shape[0] // 2, video.shape[0] - 1}) + for index in frame_indices: + frame = video[index].mul(255.0).round().to(torch.uint8).numpy() + Image.fromarray(frame).save(output_dir / f"frame_{index:04d}.png") + + metadata.update( + { + "video_shape": list(video.shape), + "audio_shape": list(audio.shape), + "audio_sample_rate": int(audio_sample_rate), + "frame_indices": frame_indices, + "mp4": "output.mp4", + } + ) + (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") + + +def _reset_accelerator_peak_memory_stats() -> None: + if torch.accelerator.is_available(): + torch.accelerator.reset_peak_memory_stats() + + +def _accelerator_peak_memory_stats() -> dict[str, float]: + if not torch.accelerator.is_available(): + return {} + torch.accelerator.synchronize() + mib = 1024**2 + return { + "peak_memory_allocated_mb": torch.accelerator.max_memory_allocated() / mib, + "peak_memory_reserved_mb": torch.accelerator.max_memory_reserved() / mib, + } + + +def _insert_official_paths(official_root: Path) -> None: + for relative_path in ("packages/ltx-core/src", "packages/ltx-pipelines/src"): + path = str((official_root / relative_path).resolve()) + if path not in sys.path: + sys.path.insert(0, path) + + +def _configure_official_sdpa(pipeline: Any) -> None: + """Pin official connector and denoiser attention to cuDNN, with MATH fallback.""" + from ltx_core.loader.attention_ops import set_attention_module_op + from ltx_core.model.transformer.attention import PytorchAttention + from torch.nn.attention import SDPBackend + + attention = PytorchAttention(priority=[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]) + module_op = set_attention_module_op( + attention=attention, + masked_attention=attention, + ) + transformer_owners = [ + owner + for attribute in ("stage", "stage_1", "stage_2") + if (owner := getattr(pipeline, attribute, None)) is not None + ] + owners_and_attributes = [ + *((owner, "_transformer_builder") for owner in transformer_owners), + (pipeline.prompt_encoder, "_embeddings_processor_builder"), + ] + for owner, attribute in owners_and_attributes: + builder = getattr(owner, attribute) + setattr( + owner, + attribute, + builder.with_module_ops((*builder.module_ops, module_op)), + ) + + +def _use_diffusers_connector_weights(prompt_encoder: Any, model: Path) -> None: + """Run official connector code with the checkpoint under test.""" + from safetensors import safe_open + + connector_root = model / "connectors" + shards = sorted(connector_root.glob("*.safetensors")) + if not shards: + raise ValueError(f"No connector safetensors found under {connector_root}") + + original_build_embeddings_processor = prompt_encoder._build_embeddings_processor + + def build_embeddings_processor(): + processor = original_build_embeddings_processor() + parameters = dict(processor.named_parameters()) + expected = {name for name in parameters if name.startswith(("video_connector.", "audio_connector."))} + loaded: set[str] = set() + with torch.no_grad(): + for shard in shards: + with safe_open(str(shard), framework="pt", device="cpu") as weights: + for source_name in weights.keys(): + if not source_name.startswith(("video_connector.", "audio_connector.")): + continue + target_name = ( + source_name.replace(".transformer_blocks.", ".transformer_1d_blocks.") + .replace(".attn1.norm_q.", ".attn1.q_norm.") + .replace(".attn1.norm_k.", ".attn1.k_norm.") + ) + parameter = parameters[target_name] + parameter.copy_(weights.get_tensor(source_name).to(parameter.device, parameter.dtype)) + loaded.add(target_name) + if loaded != expected: + missing = sorted(expected - loaded) + unexpected = sorted(loaded - expected) + raise ValueError(f"Connector weight mismatch: missing={missing}, unexpected={unexpected}") + return processor + + prompt_encoder._build_embeddings_processor = build_embeddings_processor + + +@torch.inference_mode() +def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: + if args.official_root is None: + raise ValueError("Official backend requires --official-root") + _insert_official_paths(args.official_root) + # Keep Gemma on the same SDPA path in both subprocesses. DiT attention is + # still pinned independently to the explicit cuDNN backend below. + torch.backends.cuda.enable_cudnn_sdp(False) + _reset_accelerator_peak_memory_stats() + + from ltx_pipelines.utils.args import ImageConditioningInput + from ltx_pipelines.utils.model_paths import ModelPaths + from ltx_pipelines.utils.types import OffloadMode + + offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE + image_path = request.get("image") + images = ( + [] + if image_path is None + else [ + ImageConditioningInput( + path=str(image_path), + frame_idx=0, + strength=1.0, + crf=request.get("image_crf", 18), + ) + ] + ) + + if args.transformer_path is not None: + required = { + "--text-encoder-path": args.text_encoder_path, + "--video-vae-path": args.video_vae_path, + "--audio-vae-path": args.audio_vae_path, + } + if args.official_pipeline in {"distilled_two_stage", "full_two_stage"}: + required["--spatial-upsampler-path"] = args.spatial_upsampler_path + if args.official_pipeline == "full_two_stage": + required["--distilled-lora-path"] = args.distilled_lora_path + missing = [name for name, value in required.items() if value is None] + if missing: + raise ValueError(f"Official LTX-2.5 split backend is missing: {', '.join(missing)}") + + model_paths = ModelPaths.from_split( + transformer_path=str(args.transformer_path), + text_encoder_path=str(args.text_encoder_path), + video_vae_path=str(args.video_vae_path), + audio_vae_path=str(args.audio_vae_path), + ) + if args.official_pipeline == "distilled_two_stage": + from ltx_pipelines.distilled import DistilledPipeline + + pipeline = DistilledPipeline( + model_paths=model_paths, + spatial_upsampler_path=str(args.spatial_upsampler_path), + loras=(), + offload_mode=offload_mode, + ) + if args.connector_model is not None: + _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model) + _configure_official_sdpa(pipeline) + video, audio, _, _ = pipeline( + prompt=request["prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + images=images, + ) + pipeline_name = "DistilledPipeline" + elif args.official_pipeline == "full_two_stage": + from ltx_core.components.guiders import MultiModalGuiderParams + from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps + from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline + + pipeline = TI2VidTwoStagesPipeline( + model_paths=model_paths, + distilled_lora=[ + LoraPathStrengthAndSDOps( + path=str(args.distilled_lora_path), + strength=1.0, + sd_ops=LTXV_LORA_COMFY_RENAMING_MAP, + ) + ], + spatial_upsampler_path=str(args.spatial_upsampler_path), + loras=[], + offload_mode=offload_mode, + ) + if args.connector_model is not None: + _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model) + _configure_official_sdpa(pipeline) + video, audio, _, _ = pipeline( + prompt=request["prompt"], + negative_prompt=request["negative_prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + num_inference_steps=request["num_inference_steps"], + video_guider_params=MultiModalGuiderParams( + cfg_scale=request["video_cfg_scale"], + stg_scale=request["video_stg_scale"], + rescale_scale=request["video_rescale_scale"], + modality_scale=request["video_modality_scale"], + skip_step=0, + stg_blocks=request["video_stg_blocks"], + ), + audio_guider_params=MultiModalGuiderParams( + cfg_scale=request["audio_cfg_scale"], + stg_scale=request["audio_stg_scale"], + rescale_scale=request["audio_rescale_scale"], + modality_scale=request["audio_modality_scale"], + skip_step=0, + stg_blocks=request["audio_stg_blocks"], + ), + images=images, + max_batch_size=4, + stage_1_sigmas=torch.tensor(request["stage_1_sigmas"], dtype=torch.float32), + stage_2_sigmas=torch.tensor(request["stage_2_sigmas"], dtype=torch.float32), + ) + pipeline_name = "TI2VidTwoStagesPipeline" + else: + from ltx_core.components.guiders import MultiModalGuiderParams + from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline + + pipeline = TI2VidOneStagePipeline( + model_paths=model_paths, + loras=(), + offload_mode=offload_mode, + ) + if args.connector_model is not None: + _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model) + _configure_official_sdpa(pipeline) + request_sigmas = request.get("sigmas") + video, audio, _ = pipeline( + prompt=request["prompt"], + negative_prompt=request["negative_prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + num_inference_steps=request["num_inference_steps"], + video_guider_params=MultiModalGuiderParams( + cfg_scale=request["video_cfg_scale"], + stg_scale=request["video_stg_scale"], + rescale_scale=request["video_rescale_scale"], + modality_scale=request["video_modality_scale"], + skip_step=0, + stg_blocks=request["video_stg_blocks"], + ), + audio_guider_params=MultiModalGuiderParams( + cfg_scale=request["audio_cfg_scale"], + stg_scale=request["audio_stg_scale"], + rescale_scale=request["audio_rescale_scale"], + modality_scale=request["audio_modality_scale"], + skip_step=0, + stg_blocks=request["audio_stg_blocks"], + ), + images=images, + max_batch_size=4, + sigmas=None if request_sigmas is None else torch.tensor(request_sigmas, dtype=torch.float32), + ) + pipeline_name = "TI2VidOneStagePipeline" + checkpoint_metadata = { + "transformer_path": str(args.transformer_path), + "video_vae_path": str(args.video_vae_path), + "connector_model": None if args.connector_model is None else str(args.connector_model), + "pipeline": pipeline_name, + } + else: + if args.checkpoint is None or args.gemma_root is None: + raise ValueError("Official monolith backend requires --checkpoint and --gemma-root") + + from ltx_core.components.guiders import MultiModalGuiderParams + from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline + + pipeline = TI2VidOneStagePipeline( + model_paths=ModelPaths.from_monolith(str(args.checkpoint), str(args.gemma_root)), + loras=(), + offload_mode=offload_mode, + ) + _configure_official_sdpa(pipeline) + video, audio, _ = pipeline( + prompt=request["prompt"], + negative_prompt=request["negative_prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + num_inference_steps=request["num_inference_steps"], + video_guider_params=MultiModalGuiderParams( + cfg_scale=request["video_cfg_scale"], + stg_scale=request["video_stg_scale"], + rescale_scale=request["video_rescale_scale"], + modality_scale=request["video_modality_scale"], + skip_step=0, + stg_blocks=request["video_stg_blocks"], + ), + audio_guider_params=MultiModalGuiderParams( + cfg_scale=request["audio_cfg_scale"], + stg_scale=request["audio_stg_scale"], + rescale_scale=request["audio_rescale_scale"], + modality_scale=request["audio_modality_scale"], + skip_step=0, + stg_blocks=request["audio_stg_blocks"], + ), + images=images, + max_batch_size=4, + ) + checkpoint_metadata = { + "checkpoint": str(args.checkpoint), + "pipeline": "TI2VidOneStagePipeline", + } + + video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0) + _save_outputs( + args.output_dir, + video=_canonical_video(video_tensor), + audio=audio.waveform, + audio_sample_rate=audio.sampling_rate, + fps=float(request["fps"]), + metadata={ + "backend": "official", + "attention_backend": "torch_sdpa_cudnn_bf16", + "official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"), + "seed": request["seed"], + "sigmas": request.get("sigmas"), + "stage_1_sigmas": request.get("stage_1_sigmas"), + "stage_2_sigmas": request.get("stage_2_sigmas"), + **_accelerator_peak_memory_stats(), + **checkpoint_metadata, + }, + ) + + +def _omni_worker_peak_memory_mb(output: Any) -> float: + result = output[0] if isinstance(output, list) and output else output + return float(getattr(result, "peak_memory_mb", 0.0) or 0.0) + + +def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]: + from vllm_omni.outputs import OmniRequestOutput + + audio = None + audio_sample_rate = None + frames = output[0] if isinstance(output, list) and output else output + if isinstance(frames, OmniRequestOutput): + multimodal_output = frames.multimodal_output or {} + audio = multimodal_output.get("audio") + audio_sample_rate = multimodal_output.get("audio_sample_rate") + nested_output = getattr(frames, "request_output", None) + if frames.is_pipeline_output and isinstance(nested_output, OmniRequestOutput): + frames = nested_output + multimodal_output = frames.multimodal_output or {} + audio = multimodal_output.get("audio", audio) + audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate) + if isinstance(frames, OmniRequestOutput): + if not frames.images: + raise ValueError("No video frames found in OmniRequestOutput") + frames = frames.images + + if isinstance(frames, list) and len(frames) == 1: + frames = frames[0] + if isinstance(frames, tuple) and len(frames) == 2: + frames, audio = frames + elif isinstance(frames, dict): + audio = frames.get("audio", audio) + audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate) + frames = frames.get("frames", frames.get("video")) + + if frames is None or audio is None or audio_sample_rate is None: + raise ValueError("Omni output did not contain video, audio, and audio_sample_rate") + return frames, audio, int(audio_sample_rate) + + +def _canonical_video(video: Any) -> torch.Tensor: + if isinstance(video, list): + if len(video) == 1: + return _canonical_video(video[0]) + tensors = [torch.as_tensor(np.asarray(item) if not isinstance(item, torch.Tensor) else item) for item in video] + if all(item.ndim == 3 for item in tensors): + video = torch.stack(tensors) + elif all(item.ndim == 4 for item in tensors): + video = torch.cat(tensors) + else: + raise ValueError(f"Cannot combine video items with dimensions {[item.ndim for item in tensors]}") + tensor = torch.as_tensor(np.asarray(video) if not isinstance(video, torch.Tensor) else video).detach().cpu() + if tensor.ndim == 5: + tensor = tensor[0] + if tensor.ndim != 4: + raise ValueError(f"Expected a 4D video tensor, got {tuple(tensor.shape)}") + if tensor.shape[-1] in (3, 4): + tensor = tensor[..., :3] + elif tensor.shape[1] in (3, 4): + tensor = tensor[:, :3].permute(0, 2, 3, 1) + elif tensor.shape[0] in (3, 4): + tensor = tensor[:3].permute(1, 2, 3, 0) + else: + raise ValueError(f"Cannot infer video channel dimension from {tuple(tensor.shape)}") + tensor = tensor.float() + if tensor.numel() and tensor.max() > 1: + tensor = tensor / 255.0 + elif tensor.numel() and tensor.min() < 0: + tensor = tensor.clamp(-1.0, 1.0).add(1.0).mul(0.5) + return tensor.clamp(0.0, 1.0) + + +@torch.inference_mode() +def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: + if args.model is None or args.model_class_name is None: + raise ValueError("Omni backend requires --model and --model-class-name") + + # vLLM imports can otherwise change the process-wide Gemma dispatch. + torch.backends.cuda.enable_cudnn_sdp(False) + _reset_accelerator_peak_memory_stats() + + attention_config = {"default": {"backend": "CUDNN_ATTN"}} + + from vllm_omni.diffusion.data import DiffusionParallelConfig + from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args + from vllm_omni.entrypoints.omni import Omni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + from vllm_omni.model_extras import get_extra_body_params + from vllm_omni.platforms import current_omni_platform + + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(request["seed"]) + omni = Omni( + model=args.model, + model_class_name=args.model_class_name, + enforce_eager=True, + enable_layerwise_offload=args.enable_layerwise_offload, + diffusion_attention_config=attention_config, + parallel_config=DiffusionParallelConfig(), + ) + try: + # The explicit public class argument selects the worker pipeline. The + # parent Omni proxy can only expose the checkpoint's base class, which + # is insufficient to distinguish LTX execution modes. + model_class_name = args.model_class_name + sampling_params = OmniDiffusionSamplingParams( + height=request["height"], + width=request["width"], + generator=generator, + guidance_scale=None, + num_inference_steps=request["num_inference_steps"], + num_frames=request["num_frames"], + fps=request["fps"], + frame_rate=float(request["fps"]), + output_type="np", + ) + extra_args = { + key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_") + } + for key in ("sigmas", "stage_1_sigmas", "stage_2_sigmas", "image_crf"): + if key in request: + extra_args[key] = request[key] + apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), extra_args) + prompt: dict[str, Any] = {"prompt": request["prompt"]} + if "negative_prompt" in request: + prompt["negative_prompt"] = request["negative_prompt"] + image_path = request.get("image") + if image_path is not None: + with Image.open(str(image_path)) as source_image: + image = source_image.convert("RGB") + image.load() + prompt["multi_modal_data"] = {"image": image} + output = omni.generate(prompt, sampling_params) + video, audio, audio_sample_rate = _unwrap_omni_output(output) + audio_tensor = torch.as_tensor(np.asarray(audio) if not isinstance(audio, torch.Tensor) else audio) + _save_outputs( + args.output_dir, + video=_canonical_video(video), + audio=audio_tensor, + audio_sample_rate=audio_sample_rate, + fps=float(request["fps"]), + metadata={ + "backend": "omni", + "attention_backend": "torch_sdpa_cudnn_bf16", + "seed": request["seed"], + "sigmas": request.get("sigmas"), + "model": args.model, + "stage_1_sigmas": request.get("stage_1_sigmas"), + "stage_2_sigmas": request.get("stage_2_sigmas"), + "model_class_name": model_class_name, + "worker_peak_memory_mb": _omni_worker_peak_memory_mb(output), + **_accelerator_peak_memory_stats(), + }, + ) + finally: + omni.shutdown() + + +def main() -> None: + args = _parse_args() + request = json.loads(args.request.read_text()) + if args.backend == "official": + _run_official(args, request) + else: + _run_omni(args, request) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py new file mode 100644 index 000000000..9f1623c8a --- /dev/null +++ b/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py @@ -0,0 +1,753 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""E2E accuracy guards against pinned official Lightricks LTX references. + +The LTX-2/2.3 comparison runs both runtimes through PyTorch SDPA and uses +``max_batch_size=4`` in the official reference to match Omni's fused guidance +batch. Video and audio guidance use the official non-HQ one-stage defaults; +only the generation shape and step count are reduced for CI runtime. +LTX-2.5 runs the official split-artifact pipelines with connector weights from +the same Diffusers checkpoint under test. The distilled case covers the fixed +two-stage schedule; the Full/SFT case covers raw dev weights and an explicit +shared one-stage schedule. +""" + +from __future__ import annotations + +import json +import math +import os +import shutil +import subprocess +import sys +import time +from pathlib import Path + +import numpy as np +import pytest +import torch +from huggingface_hub import hf_hub_download, snapshot_download +from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure + +from tests.e2e.accuracy.helpers import reset_artifact_dir +from tests.helpers import skip_if_gated_repo_inaccessible +from tests.helpers.mark import hardware_test + +OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git" +LTX25_OFFICIAL_REVISION = "7954dcb0d986bdc36ef272564a9789ade07fcc65" +LTX25_OMNI_MODEL_ID = "Lightricks/LTX-2.5-Diffusers" +LTX25_OMNI_MODEL_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d" +LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5" +LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1" +LTX25_OFFICIAL_COMMON_FILES = { + "text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", + "video_vae": "vae/ltx-2.5-video-vae-conv-bf16.safetensors", + "audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors", +} +LTX25_OFFICIAL_VARIANT_FILES = { + "distilled": { + "transformer": "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", + "spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", + }, + "full": { + "transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", + }, + "full_two_stage": { + "transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", + "spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", + "distilled_lora": "loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors", + }, +} +# Version selected by the pinned official source's uv.lock. Keep it isolated +# from Omni's runtime and development dependencies. +OFFICIAL_OPENIMAGEIO_VERSION = "3.1.11.0" +LTX25_PROMPT = "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside." +LTX25_I2V_IMAGE_REPO = "huggingface/documentation-images" +LTX25_I2V_IMAGE_FILENAME = "diffusers/svd/rocket.png" +LTX25_I2V_IMAGE_REVISION = "645d8364f0c7a101180b364811b5a11a362e4010" +NEGATIVE_PROMPT = ( + "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, " + "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, " + "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, " + "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of field, " + "background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent lighting " + "direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny valley effect, " + "incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, mismatched lip sync, " + "silent or muted audio, distorted voice, robotic voice, echo, background noise, off-sync audio, incorrect " + "dialogue, added dialogue, repetitive speech, jittery movement, awkward pauses, incorrect timing, unnatural " + "transitions, inconsistent framing, tilted camera, flat lighting, inconsistent tone, cinematic oversaturation, " + "stylized filters, or AI artifacts." +) + +# Release parity exercises the accelerated BF16 cuDNN path in both runtimes. +# Degenerate one-token attention falls back to PyTorch SDPA MATH. +ATTENTION_BACKEND = "torch_sdpa_cudnn_bf16" +VIDEO_SSIM_MEAN_THRESHOLD = 0.95 +VIDEO_SSIM_MIN_THRESHOLD = 0.90 +VIDEO_PSNR_MEAN_THRESHOLD = 30.0 +AUDIO_RELATIVE_L2_THRESHOLD = 0.2 +AUDIO_COSINE_THRESHOLD = 0.95 + +# Decoded-output gates for the pinned official/Omni BF16 cuDNN comparison. +# The metrics artifact preserves the exact observed values. +LTX25_VIDEO_SSIM_MEAN_THRESHOLD = 0.99 +LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD = 0.99 +LTX25_VIDEO_SSIM_MIN_THRESHOLD = 0.99 +LTX25_VIDEO_PSNR_MEAN_THRESHOLD = 40.0 +LTX25_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01 +LTX25_AUDIO_RELATIVE_L2_THRESHOLD = 0.32 +LTX25_AUDIO_COSINE_THRESHOLD = 0.95 + +# Full/SFT uses independently converted dev weights, but the release contract +# still requires decoded-video parity at the same 0.99 SSIM floor as the +# distilled paths when both runtimes use BF16 cuDNN attention. +LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD = 0.99 +LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD = 0.99 +LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD = 40.0 +LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD = 0.3 +LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01 +LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD = 0.02 +LTX25_FULL_AUDIO_COSINE_THRESHOLD = AUDIO_COSINE_THRESHOLD +LTX25_STAGE_2_SIGMAS = [0.909375, 0.725, 0.421875, 0.0] + + +def test_ltx_reference_runner_unwraps_flattened_pipeline_output() -> None: + from vllm_omni.outputs import OmniRequestOutput + + from .run_ltx25_reference import _omni_worker_peak_memory_mb, _unwrap_omni_output + + frame = object() + audio = object() + output = OmniRequestOutput( + stage_id=0, + images=[frame], + _multimodal_output={"audio": audio, "audio_sample_rate": 48_000}, + ) + + frames, actual_audio, sample_rate = _unwrap_omni_output(output) + + assert frames is frame + assert actual_audio is audio + assert sample_rate == 48_000 + output.peak_memory_mb = 12_345.0 + assert _omni_worker_peak_memory_mb([output]) == 12_345.0 + + +def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None: + start = time.perf_counter() + subprocess.run(command, env=env, timeout=timeout, check=True) + print(f"{' '.join(command[:3])} finished in {time.perf_counter() - start:.1f}s") + + +def _clone_pinned_source(root: Path, repository: str, revision: str) -> None: + root.mkdir(parents=True) + subprocess.run(["git", "init", "-q", str(root)], check=True) + subprocess.run(["git", "-C", str(root), "remote", "add", "origin", repository], check=True) + last_error: subprocess.CalledProcessError | None = None + for attempt in range(3): + try: + subprocess.run( + ["git", "-C", str(root), "fetch", "--depth", "1", "origin", revision], + check=True, + ) + last_error = None + break + except subprocess.CalledProcessError as error: + last_error = error + if attempt < 2: + time.sleep(5 * (attempt + 1)) + if last_error is not None: + raise last_error + subprocess.run(["git", "-C", str(root), "checkout", "-q", "--detach", "FETCH_HEAD"], check=True) + + +def _git_revision(root: Path) -> str | None: + try: + return subprocess.check_output( + ["git", "-C", str(root), "rev-parse", "HEAD"], + stderr=subprocess.DEVNULL, + text=True, + ).strip() + except subprocess.CalledProcessError: + return None + + +def _ltx25_official_source(artifact_root: Path) -> tuple[Path, str]: + repository = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REPOSITORY", OFFICIAL_REPOSITORY) + revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REVISION", LTX25_OFFICIAL_REVISION) + configured_root = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_ROOT") + root = Path(configured_root) if configured_root else artifact_root / f"official-source-{revision[:12]}" + actual_revision = _git_revision(root) if root.exists() else None + if actual_revision != revision and configured_root: + raise AssertionError(f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}") + if actual_revision != revision: + if root.exists(): + shutil.rmtree(root) + _clone_pinned_source(root, repository, revision) + actual_revision = _git_revision(root) + assert actual_revision == revision, f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}" + return root, revision + + +def _official_runner_prefix() -> list[str]: + """Run the reference with its missing binary dependency isolated from CI.""" + uv = shutil.which("uv") + assert uv is not None, "uv is required to run the pinned official LTX reference" + return [ + uv, + "run", + "--no-project", + "--with", + f"openimageio=={OFFICIAL_OPENIMAGEIO_VERSION}", + "--python", + sys.executable, + "python", + ] + + +def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") -> tuple[Path, str, str | None]: + if transformer_subfolder not in {"transformer", "transformer_full"}: + raise ValueError(f"Unsupported LTX-2.5 transformer subfolder: {transformer_subfolder!r}") + configured_model = os.environ.get("VLLM_TEST_LTX25_MODEL") + configured_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION") + if configured_model and Path(configured_model).exists(): + model = Path(configured_model).resolve() + snapshot_revision = model.name if model.parent.name == "snapshots" else None + assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}" + return model, configured_model, configured_revision or snapshot_revision + model_id = configured_model or LTX25_OMNI_MODEL_ID + revision = configured_revision + if revision is None and model_id == LTX25_OMNI_MODEL_ID: + revision = LTX25_OMNI_MODEL_REVISION + model = Path( + snapshot_download( + repo_id=model_id, + revision=revision, + allow_patterns=[ + "model_index.json", + "audio_vae/*", + "connectors/config.json", + "connectors/diffusion_pytorch_model.safetensors.index.json", + "connectors/diffusion_pytorch_model-*.safetensors", + "scheduler/*", + "text_encoder/config.json", + "text_encoder/generation_config.json", + "text_encoder/model*", + "tokenizer/*", + f"{transformer_subfolder}/*", + "vae/*", + "vocoder/*", + ], + ) + ).resolve() + snapshot_revision = model.name if model.parent.name == "snapshots" else None + return model, model_id, snapshot_revision or revision + + +def _resolve_ltx25_official_artifacts( + *, checkpoint_variant: str = "distilled" +) -> tuple[dict[str, Path], str, str | None]: + try: + variant_files = LTX25_OFFICIAL_VARIANT_FILES[checkpoint_variant] + except KeyError as exc: + raise ValueError(f"Unsupported official LTX-2.5 checkpoint variant: {checkpoint_variant!r}") from exc + files = {**LTX25_OFFICIAL_COMMON_FILES, **variant_files} + configured_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL") + configured_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION") + if configured_model and Path(configured_model).exists(): + root = Path(configured_model).resolve() + model_source = configured_model + snapshot_revision = root.name if root.parent.name == "snapshots" else None + resolved_revision = configured_revision or snapshot_revision + else: + model_source = configured_model or LTX25_OFFICIAL_MODEL_ID + revision = configured_revision + if revision is None and model_source == LTX25_OFFICIAL_MODEL_ID: + revision = LTX25_OFFICIAL_MODEL_REVISION + root = Path( + snapshot_download( + repo_id=model_source, + revision=revision, + allow_patterns=list(files.values()), + ) + ).resolve() + snapshot_revision = root.name if root.parent.name == "snapshots" else None + resolved_revision = snapshot_revision or revision + + artifacts = {name: root / relative_path for name, relative_path in files.items()} + missing = [f"{name}: {path}" for name, path in artifacts.items() if not path.is_file()] + assert not missing, f"Official LTX-2.5 artifacts are missing: {', '.join(missing)}" + return artifacts, model_source, resolved_revision + + +def _resolve_gemma_root(model: Path) -> Path: + configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT") + if configured_root: + root = Path(configured_root) + assert root.is_dir(), f"Gemma root not found: {root}" + return root + return model + + +def _resolve_ltx25_image() -> Path: + configured_image = os.environ.get("VLLM_TEST_LTX25_I2V_IMAGE") + if configured_image: + image = Path(configured_image) + assert image.is_file(), f"LTX-2.5 I2V conditioning image not found: {image}" + return image + return Path( + hf_hub_download( + repo_id=LTX25_I2V_IMAGE_REPO, + repo_type="dataset", + filename=LTX25_I2V_IMAGE_FILENAME, + revision=LTX25_I2V_IMAGE_REVISION, + ) + ) + + +def _require_ltx25_model_access() -> None: + configured_omni_model = os.environ.get("VLLM_TEST_LTX25_MODEL") + if not (configured_omni_model and Path(configured_omni_model).exists()): + omni_model = configured_omni_model or LTX25_OMNI_MODEL_ID + omni_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION") + if omni_revision is None and omni_model == LTX25_OMNI_MODEL_ID: + omni_revision = LTX25_OMNI_MODEL_REVISION + skip_if_gated_repo_inaccessible( + omni_model, + revision=omni_revision, + filename="model_index.json", + ) + + configured_official_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL") + if not (configured_official_model and Path(configured_official_model).exists()): + official_model = configured_official_model or LTX25_OFFICIAL_MODEL_ID + official_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION") + if official_revision is None and official_model == LTX25_OFFICIAL_MODEL_ID: + official_revision = LTX25_OFFICIAL_MODEL_REVISION + skip_if_gated_repo_inaccessible( + official_model, + revision=official_revision, + filename=LTX25_OFFICIAL_COMMON_FILES["text_encoder"], + ) + + +def _ltx25_request(image: Path | None = None) -> dict[str, object]: + request: dict[str, object] = { + "prompt": LTX25_PROMPT, + # Keep the nightly parity test small enough to run the pinned official + # pipeline and Omni sequentially on one H100. The public recipe covers + # the release-resolution 1920x1088x121 configuration. + "width": 1024, + "height": 576, + "num_frames": 25, + "fps": 24, + "num_inference_steps": 8, + "seed": 42, + } + if image is not None: + request.update(image=str(image.resolve()), image_crf=18) + return request + + +def _ltx25_full_sigmas(num_inference_steps: int) -> list[float]: + """Materialize the official LTX-2.5 one-stage schedule in FP32.""" + sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1) + base_anchor = 1024 + max_anchor = 4096 + base_shift = 0.95 + max_shift = 2.05 + slope = (max_shift - base_shift) / (max_anchor - base_anchor) + sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor) + exp_shift = math.exp(sigma_shift) + sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0) + + non_zero = sigmas != 0 + one_minus_sigmas = 1.0 - sigmas[non_zero] + scale = one_minus_sigmas[-1] / (1.0 - 0.1) + sigmas[non_zero] = 1.0 - one_minus_sigmas / scale + return [float(sigma) for sigma in sigmas] + + +def _ltx25_full_request(image: Path | None = None) -> dict[str, object]: + num_inference_steps = 30 + request: dict[str, object] = { + "prompt": LTX25_PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + # Full guidance expands to a four-way batch. Keep the decoded output + # small enough for the pinned official and Omni runs to share one H100. + "width": 512, + "height": 384, + "num_frames": 25, + "fps": 24, + "num_inference_steps": num_inference_steps, + "seed": 42, + "video_cfg_scale": 3.0, + "audio_cfg_scale": 7.0, + "video_stg_scale": 1.0, + "audio_stg_scale": 1.0, + "video_modality_scale": 3.0, + "audio_modality_scale": 3.0, + "video_rescale_scale": 0.7, + "audio_rescale_scale": 0.7, + "video_stg_blocks": [28], + "audio_stg_blocks": [28], + "sigmas": _ltx25_full_sigmas(num_inference_steps), + } + if image is not None: + request.update(image=str(image.resolve()), image_crf=18) + return request + + +def test_ltx25_full_request_pins_official_schedule() -> None: + request = _ltx25_full_request() + sigmas = request["sigmas"] + + assert isinstance(sigmas, list) + assert isinstance(request["num_inference_steps"], int) + assert len(sigmas) == request["num_inference_steps"] + 1 + assert math.isclose(sigmas[0], 1.0, rel_tol=0.0, abs_tol=1e-6) + assert sigmas[-1] == 0.0 + assert all(sigmas[index] > sigmas[index + 1] for index in range(len(sigmas) - 1)) + + +def test_ltx25_i2v_requests_pin_official_crf() -> None: + image = Path("conditioning.png") + for request in (_ltx25_request(image), _ltx25_full_request(image)): + assert request["image"] == str(image.resolve()) + assert request["image_crf"] == 18 + + +def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]: + from benchmarks.diffusion.quantization_quality import compute_lpips_video + + assert reference.shape == prediction.shape + assert reference.ndim == 4 and reference.shape[-1] == 3 + ssim_scores: list[float] = [] + psnr_scores: list[float] = [] + for reference_frame, prediction_frame in zip(reference, prediction, strict=True): + reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0) + prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0) + ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor))) + psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor))) + difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64)) + return { + "ssim_mean": float(np.mean(ssim_scores)), + "ssim_min": float(np.min(ssim_scores)), + "psnr_mean_db": float(np.mean(psnr_scores)), + "lpips_alex_mean": compute_lpips_video(reference, prediction), + "max_abs": float(difference.max()), + "mean_abs": float(difference.mean()), + } + + +def _canonical_audio(audio: np.ndarray) -> np.ndarray: + while audio.ndim > 2 and audio.shape[0] == 1: + audio = audio[0] + return audio.astype(np.float64) + + +def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float | bool]: + reference = _canonical_audio(reference) + prediction = _canonical_audio(prediction) + assert reference.shape == prediction.shape + difference = reference - prediction + reference_norm = max(float(np.linalg.norm(reference)), 1e-12) + prediction_norm = max(float(np.linalg.norm(prediction)), 1e-12) + return { + "bitwise_equal": bool(np.array_equal(reference, prediction)), + "max_abs": float(np.abs(difference).max()), + "mean_abs": float(np.abs(difference).mean()), + "relative_l2": float(np.linalg.norm(difference) / reference_norm), + "cosine_similarity": float(np.vdot(reference.ravel(), prediction.ravel()) / (reference_norm * prediction_norm)), + } + + +@pytest.mark.full_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("task", ["t2v", "i2v"]) +def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None: + """Compare Omni with the official runtime using the same checkpoint connector weights.""" + _require_ltx25_model_access() + artifact_parent = accuracy_artifact_root / "ltx_official" + output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}") + official_root, official_revision = _ltx25_official_source(artifact_parent) + official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts() + omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model() + image = _resolve_ltx25_image() if task == "i2v" else None + request_path = output_root / "request.json" + request_path.write_text(json.dumps(_ltx25_request(image), indent=2) + "\n") + + runner = Path(__file__).with_name("run_ltx25_reference.py") + runner_args = [ + str(runner), + "--request", + str(request_path), + "--enable-layerwise-offload", + ] + env = os.environ.copy() + env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision + env["PYTHONUNBUFFERED"] = "1" + repository_root = Path(__file__).resolve().parents[4] + existing_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}" + ) + + official_output = output_root / "official" + _run( + _official_runner_prefix() + + runner_args + + [ + "--backend", + "official", + "--output-dir", + str(official_output), + "--official-root", + str(official_root), + "--transformer-path", + str(official_artifacts["transformer"]), + "--text-encoder-path", + str(official_artifacts["text_encoder"]), + "--video-vae-path", + str(official_artifacts["video_vae"]), + "--audio-vae-path", + str(official_artifacts["audio_vae"]), + "--spatial-upsampler-path", + str(official_artifacts["spatial_upsampler"]), + "--connector-model", + str(omni_model), + ], + env=env, + ) + + omni_output = output_root / "omni" + _run( + [sys.executable] + + runner_args + + [ + "--backend", + "omni", + "--output-dir", + str(omni_output), + "--model", + str(omni_model), + "--model-class-name", + "LTX2DistilledTwoStagePipeline", + ], + env=env, + ) + + official_metadata = json.loads((official_output / "metadata.json").read_text()) + omni_metadata = json.loads((omni_output / "metadata.json").read_text()) + assert official_metadata["official_revision"] == official_revision + assert official_metadata["pipeline"] == "DistilledPipeline" + assert official_metadata["attention_backend"] == ATTENTION_BACKEND + assert official_metadata["connector_model"] == str(omni_model) + assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline" + assert omni_metadata["attention_backend"] == ATTENTION_BACKEND + assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"] + video_metrics = _video_metrics( + np.load(official_output / "video.npy"), + np.load(omni_output / "video.npy"), + ) + audio_metrics = _audio_metrics( + np.load(official_output / "audio.npy"), + np.load(omni_output / "audio.npy"), + ) + result = { + "case": "ltx2_5_distilled", + "task": task, + "attention_backend": ATTENTION_BACKEND, + "official_revision": official_revision, + "official_model": official_model_source, + "official_model_revision": official_model_revision, + "official_connector_model": omni_model_source, + "official_connector_revision": omni_model_revision, + "omni_model": omni_model_source, + "omni_model_revision": omni_model_revision, + "resolved_omni_model_path": str(omni_model), + "peak_memory_mb": { + "official_allocated": official_metadata["peak_memory_allocated_mb"], + "official_reserved": official_metadata["peak_memory_reserved_mb"], + "omni_worker_reported": omni_metadata["worker_peak_memory_mb"], + "omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"], + "omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"], + }, + "video": video_metrics, + "audio": audio_metrics, + } + (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + ssim_mean_threshold = LTX25_VIDEO_SSIM_MEAN_THRESHOLD if task == "t2v" else LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD + assert video_metrics["ssim_mean"] >= ssim_mean_threshold + assert video_metrics["ssim_min"] >= LTX25_VIDEO_SSIM_MIN_THRESHOLD + assert video_metrics["psnr_mean_db"] >= LTX25_VIDEO_PSNR_MEAN_THRESHOLD + assert video_metrics["lpips_alex_mean"] <= LTX25_VIDEO_LPIPS_MEAN_THRESHOLD + assert audio_metrics["cosine_similarity"] >= LTX25_AUDIO_COSINE_THRESHOLD + assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD + + +@pytest.mark.full_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("task", ["t2v", "i2v"]) +@pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"]) +def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None: + """Compare raw official Full/SFT weights with Omni's converted Full pipeline.""" + _require_ltx25_model_access() + artifact_parent = accuracy_artifact_root / "ltx_official" + output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}") + model_class_name = { + "full_one_stage": "LTX2Pipeline", + "full_two_stage": "LTX2TwoStagePipeline", + }[pipeline_mode] + official_root, official_revision = _ltx25_official_source(artifact_parent) + official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts( + checkpoint_variant="full" if pipeline_mode == "full_one_stage" else "full_two_stage" + ) + omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model( + transformer_subfolder="transformer_full" + ) + image = _resolve_ltx25_image() if task == "i2v" else None + request = _ltx25_full_request(image) + if pipeline_mode == "full_two_stage": + request["stage_1_sigmas"] = request.pop("sigmas") + request["stage_2_sigmas"] = LTX25_STAGE_2_SIGMAS + request_path = output_root / "request.json" + request_path.write_text(json.dumps(request, indent=2) + "\n") + + runner = Path(__file__).with_name("run_ltx25_reference.py") + runner_args = [ + str(runner), + "--request", + str(request_path), + "--enable-layerwise-offload", + ] + env = os.environ.copy() + env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision + env["PYTHONUNBUFFERED"] = "1" + repository_root = Path(__file__).resolve().parents[4] + existing_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}" + ) + + official_output = output_root / "official" + official_sidecar_args: list[str] = [] + if pipeline_mode == "full_two_stage": + official_sidecar_args = [ + "--spatial-upsampler-path", + str(official_artifacts["spatial_upsampler"]), + "--distilled-lora-path", + str(official_artifacts["distilled_lora"]), + ] + + _run( + _official_runner_prefix() + + runner_args + + [ + "--backend", + "official", + "--output-dir", + str(official_output), + "--official-root", + str(official_root), + "--official-pipeline", + pipeline_mode, + "--transformer-path", + str(official_artifacts["transformer"]), + "--text-encoder-path", + str(official_artifacts["text_encoder"]), + "--video-vae-path", + str(official_artifacts["video_vae"]), + "--audio-vae-path", + str(official_artifacts["audio_vae"]), + "--connector-model", + str(omni_model), + ] + + official_sidecar_args, + env=env, + ) + + omni_output = output_root / "omni" + _run( + [sys.executable] + + runner_args + + [ + "--backend", + "omni", + "--output-dir", + str(omni_output), + "--model", + str(omni_model), + "--model-class-name", + model_class_name, + ], + env=env, + ) + + official_metadata = json.loads((official_output / "metadata.json").read_text()) + omni_metadata = json.loads((omni_output / "metadata.json").read_text()) + assert official_metadata["official_revision"] == official_revision + assert ( + official_metadata["pipeline"] + == { + "full_one_stage": "TI2VidOneStagePipeline", + "full_two_stage": "TI2VidTwoStagesPipeline", + }[pipeline_mode] + ) + assert official_metadata["attention_backend"] == ATTENTION_BACKEND + assert official_metadata["connector_model"] == str(omni_model) + assert omni_metadata["model_class_name"] == model_class_name + assert omni_metadata["attention_backend"] == ATTENTION_BACKEND + assert official_metadata["seed"] == omni_metadata["seed"] + for schedule_name in ("sigmas", "stage_1_sigmas", "stage_2_sigmas"): + assert official_metadata[schedule_name] == omni_metadata[schedule_name] + assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"] + video_metrics = _video_metrics( + np.load(official_output / "video.npy"), + np.load(omni_output / "video.npy"), + ) + audio_metrics = _audio_metrics( + np.load(official_output / "audio.npy"), + np.load(omni_output / "audio.npy"), + ) + result = { + "case": f"ltx2_5_{pipeline_mode}", + "task": task, + "attention_backend": ATTENTION_BACKEND, + "official_revision": official_revision, + "official_model": official_model_source, + "official_model_revision": official_model_revision, + "official_connector_model": omni_model_source, + "official_connector_revision": omni_model_revision, + "omni_model": omni_model_source, + "omni_model_revision": omni_model_revision, + "resolved_omni_model_path": str(omni_model), + "peak_memory_mb": { + "official_allocated": official_metadata["peak_memory_allocated_mb"], + "official_reserved": official_metadata["peak_memory_reserved_mb"], + "omni_worker_reported": omni_metadata["worker_peak_memory_mb"], + "omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"], + "omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"], + }, + "video": video_metrics, + "audio": audio_metrics, + } + (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + assert video_metrics["ssim_mean"] >= LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD + assert video_metrics["ssim_min"] >= LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD + assert video_metrics["psnr_mean_db"] >= LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD + lpips_threshold = ( + LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD if task == "t2v" else LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD + ) + assert audio_metrics["relative_l2"] <= LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD + assert video_metrics["lpips_alex_mean"] <= lpips_threshold + assert audio_metrics["cosine_similarity"] >= LTX25_FULL_AUDIO_COSINE_THRESHOLD diff --git a/tests/e2e/online_serving/test_ltx25.py b/tests/e2e/online_serving/test_ltx25.py new file mode 100644 index 000000000..2fad9f16a --- /dev/null +++ b/tests/e2e/online_serving/test_ltx25.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Online serving smokes for all four canonical LTX-2.5 pipelines.""" + +import os + +import pytest + +from tests.helpers import skip_if_gated_repo_inaccessible +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +DEFAULT_MODEL = "Lightricks/LTX-2.5-Diffusers" +DEFAULT_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d" +MODEL = os.getenv("VLLM_TEST_LTX25_MODEL", DEFAULT_MODEL) +MODEL_REVISION = os.getenv("VLLM_TEST_LTX25_MODEL_REVISION", DEFAULT_REVISION if MODEL == DEFAULT_MODEL else "") +PROMPT = "A red fox walks through a snowy forest while the camera remains fixed." + +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) + + +@pytest.fixture(scope="module", autouse=True) +def require_ltx25_model_access() -> None: + if not os.path.isdir(MODEL): + skip_if_gated_repo_inaccessible( + MODEL, + revision=MODEL_REVISION or None, + filename="model_index.json", + ) + + +def _server(model_class_name: str) -> OmniServerParams: + return OmniServerParams( + model=MODEL, + server_args=[ + *(["--revision", MODEL_REVISION] if MODEL_REVISION else []), + "--model-class-name", + model_class_name, + "--enforce-eager", + "--enable-layerwise-offload", + "--diffusion-attention-backend", + "CUDNN_ATTN", + ], + ) + + +def _cases(): + return [ + pytest.param(_server("LTX2Pipeline"), 30, id="full_one_stage", marks=SINGLE_CARD_MARKS), + pytest.param(_server("LTX2TwoStagePipeline"), 30, id="full_two_stage", marks=SINGLE_CARD_MARKS), + pytest.param(_server("LTX2DistilledOneStagePipeline"), 8, id="distilled_one_stage", marks=SINGLE_CARD_MARKS), + pytest.param(_server("LTX2DistilledTwoStagePipeline"), 8, id="distilled_two_stage", marks=SINGLE_CARD_MARKS), + ] + + +@pytest.mark.parametrize(("omni_server", "num_inference_steps"), _cases(), indirect=["omni_server"]) +def test_ltx25_pipeline_entries( + omni_server: OmniServer, + num_inference_steps: int, + openai_client: OpenAIClientHandler, + subtests, +) -> None: + """Generate T2V and first-frame I2V through the synchronous video API.""" + for task in ("t2v", "i2v"): + with subtests.test(task=task): + request_config = { + "model": omni_server.model, + "form_data": { + "model": omni_server.model, + "prompt": PROMPT, + "height": 128, + "width": 128, + "num_frames": 9, + "fps": 24, + "num_inference_steps": num_inference_steps, + "seed": 42, + }, + } + if task == "i2v": + request_config["image_reference"] = ( + f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + ) + + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index bcda12e2d..b7527e233 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -32,7 +32,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( VideoParams, VideoResponse, ) -from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo +from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage from vllm_omni.entrypoints.openai.storage import LocalStorageManager from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry from vllm_omni.errors import GuardrailViolationError @@ -461,6 +461,51 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client, assert input_image.size == (96, 64) +def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch): + engine = FakeAsyncOmni() + engine.get_diffusion_od_config = lambda: SimpleNamespace( + model="org/model", + model_class_name="ExamplePipeline", + revision="pinned-revision", + ) + captured = {} + + def fake_policy(model_class_name, *, model, revision=None): + captured.update( + model_class_name=model_class_name, + model=model, + revision=revision, + ) + return True + + monkeypatch.setattr( + "vllm_omni.entrypoints.openai.serving_video.should_preserve_reference_image_size", + fake_policy, + ) + handler = OmniOpenAIServingVideo.for_diffusion( + diffusion_engine=engine, + model_name="fallback/model", + ) + image = Image.new("RGB", (48, 32)) + + asyncio.run( + handler._run_and_extract( + VideoGenerationRequest(prompt="A bear playing with yarn.", width=96, height=64), + "pipeline-owned-resize", + reference_image=ReferenceImage(image), + ) + ) + + input_image = engine.captured_prompt["multi_modal_data"]["image"] + assert isinstance(input_image, Image.Image) + assert input_image.size == (48, 32) + assert captured == { + "model_class_name": "ExamplePipeline", + "model": "org/model", + "revision": "pinned-revision", + } + + def test_i2v_extra_params_dimensions_preserve_input_image_geometry(test_client, mocker: MockerFixture): image_bytes = _make_test_image_bytes((48, 48)) mocker.patch( diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py index cc459c269..660203710 100644 --- a/tests/entrypoints/test_resolve_dreamzero_config.py +++ b/tests/entrypoints/test_resolve_dreamzero_config.py @@ -49,7 +49,9 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch): def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch): monkeypatch.setattr( "vllm.transformers_utils.config.get_hf_file_to_dict", - lambda path, _model: None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]}, + lambda path, _model, **_kwargs: ( + None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]} + ), ) od_config = OmniDiffusionConfig( diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py index 9929beab8..29cd9567e 100644 --- a/tests/helpers/__init__.py +++ b/tests/helpers/__init__.py @@ -11,11 +11,16 @@ affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based import pytest -def skip_if_gated_repo_inaccessible(repo_id: str) -> None: +def skip_if_gated_repo_inaccessible( + repo_id: str, + *, + revision: str | None = None, + filename: str = "config.json", +) -> None: """Skip the test if a gated HuggingFace repo is not accessible. - Tries to download the model's config.json via ``hf_hub_download``, - which performs an actual file-access check (unlike ``HfApi().model_info()`` + Tries to download ``filename`` via ``hf_hub_download``, which performs an + actual file-access check (unlike ``HfApi().model_info()`` that only checks metadata). If the token has metadata access but not file-download access, ``hf_hub_download`` will raise ``GatedRepoError`` and we skip cleanly. @@ -26,7 +31,7 @@ def skip_if_gated_repo_inaccessible(repo_id: str) -> None: except Exception: return try: - hf_hub_download(repo_id=repo_id, filename="config.json") + hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) except GatedRepoError as exc: pytest.skip( f"Skipping: gated HF repo {repo_id!r} inaccessible to the current " diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 557fb5f6c..573f0bd79 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -20,6 +20,7 @@ from vllm_omni.model_extras import ( get_output_tensor_range, get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, + should_preserve_reference_image_size, ) @@ -230,18 +231,62 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None: "audio_rescale_scale", "video_stg_blocks", "audio_stg_blocks", + "sigmas", + "stage_1_sigmas", + "stage_2_sigmas", + "image_crf", } ) for pipeline_name in ( "LTX2Pipeline", "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", + "LTX2DistilledTwoStagePipeline", "LTX2DistilledPipeline", ): assert get_extra_body_params(pipeline_name) == expected assert get_extra_output_params(pipeline_name) == frozenset() +@pytest.mark.parametrize( + ("model_version", "expected"), + [("2", False), ("2.3", False), ("2.5", True)], +) +def test_ltx_reference_image_size_policy(tmp_path, model_version: str, expected: bool) -> None: + (tmp_path / "model_index.json").write_text( + '{"_class_name": "LTX2Pipeline", "model_version": "' + model_version + '"}' + ) + + assert ( + should_preserve_reference_image_size( + None, + model=str(tmp_path), + ) + is expected + ) + + +def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.MonkeyPatch) -> None: + captured = {} + + def fake_policy(*, model, revision=None): + captured.update(model=model, revision=revision) + return True + + monkeypatch.setattr( + "vllm_omni.diffusion.models.ltx2.ltx2_components.preserves_reference_image_size", + fake_policy, + ) + + assert should_preserve_reference_image_size( + "LTX2Pipeline", + model="org/model", + revision="pinned-revision", + ) + assert captured == {"model": "org/model", "revision": "pinned-revision"} + + @pytest.mark.core_model @pytest.mark.cpu def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None: @@ -424,6 +469,13 @@ def test_vace_extra_registry_has_no_pipeline_params() -> None: def test_unknown_pipeline_has_empty_extra_registry() -> None: assert get_extra_body_params("UnknownPipeline") == frozenset() assert get_extra_output_params("UnknownPipeline") == frozenset() + assert ( + should_preserve_reference_image_size( + "UnknownPipeline", + model="org/model", + ) + is False + ) assert should_init_extra_args_for_non_diffusion_stages("UnknownPipeline") is False diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index f5d6c18c4..fd13fd482 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -28,6 +28,8 @@ EXCLUDED_MODELS = [ "WanPipeline", "WanVACEPipeline", "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", + "LTX2DistilledTwoStagePipeline", "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py index ffcc2fb44..ab814f713 100644 --- a/vllm_omni/diffusion/attention/backends/cudnn_attn.py +++ b/vllm_omni/diffusion/attention/backends/cudnn_attn.py @@ -85,6 +85,24 @@ class CuDNNAttentionImpl(AttentionImpl): enable_gqa = query.shape[2] != key.shape[2] query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + # cuDNN has no kernel for degenerate attention shapes when either + # sequence has one token. Route those degenerate shapes before entering the + # cuDNN-only context: under torch.compile, FakeTensor evaluation raises + # before the eager RuntimeError fallback below can run. + if query.shape[-2] == 1 or key.shape[-2] == 1: + with sdpa_kernel([SDPBackend.MATH]): + output = torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=self.causal, + scale=self.softmax_scale, + enable_gqa=enable_gqa, + ) + return output.permute(0, 2, 1, 3) + # Pin cuDNN exclusively. A priority list like [CUDNN, FLASH, MATH] hits a # PyTorch SDPA dispatch quirk: when FLASH rejects a non-None attn_mask, # cuDNN gets runtime-disabled in the same call and the dispatcher falls @@ -93,7 +111,7 @@ class CuDNNAttentionImpl(AttentionImpl): # # Fall back to the default SDPA dispatcher if cuDNN rejects the shape, # e.g. under torch.compile where Dynamo sees a symbolic head_dim and - # cuDNN's kernel selection fails (observed in LTX-2 audio attention). + # cuDNN's kernel selection fails for some symbolic attention shapes. # The unpinned dispatcher then picks EFFICIENT/MATH instead of raising. try: with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]): diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 137fb9860..a5b6d06b5 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -538,7 +538,11 @@ class DiffusionCacheConfig: raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'") -def resolve_model_class_name(model: str | None, diffusion_load_format: str = "default") -> str | None: +def resolve_model_class_name( + model: str | None, + diffusion_load_format: str = "default", + revision: str | None = None, +) -> str | None: """Resolve the diffusion pipeline class name from the model config. Read-only counterpart of ``OmniDiffusionConfig.enrich_config``, safe to call @@ -550,11 +554,10 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de if not model: return None - is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"} # Diffusers models: read _class_name from the pipeline index. - model_index = get_diffusion_model_index(model) + model_index = get_diffusion_model_index(model, revision=revision) if model_index is not None: return model_index.get("_class_name") if diffusion_load_format == "diffusers": @@ -562,7 +565,7 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de # Other models: map model_type / architecture from config.json. try: - cfg = get_hf_file_to_dict("config.json", model) or {} + cfg = get_hf_file_to_dict("config.json", model, revision=revision) or {} except Exception: cfg = {} model_type = cfg.get("model_type") @@ -1210,9 +1213,20 @@ class OmniDiffusionConfig: exc, ) else: - tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) + from vllm_omni.model_extras import get_transformer_config_subfolder + + transformer_subfolder = get_transformer_config_subfolder( + self.model_class_name, + model=self.model, + revision=self.revision, + ) + tf_config_dict = get_hf_file_to_dict( + f"{transformer_subfolder}/config.json", + self.model, + revision=self.revision, + ) if tf_config_dict is None: - tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model) + tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision) if tf_config_dict is not None: self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) else: @@ -1232,7 +1246,7 @@ class OmniDiffusionConfig: "that require additional inputs." ) else: - cfg = get_hf_file_to_dict("config.json", self.model) + cfg = get_hf_file_to_dict("config.json", self.model, revision=self.revision) if cfg is None: # Lance ships its top-level config.json one directory above # the per-checkpoint subfolders (``Lance_3B/`` or diff --git a/vllm_omni/diffusion/model_loader/hub_prefetch.py b/vllm_omni/diffusion/model_loader/hub_prefetch.py index 9a9295ab9..3594aa4e1 100644 --- a/vllm_omni/diffusion/model_loader/hub_prefetch.py +++ b/vllm_omni/diffusion/model_loader/hub_prefetch.py @@ -265,6 +265,7 @@ def prefetch_subfolders( subfolders: Iterable[str], *, local_files_only: bool | None = None, + revision: str | None = None, include_root_metadata: bool = True, ) -> None: """Materialise ``model``'s ``subfolders`` in the HF cache before loading. @@ -278,6 +279,7 @@ def prefetch_subfolders( local_files_only: When ``True``, skip the prefetch entirely. When ``None`` (default), auto-detect: skip if *model* is a local directory, run otherwise. + revision: Optional Hub revision shared with component loaders. include_root_metadata: When True, also pull ``*.json`` at the repo root so ``model_index.json`` / ``config.json`` resolution during ``from_pretrained`` also hits a warm cache. @@ -334,6 +336,7 @@ def prefetch_subfolders( with _repo_prefetch_lock(model): snapshot_download( repo_id=model, + revision=revision, allow_patterns=allow_patterns, ) logger.info("Prefetch complete for %s", model) @@ -450,6 +453,7 @@ def from_pretrained_with_prefetch( on the first failure exactly as before. """ prefetch_list = list(prefetch_list) + revision = from_pretrained_kwargs.get("revision") can_heal = not local_files_only and bool(model) and not os.path.isdir(model) last_exc: BaseException | None = None @@ -481,7 +485,7 @@ def from_pretrained_with_prefetch( # Force a fresh, verified snapshot of every component this pipeline # needs - not just ``subfolder`` - so a sibling component that was # also half-written gets repaired in the same pass. - prefetch_subfolders(model, prefetch_list, local_files_only=False) + prefetch_subfolders(model, prefetch_list, local_files_only=False, revision=revision) assert last_exc is not None # loop only exits via return or a caught exc raise last_exc diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index 20ada6adc..a493f1c0b 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -8,21 +8,25 @@ from vllm_omni.diffusion.models.ltx2.ltx2_components import ( ) from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2DistilledOneStagePipeline, LTX2I2VDMD2Pipeline, LTX2Pipeline, LTX2T2VDMD2Pipeline, ) from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( LTX2DistilledPipeline, + LTX2DistilledTwoStagePipeline, LTX2TwoStagePipeline, ) __all__ = [ "LTX2Pipeline", + "LTX2DistilledOneStagePipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", "LTX2TwoStagePipeline", "LTX2DistilledPipeline", + "LTX2DistilledTwoStagePipeline", "get_ltx2_post_process_func", "load_transformer_config", "create_transformer_from_config", diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 7f2f31a60..37c3f05dc 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -22,7 +22,7 @@ from diffusers.video_processor import VideoProcessor from huggingface_hub import hf_hub_download from safetensors import safe_open from safetensors.torch import load_file -from transformers import AutoTokenizer, Gemma3ForConditionalGeneration +from transformers import AutoModelForImageTextToText, AutoTokenizer, Gemma3ForConditionalGeneration from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as OmniAttention @@ -48,6 +48,14 @@ try: except ImportError: LTX2VocoderWithBWE = None +try: + from transformers import Gemma4UnifiedForConditionalGeneration as _Gemma4UnifiedForConditionalGeneration +except ImportError: + _Gemma4UnifiedForConditionalGeneration = None + + +_LTX25_TEXT_ENCODER_CLS = AutoModelForImageTextToText if _Gemma4UnifiedForConditionalGeneration is not None else None + _LTX_COMPONENT_SUBFOLDERS = ( "tokenizer", @@ -73,10 +81,15 @@ class LTXComponentProfile: resident_modules: tuple[str, ...] = () video_vae_cls: type = AutoencoderKLLTX2Video vocoder_cls: type = LTX2Vocoder + text_encoder_cls: type | None = Gemma3ForConditionalGeneration vocoder_fallback_cls: type | None = None artifact_repo_id: str | None = None latent_upsampler_filename: str | None = None distilled_lora_filename: str | None = None + transformer_subfolder: str = "transformer" + scheduler_use_dynamic_shifting: bool = False + scheduler_shift_terminal: float | None = None + preserve_connector_attention_mask: bool = False LTX2_COMPONENT_PROFILE = LTXComponentProfile( @@ -99,6 +112,23 @@ LTX23_COMPONENT_PROFILE = LTXComponentProfile( vocoder_fallback_cls=LTX2Vocoder, ) +LTX25_FULL_COMPONENT_PROFILE = LTXComponentProfile( + name="ltx2_5_full", + dit_modules=("transformer",), + encoder_modules=("text_encoder", "connectors"), + vae_modules=("vae", "audio_vae"), + resident_modules=("vocoder",), + video_vae_cls=DistributedAutoencoderKLLTX2Video, + vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder, + vocoder_fallback_cls=LTX2Vocoder, + text_encoder_cls=_LTX25_TEXT_ENCODER_CLS, + transformer_subfolder="transformer_full", + scheduler_use_dynamic_shifting=True, + scheduler_shift_terminal=0.1, + preserve_connector_attention_mask=True, +) + + LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile( name="ltx2_distilled", dit_modules=("transformer",), @@ -108,6 +138,32 @@ LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile( video_vae_cls=DistributedAutoencoderKLLTX2Video, ) +LTX25_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile( + name="ltx2_5_distilled", + dit_modules=("transformer",), + encoder_modules=("text_encoder", "connectors"), + vae_modules=("vae", "audio_vae"), + resident_modules=("vocoder", "latent_upsampler"), + video_vae_cls=DistributedAutoencoderKLLTX2Video, + vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder, + vocoder_fallback_cls=LTX2Vocoder, + text_encoder_cls=_LTX25_TEXT_ENCODER_CLS, + preserve_connector_attention_mask=True, +) + +LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace( + LTX2_DISTILLED_COMPONENT_PROFILE, + name="ltx2_distilled_one_stage", + resident_modules=LTX2_COMPONENT_PROFILE.resident_modules, +) + +LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace( + LTX25_DISTILLED_COMPONENT_PROFILE, + name="ltx2_5_distilled_one_stage", + resident_modules=LTX25_FULL_COMPONENT_PROFILE.resident_modules, +) + + LTX23_DISTILLED_COMPONENT_PROFILE = replace( LTX23_COMPONENT_PROFILE, name="ltx2_3_distilled", @@ -116,6 +172,12 @@ LTX23_DISTILLED_COMPONENT_PROFILE = replace( latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", ) +LTX23_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace( + LTX23_DISTILLED_COMPONENT_PROFILE, + name="ltx2_3_distilled_one_stage", + resident_modules=LTX23_COMPONENT_PROFILE.resident_modules, +) + LTX2_TWO_STAGE_COMPONENT_PROFILE = replace( LTX2_COMPONENT_PROFILE, name="ltx2_two_stage", @@ -135,18 +197,44 @@ LTX23_TWO_STAGE_COMPONENT_PROFILE = replace( ) +LTX25_TWO_STAGE_COMPONENT_PROFILE = replace( + LTX25_FULL_COMPONENT_PROFILE, + name="ltx2_5_two_stage", + resident_modules=(*LTX25_FULL_COMPONENT_PROFILE.resident_modules, "latent_upsampler"), + artifact_repo_id="Lightricks/LTX-2.5", + latent_upsampler_filename=("latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors"), + distilled_lora_filename="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors", +) + _COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = { ("one_stage", "2"): LTX2_COMPONENT_PROFILE, ("one_stage", "2.3"): LTX23_COMPONENT_PROFILE, + ("one_stage", "2.5"): LTX25_FULL_COMPONENT_PROFILE, ("two_stage", "2"): LTX2_TWO_STAGE_COMPONENT_PROFILE, ("two_stage", "2.3"): LTX23_TWO_STAGE_COMPONENT_PROFILE, + ("two_stage", "2.5"): LTX25_TWO_STAGE_COMPONENT_PROFILE, + ("distilled_one_stage", "2"): LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, + ("distilled_one_stage", "2.3"): LTX23_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, + ("distilled_one_stage", "2.5"): LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, ("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE, ("distilled_two_stage", "2.3"): LTX23_DISTILLED_COMPONENT_PROFILE, + ("distilled_two_stage", "2.5"): LTX25_DISTILLED_COMPONENT_PROFILE, ("dmd2", "2"): LTX2_COMPONENT_PROFILE, ("dmd2", "2.3"): LTX23_COMPONENT_PROFILE, } +def resolve_ltx_checkpoint_kind(pipeline_kind: str) -> LTXCheckpointKind | None: + """Derive checkpoint requirements from the execution contract.""" + if pipeline_kind in {"one_stage", "two_stage"}: + return "regular" + if pipeline_kind in {"distilled_one_stage", "distilled_two_stage"}: + return "distilled" + if pipeline_kind == "dmd2": + return None + raise ValueError(f"Unsupported LTX pipeline kind: {pipeline_kind!r}.") + + def resolve_ltx_artifact( model: str, repo_id: str, @@ -206,7 +294,7 @@ def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTX raise ValueError(f"Unsupported LTX component kind/version: {pipeline_kind!r}/{model_version!r}.") from exc -def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]: +def _load_ltx_metadata_json(model: str, filename: str, revision: str | None = None) -> dict[str, Any]: """Load small checkpoint metadata without relying on repository names.""" if os.path.isdir(model): path = os.path.join(model, filename) @@ -214,7 +302,7 @@ def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]: return {} else: try: - path = hf_hub_download(repo_id=model, filename=filename) + path = hf_hub_download(repo_id=model, filename=filename, revision=revision) except Exception: return {} try: @@ -225,17 +313,34 @@ def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]: return value if isinstance(value, dict) else {} -def detect_ltx_model_version(model: str) -> str: - """Detect LTX-2 versus LTX-2.3 from checkpoint component metadata. +def detect_ltx_model_version(model: str, revision: str | None = None) -> str: + """Detect the LTX model version from checkpoint component metadata.""" + model_index = _load_ltx_metadata_json(model, "model_index.json", revision) + model_version = str(model_index.get("model_version", "")) + if model_version.startswith("2.5"): + return "2.5" - Official checkpoints use ``model_version`` metadata. Diffusers repositories - expose the equivalent distinction through the BWE vocoder introduced by - LTX-2.3. Unknown conversions retain the official LTX-2 fallback. - """ - model_index = _load_ltx_metadata_json(model, "model_index.json") - if str(model_index.get("model_version", "")).startswith("2.3"): + text_encoder_entry = model_index.get("text_encoder") + if isinstance(text_encoder_entry, (list, tuple)) and text_encoder_entry: + text_encoder_class = str(text_encoder_entry[-1]) + elif isinstance(text_encoder_entry, dict): + text_encoder_class = str(text_encoder_entry.get("_class_name", "")) + else: + text_encoder_class = "" + if text_encoder_class == "Gemma4UnifiedForConditionalGeneration": + return "2.5" + + text_encoder_config = _load_ltx_metadata_json(model, "text_encoder/config.json", revision) + if text_encoder_config.get("model_type") in ("gemma4_unified", "gemma4"): + return "2.5" + + if model_version.startswith("2.3"): return "2.3" + transformer_config = _load_ltx_metadata_json(model, "transformer/config.json", revision) + if transformer_config.get("ff_bias") is False: + return "2.5" + vocoder_entry = model_index.get("vocoder") if isinstance(vocoder_entry, (list, tuple)) and vocoder_entry: vocoder_class = str(vocoder_entry[-1]) @@ -246,21 +351,32 @@ def detect_ltx_model_version(model: str) -> str: if vocoder_class == "LTX2VocoderWithBWE": return "2.3" - vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json") + vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json", revision) if str(vocoder_config.get("model_version", "")).startswith("2.3"): return "2.3" if vocoder_config.get("_class_name") == "LTX2VocoderWithBWE": return "2.3" - - logger.info("Using LTX-2 defaults for checkpoint %s", model) + if "bwe" in vocoder_config or "bwe_config" in vocoder_config: + return "2.3" return "2" +def preserves_reference_image_size(*, model: str | None, revision: str | None = None) -> bool: + """Preserve source geometry only for the LTX-2.5 CRF-18 I2V path.""" + return model is not None and detect_ltx_model_version(model, revision=revision) == "2.5" + + class _LTXConnectorAttnProcessor: """Preserve official connector math around Omni attention dispatch.""" - def __init__(self, *, has_learned_registers: bool) -> None: + def __init__( + self, + *, + has_learned_registers: bool = False, + preserve_learned_register_mask: bool = False, + ) -> None: self.has_learned_registers = has_learned_registers + self.preserve_learned_register_mask = preserve_learned_register_mask def __call__( self, @@ -274,12 +390,20 @@ class _LTXConnectorAttnProcessor: encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states gate_logits = attn.to_gate_logits(hidden_states) if attn.to_gate_logits is not None else None - query = attn.norm_q(attn.to_q(hidden_states)) - key = attn.norm_k(attn.to_k(encoder_hidden_states)) + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) value = attn.to_v(encoder_hidden_states) + # Offload hooks may execute affine-free Q/K norms in FP32. Restore the + # projection dtype before attention, matching the fully resident path. + query = attn.norm_q(query).to(dtype=value.dtype) + key = attn.norm_k(key).to(dtype=value.dtype) if query_rotary_emb is not None: + # Diffusers builds connector RoPE in FP32, while the official + # connector materializes it in the hidden-state dtype. + query_rotary_emb = tuple(component.to(value.dtype) for component in query_rotary_emb) key_rotary_emb = key_rotary_emb if key_rotary_emb is not None else query_rotary_emb + key_rotary_emb = tuple(component.to(value.dtype) for component in key_rotary_emb) if attn.rope_type == "interleaved": query = apply_interleaved_rotary_emb(query, query_rotary_emb) key = apply_interleaved_rotary_emb(key, key_rotary_emb) @@ -289,6 +413,10 @@ class _LTXConnectorAttnProcessor: else: raise ValueError(f"Unsupported LTX connector RoPE type: {attn.rope_type}") + # Keep Q/K in the projection dtype expected by the attention backend. + query = query.to(dtype=value.dtype) + key = key.to(dtype=value.dtype) + batch_size, _, inner_dim = query.shape head_dim = inner_dim // attn.heads kv_heads = attn.inner_kv_dim // attn.head_dim @@ -296,10 +424,10 @@ class _LTXConnectorAttnProcessor: key = key.view(batch_size, -1, kv_heads, head_dim) value = value.view(batch_size, -1, kv_heads, head_dim) - # The connector replaces padding tokens with learned registers before - # entering its blocks, so every key is valid and the old padding mask - # becomes an all-keep no-op. - if self.has_learned_registers: + # LTX-2/2.3 replace padding tokens with learned registers, so their + # old padding mask becomes a no-op. LTX-2.5 preserves that all-zero + # additive mask to match the official masked-SDPA dispatch. + if self.has_learned_registers and not self.preserve_learned_register_mask: attention_mask = None elif attention_mask is not None and attn.omni_attention.attn_backend.get_name().upper() == "FLASH_ATTN": attention_mask = to_ltx_padding_mask(attention_mask) @@ -316,7 +444,11 @@ class _LTXConnectorAttnProcessor: return attn.to_out[1](hidden_states) -def _install_connector_attention(connectors: LTX2TextConnectors) -> None: +def _install_connector_attention( + connectors: LTX2TextConnectors, + *, + preserve_learned_register_mask: bool = False, +) -> None: for connector_name in ("video_connector", "audio_connector"): connector = getattr(connectors, connector_name, None) has_learned_registers = getattr(connector, "learnable_registers", None) is not None @@ -335,15 +467,20 @@ def _install_connector_attention(connectors: LTX2TextConnectors) -> None: skip_sequence_parallel=True, disable_kv_quant=True, ) - attention.set_processor(_LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)) + attention.set_processor( + _LTXConnectorAttnProcessor( + has_learned_registers=has_learned_registers, + preserve_learned_register_mask=preserve_learned_register_mask, + ) + ) -def _detect_vocoder_output_sample_rate(model: str) -> int | None: +def _detect_vocoder_output_sample_rate(model: str, revision: str | None = None) -> int | None: """Read the generated waveform sample rate from the vocoder config.""" vocoder_config_path = os.path.join(model, "vocoder", "config.json") if not os.path.exists(vocoder_config_path): try: - vocoder_config_path = hf_hub_download(model, "vocoder/config.json") + vocoder_config_path = hf_hub_download(model, "vocoder/config.json", revision=revision) except Exception: return None try: @@ -355,7 +492,10 @@ def _detect_vocoder_output_sample_rate(model: str) -> int | None: def get_ltx2_post_process_func(od_config: Any): """Build the common LTX engine-output adapter.""" - output_sample_rate = _detect_vocoder_output_sample_rate(od_config.model) + output_sample_rate = _detect_vocoder_output_sample_rate( + od_config.model, + revision=getattr(od_config, "revision", None), + ) def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor): if not (isinstance(output, tuple) and len(output) == 2): @@ -378,6 +518,7 @@ def _load_component( *, local_files_only: bool, dtype: torch.dtype, + revision: str | None, ) -> Any: return from_pretrained_with_prefetch( component_cls.from_pretrained, @@ -385,6 +526,7 @@ def _load_component( subfolder=subfolder, prefetch_list=_LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only, + revision=revision, torch_dtype=dtype, ) @@ -411,31 +553,36 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: pipeline.device = get_local_device() dtype = getattr(od_config, "dtype", torch.bfloat16) model = od_config.model + revision = getattr(od_config, "revision", None) local_files_only = os.path.exists(model) pipeline.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=model, - subfolder="transformer", - revision=None, + subfolder=profile.transformer_subfolder, + revision=revision, prefix="transformer.", fall_back_to_pt=True, ), ] - prefetch_subfolders(model, _LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only) + prefetch_subfolders(model, _LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only, revision=revision) pipeline.tokenizer = AutoTokenizer.from_pretrained( model, subfolder="tokenizer", local_files_only=local_files_only, + revision=revision, ) + if profile.text_encoder_cls is None: + raise ImportError("LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15.") with torch.device("cpu"): pipeline.text_encoder = _load_component( - Gemma3ForConditionalGeneration, + profile.text_encoder_cls, model, "text_encoder", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) pipeline.connectors = _load_component( LTX2TextConnectors, @@ -443,14 +590,19 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: "connectors", local_files_only=local_files_only, dtype=dtype, + revision=revision, + ) + _install_connector_attention( + pipeline.connectors, + preserve_learned_register_mask=profile.preserve_connector_attention_mask, ) - _install_connector_attention(pipeline.connectors) pipeline.vae = _load_component( profile.video_vae_cls, model, "vae", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) pipeline.audio_vae = _load_component( AutoencoderKLLTX2Audio, @@ -458,6 +610,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: "audio_vae", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) try: pipeline.vocoder = _load_component( @@ -466,6 +619,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: "vocoder", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) except (TypeError, OSError, ValueError): if profile.vocoder_fallback_cls is None or profile.vocoder_fallback_cls is profile.vocoder_cls: @@ -476,6 +630,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: "vocoder", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) if "latent_upsampler" in profile.resident_modules: @@ -488,6 +643,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: "latent_upsampler", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) except (OSError, ValueError): if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None: @@ -508,7 +664,9 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: ) pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype) - transformer_config = load_transformer_config(model, "transformer", local_files_only) + transformer_config = load_transformer_config( + model, profile.transformer_subfolder, local_files_only, revision=revision + ) quant_config = getattr(od_config, "quantization_config", None) pipeline.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) _place_aux_components(pipeline) @@ -516,15 +674,17 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: model, subfolder="scheduler", local_files_only=local_files_only, + revision=revision, ) - expected_checkpoint_kind: LTXCheckpointKind | None = None - if pipeline.pipeline_kind == "two_stage": - expected_checkpoint_kind = "regular" - elif pipeline.pipeline_kind == "distilled_two_stage": - expected_checkpoint_kind = "distilled" + if profile.scheduler_use_dynamic_shifting: + pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config( + pipeline.scheduler.config, + use_dynamic_shifting=True, + shift_terminal=profile.scheduler_shift_terminal, + ) validate_ltx_checkpoint( pipeline.scheduler.config, - expected_kind=expected_checkpoint_kind, + expected_kind=resolve_ltx_checkpoint_kind(pipeline.pipeline_kind), pipeline_name=type(pipeline).__name__, ) @@ -553,6 +713,8 @@ def load_transformer_config( model_path: str, subfolder: str = "transformer", local_files_only: bool = True, + *, + revision: str | None = None, ) -> dict: """Load an LTX transformer config from a local model or the HF Hub.""" if local_files_only: @@ -563,6 +725,7 @@ def load_transformer_config( config_path = hf_hub_download( repo_id=model_path, filename=f"{subfolder}/config.json", + revision=revision, ) with open(config_path) as config_file: return json.load(config_file) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py index f4d2f26f6..71beab643 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py @@ -7,6 +7,7 @@ from __future__ import annotations import math from dataclasses import dataclass +from io import BytesIO from typing import TYPE_CHECKING, Any import numpy as np @@ -52,19 +53,18 @@ def _preprocess_i2v_pil_images( *, height: int, width: int, + crf: int = 0, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> torch.Tensor: - """Match official LTX aspect-preserving resize and center crop.""" + """Match the version-specific official LTX PIL preprocessing.""" image_list = images if isinstance(images, list) else [images] processed = [] for image in image_list: - pixels = ( - torch.from_numpy(np.array(image.convert("RGB"), copy=True)) - .permute(2, 0, 1) - .unsqueeze(0) - .to(device=device, dtype=torch.float32) - ) + image_array = np.array(image.convert("RGB"), copy=True) + if crf: + image_array = _apply_image_conditioning_crf(image_array, crf) + pixels = torch.from_numpy(image_array).permute(2, 0, 1).unsqueeze(0).to(device=device, dtype=torch.float32) source_height, source_width = pixels.shape[-2:] scale = max(height / source_height, width / source_width) resized_height = math.ceil(source_height * scale) @@ -82,6 +82,47 @@ def _preprocess_i2v_pil_images( return torch.cat(processed, dim=0) +def _apply_image_conditioning_crf(image: np.ndarray, crf: int) -> np.ndarray: + """Round-trip one RGB frame through H.264 to match LTX-2.5 training.""" + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[-1] != 3: + raise ValueError("LTX image-conditioning CRF expects a uint8 RGB image.") + if min(image.shape[:2]) < 2: + return image + try: + import av + except ImportError as exc: + raise ImportError("PyAV with a libx264 encoder is required for LTX-2.5 image conditioning (CRF 18).") from exc + + with BytesIO() as output_file: + container = av.open(output_file, "w", format="mp4") + try: + stream = container.add_stream("libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"}) + height = image.shape[0] // 2 * 2 + width = image.shape[1] // 2 * 2 + stream.height = height + stream.width = width + frame = av.VideoFrame.from_ndarray(image[:height, :width], format="rgb24").reformat(format="yuv420p") + container.mux(stream.encode(frame)) + container.mux(stream.encode()) + except Exception as exc: + raise RuntimeError( + "The installed PyAV/FFmpeg build cannot encode libx264 video required " + "for LTX-2.5 image conditioning (CRF 18)." + ) from exc + finally: + container.close() + video_bytes = output_file.getvalue() + + with BytesIO(video_bytes) as video_file: + container = av.open(video_file) + try: + stream = next(stream for stream in container.streams if stream.type == "video") + frame = next(container.decode(stream)) + finally: + container.close() + return frame.to_ndarray(format="rgb24") + + class LTXTextConditioningMixin: """Shared Gemma encoding and text connector orchestration.""" @@ -106,13 +147,30 @@ class LTXTextConditioningMixin: if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token - text_inputs = self.tokenizer( - [text.strip() for text in prompt], + bos_token_id = self.tokenizer.bos_token_id + if bos_token_id is None: + raise ValueError("Tokenizer is missing bos_token_id; LTX prompt encoding requires a leading BOS.") + + input_ids = [] + for text in prompt: + encoded = self.tokenizer( + text.strip(), + padding=False, + truncation=True, + max_length=max_sequence_length, + return_tensors=None, + ) + ids = encoded.input_ids + if not ids or ids[0] != bos_token_id: + ids = [bos_token_id, *ids][:max_sequence_length] + input_ids.append(ids) + + text_inputs = self.tokenizer.pad( + {"input_ids": input_ids}, padding="max_length", max_length=max_sequence_length, - truncation=True, - add_special_tokens=True, return_tensors="pt", + return_attention_mask=True, ) text_input_ids = text_inputs.input_ids.to(device) prompt_attention_mask = text_inputs.attention_mask.to(device) @@ -121,7 +179,6 @@ class LTXTextConditioningMixin: attention_mask=prompt_attention_mask, output_hidden_states=True, ).hidden_states - prompt_embeds = torch.stack(hidden_states, dim=-1).flatten(2, 3).to(dtype=dtype) prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt) prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) @@ -501,6 +558,7 @@ class LTXI2VConditioningMixin: vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, ) mask_shape = (batch_size, 1, num_frames, height, width) + is_ltx25 = getattr(self, "model_version", None) == "2.5" if latents is not None: unpacked_latents = latents.ndim == 5 @@ -518,6 +576,7 @@ class LTXI2VConditioningMixin: self.vae.latents_std, self.vae.config.scaling_factor, ) + clean_latents = latents if image is not None: image_latents = self._encode_i2v_image_latents( image, @@ -525,14 +584,19 @@ class LTXI2VConditioningMixin: generator=generator, dtype=dtype, ) - # Official distilled Stage 2 starts from the upsampled Stage 1 - # state, then reapplies the source image at final resolution. - latents = self._replace_first_latent_frame(latents, image_latents) + # Official Stage 2 starts from the upsampled Stage 1 state, + # then reapplies the source image at final resolution. + clean_latents = self._replace_first_latent_frame(clean_latents, image_latents) latents = latent_ops.pack_latents( latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ) + clean_latents = latent_ops.pack_latents( + clean_latents, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) else: conditioning_mask = latents.new_zeros(mask_shape) conditioning_mask[:, :, 0] = 1.0 @@ -548,14 +612,14 @@ class LTXI2VConditioningMixin: f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}." ) if unpacked_latents: - # Official LTX patchifies the latent state before drawing noise. - # Sampling the packed shape preserves its seed-to-output contract. - latents = latent_ops.create_noised_state( + latents = latent_ops.create_conditioned_noised_state( latents, - noise_scale * (1 - conditioning_mask).unsqueeze(-1), + clean_latents, + (1 - conditioning_mask.float()).unsqueeze(-1), + noise_scale, generator, ) - elif image is not None: + elif not unpacked_latents and image is not None: image_latents = self._encode_i2v_image_latents( image, batch_size=batch_size, @@ -595,13 +659,27 @@ class LTXI2VConditioningMixin: self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ) - conditioning_mask = latent_ops.pack_latents( + packed_conditioning_mask = latent_ops.pack_latents( conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ).squeeze(-1) - noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype) - latents = init_latents * conditioning_mask.unsqueeze(-1) + noise * (1 - conditioning_mask).unsqueeze(-1) + if is_ltx25: + clean_latents = init_latents + init_latents = torch.zeros_like(init_latents) + latents = latent_ops.create_conditioned_noised_state( + init_latents, + clean_latents, + (1 - packed_conditioning_mask.float()).unsqueeze(-1), + noise_scale, + generator=generator, + ) + else: + noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype) + latents = init_latents * packed_conditioning_mask.unsqueeze(-1) + noise * ( + 1 - packed_conditioning_mask + ).unsqueeze(-1) + conditioning_mask = packed_conditioning_mask return latents, conditioning_mask def _step_video_latents_i2v( @@ -665,6 +743,17 @@ class LTXI2VConditioningMixin: ) if image is not None: + default_crf = 18 if getattr(self, "model_version", "2") == "2.5" else 0 + image_crf = getattr(request_inputs, "image_crf", None) + image_crf = default_crf if image_crf is None else int(image_crf) + is_pil_input = isinstance(image, PIL.Image.Image) or ( + isinstance(image, list) and image and all(isinstance(item, PIL.Image.Image) for item in image) + ) + if image_crf != 0 and not is_pil_input: + raise ValueError( + "`image_crf` re-compression requires PIL images. Pass PIL images, or set `image_crf=0` " + "to skip re-compression." + ) if isinstance(image, torch.Tensor): if image.ndim == 3: image = image.unsqueeze(0) @@ -677,6 +766,7 @@ class LTXI2VConditioningMixin: image, height=request_inputs.height, width=request_inputs.width, + crf=image_crf, device=device, dtype=prompt_context.positive_connector_prompt_embeds.dtype, ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py index dd55f3953..7ebd94143 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -42,6 +42,7 @@ class LTXForwardContext: original_audio_num_frames: int padded_audio_num_frames: int timesteps: torch.Tensor + sampler: str audio_scheduler: Any video_audio_step_adapter: Any @@ -88,6 +89,11 @@ class LTXDenoisePipeline(Protocol): LTXDenoiseStep = Callable[[int, torch.Tensor, LTXAVState], LTXAVState] +LTX_ANCESTRAL_ETA = 1.0 +LTX_ANCESTRAL_S_NOISE = 1.0 +LTX_ANCESTRAL_NOISE_SEED_OFFSET = 10000 +LTX_OFFICIAL_DEFAULT_SEED = 10 + class LTXDenoiseExecutor: """Run the one shared LTX denoise loop. @@ -138,6 +144,9 @@ class LTXVideoAudioStepAdapter: latent_width: int, *, image_conditioned: bool, + sampler: str = "euler", + generator: torch.Generator | list[torch.Generator] | None = None, + conditioning_mask: torch.Tensor | None = None, ) -> None: self._pipeline = pipeline self._audio_scheduler = audio_scheduler @@ -145,11 +154,27 @@ class LTXVideoAudioStepAdapter: self._latent_height = latent_height self._latent_width = latent_width self._image_conditioned = image_conditioned + self._sampler = sampler + self._conditioning_mask = conditioning_mask + self._ancestral_generators = ( + _make_ancestral_generators(generator, pipeline.device) if sampler == "euler_ancestral" else None + ) self._step_index = 0 def step(self, noise_pred, t, latents, return_dict=False, generator=None): del t, return_dict, generator - if self._image_conditioned: + if self._sampler == "euler_ancestral": + video_out = _ancestral_euler_step_from_velocity( + latents[0], + noise_pred[0], + self._pipeline.scheduler.sigmas, + self._step_index, + self._ancestral_generators, + ) + if self._conditioning_mask is not None: + mask = self._conditioning_mask.unsqueeze(-1).to(video_out.dtype) + video_out = torch.lerp(video_out, latents[0], mask) + elif self._image_conditioned: video_out = self._pipeline._step_video_latents_i2v( noise_pred[0], latents[0], @@ -165,17 +190,96 @@ class LTXVideoAudioStepAdapter: self._pipeline.scheduler.sigmas, self._step_index, ) - audio_out = euler_step_from_velocity( - latents[1], - noise_pred[1], - self._audio_scheduler.sigmas, - self._step_index, - ) + if self._sampler == "euler_ancestral": + audio_out = _ancestral_euler_step_from_velocity( + latents[1], + noise_pred[1], + self._audio_scheduler.sigmas, + self._step_index, + self._ancestral_generators, + ) + else: + audio_out = euler_step_from_velocity( + latents[1], + noise_pred[1], + self._audio_scheduler.sigmas, + self._step_index, + ) self._step_index += 1 return ((video_out, audio_out),) -def _official_ltx_sigmas(scheduler: Any, steps: int, device: torch.device) -> torch.Tensor: +def _make_ancestral_generators( + generator: torch.Generator | list[torch.Generator] | None, + device: torch.device, +) -> torch.Generator | list[torch.Generator]: + """Create the official independent ``seed + 10000`` noise stream.""" + + def _offset(source: torch.Generator | None) -> torch.Generator: + seed = LTX_OFFICIAL_DEFAULT_SEED if source is None else source.initial_seed() + return torch.Generator(device=device).manual_seed(seed + LTX_ANCESTRAL_NOISE_SEED_OFFSET) + + if isinstance(generator, list): + return [_offset(item) for item in generator] + return _offset(generator) + + +def _randn_like_with_generators( + sample: torch.Tensor, + generators: torch.Generator | list[torch.Generator], +) -> torch.Tensor: + if not isinstance(generators, list): + return torch.randn(sample.shape, generator=generators, dtype=sample.dtype, device=sample.device) + if len(generators) != sample.shape[0]: + raise ValueError( + f"LTX ancestral sampling received {len(generators)} generators for batch size {sample.shape[0]}." + ) + return torch.cat( + [ + torch.randn((1, *sample.shape[1:]), generator=item, dtype=sample.dtype, device=sample.device) + for item in generators + ] + ) + + +def _ancestral_euler_step_from_velocity( + sample: torch.Tensor, + velocity: torch.Tensor, + sigmas: torch.Tensor, + step_index: int, + generators: torch.Generator | list[torch.Generator] | None, +) -> torch.Tensor: + """Apply the official LTX-2.5 rectified-flow ancestral Euler step.""" + sigma = sigmas[step_index].to(torch.float32) + sigma_next = sigmas[step_index + 1].to(torch.float32) + # Official X0Model materializes the denoised prediction in the model + # dtype before the ancestral step promotes it back to float32. + denoised_model_dtype = (sample.float() - velocity.float() * sigma).to(sample.dtype) + denoised = denoised_model_dtype.float() + if sigma_next == 0: + return denoised.to(sample.dtype) + if generators is None: + raise ValueError("LTX ancestral Euler sampling requires a noise generator.") + + downstep_ratio = 1.0 + (sigma_next / sigma - 1.0) * LTX_ANCESTRAL_ETA + sigma_down = sigma_next * downstep_ratio + sigma_down_ratio = sigma_down / sigma + x_next = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised + + alpha_next = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + renoise_coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).clamp(min=0).sqrt() + noise = _randn_like_with_generators(sample, generators) + x_next = (alpha_next / alpha_down) * x_next + noise.float() * LTX_ANCESTRAL_S_NOISE * renoise_coeff + result = x_next.to(sample.dtype) + return result + + +def _official_ltx_sigmas( + scheduler: Any, + steps: int, + device: torch.device, +) -> torch.Tensor: """Build the official LTX one-stage sigma schedule in torch fp32.""" config = scheduler.config base_anchor = config.get("base_image_seq_len", 1024) @@ -185,7 +289,8 @@ def _official_ltx_sigmas(scheduler: Any, steps: int, device: torch.device) -> to sigmas = torch.linspace(1.0, 0.0, steps + 1) slope = (max_shift - base_shift) / (max_anchor - base_anchor) - # Official LTX2 one-stage intentionally uses the max sequence anchor, so the shift stays at max_shift. + # Official LTX one-stage pipelines omit the latent when constructing + # this schedule, so the shift stays at the max sequence anchor. sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor) exp_shift = math.exp(sigma_shift) sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0) @@ -224,6 +329,9 @@ def prepare_scheduler_stage( latent_width: int, use_official_sigma_schedule: bool, image_conditioned: bool = False, + sampler: str = "euler", + generator: torch.Generator | list[torch.Generator] | None = None, + conditioning_mask: torch.Tensor | None = None, ) -> tuple[Any, Any, torch.Tensor]: if sigmas is not None and timesteps is not None: raise ValueError("Only one of `sigmas` or `timesteps` may be provided.") @@ -236,6 +344,9 @@ def prepare_scheduler_stage( latent_height, latent_width, image_conditioned=image_conditioned, + sampler=sampler, + generator=generator, + conditioning_mask=conditioning_mask, ) if sigmas is not None: scheduler_sigmas = torch.as_tensor(sigmas, dtype=torch.float32, device=device) @@ -302,6 +413,20 @@ def prepare_rope_coords_stage( return video_coords, audio_coords +def _first_frame_keyframes_mask(reference: torch.Tensor, latent_num_frames: int) -> torch.Tensor: + """Mark the first causal latent frame, matching official LTX-2.5.""" + if latent_num_frames <= 0: + raise ValueError(f"LTX latent frame count must be positive, got {latent_num_frames}.") + tokens_per_frame, remainder = divmod(reference.shape[1], latent_num_frames) + if remainder: + raise ValueError( + f"LTX video token count {reference.shape[1]} is not divisible by {latent_num_frames} latent frames." + ) + mask = reference.new_zeros((reference.shape[0], reference.shape[1], 1)) + mask[:, :tokens_per_frame] = 1 + return mask + + def build_transformer_kwargs( pipeline: Any, forward_ctx: LTXForwardContext, @@ -317,9 +442,13 @@ def build_transformer_kwargs( attention_kwargs: dict[str, Any] | None = None, ) -> dict[str, Any]: del encoder_attention_mask, audio_encoder_attention_mask + keyframes_mask = None + if getattr(pipeline.transformer.config, "use_keyframes_abs_pos_embedding", False): + keyframes_mask = _first_frame_keyframes_mask(hidden_states, forward_ctx.latent_num_frames) return { "hidden_states": hidden_states, "audio_hidden_states": audio_hidden_states, + "keyframes_mask": keyframes_mask, "encoder_hidden_states": encoder_hidden_states, "audio_encoder_hidden_states": audio_encoder_hidden_states, **pipeline._denoise_timestep_kwargs( @@ -426,6 +555,9 @@ class LTXPhaseExecutor: latent_width=latent_width, use_official_sigma_schedule=phase_recipe.use_official_sigma_schedule, image_conditioned=conditioning_mask is not None, + sampler=phase_recipe.sampler, + generator=request_inputs.generator, + conditioning_mask=conditioning_mask, ) forward_ctx = LTXForwardContext( req=req, @@ -441,6 +573,7 @@ class LTXPhaseExecutor: original_audio_num_frames=original_audio_num_frames, padded_audio_num_frames=padded_audio_num_frames, timesteps=timesteps_tensor, + sampler=phase_recipe.sampler, audio_scheduler=audio_scheduler, video_audio_step_adapter=video_audio_step_adapter, ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py index bc84e7c32..9bb8ded75 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py @@ -126,9 +126,9 @@ def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch. def velocity_from_x0(sample: torch.Tensor, x0: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: """Convert x0 back to velocity using the official fp32 arithmetic order.""" - # Official LTX converts the scalar scheduler sigma to a host value before - # division. A CUDA scalar tensor selects different kernel arithmetic and - # can cross bf16 rounding boundaries late in the denoise trajectory. + # Official LTX materializes the scalar schedule value on the host before + # dividing. CUDA tensor-scalar division can cross bfloat16 rounding bounds + # even when the scalar has the same float32 value. sigma = sigma.to(torch.float32).item() return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype) @@ -305,10 +305,17 @@ class LTXGuidanceExecutor: ts: torch.Tensor, video_token_count: int, audio_token_count: int, + *, + expand_for_sequence_parallel: bool = False, ) -> dict[str, torch.Tensor]: + timestep = ts + audio_timestep = ts + if expand_for_sequence_parallel: + timestep = ts.reshape(-1, 1).expand(-1, video_token_count) + audio_timestep = ts.reshape(-1, 1).expand(-1, audio_token_count) return { - "timestep": ts.reshape(-1, 1).expand(-1, video_token_count), - "audio_timestep": ts.reshape(-1, 1).expand(-1, audio_token_count), + "timestep": timestep, + "audio_timestep": audio_timestep, "sigma": ts, "audio_sigma": ts, } @@ -342,8 +349,11 @@ class LTXGuidanceExecutor: guidance: LTXModalityGuidance, *, model_sigma: torch.Tensor | None = None, + preserve_positive_velocity: bool = False, rescale_token_count: int | None = None, ) -> torch.Tensor: + if preserve_positive_velocity and tuple(splits) == ("cond",): + return splits["cond"] model_sigma = sigma if model_sigma is None else model_sigma # Official guidance reduces contiguous BSC tensors. The fused # transformer output may be channel-major after splitting, which @@ -368,6 +378,7 @@ class LTXGuidanceExecutor: state: LTXAVState, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, + preserve_positive_velocity: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: guidance_world_size = get_guidance_parallel_world_size() self.validate_guidance_world_size(plan, guidance_world_size) @@ -450,12 +461,14 @@ class LTXGuidanceExecutor: video_sigma, plan.spec.video, model_sigma=pipeline._video_guidance_model_sigma(video_sigma, denoise_ctx), + preserve_positive_velocity=preserve_positive_velocity, ), self._guide_modality( state.audio, audio_splits, forward_ctx.audio_scheduler.sigmas[index], plan.spec.audio, + preserve_positive_velocity=preserve_positive_velocity, rescale_token_count=forward_ctx.original_audio_num_frames, ), ) @@ -469,9 +482,19 @@ class LTXGuidanceExecutor: state: LTXAVState, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, + preserve_positive_velocity: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: if forward_ctx.guidance_parallel_ready: - return self.predict_parallel_guidance(pipeline, plan, index, timestep, state, forward_ctx, denoise_ctx) + return self.predict_parallel_guidance( + pipeline, + plan, + index, + timestep, + state, + forward_ctx, + denoise_ctx, + preserve_positive_velocity=preserve_positive_velocity, + ) prompt = forward_ctx.prompt_context video_contexts: list[torch.Tensor] = [] @@ -526,12 +549,14 @@ class LTXGuidanceExecutor: pipeline.scheduler.sigmas[index], denoise_ctx, ), + preserve_positive_velocity=preserve_positive_velocity, ), self._guide_modality( state.audio, audio_splits, forward_ctx.audio_scheduler.sigmas[index], plan.spec.audio, + preserve_positive_velocity=preserve_positive_velocity, rescale_token_count=forward_ctx.original_audio_num_frames, ), ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py index 11929d331..1b25a8e01 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py @@ -20,6 +20,11 @@ class LTXAVState: audio: torch.Tensor +def official_video_token_layout(latents: torch.Tensor) -> torch.Tensor: + """Match the token-major view produced by the official LTX patchifier.""" + return latents.transpose(1, 2).contiguous().transpose(1, 2) + + def pack_latents( latents: torch.Tensor, patch_size: int = 1, @@ -117,9 +122,23 @@ def create_noised_state( device=latents.device, dtype=latents.dtype, ) - if isinstance(noise_scale, torch.Tensor): - noise_scale = noise_scale.float() - return torch.lerp(latents.float(), noise.float(), noise_scale).to(latents.dtype) + # Official GaussianNoisier blends in fp32, then materializes the next + # latent state in the model dtype. Batched request weights may arrive in + # the model dtype, so promote tensor weights alongside both endpoints. + lerp_weight = noise_scale.float() if isinstance(noise_scale, torch.Tensor) else noise_scale + return torch.lerp(latents.float(), noise.float(), lerp_weight).to(latents.dtype) + + +def create_conditioned_noised_state( + latents: torch.Tensor, + clean_latents: torch.Tensor, + denoise_mask: torch.Tensor, + noise_scale: float | torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None = None, +) -> torch.Tensor: + """Match the official conditioned Gaussian noisier operation order.""" + noised_latents = create_noised_state(latents, noise_scale, generator) + return torch.lerp(clean_latents.float(), noised_latents.float(), denoise_mask.float()).to(latents.dtype) def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor: @@ -138,7 +157,7 @@ def unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: def clear_audio_padding(latents: torch.Tensor, num_frames: int) -> torch.Tensor: - """Keep Omni's SP-only audio padding outside the sampler state.""" + """Keep SP-only audio padding outside the logical sampler state.""" if not 0 < num_frames <= latents.shape[1]: raise ValueError(f"Audio frame count must be in [1, {latents.shape[1]}], got {num_frames}.") if num_frames == latents.shape[1]: @@ -196,7 +215,8 @@ def prepare_video_latents( ) if latents.ndim != 3: raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") - return create_noised_state(latents, noise_scale, generator).to(device=device, dtype=dtype) + latents = create_noised_state(latents, noise_scale, generator).to(device=device, dtype=dtype) + return official_video_token_layout(latents) num_frames, height, width = resolve_video_latent_shape( height, @@ -218,11 +238,7 @@ def prepare_video_latents( f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # Official LTX samples the 5D video latent first and then patchifies it, - # leaving token-major storage behind the [B, tokens, channels] view. The - # logical values are unchanged, but preserving that layout is important: - # CUDA selects a different bf16 GEMM path for a contiguous token tensor. - return latents.transpose(1, 2).contiguous().transpose(1, 2) + return official_video_token_layout(latents) def prepare_audio_latents( diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py index bb42bc0e9..0caf93f64 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py @@ -3,7 +3,7 @@ """Declarative execution recipes for the LTX model family.""" -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Literal from .ltx2_guidance import LTXGuidanceSpec, LTXModalityGuidance @@ -22,6 +22,11 @@ LTX_DEFAULT_NEGATIVE_PROMPT = ( "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." ) +LTX25_DEFAULT_NEGATIVE_PROMPT = ( + "has_subtitles, has_blurbox, transition from black, transition to black, speech_ending_short, " + + LTX_DEFAULT_NEGATIVE_PROMPT +) + LTX_DISTILLED_SIGMAS = (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0) LTX_STAGE_2_DISTILLED_SIGMAS = (0.909375, 0.725, 0.421875, 0.0) LTX_DISTILLED_ADAPTER_SLOT = "ltx_distilled" @@ -38,6 +43,7 @@ class LTXPhaseRecipe: noise_scale: float = 0.0 input_transform: Literal["initial", "spatial_upsample"] = "initial" adapter_slot: str | None = None + sampler: Literal["euler", "euler_ancestral"] = "euler" allow_guidance_override: bool = True use_official_sigma_schedule: bool = True @@ -46,6 +52,8 @@ class LTXPhaseRecipe: raise ValueError("LTX phase spatial_downscale must be positive.") if self.sigmas is not None and len(self.sigmas) < 2: raise ValueError("An explicit LTX sigma schedule must contain at least two denoise sigmas.") + if self.sampler not in ("euler", "euler_ancestral"): + raise ValueError(f"Unsupported LTX sampler: {self.sampler!r}.") @property def num_inference_steps(self) -> int | None: @@ -66,9 +74,11 @@ class LTXPipelineRecipe: video_output_phase: int = -1 audio_output_phase: int = -1 allow_request_sigmas: bool = True + allow_request_phase_sigmas: bool = False allow_request_latents: bool = True allow_negative_prompt: bool = True fixed_num_inference_steps: bool = False + supports_cache_dit: bool = False def __post_init__(self) -> None: if not self.phases: @@ -109,13 +119,29 @@ def _official_guidance(stg_block: int) -> LTXGuidanceSpec: LTX2_ONE_STAGE_RECIPE = LTXPipelineRecipe( + supports_cache_dit=True, phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(29)),), ) LTX23_ONE_STAGE_RECIPE = LTXPipelineRecipe( + supports_cache_dit=True, num_inference_steps=30, phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(28)),), ) +LTX25_FULL_RECIPE = LTXPipelineRecipe( + height=544, + width=960, + num_inference_steps=30, + negative_prompt=LTX25_DEFAULT_NEGATIVE_PROMPT, + phases=( + LTXPhaseRecipe( + name="generate", + guidance=_official_guidance(28), + noise_scale=1.0, + ), + ), +) LTX_POSITIVE_ONLY_RECIPE = LTXPipelineRecipe( + supports_cache_dit=True, phases=( LTXPhaseRecipe( name="generate", @@ -159,6 +185,41 @@ LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe( fixed_num_inference_steps=True, ) +LTX25_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe( + height=1088, + width=1920, + num_inference_steps=len(LTX_DISTILLED_SIGMAS) - 1, + negative_prompt="", + phases=( + LTXPhaseRecipe( + name="generate_lowres", + guidance=LTXGuidanceSpec.positive_only(), + spatial_downscale=2, + sigmas=LTX_DISTILLED_SIGMAS, + noise_scale=1.0, + allow_guidance_override=False, + use_official_sigma_schedule=False, + sampler="euler_ancestral", + ), + LTXPhaseRecipe( + name="refine", + guidance=LTXGuidanceSpec.positive_only(), + sigmas=LTX_STAGE_2_DISTILLED_SIGMAS, + noise_scale=LTX_STAGE_2_DISTILLED_SIGMAS[0], + input_transform="spatial_upsample", + allow_guidance_override=False, + use_official_sigma_schedule=False, + ), + ), + video_output_phase=1, + audio_output_phase=1, + allow_request_sigmas=False, + allow_request_phase_sigmas=True, + allow_request_latents=False, + allow_negative_prompt=False, + fixed_num_inference_steps=True, +) + # LTX-2.3 full-distilled uses the same official fixed 8 + 3 sigma schedules # and request contract. Its distinct component profile selects the 22B merged # distilled Transformer, BWE vocoder, and matching spatial upsampler; no @@ -166,6 +227,37 @@ LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe( LTX23_DISTILLED_TWO_STAGE_RECIPE = LTX2_DISTILLED_TWO_STAGE_RECIPE +def _distilled_one_stage_recipe( + two_stage_recipe: LTXPipelineRecipe, + *, + supports_cache_dit: bool = True, +) -> LTXPipelineRecipe: + """Run the merged-distilled checkpoint at its native Stage-1 resolution.""" + generate_phase = replace( + two_stage_recipe.phases[0], + name="generate", + spatial_downscale=1, + ) + return replace( + two_stage_recipe, + height=two_stage_recipe.height // 2, + width=two_stage_recipe.width // 2, + phases=(generate_phase,), + video_output_phase=0, + audio_output_phase=0, + allow_request_phase_sigmas=False, + supports_cache_dit=supports_cache_dit, + ) + + +LTX2_DISTILLED_ONE_STAGE_RECIPE = _distilled_one_stage_recipe(LTX2_DISTILLED_TWO_STAGE_RECIPE) +LTX23_DISTILLED_ONE_STAGE_RECIPE = _distilled_one_stage_recipe(LTX23_DISTILLED_TWO_STAGE_RECIPE) +LTX25_DISTILLED_ONE_STAGE_RECIPE = _distilled_one_stage_recipe( + LTX25_DISTILLED_TWO_STAGE_RECIPE, + supports_cache_dit=False, +) + + def _official_two_stage_recipe(one_stage_recipe: LTXPipelineRecipe) -> LTXPipelineRecipe: return LTXPipelineRecipe( height=one_stage_recipe.height * 2, @@ -197,21 +289,29 @@ def _official_two_stage_recipe(one_stage_recipe: LTXPipelineRecipe) -> LTXPipeli # discards its audio result. Decode the full-context Stage-1 audio. audio_output_phase=0, allow_request_sigmas=False, + allow_request_phase_sigmas=True, allow_request_latents=False, ) LTX2_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX2_ONE_STAGE_RECIPE) LTX23_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX23_ONE_STAGE_RECIPE) +LTX25_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX25_FULL_RECIPE) _PIPELINE_RECIPES: dict[tuple[str, str], LTXPipelineRecipe] = { ("one_stage", "2"): LTX2_ONE_STAGE_RECIPE, ("one_stage", "2.3"): LTX23_ONE_STAGE_RECIPE, + ("one_stage", "2.5"): LTX25_FULL_RECIPE, ("two_stage", "2"): LTX2_TWO_STAGE_RECIPE, ("two_stage", "2.3"): LTX23_TWO_STAGE_RECIPE, + ("two_stage", "2.5"): LTX25_TWO_STAGE_RECIPE, + ("distilled_one_stage", "2"): LTX2_DISTILLED_ONE_STAGE_RECIPE, + ("distilled_one_stage", "2.3"): LTX23_DISTILLED_ONE_STAGE_RECIPE, + ("distilled_one_stage", "2.5"): LTX25_DISTILLED_ONE_STAGE_RECIPE, ("distilled_two_stage", "2"): LTX2_DISTILLED_TWO_STAGE_RECIPE, ("distilled_two_stage", "2.3"): LTX23_DISTILLED_TWO_STAGE_RECIPE, + ("distilled_two_stage", "2.5"): LTX25_DISTILLED_TWO_STAGE_RECIPE, ("dmd2", "2"): LTX_POSITIVE_ONLY_RECIPE, ("dmd2", "2.3"): LTX_POSITIVE_ONLY_RECIPE, } diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py index 91d62943a..35e0cef0b 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -43,6 +43,7 @@ class LTXRequestInputs: output_type: str max_sequence_length: int audio_latents_normalized: bool = False + image_crf: int | None = None @property def guidance_scale(self) -> float: @@ -86,6 +87,7 @@ def validate_pipeline_request( vae_temporal_compression_ratio: int, pipeline_name: str, request_sigmas: list[float] | None = None, + request_phase_sigmas: tuple[list[float] | None, ...] | None = None, ) -> None: """Validate the request capabilities declared by an execution recipe.""" alignment = vae_spatial_compression_ratio * pipeline_recipe.max_spatial_downscale @@ -108,9 +110,28 @@ def validate_pipeline_request( if not pipeline_recipe.allow_request_sigmas and request_sigmas is not None: raise ValueError(f"{pipeline_name} uses fixed phase sigma schedules.") + if request_sigmas is not None and len(request_sigmas) < 2: + raise ValueError("A custom LTX sigma schedule must contain at least two values.") - if pipeline_recipe.fixed_num_inference_steps and ( - request_inputs.num_inference_steps != pipeline_recipe.num_inference_steps + has_phase_sigmas = request_phase_sigmas is not None and any(item is not None for item in request_phase_sigmas) + if request_sigmas is not None and has_phase_sigmas: + raise ValueError("Use either `sigmas` or per-stage sigma schedules, not both.") + if has_phase_sigmas and not pipeline_recipe.allow_request_phase_sigmas: + raise ValueError(f"{pipeline_name} does not accept per-stage sigma schedules.") + if request_phase_sigmas is not None and len(request_phase_sigmas) != len(pipeline_recipe.phases): + raise ValueError( + f"{pipeline_name} has {len(pipeline_recipe.phases)} phases but received " + f"{len(request_phase_sigmas)} per-stage sigma schedules." + ) + for index, phase_sigmas in enumerate(request_phase_sigmas or ()): + if phase_sigmas is not None and len(phase_sigmas) < 2: + raise ValueError(f"`stage_{index + 1}_sigmas` must contain at least two values.") + + if ( + request_sigmas is None + and (request_phase_sigmas is None or request_phase_sigmas[0] is None) + and pipeline_recipe.fixed_num_inference_steps + and (request_inputs.num_inference_steps != pipeline_recipe.num_inference_steps) ): raise ValueError( f"{pipeline_name} uses {pipeline_recipe.num_inference_steps} fixed Stage 1 denoise steps; " @@ -165,6 +186,16 @@ def _get_extra_arg(sampling: Any, *names: str) -> Any: return None +def _normalize_image_crf(value: Any) -> int | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError("`image_crf` must be an integer from 0 through 51.") + if not 0 <= value <= 51: + raise ValueError("`image_crf` must be an integer from 0 through 51.") + return value + + def _normalize_stg_blocks(value: Any, default: tuple[int, ...]) -> tuple[int, ...]: if value is None: return default @@ -274,6 +305,7 @@ class LTXRequestMixin: req: DiffusionRequestBatch, *, image: Any | None = None, + image_crf: int | None = None, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, height: int | None = None, @@ -282,6 +314,8 @@ class LTXRequestMixin: frame_rate: float | None = None, num_inference_steps: int | None = None, sigmas: list[float] | None = None, + stage_1_sigmas: list[float] | None = None, + stage_2_sigmas: list[float] | None = None, timesteps: list[int] | None = None, guidance_scale: float | None = None, guidance_rescale: float | None = None, @@ -315,6 +349,7 @@ class LTXRequestMixin: return self._forward_request( req, image=image, + image_crf=image_crf, prompt=prompt, negative_prompt=negative_prompt, height=height, @@ -323,6 +358,8 @@ class LTXRequestMixin: frame_rate=frame_rate, num_inference_steps=num_inference_steps, sigmas=sigmas, + stage_1_sigmas=stage_1_sigmas, + stage_2_sigmas=stage_2_sigmas, guidance_scale=guidance_scale, num_videos_per_prompt=num_videos_per_prompt, generator=generator, @@ -343,12 +380,30 @@ class LTXRequestMixin: req: DiffusionRequestBatch, fallback: list[float] | None, ) -> list[float] | None: - request_sigmas = [sampling.sigmas for sampling in req.sampling_params_list] + request_sigmas = [ + sampling.sigmas if sampling.sigmas is not None else _get_extra_arg(sampling, "sigmas") + for sampling in req.sampling_params_list + ] first = request_sigmas[0] if request_sigmas else None if any(item != first for item in request_sigmas[1:]): raise ValueError("Batched LTX requests must use identical custom sigmas.") return first if first is not None else fallback + @staticmethod + def _resolve_request_phase_sigmas( + req: DiffusionRequestBatch, + stage_1_fallback: list[float] | None, + stage_2_fallback: list[float] | None, + ) -> tuple[list[float] | None, list[float] | None] | None: + resolved: list[list[float] | None] = [] + for name, fallback in (("stage_1_sigmas", stage_1_fallback), ("stage_2_sigmas", stage_2_fallback)): + values = [_get_extra_arg(sampling, name) for sampling in req.sampling_params_list] + first = values[0] if values else None + if any(item != first for item in values[1:]): + raise ValueError(f"Batched LTX requests must use identical `{name}` values.") + resolved.append(first if first is not None else fallback) + return tuple(resolved) if any(item is not None for item in resolved) else None + def check_inputs( self, prompt: str | list[str] | None, @@ -408,6 +463,7 @@ class LTXRequestMixin: decode_noise_scale: float | list[float] | None, output_type: str, max_sequence_length: int | None, + image_crf: int | None = None, ) -> LTXRequestInputs: sampling_params_list = req.sampling_params_list for sampling in sampling_params_list: @@ -421,6 +477,12 @@ class LTXRequestMixin: if _get_extra_arg(sampling, "flow_shift") is not None: raise ValueError("LTX does not support `flow_shift`; use `sigmas` for a custom schedule.") + request_image_crfs = [_get_extra_arg(item, "image_crf") for item in sampling_params_list] + first_image_crf = request_image_crfs[0] if request_image_crfs else None + if any(item != first_image_crf for item in request_image_crfs[1:]): + raise ValueError("Batched LTX requests must use identical `image_crf` values.") + image_crf = _normalize_image_crf(first_image_crf if first_image_crf is not None else image_crf) + sampling = sampling_params_list[0] is_dummy_run = req.is_dummy_run() prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt @@ -551,4 +613,5 @@ class LTXRequestMixin: decode_noise_scale=decode_noise_scale, output_type=output_type, max_sequence_length=int(max_sequence_length), + image_crf=image_crf, ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index cb4130259..f560f678b 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -144,20 +144,19 @@ class LTXRuntime( f"{self.__class__.__name__} does not support ulysses_mode='advanced_uaa'. " "Use the default ulysses_mode='strict' for LTX sequence parallelism." ) - self.model_version = detect_ltx_model_version(od_config.model) + self.model_version = detect_ltx_model_version(od_config.model, revision=getattr(od_config, "revision", None)) self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version) self.pipeline_recipe = resolve_ltx_pipeline_recipe(self.pipeline_kind, self.model_version) - if getattr(od_config, "cache_backend", "none") == "cache_dit" and len(self.pipeline_recipe.phases) > 1: + if getattr(od_config, "cache_backend", "none") == "cache_dit" and not self.pipeline_recipe.supports_cache_dit: raise ValueError( f"{self.__class__.__name__} does not support cache_backend='cache_dit'. " - "Cache-DiT currently supports only one-stage LTX recipes; multi-stage recipes require " - "phase-aware cache enable and refresh." + "Cache-DiT is not qualified for this LTX recipe." ) self._dit_modules = list(self.component_profile.dit_modules) self._encoder_modules = list(self.component_profile.encoder_modules) self._vae_modules = list(self.component_profile.vae_modules) self._resident_modules = list(self.component_profile.resident_modules) - if self.model_version == "2.3": + if self.model_version in ("2.3", "2.5"): self.preserve_sp_padded_audio_duration = True self.reports_stage_durations = True super().__init__() @@ -173,6 +172,7 @@ class LTXRuntime( req: DiffusionRequestBatch, *, image: Any | None = None, + image_crf: int | None = None, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, height: int | None = None, @@ -181,6 +181,8 @@ class LTXRuntime( frame_rate: float | None = None, num_inference_steps: int | None = None, sigmas: list[float] | None = None, + stage_1_sigmas: list[float] | None = None, + stage_2_sigmas: list[float] | None = None, guidance_scale: float | None = None, num_videos_per_prompt: int | None = 1, generator: torch.Generator | list[torch.Generator] | None = None, @@ -217,11 +219,13 @@ class LTXRuntime( decode_noise_scale=decode_noise_scale, output_type=output_type, max_sequence_length=max_sequence_length, + image_crf=image_crf, ) image = self._resolve_request_image(req, image, request_inputs) if image is not None and not self.support_image_input: raise ValueError(f"{self.__class__.__name__} does not support `image` input.") request_sigmas = self._resolve_request_sigmas(req, sigmas) + request_phase_sigmas = self._resolve_request_phase_sigmas(req, stage_1_sigmas, stage_2_sigmas) validate_pipeline_request( request_inputs, pipeline_recipe=self.pipeline_recipe, @@ -229,6 +233,7 @@ class LTXRuntime( vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, pipeline_name=self.__class__.__name__, request_sigmas=request_sigmas, + request_phase_sigmas=request_phase_sigmas, ) phase_adapter = getattr(self, "_phase_adapter", None) if phase_adapter is not None and any( @@ -237,7 +242,13 @@ class LTXRuntime( raise ValueError( f"{self.__class__.__name__} cannot compose a request LoRA with its internal phase adapter." ) - return self._run_recipe(req, request_inputs, request_sigmas=request_sigmas, image=image) + return self._run_recipe( + req, + request_inputs, + request_sigmas=request_sigmas, + request_phase_sigmas=request_phase_sigmas, + image=image, + ) def _run_recipe( self, @@ -245,24 +256,39 @@ class LTXRuntime( request_inputs: LTXRequestInputs, *, request_sigmas: list[float] | None, + request_phase_sigmas: tuple[list[float] | None, ...] | None = None, image: Any | None = None, ) -> DiffusionOutput | list[DiffusionOutput]: """Execute one- and multi-phase recipes through the same control flow.""" phase_results: list[LTXPhaseResult] = [] prompt_context = None - - for phase_recipe in self.pipeline_recipe.phases: + for phase_index, phase_recipe in enumerate(self.pipeline_recipe.phases): + override_sigmas = None if request_phase_sigmas is None else request_phase_sigmas[phase_index] + phase_sigmas = ( + override_sigmas + if override_sigmas is not None + else ( + request_sigmas + if request_sigmas is not None + else (list(phase_recipe.sigmas) if phase_recipe.sigmas is not None else None) + ) + ) self._enter_phase(phase_recipe) phase_inputs = self._build_phase_inputs( request_inputs, phase_recipe, phase_results[-1] if phase_results else None, ) + if phase_sigmas is not None: + phase_inputs = replace(phase_inputs, num_inference_steps=len(phase_sigmas) - 1) + noise_scale = phase_recipe.noise_scale + if override_sigmas is not None and phase_recipe.input_transform == "spatial_upsample": + noise_scale = float(override_sigmas[0]) phase_result = self.run_phase( req, phase_inputs, - noise_scale=phase_recipe.noise_scale, - sigmas=list(phase_recipe.sigmas) if phase_recipe.sigmas is not None else request_sigmas, + noise_scale=noise_scale, + sigmas=phase_sigmas, timesteps=None, attention_kwargs=None, phase_recipe=phase_recipe, @@ -282,6 +308,7 @@ class LTXRuntime( def _enter_phase(self, phase: LTXPhaseRecipe) -> None: self._active_phase_name = phase.name + self._guidance_plan = LTXGuidancePlan.build(phase.guidance) phase_adapter = getattr(self, "_phase_adapter", None) if phase_adapter is None: if phase.adapter_slot is not None: @@ -645,7 +672,12 @@ class LTXRuntime( audio_token_count: int, ) -> dict[str, torch.Tensor]: del forward_ctx, denoise_ctx - return self.guidance_executor.timestep_kwargs(ts, video_token_count, audio_token_count) + return self.guidance_executor.timestep_kwargs( + ts, + video_token_count, + audio_token_count, + expand_for_sequence_parallel=True, + ) def _video_guidance_model_sigma( self, @@ -700,6 +732,7 @@ class LTXRuntime( state, forward_ctx, denoise_ctx, + preserve_positive_velocity=forward_ctx.sampler == "euler_ancestral", ) def _denoise_step( diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index a59655cd3..0b0e7fc08 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -55,6 +55,18 @@ logger = init_logger(__name__) _RMSNORM_INIT_PARAMS = inspect.signature(RMSNorm.__init__).parameters +def apply_keyframes_absolute_embedding( + hidden_states: torch.Tensor, + keyframes_mask: torch.Tensor | None, + embedding: torch.Tensor | None, +) -> torch.Tensor: + """Add the official learned marker to first-frame/keyframe tokens.""" + if keyframes_mask is None or embedding is None: + return hidden_states + mask = (keyframes_mask > 0).to(dtype=hidden_states.dtype) + return hidden_states + mask * embedding.to(device=hidden_states.device, dtype=hidden_states.dtype) + + def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) -> nn.Module: """Bridge diffusers' RMSNorm API onto vLLM's `has_weight` variant.""" if not elementwise_affine: @@ -82,15 +94,10 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) -> def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: cos, sin = freqs - if cos.dtype != x.dtype: - cos = cos.to(dtype=x.dtype) - if sin.dtype != x.dtype: - sin = sin.to(dtype=x.dtype) # Concrete pair count instead of -1 keeps SDPA shape static under torch.compile. x_real, x_imag = x.unflatten(2, (x.shape[2] // 2, 2)).unbind(-1) # [B, S, C // 2] x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) - out = x * cos + x_rotated * sin - return out + return x * cos + x_rotated * sin def apply_split_rotary_emb( @@ -101,12 +108,6 @@ def apply_split_rotary_emb( ) -> torch.Tensor: # `head_dim` is plumbed in (not inferred via `-1`) so SDPA shape stays static under torch.compile. cos, sin = freqs - if cos.dtype != x.dtype: - cos = cos.to(dtype=x.dtype) - if sin.dtype != x.dtype: - sin = sin.to(dtype=x.dtype) - - x_dtype = x.dtype needs_reshape = False if x.ndim != 4 and cos.ndim == 4: # cos is (#b, h, t, r) -> reshape x to (b, h, t, dim_per_head) @@ -142,7 +143,6 @@ def apply_split_rotary_emb( if needs_reshape: out = out.swapaxes(1, 2).reshape(b, t, h * head_dim) - out = out.to(dtype=x_dtype) return out @@ -269,6 +269,7 @@ class LTX2FeedForward(nn.Module): RowParallelLinear( inner_dim, dim_out, + bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config, @@ -517,8 +518,8 @@ class LTX2AudioVideoAttnProcessor: is_self_attention=is_self_attention, ) - query = attn.norm_q(query) - key = attn.norm_k(key) + query = attn.norm_q(query).to(dtype=value.dtype) + key = attn.norm_k(key).to(dtype=value.dtype) if query_rotary_emb is not None: query_rotary_emb = self._slice_rope_for_tp(query_rotary_emb, attn) @@ -537,6 +538,10 @@ class LTX2AudioVideoAttnProcessor: head_dim=attn.head_dim, ) + # Keep FP32 RoPE computation while satisfying attention dtype contracts. + query = query.to(dtype=value.dtype) + key = key.to(dtype=value.dtype) + query = query.unflatten(2, (attn.heads, attn.head_dim)) key = key.unflatten(2, (attn.heads, attn.head_dim)) value_heads = value.unflatten(2, (attn.heads, attn.head_dim)) @@ -846,6 +851,8 @@ class LTX2VideoTransformerBlock(nn.Module): elementwise_affine: bool = False, rope_type: str = "interleaved", perturbed_attn: bool = False, + ff_bias: bool = True, + audio_ff_bias: bool = True, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ): @@ -965,6 +972,7 @@ class LTX2VideoTransformerBlock(nn.Module): self.ff = LTX2FeedForward( dim, activation_fn=activation_fn, + bias=ff_bias, quant_config=quant_config, prefix=f"{prefix}.ff" if prefix else "ff", ) @@ -973,6 +981,7 @@ class LTX2VideoTransformerBlock(nn.Module): self.audio_ff = LTX2FeedForward( audio_dim, activation_fn=activation_fn, + bias=audio_ff_bias, quant_config=quant_config, prefix=f"{prefix}.audio_ff" if prefix else "audio_ff", ) @@ -1078,21 +1087,34 @@ class LTX2VideoTransformerBlock(nn.Module): audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa # 2. Video and Audio Cross-Attention with text embeddings (Q: Video/Audio; K,V: Text) - # LTX-2.3: compute prompt modulation params for K/V - if self.cross_attn_adaln and temb_prompt is not None: - video_prompt_ada_params = self.get_mod_params(self.prompt_scale_shift_table, temb_prompt, batch_size) - shift_text_kv, scale_text_kv = video_prompt_ada_params + if self.cross_attn_adaln: + if temb_prompt is not None: + shift_text_kv, scale_text_kv = self.get_mod_params( + self.prompt_scale_shift_table, temb_prompt, batch_size + ) + else: + shift_text_kv, scale_text_kv = ( + self.prompt_scale_shift_table[None, None] + .to(device=hidden_states.device, dtype=hidden_states.dtype) + .unbind(dim=2) + ) - audio_prompt_ada_params = self.get_mod_params( - self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size - ) - audio_shift_text_kv, audio_scale_text_kv = audio_prompt_ada_params + if temb_prompt_audio is not None: + audio_shift_text_kv, audio_scale_text_kv = self.get_mod_params( + self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size + ) + else: + audio_shift_text_kv, audio_scale_text_kv = ( + self.audio_prompt_scale_shift_table[None, None] + .to(device=audio_hidden_states.device, dtype=audio_hidden_states.dtype) + .unbind(dim=2) + ) # 2.1. Video-Text Cross-Attention norm_hidden_states = self.norm2(hidden_states) if self.video_cross_attn_adaln: norm_hidden_states = norm_hidden_states * (1 + scale_text_q) + shift_text_q - if self.cross_attn_adaln and temb_prompt is not None: + if self.cross_attn_adaln: encoder_hidden_states = encoder_hidden_states * (1 + scale_text_kv) + shift_text_kv attn_hidden_states = self.attn2( @@ -1109,7 +1131,7 @@ class LTX2VideoTransformerBlock(nn.Module): norm_audio_hidden_states = self.audio_norm2(audio_hidden_states) if self.audio_cross_attn_adaln: norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_text_q) + audio_shift_text_q - if self.cross_attn_adaln and temb_prompt is not None: + if self.cross_attn_adaln: audio_encoder_hidden_states = audio_encoder_hidden_states * (1 + audio_scale_text_kv) + audio_shift_text_kv attn_audio_hidden_states = self.audio_attn2( @@ -1508,7 +1530,10 @@ class LTX2VideoTransformer3DModel(nn.Module): block_forward_patterns={ "transformer_blocks": ForwardPattern.Pattern_0, }, - has_separate_cfg=True, + # LTX fuses all guidance passes into the batch dimension and invokes + # the transformer once per denoise step. Cache-DiT must therefore + # advance its step counter for every transformer forward. + has_separate_cfg=False, check_forward_pattern=False, ) @@ -1546,6 +1571,7 @@ class LTX2VideoTransformer3DModel(nn.Module): # Shard video/audio latents across sequence "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), "audio_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), + "keyframes_mask": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), # Shard prompt embeds across sequence "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), "audio_encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), @@ -1620,6 +1646,10 @@ class LTX2VideoTransformer3DModel(nn.Module): cross_attn_mod: bool = False, audio_gated_attn: bool = False, audio_cross_attn_mod: bool = False, + ff_bias: bool = True, + audio_ff_bias: bool = True, + use_prompt_adaln_single: bool = True, + use_keyframes_abs_pos_embedding: bool = False, quant_config: "QuantizationConfig | None" = None, ) -> None: super().__init__() @@ -1666,12 +1696,19 @@ class LTX2VideoTransformer3DModel(nn.Module): timestep_scale_multiplier=timestep_scale_multiplier, cross_attn_timestep_scale_multiplier=cross_attn_timestep_scale_multiplier, rope_type=rope_type, + ff_bias=ff_bias, + audio_ff_bias=audio_ff_bias, + use_prompt_adaln_single=use_prompt_adaln_single, + use_keyframes_abs_pos_embedding=use_keyframes_abs_pos_embedding, ) # 1. Patchification input projections self.proj_in = nn.Linear(in_channels, inner_dim) self.audio_proj_in = nn.Linear(audio_in_channels, audio_inner_dim) + if use_keyframes_abs_pos_embedding: + self.keyframes_abs_pos_embedding = nn.Parameter(torch.zeros(1, inner_dim)) + # 2. Prompt embeddings # LTX-2 (use_prompt_embeddings=True): caption projection in the transformer # LTX-2.3 (use_prompt_embeddings=False): caption projection in the connectors @@ -1695,7 +1732,7 @@ class LTX2VideoTransformer3DModel(nn.Module): # 3.3. LTX-2.3: Prompt modulation from sigma self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod - if self.prompt_modulation: + if self.prompt_modulation and use_prompt_adaln_single: self.prompt_adaln = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=2, use_additional_conditions=False) self.audio_prompt_adaln = LTX2AdaLayerNormSingle( audio_inner_dim, num_mod_params=2, use_additional_conditions=False @@ -1815,6 +1852,8 @@ class LTX2VideoTransformer3DModel(nn.Module): elementwise_affine=norm_elementwise_affine, rope_type=rope_type, perturbed_attn=perturbed_attn, + ff_bias=ff_bias, + audio_ff_bias=audio_ff_bias, quant_config=quant_config, prefix=f"transformer_blocks.{layer_idx}", ) @@ -1852,6 +1891,7 @@ class LTX2VideoTransformer3DModel(nn.Module): audio_encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, audio_timestep: torch.LongTensor | None = None, + keyframes_mask: torch.Tensor | None = None, sigma: torch.Tensor | None = None, audio_sigma: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, @@ -1952,9 +1992,7 @@ class LTX2VideoTransformer3DModel(nn.Module): video_rotary_emb = self.rope(video_coords, device=hidden_states.device, out_dtype=hidden_states.dtype) audio_rotary_emb = self.audio_rope( - audio_coords, - device=audio_hidden_states.device, - out_dtype=audio_hidden_states.dtype, + audio_coords, device=audio_hidden_states.device, out_dtype=audio_hidden_states.dtype ) video_cross_attn_rotary_emb = self.cross_attn_rope( @@ -1970,6 +2008,11 @@ class LTX2VideoTransformer3DModel(nn.Module): # 2. Patchify input projections hidden_states = self.proj_in(hidden_states) + hidden_states = apply_keyframes_absolute_embedding( + hidden_states, + keyframes_mask, + getattr(self, "keyframes_abs_pos_embedding", None), + ) audio_hidden_states = self.audio_proj_in(audio_hidden_states) # 3. Prepare timestep embeddings and modulation parameters @@ -2028,7 +2071,7 @@ class LTX2VideoTransformer3DModel(nn.Module): audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1]) # 3.3. LTX-2.3: Compute prompt modulation from sigma - if self.prompt_modulation and sigma is not None: + if self.prompt_modulation and self.config.use_prompt_adaln_single and sigma is not None: temb_prompt, _ = self.prompt_adaln(sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype) temb_prompt_audio, _ = self.audio_prompt_adaln( audio_sigma.flatten(), batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index ff66af5e1..ff2ab0457 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -10,15 +10,12 @@ from typing import ClassVar from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin -from .ltx2_components import LTX2_COMPONENT_PROFILE +from .ltx2_components import LTX2_COMPONENT_PROFILE, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE from .ltx2_components import ( get_ltx2_post_process_func as get_ltx2_post_process_func, # noqa: F401 ) from .ltx2_conditioning import LTXI2VConditioningMixin -from .ltx2_recipes import ( - LTX2_ONE_STAGE_RECIPE, - LTX_POSITIVE_ONLY_RECIPE, -) +from .ltx2_recipes import LTX2_DISTILLED_ONE_STAGE_RECIPE, LTX2_ONE_STAGE_RECIPE, LTX_POSITIVE_ONLY_RECIPE from .ltx2_runtime import LTXRuntime @@ -36,6 +33,14 @@ class LTX2Pipeline(LTXI2VConditioningMixin, LTXRuntime): supports_request_batch = True +class LTX2DistilledOneStagePipeline(LTX2Pipeline): + """Merged-distilled checkpoint at its native one-stage resolution.""" + + pipeline_kind = "distilled_one_stage" + component_profile = LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE + pipeline_recipe = LTX2_DISTILLED_ONE_STAGE_RECIPE + + class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline): """LTX2 T2V entry for FastGen DMD2-distilled models.""" diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py index 3cb6b84f1..5d50baf42 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py @@ -40,9 +40,13 @@ class LTX2TwoStagePipeline(_LTX2TwoStageBase): pipeline_recipe = LTX2_TWO_STAGE_RECIPE -class LTX2DistilledPipeline(_LTX2TwoStageBase): - """Unified LTX-2/LTX-2.3 full-distilled two-stage T2V/I2V entry.""" +class LTX2DistilledTwoStagePipeline(_LTX2TwoStageBase): + """Merged-distilled checkpoint with latent upsampling and refinement.""" pipeline_kind = "distilled_two_stage" component_profile = LTX2_DISTILLED_COMPONENT_PROFILE pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE + + +# Backwards-compatible name retained for existing LTX-2/LTX-2.3 recipes. +LTX2DistilledPipeline = LTX2DistilledTwoStagePipeline diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 0e814886b..9bac50323 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -76,6 +76,16 @@ _DIFFUSION_MODELS = { "pipeline_ltx2_two_stage", "LTX2TwoStagePipeline", ), + "LTX2DistilledOneStagePipeline": ( + "ltx2", + "pipeline_ltx2", + "LTX2DistilledOneStagePipeline", + ), + "LTX2DistilledTwoStagePipeline": ( + "ltx2", + "pipeline_ltx2_two_stage", + "LTX2DistilledTwoStagePipeline", + ), "LTX2DistilledPipeline": ( "ltx2", "pipeline_ltx2_two_stage", @@ -522,6 +532,8 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "WanVACEPipeline": "get_wan22_vace_post_process_func", "LTX2Pipeline": "get_ltx2_post_process_func", "LTX2TwoStagePipeline": "get_ltx2_post_process_func", + "LTX2DistilledOneStagePipeline": "get_ltx2_post_process_func", + "LTX2DistilledTwoStagePipeline": "get_ltx2_post_process_func", "LTX2DistilledPipeline": "get_ltx2_post_process_func", "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func", "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func", diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 9d511ba6d..2628c14f8 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -7,6 +7,7 @@ import copy import time from collections.abc import Mapping from dataclasses import dataclass +from functools import cached_property from http import HTTPStatus from typing import Any, cast @@ -30,6 +31,7 @@ from vllm_omni.entrypoints.openai.stage_params import ( from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.model_extras import should_preserve_reference_image_size from vllm_omni.outputs.output_metadata import ( DiffusionMetadataMapping, DiffusionMultimodalOutput, @@ -100,6 +102,21 @@ class OmniOpenAIServingVideo: if self._stage_configs is None and stage_configs is not None: self._stage_configs = stage_configs + @cached_property + def preserves_reference_image_size(self) -> bool: + """Return whether the active pipeline owns reference-image resizing.""" + get_od_config = getattr(self._engine_client, "get_diffusion_od_config", None) + od_config = get_od_config() if callable(get_od_config) else getattr(self._engine_client, "od_config", None) + model_class_name = None if od_config is None else getattr(od_config, "model_class_name", None) + model = getattr(od_config, "model", None) if od_config is not None else None + model = model or self.model_name + revision = getattr(od_config, "revision", None) if od_config is not None else None + return should_preserve_reference_image_size( + model_class_name, + model=None if model is None else str(model), + revision=revision, + ) + @property def supports_mixed_reference_inputs(self) -> bool: """Return whether the configured diffusion model accepts mixed refs.""" @@ -154,7 +171,12 @@ class OmniOpenAIServingVideo: provided_fields = request.model_fields_set fps_provided = self._request_fps_provided(request) vp = request.resolve_video_params() - if input_image is not None and vp.width is not None and vp.height is not None: + if ( + input_image is not None + and vp.width is not None + and vp.height is not None + and not self.preserves_reference_image_size + ): target_size = (vp.width, vp.height) image_items = input_image if isinstance(input_image, list) else [input_image] resized_images = [ diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index b11ceadec..f10ca9806 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/vllm_omni/model_extras/__init__.py @@ -10,8 +10,10 @@ from vllm_omni.model_extras.registry import ( get_extra_output_params, get_model_class_name, get_output_tensor_range, + get_transformer_config_subfolder, get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, + should_preserve_reference_image_size, ) __all__ = [ @@ -23,6 +25,8 @@ __all__ = [ "get_extra_output_params", "get_model_class_name", "get_output_tensor_range", + "get_transformer_config_subfolder", "get_x_to_text_model_family", "should_init_extra_args_for_non_diffusion_stages", + "should_preserve_reference_image_size", ] diff --git a/vllm_omni/model_extras/ltx2.py b/vllm_omni/model_extras/ltx2.py index 5bac58469..7516a55f2 100644 --- a/vllm_omni/model_extras/ltx2.py +++ b/vllm_omni/model_extras/ltx2.py @@ -21,7 +21,35 @@ LTX_EXTRA_BODY_PARAMS = frozenset( "audio_rescale_scale", "video_stg_blocks", "audio_stg_blocks", + "sigmas", + "stage_1_sigmas", + "stage_2_sigmas", + "image_crf", } ) LTX_EXTRA_OUTPUT_PARAMS = frozenset() + + +def ltx_preserves_reference_image_size(*, model: str | None, revision: str | None = None) -> bool: + """Keep source geometry until LTX-2.5 applies conditioning compression.""" + if model is None: + return False + from vllm_omni.diffusion.models.ltx2.ltx2_components import ( + preserves_reference_image_size as ltx_preserves_reference_image_size, + ) + + return ltx_preserves_reference_image_size(model=model, revision=revision) + + +def ltx_transformer_config_subfolder( + *, + model: str | None, + revision: str | None = None, +) -> str: + """Select the transformer config used by the requested LTX checkpoint.""" + if model is None: + return "transformer" + from vllm_omni.diffusion.models.ltx2.ltx2_components import detect_ltx_model_version + + return "transformer_full" if detect_ltx_model_version(model, revision=revision) == "2.5" else "transformer" diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index bc923ea86..c394b8999 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -4,7 +4,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping -from typing import Any, Literal +from typing import Any, Literal, Protocol from PIL import Image @@ -34,7 +34,12 @@ from vllm_omni.model_extras.helios import ( ) from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as build_hunyuan_x_to_text_prompt from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS -from vllm_omni.model_extras.ltx2 import LTX_EXTRA_BODY_PARAMS, LTX_EXTRA_OUTPUT_PARAMS +from vllm_omni.model_extras.ltx2 import ( + LTX_EXTRA_BODY_PARAMS, + LTX_EXTRA_OUTPUT_PARAMS, + ltx_preserves_reference_image_size, + ltx_transformer_config_subfolder, +) from vllm_omni.model_extras.magi_human import ( MAGI_HUMAN_EXTRA_BODY_PARAMS, MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, @@ -96,6 +101,24 @@ XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int OutputTensorRange = Literal["negative_one_to_one", "zero_to_one"] +class ReferenceImageSizeResolver(Protocol): + def __call__( + self, + *, + model: str | None, + revision: str | None = None, + ) -> bool: ... + + +class TransformerConfigSubfolderResolver(Protocol): + def __call__( + self, + *, + model: str | None, + revision: str | None = None, + ) -> str: ... + + def default_x_to_text_prompt( model: str, prompt: str, @@ -202,11 +225,14 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { model_class_name: { "extra_body_params": LTX_EXTRA_BODY_PARAMS, "extra_output_params": LTX_EXTRA_OUTPUT_PARAMS, + "reference_image_size_resolver": ltx_preserves_reference_image_size, } for model_class_name in ( "LTX2Pipeline", "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", "LTX2DistilledPipeline", + "LTX2DistilledTwoStagePipeline", ) }, "WanVACEPipeline": { @@ -229,6 +255,10 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { }, } +for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"): + _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder + + # Multi-stage discovery reports the top-level wrapper rather than its DiT # submodule, so both names must resolve to the same request builders. _EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"] @@ -277,6 +307,36 @@ def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange: return spec.get("output_tensor_range", "negative_one_to_one") +def get_transformer_config_subfolder( + model_class_name: str | None, + *, + model: str | None, + revision: str | None = None, +) -> str: + """Return the model-declared DiT config subfolder, or the standard default.""" + spec = _get_spec(model_class_name) + resolver: TransformerConfigSubfolderResolver | None = ( + spec.get("transformer_config_subfolder_resolver") if spec else None + ) + return resolver(model=model, revision=revision) if resolver else "transformer" + + +def should_preserve_reference_image_size( + model_class_name: str | None, + *, + model: str | None, + revision: str | None = None, +) -> bool: + """Return whether the selected pipeline owns reference-image resizing.""" + if model_class_name is None and model is not None: + from vllm_omni.diffusion.data import resolve_model_class_name + + model_class_name = resolve_model_class_name(model, revision=revision) + spec = _get_spec(model_class_name) + resolver: ReferenceImageSizeResolver | None = spec.get("reference_image_size_resolver") if spec else None + return bool(resolver and resolver(model=model, revision=revision)) + + def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool: spec = _get_spec(model_class_name) return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False))