diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 1ab76588f..c1dac6ca7 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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` | ✅︎ | ✅︎ | | | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index d3ed3bd67..c34c1338c 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -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** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/docs/user_guide/examples/offline_inference/image_to_video.md b/docs/user_guide/examples/offline_inference/image_to_video.md index dc1d3d6b2..b3d76ac46 100644 --- a/docs/user_guide/examples/offline_inference/image_to_video.md +++ b/docs/user_guide/examples/offline_inference/image_to_video.md @@ -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. diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md index 5d89eac34..61c498245 100644 --- a/docs/user_guide/examples/offline_inference/text_to_video.md +++ b/docs/user_guide/examples/offline_inference/text_to_video.md @@ -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. diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md index d4d520d5e..0f1f33a25 100644 --- a/docs/user_guide/examples/online_serving/text_to_video.md +++ b/docs/user_guide/examples/online_serving/text_to_video.md @@ -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" diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index c161022a5..4604a64bd 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -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. diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index d64367bc2..164290d43 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -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 \ diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index edc75bb26..876b8429b 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -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. diff --git a/recipes/LTX/LTX-2.3.md b/recipes/LTX/LTX-2.3.md index 7a4f1507c..471ee7ae0 100644 --- a/recipes/LTX/LTX-2.3.md +++ b/recipes/LTX/LTX-2.3.md @@ -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: -- Diffusers-format checkpoint: +- Diffusers-format checkpoint: - 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: diff --git a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py index 21679b5b7..5f8de588e 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py +++ b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py @@ -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 # --------------------------------------------------------------------------- diff --git a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py index 0b6ca3a9e..a395af959 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py +++ b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py @@ -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 diff --git a/tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py b/tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py deleted file mode 100644 index bbfe63dfa..000000000 --- a/tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py +++ /dev/null @@ -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) diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py similarity index 71% rename from tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py rename to tests/diffusion/models/ltx2/test_ltx2_conditioning.py index 9b883ff34..43737f8c1 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py @@ -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 diff --git a/tests/diffusion/models/ltx2/test_ltx2_hsdp.py b/tests/diffusion/models/ltx2/test_ltx2_hsdp.py deleted file mode 100644 index 4dd07e1bf..000000000 --- a/tests/diffusion/models/ltx2/test_ltx2_hsdp.py +++ /dev/null @@ -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"] diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py new file mode 100644 index 000000000..74cad09d1 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py @@ -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, + ) diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py similarity index 88% rename from tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py rename to tests/diffusion/models/ltx2/test_ltx2_parallel.py index 3e592f14b..0c6b12fbe 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py +++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py @@ -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, diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py similarity index 57% rename from tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py rename to tests/diffusion/models/ltx2/test_ltx2_pipeline.py index 25aae4510..2f9155fe0 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -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, - ) diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index e71253297..01a1fa3cf 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -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"] diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py similarity index 85% rename from tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py rename to tests/diffusion/models/ltx2/test_ltx2_vae.py index 9b53bb4f8..4ae7a0721 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -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: diff --git a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py b/tests/e2e/accuracy/test_ltx2_3_video_similarity.py index 7fb2297a2..a169817a2 100644 --- a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py +++ b/tests/e2e/accuracy/test_ltx2_3_video_similarity.py @@ -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) diff --git a/vllm_omni/diffusion/models/dmd2/mixin.py b/vllm_omni/diffusion/models/dmd2/mixin.py index 873120fc1..9c98e96b4 100644 --- a/vllm_omni/diffusion/models/dmd2/mixin.py +++ b/vllm_omni/diffusion/models/dmd2/mixin.py @@ -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] diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index 76a26b5b0..ca34e0390 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -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", diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py new file mode 100644 index 000000000..cfe990a4e --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -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) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py new file mode 100644 index 000000000..48d319f6d --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py @@ -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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py new file mode 100644 index 000000000..af8df4b41 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -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) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py new file mode 100644 index 000000000..395f8e7af --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py @@ -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() diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py new file mode 100644 index 000000000..4da3edf92 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py @@ -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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py new file mode 100644 index 000000000..08bf20330 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py @@ -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) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py new file mode 100644 index 000000000..6491b4a5f --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py @@ -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() diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py new file mode 100644 index 000000000..02cab34d2 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -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), + ) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index bce0813e9..c32a6918a 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -1,749 +1,66 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""One-stage entry points for the LTX model family.""" + from __future__ import annotations -import copy -import inspect -import json -import os -from collections.abc import Iterable -from contextlib import nullcontext -from typing import TYPE_CHECKING, Any, ClassVar +from typing import Any, ClassVar -import numpy as np import torch -if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.base_config import QuantizationConfig -from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler -from diffusers.pipelines.ltx2 import LTX2TextConnectors -from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES -from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder -from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg, retrieve_timesteps -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -from torch import nn -from transformers import AutoTokenizer, Gemma3ForConditionalGeneration -from vllm.logger import init_logger -from vllm.model_executor.models.utils import AutoWeightsLoader - from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin -from vllm_omni.diffusion.distributed.parallel_state import ( - get_classifier_free_guidance_world_size, -) -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.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -from vllm_omni.lora.request import LoRARequest -from .ltx2_transformer import LTX2VideoTransformer3DModel -from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline - -logger = init_logger(__name__) +from .ltx2_components import ( + LTX2_COMPONENT_PROFILE, + LTX23_COMPONENT_PROFILE, + LTXComponentProfile, +) +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, + LTX_OFFICIAL_X0_GUIDANCE, +) +from .ltx2_pipeline_runtime import LTXPipelineRuntime +from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE, LTXOneStageRecipe +from .ltx2_request import LTXRequestInputs -def load_transformer_config(model_path: str, subfolder: str = "transformer", local_files_only: bool = True) -> dict: - """Load transformer config from model directory or 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 f: - return json.load(f) - else: - try: - from huggingface_hub import hf_hub_download +class LTXOneStagePipeline(LTXPipelineRuntime): + """Single execution path configured by model-version and task entries.""" - config_path = hf_hub_download( - repo_id=model_path, - filename=f"{subfolder}/config.json", - ) - with open(config_path) as f: - return json.load(f) - except Exception: - pass - return {} + component_profile: ClassVar[LTXComponentProfile] + one_stage_recipe: ClassVar[LTXOneStageRecipe] + supports_request_batch = True + supports_guidance_rescale = False -def create_transformer_from_config( - config: dict, - quant_config: QuantizationConfig | None = None, -) -> LTX2VideoTransformer3DModel: - """Create LTX2VideoTransformer3DModel from config dict.""" - if not config and quant_config is None: - return LTX2VideoTransformer3DModel() - - signature = inspect.signature(LTX2VideoTransformer3DModel.__init__) - allowed_keys = set(signature.parameters.keys()) - kwargs = {k: v for k, v in config.items() if k in allowed_keys} - if quant_config is not None: - kwargs["quant_config"] = quant_config - - return LTX2VideoTransformer3DModel(**kwargs) - - -def get_ltx2_post_process_func( - od_config: OmniDiffusionConfig, -): - def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor): - if isinstance(output, tuple) and len(output) == 2: - video, audio = output - if isinstance(audio, torch.Tensor): - audio = audio.detach().cpu() - return {"video": video, "audio": audio} - return output - - return post_process_func - - -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, key: str) -> Any: - if isinstance(prompt, str): - return None - value = prompt.get(key) - if value is None: - additional = prompt.get("additional_information") - if isinstance(additional, dict): - value = additional.get(key) - return _unwrap_request_tensor(value) - - -def calculate_shift( - image_seq_len, - base_seq_len: int = 256, - max_seq_len: int = 4096, - base_shift: float = 0.5, - max_shift: float = 1.15, -): - m = (max_shift - base_shift) / (max_seq_len - base_seq_len) - b = base_shift - m * base_seq_len - mu = image_seq_len * m + b - return mu - - -class _VideoAudioScheduler: - """Composite scheduler dispatching to video and audio schedulers.""" - - 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 LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsComponentDiscovery): - supports_request_batch = False - - _dit_modules: ClassVar[list[str]] = ["transformer"] - _encoder_modules: ClassVar[list[str]] = ["text_encoder"] - _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] - - # Audio is diffused jointly with video; warmup must size audio tokens. - dummy_run_num_frames = 2 - - def __init__( + def _forward_impl( self, + req: DiffusionRequestBatch, + request_inputs: LTXRequestInputs, *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - - self.device = get_local_device() - dtype = getattr(od_config, "dtype", torch.bfloat16) - - model = od_config.model - local_files_only = os.path.exists(model) - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=od_config.model, - subfolder="transformer", - revision=None, - prefix="transformer.", - fall_back_to_pt=True, - ), - ] - - # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder - # race; prefetch the whole component set before any from_pretrained. - ltx2_subfolders = [ - "tokenizer", - "text_encoder", - "connectors", - "vae", - "audio_vae", - "vocoder", - "scheduler", - ] - prefetch_subfolders(model, ltx2_subfolders, local_files_only=local_files_only) - - self.tokenizer = AutoTokenizer.from_pretrained( - model, - subfolder="tokenizer", - local_files_only=local_files_only, + noise_scale: float, + sigmas: list[float] | None, + timesteps: list[int] | None, + attention_kwargs: dict[str, Any] | None, + image: Any | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + phase = self.run_phase( + req, + request_inputs, + noise_scale=noise_scale, + sigmas=sigmas, + timesteps=timesteps, + attention_kwargs=attention_kwargs, + image=image, ) - # prefer mmap loading as default device is cuda, and the output of text encoder - # could be deterministic. - with torch.device("cpu"): - self.text_encoder = from_pretrained_with_prefetch( - Gemma3ForConditionalGeneration.from_pretrained, - model, - subfolder="text_encoder", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - self.connectors = from_pretrained_with_prefetch( - LTX2TextConnectors.from_pretrained, - model, - subfolder="connectors", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - - self.vae = from_pretrained_with_prefetch( - AutoencoderKLLTX2Video.from_pretrained, - model, - subfolder="vae", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - self.audio_vae = from_pretrained_with_prefetch( - AutoencoderKLLTX2Audio.from_pretrained, - model, - subfolder="audio_vae", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - self.vocoder = from_pretrained_with_prefetch( - LTX2Vocoder.from_pretrained, - model, - subfolder="vocoder", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - - transformer_config = load_transformer_config(model, "transformer", local_files_only) - quant_config = getattr(self.od_config, "quantization_config", None) - self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) - - self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( - model, subfolder="scheduler", local_files_only=local_files_only - ) - - self.vae_spatial_compression_ratio = ( - self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 - ) - self.vae_temporal_compression_ratio = ( - self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 - ) - self.audio_vae_mel_compression_ratio = ( - self.audio_vae.mel_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 - ) - self.audio_vae_temporal_compression_ratio = ( - self.audio_vae.temporal_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 - ) - self.transformer_spatial_patch_size = ( - self.transformer.config.patch_size if getattr(self, "transformer", None) is not None else 1 - ) - self.transformer_temporal_patch_size = ( - self.transformer.config.patch_size_t if getattr(self, "transformer", None) is not None else 1 - ) - - self.audio_sampling_rate = ( - self.audio_vae.config.sample_rate if getattr(self, "audio_vae", None) is not None else 16000 - ) - self.audio_hop_length = ( - self.audio_vae.config.mel_hop_length if getattr(self, "audio_vae", None) is not None else 160 - ) - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - tokenizer_max_length = 1024 - if getattr(self, "tokenizer", None) is not None: - tokenizer_max_length = self.tokenizer.model_max_length - if tokenizer_max_length is None or tokenizer_max_length > 100000: - encoder_config = getattr(self.text_encoder, "config", None) - config_max_len = getattr(encoder_config, "max_position_embeddings", None) - if config_max_len is None: - config_max_len = getattr(encoder_config, "max_seq_len", None) - tokenizer_max_length = config_max_len or 1024 - self.tokenizer_max_length = int(tokenizer_max_length) - - self._guidance_scale = None - self._guidance_rescale = None - self._attention_kwargs = None - self._interrupt = False - self._num_timesteps = None - self._current_timestep = None - - 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, - ): - device = device or self.device - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - if getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "left" - if self.tokenizer.pad_token is None: - self.tokenizer.pad_token = self.tokenizer.eos_token - - prompt = [p.strip() for p in prompt] - text_inputs = self.tokenizer( - 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 - prompt_attention_mask = text_inputs.attention_mask - text_input_ids = text_input_ids.to(device) - prompt_attention_mask = prompt_attention_mask.to(device) - - text_encoder_outputs = self.text_encoder( - input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True - ) - text_encoder_hidden_states = text_encoder_outputs.hidden_states - text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1) - # `diffusers>=0.38` moved per_layer_masked_mean_norm inside the connector, - # so pack to 3D here and let the connector handle normalization. - prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(dtype=dtype) - - _, seq_len, _ = prompt_embeds.shape - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - - prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) - prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) - - return prompt_embeds, prompt_attention_mask - - def encode_prompt( - self, - prompt: str | list[str], - 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, - ): - device = device or self.device - - prompt = [prompt] if isinstance(prompt, str) else prompt - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - 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, - ) - - 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 to `prompt`, but got {type(negative_prompt)} !=" - f" {type(prompt)}." - ) - elif batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`." - ) - - 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, - ) - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - def check_inputs( - self, - prompt, - height, - width, - 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 prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two." - ) - elif prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." - ) - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if prompt_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 got: `prompt_attention_mask` " - f"{prompt_attention_mask.shape} != `negative_prompt_attention_mask` " - f"{negative_prompt_attention_mask.shape}." - ) - - @staticmethod - def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: - batch_size, num_channels, 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, - ) - latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) - return latents - - @staticmethod - 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) - latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) - return latents - - @staticmethod - def _normalize_latents( - latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 - ) -> torch.Tensor: - # Normalize latents across the channel dimension [B, C, F, H, W] - 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 - - @staticmethod - def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: 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 - - @staticmethod - 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) - latents = latents * latents_std / scaling_factor + latents_mean - return latents - - @staticmethod - def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: 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 - - @staticmethod - def _create_noised_state( - latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None - ): - noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) - noised_latents = noise_scale * noise + (1 - noise_scale) * latents - return noised_latents - - @staticmethod - 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, num_channels, 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 - ) - latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) - else: - latents = latents.transpose(1, 2).flatten(2, 3) - return latents - - @staticmethod - 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) - latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) - else: - latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) - return latents - - @staticmethod - def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: - return latents[:, :num_frames] - - 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 | None = None, - latents: torch.Tensor | None = None, - ) -> torch.Tensor: - if latents is not None: - if latents.ndim == 5: - latents = self._normalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - # 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 - ) - if latents.ndim != 3: - raise ValueError( - f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]." # noqa - ) - latents = self._create_noised_state(latents, noise_scale, generator) - return latents.to(device=device, dtype=dtype) - - 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) - - 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) - latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size) - - return latents - - def prepare_audio_latents( - self, - batch_size: int = 1, - num_channels_latents: int = 8, - audio_latent_length: int = 1, # 1 is just a dummy value - 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 - padded_latent_length = original_latent_length - - latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio - - sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) - if sp_size > 1: - padded_latent_length += (sp_size - (original_latent_length % sp_size)) % sp_size - - if latents is not None: - if latents.ndim == 4: - # latents are of shape [B, C, L, M], need to be packed - latents = self._pack_audio_latents(latents) - if latents.ndim != 3: - raise ValueError( - f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is " - "[batch_size, num_seq, num_features] or [batch_size, num_channels, audio_length, mel_bins]." - ) - latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) - latents = self._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) - latents = self._pack_audio_latents(latents) - return latents, original_latent_length, padded_latent_length - - @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 > 1.0 - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def current_timestep(self): - return self._current_timestep - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @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): - """Per-element CFG combine with guidance_rescale support.""" - (video_pos, audio_pos) = positive_noise_pred - (video_neg, audio_neg) = negative_noise_pred - video_combined = super().combine_cfg_noise(video_pos, video_neg, true_cfg_scale, cfg_normalize) - audio_combined = super().combine_cfg_noise(audio_pos, audio_neg, true_cfg_scale, cfg_normalize) - if self._guidance_rescale and self._guidance_rescale > 0: - video_combined = rescale_noise_cfg(video_combined, video_pos, guidance_rescale=self._guidance_rescale) - audio_combined = rescale_noise_cfg(audio_combined, audio_pos, guidance_rescale=self._guidance_rescale) - return (video_combined, audio_combined) - - 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 - - # Without this sync, CUDA async execution causes non-deterministic - # numerical drift across denoising steps in CFG parallel mode, - # producing different video outputs across runs. - 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 + return self.decode_phase(phase) @torch.no_grad() def forward( @@ -775,428 +92,79 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompon 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 - - self.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, - ) - - 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( + *, + image: Any | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + del return_dict + request_inputs = self._resolve_request_inputs( + req, prompt=prompt, negative_prompt=negative_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + num_inference_steps=num_inference_steps, + timesteps=timesteps, + guidance_scale=self.one_stage_recipe.guidance_scale 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, - 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 + image = self._resolve_request_image(req, image, request_inputs) + forward_kwargs = { + "noise_scale": noise_scale, + "sigmas": sigmas, + "timesteps": timesteps, + "attention_kwargs": attention_kwargs, + } + if self.support_image_input: + forward_kwargs["image"] = image + return self._forward_impl( + req, + request_inputs, + **forward_kwargs, ) - # 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 +class LTX2Pipeline(LTXOneStagePipeline): + """LTX2 one-stage text-to-video entry.""" - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - 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) - video_audio_scheduler = _VideoAudioScheduler(self.scheduler, audio_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 - ) - # No coord duplication needed: mixin handles CFG via separate forward calls, - # not batch=2. Each forward gets batch=1 coords directly. - - 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]) - 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": 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=video_audio_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)) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + supports_guidance_rescale = True + component_profile = LTX2_COMPONENT_PROFILE + guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE + one_stage_recipe = LTX2_ONE_STAGE_RECIPE + _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) -class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery): - """LTX2TwoStagesPipeline is for two stages image to video generation""" +class LTX23Pipeline(LTXOneStagePipeline): + """LTX2.3 one-stage text-to-video entry.""" - dummy_run_num_frames = 2 - supports_request_batch = False - - _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 = LTX2Pipeline(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, - ), - ] + connector_batches_cfg = True + preserve_sp_padded_audio_duration = True + reports_stage_durations = True + component_profile = LTX23_COMPONENT_PROFILE + guidance_strategy = LTX_OFFICIAL_X0_GUIDANCE + one_stage_recipe = LTX23_ONE_STAGE_RECIPE + _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) + @torch.no_grad() def forward( self, req: DiffusionRequestBatch, @@ -1207,9 +175,9 @@ class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery): 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, @@ -1225,9 +193,12 @@ class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery): 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: Any | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + """Preserve the pre-refactor LTX2.3 positional argument order.""" + return super().forward( + req, prompt=prompt, negative_prompt=negative_prompt, height=height, @@ -1235,10 +206,10 @@ class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery): num_frames=num_frames, frame_rate=frame_rate, num_inference_steps=num_inference_steps, - sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None, + sigmas=sigmas, timesteps=timesteps, guidance_scale=guidance_scale, - guidance_rescale=guidance_rescale, + guidance_rescale=0.0, noise_scale=noise_scale, num_videos_per_prompt=num_videos_per_prompt, generator=generator, @@ -1250,71 +221,32 @@ class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery): negative_prompt_attention_mask=negative_prompt_attention_mask, decode_timestep=decode_timestep, decode_noise_scale=decode_noise_scale, - output_type="latent", + output_type=output_type, return_dict=return_dict, attention_kwargs=attention_kwargs, max_sequence_length=max_sequence_length, + image=image, ) - 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) +class LTX2ImageToVideoPipeline(LTXI2VConditioningMixin, LTX2Pipeline): + """LTX2 one-stage image-to-video entry.""" - # 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 - # We only want to change num_inference_steps here, so no need - # to deep copy the whole request. `req` is a DiffusionRequestBatch - # whose `sampling_params` is a read-only property, so clone the - # underlying request(s) and override their sampling params instead. - 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 LTX23ImageToVideoPipeline(LTXI2VConditioningMixin, LTX23Pipeline): + """LTX2.3 one-stage image-to-video entry.""" class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline): - """LTX-2 T2V pipeline for FastGen DMD2-distilled models.""" + """LTX2 T2V entry for FastGen DMD2-distilled models.""" + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__(od_config=od_config, prefix=prefix) + self.__init_dmd2__() + + +class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2ImageToVideoPipeline): + """LTX2 I2V entry for FastGen DMD2-distilled models.""" def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py deleted file mode 100644 index 45be23296..000000000 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ /dev/null @@ -1,1832 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -Fully independent LTX-2.3 pipeline for vLLM-Omni. - -This pipeline does NOT inherit from LTX2Pipeline because: -- LTX-2.3 connectors run per_token_rms_norm + per-modality video/audio - projection internally (per_modality_projections=True), - versus LTX-2's per_layer_masked_mean_norm + shared projection path -- LTX-2.3 uses a BWE vocoder outputting 48kHz audio (not 16kHz) -- LTX-2.3 transformer requires the sigma parameter for prompt modulation -""" - -from __future__ import annotations - -import copy -import json -import os -from collections.abc import Iterable -from contextlib import nullcontext -from dataclasses import dataclass -from typing import Any, ClassVar - -import numpy as np -import torch -from diffusers import AutoencoderKLLTX2Audio, FlowMatchEulerDiscreteScheduler -from diffusers.pipelines.ltx2 import LTX2TextConnectors -from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder -from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -from huggingface_hub import hf_hub_download -from torch import nn -from transformers import AutoTokenizer, Gemma3ForConditionalGeneration -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.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video -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, -) -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin -from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery -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 .pipeline_ltx2 import ( - _VideoAudioScheduler, - calculate_shift, - create_transformer_from_config, - load_transformer_config, -) - -logger = init_logger(__name__) - - -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") - - -def _is_output_rank() -> bool: - return not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 - - -def _vae_decode_needs_all_ranks(vae: Any) -> bool: - if not torch.distributed.is_initialized(): - return False - is_distributed_enabled = getattr(vae, "is_distributed_enabled", None) - if not callable(is_distributed_enabled): - return False - try: - return bool(is_distributed_enabled()) - except Exception: - return False - - -def _should_decode_video_on_rank(vae: Any) -> bool: - return _is_output_rank() or _vae_decode_needs_all_ranks(vae) - - -# Try to import LTX2VocoderWithBWE (diffusers >= 0.38.0) -try: - from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE -except ImportError: - LTX2VocoderWithBWE = None - - -@dataclass -class _LTX23RequestInputs: - 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 - 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 - - -@dataclass -class _LTX23PromptContext: - 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 - - -@dataclass -class _LTX23ForwardContext: - req: DiffusionRequestBatch - request_inputs: _LTX23RequestInputs - prompt_context: _LTX23PromptContext - 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 _LTX23DenoiseContext: - 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 - - -def _repeat_prompt_tensor_for_outputs(tensor: torch.Tensor, num_videos_per_prompt: int) -> torch.Tensor: - if num_videos_per_prompt == 1: - return tensor - return tensor.repeat_interleave(num_videos_per_prompt, dim=0) - - -def _detect_vocoder_output_sample_rate(model: str) -> int | None: - """Detect the vocoder output sample rate from vocoder/config.json. - - This runs at factory time (engine process) so the rate is captured in - the post-process closure and doesn't need cross-process communication. - - Returns: - Output sample rate (e.g. 48000 for LTX-2.3 BWE vocoder) or None. - """ - 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 f: - cfg = json.load(f) - return cfg.get("output_sampling_rate") - except Exception: - return None - - -def get_ltx2_post_process_func(od_config: OmniDiffusionConfig): - """Factory for the LTX-2.3 post-process function. - - Detects the vocoder output sample rate at factory time and captures it - in the closure so that the audio_sample_rate flows through - DiffusionEngine -> OmniRequestOutput -> serving_video. - """ - output_sr = _detect_vocoder_output_sample_rate(od_config.model) - - def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor): - if isinstance(output, tuple) and len(output) == 2: - video, audio = output - if isinstance(audio, torch.Tensor): - audio = audio.detach().cpu() - result: dict[str, Any] = {"video": video, "audio": audio} - if output_sr is not None: - result["audio_sample_rate"] = output_sr - return result - return output - - return post_process_func - - -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", - ) - if decode_noise_scale is None: - decode_noise_scale_values = decode_timestep_values - else: - decode_noise_scale_values = _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 LTX23Pipeline( - nn.Module, - CFGParallelMixin, - ProgressBarMixin, - SupportsComponentDiscovery, - DiffusionPipelineProfilerMixin, -): - """Fully independent LTX-2.3 pipeline. - - Key differences from LTX2Pipeline: - - Text encoding: uses ALL 49 hidden states from Gemma-3-12B, flattened - - Connectors: uses padding_side API (not additive_mask) - - Vocoder: uses LTX2VocoderWithBWE (48kHz output) - - Transformer: passes sigma for prompt_adaln - """ - - supports_request_batch = True - # Audio is diffused jointly with video; warmup must size audio tokens. - dummy_run_num_frames = 2 - _dit_modules: ClassVar[list[str]] = ["transformer"] - _encoder_modules: ClassVar[list[str]] = ["text_encoder", "connectors"] - _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] - _resident_modules: ClassVar[list[str]] = ["vocoder"] - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - self.device = get_local_device() - dtype = getattr(od_config, "dtype", torch.bfloat16) - model = od_config.model - local_files_only = os.path.exists(model) - - # Weight sources for transformer (loaded via AutoWeightsLoader) - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=od_config.model, - subfolder="transformer", - revision=None, - prefix="transformer.", - fall_back_to_pt=True, - ), - ] - - # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder - # race; prefetch the whole component set before any from_pretrained. - ltx2_subfolders = [ - "tokenizer", - "text_encoder", - "connectors", - "vae", - "audio_vae", - "vocoder", - "scheduler", - ] - prefetch_subfolders(model, ltx2_subfolders, local_files_only=local_files_only) - - # --- Tokenizer (lightweight, stays wherever) --- - self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - - # --- Text encoder --- - with torch.device("cpu"): - self.text_encoder = from_pretrained_with_prefetch( - Gemma3ForConditionalGeneration.from_pretrained, - model, - subfolder="text_encoder", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ) - - # --- Connectors (LTX-2.3 connectors include caption projection) --- - self.connectors = from_pretrained_with_prefetch( - LTX2TextConnectors.from_pretrained, - model, - subfolder="connectors", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ) - - # --- VAE, Audio VAE --- - self.vae = from_pretrained_with_prefetch( - DistributedAutoencoderKLLTX2Video.from_pretrained, - model, - subfolder="vae", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ) - self.audio_vae = from_pretrained_with_prefetch( - AutoencoderKLLTX2Audio.from_pretrained, - model, - subfolder="audio_vae", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ) - - # --- Vocoder: prefer BWE vocoder (48kHz) for LTX-2.3 --- - vocoder_cls = LTX2VocoderWithBWE or LTX2Vocoder - try: - self.vocoder = vocoder_cls.from_pretrained( - model, subfolder="vocoder", torch_dtype=dtype, local_files_only=local_files_only - ) - except (TypeError, OSError, ValueError): - self.vocoder = LTX2Vocoder.from_pretrained( - model, subfolder="vocoder", torch_dtype=dtype, local_files_only=local_files_only - ) - - # --- Transformer: created empty, weights loaded via AutoWeightsLoader --- - transformer_config = load_transformer_config(model, "transformer", local_files_only) - quant_config = getattr(self.od_config, "quantization_config", None) - self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) - self._place_aux_components() - - # --- Scheduler --- - self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( - model, subfolder="scheduler", local_files_only=local_files_only - ) - - # --- Derived compression ratios --- - self.vae_spatial_compression_ratio = self.vae.spatial_compression_ratio if self.vae is not None else 32 - self.vae_temporal_compression_ratio = self.vae.temporal_compression_ratio if self.vae is not None else 8 - self.audio_vae_mel_compression_ratio = self.audio_vae.mel_compression_ratio if self.audio_vae is not None else 4 - self.audio_vae_temporal_compression_ratio = ( - self.audio_vae.temporal_compression_ratio if self.audio_vae is not None else 4 - ) - self.transformer_spatial_patch_size = self.transformer.config.patch_size if self.transformer is not None else 1 - self.transformer_temporal_patch_size = ( - self.transformer.config.patch_size_t if self.transformer is not None else 1 - ) - self.audio_sampling_rate = self.audio_vae.config.sample_rate if self.audio_vae is not None else 16000 - self.audio_hop_length = self.audio_vae.config.mel_hop_length if self.audio_vae is not None else 160 - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - - # Tokenizer max length - tokenizer_max_length = 1024 - if self.tokenizer is not None: - tokenizer_max_length = self.tokenizer.model_max_length - if tokenizer_max_length is None or tokenizer_max_length > 100000: - encoder_config = getattr(self.text_encoder, "config", None) - config_max_len = getattr(encoder_config, "max_position_embeddings", None) - if config_max_len is None: - config_max_len = getattr(encoder_config, "max_seq_len", None) - tokenizer_max_length = config_max_len or 1024 - self.tokenizer_max_length = int(tokenizer_max_length) - - # Pipeline state - self._guidance_scale = None - self._attention_kwargs = None - self._interrupt = False - self._num_timesteps = None - self._current_timestep = None - - self.setup_diffusion_pipeline_profiler( - enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler - ) - - def _place_aux_components(self) -> None: - parallel_config = getattr(self.od_config, "parallel_config", None) - use_managed_placement = bool( - getattr(self.od_config, "enable_cpu_offload", False) - or getattr(self.od_config, "enable_layerwise_offload", False) - or getattr(parallel_config, "use_hsdp", False) - ) - if use_managed_placement: - return - - modules = ModuleDiscovery.discover(self) - for module in (*modules.encoders, *modules.vaes, *modules.resident_modules): - module.to(self.device) - - # ------------------------------------------------------------------ - # Text Encoding (LTX-2.3 specific) - # ------------------------------------------------------------------ - - def _get_gemma_prompt_embeds( - self, - prompt: str | list[str], - num_videos_per_prompt: int = 1, - max_sequence_length: int = 1024, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - """Encode prompts using Gemma-3-12B, returning ALL 49 hidden states flattened. - - Stacks all 49 hidden states and flattens to [B, seq, hidden * 49]. The - connectors unflatten, apply per_token_rms_norm, and project internally - (same shape contract as LTX-2 since the `diffusers==0.38` connector - migration; the two differ only in the connector's internal norm path). - """ - device = device or self.device - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - if self.tokenizer is not None: - self.tokenizer.padding_side = "left" - if self.tokenizer.pad_token is None: - self.tokenizer.pad_token = self.tokenizer.eos_token - - prompt = [p.strip() for p in prompt] - text_inputs = self.tokenizer( - 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) - - text_encoder_outputs = self.text_encoder( - input_ids=text_input_ids, - attention_mask=prompt_attention_mask, - output_hidden_states=True, - ) - - hidden_states = text_encoder_outputs.hidden_states - - # LTX-2.3: Stack ALL 49 hidden states and flatten - # [49 x (B, seq, 3840)] -> [B, seq, 3840, 49] -> [B, seq, 188160] - prompt_embeds = torch.stack(hidden_states, dim=-1).flatten(2, 3).to(dtype=dtype) - - prompt_attention_mask = prompt_attention_mask.view(batch_size, -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) - - return prompt_embeds, prompt_attention_mask - - def encode_prompt( - self, - prompt: str | list[str], - 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, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - device = device or self.device - - prompt = [prompt] if isinstance(prompt, str) else prompt - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = 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, - 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`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`." - ) - - 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, - 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 - - # ------------------------------------------------------------------ - # Latent utilities (shared with LTX2Pipeline) - # ------------------------------------------------------------------ - - @staticmethod - def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: - batch_size, num_channels, 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, - ) - latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) - return latents - - @staticmethod - 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) - latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) - return latents - - @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 - - @staticmethod - def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: 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 - - @staticmethod - 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) - latents = latents * latents_std / scaling_factor + latents_mean - return latents - - @staticmethod - def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: 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 - - @staticmethod - 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, num_channels, 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 - ) - latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) - else: - latents = latents.transpose(1, 2).flatten(2, 3) - return latents - - @staticmethod - 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) - latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) - else: - latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) - return latents - - @staticmethod - def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: - return latents[:, :num_frames] - - @staticmethod - 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_audio_latent_length(self, audio_latent_length: int, audio_latents: torch.Tensor | None) -> int: - if audio_latents is None or audio_latents.ndim != 4: - return audio_latent_length - - provided_latent_length = audio_latents.shape[2] - sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1 - padded_latent_length = self._get_sp_padded_audio_latent_length(audio_latent_length, int(sp_size)) - - # Keep requested duration semantics when callers pass 4D latents that - # are already padded for SP; other 4D lengths retain shape inference. - if provided_latent_length in {audio_latent_length, padded_latent_length}: - return audio_latent_length - return provided_latent_length - - 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 DiffusionOutput( - output=(latents, audio_latents), - stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, - ) - - 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 - - if _should_decode_video_on_rank(self.vae): - latents = latents.to(self.vae.dtype) - video = self.vae.decode(latents, timestep_decode, return_dict=False)[0] - else: - video = torch.empty(0, device=latents.device, dtype=latents.dtype) - - if not _is_output_rank(): - return DiffusionOutput( - output=( - torch.empty(0, device=video.device, dtype=video.dtype), - torch.empty(0, device=audio_latents.device, dtype=audio_latents.dtype), - ), - stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, - ) - - if video.numel() > 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), - stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, - ) - - # ------------------------------------------------------------------ - # Latent preparation - # ------------------------------------------------------------------ - - 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 | None = None, - latents: torch.Tensor | None = None, - ) -> torch.Tensor: - if latents is not None: - if latents.ndim == 5: - latents = self._normalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - latents = self._pack_latents( - latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size - ) - if latents.ndim != 3: - raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") - noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) - latents = noise_scale * noise + (1 - noise_scale) * latents - return latents.to(device=device, dtype=dtype) - - 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) - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size) - return 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]: - original_latent_length = audio_latent_length - latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio - - sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1 - padded_latent_length = self._get_sp_padded_audio_latent_length(original_latent_length, int(sp_size)) - - if latents is not None: - if latents.ndim == 4: - latents = self._pack_audio_latents(latents) - if latents.ndim != 3: - raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") - latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) - noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) - latents = noise_scale * noise + (1 - noise_scale) * latents - - 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) - latents = self._pack_audio_latents(latents) - return latents, original_latent_length, padded_latent_length - - # ------------------------------------------------------------------ - # Properties - # ------------------------------------------------------------------ - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale is not None and self._guidance_scale > 1.0 - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def current_timestep(self): - return self._current_timestep - - @property - def interrupt(self): - return self._interrupt - - # ------------------------------------------------------------------ - # Input validation - # ------------------------------------------------------------------ - - def check_inputs( - self, - prompt, - height, - width, - 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` 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`.") - elif prompt is None and prompt_embeds is None: - raise ValueError("Provide either `prompt` or `prompt_embeds`.") - elif 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` {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 got: `prompt_attention_mask` " - f"{prompt_attention_mask.shape} != `negative_prompt_attention_mask` " - f"{negative_prompt_attention_mask.shape}." - ) - - # ------------------------------------------------------------------ - # Cache context - # ------------------------------------------------------------------ - - 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() - - # ------------------------------------------------------------------ - # CFG helpers - # ------------------------------------------------------------------ - - @staticmethod - def _combine_x0_space_cfg( - sample: torch.Tensor, - positive_noise_pred: torch.Tensor, - negative_noise_pred: torch.Tensor, - sigma: torch.Tensor, - guidance_scale: float, - ) -> torch.Tensor: - x0_cond = sample - positive_noise_pred * sigma - x0_uncond = sample - negative_noise_pred * sigma - x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond) - return (sample - x0_guided) / sigma - - 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, - *, - video_latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | None = None, - video_sigma: torch.Tensor | None = None, - audio_sigma: torch.Tensor | None = None, - ): - if video_latents is None or audio_latents is None or video_sigma is None or audio_sigma is None: - raise ValueError("LTX23Pipeline applies CFG in x0-space and requires video/audio latents and sigmas.") - - video_pos, audio_pos = positive_noise_pred - video_neg, audio_neg = negative_noise_pred - video_combined = self._combine_x0_space_cfg( - video_latents, - video_pos, - video_neg, - video_sigma, - true_cfg_scale, - ) - audio_combined = self._combine_x0_space_cfg( - audio_latents, - audio_pos, - audio_neg, - audio_sigma, - true_cfg_scale, - ) - if cfg_normalize: - video_combined = self.cfg_normalize_function(video_pos, video_combined) - audio_combined = self.cfg_normalize_function(audio_pos, audio_combined) - return video_combined, audio_combined - - def predict_noise_with_parallel_cfg( - self, - true_cfg_scale: float, - positive_kwargs: dict[str, Any], - negative_kwargs: dict[str, Any], - cfg_normalize: bool = True, - output_slice: int | None = None, - *, - video_latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | None = None, - video_sigma: torch.Tensor | None = None, - audio_sigma: torch.Tensor | None = None, - ) -> 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() - if cfg_world_size != 2: - raise ValueError(f"LTX23Pipeline parallel CFG requires cfg_parallel_size 2, but got {cfg_world_size}.") - - cfg_group = get_cfg_group() - 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(self.predict_noise(**branch_kwargs)) - - gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True) - gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True) - positive_noise_pred = (gathered_video[0], gathered_audio[0]) - negative_noise_pred = (gathered_video[1], gathered_audio[1]) - - return self.combine_cfg_noise( - positive_noise_pred, - negative_noise_pred, - true_cfg_scale, - cfg_normalize, - video_latents=video_latents, - audio_latents=audio_latents, - video_sigma=video_sigma, - audio_sigma=audio_sigma, - ) - - 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 - - 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 _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, - ) -> _LTX23RequestInputs: - sampling_params_list = req.sampling_params_list - common_sampling_params = sampling_params_list[0] - 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 = common_sampling_params.height or height or 512 - width = common_sampling_params.width or width or 768 - num_frames = common_sampling_params.num_frames or num_frames or 121 - frame_rate = common_sampling_params.resolved_frame_rate or frame_rate or 24.0 - num_inference_steps = common_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 = ( - common_sampling_params.num_outputs_per_prompt - if common_sampling_params.num_outputs_per_prompt > 0 - else num_videos_per_prompt or 1 - ) - max_sequence_length = ( - common_sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length - ) - - if common_sampling_params.guidance_scale_provided: - guidance_scale = common_sampling_params.guidance_scale - - 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(sampling) for sampling 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 - - if common_sampling_params.decode_timestep is not None: - decode_timestep = common_sampling_params.decode_timestep - if common_sampling_params.decode_noise_scale is not None: - decode_noise_scale = common_sampling_params.decode_noise_scale - if common_sampling_params.output_type is not None: - output_type = common_sampling_params.output_type - - return _LTX23RequestInputs( - 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, - 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), - ) - - 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, - ) -> _LTX23PromptContext: - 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] - - 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, - ) - ) - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) - - connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( - prompt_embeds, - prompt_attention_mask, - padding_side=getattr(self.tokenizer, "padding_side", "left"), - ) - - positive_connector_prompt_embeds = connector_prompt_embeds - positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds - positive_connector_attention_mask = connector_attention_mask - negative_connector_prompt_embeds = None - negative_connector_audio_prompt_embeds = None - negative_connector_attention_mask = None - if self.do_classifier_free_guidance: - split_batch = batch_size * num_videos_per_prompt - negative_connector_prompt_embeds = connector_prompt_embeds[:split_batch] - positive_connector_prompt_embeds = connector_prompt_embeds[split_batch:] - negative_connector_audio_prompt_embeds = connector_audio_prompt_embeds[:split_batch] - positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds[split_batch:] - negative_connector_attention_mask = connector_attention_mask[:split_batch] - positive_connector_attention_mask = connector_attention_mask[split_batch:] - - return _LTX23PromptContext( - batch_size=batch_size, - connector_prompt_embeds=connector_prompt_embeds, - connector_audio_prompt_embeds=connector_audio_prompt_embeds, - connector_attention_mask=connector_attention_mask, - positive_connector_prompt_embeds=positive_connector_prompt_embeds, - positive_connector_audio_prompt_embeds=positive_connector_audio_prompt_embeds, - positive_connector_attention_mask=positive_connector_attention_mask, - negative_connector_prompt_embeds=negative_connector_prompt_embeds, - negative_connector_audio_prompt_embeds=negative_connector_audio_prompt_embeds, - negative_connector_attention_mask=negative_connector_attention_mask, - ) - - def _setup_forward_runtime( - self, - request_inputs: _LTX23RequestInputs, - attention_kwargs: dict[str, Any] | None, - ) -> bool: - self._guidance_scale = request_inputs.guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - self._current_timestep = None - cfg_world_size = get_classifier_free_guidance_world_size() - if self.do_classifier_free_guidance and cfg_world_size not in (1, 2): - raise ValueError( - f"LTX23Pipeline supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}." - ) - return self.do_classifier_free_guidance and cfg_world_size > 1 - - def _check_forward_inputs( - self, - request_inputs: _LTX23RequestInputs, - 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_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 and request_inputs.latents.ndim == 5: - _, _, latent_num_frames, latent_height, latent_width = request_inputs.latents.shape - 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]: - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - 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, - ) - return latents, None - - def _prepare_audio_latents_stage( - self, - request_inputs: _LTX23RequestInputs, - prompt_context: _LTX23PromptContext, - *, - 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_latents_audio = self.audio_vae.config.latent_channels if self.audio_vae is not None else 8 - audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents( - prompt_context.batch_size * request_inputs.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=request_inputs.generator, - latents=request_inputs.audio_latents, - ) - return audio_latents, original_audio_num_frames, padded_audio_num_frames, latent_mel_bins - - 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_scheduler_stage( - self, - request_inputs: _LTX23RequestInputs, - *, - 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 - ) - # Use max_image_seq_len (not actual video_sequence_length) for mu calculation, - # matching diffusers' LTX2Pipeline which hardcodes this value. - mu = calculate_shift( - self.scheduler.config.get("max_image_seq_len", 4096), - 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) - video_audio_scheduler = self._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, _ = retrieve_timesteps( - self.scheduler, - request_inputs.num_inference_steps, - device, - timesteps, - sigmas=sigmas, - mu=mu, - ) - self._num_timesteps = len(timesteps) - return audio_scheduler, video_audio_scheduler, timesteps - - def _prepare_rope_coords_stage( - self, - forward_ctx: _LTX23ForwardContext, - latents: torch.Tensor, - audio_latents: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - request_inputs = forward_ctx.request_inputs - video_coords = self.transformer.rope.prepare_video_coords( - latents.shape[0], - forward_ctx.latent_num_frames, - forward_ctx.latent_height, - forward_ctx.latent_width, - latents.device, - fps=request_inputs.frame_rate, - ) - audio_coords = self.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 _prepare_denoise_context_for_cfg( - self, - forward_ctx: _LTX23ForwardContext, - denoise_ctx: _LTX23DenoiseContext, - ) -> _LTX23DenoiseContext: - if self.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 _denoise_timestep_kwargs( - self, - ts: torch.Tensor, - forward_ctx: _LTX23ForwardContext, - denoise_ctx: _LTX23DenoiseContext, - ) -> dict[str, torch.Tensor]: - return {"timestep": ts, "sigma": ts} - - def _build_transformer_kwargs( - self, - forward_ctx: _LTX23ForwardContext, - denoise_ctx: _LTX23DenoiseContext, - *, - 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, - audio_encoder_attention_mask: torch.Tensor, - ts: torch.Tensor, - ) -> dict[str, Any]: - request_inputs = forward_ctx.request_inputs - 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, - **self._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": 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( - 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]: - if forward_ctx.cfg_parallel_ready: - 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, - ) - - latents = self.scheduler.step(noise_pred_video, t, denoise_ctx.latents, return_dict=False)[0] - audio_latents = forward_ctx.audio_scheduler.step( - noise_pred_audio, - t, - denoise_ctx.audio_latents, - return_dict=False, - )[0] - return latents, audio_latents - - def _denoise_loop( - self, - forward_ctx: _LTX23ForwardContext, - denoise_ctx: _LTX23DenoiseContext, - ) -> tuple[torch.Tensor, torch.Tensor]: - request_inputs = forward_ctx.request_inputs - prompt_context = forward_ctx.prompt_context - guidance_scale = request_inputs.guidance_scale - audio_scheduler = forward_ctx.audio_scheduler - - with self.progress_bar(total=len(forward_ctx.timesteps)) as pbar: - for i, t in enumerate(forward_ctx.timesteps): - if self.interrupt: - continue - - self._current_timestep = t - - if forward_ctx.cfg_parallel_ready: - latent_model_input = denoise_ctx.latents.to(prompt_context.positive_connector_prompt_embeds.dtype) - audio_latent_model_input = denoise_ctx.audio_latents.to( - prompt_context.positive_connector_prompt_embeds.dtype - ) - ts = t.expand(latent_model_input.shape[0]) - positive_kwargs = self._build_transformer_kwargs( - forward_ctx, - denoise_ctx, - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_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, - } - noise_pred_video, noise_pred_audio = self.predict_noise_with_parallel_cfg( - true_cfg_scale=guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - video_latents=denoise_ctx.latents, - audio_latents=denoise_ctx.audio_latents, - video_sigma=self.scheduler.sigmas[i], - audio_sigma=audio_scheduler.sigmas[i], - ) - else: - latent_model_input = ( - torch.cat([denoise_ctx.latents] * 2) - if self.do_classifier_free_guidance - else denoise_ctx.latents - ) - latent_model_input = latent_model_input.to(prompt_context.connector_prompt_embeds.dtype) - audio_latent_model_input = ( - torch.cat([denoise_ctx.audio_latents] * 2) - if self.do_classifier_free_guidance - else denoise_ctx.audio_latents - ) - audio_latent_model_input = audio_latent_model_input.to(prompt_context.connector_prompt_embeds.dtype) - ts = t.expand(latent_model_input.shape[0]) - - transformer_kwargs = self._build_transformer_kwargs( - forward_ctx, - denoise_ctx, - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_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 self._transformer_cache_context("cond_uncond"): - noise_pred_video, noise_pred_audio = self.transformer(**transformer_kwargs) - - noise_pred_video = noise_pred_video.float() - noise_pred_audio = noise_pred_audio.float() - - if self.do_classifier_free_guidance: - noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2) - noise_pred_video = self._combine_x0_space_cfg( - denoise_ctx.latents, - noise_pred_video_cond, - noise_pred_video_uncond, - self.scheduler.sigmas[i], - guidance_scale, - ) - - noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2) - noise_pred_audio = self._combine_x0_space_cfg( - denoise_ctx.audio_latents, - noise_pred_audio_cond, - noise_pred_audio_uncond, - audio_scheduler.sigmas[i], - guidance_scale, - ) - - denoise_ctx.latents, denoise_ctx.audio_latents = self._step_denoised_latents( - forward_ctx, - denoise_ctx, - noise_pred_video, - noise_pred_audio, - t, - ) - pbar.update() - - return denoise_ctx.latents, denoise_ctx.audio_latents - - def _unpack_and_denormalize_stage( - self, - forward_ctx: _LTX23ForwardContext, - latents: torch.Tensor, - audio_latents: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - latents = self._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 = 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, forward_ctx.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, - forward_ctx.original_audio_num_frames, - num_mel_bins=forward_ctx.latent_mel_bins, - ) - return latents, audio_latents - - def _decode_and_split( - self, - forward_ctx: _LTX23ForwardContext, - latents: torch.Tensor, - audio_latents: torch.Tensor, - ) -> list[DiffusionOutput]: - request_inputs = forward_ctx.request_inputs - return split_diffusion_output_by_request( - self._decode_output( - latents=latents, - audio_latents=audio_latents, - 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, - ), - forward_ctx.req, - num_outputs_per_prompt=forward_ctx.num_videos_per_prompt, - ) - - def _forward_impl( - self, - req: DiffusionRequestBatch, - request_inputs: _LTX23RequestInputs, - *, - noise_scale: float, - sigmas: list[float] | None, - timesteps: list[int] | None, - attention_kwargs: dict[str, Any] | None, - image: Any | None = None, - ) -> list[DiffusionOutput]: - self._check_forward_inputs(request_inputs, image=image) - cfg_parallel_ready = self._setup_forward_runtime(request_inputs, attention_kwargs) - device = self.device - prompt_context = self._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 = self._resolve_video_latent_dimensions(request_inputs) - latents, conditioning_mask = self._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 = ( - self._prepare_audio_latents_stage( - request_inputs, - prompt_context, - device=device, - noise_scale=noise_scale, - ) - ) - audio_scheduler, video_audio_scheduler, timesteps_tensor = self._prepare_scheduler_stage( - request_inputs, - device=device, - sigmas=sigmas, - timesteps=timesteps, - latent_num_frames=latent_num_frames, - latent_height=latent_height, - latent_width=latent_width, - ) - forward_ctx = _LTX23ForwardContext( - 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 = self._prepare_rope_coords_stage(forward_ctx, latents, audio_latents) - denoise_ctx = _LTX23DenoiseContext( - latents=latents, - audio_latents=audio_latents, - video_coords=video_coords, - audio_coords=audio_coords, - conditioning_mask=conditioning_mask, - ) - denoise_ctx = self._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx) - latents, audio_latents = self._denoise_loop(forward_ctx, denoise_ctx) - latents, audio_latents = self._unpack_and_denormalize_stage(forward_ctx, latents, audio_latents) - return self._decode_and_split(forward_ctx, latents, audio_latents) - - # ------------------------------------------------------------------ - # Forward pass - # ------------------------------------------------------------------ - - @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, - 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, - ) - return self._forward_impl( - req, - request_inputs, - noise_scale=noise_scale, - sigmas=sigmas, - timesteps=timesteps, - attention_kwargs=attention_kwargs, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py deleted file mode 100644 index bcccd7596..000000000 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py +++ /dev/null @@ -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", -] diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py deleted file mode 100644 index a1caf6ce8..000000000 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ /dev/null @@ -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__() diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py new file mode 100644 index 000000000..8b786b817 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py @@ -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, + ) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 5b8b2fba1..22c36f32b 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -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": (