anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Refactor] Unify the LTX-2 and LTX-2.3 pipeline runtime (#5147)

Browse Source 
Signed-off-by: mglyn <1203789601@qq.com>
This commit is contained in:
Mu GuanLin authored and GitHub committed 2026-07-20 20:02:30 +08:00
1 parent 6cd1ff7e64
commit a06664a243
36 files changed
+3886 -4835
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/models/supported_models.md
+2 -2
View File
@@ -45,8 +45,8 @@ th {
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX23Pipeline` | LTX-2.3-T2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
docs/user_guide/diffusion_features.md
+2 -2
View File
@@ -154,8 +154,8 @@ The following tables show which models support each feature:
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
docs/user_guide/examples/offline_inference/image_to_video.md
+2 -2
View File
@@ -57,7 +57,7 @@ python image_to_video.py \
```bash
python image_to_video.py \
--model dg845/LTX-2.3-Diffusers \
--model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23ImageToVideoPipeline \
--image cherry_blossom.jpg \
--prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
@@ -72,7 +72,7 @@ python image_to_video.py \
--output ltx23_i2v_output.mp4
```
Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
the LTX-2.3 image-to-video pipeline.
docs/user_guide/examples/offline_inference/text_to_video.md
+3 -3
View File
@@ -12,7 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|---|---|---|---|---|---|
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `dg845/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
@@ -55,7 +55,7 @@ python text_to_video.py \
```bash
python text_to_video.py \
--model dg845/LTX-2.3-Diffusers \
--model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23Pipeline \
--prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
@@ -70,7 +70,7 @@ python text_to_video.py \
--output ltx23_t2v_output.mp4
```
Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
text-to-video pipeline explicitly.
docs/user_guide/examples/online_serving/text_to_video.md
+12
View File
@@ -13,6 +13,7 @@ This example demonstrates how to deploy text-to-video models for online video ge
| Wan2.1 T2V (14B) | `Wan-AI/Wan2.1-T2V-14B-Diffusers` |
| Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
| LTX-2 | `Lightricks/LTX-2` |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` |
## Wan2.2 T2V
@@ -374,6 +375,17 @@ curl -sS -X POST http://localhost:8098/v1/videos \
-F "seed=42"
```
## LTX-2.3
LTX-2.3 uses a Diffusers-format checkpoint and a dedicated pipeline class:
```bash
vllm serve diffusers/LTX-2.3-Diffusers --omni --port 8098 \
--model-class-name LTX23Pipeline
```
Send requests through the same `/v1/videos` API shown in the LTX-2 example.
## Example materials
??? abstract "response.json"
examples/offline_inference/image_to_video/README.md
+2 -2
View File
@@ -237,7 +237,7 @@ python image_to_video.py \
```bash
python image_to_video.py \
--model dg845/LTX-2.3-Diffusers \
--model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23ImageToVideoPipeline \
--image cherry_blossom.jpg \
--prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
@@ -252,7 +252,7 @@ python image_to_video.py \
--output ltx23_i2v_output.mp4
```
Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
the LTX-2.3 image-to-video pipeline.
examples/offline_inference/image_to_video/image_to_video.py
+1 -1
View File
@@ -29,7 +29,7 @@ Usage:
--frame-rate 24 --fps 24 --output ltx2_i2v.mp4
# LTX-2.3 image-to-video
python image_to_video.py --model dg845/LTX-2.3-Diffusers \
python image_to_video.py --model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23ImageToVideoPipeline \
--image input.jpg --prompt "A cinematic dolly shot of a boat" \
--height 384 --width 512 --num-frames 25 --num-inference-steps 20 \
examples/offline_inference/text_to_video/text_to_video.md
+3 -3
View File
@@ -8,7 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
|---|---|---|---|---|---|
| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `dg845/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
| `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
@@ -75,7 +75,7 @@ python text_to_video.py \
```bash
python text_to_video.py \
--model dg845/LTX-2.3-Diffusers \
--model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23Pipeline \
--prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
@@ -90,7 +90,7 @@ python text_to_video.py \
--output ltx23_t2v_output.mp4
```
Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
text-to-video pipeline explicitly.
recipes/LTX/LTX-2.3.md
+9 -9
View File
@@ -5,7 +5,7 @@
## Summary
- Vendor: Lightricks
- Model: `dg845/LTX-2.3-Diffusers`
- Model: `diffusers/LTX-2.3-Diffusers`
- Task: Text-to-video and image-to-video with synchronized audio generation
- Mode: Online serving (pure diffusion)
- Maintainer: @oglok
@@ -22,7 +22,7 @@ encoder, VAE, and vocoder components.
## References
- Upstream raw checkpoints: <https://huggingface.co/Lightricks/LTX-2.3>
- Diffusers-format checkpoint: <https://huggingface.co/dg845/LTX-2.3-Diffusers>
- Diffusers-format checkpoint: <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
- Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`)
## Serving
@@ -30,7 +30,7 @@ encoder, VAE, and vocoder components.
### Command
```bash
vllm serve dg845/LTX-2.3-Diffusers \
vllm serve diffusers/LTX-2.3-Diffusers \
--omni \
--model-class-name LTX23Pipeline \
--stage-init-timeout 600
@@ -39,7 +39,7 @@ vllm serve dg845/LTX-2.3-Diffusers \
For image-to-video:
```bash
vllm serve dg845/LTX-2.3-Diffusers \
vllm serve diffusers/LTX-2.3-Diffusers \
--omni \
--model-class-name LTX23ImageToVideoPipeline \
--stage-init-timeout 600
@@ -55,7 +55,7 @@ curl http://localhost:8000/health
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=A majestic bald eagle soaring over a misty mountain valley at dawn, golden sunlight breaking through clouds" \
-F "negative_prompt=blurry, low quality, distorted, watermark" \
-F "model=dg845/LTX-2.3-Diffusers" \
-F "model=diffusers/LTX-2.3-Diffusers" \
-F "num_frames=81" \
-F "fps=24" \
-F "size=768x512" \
@@ -66,7 +66,7 @@ curl -X POST http://localhost:8000/v1/videos \
# Generate a 10-second video (241 frames)
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=A cozy Japanese ramen shop at night in the rain, steam rising from bowls, neon signs reflecting on wet cobblestone streets" \
-F "model=dg845/LTX-2.3-Diffusers" \
-F "model=diffusers/LTX-2.3-Diffusers" \
-F "num_frames=241" \
-F "fps=24" \
-F "size=768x512" \
@@ -76,7 +76,7 @@ curl -X POST http://localhost:8000/v1/videos \
# Generate a 20-second video (481 frames)
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=An underwater coral reef teeming with tropical fish, sea turtles gliding gracefully, National Geographic documentary style" \
-F "model=dg845/LTX-2.3-Diffusers" \
-F "model=diffusers/LTX-2.3-Diffusers" \
-F "num_frames=481" \
-F "fps=24" \
-F "size=768x512" \
@@ -96,7 +96,7 @@ curl http://localhost:8000/health
curl -X POST http://localhost:8000/v1/videos/sync \
-F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
-F "input_reference=@/absolute/path/to/reference.png" \
-F "model=dg845/LTX-2.3-Diffusers" \
-F "model=diffusers/LTX-2.3-Diffusers" \
-F "num_frames=81" \
-F "fps=24" \
-F "size=768x512" \
@@ -111,7 +111,7 @@ curl -X POST http://localhost:8000/v1/videos/sync \
- Memory usage: run a latest-head serving smoke or DFX benchmark on the target
GPU before publishing load/peak VRAM numbers.
- Checkpoint format: use a Diffusers-format checkpoint such as
`dg845/LTX-2.3-Diffusers`; the upstream `Lightricks/LTX-2.3` repository ships
`diffusers/LTX-2.3-Diffusers`; the upstream `Lightricks/LTX-2.3` repository ships
raw safetensors and does not contain the subfolder configs required by this
pipeline loader.
- Key flags:
tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py
+29 -2
View File
@@ -6,12 +6,17 @@ import pytest
import torch
from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2ImageToVideoPipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImageDMD2Pipeline, QwenImagePipeline
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -152,6 +157,28 @@ def test_string_prompt_not_mutated(pipeline):
assert req.prompt == "a cat dancing"
def test_request_batch_sanitizes_each_request(pipeline):
requests = [
_make_request(
prompt={"prompt": f"prompt-{index}", "negative_prompt": "blurry"},
num_inference_steps=40 + index,
guidance_scale=5.0 + index,
guidance_scale_provided=True,
)
for index in range(2)
]
pipeline._sanitize_dmd2_request(DiffusionRequestBatch(requests))
for request in requests:
assert request.sampling_params.num_inference_steps == pipeline.dmd2_config.num_inference_steps
assert request.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale
assert request.sampling_params.guidance_scale_provided is False
assert request.sampling_params.do_classifier_free_guidance is False
assert request.sampling_params.is_cfg_negative is False
assert "negative_prompt" not in request.prompt
# ---------------------------------------------------------------------------
# Clean request — nothing changes
# ---------------------------------------------------------------------------
tests/diffusion/models/dmd2/test_dmd2_scheduler.py
+6 -2
View File
@@ -6,8 +6,12 @@ import pytest
import torch
from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2ImageToVideoPipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImageDMD2Pipeline, QwenImagePipeline
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline
tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py
-58
View File
@@ -1,58 +0,0 @@
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_pipeline(sequence_parallel_size: int = 1) -> LTX2Pipeline:
pipeline = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.audio_vae_temporal_compression_ratio = 4
pipeline.audio_vae_mel_compression_ratio = 4
pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
# Mock audio_vae with identity normalization (mean=0, std=1) so
# _normalize_audio_latents is a no-op and test values are preserved.
pipeline.audio_vae = SimpleNamespace(
latents_mean=torch.tensor(0.0),
latents_std=torch.tensor(1.0),
)
return pipeline
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
pipeline = _make_pipeline(sequence_parallel_size=4)
latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=10,
dtype=torch.float32,
device=torch.device("cpu"),
latents=latents,
)
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 4)
torch.testing.assert_close(padded[:, :10], latents)
torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
def test_unpad_audio_latents_restores_original_frames_before_unpack():
pipeline = _make_pipeline()
original = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
unpadded = pipeline._unpad_audio_latents(padded, 10)
unpacked = pipeline._unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
expected = pipeline._unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
assert unpacked.shape == (1, 2, 10, 2)
assert not (unpacked == 999.0).any()
torch.testing.assert_close(unpacked, expected)
tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py → tests/diffusion/models/ltx2/test_ltx2_conditioning.py
+20 -45
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the LTX-2.3 image-to-video pipeline."""
"""Unit tests for LTX image-to-video input and conditioning behavior."""
from types import SimpleNamespace
@@ -19,9 +19,9 @@ def _make_ltx23_request_pipe(cls):
return pipe
class TestLTX23ImageToVideoForwardStages:
class TestLTXImageToVideoForwardStages:
def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -66,7 +66,7 @@ class TestLTX23ImageToVideoForwardStages:
assert seen["kwargs"]["noise_scale"] == 0.5
def test_denoise_timestep_kwargs_masks_video_only(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
ts = torch.tensor([2.0, 4.0])
@@ -86,16 +86,14 @@ class TestLTX23ImageToVideoForwardStages:
torch.testing.assert_close(kwargs["sigma"], ts)
class TestLTX23ImageToVideoPipeline:
def test_ltx23_i2v_pipeline_reuses_ltx23_semantics(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
class TestLTXImageToVideoConditioning:
def test_ltx23_i2v_supports_image_input(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
assert issubclass(LTX23ImageToVideoPipeline, LTX23Pipeline)
assert LTX23ImageToVideoPipeline.support_image_input is True
def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
image = object()
@@ -104,7 +102,7 @@ class TestLTX23ImageToVideoPipeline:
LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()])
def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
image = torch.zeros(1, 3, 4, 4)
additional = {
@@ -116,8 +114,9 @@ class TestLTX23ImageToVideoPipeline:
assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image
def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
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 LTX23ImageToVideoPipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
torch.nn.Module.__init__(pipe)
@@ -129,7 +128,8 @@ class TestLTX23ImageToVideoPipeline:
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
raise AssertionError("packed I2V latents should not be noised")
monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]])
@@ -150,9 +150,9 @@ class TestLTX23ImageToVideoPipeline:
torch.testing.assert_close(out, latents)
def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx2
import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
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 LTX23ImageToVideoPipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
torch.nn.Module.__init__(pipe)
@@ -169,8 +169,8 @@ class TestLTX23ImageToVideoPipeline:
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
return torch.ones(shape, device=device, dtype=dtype)
monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
monkeypatch.setattr(ltx2, "randn_tensor", fake_randn_tensor)
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
@@ -191,7 +191,7 @@ class TestLTX23ImageToVideoPipeline:
torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
torch.nn.Module.__init__(pipe)
@@ -217,28 +217,3 @@ class TestLTX23ImageToVideoPipeline:
torch.testing.assert_close(out[:, :1], latents[:, :1])
torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5)
class TestLTX23ImageToVideoModule:
"""Test that pipeline_ltx2_3_image2video.py exposes I2V entry points."""
def test_i2v_classes_importable(self):
"""I2V classes must be importable from the implementation module."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
assert LTX23ImageToVideoPipeline is not None
def test_post_process_func_importable(self):
"""get_ltx2_post_process_func must be importable for registry lookup."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func
assert callable(get_ltx2_post_process_func)
def test_i2v_class_matches_package_export(self):
"""Package export must point at the I2V implementation module."""
from vllm_omni.diffusion.models.ltx2 import LTX23ImageToVideoPipeline as PackageExported
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import (
LTX23ImageToVideoPipeline as ModuleExported,
)
assert PackageExported is ModuleExported
tests/diffusion/models/ltx2/test_ltx2_hsdp.py
-25
View File
@@ -1,25 +0,0 @@
import pytest
import torch.nn as nn
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_ltx2_exposes_hsdp_shard_conditions_for_transformer_blocks():
model = object.__new__(LTX2VideoTransformer3DModel)
nn.Module.__init__(model)
model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
model.norm_out = nn.LayerNorm(4)
conditions = getattr(model, "_hsdp_shard_conditions", None)
assert conditions is not None
assert len(conditions) == 1
matched = []
for name, module in model.named_modules():
if any(cond(name, module) for cond in conditions):
matched.append(name)
assert matched == ["transformer_blocks.0", "transformer_blocks.1"]
tests/diffusion/models/ltx2/test_ltx2_latents.py
+121
View File
@@ -0,0 +1,121 @@
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.models.ltx2 import ltx2_latents as latent_ops
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX23Pipeline
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.audio_vae_temporal_compression_ratio = 4
pipeline.audio_vae_mel_compression_ratio = 4
pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
# Mock audio_vae with identity normalization (mean=0, std=1).
pipeline.audio_vae = SimpleNamespace(
latents_mean=torch.tensor(0.0),
latents_std=torch.tensor(1.0),
)
return pipeline
def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=2)
latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=8,
num_mel_bins=64,
audio_latent_length=1,
dtype=torch.float32,
device=torch.device("cpu"),
)
assert original_num_frames == 1
assert padded_num_frames == 2
assert latents.shape == (1, 2, 128)
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=10,
dtype=torch.float32,
device=torch.device("cpu"),
latents=latents,
)
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 4)
torch.testing.assert_close(padded[:, :10], latents)
torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
def test_unpad_audio_latents_restores_original_frames_before_unpack():
original = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
unpadded = latent_ops.unpad_audio_latents(padded, 10)
unpacked = latent_ops.unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
expected = latent_ops.unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
assert unpacked.shape == (1, 2, 10, 2)
assert not (unpacked == 999.0).any()
torch.testing.assert_close(unpacked, expected)
def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=16,
audio_latent_length=audio_latent_length,
dtype=torch.float32,
device=torch.device("cpu"),
latents=latents,
)
assert audio_latent_length == 10
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 8)
torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents))
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
latents = torch.zeros(1, 2, 13, 4)
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
assert audio_latent_length == 13
def test_prepare_audio_latents_rejects_incompatible_provided_length():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
latents = torch.zeros(1, 11, 4)
with pytest.raises(ValueError, match="incompatible audio frame count"):
pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=10,
dtype=torch.float32,
device=torch.device("cpu"),
latents=latents,
)
tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py → tests/diffusion/models/ltx2/test_ltx2_parallel.py
+25 -37
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for LTX-2.3 forward-parallel behavior."""
"""Unit tests for LTX guidance and forward-parallel behavior."""
from types import SimpleNamespace
@@ -11,29 +11,11 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestCFGParallelIntegration:
def test_ltx23_pipeline_has_cfg_parallel_mixin(self):
"""LTX23Pipeline must use the shared CFG parallel implementation."""
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
assert issubclass(LTX23Pipeline, CFGParallelMixin)
def test_ltx2_transformer_has_dit_cache_config(self):
"""Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG."""
from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
assert isinstance(adapter_config, CacheDiTAdapterConfig)
assert adapter_config.has_separate_cfg
class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
def test_combine_cfg_noise_matches_x0_space_formula(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
pipe = object.__new__(LTX23Pipeline)
video_sample = torch.tensor([[[1.0, -2.0]]])
@@ -69,7 +51,8 @@ class TestCFGParallelHelpers:
assert torch.allclose(audio_combined, expected_audio)
def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
from vllm_omni.diffusion.models.ltx2 import ltx2_guidance
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
pipe = object.__new__(ltx23.LTX23Pipeline)
video_sample = torch.tensor([[[1.0, -2.0]]])
@@ -89,17 +72,17 @@ class TestCFGParallelHelpers:
return [video_pos, video_neg]
return [audio_pos, audio_neg]
monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg(
expected_video = ltx2_guidance.combine_velocity_via_x0(
video_sample,
video_pos,
video_neg,
video_sigma,
scale,
)
expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg(
expected_audio = ltx2_guidance.combine_velocity_via_x0(
audio_sample,
audio_pos,
audio_neg,
@@ -109,7 +92,7 @@ class TestCFGParallelHelpers:
for rank, expected_branch in ((0, "positive"), (1, "negative")):
calls = []
monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
def fake_predict_noise(**kwargs):
calls.append(kwargs["branch"])
@@ -140,17 +123,20 @@ class TestCFGParallelHelpers:
class TestCFGParallelForwardPath:
"""Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
@pytest.mark.parametrize("pipeline_cls_name", ["LTX2Pipeline", "LTX23Pipeline"])
def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch, pipeline_cls_name):
from vllm_omni.diffusion.models.ltx2 import ltx2_pipeline_runtime
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pipe = object.__new__(ltx23.LTX23Pipeline)
pipeline_cls = getattr(ltx23, pipeline_cls_name)
pipe = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipe)
pipe.device = torch.device("cpu")
pipe.tokenizer_max_length = 4
monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 1)
monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
class StopAtEncodePromptError(Exception):
pass
@@ -247,7 +233,8 @@ class TestCFGParallelForwardPath:
audio_sampling_rate,
expected_frame_rate,
):
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_pipeline_runtime
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -291,15 +278,16 @@ class TestCFGParallelForwardPath:
return [audio_pos, audio_neg]
raise AssertionError(f"Unexpected gathered tensor: {tensor}")
monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank)
monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank)
monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)
return torch.tensor([1.0, 0.5], device=device), 2
monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps)
monkeypatch.setattr(ltx2_denoise, "retrieve_timesteps", fake_retrieve_timesteps)
class FakeScheduler:
def __init__(self, name="video", calls=None):
@@ -402,14 +390,14 @@ class TestCFGParallelForwardPath:
output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0]
expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
expected_video_noise = ltx2_guidance.combine_velocity_via_x0(
video_latents,
video_pos,
video_neg,
pipe.scheduler.sigmas[0],
4.0,
)
expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
expected_audio_noise = ltx2_guidance.combine_velocity_via_x0(
audio_latents,
audio_pos,
audio_neg,
tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py → tests/diffusion/models/ltx2/test_ltx2_pipeline.py
+340 -193
View File
@@ -1,16 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for LTX-2.3 pipeline integration.
These tests verify:
- Pipeline is properly registered in the diffusion registry
- Post-process function is registered
- Cache-DiT enablers are registered
- Pipeline does NOT inherit from LTX2Pipeline
- Vocoder sample rate detection logic
- Re-export module works correctly
"""
"""Unit tests for the shared LTX pipeline runtime and public contracts."""
import json
import os
@@ -20,25 +11,37 @@ from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
LTX2_COMPONENT_PROFILE,
LTX23_COMPONENT_PROFILE,
)
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseExecutor, LTXPhaseResult
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTX_LEGACY_VELOCITY_GUIDANCE,
LTX_OFFICIAL_X0_GUIDANCE,
)
from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import LTXPipelineRuntime
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE
from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2ImageToVideoPipeline,
LTX2Pipeline,
LTX23ImageToVideoPipeline,
LTX23Pipeline,
LTXOneStagePipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
LTX2ImageToVideoTwoStagesPipeline,
LTX2TwoStagesPipeline,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_ltx23_pipeline(sequence_parallel_size: int = 1):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
pipeline = object.__new__(LTX23Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.audio_vae_temporal_compression_ratio = 4
pipeline.audio_vae_mel_compression_ratio = 4
pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
pipeline.audio_vae = SimpleNamespace(
latents_mean=torch.tensor(0.0),
latents_std=torch.tensor(1.0),
)
return pipeline
def _make_ltx23_request_pipe(cls):
def _make_ltx_request_pipe(cls):
pipe = object.__new__(cls)
torch.nn.Module.__init__(pipe)
pipe.device = torch.device("cpu")
@@ -73,10 +76,263 @@ def _resolve_request_inputs_for_test(pipe, req):
)
class TestLTX23RequestParsing:
def test_t2v_and_i2v_share_request_input_resolution(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
def test_ltx_versions_share_runtime_without_cross_version_inheritance():
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
DistributedAutoencoderKLLTX2Video,
)
assert issubclass(LTX2Pipeline, LTXPipelineRuntime)
assert issubclass(LTX23Pipeline, LTXPipelineRuntime)
assert issubclass(LTX2TwoStagesPipeline, LTXPipelineRuntime)
assert not issubclass(LTX23Pipeline, LTX2Pipeline)
assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
assert LTX23Pipeline.component_profile is LTX23_COMPONENT_PROFILE
assert LTX2Pipeline.one_stage_recipe is LTX2_ONE_STAGE_RECIPE
assert LTX23Pipeline.one_stage_recipe is LTX23_ONE_STAGE_RECIPE
assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
def test_ltx_one_stage_variants_expose_shared_request_batch_and_distributed_decode_contracts():
one_stage_variants = (
LTX2Pipeline,
LTX2ImageToVideoPipeline,
LTX23Pipeline,
LTX23ImageToVideoPipeline,
)
assert all(pipeline_cls.supports_request_batch for pipeline_cls in one_stage_variants)
assert all(pipeline_cls.distributed_video_decode for pipeline_cls in one_stage_variants)
assert not LTX2TwoStagesPipeline.supports_request_batch
assert not LTX2ImageToVideoTwoStagesPipeline.supports_request_batch
@pytest.mark.parametrize("pipeline_cls", [LTX2Pipeline, LTX23Pipeline])
def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cls):
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
requests = [
OmniDiffusionRequest(
prompt=f"prompt-{index}",
sampling_params=OmniDiffusionSamplingParams(),
request_id=f"request-{index}",
)
for index in range(2)
]
request_batch = DiffusionRequestBatch(requests)
video = torch.tensor([[1.0], [2.0]])
audio = torch.tensor([[3.0], [4.0]])
pipe = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipe)
object.__setattr__(pipe, "_decode_output", lambda **_kwargs: DiffusionOutput(output=(video, audio)))
forward_ctx = SimpleNamespace(
req=request_batch,
request_inputs=SimpleNamespace(
output_type="np",
generator=None,
decode_timestep=0.0,
decode_noise_scale=None,
),
prompt_context=SimpleNamespace(connector_prompt_embeds=torch.empty(2, 1)),
device=torch.device("cpu"),
batch_size=2,
num_videos_per_prompt=1,
)
outputs = pipe.decode_phase(
LTXPhaseResult(
forward_context=forward_ctx,
video=video,
audio=audio,
)
)
assert len(outputs) == 2
torch.testing.assert_close(outputs[0].output[0], video[:1])
torch.testing.assert_close(outputs[0].output[1], audio[:1])
torch.testing.assert_close(outputs[1].output[0], video[1:])
torch.testing.assert_close(outputs[1].output[1], audio[1:])
def test_ltx_one_stage_variants_share_forward_template():
assert LTX2Pipeline._forward_impl is LTXOneStagePipeline._forward_impl
assert LTX2ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl
assert LTX23Pipeline._forward_impl is LTXOneStagePipeline._forward_impl
assert LTX23ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl
assert LTX2ImageToVideoPipeline.forward is LTX23ImageToVideoPipeline.forward
def test_ltx_versions_share_request_prompt_and_step_templates():
shared_methods = (
"_get_gemma_prompt_embeds",
"encode_prompt",
"check_inputs",
"_resolve_request_inputs",
"_prepare_prompt_context",
"_denoise_step",
)
for method_name in shared_methods:
base_method = getattr(LTXPipelineRuntime, method_name)
assert getattr(LTX2Pipeline, method_name) is base_method
assert getattr(LTX23Pipeline, method_name) is base_method
def test_ltx_versions_select_guidance_without_overriding_control_flow():
assert LTX2Pipeline.guidance_strategy is LTX_LEGACY_VELOCITY_GUIDANCE
assert LTX23Pipeline.guidance_strategy is LTX_OFFICIAL_X0_GUIDANCE
assert LTX2Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step
assert LTX23Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step
assert LTX2Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise
assert LTX23Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise
def test_ltx2_two_stage_variants_share_stage_orchestration():
assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTX2TwoStagesPipeline)
assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTXI2VConditioningMixin)
assert LTX2ImageToVideoTwoStagesPipeline._run_two_stage is LTX2TwoStagesPipeline._run_two_stage
def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
assert "_denoise_loop" not in LTX2Pipeline.__dict__
assert "_denoise_loop" not in LTX23Pipeline.__dict__
assert issubclass(LTX2ImageToVideoPipeline, LTXI2VConditioningMixin)
assert issubclass(LTX23ImageToVideoPipeline, LTXI2VConditioningMixin)
assert LTX2ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents
assert LTX23ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents
assert LTX2Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents
assert LTX23Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents
assert LTX2Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents
assert LTX23Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents
assert LTX2Pipeline.decode_phase is LTXPipelineRuntime.decode_phase
assert LTX23Pipeline.decode_phase is LTXPipelineRuntime.decode_phase
def test_denoise_executor_owns_progress_and_interrupt():
updates = []
class Progress:
def __enter__(self):
return self
def __exit__(self, exc_type, exc, traceback):
return False
def update(self):
updates.append(True)
pipeline = SimpleNamespace(
interrupt=False,
progress_bar=lambda total: Progress(),
)
seen: list[tuple[int, float]] = []
timesteps = torch.tensor([3.0, 2.0, 1.0])
initial_state = LTXAVState(video=torch.tensor(0.0), audio=torch.tensor(10.0))
def step(index, timestep, state):
seen.append((index, timestep.item()))
pipeline.interrupt = True
return LTXAVState(video=state.video + 1, audio=state.audio + 1)
state = LTXDenoiseExecutor.run(pipeline, initial_state, timesteps, step)
assert seen == [(0, 3.0)]
assert updates == [True]
torch.testing.assert_close(state.video, torch.tensor(1.0))
torch.testing.assert_close(state.audio, torch.tensor(11.0))
def test_ltx2_two_stage_reuses_prompt_context_between_phases():
request_inputs = LTXRequestInputs(
prompt="prompt",
negative_prompt="negative",
height=32,
width=32,
num_frames=1,
frame_rate=24.0,
num_inference_steps=4,
guidance_scale=4.0,
guidance_rescale=0.0,
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_calls.append((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
video = torch.tensor([1.0])
audio = torch.tensor([2.0])
else:
assert prompt_context is prompt_context_sentinel
torch.testing.assert_close(inputs.latents, torch.tensor([11.0]))
torch.testing.assert_close(inputs.audio_latents, torch.tensor([2.0]))
assert inputs.guidance_scale == 1.0
assert inputs.num_inference_steps == 3
assert inputs.decode_timestep == 0.0
assert inputs.decode_noise_scale is None
context = prompt_context
video = torch.tensor([3.0])
audio = torch.tensor([4.0])
return LTXPhaseResult(
forward_context=SimpleNamespace(prompt_context=context),
video=video,
audio=audio,
)
def decode_phase(phase):
return DiffusionOutput(output=(phase.video, phase.audio))
class FakeUpsampler(torch.nn.Module):
def forward(self, *, latents, output_type, return_dict):
assert output_type == "latent"
assert not return_dict
return (latents + 10,)
prompt_context_sentinel = prompt_context
pipeline = object.__new__(LTX2TwoStagesPipeline)
torch.nn.Module.__init__(pipeline)
pipeline.distilled = True
pipeline.upsample_pipe = FakeUpsampler()
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
object.__setattr__(pipeline, "run_phase", run_phase)
object.__setattr__(pipeline, "decode_phase", decode_phase)
output = pipeline.forward(SimpleNamespace())
assert len(phase_calls) == 2
assert phase_calls[1][1] is prompt_context_sentinel
torch.testing.assert_close(output.output[0], torch.tensor([3.0]))
torch.testing.assert_close(output.output[1], torch.tensor([4.0]))
class TestLTXRequestParsing:
@pytest.mark.parametrize(
("t2v_cls", "i2v_cls"),
[
(LTX2Pipeline, LTX2ImageToVideoPipeline),
(LTX23Pipeline, LTX23ImageToVideoPipeline),
],
)
def test_t2v_and_i2v_share_request_input_resolution(self, t2v_cls, i2v_cls):
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -124,11 +380,11 @@ class TestLTX23RequestParsing:
)
resolved_t2v = _resolve_request_inputs_for_test(
_make_ltx23_request_pipe(LTX23Pipeline),
_make_ltx_request_pipe(t2v_cls),
req,
)
resolved_i2v = _resolve_request_inputs_for_test(
_make_ltx23_request_pipe(LTX23ImageToVideoPipeline),
_make_ltx_request_pipe(i2v_cls),
req,
)
@@ -169,12 +425,12 @@ class TestLTX23RequestParsing:
)
def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pipe = _make_ltx23_request_pipe(LTX23Pipeline)
pipe = _make_ltx_request_pipe(LTX23Pipeline)
req = DiffusionRequestBatch(
[
OmniDiffusionRequest(
@@ -209,9 +465,9 @@ class TestLTX23RequestParsing:
class TestLTX23ForwardStages:
def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
pipe = _make_ltx23_request_pipe(LTX23Pipeline)
pipe = _make_ltx_request_pipe(LTX23Pipeline)
prompt_embeds = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]])
negative_prompt_embeds = torch.tensor([[[-1.0], [-2.0]], [[-3.0], [-4.0]]])
prompt_attention_mask = torch.tensor([[True, False], [False, True]])
@@ -251,12 +507,12 @@ class TestLTX23ForwardStages:
)
def test_t2v_forward_delegates_to_shared_forward_impl(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pipe = _make_ltx23_request_pipe(LTX23Pipeline)
pipe = _make_ltx_request_pipe(LTX23Pipeline)
req = DiffusionRequestBatch(
[
OmniDiffusionRequest(
@@ -311,36 +567,10 @@ class TestLTX23ForwardStages:
}
class TestPipelineIndependence:
"""Verify LTX23Pipeline is fully independent from LTX2Pipeline."""
def test_ltx23_pipeline_does_not_inherit_from_ltx2(self):
"""LTX23Pipeline must NOT inherit from LTX2Pipeline."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
assert not issubclass(LTX23Pipeline, LTX2Pipeline), (
"LTX23Pipeline should be fully independent and not inherit from LTX2Pipeline"
)
def test_ltx23_pipeline_is_nn_module(self):
"""LTX23Pipeline must be an nn.Module."""
import torch.nn as nn
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
assert issubclass(LTX23Pipeline, nn.Module)
def test_ltx23_pipeline_has_progress_bar(self):
"""LTX23Pipeline must mix in ProgressBarMixin."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
assert issubclass(LTX23Pipeline, ProgressBarMixin)
class TestPipelineComponents:
def test_ltx23_pipeline_declares_offload_components(self):
"""LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
pipe = object.__new__(LTX23Pipeline)
@@ -363,17 +593,10 @@ class TestPipelineIndependence:
assert modules.resident_names == ["vocoder"]
assert len(modules.vaes) == 2
def test_ltx23_pipeline_has_diffusion_pipeline_profiler_mixin(self):
"""LTX23Pipeline must support lightweight stage timing."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin)
class TestLTX23DecodeConditioning:
def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value
from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value
assert _expand_per_prompt_decode_value(
[0.1, 0.2],
@@ -395,7 +618,7 @@ class TestLTX23DecodeConditioning:
) == [0.1, 0.2, 0.3, 0.4]
def test_decode_conditioning_rejects_ambiguous_lengths(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value
from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value
with pytest.raises(ValueError, match="decode_timestep"):
_expand_per_prompt_decode_value(
@@ -409,28 +632,15 @@ class TestLTX23DecodeConditioning:
class TestRegistryIntegration:
"""Verify all LTX-2.3 pipeline variants are registered."""
def test_pipeline_models_registered(self):
"""LTX-2.3 pipeline variants must be in _DIFFUSION_MODELS."""
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
expected = [
"LTX23Pipeline",
"LTX23ImageToVideoPipeline",
]
for name in expected:
assert name in _DIFFUSION_MODELS, f"{name} not found in _DIFFUSION_MODELS"
def test_pipeline_module_paths(self):
"""Registry entries must point to the correct modules."""
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
# T2V -> pipeline_ltx2_3
assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2_3", "LTX23Pipeline")
assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX23Pipeline")
# I2V -> pipeline_ltx2_3_image2video
assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == (
"ltx2",
"pipeline_ltx2_3_image2video",
"pipeline_ltx2",
"LTX23ImageToVideoPipeline",
)
@@ -439,6 +649,12 @@ class TestRegistryIntegration:
from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS
expected = [
"LTX2Pipeline",
"LTX2ImageToVideoPipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
"LTX23Pipeline",
"LTX23ImageToVideoPipeline",
]
@@ -446,6 +662,35 @@ class TestRegistryIntegration:
assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS"
assert _DIFFUSION_POST_PROCESS_FUNCS[name] == "get_ltx2_post_process_func"
@pytest.mark.parametrize(
"model_class_name",
[
"LTX2Pipeline",
"LTX2ImageToVideoPipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
"LTX23Pipeline",
"LTX23ImageToVideoPipeline",
],
)
def test_post_process_func_resolves_from_every_entry_module(self, model_class_name):
from vllm_omni.diffusion.registry import get_diffusion_post_process_func
with tempfile.TemporaryDirectory() as tmpdir:
vocoder_dir = os.path.join(tmpdir, "vocoder")
os.makedirs(vocoder_dir)
with open(os.path.join(vocoder_dir, "config.json"), "w") as config_file:
json.dump({"output_sampling_rate": 48000}, config_file)
post_process = get_diffusion_post_process_func(
SimpleNamespace(model_class_name=model_class_name, model=tmpdir)
)
result = post_process((torch.zeros(1), torch.zeros(1)))
assert result["audio_sample_rate"] == 48000
def test_cache_dit_for_ltx2_does_not_have_custom_enablers_registered(self):
"""Pipeline variants are *not* registered in CUSTOM_DIT_ENABLERS."""
from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
@@ -465,7 +710,7 @@ class TestVocoderSampleRateDetection:
def test_detects_48khz_from_config(self):
"""Should detect output_sampling_rate=48000 from vocoder/config.json."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate
from vllm_omni.diffusion.models.ltx2.ltx2_components import _detect_vocoder_output_sample_rate
with tempfile.TemporaryDirectory() as tmpdir:
vocoder_dir = os.path.join(tmpdir, "vocoder")
@@ -478,7 +723,7 @@ class TestVocoderSampleRateDetection:
def test_returns_none_for_no_output_sr(self):
"""Should return None if vocoder config has no output_sampling_rate."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate
from vllm_omni.diffusion.models.ltx2.ltx2_components import _detect_vocoder_output_sample_rate
with tempfile.TemporaryDirectory() as tmpdir:
vocoder_dir = os.path.join(tmpdir, "vocoder")
@@ -491,7 +736,7 @@ class TestVocoderSampleRateDetection:
def test_returns_none_for_missing_directory(self):
"""Should return None if vocoder directory doesn't exist."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate
from vllm_omni.diffusion.models.ltx2.ltx2_components import _detect_vocoder_output_sample_rate
result = _detect_vocoder_output_sample_rate("/nonexistent/path")
assert result is None
@@ -504,7 +749,7 @@ class TestPostProcessFunction:
"""Post-process func should include audio_sample_rate when detected."""
import torch
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import get_ltx2_post_process_func
with tempfile.TemporaryDirectory() as tmpdir:
vocoder_dir = os.path.join(tmpdir, "vocoder")
@@ -531,7 +776,7 @@ class TestPostProcessFunction:
"""Post-process func should work without vocoder config (no audio_sample_rate key)."""
import torch
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import get_ltx2_post_process_func
class MockConfig:
model = "/nonexistent/path"
@@ -546,101 +791,3 @@ class TestPostProcessFunction:
assert "video" in result
assert "audio" in result
assert "audio_sample_rate" not in result
class TestInitExports:
"""Test that __init__.py exports all LTX-2.3 classes."""
def test_all_ltx23_classes_exported(self):
"""All LTX-2.3 pipeline classes must be in the ltx2 package __all__."""
from vllm_omni.diffusion.models import ltx2
expected_classes = [
"LTX23Pipeline",
"LTX23ImageToVideoPipeline",
]
for name in expected_classes:
assert hasattr(ltx2, name), f"{name} not exported from ltx2 package"
assert name in ltx2.__all__, f"{name} not in ltx2.__all__"
class TestAudioLatentSPPadding:
def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(self):
pipeline = _make_ltx23_pipeline(sequence_parallel_size=2)
latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=8,
num_mel_bins=64,
audio_latent_length=1,
dtype=torch.float32,
device=torch.device("cpu"),
)
assert original_num_frames == 1
assert padded_num_frames == 2
assert latents.shape == (1, 2, 128)
def test_prepare_audio_latents_pads_provided_packed_sequence_dim_for_sp(self):
pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=10,
dtype=torch.float32,
device=torch.device("cpu"),
latents=latents,
)
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 4)
torch.testing.assert_close(padded[:, :10], latents)
torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(self):
pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=16,
audio_latent_length=audio_latent_length,
dtype=torch.float32,
device=torch.device("cpu"),
latents=latents,
)
assert audio_latent_length == 10
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 8)
torch.testing.assert_close(padded, pipeline._pack_audio_latents(latents))
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(self):
pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
latents = torch.zeros(1, 2, 13, 4)
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
assert audio_latent_length == 13
def test_prepare_audio_latents_rejects_incompatible_provided_length(self):
pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
latents = torch.zeros(1, 11, 4)
with pytest.raises(ValueError, match="incompatible audio frame count"):
pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=10,
dtype=torch.float32,
device=torch.device("cpu"),
latents=latents,
)
tests/diffusion/models/ltx2/test_ltx2_transformer.py
+31 -1
View File
@@ -1,9 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from torch import nn
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm
from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
@@ -20,3 +24,29 @@ def test_ltx_rms_norm_affine_weight_remains_parameter():
assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter)
assert "weight" not in dict(norm.named_buffers())
assert "weight" in norm.state_dict()
def test_ltx_transformer_has_separate_cfg_cache_dit_config():
adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
assert isinstance(adapter_config, CacheDiTAdapterConfig)
assert adapter_config.has_separate_cfg
def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
model = object.__new__(LTX2VideoTransformer3DModel)
nn.Module.__init__(model)
model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
model.norm_out = nn.LayerNorm(4)
conditions = getattr(model, "_hsdp_shard_conditions", None)
assert conditions is not None
assert len(conditions) == 1
matched = []
for name, module in model.named_modules():
if any(condition(name, module) for condition in conditions):
matched.append(name)
assert matched == ["transformer_blocks.0", "transformer_blocks.1"]
tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py → tests/diffusion/models/ltx2/test_ltx2_vae.py
+45 -13
View File
@@ -11,24 +11,56 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestLTX23OutputRank:
def test_single_process_rank_is_output_rank(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
class TestLTXOutputRank:
@pytest.mark.parametrize(
("distributed_vae_state", "expected_decode_calls"),
[(False, 0), (True, 1), (RuntimeError("unavailable"), 0)],
)
def test_non_output_rank_only_enters_collective_vae_decode(
self,
monkeypatch,
distributed_vae_state,
expected_decode_calls,
):
import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx
monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: False)
class FakeVae:
dtype = torch.float32
config = SimpleNamespace(timestep_conditioning=False)
assert ltx23._is_output_rank() is True
assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is True
def __init__(self):
self.decode_calls = 0
def test_non_output_rank_skips_decode_unless_vae_decode_is_distributed(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
def is_distributed_enabled(self):
if isinstance(distributed_vae_state, Exception):
raise distributed_vae_state
return distributed_vae_state
monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx23.torch.distributed, "get_rank", lambda: 1)
def decode(self, *_args, **_kwargs):
self.decode_calls += 1
return (torch.ones(1, 1),)
assert ltx23._is_output_rank() is False
assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is False
assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: True)) is True
monkeypatch.setattr(ltx.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx.torch.distributed, "get_rank", lambda: 1)
pipe = object.__new__(ltx.LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.vae = FakeVae()
output = pipe._decode_output(
latents=torch.ones(1, 1),
audio_latents=torch.ones(1, 1),
output_type="np",
connector_prompt_embeds=torch.ones(1, 1),
generator=None,
device=torch.device("cpu"),
decode_timestep=0.0,
decode_noise_scale=None,
prompt_batch_size=1,
)
assert pipe.vae.decode_calls == expected_decode_calls
assert output.output[0].numel() == 0
assert output.output[1].numel() == 0
class TestLTX23VaeDistributedDecode:
tests/e2e/accuracy/test_ltx2_3_video_similarity.py
+5 -2
View File
@@ -40,7 +40,7 @@ from tests.helpers.runtime import OmniServer
_DIFFUSERS_VERSION = tuple(int(x) for x in diffusers.__version__.split(".")[:2] if x.isdigit())
_DIFFUSERS_038 = _DIFFUSERS_VERSION >= (0, 38)
MODEL_ID = "dg845/LTX-2.3-Diffusers"
MODEL_ID = "diffusers/LTX-2.3-Diffusers"
MODEL_ENV_VAR = "VLLM_TEST_LTX23_MODEL"
PROMPT = "A lighthouse on a rocky cliff at sunset, waves crashing below, golden hour lighting"
NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark"
@@ -216,7 +216,10 @@ def _run_with_custom_transformer(model: str, output_dir: Path) -> list[Image.Ima
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.distributed.parallel_state import init_distributed_environment, initialize_model_parallel
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import create_transformer_from_config, load_transformer_config
from vllm_omni.diffusion.models.ltx2 import (
create_transformer_from_config,
load_transformer_config,
)
vllm_config = VllmConfig()
ctx = set_current_vllm_config(vllm_config)
vllm_omni/diffusion/models/dmd2/mixin.py
+31 -32
View File
@@ -37,43 +37,42 @@ class DMD2PipelineMixin:
def _sanitize_dmd2_request(self, req) -> None:
"""Sanitize CFG-related fields in-place. Works with both OmniDiffusionRequest and DiffusionRequestBatch."""
sp = req.sampling_params
sampling_params = req.sampling_params_list if isinstance(req, DiffusionRequestBatch) else [req.sampling_params]
for sp in sampling_params:
if sp.num_inference_steps and sp.num_inference_steps != self.dmd2_config.num_inference_steps:
logger.warning(
"DMD2: ignoring num_inference_steps=%d, forcing %d.",
sp.num_inference_steps,
self.dmd2_config.num_inference_steps,
)
sp.num_inference_steps = self.dmd2_config.num_inference_steps
if sp.num_inference_steps and sp.num_inference_steps != self.dmd2_config.num_inference_steps:
logger.warning(
"DMD2: ignoring num_inference_steps=%d, forcing %d.",
sp.num_inference_steps,
self.dmd2_config.num_inference_steps,
)
sp.num_inference_steps = self.dmd2_config.num_inference_steps
if sp.guidance_scale_provided and sp.guidance_scale != self.dmd2_config.guidance_scale:
logger.warning(
"DMD2: ignoring guidance_scale=%.2f, forcing %.2f.",
sp.guidance_scale,
self.dmd2_config.guidance_scale,
)
sp.guidance_scale = self.dmd2_config.guidance_scale
sp.guidance_scale_provided = False
if sp.guidance_scale_provided and sp.guidance_scale != self.dmd2_config.guidance_scale:
logger.warning(
"DMD2: ignoring guidance_scale=%.2f, forcing %.2f.",
sp.guidance_scale,
self.dmd2_config.guidance_scale,
)
sp.guidance_scale = self.dmd2_config.guidance_scale
sp.guidance_scale_provided = False
if sp.guidance_scale_2 is not None:
logger.warning("DMD2: ignoring guidance_scale_2.")
sp.guidance_scale_2 = None
if sp.guidance_scale_2 is not None:
logger.warning("DMD2: ignoring guidance_scale_2.")
sp.guidance_scale_2 = None
if sp.true_cfg_scale is not None:
logger.warning("DMD2: ignoring true_cfg_scale.")
sp.true_cfg_scale = None
if sp.true_cfg_scale is not None:
logger.warning("DMD2: ignoring true_cfg_scale.")
sp.true_cfg_scale = None
sp.do_classifier_free_guidance = False
sp.is_cfg_negative = False
sp.do_classifier_free_guidance = False
sp.is_cfg_negative = False
# defense: strip scheduler-override extra_args that would let the base pipeline
# (e.g. Wan22Pipeline.forward) rebuild self.scheduler mid-forward and clobber DMD2EulerScheduler.
extra_args = getattr(sp, "extra_args", None) or {}
for key in ("sample_solver", "flow_shift"):
if key in extra_args:
logger.warning("DMD2: ignoring extra_args.%s.", key)
extra_args.pop(key)
# Prevent base pipelines from replacing the DMD2 scheduler.
extra_args = getattr(sp, "extra_args", None) or {}
for key in ("sample_solver", "flow_shift"):
if key in extra_args:
logger.warning("DMD2: ignoring extra_args.%s.", key)
extra_args.pop(key)
# Strip negative_prompt from each request's prompt in-place.
requests = req.requests if hasattr(req, "requests") else [req]
vllm_omni/diffusion/models/ltx2/__init__.py
+11 -9
View File
@@ -1,23 +1,25 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
LTX2TwoStagesPipeline,
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
create_transformer_from_config,
get_ltx2_post_process_func,
load_transformer_config,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import (
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2ImageToVideoPipeline,
LTX2ImageToVideoTwoStagesPipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
LTX23ImageToVideoPipeline,
LTX23Pipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
LTX2ImageToVideoTwoStagesPipeline,
LTX2TwoStagesPipeline,
)
__all__ = [
"LTX2Pipeline",
vllm_omni/diffusion/models/ltx2/ltx2_components.py
+297
View File
@@ -0,0 +1,297 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared component construction helpers for the LTX model family."""
from __future__ import annotations
import inspect
import json
import os
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import torch
from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.ltx2 import LTX2TextConnectors
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from diffusers.video_processor import VideoProcessor
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from .ltx2_transformer import LTX2VideoTransformer3DModel
try:
from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
except ImportError:
LTX2VocoderWithBWE = None
_LTX_COMPONENT_SUBFOLDERS = (
"tokenizer",
"text_encoder",
"connectors",
"vae",
"audio_vae",
"vocoder",
"scheduler",
)
@dataclass(frozen=True)
class LTXComponentProfile:
"""Component construction and discovery contract for an LTX variant."""
name: str
dit_modules: tuple[str, ...]
encoder_modules: tuple[str, ...]
vae_modules: tuple[str, ...]
resident_modules: tuple[str, ...] = ()
video_vae_cls: type = AutoencoderKLLTX2Video
vocoder_cls: type = LTX2Vocoder
vocoder_fallback_cls: type | None = None
LTX2_COMPONENT_PROFILE = LTXComponentProfile(
name="ltx2",
dit_modules=("transformer",),
encoder_modules=("text_encoder", "connectors"),
vae_modules=("vae", "audio_vae"),
resident_modules=("vocoder",),
video_vae_cls=DistributedAutoencoderKLLTX2Video,
)
LTX23_COMPONENT_PROFILE = LTXComponentProfile(
name="ltx2_3",
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,
)
def _detect_vocoder_output_sample_rate(model: str) -> 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")
except Exception:
return None
try:
with open(vocoder_config_path) as config_file:
return json.load(config_file).get("output_sampling_rate")
except Exception:
return 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)
def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor):
if not (isinstance(output, tuple) and len(output) == 2):
return output
video, audio = output
if isinstance(audio, torch.Tensor):
audio = audio.detach().cpu()
result: dict[str, Any] = {"video": video, "audio": audio}
if output_sample_rate is not None:
result["audio_sample_rate"] = output_sample_rate
return result
return post_process_func
def _load_component(
component_cls: type,
model: str,
subfolder: str,
*,
local_files_only: bool,
dtype: torch.dtype,
) -> Any:
return from_pretrained_with_prefetch(
component_cls.from_pretrained,
model,
subfolder=subfolder,
prefetch_list=_LTX_COMPONENT_SUBFOLDERS,
local_files_only=local_files_only,
torch_dtype=dtype,
)
def _place_aux_components(pipeline: Any) -> None:
parallel_config = getattr(pipeline.od_config, "parallel_config", None)
use_managed_placement = bool(
getattr(pipeline.od_config, "enable_cpu_offload", False)
or getattr(pipeline.od_config, "enable_layerwise_offload", False)
or getattr(parallel_config, "use_hsdp", False)
)
if use_managed_placement:
return
modules = ModuleDiscovery.discover(pipeline)
for module in (*modules.encoders, *modules.vaes, *modules.resident_modules):
module.to(pipeline.device)
def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
"""Build the common LTX component graph selected by ``component_profile``."""
profile: LTXComponentProfile = pipeline.component_profile
pipeline.od_config = od_config
pipeline.device = get_local_device()
dtype = getattr(od_config, "dtype", torch.bfloat16)
model = od_config.model
local_files_only = os.path.exists(model)
pipeline.weights_sources = [
DiffusersPipelineLoader.ComponentSource(
model_or_path=model,
subfolder="transformer",
revision=None,
prefix="transformer.",
fall_back_to_pt=True,
),
]
prefetch_subfolders(model, _LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only)
pipeline.tokenizer = AutoTokenizer.from_pretrained(
model,
subfolder="tokenizer",
local_files_only=local_files_only,
)
with torch.device("cpu"):
pipeline.text_encoder = _load_component(
Gemma3ForConditionalGeneration,
model,
"text_encoder",
local_files_only=local_files_only,
dtype=dtype,
)
pipeline.connectors = _load_component(
LTX2TextConnectors,
model,
"connectors",
local_files_only=local_files_only,
dtype=dtype,
)
pipeline.vae = _load_component(
profile.video_vae_cls,
model,
"vae",
local_files_only=local_files_only,
dtype=dtype,
)
pipeline.audio_vae = _load_component(
AutoencoderKLLTX2Audio,
model,
"audio_vae",
local_files_only=local_files_only,
dtype=dtype,
)
try:
pipeline.vocoder = _load_component(
profile.vocoder_cls,
model,
"vocoder",
local_files_only=local_files_only,
dtype=dtype,
)
except (TypeError, OSError, ValueError):
if profile.vocoder_fallback_cls is None or profile.vocoder_fallback_cls is profile.vocoder_cls:
raise
pipeline.vocoder = _load_component(
profile.vocoder_fallback_cls,
model,
"vocoder",
local_files_only=local_files_only,
dtype=dtype,
)
transformer_config = load_transformer_config(model, "transformer", local_files_only)
quant_config = getattr(od_config, "quantization_config", None)
pipeline.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
_place_aux_components(pipeline)
pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model,
subfolder="scheduler",
local_files_only=local_files_only,
)
pipeline.vae_spatial_compression_ratio = pipeline.vae.spatial_compression_ratio
pipeline.vae_temporal_compression_ratio = pipeline.vae.temporal_compression_ratio
pipeline.audio_vae_mel_compression_ratio = pipeline.audio_vae.mel_compression_ratio
pipeline.audio_vae_temporal_compression_ratio = pipeline.audio_vae.temporal_compression_ratio
pipeline.transformer_spatial_patch_size = pipeline.transformer.config.patch_size
pipeline.transformer_temporal_patch_size = pipeline.transformer.config.patch_size_t
pipeline.audio_sampling_rate = pipeline.audio_vae.config.sample_rate
pipeline.audio_hop_length = pipeline.audio_vae.config.mel_hop_length
pipeline.video_processor = VideoProcessor(vae_scale_factor=pipeline.vae_spatial_compression_ratio)
tokenizer_max_length = pipeline.tokenizer.model_max_length
if tokenizer_max_length is None or tokenizer_max_length > 100000:
encoder_config = getattr(pipeline.text_encoder, "config", None)
tokenizer_max_length = getattr(encoder_config, "max_position_embeddings", None)
if tokenizer_max_length is None:
tokenizer_max_length = getattr(encoder_config, "max_seq_len", None)
pipeline.tokenizer_max_length = int(tokenizer_max_length or 1024)
pipeline._guidance_scale = None
pipeline._guidance_rescale = None
pipeline._interrupt = False
def load_transformer_config(
model_path: str,
subfolder: str = "transformer",
local_files_only: bool = True,
) -> dict:
"""Load an LTX transformer config from a local model or the HF Hub."""
if local_files_only:
config_path = os.path.join(model_path, subfolder, "config.json")
if os.path.exists(config_path):
with open(config_path) as config_file:
return json.load(config_file)
else:
try:
from huggingface_hub import hf_hub_download
config_path = hf_hub_download(
repo_id=model_path,
filename=f"{subfolder}/config.json",
)
with open(config_path) as config_file:
return json.load(config_file)
except Exception:
pass
return {}
def create_transformer_from_config(
config: dict,
quant_config: QuantizationConfig | None = None,
) -> LTX2VideoTransformer3DModel:
"""Construct the shared LTX transformer from a Diffusers config."""
if not config and quant_config is None:
return LTX2VideoTransformer3DModel()
signature = inspect.signature(LTX2VideoTransformer3DModel.__init__)
allowed_keys = set(signature.parameters)
kwargs = {key: value for key, value in config.items() if key in allowed_keys}
if quant_config is not None:
kwargs["quant_config"] = quant_config
return LTX2VideoTransformer3DModel(**kwargs)
vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py
+578
View File
@@ -0,0 +1,578 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Task-specific conditioning shared by LTX model versions."""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import PIL.Image
import torch
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from . import ltx2_latents as latent_ops
from .ltx2_denoise import I2VVideoAudioScheduler
if TYPE_CHECKING:
from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext
from .ltx2_request import LTXRequestInputs
@dataclass
class LTXPromptContext:
"""Connector outputs consumed by an LTX denoise phase."""
batch_size: int
connector_prompt_embeds: torch.Tensor
connector_audio_prompt_embeds: torch.Tensor
connector_attention_mask: torch.Tensor
positive_connector_prompt_embeds: torch.Tensor
positive_connector_audio_prompt_embeds: torch.Tensor
positive_connector_attention_mask: torch.Tensor
negative_connector_prompt_embeds: torch.Tensor | None
negative_connector_audio_prompt_embeds: torch.Tensor | None
negative_connector_attention_mask: torch.Tensor | None
def _repeat_prompt_tensor_for_outputs(tensor: torch.Tensor, num_outputs: int) -> torch.Tensor:
if num_outputs == 1:
return tensor
return tensor.repeat_interleave(num_outputs, dim=0)
class LTXTextConditioningMixin:
"""Shared Gemma encoding and text connector orchestration."""
connector_batches_cfg = False
def _get_gemma_prompt_embeds(
self,
prompt: str | list[str],
num_videos_per_prompt: int = 1,
max_sequence_length: int = 1024,
scale_factor: int = 8,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
del scale_factor
device = device or self.device
dtype = dtype or self.text_encoder.dtype
prompt = [prompt] if isinstance(prompt, str) else prompt
batch_size = len(prompt)
self.tokenizer.padding_side = "left"
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],
padding="max_length",
max_length=max_sequence_length,
truncation=True,
add_special_tokens=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids.to(device)
prompt_attention_mask = text_inputs.attention_mask.to(device)
hidden_states = self.text_encoder(
input_ids=text_input_ids,
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)
prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
prompt_attention_mask,
num_videos_per_prompt,
)
return prompt_embeds, prompt_attention_mask
def encode_prompt(
self,
prompt: str | list[str] | None,
negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
max_sequence_length: int = 1024,
scale_factor: int = 8,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
device = device or self.device
prompt = [prompt] if isinstance(prompt, str) else prompt
batch_size = len(prompt) if prompt is not None else prompt_embeds.shape[0]
negative_prompt_embeds_provided = negative_prompt_embeds is not None
if prompt_embeds is None:
prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds(
prompt=prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
scale_factor=scale_factor,
device=device,
dtype=dtype,
)
elif num_videos_per_prompt > 1:
prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt)
prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
prompt_attention_mask,
num_videos_per_prompt,
)
if do_classifier_free_guidance and negative_prompt_embeds is None:
negative_prompt = negative_prompt or ""
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type as `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
if isinstance(negative_prompt, list) and batch_size != len(negative_prompt):
raise ValueError(
f"`negative_prompt` has batch size {len(negative_prompt)}, but `prompt` has batch size"
f" {batch_size}."
)
negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds(
prompt=negative_prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
scale_factor=scale_factor,
device=device,
dtype=dtype,
)
elif do_classifier_free_guidance and negative_prompt_embeds_provided and num_videos_per_prompt > 1:
negative_prompt_embeds = _repeat_prompt_tensor_for_outputs(
negative_prompt_embeds,
num_videos_per_prompt,
)
negative_prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
negative_prompt_attention_mask,
num_videos_per_prompt,
)
return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
def _prepare_prompt_context(
self,
*,
prompt: str | list[str] | None,
negative_prompt: str | list[str] | None,
prompt_embeds: torch.Tensor | None,
negative_prompt_embeds: torch.Tensor | None,
prompt_attention_mask: torch.Tensor | None,
negative_prompt_attention_mask: torch.Tensor | None,
num_videos_per_prompt: int,
max_sequence_length: int,
) -> LTXPromptContext:
if isinstance(prompt, str):
batch_size = 1
elif isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = (
self.encode_prompt(
prompt=prompt,
negative_prompt=negative_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
num_videos_per_prompt=num_videos_per_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
max_sequence_length=max_sequence_length,
device=self.device,
)
)
padding_side = getattr(self.tokenizer, "padding_side", "left")
if self.do_classifier_free_guidance and self.connector_batches_cfg:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
video_context, audio_context, attention_mask = self.connectors(
prompt_embeds,
prompt_attention_mask,
padding_side=padding_side,
)
positive_video_context = video_context
positive_audio_context = audio_context
positive_attention_mask = attention_mask
negative_video_context = None
negative_audio_context = None
negative_attention_mask = None
if self.do_classifier_free_guidance and self.connector_batches_cfg:
split_batch = batch_size * num_videos_per_prompt
negative_video_context = video_context[:split_batch]
positive_video_context = video_context[split_batch:]
negative_audio_context = audio_context[:split_batch]
positive_audio_context = audio_context[split_batch:]
negative_attention_mask = attention_mask[:split_batch]
positive_attention_mask = attention_mask[split_batch:]
elif self.do_classifier_free_guidance:
negative_video_context, negative_audio_context, negative_attention_mask = self.connectors(
negative_prompt_embeds,
negative_prompt_attention_mask,
padding_side=padding_side,
)
return LTXPromptContext(
batch_size=batch_size,
connector_prompt_embeds=video_context,
connector_audio_prompt_embeds=audio_context,
connector_attention_mask=attention_mask,
positive_connector_prompt_embeds=positive_video_context,
positive_connector_audio_prompt_embeds=positive_audio_context,
positive_connector_attention_mask=positive_attention_mask,
negative_connector_prompt_embeds=negative_video_context,
negative_connector_audio_prompt_embeds=negative_audio_context,
negative_connector_attention_mask=negative_attention_mask,
)
class LTXI2VConditioningMixin:
"""First-frame conditioning behavior common to LTX2 and LTX2.3."""
support_image_input = True
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.video_processor = VideoProcessor(
vae_scale_factor=self.vae_spatial_compression_ratio,
resample="bilinear",
)
def forward(self, req: Any, image: Any | None = None, *args: Any, **kwargs: Any) -> Any:
"""Preserve the public I2V positional API while sharing the T2V runner."""
return super().forward(req, *args, image=image, **kwargs)
@staticmethod
def _resolve_single_prompt_image(raw_image: Any) -> Any:
if isinstance(raw_image, list):
if len(raw_image) != 1:
raise ValueError(
"LTX I2V prompt dictionaries support exactly one image per prompt. "
"Pass one image per prompt for batched I2V requests."
)
return raw_image[0]
return raw_image
@staticmethod
def _resolve_additional_image(additional: dict[str, Any]) -> Any:
for field_name in ("preprocessed_image", "pixel_values", "image"):
raw_image = additional.get(field_name)
if raw_image is not None:
return raw_image
return None
def _resolve_request_image(
self,
req: Any,
image: Any | None,
request_inputs: LTXRequestInputs,
) -> Any | None:
if image is not None or not req.prompts:
return image
raw_images = []
for prompt_item in req.prompts:
if isinstance(prompt_item, str):
raw_image = None
else:
multi_modal_data = prompt_item.get("multi_modal_data") or {}
raw_image = multi_modal_data.get("image")
if raw_image is None:
raw_image = self._resolve_additional_image(prompt_item.get("additional_information") or {})
raw_image = self._resolve_single_prompt_image(raw_image)
if isinstance(raw_image, str):
raw_image = PIL.Image.open(raw_image).convert("RGB")
raw_images.append(raw_image)
if any(raw_image is None for raw_image in raw_images) and request_inputs.latents is None:
raise ValueError("Image is required for LTX I2V generation.")
if len(raw_images) == 1:
return raw_images[0]
return raw_images or image
def _check_forward_inputs(
self,
request_inputs: LTXRequestInputs,
image: Any | None = None,
) -> None:
if image is None and request_inputs.latents is None:
raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.")
super()._check_forward_inputs(request_inputs, image=image)
def prepare_latents(
self,
image: torch.Tensor | None = None,
batch_size: int = 1,
num_channels_latents: int = 128,
height: int = 512,
width: int = 768,
num_frames: int = 121,
noise_scale: float = 0.0,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
num_frames, height, width = latent_ops.resolve_video_latent_shape(
height,
width,
num_frames,
vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
)
shape = (batch_size, num_channels_latents, num_frames, height, width)
mask_shape = (batch_size, 1, num_frames, height, width)
if latents is not None:
if latents.ndim == 5:
batch_size, _, num_frames, height, width = latents.shape
mask_shape = (batch_size, 1, num_frames, height, width)
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
latents = latent_ops.normalize_latents(
latents,
self.vae.latents_mean,
self.vae.latents_std,
self.vae.config.scaling_factor,
)
latents = latent_ops.create_noised_state(
latents,
noise_scale * (1 - conditioning_mask),
generator,
)
latents = latent_ops.pack_latents(
latents,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
else:
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
conditioning_mask = latent_ops.pack_latents(
conditioning_mask,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
).squeeze(-1)
if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
raise ValueError(
"Provided `latents` tensor has shape"
f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
)
return latents.to(device=device, dtype=dtype), conditioning_mask
if image is None:
raise ValueError("`image` must be provided when `latents` is None.")
image_batch_size = image.shape[0]
if image_batch_size == 0:
raise ValueError("`image` batch is empty.")
if batch_size % image_batch_size != 0:
raise ValueError(
f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) "
"for image-to-video outputs."
)
num_videos_per_prompt = batch_size // image_batch_size
if isinstance(generator, list):
if len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective"
f" batch size of {batch_size}. Make sure the batch size matches the length of the generators."
)
image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)]
init_latents = [
retrieve_latents(
self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)),
image_generators[i],
"argmax",
)
for i in range(image_batch_size)
]
else:
init_latents = [
retrieve_latents(
self.vae.encode(img.unsqueeze(0).unsqueeze(2)),
generator,
"argmax",
)
for img in image
]
init_latents = torch.cat(init_latents, dim=0).to(dtype)
if num_videos_per_prompt > 1:
init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0)
init_latents = latent_ops.normalize_latents(
init_latents,
self.vae.latents_mean,
self.vae.latents_std,
)
init_latents = init_latents.repeat(1, 1, num_frames, 1, 1)
conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype)
conditioning_mask[:, :, 0] = 1.0
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)
conditioning_mask = latent_ops.pack_latents(
conditioning_mask,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
).squeeze(-1)
latents = latent_ops.pack_latents(
latents,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
return latents, conditioning_mask
def _step_video_latents_i2v(
self,
noise_pred_video: torch.Tensor,
latents: torch.Tensor,
timestep: torch.Tensor,
latent_num_frames: int,
latent_height: int,
latent_width: int,
) -> torch.Tensor:
noise_pred_video = latent_ops.unpack_latents(
noise_pred_video,
latent_num_frames,
latent_height,
latent_width,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
latents_unpacked = latent_ops.unpack_latents(
latents,
latent_num_frames,
latent_height,
latent_width,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
noise_pred_video = noise_pred_video[:, :, 1:]
noise_latents = latents_unpacked[:, :, 1:]
pred_latents = self.scheduler.step(
noise_pred_video,
timestep,
noise_latents,
return_dict=False,
)[0]
latents_unpacked = torch.cat([latents_unpacked[:, :, :1], pred_latents], dim=2)
return latent_ops.pack_latents(
latents_unpacked,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
def _prepare_video_latents_stage(
self,
request_inputs: LTXRequestInputs,
prompt_context: LTXPromptContext,
*,
device: torch.device,
noise_scale: float,
image: Any | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
if request_inputs.latents is None:
if isinstance(image, torch.Tensor):
if image.ndim == 3:
image = image.unsqueeze(0)
elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor):
image = torch.stack(image, dim=0)
else:
image = self.video_processor.preprocess(
image,
height=request_inputs.height,
width=request_inputs.width,
)
image = image.to(
device=device,
dtype=prompt_context.positive_connector_prompt_embeds.dtype,
)
return self.prepare_latents(
image,
prompt_context.batch_size * request_inputs.num_videos_per_prompt,
self.transformer.config.in_channels,
request_inputs.height,
request_inputs.width,
request_inputs.num_frames,
noise_scale,
torch.float32,
device,
request_inputs.generator,
request_inputs.latents,
)
def _make_video_audio_scheduler(
self,
audio_scheduler: Any,
latent_num_frames: int,
latent_height: int,
latent_width: int,
) -> Any:
return I2VVideoAudioScheduler(
self,
audio_scheduler,
latent_num_frames,
latent_height,
latent_width,
)
def _prepare_denoise_context_for_cfg(
self,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> LTXDenoiseContext:
denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
if denoise_ctx.conditioning_mask is None:
raise ValueError("LTX I2V denoising requires a conditioning mask.")
mask_batch = denoise_ctx.conditioning_mask.shape[0]
model_batch = denoise_ctx.video_coords.shape[0]
if model_batch % mask_batch != 0:
raise ValueError(
"I2V conditioning-mask batch must divide the Transformer input batch, "
f"but got {mask_batch} and {model_batch}."
)
repeats = model_batch // mask_batch
denoise_ctx.conditioning_mask_for_model = (
denoise_ctx.conditioning_mask if repeats == 1 else torch.cat([denoise_ctx.conditioning_mask] * repeats)
)
return denoise_ctx
def _denoise_timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> dict[str, torch.Tensor]:
kwargs = super()._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx)
conditioning_mask = (
denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model
)
if conditioning_mask is None:
raise ValueError("LTX I2V denoising requires a conditioning mask.")
kwargs.update(
timestep=ts.unsqueeze(-1) * (1 - conditioning_mask),
audio_timestep=ts,
)
return kwargs
vllm_omni/diffusion/models/ltx2/ltx2_denoise.py
+407
View File
@@ -0,0 +1,407 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared denoise execution primitives for LTX pipelines."""
from __future__ import annotations
import copy
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Protocol
import numpy as np
import torch
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from .ltx2_latents import LTXAVState
if TYPE_CHECKING:
from .ltx2_conditioning import LTXPromptContext
from .ltx2_request import LTXRequestInputs
@dataclass
class LTXForwardContext:
"""Immutable metadata and schedulers for one LTX denoise phase."""
req: DiffusionRequestBatch
request_inputs: LTXRequestInputs
prompt_context: LTXPromptContext
device: torch.device
cfg_parallel_ready: bool
attention_kwargs: dict[str, Any] | None
latent_num_frames: int
latent_height: int
latent_width: int
latent_mel_bins: int
original_audio_num_frames: int
padded_audio_num_frames: int
timesteps: torch.Tensor
audio_scheduler: Any
video_audio_scheduler: Any
@property
def batch_size(self) -> int:
return self.prompt_context.batch_size
@property
def num_videos_per_prompt(self) -> int:
return self.request_inputs.num_videos_per_prompt
@dataclass
class LTXDenoiseContext:
"""Mutable AV state and positional metadata for a denoise phase."""
latents: torch.Tensor
audio_latents: torch.Tensor
video_coords: torch.Tensor
audio_coords: torch.Tensor
conditioning_mask: torch.Tensor | None = None
conditioning_mask_for_model: torch.Tensor | None = None
@dataclass
class LTXPhaseResult:
"""Denoised, unpacked AV latents and the context used to produce them."""
forward_context: LTXForwardContext
video: torch.Tensor
audio: torch.Tensor
class LTXDenoisePipeline(Protocol):
"""Pipeline state required by :class:`LTXDenoiseExecutor`."""
@property
def interrupt(self) -> bool: ...
def progress_bar(self, iterable=None, total=None): ...
LTXDenoiseStep = Callable[[int, torch.Tensor, LTXAVState], LTXAVState]
class LTXDenoiseExecutor:
"""Run the one shared LTX denoise loop.
Prediction and scheduler math remain injectable so structural refactors do
not change the existing LTX2/LTX2.3 numerical paths. Guidance will replace
that step policy independently.
"""
@staticmethod
def run(
pipeline: LTXDenoisePipeline,
state: LTXAVState,
timesteps: Iterable[torch.Tensor],
step: LTXDenoiseStep,
) -> LTXAVState:
timesteps = tuple(timesteps)
with pipeline.progress_bar(total=len(timesteps)) as progress_bar:
for index, timestep in enumerate(timesteps):
if pipeline.interrupt:
continue
state = step(index, timestep, state)
progress_bar.update()
return state
def calculate_shift(
image_seq_len: int,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.15,
) -> float:
slope = (max_shift - base_shift) / (max_seq_len - base_seq_len)
intercept = base_shift - slope * base_seq_len
return image_seq_len * slope + intercept
class VideoAudioScheduler:
"""Composite scheduler dispatching video and audio updates."""
def __init__(self, video_scheduler, audio_scheduler):
self.video_scheduler = video_scheduler
self.audio_scheduler = audio_scheduler
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
video_out = self.video_scheduler.step(
noise_pred[0],
t[0],
latents[0],
return_dict=False,
generator=generator,
)[0]
audio_out = self.audio_scheduler.step(
noise_pred[1],
t[1],
latents[1],
return_dict=False,
generator=generator,
)[0]
return ((video_out, audio_out),)
class I2VVideoAudioScheduler:
"""Update the unconditioned video frames and the full audio state."""
def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width):
self.video_scheduler = pipeline.scheduler
self.audio_scheduler = audio_scheduler
self._pipeline = pipeline
self._latent_num_frames = latent_num_frames
self._latent_height = latent_height
self._latent_width = latent_width
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
video_out = self._pipeline._step_video_latents_i2v(
noise_pred[0],
latents[0],
t[0],
self._latent_num_frames,
self._latent_height,
self._latent_width,
)
audio_out = self.audio_scheduler.step(
noise_pred[1],
t[1],
latents[1],
return_dict=False,
generator=generator,
)[0]
return ((video_out, audio_out),)
def prepare_scheduler_stage(
pipeline: Any,
request_inputs: LTXRequestInputs,
*,
device: torch.device,
sigmas: list[float] | None,
timesteps: list[int] | None,
latent_num_frames: int,
latent_height: int,
latent_width: int,
) -> tuple[Any, Any, torch.Tensor]:
sigmas = (
np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps)
if sigmas is None
else sigmas
)
mu = calculate_shift(
pipeline.scheduler.config.get("max_image_seq_len", 4096),
pipeline.scheduler.config.get("base_image_seq_len", 1024),
pipeline.scheduler.config.get("max_image_seq_len", 4096),
pipeline.scheduler.config.get("base_shift", 0.95),
pipeline.scheduler.config.get("max_shift", 2.05),
)
audio_scheduler = copy.deepcopy(pipeline.scheduler)
video_audio_scheduler = pipeline._make_video_audio_scheduler(
audio_scheduler,
latent_num_frames,
latent_height,
latent_width,
)
retrieve_timesteps(
audio_scheduler,
request_inputs.num_inference_steps,
device,
timesteps,
sigmas=sigmas,
mu=mu,
)
timesteps_tensor, _ = retrieve_timesteps(
pipeline.scheduler,
request_inputs.num_inference_steps,
device,
timesteps,
sigmas=sigmas,
mu=mu,
)
return audio_scheduler, video_audio_scheduler, timesteps_tensor
def prepare_rope_coords_stage(
pipeline: Any,
forward_ctx: LTXForwardContext,
latents: torch.Tensor,
audio_latents: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
video_coords = pipeline.transformer.rope.prepare_video_coords(
latents.shape[0],
forward_ctx.latent_num_frames,
forward_ctx.latent_height,
forward_ctx.latent_width,
latents.device,
fps=forward_ctx.request_inputs.frame_rate,
)
audio_coords = pipeline.transformer.audio_rope.prepare_audio_coords(
audio_latents.shape[0],
forward_ctx.padded_audio_num_frames,
audio_latents.device,
)
return video_coords, audio_coords
def build_transformer_kwargs(
pipeline: Any,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
*,
hidden_states: torch.Tensor,
audio_hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
audio_encoder_hidden_states: torch.Tensor,
encoder_attention_mask: torch.Tensor | None,
audio_encoder_attention_mask: torch.Tensor | None,
ts: torch.Tensor,
) -> dict[str, Any]:
return {
"hidden_states": hidden_states,
"audio_hidden_states": audio_hidden_states,
"encoder_hidden_states": encoder_hidden_states,
"audio_encoder_hidden_states": audio_encoder_hidden_states,
**pipeline._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx),
"encoder_attention_mask": encoder_attention_mask,
"audio_encoder_attention_mask": audio_encoder_attention_mask,
"num_frames": forward_ctx.latent_num_frames,
"height": forward_ctx.latent_height,
"width": forward_ctx.latent_width,
"fps": forward_ctx.request_inputs.frame_rate,
"audio_num_frames": forward_ctx.padded_audio_num_frames,
"video_coords": denoise_ctx.video_coords,
"audio_coords": denoise_ctx.audio_coords,
"attention_kwargs": forward_ctx.attention_kwargs,
"return_dict": False,
}
def step_denoised_latents(
pipeline: Any,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
noise_pred_video: torch.Tensor,
noise_pred_audio: torch.Tensor,
timestep: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
latents = pipeline.scheduler_step_maybe_with_cfg(
(noise_pred_video, noise_pred_audio),
(timestep, timestep),
(denoise_ctx.latents, denoise_ctx.audio_latents),
do_true_cfg=pipeline.do_classifier_free_guidance,
per_request_scheduler=forward_ctx.video_audio_scheduler,
)
return pipeline._synchronize_cfg_parallel_step_output(
latents,
do_true_cfg=pipeline.do_classifier_free_guidance,
)
class LTXPhaseExecutor:
"""Prepare and execute one LTX phase without owning model modules."""
@staticmethod
def run(
pipeline: Any,
req: DiffusionRequestBatch,
request_inputs: LTXRequestInputs,
*,
noise_scale: float,
sigmas: list[float] | None,
timesteps: list[int] | None,
attention_kwargs: dict[str, Any] | None,
image: Any | None = None,
prompt_context: LTXPromptContext | None = None,
) -> LTXPhaseResult:
pipeline._check_forward_inputs(request_inputs, image=image)
cfg_parallel_ready = pipeline._setup_forward_runtime(request_inputs, attention_kwargs)
device = pipeline.device
if prompt_context is None:
prompt_context = pipeline._prepare_prompt_context(
prompt=request_inputs.prompt,
negative_prompt=request_inputs.negative_prompt,
prompt_embeds=request_inputs.prompt_embeds,
negative_prompt_embeds=request_inputs.negative_prompt_embeds,
prompt_attention_mask=request_inputs.prompt_attention_mask,
negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask,
num_videos_per_prompt=request_inputs.num_videos_per_prompt,
max_sequence_length=request_inputs.max_sequence_length,
)
latent_num_frames, latent_height, latent_width = pipeline._resolve_video_latent_dimensions(request_inputs)
latents, conditioning_mask = pipeline._prepare_video_latents_stage(
request_inputs,
prompt_context,
device=device,
noise_scale=noise_scale,
image=image,
)
audio_latents, original_audio_num_frames, padded_audio_num_frames, latent_mel_bins = (
pipeline._prepare_audio_latents_stage(
request_inputs,
prompt_context,
device=device,
noise_scale=noise_scale,
)
)
audio_scheduler, video_audio_scheduler, timesteps_tensor = prepare_scheduler_stage(
pipeline,
request_inputs,
device=device,
sigmas=sigmas,
timesteps=timesteps,
latent_num_frames=latent_num_frames,
latent_height=latent_height,
latent_width=latent_width,
)
forward_ctx = LTXForwardContext(
req=req,
request_inputs=request_inputs,
prompt_context=prompt_context,
device=device,
cfg_parallel_ready=cfg_parallel_ready,
attention_kwargs=attention_kwargs,
latent_num_frames=latent_num_frames,
latent_height=latent_height,
latent_width=latent_width,
latent_mel_bins=latent_mel_bins,
original_audio_num_frames=original_audio_num_frames,
padded_audio_num_frames=padded_audio_num_frames,
timesteps=timesteps_tensor,
audio_scheduler=audio_scheduler,
video_audio_scheduler=video_audio_scheduler,
)
video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents)
denoise_ctx = LTXDenoiseContext(
latents=latents,
audio_latents=audio_latents,
video_coords=video_coords,
audio_coords=audio_coords,
conditioning_mask=conditioning_mask,
)
denoise_ctx = pipeline._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
state = LTXDenoiseExecutor.run(
pipeline,
LTXAVState(video=denoise_ctx.latents, audio=denoise_ctx.audio_latents),
forward_ctx.timesteps,
lambda index, timestep, state: pipeline._denoise_step(
index,
timestep,
state,
forward_ctx,
denoise_ctx,
),
)
denoise_ctx.latents = state.video
denoise_ctx.audio_latents = state.audio
latents, audio_latents = pipeline._unpack_and_denormalize_stage(
forward_ctx,
state.video,
state.audio,
)
return LTXPhaseResult(forward_context=forward_ctx, video=latents, audio=audio_latents)
vllm_omni/diffusion/models/ltx2/ltx2_guidance.py
+379
View File
@@ -0,0 +1,379 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Guidance strategies for the LTX model family."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Protocol
import torch
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
get_cfg_group,
get_classifier_free_guidance_rank,
get_classifier_free_guidance_world_size,
)
if TYPE_CHECKING:
from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext
from .ltx2_latents import LTXAVState
class LTXGuidanceStrategy(Protocol):
"""Model-version guidance behavior consumed by the shared pipeline."""
def validate_cfg_world_size(self, cfg_world_size: int) -> None: ...
def prepare_denoise_context(
self,
pipeline: Any,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> LTXDenoiseContext: ...
def timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> dict[str, torch.Tensor]: ...
def combine_cfg_noise(
self,
pipeline: Any,
positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
true_cfg_scale: float,
cfg_normalize: bool,
context: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]: ...
def predict_noise(
self,
pipeline: Any,
index: int,
timestep: torch.Tensor,
state: LTXAVState,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> tuple[torch.Tensor, torch.Tensor]: ...
class _LTXGuidanceStrategyBase:
def validate_cfg_world_size(self, cfg_world_size: int) -> None:
del cfg_world_size
def prepare_denoise_context(
self,
pipeline: Any,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> LTXDenoiseContext:
del pipeline, forward_ctx
return denoise_ctx
def timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> dict[str, torch.Tensor]:
del forward_ctx, denoise_ctx
return {"timestep": ts}
def combine_velocity_via_x0(
sample: torch.Tensor,
positive_velocity: torch.Tensor,
negative_velocity: torch.Tensor,
sigma: torch.Tensor,
guidance_scale: float,
) -> torch.Tensor:
"""Apply CFG in x0 space and convert the result back to velocity."""
x0_cond = sample - positive_velocity * sigma
x0_uncond = sample - negative_velocity * sigma
x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond)
return (sample - x0_guided) / sigma
class LTXLegacyVelocityGuidance(_LTXGuidanceStrategyBase):
"""Existing LTX2 Diffusers-style velocity-space CFG behavior."""
def combine_cfg_noise(
self,
pipeline: Any,
positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
true_cfg_scale: float,
cfg_normalize: bool,
context: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
del context
video_pos, audio_pos = positive_noise_pred
video_neg, audio_neg = negative_noise_pred
video_combined = CFGParallelMixin.combine_cfg_noise(
pipeline,
video_pos,
video_neg,
true_cfg_scale,
cfg_normalize,
)
audio_combined = CFGParallelMixin.combine_cfg_noise(
pipeline,
audio_pos,
audio_neg,
true_cfg_scale,
cfg_normalize,
)
if pipeline.guidance_rescale and pipeline.guidance_rescale > 0:
video_combined = rescale_noise_cfg(
video_combined,
video_pos,
guidance_rescale=pipeline.guidance_rescale,
)
audio_combined = rescale_noise_cfg(
audio_combined,
audio_pos,
guidance_rescale=pipeline.guidance_rescale,
)
return video_combined, audio_combined
def predict_noise(
self,
pipeline: Any,
index: int,
timestep: torch.Tensor,
state: LTXAVState,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> tuple[torch.Tensor, torch.Tensor]:
del index
prompt_context = forward_ctx.prompt_context
video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype)
audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype)
expanded_timestep = timestep.expand(video_input.shape[0])
positive_kwargs = pipeline._build_transformer_kwargs(
forward_ctx,
denoise_ctx,
hidden_states=video_input,
audio_hidden_states=audio_input,
encoder_hidden_states=prompt_context.positive_connector_prompt_embeds,
audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds,
encoder_attention_mask=prompt_context.positive_connector_attention_mask,
audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask,
ts=expanded_timestep,
)
negative_kwargs = (
{
**positive_kwargs,
"encoder_hidden_states": prompt_context.negative_connector_prompt_embeds,
"audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds,
"encoder_attention_mask": prompt_context.negative_connector_attention_mask,
"audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask,
}
if pipeline.do_classifier_free_guidance
else None
)
return pipeline.predict_noise_maybe_with_cfg(
do_true_cfg=pipeline.do_classifier_free_guidance,
true_cfg_scale=forward_ctx.request_inputs.guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
class LTXOfficialX0Guidance(_LTXGuidanceStrategyBase):
"""Official LTX x0-space CFG behavior currently used by LTX2.3."""
def validate_cfg_world_size(self, cfg_world_size: int) -> None:
if cfg_world_size not in (1, 2):
raise ValueError(
f"LTX x0 guidance supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}."
)
def prepare_denoise_context(
self,
pipeline: Any,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> LTXDenoiseContext:
if pipeline.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready:
denoise_ctx.video_coords = denoise_ctx.video_coords.repeat(
(2,) + (1,) * (denoise_ctx.video_coords.ndim - 1)
)
denoise_ctx.audio_coords = denoise_ctx.audio_coords.repeat(
(2,) + (1,) * (denoise_ctx.audio_coords.ndim - 1)
)
return denoise_ctx
def timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> dict[str, torch.Tensor]:
del forward_ctx, denoise_ctx
return {"timestep": ts, "sigma": ts}
def combine_cfg_noise(
self,
pipeline: Any,
positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
true_cfg_scale: float,
cfg_normalize: bool,
context: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
required = ("video_latents", "audio_latents", "video_sigma", "audio_sigma")
if any(context.get(name) is None for name in required):
raise ValueError("LTX x0-space CFG requires video/audio latents and sigmas.")
video_pos, audio_pos = positive_noise_pred
video_neg, audio_neg = negative_noise_pred
video_combined = combine_velocity_via_x0(
context["video_latents"],
video_pos,
video_neg,
context["video_sigma"],
true_cfg_scale,
)
audio_combined = combine_velocity_via_x0(
context["audio_latents"],
audio_pos,
audio_neg,
context["audio_sigma"],
true_cfg_scale,
)
if cfg_normalize:
video_combined = pipeline.cfg_normalize_function(video_pos, video_combined)
audio_combined = pipeline.cfg_normalize_function(audio_pos, audio_combined)
return video_combined, audio_combined
def predict_parallel_cfg(
self,
pipeline: Any,
true_cfg_scale: float,
positive_kwargs: dict[str, Any],
negative_kwargs: dict[str, Any],
cfg_normalize: bool = True,
output_slice: int | None = None,
**context: Any,
) -> tuple[torch.Tensor, torch.Tensor]:
def maybe_slice(pred: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
if output_slice is None:
return pred
return pred[0][:, :output_slice], pred[1][:, :output_slice]
cfg_world_size = get_classifier_free_guidance_world_size()
self.validate_cfg_world_size(cfg_world_size)
if cfg_world_size != 2:
raise ValueError(f"Parallel LTX x0 guidance requires cfg_parallel_size 2, but got {cfg_world_size}.")
cfg_rank = get_classifier_free_guidance_rank()
branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
local_video_pred, local_audio_pred = maybe_slice(pipeline.predict_noise(**branch_kwargs))
cfg_group = get_cfg_group()
gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True)
gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True)
return self.combine_cfg_noise(
pipeline,
(gathered_video[0], gathered_audio[0]),
(gathered_video[1], gathered_audio[1]),
true_cfg_scale,
cfg_normalize,
context,
)
def predict_noise(
self,
pipeline: Any,
index: int,
timestep: torch.Tensor,
state: LTXAVState,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> tuple[torch.Tensor, torch.Tensor]:
prompt_context = forward_ctx.prompt_context
guidance_scale = forward_ctx.request_inputs.guidance_scale
audio_scheduler = forward_ctx.audio_scheduler
if forward_ctx.cfg_parallel_ready:
video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype)
audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype)
ts = timestep.expand(video_input.shape[0])
positive_kwargs = pipeline._build_transformer_kwargs(
forward_ctx,
denoise_ctx,
hidden_states=video_input,
audio_hidden_states=audio_input,
encoder_hidden_states=prompt_context.positive_connector_prompt_embeds,
audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds,
encoder_attention_mask=prompt_context.positive_connector_attention_mask,
audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask,
ts=ts,
)
negative_kwargs = {
**positive_kwargs,
"encoder_hidden_states": prompt_context.negative_connector_prompt_embeds,
"audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds,
"encoder_attention_mask": prompt_context.negative_connector_attention_mask,
"audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask,
}
return self.predict_parallel_cfg(
pipeline,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
video_latents=state.video,
audio_latents=state.audio,
video_sigma=pipeline.scheduler.sigmas[index],
audio_sigma=audio_scheduler.sigmas[index],
)
video_input = torch.cat([state.video] * 2) if pipeline.do_classifier_free_guidance else state.video
video_input = video_input.to(prompt_context.connector_prompt_embeds.dtype)
audio_input = torch.cat([state.audio] * 2) if pipeline.do_classifier_free_guidance else state.audio
audio_input = audio_input.to(prompt_context.connector_prompt_embeds.dtype)
ts = timestep.expand(video_input.shape[0])
transformer_kwargs = pipeline._build_transformer_kwargs(
forward_ctx,
denoise_ctx,
hidden_states=video_input,
audio_hidden_states=audio_input,
encoder_hidden_states=prompt_context.connector_prompt_embeds,
audio_encoder_hidden_states=prompt_context.connector_audio_prompt_embeds,
encoder_attention_mask=prompt_context.connector_attention_mask,
audio_encoder_attention_mask=prompt_context.connector_attention_mask,
ts=ts,
)
with pipeline._transformer_cache_context("cond_uncond"):
video_pred, audio_pred = pipeline.transformer(**transformer_kwargs)
video_pred = video_pred.float()
audio_pred = audio_pred.float()
if pipeline.do_classifier_free_guidance:
video_uncond, video_cond = video_pred.chunk(2)
video_pred = combine_velocity_via_x0(
state.video,
video_cond,
video_uncond,
pipeline.scheduler.sigmas[index],
guidance_scale,
)
audio_uncond, audio_cond = audio_pred.chunk(2)
audio_pred = combine_velocity_via_x0(
state.audio,
audio_cond,
audio_uncond,
audio_scheduler.sigmas[index],
guidance_scale,
)
return video_pred, audio_pred
LTX_LEGACY_VELOCITY_GUIDANCE = LTXLegacyVelocityGuidance()
LTX_OFFICIAL_X0_GUIDANCE = LTXOfficialX0Guidance()
vllm_omni/diffusion/models/ltx2/ltx2_latents.py
+293
View File
@@ -0,0 +1,293 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared latent layout and normalization primitives for LTX pipelines."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
import torch
from diffusers.utils.torch_utils import randn_tensor
@dataclass
class LTXAVState:
"""Packed video and audio latents carried between denoise steps."""
video: torch.Tensor
audio: torch.Tensor
def pack_latents(
latents: torch.Tensor,
patch_size: int = 1,
patch_size_t: int = 1,
) -> torch.Tensor:
batch_size, _, num_frames, height, width = latents.shape
post_patch_num_frames = num_frames // patch_size_t
post_patch_height = height // patch_size
post_patch_width = width // patch_size
latents = latents.reshape(
batch_size,
-1,
post_patch_num_frames,
patch_size_t,
post_patch_height,
patch_size,
post_patch_width,
patch_size,
)
return latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
def unpack_latents(
latents: torch.Tensor,
num_frames: int,
height: int,
width: int,
patch_size: int = 1,
patch_size_t: int = 1,
) -> torch.Tensor:
batch_size = latents.size(0)
latents = latents.reshape(
batch_size,
num_frames,
height,
width,
-1,
patch_size_t,
patch_size,
patch_size,
)
return latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
def normalize_latents(
latents: torch.Tensor,
latents_mean: torch.Tensor,
latents_std: torch.Tensor,
scaling_factor: float = 1.0,
) -> torch.Tensor:
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
return (latents - latents_mean) * scaling_factor / latents_std
def normalize_audio_latents(
latents: torch.Tensor,
latents_mean: torch.Tensor,
latents_std: torch.Tensor,
) -> torch.Tensor:
latents_mean = latents_mean.to(latents.device, latents.dtype)
latents_std = latents_std.to(latents.device, latents.dtype)
return (latents - latents_mean) / latents_std
def denormalize_latents(
latents: torch.Tensor,
latents_mean: torch.Tensor,
latents_std: torch.Tensor,
scaling_factor: float = 1.0,
) -> torch.Tensor:
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
return latents * latents_std / scaling_factor + latents_mean
def denormalize_audio_latents(
latents: torch.Tensor,
latents_mean: torch.Tensor,
latents_std: torch.Tensor,
) -> torch.Tensor:
latents_mean = latents_mean.to(latents.device, latents.dtype)
latents_std = latents_std.to(latents.device, latents.dtype)
return latents * latents_std + latents_mean
def create_noised_state(
latents: torch.Tensor,
noise_scale: float | torch.Tensor,
generator: torch.Generator | list[torch.Generator] | None = None,
) -> torch.Tensor:
noise = randn_tensor(
latents.shape,
generator=generator,
device=latents.device,
dtype=latents.dtype,
)
return noise_scale * noise + (1 - noise_scale) * latents
def pack_audio_latents(
latents: torch.Tensor,
patch_size: int | None = None,
patch_size_t: int | None = None,
) -> torch.Tensor:
if patch_size is not None and patch_size_t is not None:
batch_size, _, latent_length, latent_mel_bins = latents.shape
post_patch_latent_length = latent_length / patch_size_t
post_patch_mel_bins = latent_mel_bins / patch_size
latents = latents.reshape(
batch_size,
-1,
post_patch_latent_length,
patch_size_t,
post_patch_mel_bins,
patch_size,
)
return latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
return latents.transpose(1, 2).flatten(2, 3)
def unpack_audio_latents(
latents: torch.Tensor,
latent_length: int,
num_mel_bins: int,
patch_size: int | None = None,
patch_size_t: int | None = None,
) -> torch.Tensor:
if patch_size is not None and patch_size_t is not None:
batch_size = latents.size(0)
latents = latents.reshape(
batch_size,
latent_length,
num_mel_bins,
-1,
patch_size_t,
patch_size,
)
return latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
return latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2)
def unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
return latents[:, :num_frames]
def get_sp_padded_audio_latent_length(audio_latent_length: int, sp_size: int) -> int:
if sp_size > 1:
audio_latent_length += (sp_size - (audio_latent_length % sp_size)) % sp_size
return audio_latent_length
def resolve_video_latent_shape(
height: int,
width: int,
num_frames: int,
*,
vae_spatial_compression_ratio: int,
vae_temporal_compression_ratio: int,
) -> tuple[int, int, int]:
return (
(num_frames - 1) // vae_temporal_compression_ratio + 1,
height // vae_spatial_compression_ratio,
width // vae_spatial_compression_ratio,
)
def prepare_video_latents(
pipeline: Any,
batch_size: int = 1,
num_channels_latents: int = 128,
height: int = 512,
width: int = 768,
num_frames: int = 121,
noise_scale: float = 0.0,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
if latents.ndim == 5:
latents = normalize_latents(
latents,
pipeline.vae.latents_mean,
pipeline.vae.latents_std,
pipeline.vae.config.scaling_factor,
)
latents = pack_latents(
latents,
pipeline.transformer_spatial_patch_size,
pipeline.transformer_temporal_patch_size,
)
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)
num_frames, height, width = resolve_video_latent_shape(
height,
width,
num_frames,
vae_spatial_compression_ratio=pipeline.vae_spatial_compression_ratio,
vae_temporal_compression_ratio=pipeline.vae_temporal_compression_ratio,
)
shape = (batch_size, num_channels_latents, num_frames, height, width)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
return pack_latents(
latents,
pipeline.transformer_spatial_patch_size,
pipeline.transformer_temporal_patch_size,
)
def prepare_audio_latents(
pipeline: Any,
batch_size: int = 1,
num_channels_latents: int = 8,
audio_latent_length: int = 1,
num_mel_bins: int = 64,
noise_scale: float = 0.0,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
) -> tuple[torch.Tensor, int, int]:
original_latent_length = audio_latent_length
latent_mel_bins = num_mel_bins // pipeline.audio_vae_mel_compression_ratio
sp_size = getattr(pipeline.od_config.parallel_config, "sequence_parallel_size", 1) or 1
padded_latent_length = get_sp_padded_audio_latent_length(original_latent_length, int(sp_size))
if latents is not None:
if latents.ndim == 4:
latents = pack_audio_latents(latents)
if latents.ndim != 3:
raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].")
latents = normalize_audio_latents(
latents,
pipeline.audio_vae.latents_mean,
pipeline.audio_vae.latents_std,
)
latents = create_noised_state(latents, noise_scale, generator)
if latents.shape[1] not in {original_latent_length, padded_latent_length}:
raise ValueError(
"Provided `audio_latents` has incompatible audio frame count "
f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}."
)
if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length:
padding = torch.zeros(
latents.shape[0],
padded_latent_length - original_latent_length,
latents.shape[2],
dtype=latents.dtype,
device=latents.device,
)
latents = torch.cat([latents, padding], dim=1)
return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length
shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
return pack_audio_latents(latents), original_latent_length, padded_latent_length
vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py
+619
View File
@@ -0,0 +1,619 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared runtime surface for LTX pipeline variants."""
from __future__ import annotations
from collections.abc import Iterable
from contextlib import nullcontext
from typing import Any, ClassVar
import torch
from diffusers.utils.torch_utils import randn_tensor
from torch import nn
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
from . import ltx2_latents as latent_ops
from .ltx2_components import LTXComponentProfile, initialize_pipeline_components
from .ltx2_conditioning import LTXPromptContext, LTXTextConditioningMixin
from .ltx2_denoise import (
LTXDenoiseContext,
LTXForwardContext,
LTXPhaseExecutor,
LTXPhaseResult,
VideoAudioScheduler,
build_transformer_kwargs,
step_denoised_latents,
)
from .ltx2_guidance import LTXGuidanceStrategy
from .ltx2_recipes import LTXOneStageRecipe
from .ltx2_request import LTXRequestInputs, LTXRequestMixin
def _expand_per_prompt_decode_value(
value: float | list[float],
*,
prompt_batch_size: int,
effective_batch_size: int,
field_name: str,
) -> list[float]:
if not isinstance(value, list):
return [value] * effective_batch_size
if len(value) == 1:
return value * effective_batch_size
if len(value) == effective_batch_size:
return value
if prompt_batch_size > 0 and len(value) == prompt_batch_size and effective_batch_size % prompt_batch_size == 0:
repeats = effective_batch_size // prompt_batch_size
return [item for item in value for _ in range(repeats)]
raise ValueError(
f"`{field_name}` must have length 1, prompt batch size ({prompt_batch_size}), or effective batch size"
f" ({effective_batch_size}); got {len(value)}."
)
def _prepare_decode_timestep_conditioning(
*,
decode_timestep: float | list[float],
decode_noise_scale: float | list[float] | None,
prompt_batch_size: int,
effective_batch_size: int,
device: torch.device,
dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
decode_timestep_values = _expand_per_prompt_decode_value(
decode_timestep,
prompt_batch_size=prompt_batch_size,
effective_batch_size=effective_batch_size,
field_name="decode_timestep",
)
decode_noise_scale_values = (
decode_timestep_values
if decode_noise_scale is None
else _expand_per_prompt_decode_value(
decode_noise_scale,
prompt_batch_size=prompt_batch_size,
effective_batch_size=effective_batch_size,
field_name="decode_noise_scale",
)
)
return (
torch.tensor(decode_timestep_values, device=device, dtype=dtype),
torch.tensor(decode_noise_scale_values, device=device, dtype=dtype)[:, None, None, None, None],
)
class LTXPipelineRuntime(
LTXRequestMixin,
LTXTextConditioningMixin,
nn.Module,
CFGParallelMixin,
ProgressBarMixin,
SupportsComponentDiscovery,
DiffusionPipelineProfilerMixin,
):
"""Shared Omni runtime for explicitly composed LTX denoise phases."""
component_profile: ClassVar[LTXComponentProfile]
guidance_strategy: ClassVar[LTXGuidanceStrategy]
one_stage_recipe: ClassVar[LTXOneStageRecipe]
supports_request_batch = False
connector_batches_cfg = False
distributed_video_decode = True
support_image_input = False
dummy_run_num_frames = 2
preserve_sp_padded_audio_duration = False
reports_stage_durations = False
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
del prefix
super().__init__()
initialize_pipeline_components(self, od_config)
self.setup_diffusion_pipeline_profiler(
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
)
def prepare_latents(
self,
batch_size: int = 1,
num_channels_latents: int = 128,
height: int = 512,
width: int = 768,
num_frames: int = 121,
noise_scale: float = 0.0,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
) -> torch.Tensor:
return latent_ops.prepare_video_latents(
self,
batch_size,
num_channels_latents,
height,
width,
num_frames,
noise_scale,
dtype,
device,
generator,
latents,
)
def prepare_audio_latents(
self,
batch_size: int = 1,
num_channels_latents: int = 8,
audio_latent_length: int = 1,
num_mel_bins: int = 64,
noise_scale: float = 0.0,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
) -> tuple[torch.Tensor, int, int]:
return latent_ops.prepare_audio_latents(
self,
batch_size,
num_channels_latents,
audio_latent_length,
num_mel_bins,
noise_scale,
dtype,
device,
generator,
latents,
)
@property
def guidance_scale(self):
return self._guidance_scale
@property
def guidance_rescale(self):
return self._guidance_rescale
@property
def do_classifier_free_guidance(self):
return self._guidance_scale is not None and self._guidance_scale > 1.0
@property
def interrupt(self):
return self._interrupt
def _transformer_cache_context(self, context_name: str):
cache_context = getattr(self.transformer, "cache_context", None)
if callable(cache_context):
return cache_context(context_name)
return nullcontext()
def predict_noise(self, **kwargs):
with self._transformer_cache_context("cond_uncond"):
noise_pred_video, noise_pred_audio = self.transformer(**kwargs)
return noise_pred_video.float(), noise_pred_audio.float()
def combine_cfg_noise(
self,
positive_noise_pred,
negative_noise_pred,
true_cfg_scale,
cfg_normalize=False,
kwargs: dict[str, Any] | None = None,
**context: Any,
):
if kwargs is not None:
context = {**kwargs, **context}
return self.guidance_strategy.combine_cfg_noise(
self,
positive_noise_pred,
negative_noise_pred,
true_cfg_scale,
cfg_normalize,
context,
)
def predict_noise_with_parallel_cfg(self, *args, **kwargs):
predict_parallel_cfg = getattr(self.guidance_strategy, "predict_parallel_cfg", None)
if predict_parallel_cfg is None:
raise NotImplementedError("The selected LTX guidance strategy does not implement parallel CFG.")
return predict_parallel_cfg(self, *args, **kwargs)
def _synchronize_cfg_parallel_step_output(
self,
latents: tuple[torch.Tensor, torch.Tensor],
do_true_cfg: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1):
return latents
# CUDA async execution otherwise permits numerical drift to accumulate
# across CFG-parallel denoise steps.
latents = tuple(tensor.contiguous() for tensor in latents)
device = next((tensor.device for tensor in latents if tensor.is_cuda), None)
if device is not None:
torch.cuda.current_stream(device).synchronize()
return latents
def _setup_forward_runtime(
self,
request_inputs: LTXRequestInputs,
attention_kwargs: dict[str, Any] | None,
) -> bool:
self._guidance_scale = request_inputs.guidance_scale
self._guidance_rescale = request_inputs.guidance_rescale
del attention_kwargs
self._interrupt = False
cfg_world_size = get_classifier_free_guidance_world_size()
if self.do_classifier_free_guidance:
self.guidance_strategy.validate_cfg_world_size(cfg_world_size)
return self.do_classifier_free_guidance and cfg_world_size > 1
def _check_forward_inputs(
self,
request_inputs: LTXRequestInputs,
image: Any | None = None,
) -> None:
self.check_inputs(
prompt=request_inputs.prompt,
height=request_inputs.height,
width=request_inputs.width,
prompt_embeds=request_inputs.prompt_embeds,
negative_prompt_embeds=request_inputs.negative_prompt_embeds,
prompt_attention_mask=request_inputs.prompt_attention_mask,
negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask,
)
def _resolve_request_image(
self,
req: DiffusionRequestBatch,
image: Any | None,
request_inputs: LTXRequestInputs,
) -> Any | None:
del req, request_inputs
return image
def _make_output(self, output: tuple[torch.Tensor, torch.Tensor]) -> DiffusionOutput:
if self.reports_stage_durations:
return DiffusionOutput(
output=output,
stage_durations=getattr(self, "stage_durations", None),
)
return DiffusionOutput(output=output)
def _decode_output(
self,
*,
latents: torch.Tensor,
audio_latents: torch.Tensor,
output_type: str,
connector_prompt_embeds: torch.Tensor,
generator: torch.Generator | list[torch.Generator] | None,
device: torch.device,
decode_timestep: float | list[float],
decode_noise_scale: float | list[float] | None,
prompt_batch_size: int,
) -> DiffusionOutput:
if output_type == "latent":
return self._make_output((latents, audio_latents))
latents = latents.to(connector_prompt_embeds.dtype)
if not self.vae.config.timestep_conditioning:
timestep_decode = None
else:
noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
timestep_decode, decode_noise_scale_t = _prepare_decode_timestep_conditioning(
decode_timestep=decode_timestep,
decode_noise_scale=decode_noise_scale,
prompt_batch_size=prompt_batch_size,
effective_batch_size=latents.shape[0],
device=device,
dtype=latents.dtype,
)
latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise
dist_initialized = torch.distributed.is_initialized()
is_output_rank = not dist_initialized or torch.distributed.get_rank() == 0
vae_decode_needs_all_ranks = False
is_distributed_vae_enabled = getattr(self.vae, "is_distributed_enabled", None)
if self.distributed_video_decode and dist_initialized and callable(is_distributed_vae_enabled):
try:
# Distributed tiled decode is collective, so every rank must enter it.
vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled())
except Exception:
pass
should_decode_video = not self.distributed_video_decode or is_output_rank or vae_decode_needs_all_ranks
if should_decode_video:
video = self.vae.decode(latents.to(self.vae.dtype), timestep_decode, return_dict=False)[0]
else:
video = torch.empty(0, device=latents.device, dtype=latents.dtype)
if self.distributed_video_decode and not is_output_rank:
return self._make_output(
(
torch.empty(0, device=video.device, dtype=video.dtype),
torch.empty(0, device=audio_latents.device, dtype=audio_latents.dtype),
)
)
if video.numel() > 0:
video = self.video_processor.postprocess_video(video, output_type=output_type)
generated_mel = self.audio_vae.decode(audio_latents.to(self.audio_vae.dtype), return_dict=False)[0]
audio = self.vocoder(generated_mel)
return self._make_output((video, audio))
def _prepare_video_latents_stage(
self,
request_inputs: LTXRequestInputs,
prompt_context: LTXPromptContext,
*,
device: torch.device,
noise_scale: float,
image: Any | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
latents = self.prepare_latents(
prompt_context.batch_size * request_inputs.num_videos_per_prompt,
self.transformer.config.in_channels,
request_inputs.height,
request_inputs.width,
request_inputs.num_frames,
noise_scale,
torch.float32,
device,
request_inputs.generator,
request_inputs.latents,
)
return latents, None
def _resolve_video_latent_dimensions(self, request_inputs: LTXRequestInputs) -> tuple[int, int, int]:
latent_num_frames, latent_height, latent_width = latent_ops.resolve_video_latent_shape(
request_inputs.height,
request_inputs.width,
request_inputs.num_frames,
vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
)
latents = request_inputs.latents
if latents is not None:
if latents.ndim == 5:
_, _, latent_num_frames, latent_height, latent_width = latents.shape
elif latents.ndim != 3:
raise ValueError(
f"Provided `latents` tensor has shape {latents.shape}, expected a packed 3D or unpacked 5D tensor."
)
return latent_num_frames, latent_height, latent_width
def _prepare_audio_latents_stage(
self,
request_inputs: LTXRequestInputs,
prompt_context: LTXPromptContext,
*,
device: torch.device,
noise_scale: float,
) -> tuple[torch.Tensor, int, int, int]:
duration_s = request_inputs.num_frames / request_inputs.frame_rate
audio_latents_per_second = (
self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio)
)
audio_num_frames = round(duration_s * audio_latents_per_second)
audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, request_inputs.audio_latents)
num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
num_channels = self.audio_vae.config.latent_channels if self.audio_vae is not None else 8
audio_latents, original_num_frames, padded_num_frames = self.prepare_audio_latents(
prompt_context.batch_size * request_inputs.num_videos_per_prompt,
num_channels_latents=num_channels,
audio_latent_length=audio_num_frames,
num_mel_bins=num_mel_bins,
noise_scale=noise_scale,
dtype=torch.float32,
device=device,
generator=request_inputs.generator,
latents=request_inputs.audio_latents,
)
return audio_latents, original_num_frames, padded_num_frames, latent_mel_bins
def _resolve_audio_latent_length(
self,
requested_length: int,
audio_latents: torch.Tensor | None,
) -> int:
if audio_latents is None or audio_latents.ndim != 4:
return requested_length
provided_length = audio_latents.shape[2]
if not self.preserve_sp_padded_audio_duration:
return provided_length
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1
padded_length = latent_ops.get_sp_padded_audio_latent_length(requested_length, int(sp_size))
return requested_length if provided_length in {requested_length, padded_length} else provided_length
def _make_video_audio_scheduler(
self,
audio_scheduler: Any,
latent_num_frames: int,
latent_height: int,
latent_width: int,
) -> Any:
return VideoAudioScheduler(self.scheduler, audio_scheduler)
def _prepare_denoise_context_for_cfg(
self,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> LTXDenoiseContext:
return self.guidance_strategy.prepare_denoise_context(self, forward_ctx, denoise_ctx)
def _denoise_timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> dict[str, torch.Tensor]:
return self.guidance_strategy.timestep_kwargs(ts, forward_ctx, denoise_ctx)
def _build_transformer_kwargs(
self,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
*,
hidden_states: torch.Tensor,
audio_hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
audio_encoder_hidden_states: torch.Tensor,
encoder_attention_mask: torch.Tensor | None,
audio_encoder_attention_mask: torch.Tensor | None,
ts: torch.Tensor,
) -> dict[str, Any]:
return build_transformer_kwargs(
self,
forward_ctx,
denoise_ctx,
hidden_states=hidden_states,
audio_hidden_states=audio_hidden_states,
encoder_hidden_states=encoder_hidden_states,
audio_encoder_hidden_states=audio_encoder_hidden_states,
encoder_attention_mask=encoder_attention_mask,
audio_encoder_attention_mask=audio_encoder_attention_mask,
ts=ts,
)
def _predict_noise_for_step(
self,
index: int,
timestep: torch.Tensor,
state: latent_ops.LTXAVState,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> tuple[torch.Tensor, torch.Tensor]:
return self.guidance_strategy.predict_noise(
self,
index,
timestep,
state,
forward_ctx,
denoise_ctx,
)
def _denoise_step(
self,
index: int,
timestep: torch.Tensor,
state: latent_ops.LTXAVState,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> latent_ops.LTXAVState:
denoise_ctx.latents = state.video
denoise_ctx.audio_latents = state.audio
noise_pred_video, noise_pred_audio = self._predict_noise_for_step(
index,
timestep,
state,
forward_ctx,
denoise_ctx,
)
video, audio = step_denoised_latents(
self,
forward_ctx,
denoise_ctx,
noise_pred_video,
noise_pred_audio,
timestep,
)
return latent_ops.LTXAVState(video=video, audio=audio)
def _unpack_and_denormalize_stage(
self,
forward_ctx: LTXForwardContext,
latents: torch.Tensor,
audio_latents: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
latents = latent_ops.unpack_latents(
latents,
forward_ctx.latent_num_frames,
forward_ctx.latent_height,
forward_ctx.latent_width,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
latents = latent_ops.denormalize_latents(
latents,
self.vae.latents_mean,
self.vae.latents_std,
self.vae.config.scaling_factor,
)
audio_latents = latent_ops.unpad_audio_latents(audio_latents, forward_ctx.original_audio_num_frames)
audio_latents = latent_ops.denormalize_audio_latents(
audio_latents,
self.audio_vae.latents_mean,
self.audio_vae.latents_std,
)
audio_latents = latent_ops.unpack_audio_latents(
audio_latents,
forward_ctx.original_audio_num_frames,
num_mel_bins=forward_ctx.latent_mel_bins,
)
return latents, audio_latents
def run_phase(
self,
req: DiffusionRequestBatch,
request_inputs: LTXRequestInputs,
*,
noise_scale: float,
sigmas: list[float] | None,
timesteps: list[int] | None,
attention_kwargs: dict[str, Any] | None,
image: Any | None = None,
prompt_context: LTXPromptContext | None = None,
) -> LTXPhaseResult:
"""Prepare and execute one phase without decoding its output."""
return LTXPhaseExecutor.run(
self,
req,
request_inputs,
noise_scale=noise_scale,
sigmas=sigmas,
timesteps=timesteps,
attention_kwargs=attention_kwargs,
image=image,
prompt_context=prompt_context,
)
def decode_phase(self, phase: LTXPhaseResult) -> DiffusionOutput | list[DiffusionOutput]:
"""Decode one completed phase and restore per-request outputs."""
forward_ctx = phase.forward_context
request_inputs = forward_ctx.request_inputs
output = self._decode_output(
latents=phase.video,
audio_latents=phase.audio,
output_type=request_inputs.output_type,
connector_prompt_embeds=forward_ctx.prompt_context.connector_prompt_embeds,
generator=request_inputs.generator,
device=forward_ctx.device,
decode_timestep=request_inputs.decode_timestep,
decode_noise_scale=request_inputs.decode_noise_scale,
prompt_batch_size=forward_ctx.batch_size,
)
if not self.supports_request_batch:
return output
return split_diffusion_output_by_request(
output,
forward_ctx.req,
num_outputs_per_prompt=forward_ctx.num_videos_per_prompt,
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
return AutoWeightsLoader(self).load_weights(weights)
vllm_omni/diffusion/models/ltx2/ltx2_recipes.py
+20
View File
@@ -0,0 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""User-visible one-stage defaults for the LTX model family."""
from dataclasses import dataclass
@dataclass(frozen=True)
class LTXOneStageRecipe:
height: int = 512
width: int = 768
num_frames: int = 121
frame_rate: float = 24.0
num_inference_steps: int = 40
guidance_scale: float = 4.0
LTX2_ONE_STAGE_RECIPE = LTXOneStageRecipe()
LTX23_ONE_STAGE_RECIPE = LTXOneStageRecipe()
vllm_omni/diffusion/models/ltx2/ltx2_request.py
+257
View File
@@ -0,0 +1,257 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Request normalization shared by LTX pipeline variants."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
import torch
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@dataclass
class LTXRequestInputs:
"""Resolved request values shared by all LTX execution variants."""
prompt: str | list[str] | None
negative_prompt: str | list[str] | None
height: int
width: int
num_frames: int
frame_rate: float
num_inference_steps: int
guidance_scale: float
guidance_rescale: float
num_videos_per_prompt: int
generator: torch.Generator | list[torch.Generator] | None
latents: torch.Tensor | None
audio_latents: torch.Tensor | None
prompt_embeds: torch.Tensor | None
negative_prompt_embeds: torch.Tensor | None
prompt_attention_mask: torch.Tensor | None
negative_prompt_attention_mask: torch.Tensor | None
decode_timestep: float | list[float]
decode_noise_scale: float | list[float] | None
output_type: str
max_sequence_length: int
def _unwrap_request_tensor(value: Any) -> Any:
if isinstance(value, list):
return value[0] if value else None
return value
def _get_prompt_field(prompt: Any, *keys: str) -> Any:
if isinstance(prompt, str):
return None
additional = prompt.get("additional_information")
for key in keys:
value = prompt.get(key)
if value is None and isinstance(additional, dict):
value = additional.get(key)
if value is not None:
return _unwrap_request_tensor(value)
return None
def _get_audio_latents_from_sampling(sampling: Any) -> torch.Tensor | None:
if sampling.audio_latents is not None:
return sampling.audio_latents
return sampling.extra_args.get("audio_latents")
class LTXRequestMixin:
"""Normalize serving requests without coupling them to a model version."""
supports_request_batch = False
supports_guidance_rescale = False
def check_inputs(
self,
prompt: str | list[str] | None,
height: int,
width: int,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
) -> None:
if height % 32 != 0 or width % 32 != 0:
raise ValueError(f"`height` and `width` must be divisible by 32 but are {height} and {width}.")
if prompt is not None and prompt_embeds is not None:
raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.")
if prompt is None and prompt_embeds is None:
raise ValueError("Provide either `prompt` or `prompt_embeds`.")
if prompt is not None and not isinstance(prompt, (str, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
if prompt_embeds is not None and prompt_attention_mask is None:
raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
if prompt_embeds is not None and negative_prompt_embeds is not None:
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got {prompt_embeds.shape} and {negative_prompt_embeds.shape}."
)
if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
raise ValueError(
"`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when"
f" passed directly, but got {prompt_attention_mask.shape} and"
f" {negative_prompt_attention_mask.shape}."
)
def _resolve_request_inputs(
self,
req: DiffusionRequestBatch,
*,
prompt: str | list[str] | None,
negative_prompt: str | list[str] | None,
height: int | None,
width: int | None,
num_frames: int | None,
frame_rate: float | None,
num_inference_steps: int | None,
timesteps: list[int] | None,
guidance_scale: float,
num_videos_per_prompt: int | None,
generator: torch.Generator | list[torch.Generator] | None,
latents: torch.Tensor | None,
audio_latents: torch.Tensor | None,
prompt_embeds: torch.Tensor | None,
negative_prompt_embeds: torch.Tensor | None,
prompt_attention_mask: torch.Tensor | None,
negative_prompt_attention_mask: torch.Tensor | None,
decode_timestep: float | list[float],
decode_noise_scale: float | list[float] | None,
output_type: str,
max_sequence_length: int | None,
guidance_rescale: float = 0.0,
) -> LTXRequestInputs:
sampling_params_list = req.sampling_params_list
sampling = sampling_params_list[0]
prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt
if all(isinstance(item, str) or item.get("negative_prompt") is None for item in req.prompts):
negative_prompt = None
elif req.prompts:
negative_prompt = [
"" if isinstance(item, str) else (item.get("negative_prompt") or "") for item in req.prompts
]
height = sampling.height or height or self.one_stage_recipe.height
width = sampling.width or width or self.one_stage_recipe.width
num_frames = sampling.num_frames or num_frames or self.one_stage_recipe.num_frames
frame_rate = sampling.resolved_frame_rate or frame_rate or self.one_stage_recipe.frame_rate
num_inference_steps = (
sampling.num_inference_steps or num_inference_steps or self.one_stage_recipe.num_inference_steps
)
if timesteps is None:
num_inference_steps = max(int(num_inference_steps), 2)
elif len(timesteps) < 2:
raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.")
num_videos_per_prompt = (
sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else num_videos_per_prompt or 1
)
max_sequence_length = sampling.max_sequence_length or max_sequence_length or self.tokenizer_max_length
if sampling.guidance_scale_provided:
guidance_scale = sampling.guidance_scale
if self.supports_guidance_rescale and sampling.guidance_rescale is not None:
guidance_rescale = sampling.guidance_rescale
if self.supports_request_batch:
if generator is None:
generator = req.collate_request_generators(num_videos_per_prompt, generator)
latents = req.collate_request_tensors("latents", latents)
audio_latents = DiffusionRequestBatch.collate_tensors(
[_get_audio_latents_from_sampling(item) for item in sampling_params_list],
"audio_latents",
audio_latents,
)
prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
req.prompts,
{
"prompt_embeds": prompt_embeds,
"negative_prompt_embeds": negative_prompt_embeds,
"prompt_attention_mask": prompt_attention_mask,
"negative_prompt_attention_mask": negative_prompt_attention_mask,
},
field_aliases={
"prompt_attention_mask": ("prompt_attention_mask", "attention_mask"),
"negative_prompt_attention_mask": (
"negative_prompt_attention_mask",
"negative_attention_mask",
),
},
)
prompt_embeds = prompt_fields["prompt_embeds"]
negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
prompt_attention_mask = prompt_fields["prompt_attention_mask"]
negative_prompt_attention_mask = prompt_fields["negative_prompt_attention_mask"]
if prompt_embeds is not None:
prompt = None
if negative_prompt_embeds is not None:
negative_prompt = None
else:
if generator is None:
generator = sampling.generator
if generator is None and sampling.seed is not None:
generator = torch.Generator(device=self.device).manual_seed(sampling.seed)
latents = sampling.latents if sampling.latents is not None else latents
sampling_audio_latents = _get_audio_latents_from_sampling(sampling)
if sampling_audio_latents is not None:
audio_latents = sampling_audio_latents
request_prompt_embeds = [_get_prompt_field(item, "prompt_embeds") for item in req.prompts]
if any(value is not None for value in request_prompt_embeds):
prompt_embeds = torch.stack(request_prompt_embeds) # type: ignore[arg-type]
request_negative_embeds = [_get_prompt_field(item, "negative_prompt_embeds") for item in req.prompts]
if any(value is not None for value in request_negative_embeds):
negative_prompt_embeds = torch.stack(request_negative_embeds) # type: ignore[arg-type]
request_prompt_masks = [
_get_prompt_field(item, "prompt_attention_mask", "attention_mask") for item in req.prompts
]
if any(value is not None for value in request_prompt_masks):
prompt_attention_mask = torch.stack(request_prompt_masks) # type: ignore[arg-type]
request_negative_masks = [
_get_prompt_field(item, "negative_prompt_attention_mask", "negative_attention_mask")
for item in req.prompts
]
if any(value is not None for value in request_negative_masks):
negative_prompt_attention_mask = torch.stack(request_negative_masks) # type: ignore[arg-type]
if sampling.decode_timestep is not None:
decode_timestep = sampling.decode_timestep
if sampling.decode_noise_scale is not None:
decode_noise_scale = sampling.decode_noise_scale
if sampling.output_type is not None:
output_type = sampling.output_type
return LTXRequestInputs(
prompt=prompt,
negative_prompt=negative_prompt,
height=int(height),
width=int(width),
num_frames=int(num_frames),
frame_rate=float(frame_rate),
num_inference_steps=int(num_inference_steps),
guidance_scale=guidance_scale,
guidance_rescale=guidance_rescale,
num_videos_per_prompt=int(num_videos_per_prompt),
generator=generator,
latents=latents,
audio_latents=audio_latents,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
decode_timestep=decode_timestep,
decode_noise_scale=decode_noise_scale,
output_type=output_type,
max_sequence_length=int(max_sequence_length),
)
vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py
+123 -1191
View File
File diff suppressed because it is too large. Load diff
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py
-1832
View File
File diff suppressed because it is too large. Load diff
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py
-454
View File
@@ -1,454 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""LTX-2.3 image-to-video pipeline."""
from __future__ import annotations
from typing import Any
import PIL.Image
import torch
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from .pipeline_ltx2_3 import (
LTX23Pipeline,
_LTX23DenoiseContext,
_LTX23ForwardContext,
_LTX23PromptContext,
_LTX23RequestInputs,
get_ltx2_post_process_func,
)
from .pipeline_ltx2_image2video import LTX2ImageToVideoPipeline, _I2VVideoAudioScheduler
class LTX23ImageToVideoPipeline(LTX23Pipeline):
"""LTX-2.3 image-to-video pipeline.
This keeps the LTX-2.3 prompt connector, x0-space CFG, sigma prompt
modulation, and audio branch semantics from ``LTX23Pipeline`` while
reusing the existing LTX image-conditioning contract: the first video
latent frame is encoded from the input image and remains fixed during
denoising.
"""
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__(od_config=od_config, prefix=prefix)
self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear")
support_image_input = True
_normalize_latents = staticmethod(LTX2ImageToVideoPipeline._normalize_latents)
_create_noised_state = staticmethod(LTX2ImageToVideoPipeline._create_noised_state)
@staticmethod
def _resolve_single_prompt_image(raw_image: Any) -> Any:
if isinstance(raw_image, list):
if len(raw_image) != 1:
raise ValueError(
"LTX-2.3 I2V prompt dictionaries support exactly one image per prompt. "
"Pass one image per prompt for batched I2V requests."
)
return raw_image[0]
return raw_image
@staticmethod
def _resolve_additional_image(additional: dict[str, Any]) -> Any:
raw_image = additional.get("preprocessed_image")
if raw_image is None:
raw_image = additional.get("pixel_values")
if raw_image is None:
raw_image = additional.get("image")
return raw_image
def prepare_latents(
self,
image: torch.Tensor | None = None,
batch_size: int = 1,
num_channels_latents: int = 128,
height: int = 512,
width: int = 768,
num_frames: int = 121,
noise_scale: float = 0.0,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Prepare I2V latents and the first-frame conditioning mask.
If caller-provided latents are used without an image, the latents must
already represent the full video state including the conditioning first
frame. Packed 3D latents are assumed to be in transformer token layout.
"""
height = height // self.vae_spatial_compression_ratio
width = width // self.vae_spatial_compression_ratio
num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
shape = (batch_size, num_channels_latents, num_frames, height, width)
mask_shape = (batch_size, 1, num_frames, height, width)
if latents is not None:
if latents.ndim == 5:
batch_size, _, num_frames, height, width = latents.shape
mask_shape = (batch_size, 1, num_frames, height, width)
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
latents = self._normalize_latents(
latents,
self.vae.latents_mean,
self.vae.latents_std,
self.vae.config.scaling_factor,
)
latents = self._create_noised_state(latents, noise_scale * (1 - conditioning_mask), generator)
latents = self._pack_latents(
latents,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
else:
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
conditioning_mask = self._pack_latents(
conditioning_mask,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
).squeeze(-1)
if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
raise ValueError(
"Provided `latents` tensor has shape"
f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
)
return latents.to(device=device, dtype=dtype), conditioning_mask
if image is None:
raise ValueError("`image` must be provided when `latents` is None.")
image_batch_size = image.shape[0]
if image_batch_size == 0:
raise ValueError("`image` batch is empty.")
if batch_size % image_batch_size != 0:
raise ValueError(
f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) "
"for image-to-video outputs."
)
num_videos_per_prompt = batch_size // image_batch_size
if isinstance(generator, list):
if len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective"
f" batch size of {batch_size}. Make sure the batch size matches the length of the generators."
)
image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)]
init_latents = [
retrieve_latents(self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), image_generators[i], "argmax")
for i in range(image_batch_size)
]
else:
init_latents = [
retrieve_latents(self.vae.encode(img.unsqueeze(0).unsqueeze(2)), generator, "argmax") for img in image
]
init_latents = torch.cat(init_latents, dim=0).to(dtype)
if num_videos_per_prompt > 1:
init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0)
init_latents = self._normalize_latents(
init_latents,
self.vae.latents_mean,
self.vae.latents_std,
)
init_latents = init_latents.repeat(1, 1, num_frames, 1, 1)
conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype)
conditioning_mask[:, :, 0] = 1.0
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)
conditioning_mask = self._pack_latents(
conditioning_mask,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
).squeeze(-1)
latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size)
return latents, conditioning_mask
def check_inputs(
self,
image,
height,
width,
prompt,
latents=None,
prompt_embeds=None,
negative_prompt_embeds=None,
prompt_attention_mask=None,
negative_prompt_attention_mask=None,
):
if image is None and latents is None:
raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.")
super().check_inputs(
prompt=prompt,
height=height,
width=width,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
)
_step_video_latents_i2v = LTX2ImageToVideoPipeline._step_video_latents_i2v
def _resolve_request_image(
self,
req: DiffusionRequestBatch,
image: PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None,
request_inputs: _LTX23RequestInputs,
) -> PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None:
if image is not None or not req.prompts:
return image
raw_images = []
for prompt_item in req.prompts:
if isinstance(prompt_item, str):
raw_image = None
else:
multi_modal_data = prompt_item.get("multi_modal_data") or {}
raw_image = multi_modal_data.get("image")
if raw_image is None:
additional = prompt_item.get("additional_information") or {}
raw_image = self._resolve_additional_image(additional)
raw_image = self._resolve_single_prompt_image(raw_image)
if isinstance(raw_image, str):
raw_image = PIL.Image.open(raw_image).convert("RGB")
raw_images.append(raw_image)
if any(img is None for img in raw_images) and request_inputs.latents is None:
raise ValueError("Image is required for LTX-2.3 I2V generation.")
if len(raw_images) == 1:
return raw_images[0]
if raw_images:
return raw_images
return image
def _check_forward_inputs(
self,
request_inputs: _LTX23RequestInputs,
image: Any | None = None,
) -> None:
self.check_inputs(
image=image,
height=request_inputs.height,
width=request_inputs.width,
prompt=request_inputs.prompt,
latents=request_inputs.latents,
prompt_embeds=request_inputs.prompt_embeds,
negative_prompt_embeds=request_inputs.negative_prompt_embeds,
prompt_attention_mask=request_inputs.prompt_attention_mask,
negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask,
)
def _resolve_video_latent_dimensions(
self,
request_inputs: _LTX23RequestInputs,
) -> tuple[int, int, int]:
latent_num_frames = (request_inputs.num_frames - 1) // self.vae_temporal_compression_ratio + 1
latent_height = request_inputs.height // self.vae_spatial_compression_ratio
latent_width = request_inputs.width // self.vae_spatial_compression_ratio
if request_inputs.latents is not None:
if request_inputs.latents.ndim == 5:
_, _, latent_num_frames, latent_height, latent_width = request_inputs.latents.shape
elif request_inputs.latents.ndim != 3:
raise ValueError(
f"Provided `latents` tensor has shape {request_inputs.latents.shape}, but the expected shape is "
"either [batch_size, seq_len, num_features] or "
"[batch_size, latent_dim, latent_frames, latent_height, latent_width]."
)
return latent_num_frames, latent_height, latent_width
def _prepare_video_latents_stage(
self,
request_inputs: _LTX23RequestInputs,
prompt_context: _LTX23PromptContext,
*,
device: torch.device,
noise_scale: float,
image: Any | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
if request_inputs.latents is None:
if isinstance(image, torch.Tensor):
if image.ndim == 3:
image = image.unsqueeze(0)
elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor):
image = torch.stack(image, dim=0)
else:
image = self.video_processor.preprocess(
image,
height=request_inputs.height,
width=request_inputs.width,
)
image = image.to(device=device, dtype=prompt_context.positive_connector_prompt_embeds.dtype)
num_channels_latents = self.transformer.config.in_channels
return self.prepare_latents(
image,
prompt_context.batch_size * request_inputs.num_videos_per_prompt,
num_channels_latents,
request_inputs.height,
request_inputs.width,
request_inputs.num_frames,
noise_scale,
torch.float32,
device,
request_inputs.generator,
request_inputs.latents,
)
def _make_video_audio_scheduler(
self,
audio_scheduler: Any,
latent_num_frames: int,
latent_height: int,
latent_width: int,
) -> Any:
return _I2VVideoAudioScheduler(
self,
audio_scheduler,
latent_num_frames,
latent_height,
latent_width,
)
def _prepare_denoise_context_for_cfg(
self,
forward_ctx: _LTX23ForwardContext,
denoise_ctx: _LTX23DenoiseContext,
) -> _LTX23DenoiseContext:
denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
if denoise_ctx.conditioning_mask is None:
raise ValueError("I2V denoising requires a conditioning mask.")
if self.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready:
denoise_ctx.conditioning_mask_for_model = torch.cat(
[denoise_ctx.conditioning_mask, denoise_ctx.conditioning_mask]
)
else:
denoise_ctx.conditioning_mask_for_model = denoise_ctx.conditioning_mask
return denoise_ctx
def _denoise_timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: _LTX23ForwardContext,
denoise_ctx: _LTX23DenoiseContext,
) -> dict[str, torch.Tensor]:
conditioning_mask = (
denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model
)
if conditioning_mask is None:
raise ValueError("I2V denoising requires a conditioning mask.")
return {
"timestep": ts.unsqueeze(-1) * (1 - conditioning_mask),
"audio_timestep": ts,
"sigma": ts,
}
def _step_denoised_latents(
self,
forward_ctx: _LTX23ForwardContext,
denoise_ctx: _LTX23DenoiseContext,
noise_pred_video: torch.Tensor,
noise_pred_audio: torch.Tensor,
t: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
latents, audio_latents = self.scheduler_step_maybe_with_cfg(
(noise_pred_video, noise_pred_audio),
(t, t),
(denoise_ctx.latents, denoise_ctx.audio_latents),
do_true_cfg=self.do_classifier_free_guidance,
per_request_scheduler=forward_ctx.video_audio_scheduler,
)
return self._synchronize_cfg_parallel_step_output(
(latents, audio_latents),
do_true_cfg=self.do_classifier_free_guidance,
)
@torch.no_grad()
def forward(
self,
req: DiffusionRequestBatch,
image: PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None = None,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
height: int | None = None,
width: int | None = None,
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
sigmas: list[float] | None = None,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
audio_latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
output_type: str = "np",
return_dict: bool = True,
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
) -> list[DiffusionOutput]:
request_inputs = self._resolve_request_inputs(
req,
prompt=prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=num_inference_steps,
timesteps=timesteps,
guidance_scale=guidance_scale,
num_videos_per_prompt=num_videos_per_prompt,
generator=generator,
latents=latents,
audio_latents=audio_latents,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
decode_timestep=decode_timestep,
decode_noise_scale=decode_noise_scale,
output_type=output_type,
max_sequence_length=max_sequence_length,
)
image = self._resolve_request_image(req, image, request_inputs)
return self._forward_impl(
req,
request_inputs,
noise_scale=noise_scale,
sigmas=sigmas,
timesteps=timesteps,
attention_kwargs=attention_kwargs,
image=image,
)
__all__ = [
"LTX23ImageToVideoPipeline",
"get_ltx2_post_process_func",
]
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py
-909
View File
@@ -1,909 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import copy
import os
from collections.abc import Iterable
from typing import Any, ClassVar
import numpy as np
import PIL.Image
import torch
import torch.nn as nn
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from vllm.logger import init_logger
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.lora.request import LoRARequest
from .pipeline_ltx2 import (
LTX2Pipeline,
_get_prompt_field,
calculate_shift,
)
from .pipeline_ltx2 import (
get_ltx2_post_process_func as _get_ltx2_post_process_func,
)
from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
logger = init_logger(__name__)
def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
return _get_ltx2_post_process_func(od_config)
class _I2VVideoAudioScheduler:
"""Composite scheduler for I2V: uses _step_video_latents_i2v for video, standard step for audio."""
def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width):
self.video_scheduler = pipeline.scheduler
self.audio_scheduler = audio_scheduler
self._pipeline = pipeline
self._latent_num_frames = latent_num_frames
self._latent_height = latent_height
self._latent_width = latent_width
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
video_out = self._pipeline._step_video_latents_i2v(
noise_pred[0],
latents[0],
t[0],
self._latent_num_frames,
self._latent_height,
self._latent_width,
)
audio_out = self.audio_scheduler.step(noise_pred[1], t[1], latents[1], return_dict=False, generator=generator)[
0
]
return ((video_out, audio_out),)
class LTX2ImageToVideoPipeline(LTX2Pipeline):
supports_request_batch = False
support_image_input = True
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
):
super().__init__(od_config=od_config, prefix=prefix)
self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear")
@staticmethod
def _normalize_latents(
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
) -> torch.Tensor:
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents = (latents - latents_mean) * scaling_factor / latents_std
return latents
def prepare_latents(
self,
image: torch.Tensor | None = None,
batch_size: int = 1,
num_channels_latents: int = 128,
height: int = 512,
width: int = 768,
num_frames: int = 121,
noise_scale: float = 0.0,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
height = height // self.vae_spatial_compression_ratio
width = width // self.vae_spatial_compression_ratio
num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
shape = (batch_size, num_channels_latents, num_frames, height, width)
mask_shape = (batch_size, 1, num_frames, height, width)
if latents is not None:
if latents.ndim == 5:
# conditioning_mask needs to the same shape as latents in two stages generation.
batch_size, _, num_frames, height, width = latents.shape
mask_shape = (batch_size, 1, num_frames, height, width)
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
latents = self._normalize_latents(
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
)
latents = self._create_noised_state(latents, noise_scale * (1 - conditioning_mask), generator)
# latents are of shape [B, C, F, H, W], need to be packed
latents = self._pack_latents(
latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
)
else:
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
conditioning_mask = self._pack_latents(
conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
).squeeze(-1)
if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
raise ValueError(
"Provided `latents` tensor has shape"
f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
)
return latents.to(device=device, dtype=dtype), conditioning_mask
if image is None:
raise ValueError("`image` must be provided when `latents` is None.")
image_batch_size = image.shape[0]
if image_batch_size == 0:
raise ValueError("`image` batch is empty.")
if batch_size % image_batch_size != 0:
raise ValueError(
f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) "
"for image-to-video outputs."
)
num_videos_per_prompt = batch_size // image_batch_size
if isinstance(generator, list):
if len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective"
f" batch size of {batch_size}. Make sure the batch size matches the length of the generators."
)
image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)]
init_latents = [
retrieve_latents(self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), image_generators[i], "argmax")
for i in range(image_batch_size)
]
else:
init_latents = [
retrieve_latents(self.vae.encode(img.unsqueeze(0).unsqueeze(2)), generator, "argmax") for img in image
]
init_latents = torch.cat(init_latents, dim=0).to(dtype)
if num_videos_per_prompt > 1:
init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0)
init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std)
init_latents = init_latents.repeat(1, 1, num_frames, 1, 1)
conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype)
conditioning_mask[:, :, 0] = 1.0
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)
conditioning_mask = self._pack_latents(
conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
).squeeze(-1)
latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size)
return latents, conditioning_mask
def check_inputs(
self,
image,
height,
width,
prompt,
latents=None,
prompt_embeds=None,
negative_prompt_embeds=None,
prompt_attention_mask=None,
negative_prompt_attention_mask=None,
):
if height % 32 != 0 or width % 32 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
if image is None and latents is None:
raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.")
if prompt is not None and prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
if prompt is None and prompt_embeds is None:
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
if prompt_embeds is not None and prompt_attention_mask is None:
raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
if prompt_embeds is not None and negative_prompt_embeds is not None:
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
raise ValueError(
"`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when"
" passed directly, but"
f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
f" {negative_prompt_attention_mask.shape}."
)
def _step_video_latents_i2v(
self,
noise_pred_video: torch.Tensor,
latents: torch.Tensor,
t: torch.Tensor,
latent_num_frames: int,
latent_height: int,
latent_width: int,
) -> torch.Tensor:
noise_pred_video = self._unpack_latents(
noise_pred_video,
latent_num_frames,
latent_height,
latent_width,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
latents_unpacked = self._unpack_latents(
latents,
latent_num_frames,
latent_height,
latent_width,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
noise_pred_video = noise_pred_video[:, :, 1:]
noise_latents = latents_unpacked[:, :, 1:]
pred_latents = self.scheduler.step(noise_pred_video, t, noise_latents, return_dict=False)[0]
latents_unpacked = torch.cat([latents_unpacked[:, :, :1], pred_latents], dim=2)
latents = self._pack_latents(
latents_unpacked, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
)
return latents
@torch.no_grad()
def forward(
self,
req: DiffusionRequestBatch,
image: PIL.Image.Image | torch.Tensor | None = None,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
height: int | None = None,
width: int | None = None,
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
sigmas: list[float] | None = None,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
audio_latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
output_type: str = "np",
return_dict: bool = True,
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
) -> DiffusionOutput:
# Extract prompt/negative_prompt from request.
# Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys.
prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
negative_prompt = None
elif req.prompts:
negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
height = req.sampling_params.height or height or 512
width = req.sampling_params.width or width or 768
num_frames = req.sampling_params.num_frames or num_frames or 121
frame_rate = req.sampling_params.resolved_frame_rate or frame_rate or 24.0
num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps or 40
if timesteps is None:
num_inference_steps = max(int(num_inference_steps), 2)
elif len(timesteps) < 2:
raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.")
num_videos_per_prompt = (
req.sampling_params.num_outputs_per_prompt
if req.sampling_params.num_outputs_per_prompt > 0
else num_videos_per_prompt or 1
)
max_sequence_length = (
req.sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length
)
if req.sampling_params.guidance_scale_provided:
guidance_scale = req.sampling_params.guidance_scale
if req.sampling_params.guidance_rescale is not None:
guidance_rescale = req.sampling_params.guidance_rescale
if generator is None:
generator = req.sampling_params.generator
if generator is None and req.sampling_params.seed is not None:
generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
latents = req.sampling_params.latents if req.sampling_params.latents is not None else latents
audio_latents = (
req.sampling_params.audio_latents
if req.sampling_params.audio_latents is not None
else req.sampling_params.extra_args.get("audio_latents", audio_latents)
)
# Override with pre-computed embeddings if provided in request.
req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts]
if any(p is not None for p in req_prompt_embeds):
prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore[arg-type]
req_negative_prompt_embeds = [_get_prompt_field(p, "negative_prompt_embeds") for p in req.prompts]
if any(p is not None for p in req_negative_prompt_embeds):
negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) # type: ignore[arg-type]
req_prompt_attention_masks = [
_get_prompt_field(p, "prompt_attention_mask") or _get_prompt_field(p, "attention_mask") for p in req.prompts
]
if any(m is not None for m in req_prompt_attention_masks):
prompt_attention_mask = torch.stack(req_prompt_attention_masks) # type: ignore[arg-type]
req_negative_attention_masks = [
_get_prompt_field(p, "negative_prompt_attention_mask") or _get_prompt_field(p, "negative_attention_mask")
for p in req.prompts
]
if any(m is not None for m in req_negative_attention_masks):
negative_prompt_attention_mask = torch.stack(req_negative_attention_masks) # type: ignore[arg-type]
if req.sampling_params.decode_timestep is not None:
decode_timestep = req.sampling_params.decode_timestep
if req.sampling_params.decode_noise_scale is not None:
decode_noise_scale = req.sampling_params.decode_noise_scale
if req.sampling_params.output_type is not None:
output_type = req.sampling_params.output_type
if image is None and req.prompts:
raw_images = []
for prompt_item in req.prompts:
if isinstance(prompt_item, str):
raw_image = None
else:
multi_modal_data = prompt_item.get("multi_modal_data") or {}
raw_image = multi_modal_data.get("image")
if raw_image is None:
additional = prompt_item.get("additional_information") or {}
raw_image = (
additional.get("preprocessed_image")
or additional.get("pixel_values")
or additional.get("image")
)
if isinstance(raw_image, list):
raw_image = raw_image[0] if raw_image else None
if isinstance(raw_image, str):
raw_image = PIL.Image.open(raw_image).convert("RGB")
raw_images.append(raw_image)
if any(img is None for img in raw_images):
if latents is None:
raise ValueError("Image is required for LTX2 I2V generation.")
if len(raw_images) == 1:
image = raw_images[0]
elif raw_images:
image = raw_images
self.check_inputs(
image=image,
height=height,
width=width,
prompt=prompt,
latents=latents,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
)
self._guidance_scale = guidance_scale
self._guidance_rescale = guidance_rescale
self._attention_kwargs = attention_kwargs
self._interrupt = False
self._current_timestep = None
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self.device
(
prompt_embeds,
prompt_attention_mask,
negative_prompt_embeds,
negative_prompt_attention_mask,
) = self.encode_prompt(
prompt=prompt,
negative_prompt=negative_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
num_videos_per_prompt=num_videos_per_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
max_sequence_length=max_sequence_length,
device=device,
)
# Compute positive prompt connectors
tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left")
connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side
)
# Compute negative prompt connectors when CFG is enabled
negative_connector_prompt_embeds = None
negative_connector_audio_prompt_embeds = None
negative_connector_attention_mask = None
if self.do_classifier_free_guidance:
(
negative_connector_prompt_embeds,
negative_connector_audio_prompt_embeds,
negative_connector_attention_mask,
) = self.connectors(
negative_prompt_embeds,
negative_prompt_attention_mask,
padding_side=tokenizer_padding_side,
)
latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
latent_height = height // self.vae_spatial_compression_ratio
latent_width = width // self.vae_spatial_compression_ratio
if latents is not None:
if latents.ndim == 5:
logger.info(
"Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred." # noqa
)
_, _, latent_num_frames, latent_height, latent_width = latents.shape # [B, C, F, H, W]
elif latents.ndim == 3:
logger.warning(
f"You have supplied packed `latents` of shape {latents.shape}, so the latent dims cannot be"
f" inferred. Make sure the supplied `height`, `width`, and `num_frames` are correct."
)
else:
raise ValueError(
f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]." # noqa
)
video_sequence_length = latent_num_frames * latent_height * latent_width
if latents is None:
if isinstance(image, torch.Tensor):
if image.ndim == 3:
image = image.unsqueeze(0)
elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor):
image = torch.stack(image, dim=0)
else:
image = self.video_processor.preprocess(image, height=height, width=width)
image = image.to(device=device, dtype=prompt_embeds.dtype)
num_channels_latents = self.transformer.config.in_channels
latents, conditioning_mask = self.prepare_latents(
image,
batch_size * num_videos_per_prompt,
num_channels_latents,
height,
width,
num_frames,
noise_scale,
torch.float32,
device,
generator,
latents,
)
duration_s = num_frames / frame_rate
audio_latents_per_second = (
self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio)
)
audio_num_frames = round(duration_s * audio_latents_per_second)
if audio_latents is not None:
if audio_latents.ndim == 4:
logger.info(
"Got audio_latents of shape [batch_size, num_channels, audio_length, mel_bins], `audio_num_frames` will be inferred." # noqa
)
_, _, audio_num_frames, _ = audio_latents.shape # [B, C, L, M]
elif audio_latents.ndim == 3:
logger.warning(
f"You have supplied packed `audio_latents` of shape {audio_latents.shape}, so the latent dims"
f" cannot be inferred. Make sure the supplied `num_frames` and `frame_rate` are correct."
)
else:
raise ValueError(
f"Provided `audio_latents` tensor has shape {audio_latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, num_channels, audio_length, mel_bins]." # noqa
)
num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
num_channels_latents_audio = (
self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
)
audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents(
batch_size * num_videos_per_prompt,
num_channels_latents=num_channels_latents_audio,
audio_latent_length=audio_num_frames,
num_mel_bins=num_mel_bins,
noise_scale=noise_scale,
dtype=torch.float32,
device=device,
generator=generator,
latents=audio_latents,
)
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
mu = calculate_shift(
video_sequence_length,
self.scheduler.config.get("base_image_seq_len", 1024),
self.scheduler.config.get("max_image_seq_len", 4096),
self.scheduler.config.get("base_shift", 0.95),
self.scheduler.config.get("max_shift", 2.05),
)
audio_scheduler = copy.deepcopy(self.scheduler)
_ = retrieve_timesteps(
audio_scheduler,
num_inference_steps,
device,
timesteps,
sigmas=sigmas,
mu=mu,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
num_inference_steps,
device,
timesteps,
sigmas=sigmas,
mu=mu,
)
self._num_timesteps = len(timesteps)
video_coords = self.transformer.rope.prepare_video_coords(
latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate
)
audio_coords = self.transformer.audio_rope.prepare_audio_coords(
audio_latents.shape[0], padded_audio_num_frames, audio_latents.device
)
i2v_scheduler = _I2VVideoAudioScheduler(
pipeline=self,
audio_scheduler=audio_scheduler,
latent_num_frames=latent_num_frames,
latent_height=latent_height,
latent_width=latent_width,
)
# No coord duplication needed: mixin handles CFG via separate forward calls.
with self.progress_bar(total=len(timesteps)) as pbar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t
latent_model_input = latents.to(prompt_embeds.dtype)
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
do_true_cfg = self.do_classifier_free_guidance
positive_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": connector_prompt_embeds,
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
"timestep": video_timestep,
"audio_timestep": timestep,
"encoder_attention_mask": connector_attention_mask,
"audio_encoder_attention_mask": connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": padded_audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
negative_kwargs = (
{
**positive_kwargs,
"encoder_hidden_states": negative_connector_prompt_embeds,
"audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
"encoder_attention_mask": negative_connector_attention_mask,
"audio_encoder_attention_mask": negative_connector_attention_mask,
}
if do_true_cfg
else None
)
noise_pred_video, noise_pred_audio = self.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
latents, audio_latents = self.scheduler_step_maybe_with_cfg(
(noise_pred_video, noise_pred_audio),
(t, t),
(latents, audio_latents),
do_true_cfg=do_true_cfg,
per_request_scheduler=i2v_scheduler,
)
latents, audio_latents = self._synchronize_cfg_parallel_step_output(
(latents, audio_latents),
do_true_cfg=do_true_cfg,
)
pbar.update()
latents = self._unpack_latents(
latents,
latent_num_frames,
latent_height,
latent_width,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
latents = self._denormalize_latents(
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
)
audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames)
audio_latents = self._denormalize_audio_latents(
audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std
)
audio_latents = self._unpack_audio_latents(
audio_latents,
original_audio_num_frames,
num_mel_bins=latent_mel_bins,
)
if output_type == "latent":
video = latents
audio = audio_latents
else:
latents = latents.to(prompt_embeds.dtype)
if not self.vae.config.timestep_conditioning:
timestep = None
else:
noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
if not isinstance(decode_timestep, list):
decode_timestep = [decode_timestep] * batch_size
if decode_noise_scale is None:
decode_noise_scale = decode_timestep
elif not isinstance(decode_noise_scale, list):
decode_noise_scale = [decode_noise_scale] * batch_size
timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
:, None, None, None, None
]
latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
latents = latents.to(self.vae.dtype)
video = self.vae.decode(latents, timestep, return_dict=False)[0]
video = self.video_processor.postprocess_video(video, output_type=output_type)
audio_latents = audio_latents.to(self.audio_vae.dtype)
generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0]
audio = self.vocoder(generated_mel_spectrograms)
return DiffusionOutput(output=(video, audio))
class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsComponentDiscovery):
"""LTXImageToVideoTwoStagesPipeline is for two stages image to video generation"""
supports_request_batch = False
support_image_input = True
dummy_run_num_frames = 2
_dit_modules: ClassVar[list[str]] = ["pipe.transformer"]
_encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"]
_vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"]
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
):
super().__init__()
self.device = get_local_device()
self.dtype = getattr(od_config, "dtype", torch.bfloat16)
self.model_path = od_config.model
self.distilled = False
# User provided model path may contain '/' in the end and basename function
# will not return the expected directory name, so we need to remove it by normpath
if "distilled" in os.path.basename(os.path.normpath(self.model_path)):
self.distilled = True
else:
raise NotImplementedError(f"{self.model_path} is not supported for {self.__class__.__name__}.")
self.pipe = LTX2ImageToVideoPipeline(od_config=od_config, prefix=prefix)
self.upsample_pipe = LTX2LatentUpsamplePipeline(
vae=self.pipe.vae,
od_config=od_config,
)
self.lora_manager = DiffusionLoRAManager(
pipeline=self.pipe,
device=self.device,
dtype=self.dtype,
max_cached_adapters=od_config.max_cpu_loras,
)
self.weights_sources = [
DiffusersPipelineLoader.ComponentSource(
model_or_path=od_config.model,
subfolder="transformer",
revision=None,
prefix="pipe.transformer.",
fall_back_to_pt=True,
),
]
@torch.no_grad()
def forward(
self,
req: DiffusionRequestBatch,
image: PIL.Image.Image | torch.Tensor | None = None,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
height: int | None = None,
width: int | None = None,
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
sigmas: list[float] | None = None,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
audio_latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
output_type: str = "np",
return_dict: bool = True,
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
) -> DiffusionOutput:
stage1_output = self.pipe(
req=req,
image=image,
prompt=prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=num_inference_steps,
sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None,
timesteps=timesteps,
guidance_scale=guidance_scale,
guidance_rescale=guidance_rescale,
noise_scale=noise_scale,
num_videos_per_prompt=num_videos_per_prompt,
generator=generator,
latents=latents,
audio_latents=audio_latents,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
decode_timestep=decode_timestep,
decode_noise_scale=decode_noise_scale,
output_type="latent",
return_dict=return_dict,
attention_kwargs=attention_kwargs,
max_sequence_length=max_sequence_length,
)
video_latent, audio_latent = stage1_output.output
upscaled_video_latent = self.upsample_pipe(
latents=video_latent,
output_type="latent",
return_dict=False,
)[0]
if not self.distilled:
# Load Stage 2 distilled LoRA
lora_path = f"{self.model_path}/ltx-2-19b-distilled-lora-384.safetensors"
lora_request = LoRARequest(
lora_name="stage_2_distilled",
lora_int_id=1,
lora_path=lora_path,
)
self.lora_manager.set_active_adapter(lora_request, lora_scale=1.0)
# Change scheduler to use Stage 2 distilled sigmas as is
new_scheduler = FlowMatchEulerDiscreteScheduler.from_config(
self.pipe.scheduler.config,
use_dynamic_shifting=False,
shift_terminal=None,
)
self.pipe.scheduler = new_scheduler
# `req` is a DiffusionRequestBatch whose `sampling_params` is a
# read-only property; clone the underlying request(s) and override
# their sampling params instead of assigning to the batch.
stage_2_req = copy.copy(req)
stage_2_req.requests = [copy.copy(r) for r in req.requests]
for stage_2_request in stage_2_req.requests:
stage_2_request.sampling_params = stage_2_request.sampling_params.clone()
stage_2_request.sampling_params.num_inference_steps = 3
stage2_output = self.pipe(
req=stage_2_req,
latents=upscaled_video_latent,
audio_latents=audio_latent,
prompt=prompt,
negative_prompt=negative_prompt,
noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
guidance_scale=1.0,
generator=generator,
output_type="np",
return_dict=False,
)
video, audio = stage2_output.output
return DiffusionOutput(output=(video, audio))
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
return loader.load_weights(weights)
class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2ImageToVideoPipeline):
"""LTX-2 I2V pipeline for FastGen DMD2-distilled models."""
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__(od_config=od_config, prefix=prefix)
self.__init_dmd2__()
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py
+207
View File
@@ -0,0 +1,207 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Two-stage entry points for the LTX model family."""
from __future__ import annotations
import os
from dataclasses import replace
from typing import Any, ClassVar
import torch
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from .ltx2_components import (
LTX2_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_guidance import LTX_LEGACY_VELOCITY_GUIDANCE
from .ltx2_pipeline_runtime import LTXPipelineRuntime
from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE
from .ltx2_request import LTXRequestInputs
from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
class LTX2TwoStagesPipeline(LTXPipelineRuntime):
"""Legacy distilled-only LTX2 two-stage compatibility entry."""
component_profile = LTX2_COMPONENT_PROFILE
guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE
one_stage_recipe = LTX2_ONE_STAGE_RECIPE
supports_request_batch = False
support_image_input = False
_dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules)
_encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules)
_vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules)
_resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules)
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
model_path = od_config.model
self.distilled = "distilled" in os.path.basename(os.path.normpath(model_path))
if not self.distilled:
raise NotImplementedError(f"{model_path} is not supported for {self.__class__.__name__}.")
super().__init__(od_config=od_config, prefix=prefix)
self.upsample_pipe = LTX2LatentUpsamplePipeline(vae=self.vae, od_config=od_config)
def _run_two_stage(
self,
req: DiffusionRequestBatch,
request_inputs: LTXRequestInputs,
*,
noise_scale: float,
timesteps: list[int] | None,
attention_kwargs: dict[str, Any] | None,
image: Any | None = None,
) -> DiffusionOutput | list[DiffusionOutput]:
stage1 = self.run_phase(
req,
request_inputs,
noise_scale=noise_scale,
sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None,
timesteps=timesteps,
attention_kwargs=attention_kwargs,
image=image,
)
upscaled_video_latent = self.upsample_pipe(
latents=stage1.video,
output_type="latent",
return_dict=False,
)[0]
stage2_inputs = replace(
request_inputs,
num_inference_steps=3,
guidance_scale=1.0,
latents=upscaled_video_latent,
audio_latents=stage1.audio,
decode_timestep=0.0,
decode_noise_scale=None,
output_type="np",
)
stage2 = self.run_phase(
req,
stage2_inputs,
noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
timesteps=None,
attention_kwargs=attention_kwargs,
prompt_context=stage1.forward_context.prompt_context,
)
return self.decode_phase(stage2)
@torch.no_grad()
def forward(
self,
req: DiffusionRequestBatch,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
height: int | None = None,
width: int | None = None,
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
audio_latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
output_type: str = "np",
return_dict: bool = True,
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
*,
image: Any | None = None,
) -> DiffusionOutput | list[DiffusionOutput]:
del return_dict
request_inputs = self._resolve_request_inputs(
req,
prompt=prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=num_inference_steps,
timesteps=timesteps,
guidance_scale=(
getattr(getattr(self, "one_stage_recipe", None), "guidance_scale", 4.0)
if guidance_scale is None
else guidance_scale
),
guidance_rescale=guidance_rescale,
num_videos_per_prompt=num_videos_per_prompt,
generator=generator,
latents=latents,
audio_latents=audio_latents,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
decode_timestep=decode_timestep,
decode_noise_scale=decode_noise_scale,
output_type=output_type,
max_sequence_length=max_sequence_length,
)
if self.support_image_input:
image = self._resolve_request_image(req, image, request_inputs)
return self._run_two_stage(
req,
request_inputs,
noise_scale=noise_scale,
timesteps=timesteps,
attention_kwargs=attention_kwargs,
image=image,
)
class LTX2ImageToVideoTwoStagesPipeline(LTXI2VConditioningMixin, LTX2TwoStagesPipeline):
"""LTX2 two-stage image-to-video entry."""
def forward(
self,
req: DiffusionRequestBatch,
image: Any | None = None,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
height: int | None = None,
width: int | None = None,
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
sigmas: list[float] | None = None,
*args: Any,
**kwargs: Any,
) -> DiffusionOutput | list[DiffusionOutput]:
"""Preserve the legacy I2V sigma slot, which was ignored by this pipeline."""
del sigmas
return super().forward(
req,
image,
prompt,
negative_prompt,
height,
width,
num_frames,
frame_rate,
num_inference_steps,
*args,
**kwargs,
)
vllm_omni/diffusion/registry.py
+6 -6
View File
@@ -73,17 +73,17 @@ _DIFFUSION_MODELS = {
),
"LTX2ImageToVideoPipeline": (
"ltx2",
"pipeline_ltx2_image2video",
"pipeline_ltx2",
"LTX2ImageToVideoPipeline",
),
"LTX2TwoStagesPipeline": (
"ltx2",
"pipeline_ltx2",
"pipeline_ltx2_two_stage",
"LTX2TwoStagesPipeline",
),
"LTX2ImageToVideoTwoStagesPipeline": (
"ltx2",
"pipeline_ltx2_image2video",
"pipeline_ltx2_two_stage",
"LTX2ImageToVideoTwoStagesPipeline",
),
"LTX2T2VDMD2Pipeline": (
@@ -93,17 +93,17 @@ _DIFFUSION_MODELS = {
),
"LTX2I2VDMD2Pipeline": (
"ltx2",
"pipeline_ltx2_image2video",
"pipeline_ltx2",
"LTX2I2VDMD2Pipeline",
),
"LTX23Pipeline": (
"ltx2",
"pipeline_ltx2_3",
"pipeline_ltx2",
"LTX23Pipeline",
),
"LTX23ImageToVideoPipeline": (
"ltx2",
"pipeline_ltx2_3_image2video",
"pipeline_ltx2",
"LTX23ImageToVideoPipeline",
),
"StableAudioPipeline": (