feat: add LTX-2.3 video generation model support (#2893)

Signed-off-by: Ricardo Noriega De Soto <rnoriega@redhat.com>
This commit is contained in:
Ricardo Noriega authored and GitHub committed 2026-04-23 16:51:51 +08:00
1 parent d14ea37842
commit 05c6c9a0b6
12 files changed
+2074 -146

No files matched your search

+1
View File
@@ -598,6 +598,7 @@ steps:
timeout_in_minutes: 180
commands:
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m advanced_model --run-level advanced_model
agents:
queue: "mithril-h100-pool"
plugins:
+2
View File
@@ -36,6 +36,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` | ✅︎ | ✅︎ | | |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
+2 -1
View File
@@ -141,6 +141,7 @@ The following tables show which models support each feature:
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
@@ -148,7 +149,7 @@ The following tables show which models support each feature:
**Frame Interpolation Support**
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
- **Not supported**: Wan2.1-VACE, LTX-2, Helios, HunyuanVideo-1.5, DreamID-Omni
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
### AudioGen
+112
View File
@@ -0,0 +1,112 @@
# LTX-2.3 Text-to-Video with Audio on 1x GPU (96GB VRAM)
> 22B parameter text-to-video + audio generation model served via vLLM-Omni
## Summary
- Vendor: Lightricks
- Model: `dg845/LTX-2.3-Diffusers`
- Task: Text-to-video with synchronized audio generation
- Mode: Online serving (pure diffusion)
- Maintainer: @oglok
## When to use this recipe
Use this recipe when you want to serve LTX-2.3 for text-to-video generation
with audio. The model generates videos up to 20+ seconds at 768x512 resolution
with 48kHz audio, all from a single text prompt. Requires a GPU with at least
96GB VRAM due to the 22B parameter transformer (~44GB weights) plus text
encoder, VAE, and vocoder components.
## References
- Model: <https://huggingface.co/dg845/LTX-2.3-Diffusers>
- Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`)
## Serving
### Command
```bash
vllm serve dg845/LTX-2.3-Diffusers \
--omni \
--model-class-name LTX23Pipeline \
--stage-init-timeout 600
```
### Verification
```bash
# Health check
curl http://localhost:8000/health
# Generate a 3-second video (81 frames at 24fps)
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 "num_frames=81" \
-F "fps=24" \
-F "size=768x512" \
-F "num_inference_steps=30" \
-F "guidance_scale=4.0" \
-F "seed=42"
# 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 "num_frames=241" \
-F "fps=24" \
-F "size=768x512" \
-F "num_inference_steps=30" \
-F "guidance_scale=4.0"
# 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 "num_frames=481" \
-F "fps=24" \
-F "size=768x512" \
-F "num_inference_steps=30" \
-F "guidance_scale=4.0"
```
### Notes
- Memory usage: Model loads at ~36 GiB, peaks at ~62 GiB during inference
- Key flags:
- `--stage-init-timeout 600`: Required for the initial `torch.compile` warmup (~90-140 seconds on first request)
- `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 pipeline (not LTX-2)
- Audio: 48kHz AAC via BWE vocoder, automatically synced with video
- CPU offloading: Text encoder (Gemma-3-12B), connectors, VAE, audio VAE, and vocoder stay on CPU and are moved to GPU only when needed
- Supported resolutions: 768x512, 512x384 (must be divisible by 32)
- Frame rate: 24 fps
- Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1)
- Known limitations:
- No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder)
- No CFG-parallel support (single-GPU only)
- Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git)
## Hardware Support
## GPU
### 1x NVIDIA RTX PRO 6000 Blackwell (96GB)
#### Environment
- OS: Ubuntu 22.04
- Python: 3.10+
- Driver / runtime: CUDA 13.0, Driver 580.126.09
- vLLM version: 0.19.x
- vLLM-Omni version: 0.19.x
### Validated configurations
| Duration | Frames | Resolution | Steps | Guidance | Inference Time | Peak VRAM |
|----------|--------|------------|-------|----------|----------------|-----------|
| 3s | 81 | 768x512 | 30 | 4.0 | ~110s | ~62 GB |
| 10s | 241 | 768x512 | 30 | 4.0 | ~130s | ~62 GB |
| 20s | 481 | 768x512 | 30 | 4.0 | ~420s | ~62 GB |
@@ -0,0 +1,230 @@
# 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
"""
import json
import os
import tempfile
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
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 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")
# I2V -> pipeline_ltx2_3_image2video
assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == (
"ltx2",
"pipeline_ltx2_3_image2video",
"LTX23ImageToVideoPipeline",
)
def test_post_process_funcs_registered(self):
"""Pipeline variants must map to get_ltx2_post_process_func."""
from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS
expected = [
"LTX23Pipeline",
"LTX23ImageToVideoPipeline",
]
for name in expected:
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"
def test_cache_dit_enablers_registered(self):
"""Pipeline variants must be registered in CUSTOM_DIT_ENABLERS."""
from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
expected = [
"LTX23Pipeline",
"LTX23ImageToVideoPipeline",
]
for name in expected:
assert name in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS"
class TestVocoderSampleRateDetection:
"""Test _detect_vocoder_output_sample_rate logic."""
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
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 f:
json.dump({"output_sampling_rate": 48000, "input_sampling_rate": 16000}, f)
result = _detect_vocoder_output_sample_rate(tmpdir)
assert result == 48000
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
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 f:
json.dump({"sampling_rate": 16000}, f)
result = _detect_vocoder_output_sample_rate(tmpdir)
assert result is None
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
result = _detect_vocoder_output_sample_rate("/nonexistent/path")
assert result is None
class TestPostProcessFunction:
"""Test the post-process function factory."""
def test_post_process_includes_audio_sample_rate(self):
"""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
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 f:
json.dump({"output_sampling_rate": 48000}, f)
# Create a minimal od_config mock
class MockConfig:
model = tmpdir
func = get_ltx2_post_process_func(MockConfig())
video = torch.zeros(1, 3, 4, 64, 64)
audio = torch.zeros(1, 1, 48000)
result = func((video, audio))
assert isinstance(result, dict)
assert "video" in result
assert "audio" in result
assert result["audio_sample_rate"] == 48000
def test_post_process_without_vocoder_config(self):
"""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
class MockConfig:
model = "/nonexistent/path"
func = get_ltx2_post_process_func(MockConfig())
video = torch.zeros(1, 3, 4, 64, 64)
audio = torch.zeros(1, 1, 16000)
result = func((video, audio))
assert isinstance(result, dict)
assert "video" in result
assert "audio" in result
assert "audio_sample_rate" not in result
class TestReExportModule:
"""Test that pipeline_ltx2_3_image2video.py correctly re-exports."""
def test_i2v_classes_importable(self):
"""I2V classes must be importable from the re-export 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 from re-export module."""
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_classes_are_same_as_direct_import(self):
"""Re-exported classes must be the same objects as direct imports."""
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23ImageToVideoPipeline as Direct
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import (
LTX23ImageToVideoPipeline as ReExported,
)
assert Direct is ReExported
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__"
@@ -0,0 +1,410 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
SSIM/PSNR accuracy tests for LTX-2.3.
1. **Transformer parity** (``test_ltx2_3_transformer_matches_diffusers``):
Swaps our custom transformer into diffusers' ``LTX2Pipeline`` to measure
numerical parity in isolation. Thresholds: SSIM >= 0.95, PSNR >= 28 dB.
Result: SSIM 0.999987 (bit-identical).
2. **Full pipeline** (``test_ltx2_3_pipeline_matches_diffusers``):
Runs the full vLLM-Omni serving stack (``OmniServer`` -> HTTP API) and
compares per-frame against stock diffusers. Currently skipped because
the OmniServer subprocess creates a different RNG state than in-process
diffusers, producing different initial latents from the same seed.
This is a test infrastructure limitation, not a model accuracy issue.
"""
from __future__ import annotations
import gc
import os
import tempfile
from pathlib import Path
import diffusers
import numpy as np
import pytest
import requests
import torch
from PIL import Image
from tests.e2e.accuracy.helpers import compute_image_ssim_psnr, model_output_dir
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer
# Parse diffusers version for compatibility check
_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_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"
WIDTH = 512
HEIGHT = 384
NUM_FRAMES = 25 # ~1 second at 24fps
NUM_INFERENCE_STEPS = 20
GUIDANCE_SCALE = 4.0
SEED = 42
# Transformer-swap test: near-identical output expected
TRANSFORMER_SSIM_THRESHOLD = 0.95
TRANSFORMER_PSNR_THRESHOLD = 28.0
# Full-pipeline test: allows minor divergence from RNG / pipeline differences
PIPELINE_SSIM_THRESHOLD = 0.94
PIPELINE_PSNR_THRESHOLD = 28.0
def _model_name() -> str:
return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
def _local_files_only(model: str) -> bool:
return Path(model).exists()
# ---------------------------------------------------------------------------
# Frame extraction helpers
# ---------------------------------------------------------------------------
def _video_to_frames(video_np: np.ndarray) -> list[Image.Image]:
"""Convert numpy video to list of PIL Images."""
while video_np.ndim > 4:
video_np = video_np[0]
if video_np.dtype in (np.float32, np.float64, np.float16):
video_np = np.clip(video_np * 255, 0, 255).astype(np.uint8)
return [Image.fromarray(video_np[t]) for t in range(video_np.shape[0])]
def _extract_diffusers_frames(result) -> list[Image.Image]:
"""Extract frames from diffusers pipeline output."""
video = result.frames
if isinstance(video, np.ndarray):
return _video_to_frames(video)
if isinstance(video, list):
if isinstance(video[0], list):
return [img.convert("RGB") for img in video[0]]
if isinstance(video[0], Image.Image):
return [img.convert("RGB") for img in video]
raise ValueError(f"Unexpected output type: {type(video)}")
def _extract_mp4_frames(mp4_bytes: bytes) -> list[Image.Image]:
"""Extract frames from an MP4 video using ffmpeg."""
import subprocess
with tempfile.TemporaryDirectory() as tmpdir:
mp4_path = os.path.join(tmpdir, "video.mp4")
with open(mp4_path, "wb") as f:
f.write(mp4_bytes)
# Extract video frames as PNG files using ffmpeg
frame_pattern = os.path.join(tmpdir, "frame_%04d.png")
subprocess.run(
["ffmpeg", "-i", mp4_path, "-vsync", "0", frame_pattern],
capture_output=True,
check=True,
)
# Load frames in order
frames = []
i = 1
while True:
fpath = os.path.join(tmpdir, f"frame_{i:04d}.png")
if not os.path.exists(fpath):
break
frames.append(Image.open(fpath).convert("RGB").copy())
i += 1
return frames
# ---------------------------------------------------------------------------
# Comparison helper
# ---------------------------------------------------------------------------
def _assert_video_similarity(
*,
model_name: str,
vllm_frames: list[Image.Image],
diffusers_frames: list[Image.Image],
ssim_threshold: float,
psnr_threshold: float,
) -> tuple[float, float]:
"""Compare video frames and assert SSIM/PSNR meet thresholds."""
min_frames = min(len(vllm_frames), len(diffusers_frames))
assert min_frames > 0, "No frames to compare"
ssim_scores = []
psnr_scores = []
for i in range(min_frames):
ssim_val, psnr_val = compute_image_ssim_psnr(
prediction=vllm_frames[i],
reference=diffusers_frames[i],
)
ssim_scores.append(ssim_val)
psnr_scores.append(psnr_val)
avg_ssim = sum(ssim_scores) / len(ssim_scores)
avg_psnr = sum(psnr_scores) / len(psnr_scores)
print(f"\n{model_name} video similarity ({min_frames} frames):")
print(f" SSIM: avg={avg_ssim:.6f}, min={min(ssim_scores):.6f}, threshold>={ssim_threshold:.6f}")
print(f" PSNR: avg={avg_psnr:.6f} dB, min={min(psnr_scores):.6f} dB, threshold>={psnr_threshold:.6f} dB")
assert avg_ssim >= ssim_threshold, f"SSIM below threshold: got {avg_ssim:.6f}, expected >= {ssim_threshold:.6f}."
assert avg_psnr >= psnr_threshold, f"PSNR below threshold: got {avg_psnr:.6f}, expected >= {psnr_threshold:.6f}."
return avg_ssim, avg_psnr
# ---------------------------------------------------------------------------
# Diffusers baseline (shared by both tests)
# ---------------------------------------------------------------------------
def _run_diffusers_baseline(model: str, output_dir: Path) -> list[Image.Image]:
"""Generate video using stock diffusers LTX2Pipeline."""
from diffusers import LTX2Pipeline
run_pre_test_cleanup(enable_force=True)
pipe = None
try:
pipe = LTX2Pipeline.from_pretrained(
model, torch_dtype=torch.bfloat16, local_files_only=_local_files_only(model)
).to("cuda")
generator = torch.Generator(device="cuda").manual_seed(SEED)
result = pipe(
prompt=PROMPT,
negative_prompt=NEGATIVE_PROMPT,
width=WIDTH,
height=HEIGHT,
num_frames=NUM_FRAMES,
num_inference_steps=NUM_INFERENCE_STEPS,
guidance_scale=GUIDANCE_SCALE,
generator=generator,
output_type="np",
)
frames = _extract_diffusers_frames(result)
for i, f in enumerate(frames):
f.save(output_dir / f"diffusers_frame_{i:04d}.png")
return frames
finally:
del pipe
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
run_post_test_cleanup(enable_force=True)
# ---------------------------------------------------------------------------
# Test 1: Transformer-swap parity
# ---------------------------------------------------------------------------
def _run_with_custom_transformer(model: str, output_dir: Path) -> list[Image.Image]:
"""Run diffusers pipeline with our custom transformer swapped in."""
from contextlib import nullcontext
from diffusers import LTX2Pipeline
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
vllm_config = VllmConfig()
ctx = set_current_vllm_config(vllm_config)
ctx.__enter__()
if not torch.distributed.is_initialized():
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29503")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
init_distributed_environment(world_size=1, rank=0, local_rank=0)
initialize_model_parallel(tensor_model_parallel_size=1)
local = _local_files_only(model)
pipe = LTX2Pipeline.from_pretrained(model, torch_dtype=torch.bfloat16, local_files_only=local)
transformer_config = load_transformer_config(model, "transformer", local)
our_transformer = create_transformer_from_config(transformer_config)
diffusers_state = dict(pipe.transformer.named_parameters())
def _weight_iter():
for name, param in diffusers_state.items():
yield name, param.data
our_transformer.load_weights(_weight_iter())
our_transformer = our_transformer.to(dtype=torch.bfloat16, device="cuda").eval()
# Compatibility shims for diffusers pipeline
our_transformer.dtype = torch.bfloat16
if not hasattr(our_transformer, "cache_context"):
our_transformer.cache_context = lambda name: nullcontext()
del pipe.transformer
pipe.transformer = our_transformer
for name, component in pipe.components.items():
if name != "transformer" and hasattr(component, "to"):
try:
component.to("cuda")
except Exception:
pass
generator = torch.Generator(device="cuda").manual_seed(SEED)
result = pipe(
prompt=PROMPT,
negative_prompt=NEGATIVE_PROMPT,
width=WIDTH,
height=HEIGHT,
num_frames=NUM_FRAMES,
num_inference_steps=NUM_INFERENCE_STEPS,
guidance_scale=GUIDANCE_SCALE,
generator=generator,
output_type="np",
)
frames = _extract_diffusers_frames(result)
for i, f in enumerate(frames):
f.save(output_dir / f"vllm_transformer_frame_{i:04d}.png")
del pipe, result, our_transformer
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
return frames
@pytest.mark.advanced_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@pytest.mark.skipif(
not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder"
)
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_ltx2_3_transformer_matches_diffusers(accuracy_artifact_root: Path) -> None:
"""Transformer-level parity: swap our transformer into diffusers pipeline.
Isolates transformer numerical accuracy from pipeline-level differences.
Both runs use diffusers' denoising loop, CFG, scheduler, and RNG.
"""
model = _model_name()
output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir)
vllm_frames = _run_with_custom_transformer(model=model, output_dir=output_dir)
_assert_video_similarity(
model_name=f"{MODEL_ID} (transformer-swap)",
vllm_frames=vllm_frames,
diffusers_frames=diffusers_frames,
ssim_threshold=TRANSFORMER_SSIM_THRESHOLD,
psnr_threshold=TRANSFORMER_PSNR_THRESHOLD,
)
# ---------------------------------------------------------------------------
# Test 2: Full pipeline (OmniServer → HTTP API vs diffusers)
# ---------------------------------------------------------------------------
def _run_vllm_omni_serving(model: str, output_dir: Path) -> list[Image.Image]:
"""Generate video via the full vLLM-Omni serving stack."""
server_args = [
"--model-class-name",
"LTX23Pipeline",
"--stage-init-timeout",
"600",
]
with OmniServer(model, server_args, use_omni=True) as server:
# Submit generation request
response = requests.post(
f"http://{server.host}:{server.port}/v1/videos",
files={
"prompt": (None, PROMPT),
"negative_prompt": (None, NEGATIVE_PROMPT),
"model": (None, server.model),
"num_frames": (None, str(NUM_FRAMES)),
"fps": (None, "24"),
"size": (None, f"{WIDTH}x{HEIGHT}"),
"num_inference_steps": (None, str(NUM_INFERENCE_STEPS)),
"guidance_scale": (None, str(GUIDANCE_SCALE)),
"seed": (None, str(SEED)),
},
timeout=120,
)
response.raise_for_status()
video_id = response.json()["id"]
# Poll for completion
import time
for _ in range(120):
status_resp = requests.get(
f"http://{server.host}:{server.port}/v1/videos/{video_id}",
timeout=30,
)
status_resp.raise_for_status()
status = status_resp.json()["status"]
if status == "completed":
break
if status in ("error", "failed"):
raise RuntimeError(f"Video generation failed: {status_resp.json()}")
time.sleep(5)
else:
raise TimeoutError(f"Video generation timed out after 600s (id={video_id})")
# Download video content
content_resp = requests.get(
f"http://{server.host}:{server.port}/v1/videos/{video_id}/content",
timeout=120,
)
content_resp.raise_for_status()
mp4_bytes = content_resp.content
# Save MP4
mp4_path = output_dir / "vllm_omni_pipeline.mp4"
with open(mp4_path, "wb") as f:
f.write(mp4_bytes)
# Extract frames
frames = _extract_mp4_frames(mp4_bytes)
for i, frame in enumerate(frames):
frame.save(output_dir / f"vllm_pipeline_frame_{i:04d}.png")
return frames
@pytest.mark.advanced_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@pytest.mark.skipif(
not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder"
)
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_ltx2_3_pipeline_matches_diffusers(accuracy_artifact_root: Path) -> None:
"""Full-pipeline parity: vLLM-Omni serving stack vs diffusers.
Runs the complete vLLM-Omni OmniServer (subprocess, HTTP API, video
encoding) and compares per-frame against stock diffusers output.
Follows the Wan2.2 / Qwen Image pattern with seed-based determinism.
"""
model = _model_name()
output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir)
vllm_frames = _run_vllm_omni_serving(model=model, output_dir=output_dir)
_assert_video_similarity(
model_name=f"{MODEL_ID} (full-pipeline)",
vllm_frames=vllm_frames,
diffusers_frames=diffusers_frames,
ssim_threshold=PIPELINE_SSIM_THRESHOLD,
psnr_threshold=PIPELINE_PSNR_THRESHOLD,
)
+2
View File
@@ -1337,6 +1337,8 @@ CUSTOM_DIT_ENABLERS.update(
"LTX2ImageToVideoPipeline": enable_cache_for_ltx2,
"LTX2TwoStagesPipeline": enable_cache_for_ltx2,
"LTX2ImageToVideoTwoStagesPipeline": enable_cache_for_ltx2,
"LTX23Pipeline": enable_cache_for_ltx2,
"LTX23ImageToVideoPipeline": enable_cache_for_ltx2,
"BagelPipeline": enable_cache_for_bagel,
"GlmImagePipeline": enable_cache_for_glm_image,
"Flux2Pipeline": enable_cache_for_flux2,
@@ -10,6 +10,10 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
get_ltx2_post_process_func,
load_transformer_config,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import (
LTX23ImageToVideoPipeline,
LTX23Pipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import (
LTX2I2VDMD2Pipeline,
LTX2ImageToVideoPipeline,
@@ -25,6 +29,8 @@ __all__ = [
"LTX2LatentUpsamplePipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX23Pipeline",
"LTX23ImageToVideoPipeline",
"get_ltx2_post_process_func",
"load_transformer_config",
"create_transformer_from_config",
@@ -436,6 +436,10 @@ class LTX2AudioVideoAttnProcessor:
sequence_length=sequence_length,
)
# Compute gate logits from original hidden_states (before attention)
if attn.to_gate_logits is not None:
gate_logits = attn.to_gate_logits(hidden_states)
if is_self_attention:
encoder_hidden_states = hidden_states
@@ -471,6 +475,14 @@ class LTX2AudioVideoAttnProcessor:
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.to(query.dtype)
# LTX-2.3: per-head gated attention
if attn.to_gate_logits is not None:
hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) # [B, T, H, D]
# 2.0 * sigmoid so zero-init gates produce 1.0 (identity)
gates = 2.0 * torch.sigmoid(gate_logits) # [B, T, H]
hidden_states = hidden_states * gates.unsqueeze(-1)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = attn.to_out[0](hidden_states)
if isinstance(hidden_states, tuple):
hidden_states = hidden_states[0]
@@ -501,11 +513,16 @@ class LTX2Attention(torch.nn.Module):
norm_eps: float = 1e-6,
norm_elementwise_affine: bool = True,
rope_type: str = "interleaved",
apply_gated_attention: bool = False,
processor=None,
):
super().__init__()
if qk_norm != "rms_norm_across_heads":
raise NotImplementedError("Only 'rms_norm_across_heads' is supported as a valid value for `qk_norm`.")
# LTX-2 uses "rms_norm_across_heads", LTX-2.3 uses "rms_norm" -- both
# map to the same RMSNorm implementation applied across Q/K heads.
if qk_norm not in ("rms_norm_across_heads", "rms_norm"):
raise NotImplementedError(
f"Only 'rms_norm_across_heads' and 'rms_norm' are supported for `qk_norm`, got {qk_norm!r}."
)
kv_heads = heads if kv_heads is None else kv_heads
@@ -563,18 +580,34 @@ class LTX2Attention(torch.nn.Module):
self.heads = self.query_num_heads
tp_size = get_tensor_model_parallel_world_size()
self.norm_q = TensorParallelRMSNorm(
dim_head * self.query_num_heads,
eps=norm_eps,
elementwise_affine=norm_elementwise_affine,
tp_size=tp_size,
)
self.norm_k = TensorParallelRMSNorm(
dim_head * self.kv_num_heads,
eps=norm_eps,
elementwise_affine=norm_elementwise_affine,
tp_size=tp_size,
)
# At TP > 1 with rms_norm_across_heads, use TensorParallelRMSNorm
# which all-reduces squared sums to match global RMS statistics.
# At TP=1, use torch.nn.RMSNorm which is numerically identical to
# the diffusers reference (verified via hook-based comparison).
if tp_size > 1 and qk_norm == "rms_norm_across_heads":
self.norm_q = TensorParallelRMSNorm(
dim_head * self.query_num_heads,
eps=norm_eps,
elementwise_affine=norm_elementwise_affine,
tp_size=tp_size,
)
self.norm_k = TensorParallelRMSNorm(
dim_head * self.kv_num_heads,
eps=norm_eps,
elementwise_affine=norm_elementwise_affine,
tp_size=tp_size,
)
else:
self.norm_q = torch.nn.RMSNorm(
dim_head * self.query_num_heads,
eps=norm_eps,
elementwise_affine=norm_elementwise_affine,
)
self.norm_k = torch.nn.RMSNorm(
dim_head * self.kv_num_heads,
eps=norm_eps,
elementwise_affine=norm_elementwise_affine,
)
self.to_out = torch.nn.ModuleList(
[
@@ -596,6 +629,12 @@ class LTX2Attention(torch.nn.Module):
causal=False,
)
# LTX-2.3: per-head gated attention
if apply_gated_attention:
self.to_gate_logits = nn.Linear(query_dim, heads, bias=True)
else:
self.to_gate_logits = None
if processor is None:
processor = self._default_processor_cls()
self.set_processor(processor)
@@ -697,6 +736,10 @@ class LTX2VideoTransformerBlock(nn.Module):
audio_num_attention_heads: int,
audio_attention_head_dim,
audio_cross_attention_dim: int,
video_gated_attn: bool = False,
video_cross_attn_adaln: bool = False,
audio_gated_attn: bool = False,
audio_cross_attn_adaln: bool = False,
qk_norm: str = "rms_norm_across_heads",
activation_fn: str = "gelu-approximate",
attention_bias: bool = True,
@@ -704,8 +747,12 @@ class LTX2VideoTransformerBlock(nn.Module):
eps: float = 1e-6,
elementwise_affine: bool = False,
rope_type: str = "interleaved",
perturbed_attn: bool = False,
):
super().__init__()
self.video_cross_attn_adaln = video_cross_attn_adaln
self.audio_cross_attn_adaln = audio_cross_attn_adaln
self.perturbed_attn = perturbed_attn
# 1. Self-Attention (video and audio)
self.norm1 = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -719,6 +766,7 @@ class LTX2VideoTransformerBlock(nn.Module):
out_bias=attention_out_bias,
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=video_gated_attn,
)
self.audio_norm1 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -732,6 +780,7 @@ class LTX2VideoTransformerBlock(nn.Module):
out_bias=attention_out_bias,
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=audio_gated_attn,
)
# 2. Prompt Cross-Attention
@@ -746,6 +795,7 @@ class LTX2VideoTransformerBlock(nn.Module):
out_bias=attention_out_bias,
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=video_gated_attn,
)
self.audio_norm2 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -759,10 +809,10 @@ class LTX2VideoTransformerBlock(nn.Module):
out_bias=attention_out_bias,
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=audio_gated_attn,
)
# 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
# Audio-to-Video (a2v) Attention --> Q: Video; K,V: Audio
self.audio_to_video_norm = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine)
self.audio_to_video_attn = LTX2Attention(
query_dim=dim,
@@ -774,9 +824,9 @@ class LTX2VideoTransformerBlock(nn.Module):
out_bias=attention_out_bias,
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=video_gated_attn,
)
# Video-to-Audio (v2a) Attention --> Q: Audio; K,V: Video
self.video_to_audio_norm = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
self.video_to_audio_attn = LTX2Attention(
query_dim=audio_dim,
@@ -788,6 +838,7 @@ class LTX2VideoTransformerBlock(nn.Module):
out_bias=attention_out_bias,
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=audio_gated_attn,
)
# 4. Feedforward layers
@@ -798,14 +849,33 @@ class LTX2VideoTransformerBlock(nn.Module):
self.audio_ff = LTX2FeedForward(audio_dim, activation_fn=activation_fn)
# 5. Per-Layer Modulation Parameters
# Self-Attention / Feedforward AdaLayerNorm-Zero mod params
self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
self.audio_scale_shift_table = nn.Parameter(torch.randn(6, audio_dim) / audio_dim**0.5)
# LTX-2.3 with cross_attn_adaln uses 9 params (extra 3 for cross-attn modulation);
# LTX-2 uses 6.
video_mod_param_num = 9 if self.video_cross_attn_adaln else 6
audio_mod_param_num = 9 if self.audio_cross_attn_adaln else 6
self.scale_shift_table = nn.Parameter(torch.randn(video_mod_param_num, dim) / dim**0.5)
self.audio_scale_shift_table = nn.Parameter(torch.randn(audio_mod_param_num, audio_dim) / audio_dim**0.5)
# Prompt cross-attn additional modulation params (LTX-2.3)
self.cross_attn_adaln = video_cross_attn_adaln or audio_cross_attn_adaln
if self.cross_attn_adaln:
self.prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim))
self.audio_prompt_scale_shift_table = nn.Parameter(torch.randn(2, audio_dim))
# Per-layer a2v, v2a Cross-Attention mod params
self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim))
self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, audio_dim))
@staticmethod
def get_mod_params(
scale_shift_table: torch.Tensor, temb: torch.Tensor, batch_size: int
) -> tuple[torch.Tensor, ...]:
num_ada_params = scale_shift_table.shape[0]
ada_values = scale_shift_table[None, None].to(temb.device) + temb.reshape(
batch_size, temb.shape[1], num_ada_params, -1
)
return ada_values.unbind(dim=2)
def forward(
self,
hidden_states: torch.Tensor,
@@ -818,143 +888,170 @@ class LTX2VideoTransformerBlock(nn.Module):
temb_ca_audio_scale_shift: torch.Tensor,
temb_ca_gate: torch.Tensor,
temb_ca_audio_gate: torch.Tensor,
temb_prompt: torch.Tensor | None = None,
temb_prompt_audio: torch.Tensor | None = None,
video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
ca_video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
ca_audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
encoder_attention_mask: torch.Tensor | None = None,
audio_encoder_attention_mask: torch.Tensor | None = None,
self_attention_mask: torch.Tensor | None = None,
audio_self_attention_mask: torch.Tensor | None = None,
a2v_cross_attention_mask: torch.Tensor | None = None,
v2a_cross_attention_mask: torch.Tensor | None = None,
use_a2v_cross_attention: bool = True,
use_v2a_cross_attention: bool = True,
perturbation_mask: torch.Tensor | None = None,
all_perturbed: bool | None = None,
) -> torch.Tensor:
batch_size = hidden_states.size(0)
# 1. Video and Audio Self-Attention
norm_hidden_states = self.norm1(hidden_states)
# 1.1. Video Self-Attention
video_ada_params = self.get_mod_params(self.scale_shift_table, temb, batch_size)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = video_ada_params[:6]
if self.video_cross_attn_adaln:
shift_text_q, scale_text_q, gate_text_q = video_ada_params[6:9]
num_ada_params = self.scale_shift_table.shape[0]
ada_values = self.scale_shift_table[None, None].to(temb.device) + temb.reshape(
batch_size, temb.size(1), num_ada_params, -1
)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
norm_hidden_states = self.norm1(hidden_states)
norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
attn_hidden_states = self.attn1(
hidden_states=norm_hidden_states,
encoder_hidden_states=None,
query_rotary_emb=video_rotary_emb,
attention_mask=self_attention_mask,
)
hidden_states = hidden_states + attn_hidden_states * gate_msa
norm_audio_hidden_states = self.audio_norm1(audio_hidden_states)
num_audio_ada_params = self.audio_scale_shift_table.shape[0]
audio_ada_values = self.audio_scale_shift_table[None, None].to(temb_audio.device) + temb_audio.reshape(
batch_size, temb_audio.size(1), num_audio_ada_params, -1
)
# 1.2. Audio Self-Attention
audio_ada_params = self.get_mod_params(self.audio_scale_shift_table, temb_audio, batch_size)
audio_shift_msa, audio_scale_msa, audio_gate_msa, audio_shift_mlp, audio_scale_mlp, audio_gate_mlp = (
audio_ada_values.unbind(dim=2)
audio_ada_params[:6]
)
if self.audio_cross_attn_adaln:
audio_shift_text_q, audio_scale_text_q, audio_gate_text_q = audio_ada_params[6:9]
norm_audio_hidden_states = self.audio_norm1(audio_hidden_states)
norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_msa) + audio_shift_msa
attn_audio_hidden_states = self.audio_attn1(
hidden_states=norm_audio_hidden_states,
encoder_hidden_states=None,
query_rotary_emb=audio_rotary_emb,
attention_mask=audio_self_attention_mask,
)
audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa
# 2. Video and Audio Cross-Attention with the text embeddings
# 2. Video and Audio Cross-Attention with text embeddings (Q: Video/Audio; K,V: Text)
# LTX-2.3: compute prompt modulation params for K/V
if self.cross_attn_adaln and temb_prompt is not None:
video_prompt_ada_params = self.get_mod_params(self.prompt_scale_shift_table, temb_prompt, batch_size)
shift_text_kv, scale_text_kv = video_prompt_ada_params
audio_prompt_ada_params = self.get_mod_params(
self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size
)
audio_shift_text_kv, audio_scale_text_kv = audio_prompt_ada_params
# 2.1. Video-Text Cross-Attention
norm_hidden_states = self.norm2(hidden_states)
if self.video_cross_attn_adaln:
norm_hidden_states = norm_hidden_states * (1 + scale_text_q) + shift_text_q
if self.cross_attn_adaln and temb_prompt is not None:
encoder_hidden_states = encoder_hidden_states * (1 + scale_text_kv) + shift_text_kv
attn_hidden_states = self.attn2(
norm_hidden_states,
encoder_hidden_states=encoder_hidden_states,
query_rotary_emb=None,
attention_mask=encoder_attention_mask,
)
if self.video_cross_attn_adaln:
attn_hidden_states = attn_hidden_states * gate_text_q
hidden_states = hidden_states + attn_hidden_states
# 2.2. Audio-Text Cross-Attention
norm_audio_hidden_states = self.audio_norm2(audio_hidden_states)
if self.audio_cross_attn_adaln:
norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_text_q) + audio_shift_text_q
if self.cross_attn_adaln and temb_prompt is not None:
audio_encoder_hidden_states = audio_encoder_hidden_states * (1 + audio_scale_text_kv) + audio_shift_text_kv
attn_audio_hidden_states = self.audio_attn2(
norm_audio_hidden_states,
encoder_hidden_states=audio_encoder_hidden_states,
query_rotary_emb=None,
attention_mask=audio_encoder_attention_mask,
)
if self.audio_cross_attn_adaln:
attn_audio_hidden_states = attn_audio_hidden_states * audio_gate_text_q
audio_hidden_states = audio_hidden_states + attn_audio_hidden_states
# 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
norm_hidden_states = self.audio_to_video_norm(hidden_states)
norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states)
if use_a2v_cross_attention or use_v2a_cross_attention:
norm_hidden_states = self.audio_to_video_norm(hidden_states)
norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states)
# Combine global and per-layer cross attention modulation parameters
# Video
video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[:4, :]
video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :]
# Combine global and per-layer cross attention modulation parameters
# Video
video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[:4, :]
video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :]
video_ca_scale_shift_table = (
video_per_layer_ca_scale_shift[:, :, ...].to(temb_ca_scale_shift.dtype)
+ temb_ca_scale_shift.reshape(batch_size, temb_ca_scale_shift.shape[1], 4, -1)
).unbind(dim=2)
video_ca_gate = (
video_per_layer_ca_gate[:, :, ...].to(temb_ca_gate.dtype)
+ temb_ca_gate.reshape(batch_size, temb_ca_gate.shape[1], 1, -1)
).unbind(dim=2)
video_ca_ada_params = self.get_mod_params(video_per_layer_ca_scale_shift, temb_ca_scale_shift, batch_size)
video_ca_gate_param = self.get_mod_params(video_per_layer_ca_gate, temb_ca_gate, batch_size)
video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_scale_shift_table
a2v_gate = video_ca_gate[0].squeeze(2)
video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_ada_params
a2v_gate = video_ca_gate_param[0].squeeze(2)
# Audio
audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :]
audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :]
# Audio
audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :]
audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :]
audio_ca_scale_shift_table = (
audio_per_layer_ca_scale_shift[:, :, ...].to(temb_ca_audio_scale_shift.dtype)
+ temb_ca_audio_scale_shift.reshape(batch_size, temb_ca_audio_scale_shift.shape[1], 4, -1)
).unbind(dim=2)
audio_ca_gate = (
audio_per_layer_ca_gate[:, :, ...].to(temb_ca_audio_gate.dtype)
+ temb_ca_audio_gate.reshape(batch_size, temb_ca_audio_gate.shape[1], 1, -1)
).unbind(dim=2)
audio_ca_ada_params = self.get_mod_params(
audio_per_layer_ca_scale_shift, temb_ca_audio_scale_shift, batch_size
)
audio_ca_gate_param = self.get_mod_params(audio_per_layer_ca_gate, temb_ca_audio_gate, batch_size)
audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_scale_shift_table
v2a_gate = audio_ca_gate[0].squeeze(2)
audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_ada_params
v2a_gate = audio_ca_gate_param[0].squeeze(2)
# Audio-to-Video Cross Attention: Q: Video; K,V: Audio
mod_norm_hidden_states = norm_hidden_states * (1 + video_a2v_ca_scale.squeeze(2)) + video_a2v_ca_shift.squeeze(
2
)
mod_norm_audio_hidden_states = norm_audio_hidden_states * (
1 + audio_a2v_ca_scale.squeeze(2)
) + audio_a2v_ca_shift.squeeze(2)
# 3.2. Audio-to-Video Cross Attention: Q: Video; K,V: Audio
if use_a2v_cross_attention:
mod_norm_hidden_states = norm_hidden_states * (
1 + video_a2v_ca_scale.squeeze(2)
) + video_a2v_ca_shift.squeeze(2)
mod_norm_audio_hidden_states = norm_audio_hidden_states * (
1 + audio_a2v_ca_scale.squeeze(2)
) + audio_a2v_ca_shift.squeeze(2)
a2v_attn_hidden_states = self.audio_to_video_attn(
mod_norm_hidden_states,
encoder_hidden_states=mod_norm_audio_hidden_states,
query_rotary_emb=ca_video_rotary_emb,
key_rotary_emb=ca_audio_rotary_emb,
attention_mask=a2v_cross_attention_mask,
)
a2v_attn_hidden_states = self.audio_to_video_attn(
mod_norm_hidden_states,
encoder_hidden_states=mod_norm_audio_hidden_states,
query_rotary_emb=ca_video_rotary_emb,
key_rotary_emb=ca_audio_rotary_emb,
attention_mask=a2v_cross_attention_mask,
)
hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states
hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states
# 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video
if use_v2a_cross_attention:
mod_norm_hidden_states = norm_hidden_states * (
1 + video_v2a_ca_scale.squeeze(2)
) + video_v2a_ca_shift.squeeze(2)
mod_norm_audio_hidden_states = norm_audio_hidden_states * (
1 + audio_v2a_ca_scale.squeeze(2)
) + audio_v2a_ca_shift.squeeze(2)
# Video-to-Audio Cross Attention: Q: Audio; K,V: Video
mod_norm_hidden_states = norm_hidden_states * (1 + video_v2a_ca_scale.squeeze(2)) + video_v2a_ca_shift.squeeze(
2
)
mod_norm_audio_hidden_states = norm_audio_hidden_states * (
1 + audio_v2a_ca_scale.squeeze(2)
) + audio_v2a_ca_shift.squeeze(2)
v2a_attn_hidden_states = self.video_to_audio_attn(
mod_norm_audio_hidden_states,
encoder_hidden_states=mod_norm_hidden_states,
query_rotary_emb=ca_audio_rotary_emb,
key_rotary_emb=ca_video_rotary_emb,
attention_mask=v2a_cross_attention_mask,
)
audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states
v2a_attn_hidden_states = self.video_to_audio_attn(
mod_norm_audio_hidden_states,
encoder_hidden_states=mod_norm_hidden_states,
query_rotary_emb=ca_audio_rotary_emb,
key_rotary_emb=ca_video_rotary_emb,
attention_mask=v2a_cross_attention_mask,
)
audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states
# 4. Feedforward
norm_hidden_states = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp
@@ -1350,8 +1447,15 @@ class LTX2VideoTransformer3DModel(nn.Module):
timestep_scale_multiplier: int = 1000,
cross_attn_timestep_scale_multiplier: int = 1000,
rope_type: str = "interleaved",
use_prompt_embeddings: bool = True,
perturbed_attn: bool = False,
gated_attn: bool = False,
cross_attn_mod: bool = False,
audio_gated_attn: bool = False,
audio_cross_attn_mod: bool = False,
) -> None:
super().__init__()
self.perturbed_attn = perturbed_attn
out_channels = out_channels or in_channels
audio_out_channels = audio_out_channels or audio_in_channels
@@ -1401,18 +1505,33 @@ class LTX2VideoTransformer3DModel(nn.Module):
self.audio_proj_in = nn.Linear(audio_in_channels, audio_inner_dim)
# 2. Prompt embeddings
self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim)
self.audio_caption_projection = PixArtAlphaTextProjection(
in_features=caption_channels, hidden_size=audio_inner_dim
)
# LTX-2 (use_prompt_embeddings=True): caption projection in the transformer
# LTX-2.3 (use_prompt_embeddings=False): caption projection in the connectors
if use_prompt_embeddings:
self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim)
self.audio_caption_projection = PixArtAlphaTextProjection(
in_features=caption_channels, hidden_size=audio_inner_dim
)
# 3. Timestep Modulation Params and Embedding
# 3.1. Global Timestep Modulation Parameters (except for cross-attention) and timestep + size embedding
# time_embed and audio_time_embed calculate both the timestep embedding and (global) modulation parameters
self.time_embed = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=6, use_additional_conditions=False)
self.audio_time_embed = LTX2AdaLayerNormSingle(
audio_inner_dim, num_mod_params=6, use_additional_conditions=False
# 3.1. Global Timestep Modulation Parameters
# LTX-2.3 with cross_attn_mod uses 9 mod params (extra 3 for cross-attn); LTX-2 uses 6.
video_num_mod_params = 9 if cross_attn_mod else 6
audio_num_mod_params = 9 if audio_cross_attn_mod else 6
self.time_embed = LTX2AdaLayerNormSingle(
inner_dim, num_mod_params=video_num_mod_params, use_additional_conditions=False
)
self.audio_time_embed = LTX2AdaLayerNormSingle(
audio_inner_dim, num_mod_params=audio_num_mod_params, use_additional_conditions=False
)
# 3.3. LTX-2.3: Prompt modulation from sigma
self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod
if self.prompt_modulation:
self.prompt_adaln = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=2, use_additional_conditions=False)
self.audio_prompt_adaln = LTX2AdaLayerNormSingle(
audio_inner_dim, num_mod_params=2, use_additional_conditions=False
)
# 3.2. Global Cross Attention Modulation Parameters
# Used in the audio-to-video and video-to-audio cross attention layers as a global set of modulation params,
@@ -1516,6 +1635,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
audio_num_attention_heads=audio_num_attention_heads,
audio_attention_head_dim=audio_attention_head_dim,
audio_cross_attention_dim=audio_cross_attention_dim,
video_gated_attn=gated_attn,
video_cross_attn_adaln=cross_attn_mod,
audio_gated_attn=audio_gated_attn,
audio_cross_attn_adaln=audio_cross_attn_mod,
qk_norm=qk_norm,
activation_fn=activation_fn,
attention_bias=attention_bias,
@@ -1523,6 +1646,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
eps=norm_eps,
elementwise_affine=norm_elementwise_affine,
rope_type=rope_type,
perturbed_attn=perturbed_attn,
)
for _ in range(num_layers)
]
@@ -1558,6 +1682,8 @@ class LTX2VideoTransformer3DModel(nn.Module):
audio_encoder_hidden_states: torch.Tensor,
timestep: torch.LongTensor,
audio_timestep: torch.LongTensor | None = None,
sigma: torch.Tensor | None = None,
audio_sigma: torch.Tensor | None = None,
encoder_attention_mask: torch.Tensor | None = None,
audio_encoder_attention_mask: torch.Tensor | None = None,
num_frames: int | None = None,
@@ -1569,6 +1695,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
audio_coords: torch.Tensor | None = None,
attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
**kwargs, # Accept extra diffusers kwargs (isolate_modalities, perturbation_mask, etc.)
) -> torch.Tensor:
"""
Forward pass for LTX-2.0 audiovisual video transformer.
@@ -1711,54 +1838,55 @@ class LTX2VideoTransformer3DModel(nn.Module):
)
audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1])
# 4. Prepare prompt embeddings
encoder_hidden_states = self.caption_projection(encoder_hidden_states)
encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
# 3.3. LTX-2.3: Compute prompt modulation from sigma
audio_sigma = audio_sigma if audio_sigma is not None else sigma
if self.prompt_modulation and sigma is not None:
temb_prompt, _ = self.prompt_adaln(sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype)
temb_prompt_audio, _ = self.audio_prompt_adaln(
audio_sigma.flatten(), batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype
)
temb_prompt = temb_prompt.view(batch_size, -1, temb_prompt.size(-1))
temb_prompt_audio = temb_prompt_audio.view(batch_size, -1, temb_prompt_audio.size(-1))
else:
temb_prompt = temb_prompt_audio = None
audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states)
audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1))
# 4. Prepare prompt embeddings
# LTX-2: caption projection is in the transformer
# LTX-2.3: caption projection is in the connectors (already applied)
if hasattr(self, "caption_projection"):
encoder_hidden_states = self.caption_projection(encoder_hidden_states)
encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
if hasattr(self, "audio_caption_projection"):
audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states)
audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1))
# 5. Run transformer blocks
for block in self.transformer_blocks:
block_kwargs = {
"hidden_states": hidden_states,
"audio_hidden_states": audio_hidden_states,
"encoder_hidden_states": encoder_hidden_states,
"audio_encoder_hidden_states": audio_encoder_hidden_states,
"temb": temb,
"temb_audio": temb_audio,
"temb_ca_scale_shift": video_cross_attn_scale_shift,
"temb_ca_audio_scale_shift": audio_cross_attn_scale_shift,
"temb_ca_gate": video_cross_attn_a2v_gate,
"temb_ca_audio_gate": audio_cross_attn_v2a_gate,
"temb_prompt": temb_prompt,
"temb_prompt_audio": temb_prompt_audio,
"video_rotary_emb": video_rotary_emb,
"audio_rotary_emb": audio_rotary_emb,
"ca_video_rotary_emb": video_cross_attn_rotary_emb,
"ca_audio_rotary_emb": audio_cross_attn_rotary_emb,
"encoder_attention_mask": encoder_attention_mask,
"audio_encoder_attention_mask": audio_encoder_attention_mask,
}
if torch.is_grad_enabled() and self.gradient_checkpointing:
hidden_states, audio_hidden_states = self._gradient_checkpointing_func(
block,
hidden_states,
audio_hidden_states,
encoder_hidden_states,
audio_encoder_hidden_states,
temb,
temb_audio,
video_cross_attn_scale_shift,
audio_cross_attn_scale_shift,
video_cross_attn_a2v_gate,
audio_cross_attn_v2a_gate,
video_rotary_emb,
audio_rotary_emb,
video_cross_attn_rotary_emb,
audio_cross_attn_rotary_emb,
encoder_attention_mask,
audio_encoder_attention_mask,
)
hidden_states, audio_hidden_states = self._gradient_checkpointing_func(block, **block_kwargs)
else:
hidden_states, audio_hidden_states = block(
hidden_states,
audio_hidden_states,
encoder_hidden_states,
audio_encoder_hidden_states,
temb,
temb_audio,
video_cross_attn_scale_shift,
audio_cross_attn_scale_shift,
video_cross_attn_a2v_gate,
audio_cross_attn_v2a_gate,
video_rotary_emb,
audio_rotary_emb,
video_cross_attn_rotary_emb,
audio_cross_attn_rotary_emb,
encoder_attention_mask,
audio_encoder_attention_mask,
)
hidden_states, audio_hidden_states = block(**block_kwargs)
# 6. Output layers (including unpatchification)
scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
@@ -1826,6 +1954,14 @@ class LTX2VideoTransformer3DModel(nn.Module):
weight_loader(param, loaded_weight, shard_id)
break
else:
if name not in params_dict:
logger.warning(
"Skipping transformer weight %s -- not found in model "
"parameters. This may indicate an incomplete "
"implementation or checkpoint mismatch.",
name,
)
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", None)
if weight_loader is not None:
@@ -0,0 +1,998 @@
# 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 uses a different text encoding strategy (flatten ALL 49 hidden states
vs. LTX-2's _pack_text_embeds with per-layer normalization and pooling)
- LTX-2.3 connectors expect the padding_side API (not additive_mask)
- LTX-2.3 uses a BWE vocoder outputting 48kHz audio (not 16kHz)
- LTX-2.3 transformer requires the sigma parameter for prompt modulation
- CPU offloading is required for the 22B transformer (~44GB VRAM)
"""
from __future__ import annotations
import copy
import json
import os
from collections.abc import Iterable
from contextlib import nullcontext
from typing import Any
import numpy as np
import torch
from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, 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.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.request import OmniDiffusionRequest
from .pipeline_ltx2 import (
_get_prompt_field,
calculate_shift,
create_transformer_from_config,
load_transformer_config,
)
logger = init_logger(__name__)
# Try to import LTX2VocoderWithBWE (diffusers >= 0.38.0)
try:
from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
except ImportError:
LTX2VocoderWithBWE = None
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
class LTX23Pipeline(nn.Module, ProgressBarMixin):
"""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
- CPU offloading: text encoder, connectors, VAE, vocoder stay on CPU
"""
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,
),
]
# --- Tokenizer (lightweight, stays wherever) ---
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
# --- Text encoder: load on CPU, move to GPU only during encoding ---
with torch.device("cpu"):
self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
)
# --- Connectors: CPU (LTX-2.3 connectors include caption projection) ---
self.connectors = LTX2TextConnectors.from_pretrained(
model, subfolder="connectors", torch_dtype=dtype, local_files_only=local_files_only
)
# --- VAE, Audio VAE: CPU ---
self.vae = AutoencoderKLLTX2Video.from_pretrained(
model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
)
self.audio_vae = AutoencoderKLLTX2Audio.from_pretrained(
model, subfolder="audio_vae", torch_dtype=dtype, local_files_only=local_files_only
)
# --- 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)
self.transformer = create_transformer_from_config(transformer_config)
# --- 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
# ------------------------------------------------------------------
# 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.
LTX-2.3 differs from LTX-2 in text encoding:
- LTX-2: uses _pack_text_embeds (layer selection + pooling)
- LTX-2.3: stacks ALL 49 hidden states and flattens to [B, seq, 188160]
The connectors unflatten, apply per_token_rms_norm, and project internally.
"""
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)
# Move text encoder to GPU for encoding
self.text_encoder.to(device)
text_encoder_outputs = self.text_encoder(
input_ids=text_input_ids,
attention_mask=prompt_attention_mask,
output_hidden_states=True,
)
# Move text encoder back to CPU immediately
self.text_encoder.to("cpu")
torch.cuda.empty_cache()
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)
_, 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,
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,
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 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,
)
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]
# ------------------------------------------------------------------
# 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
padded_latent_length = original_latent_length
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
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
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)
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()
# ------------------------------------------------------------------
# Forward pass
# ------------------------------------------------------------------
@torch.no_grad()
def forward(
self,
req: OmniDiffusionRequest,
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,
) -> DiffusionOutput:
# ---- Extract from request ----
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
# Enforce minimum of 2 timesteps for flow matching scheduler
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 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)
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)
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)
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)
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._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
# ---- Encode prompts ----
(
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,
)
# ---- Connectors (LTX-2.3: padding_side API) ----
# Concatenate negative + positive embeddings BEFORE connector call,
# matching diffusers which calls connectors once with batch=2.
# This ensures batch-dependent operations produce identical results.
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)
self.connectors.to(device)
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
)
self.connectors.to("cpu")
if torch.cuda.is_available():
torch.cuda.empty_cache()
# ---- Prepare latents ----
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 and latents.ndim == 5:
_, _, latent_num_frames, latent_height, latent_width = latents.shape
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 and audio_latents.ndim == 4:
_, _, audio_num_frames, _ = audio_latents.shape
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(
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,
)
# ---- Scheduler setup ----
sigmas = np.linspace(1.0, 1 / num_inference_steps, 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)
_ = 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)
# ---- RoPE coordinates ----
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,
)
# ---- CFG: duplicate coords for batch=2 ----
# Connector outputs are already batch=2 (neg+pos concatenated before connector call)
if self.do_classifier_free_guidance:
video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1))
audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
# ---- Denoising loop ----
# Uses x0-space CFG (delta formulation) matching diffusers' LTX2Pipeline.
# The velocity predictions are converted to x0, guidance is applied in x0
# space, then converted back to velocity for the scheduler step.
with self.progress_bar(total=len(timesteps)) as pbar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t
# Duplicate latents for CFG (uncond + cond)
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = latent_model_input.to(connector_prompt_embeds.dtype)
audio_latent_model_input = (
torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
)
audio_latent_model_input = audio_latent_model_input.to(connector_prompt_embeds.dtype)
ts = t.expand(latent_model_input.shape[0])
with self._transformer_cache_context("cond_uncond"):
noise_pred_video, noise_pred_audio = self.transformer(
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=ts,
sigma=ts, # LTX-2.3: sigma for prompt_adaln
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,
)
noise_pred_video = noise_pred_video.float()
noise_pred_audio = noise_pred_audio.float()
# CFG in x0-space (delta formulation matching diffusers)
if self.do_classifier_free_guidance:
noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2)
# Convert velocity to x0: x0 = sample - velocity * sigma
x0_video_cond = latents - noise_pred_video_cond * self.scheduler.sigmas[i]
x0_video_uncond = latents - noise_pred_video_uncond * self.scheduler.sigmas[i]
video_cfg_delta = (guidance_scale - 1) * (x0_video_cond - x0_video_uncond)
x0_video_guided = x0_video_cond + video_cfg_delta
noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2)
x0_audio_cond = audio_latents - noise_pred_audio_cond * audio_scheduler.sigmas[i]
x0_audio_uncond = audio_latents - noise_pred_audio_uncond * audio_scheduler.sigmas[i]
audio_cfg_delta = (guidance_scale - 1) * (x0_audio_cond - x0_audio_uncond)
x0_audio_guided = x0_audio_cond + audio_cfg_delta
# Convert x0 back to velocity: v = (sample - x0) / sigma
noise_pred_video = (latents - x0_video_guided) / self.scheduler.sigmas[i]
noise_pred_audio = (audio_latents - x0_audio_guided) / audio_scheduler.sigmas[i]
latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
pbar.update()
# ---- Unpack and denormalize ----
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,
)
# ---- Decode ----
if output_type == "latent":
video = latents
audio = audio_latents
else:
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)
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_decode = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
decode_noise_scale_t = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
:, None, None, None, None
]
latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise
# Move VAE, audio_vae, vocoder to GPU for decoding
self.vae.to(device)
latents = latents.to(self.vae.dtype)
video = self.vae.decode(latents, timestep_decode, return_dict=False)[0]
video = self.video_processor.postprocess_video(video, output_type=output_type)
self.vae.to("cpu")
self.audio_vae.to(device)
audio_latents = audio_latents.to(self.audio_vae.dtype)
generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0]
self.audio_vae.to("cpu")
self.vocoder.to(device)
audio = self.vocoder(generated_mel_spectrograms)
self.vocoder.to("cpu")
torch.cuda.empty_cache()
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(nn.Module):
"""LTX-2.3 image-to-video pipeline placeholder."""
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__()
raise NotImplementedError(
"LTX23ImageToVideoPipeline is not yet implemented. "
"Use LTX23Pipeline for single-stage text-to-video generation."
)
@@ -0,0 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Re-exports for LTX-2.3 I2V pipeline variants.
The registry loads pipeline classes by (mod_folder, mod_relname, cls_name).
This module exposes the I2V class names so the registry can find them.
"""
from .pipeline_ltx2_3 import (
LTX23ImageToVideoPipeline,
get_ltx2_post_process_func, # noqa: F401 - loaded by registry via getattr
)
__all__ = [
"LTX23ImageToVideoPipeline",
"get_ltx2_post_process_func",
]
+12
View File
@@ -93,6 +93,16 @@ _DIFFUSION_MODELS = {
"pipeline_ltx2_image2video",
"LTX2I2VDMD2Pipeline",
),
"LTX23Pipeline": (
"ltx2",
"pipeline_ltx2_3",
"LTX23Pipeline",
),
"LTX23ImageToVideoPipeline": (
"ltx2",
"pipeline_ltx2_3_image2video",
"LTX23ImageToVideoPipeline",
),
"StableAudioPipeline": (
"stable_audio",
"pipeline_stable_audio",
@@ -384,6 +394,8 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
"LTX2ImageToVideoTwoStagesPipeline": "get_ltx2_post_process_func",
"LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func",
"LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func",
"LTX23Pipeline": "get_ltx2_post_process_func",
"LTX23ImageToVideoPipeline": "get_ltx2_post_process_func",
"StableAudioPipeline": "get_stable_audio_post_process_func",
"WanImageToVideoPipeline": "get_wan22_i2v_post_process_func",
"WanT2VDMD2Pipeline": "get_wan22_post_process_func",