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

[Diffusion] Refactor LTX2 to use unified CFG parallel framework (#2160)

Browse Source 
Signed-off-by: Yangshen Deng <yangshen.d@outlook.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangshen DengandClaude Opus 4.6 authored and GitHub committed 2026-04-08 11:44:04 +08:00
1 parent ec082add35
commit 1cd5210440
5 files changed
+764 -424
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/dfx/perf/tests/test_ltx2_vllm_omni.json
+217
View File
@@ -0,0 +1,217 @@
[
{
"test_name": "test_ltx2_baseline_eager",
"description": "Single-device baseline with enforce-eager (no torch.compile)",
"server_type": "vllm-omni",
"server_params": {
"model": "Lightricks/LTX-2",
"serve_args": {
"enforce-eager": true,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "256x256_145f_steps6",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 256,
"height": 256,
"num-frames": 145,
"fps": 24,
"num-inference-steps": 6,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
},
{
"name": "480x768_41f_steps20",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 768,
"height": 480,
"num-frames": 41,
"fps": 24,
"num-inference-steps": 20,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
}
]
},
{
"test_name": "test_ltx2_torch_compile",
"description": "Single-device with torch.compile (default, no enforce-eager)",
"server_type": "vllm-omni",
"server_params": {
"model": "Lightricks/LTX-2",
"serve_args": {
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "256x256_145f_steps6",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 256,
"height": 256,
"num-frames": 145,
"fps": 24,
"num-inference-steps": 6,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
},
{
"name": "480x768_41f_steps20",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 768,
"height": 480,
"num-frames": 41,
"fps": 24,
"num-inference-steps": 20,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
}
]
},
{
"test_name": "test_ltx2_cfg2_eager",
"description": "CFG-parallel=2 with enforce-eager",
"server_type": "vllm-omni",
"server_params": {
"model": "Lightricks/LTX-2",
"serve_args": {
"cfg-parallel-size": 2,
"enforce-eager": true,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "256x256_145f_steps6",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 256,
"height": 256,
"num-frames": 145,
"fps": 24,
"num-inference-steps": 6,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
},
{
"name": "480x768_41f_steps20",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 768,
"height": 480,
"num-frames": 41,
"fps": 24,
"num-inference-steps": 20,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
}
]
},
{
"test_name": "test_ltx2_cfg2_compile",
"description": "CFG-parallel=2 with torch.compile",
"server_type": "vllm-omni",
"server_params": {
"model": "Lightricks/LTX-2",
"serve_args": {
"cfg-parallel-size": 2,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "256x256_145f_steps6",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 256,
"height": 256,
"num-frames": 145,
"fps": 24,
"num-inference-steps": 6,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
},
{
"name": "480x768_41f_steps20",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 768,
"height": 480,
"num-frames": 41,
"fps": 24,
"num-inference-steps": 20,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
}
]
},
{
"test_name": "test_ltx2_cache_dit_eager",
"description": "CacheDiT with enforce-eager",
"server_type": "vllm-omni",
"server_params": {
"model": "Lightricks/LTX-2",
"serve_args": {
"cache-backend": "cache_dit",
"enforce-eager": true,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "256x256_145f_steps6",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 256,
"height": 256,
"num-frames": 145,
"fps": 24,
"num-inference-steps": 6,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
},
{
"name": "480x768_41f_steps20",
"dataset": "random",
"task": "t2v",
"backend": "v1/videos",
"width": 768,
"height": 480,
"num-frames": 41,
"fps": 24,
"num-inference-steps": 20,
"num-prompts": 3,
"max-concurrency": 1,
"enable-negative-prompt": true
}
]
}
]
tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py
+58
View File
@@ -0,0 +1,58 @@
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_pipeline(sequence_parallel_size: int = 1) -> LTX2Pipeline:
pipeline = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.audio_vae_temporal_compression_ratio = 4
pipeline.audio_vae_mel_compression_ratio = 4
pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
# Mock audio_vae with identity normalization (mean=0, std=1) so
# _normalize_audio_latents is a no-op and test values are preserved.
pipeline.audio_vae = SimpleNamespace(
latents_mean=torch.tensor(0.0),
latents_std=torch.tensor(1.0),
)
return pipeline
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
pipeline = _make_pipeline(sequence_parallel_size=4)
latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=10,
dtype=torch.float32,
device=torch.device("cpu"),
latents=latents,
)
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 4)
torch.testing.assert_close(padded[:, :10], latents)
torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
def test_unpad_audio_latents_restores_original_frames_before_unpack():
pipeline = _make_pipeline()
original = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
unpadded = pipeline._unpad_audio_latents(padded, 10)
unpacked = pipeline._unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
expected = pipeline._unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
assert unpacked.shape == (1, 2, 10, 2)
assert not (unpacked == 999.0).any()
torch.testing.assert_close(unpacked, expected)
tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py
+243
View File
@@ -0,0 +1,243 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import hashlib
import os
import subprocess
import sys
from pathlib import Path
import numpy as np
import pytest
from PIL import Image
from tests.utils import hardware_test
REPO_ROOT = Path(__file__).resolve().parents[3]
T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"
I2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py"
T2V_PROMPT = (
"At sunrise, a glowing paper lantern boat drifts through a narrow canal between mossy stone walls, "
"soft fog above the water, the camera slowly gliding forward as golden reflections shimmer across "
"the ripples, cinematic, realistic, highly detailed."
)
T2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion, distorted, oversaturated, artifacts"
I2V_PROMPT = "A cinematic dolly shot of a boat drifting on calm water at sunset"
I2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def _get_ltx2_model() -> str:
return os.environ.get("VLLM_TEST_LTX2_MODEL", "Lightricks/LTX-2")
def _md5(path: Path) -> str:
digest = hashlib.md5(usedforsecurity=False)
with path.open("rb") as f:
for chunk in iter(lambda: f.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _make_deterministic_test_image(path: Path) -> None:
"""Create a deterministic 256x256 test image for I2V tests."""
rng = np.random.RandomState(42)
img = Image.fromarray(rng.randint(0, 255, (256, 256, 3), dtype=np.uint8))
img.save(path)
def _run_and_check(cmd: list[str], env: dict, output_path: Path, expected_md5: str) -> None:
result = subprocess.run(cmd, cwd=REPO_ROOT, env=env, capture_output=True, text=True, check=False)
assert result.returncode == 0, (
f"Command failed (exit {result.returncode}).\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
)
generated_md5 = _md5(output_path)
assert generated_md5 == expected_md5, (
f"Unexpected output md5: {generated_md5} != {expected_md5}.\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
)
# ── T2V tests ──
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.slow
@hardware_test(res={"cuda": "L4"}, num_cards=2)
def test_ltx2_t2v_cfg_parallel(tmp_path: Path):
"""T2V with CFG=4.0, cfg-parallel-size=2."""
output = tmp_path / "t2v_cfg4.mp4"
env = os.environ.copy()
env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
cmd = [
sys.executable,
str(T2V_EXAMPLE),
"--model",
_get_ltx2_model(),
"--prompt",
T2V_PROMPT,
"--negative-prompt",
T2V_NEGATIVE_PROMPT,
"--height",
"256",
"--width",
"256",
"--num-frames",
"145",
"--num-inference-steps",
"6",
"--guidance-scale",
"4.0",
"--frame-rate",
"24",
"--fps",
"24",
"--seed",
"42",
"--cfg-parallel-size",
"2",
"--enforce-eager",
"--output",
str(output),
]
_run_and_check(cmd, env, output, expected_md5="08e606b9c522fee4b6f30cee8b77db40")
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.slow
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ltx2_t2v_no_cfg(tmp_path: Path):
"""T2V with CFG=1.0 (no classifier-free guidance)."""
output = tmp_path / "t2v_nocfg.mp4"
env = os.environ.copy()
env.setdefault("CUDA_VISIBLE_DEVICES", "0")
cmd = [
sys.executable,
str(T2V_EXAMPLE),
"--model",
_get_ltx2_model(),
"--prompt",
T2V_PROMPT,
"--height",
"256",
"--width",
"256",
"--num-frames",
"145",
"--num-inference-steps",
"6",
"--guidance-scale",
"1.0",
"--frame-rate",
"24",
"--fps",
"24",
"--seed",
"42",
"--enforce-eager",
"--output",
str(output),
]
_run_and_check(cmd, env, output, expected_md5="a83994b94b6e67c54a524e0383c45ce8")
# ── I2V tests ──
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.slow
@hardware_test(res={"cuda": "L4"}, num_cards=2)
def test_ltx2_i2v_cfg_parallel(tmp_path: Path):
"""I2V with CFG=4.0, cfg-parallel-size=2."""
test_image = tmp_path / "test_input.png"
_make_deterministic_test_image(test_image)
output = tmp_path / "i2v_cfg4.mp4"
env = os.environ.copy()
env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
cmd = [
sys.executable,
str(I2V_EXAMPLE),
"--model",
_get_ltx2_model(),
"--model-class-name",
"LTX2ImageToVideoPipeline",
"--image",
str(test_image),
"--prompt",
I2V_PROMPT,
"--negative-prompt",
I2V_NEGATIVE_PROMPT,
"--height",
"256",
"--width",
"256",
"--num-frames",
"73",
"--num-inference-steps",
"6",
"--guidance-scale",
"4.0",
"--frame-rate",
"24",
"--fps",
"24",
"--seed",
"42",
"--cfg-parallel-size",
"2",
"--enforce-eager",
"--output",
str(output),
]
_run_and_check(cmd, env, output, expected_md5="aed7e56084b36373244d8f839b16d115")
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.slow
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ltx2_i2v_no_cfg(tmp_path: Path):
"""I2V with CFG=1.0 (no classifier-free guidance)."""
test_image = tmp_path / "test_input.png"
_make_deterministic_test_image(test_image)
output = tmp_path / "i2v_nocfg.mp4"
env = os.environ.copy()
env.setdefault("CUDA_VISIBLE_DEVICES", "0")
cmd = [
sys.executable,
str(I2V_EXAMPLE),
"--model",
_get_ltx2_model(),
"--model-class-name",
"LTX2ImageToVideoPipeline",
"--image",
str(test_image),
"--prompt",
I2V_PROMPT,
"--height",
"256",
"--width",
"256",
"--num-frames",
"73",
"--num-inference-steps",
"6",
"--guidance-scale",
"1.0",
"--frame-rate",
"24",
"--fps",
"24",
"--seed",
"42",
"--enforce-eager",
"--output",
str(output),
]
_run_and_check(cmd, env, output, expected_md5="81b21ede12753e9e14a357a6c548b666")
vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py
+148 -268
View File
@@ -28,8 +28,6 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
get_cfg_group,
get_classifier_free_guidance_rank,
get_classifier_free_guidance_world_size,
)
from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -122,6 +120,31 @@ def calculate_shift(
return mu
class _VideoAudioScheduler:
"""Composite scheduler dispatching to video and audio schedulers."""
def __init__(self, video_scheduler, audio_scheduler):
self.video_scheduler = video_scheduler
self.audio_scheduler = audio_scheduler
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
video_out = self.video_scheduler.step(
noise_pred[0],
t[0],
latents[0],
return_dict=False,
generator=generator,
)[0]
audio_out = self.audio_scheduler.step(
noise_pred[1],
t[1],
latents[1],
return_dict=False,
generator=generator,
)[0]
return ((video_out, audio_out),)
class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
def __init__(
self,
@@ -542,6 +565,10 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2)
return latents
@staticmethod
def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
return latents[:, :num_frames]
def prepare_latents(
self,
batch_size: int = 1,
@@ -597,25 +624,49 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
noise_scale: float = 0.0,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
) -> tuple[torch.Tensor, int]:
) -> tuple[torch.Tensor, int, int]:
original_latent_length = audio_latent_length
padded_latent_length = original_latent_length
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
if sp_size > 1:
padded_latent_length += (sp_size - (original_latent_length % sp_size)) % sp_size
if latents is not None:
if latents.ndim == 4:
# latents are of shape [B, C, L, M], need to be packed
latents = self._pack_audio_latents(latents)
if latents.ndim != 3:
raise ValueError(
f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]." # noqa
f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is "
"[batch_size, num_seq, num_features] or [batch_size, num_channels, audio_length, mel_bins]."
)
latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std)
latents = self._create_noised_state(latents, noise_scale, generator)
return latents.to(device=device, dtype=dtype)
# TODO: confirm whether this logic is correct
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
if latents.shape[1] not in {original_latent_length, padded_latent_length}:
raise ValueError(
"Provided `audio_latents` has incompatible audio frame count "
f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}."
)
shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins)
if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length:
padding = torch.zeros(
latents.shape[0],
padded_latent_length - original_latent_length,
latents.shape[2],
dtype=latents.dtype,
device=latents.device,
)
latents = torch.cat([latents, padding], dim=1)
return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length
shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -625,7 +676,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
latents = self._pack_audio_latents(latents)
return latents
return latents, original_latent_length, padded_latent_length
@property
def guidance_scale(self):
@@ -655,147 +706,44 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
def interrupt(self):
return self._interrupt
def _is_cfg_parallel_enabled(self, do_true_cfg: bool) -> bool:
return do_true_cfg and get_classifier_free_guidance_world_size() > 1
def _transformer_cache_context(self, context_name: str):
cache_context = getattr(self.transformer, "cache_context", None)
if callable(cache_context):
return cache_context(context_name)
return nullcontext()
def _predict_noise_av(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
def predict_noise(self, **kwargs):
with self._transformer_cache_context("cond_uncond"):
noise_pred_video, noise_pred_audio = self.transformer(**kwargs)
return noise_pred_video, noise_pred_audio
def predict_noise_av_maybe_with_cfg(
self,
do_true_cfg: bool,
true_cfg_scale: float,
positive_kwargs: dict[str, Any],
negative_kwargs: dict[str, Any] | None,
guidance_rescale: float = 0.0,
cfg_normalize: bool = False,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
if do_true_cfg:
cfg_parallel_ready = get_classifier_free_guidance_world_size() > 1
if cfg_parallel_ready:
cfg_group = get_cfg_group()
cfg_rank = get_classifier_free_guidance_rank()
if cfg_rank == 0:
noise_pred_video, noise_pred_audio = self._predict_noise_av(**positive_kwargs)
else:
noise_pred_video, noise_pred_audio = self._predict_noise_av(**negative_kwargs)
noise_pred_video = noise_pred_video.float()
noise_pred_audio = noise_pred_audio.float()
gathered_video = cfg_group.all_gather(noise_pred_video, separate_tensors=True)
gathered_audio = cfg_group.all_gather(noise_pred_audio, separate_tensors=True)
if cfg_rank == 0:
noise_pred_video_text = gathered_video[0]
noise_pred_video_uncond = gathered_video[1]
noise_pred_audio_text = gathered_audio[0]
noise_pred_audio_uncond = gathered_audio[1]
noise_pred_video = self.combine_cfg_noise(
noise_pred_video_text,
noise_pred_video_uncond,
true_cfg_scale,
cfg_normalize,
)
noise_pred_audio = self.combine_cfg_noise(
noise_pred_audio_text,
noise_pred_audio_uncond,
true_cfg_scale,
cfg_normalize,
)
if guidance_rescale > 0:
noise_pred_video = rescale_noise_cfg(
noise_pred_video,
noise_pred_video_text,
guidance_rescale=guidance_rescale,
)
noise_pred_audio = rescale_noise_cfg(
noise_pred_audio,
noise_pred_audio_text,
guidance_rescale=guidance_rescale,
)
return noise_pred_video, noise_pred_audio
return None, None
noise_pred_video_text, noise_pred_audio_text = self._predict_noise_av(**positive_kwargs)
noise_pred_video_uncond, noise_pred_audio_uncond = self._predict_noise_av(**negative_kwargs)
noise_pred_video_text = noise_pred_video_text.float()
noise_pred_audio_text = noise_pred_audio_text.float()
noise_pred_video_uncond = noise_pred_video_uncond.float()
noise_pred_audio_uncond = noise_pred_audio_uncond.float()
noise_pred_video = self.combine_cfg_noise(
noise_pred_video_text,
noise_pred_video_uncond,
true_cfg_scale,
cfg_normalize,
)
noise_pred_audio = self.combine_cfg_noise(
noise_pred_audio_text,
noise_pred_audio_uncond,
true_cfg_scale,
cfg_normalize,
)
if guidance_rescale > 0:
noise_pred_video = rescale_noise_cfg(
noise_pred_video,
noise_pred_video_text,
guidance_rescale=guidance_rescale,
)
noise_pred_audio = rescale_noise_cfg(
noise_pred_audio,
noise_pred_audio_text,
guidance_rescale=guidance_rescale,
)
return noise_pred_video, noise_pred_audio
noise_pred_video, noise_pred_audio = self._predict_noise_av(**positive_kwargs)
return noise_pred_video.float(), noise_pred_audio.float()
def _scheduler_step_video_audio_maybe_with_cfg(
def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, true_cfg_scale, cfg_normalize=False):
"""Per-element CFG combine with guidance_rescale support."""
(video_pos, audio_pos) = positive_noise_pred
(video_neg, audio_neg) = negative_noise_pred
video_combined = super().combine_cfg_noise(video_pos, video_neg, true_cfg_scale, cfg_normalize)
audio_combined = super().combine_cfg_noise(audio_pos, audio_neg, true_cfg_scale, cfg_normalize)
if self._guidance_rescale and self._guidance_rescale > 0:
video_combined = rescale_noise_cfg(video_combined, video_pos, guidance_rescale=self._guidance_rescale)
audio_combined = rescale_noise_cfg(audio_combined, audio_pos, guidance_rescale=self._guidance_rescale)
return (video_combined, audio_combined)
def _synchronize_cfg_parallel_step_output(
self,
noise_pred_video: torch.Tensor | None,
noise_pred_audio: torch.Tensor | None,
t: torch.Tensor,
latents: torch.Tensor,
audio_latents: torch.Tensor,
audio_scheduler: FlowMatchEulerDiscreteScheduler,
latents: tuple[torch.Tensor, torch.Tensor],
do_true_cfg: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
cfg_parallel_ready = self._is_cfg_parallel_enabled(do_true_cfg)
if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1):
return latents
if cfg_parallel_ready:
cfg_group = get_cfg_group()
cfg_rank = get_classifier_free_guidance_rank()
if cfg_rank == 0:
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]
latents = latents.contiguous()
audio_latents = audio_latents.contiguous()
cfg_group.broadcast(latents, src=0)
cfg_group.broadcast(audio_latents, src=0)
return latents, audio_latents
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]
return latents, audio_latents
# Without this sync, CUDA async execution causes non-deterministic
# numerical drift across denoising steps in CFG parallel mode,
# producing different video outputs across runs.
latents = tuple(tensor.contiguous() for tensor in latents)
device = next((tensor.device for tensor in latents if tensor.is_cuda), None)
if device is not None:
torch.cuda.current_stream(device).synchronize()
return latents
@torch.no_grad()
def forward(
@@ -828,6 +776,8 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
) -> DiffusionOutput:
# Extract prompt/negative_prompt from request.
# Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys.
prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
negative_prompt = None
@@ -869,6 +819,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
else req.sampling_params.extra_args.get("audio_latents", audio_latents)
)
# Override with pre-computed embeddings if provided in request.
req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts]
if any(p is not None for p in req_prompt_embeds):
prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore[arg-type]
@@ -939,20 +890,17 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
max_sequence_length=max_sequence_length,
device=device,
)
cfg_parallel_ready = self._is_cfg_parallel_enabled(self.do_classifier_free_guidance)
if self.do_classifier_free_guidance and not cfg_parallel_ready:
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)
# Compute positive prompt connectors
additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
prompt_embeds, additive_attention_mask, additive_mask=True
)
# Compute negative prompt connectors when CFG is enabled
negative_connector_prompt_embeds = None
negative_connector_audio_prompt_embeds = None
negative_connector_attention_mask = None
if cfg_parallel_ready:
if self.do_classifier_free_guidance:
negative_additive_attention_mask = (
1 - negative_prompt_attention_mask.to(negative_prompt_embeds.dtype)
) * -1000000.0
@@ -1027,20 +975,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
num_channels_latents_audio = (
self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
)
# padding audio_latents if needed
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
if sp_size > 1:
pad_len = (sp_size - (audio_num_frames % sp_size)) % sp_size
if pad_len > 0:
if audio_latents is not None:
pad_shape = list(audio_latents.shape)
pad_shape[2] = pad_len
padding = torch.zeros(pad_shape, dtype=audio_latents.dtype, device=audio_latents.device)
audio_latents = torch.cat([audio_latents, padding], dim=2)
audio_num_frames += pad_len
audio_latents = self.prepare_audio_latents(
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,
@@ -1061,6 +996,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
self.scheduler.config.get("max_shift", 2.05),
)
audio_scheduler = copy.deepcopy(self.scheduler)
video_audio_scheduler = _VideoAudioScheduler(self.scheduler, audio_scheduler)
_ = retrieve_timesteps(
audio_scheduler,
num_inference_steps,
@@ -1083,12 +1019,10 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
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], audio_num_frames, audio_latents.device
audio_latents.shape[0], padded_audio_num_frames, audio_latents.device
)
# Duplicate the positional ids as well if using CFG
if self.do_classifier_free_guidance and not cfg_parallel_ready:
video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim
audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
# No coord duplication needed: mixin handles CFG via separate forward calls,
# not batch=2. Each forward gets batch=1 coords directly.
with self.progress_bar(total=len(timesteps)) as pbar:
for i, t in enumerate(timesteps):
@@ -1097,119 +1031,60 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
self._current_timestep = t
if cfg_parallel_ready:
latent_model_input = latents.to(prompt_embeds.dtype)
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
latent_model_input = latents.to(prompt_embeds.dtype)
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
do_true_cfg = self.do_classifier_free_guidance
positive_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": connector_prompt_embeds,
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
"timestep": timestep,
"encoder_attention_mask": connector_attention_mask,
"audio_encoder_attention_mask": connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
negative_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
positive_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": connector_prompt_embeds,
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
"timestep": timestep,
"encoder_attention_mask": connector_attention_mask,
"audio_encoder_attention_mask": connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": padded_audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
negative_kwargs = (
{
**positive_kwargs,
"encoder_hidden_states": negative_connector_prompt_embeds,
"audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
"timestep": timestep,
"encoder_attention_mask": negative_connector_attention_mask,
"audio_encoder_attention_mask": negative_connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
if do_true_cfg
else None
)
noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
do_true_cfg=True,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
guidance_rescale=guidance_rescale,
cfg_normalize=False,
)
noise_pred_video, noise_pred_audio = self.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
latents, audio_latents = self._scheduler_step_video_audio_maybe_with_cfg(
noise_pred_video,
noise_pred_audio,
t,
latents,
audio_latents,
audio_scheduler,
do_true_cfg=True,
)
else:
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = latent_model_input.to(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(prompt_embeds.dtype)
timestep = 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=timestep,
encoder_attention_mask=connector_attention_mask,
audio_encoder_attention_mask=connector_attention_mask,
num_frames=latent_num_frames,
height=latent_height,
width=latent_width,
fps=frame_rate,
audio_num_frames=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()
if self.do_classifier_free_guidance:
noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
noise_pred_video = noise_pred_video_uncond + guidance_scale * (
noise_pred_video_text - noise_pred_video_uncond
)
noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
noise_pred_audio_text - noise_pred_audio_uncond
)
if guidance_rescale > 0:
noise_pred_video = rescale_noise_cfg(
noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
)
noise_pred_audio = rescale_noise_cfg(
noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
)
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]
latents, audio_latents = self.scheduler_step_maybe_with_cfg(
(noise_pred_video, noise_pred_audio),
(t, t),
(latents, audio_latents),
do_true_cfg=do_true_cfg,
per_request_scheduler=video_audio_scheduler,
)
latents, audio_latents = self._synchronize_cfg_parallel_step_output(
(latents, audio_latents),
do_true_cfg=do_true_cfg,
)
pbar.update()
@@ -1225,10 +1100,15 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
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, audio_num_frames, num_mel_bins=latent_mel_bins)
audio_latents = self._unpack_audio_latents(
audio_latents,
original_audio_num_frames,
num_mel_bins=latent_mel_bins,
)
if output_type == "latent":
video = latents
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py
+98 -156
View File
@@ -14,7 +14,7 @@ import torch
import torch.nn as nn
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg, retrieve_timesteps
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
@@ -22,7 +22,6 @@ 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.parallel_state import get_cfg_group, get_classifier_free_guidance_rank
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -46,6 +45,32 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
return _get_ltx2_post_process_func(od_config)
class _I2VVideoAudioScheduler:
"""Composite scheduler for I2V: uses _step_video_latents_i2v for video, standard step for audio."""
def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width):
self.video_scheduler = pipeline.scheduler
self.audio_scheduler = audio_scheduler
self._pipeline = pipeline
self._latent_num_frames = latent_num_frames
self._latent_height = latent_height
self._latent_width = latent_width
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
video_out = self._pipeline._step_video_latents_i2v(
noise_pred[0],
latents[0],
t[0],
self._latent_num_frames,
self._latent_height,
self._latent_width,
)
audio_out = self.audio_scheduler.step(noise_pred[1], t[1], latents[1], return_dict=False, generator=generator)[
0
]
return ((video_out, audio_out),)
class LTX2ImageToVideoPipeline(LTX2Pipeline):
support_image_input = True
@@ -287,6 +312,8 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
) -> DiffusionOutput:
# Extract prompt/negative_prompt from request.
# Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys.
prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
negative_prompt = None
@@ -328,6 +355,7 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
else req.sampling_params.extra_args.get("audio_latents", audio_latents)
)
# Override with pre-computed embeddings if provided in request.
req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts]
if any(p is not None for p in req_prompt_embeds):
prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore[arg-type]
@@ -429,20 +457,17 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
max_sequence_length=max_sequence_length,
device=device,
)
cfg_parallel_ready = self._is_cfg_parallel_enabled(self.do_classifier_free_guidance)
if self.do_classifier_free_guidance and not cfg_parallel_ready:
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)
# Compute positive prompt connectors
additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
prompt_embeds, additive_attention_mask, additive_mask=True
)
# Compute negative prompt connectors when CFG is enabled
negative_connector_prompt_embeds = None
negative_connector_audio_prompt_embeds = None
negative_connector_attention_mask = None
if cfg_parallel_ready:
if self.do_classifier_free_guidance:
negative_additive_attention_mask = (
1 - negative_prompt_attention_mask.to(negative_prompt_embeds.dtype)
) * -1000000.0
@@ -500,8 +525,6 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
generator,
latents,
)
if self.do_classifier_free_guidance and not cfg_parallel_ready:
conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])
duration_s = num_frames / frame_rate
audio_latents_per_second = (
@@ -529,20 +552,7 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
num_channels_latents_audio = (
self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
)
# padding audio_latents if needed
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
if sp_size > 1:
pad_len = (sp_size - (audio_num_frames % sp_size)) % sp_size
if pad_len > 0:
if audio_latents is not None:
pad_shape = list(audio_latents.shape)
pad_shape[2] = pad_len
padding = torch.zeros(pad_shape, dtype=audio_latents.dtype, device=audio_latents.device)
audio_latents = torch.cat([audio_latents, padding], dim=2)
audio_num_frames += pad_len
audio_latents = self.prepare_audio_latents(
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,
@@ -585,12 +595,17 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
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], audio_num_frames, audio_latents.device
audio_latents.shape[0], padded_audio_num_frames, audio_latents.device
)
# Duplicate the positional ids as well if using CFG
if self.do_classifier_free_guidance and not cfg_parallel_ready:
video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim
audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
i2v_scheduler = _I2VVideoAudioScheduler(
pipeline=self,
audio_scheduler=audio_scheduler,
latent_num_frames=latent_num_frames,
latent_height=latent_height,
latent_width=latent_width,
)
# No coord duplication needed: mixin handles CFG via separate forward calls.
with self.progress_bar(total=len(timesteps)) as pbar:
for i, t in enumerate(timesteps):
@@ -599,140 +614,62 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
self._current_timestep = t
if cfg_parallel_ready:
latent_model_input = latents.to(prompt_embeds.dtype)
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
latent_model_input = latents.to(prompt_embeds.dtype)
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
do_true_cfg = self.do_classifier_free_guidance
timestep = t.expand(latent_model_input.shape[0])
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
positive_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": connector_prompt_embeds,
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
"timestep": video_timestep,
"audio_timestep": timestep,
"encoder_attention_mask": connector_attention_mask,
"audio_encoder_attention_mask": connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
negative_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
positive_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": connector_prompt_embeds,
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
"timestep": video_timestep,
"audio_timestep": timestep,
"encoder_attention_mask": connector_attention_mask,
"audio_encoder_attention_mask": connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": padded_audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
negative_kwargs = (
{
**positive_kwargs,
"encoder_hidden_states": negative_connector_prompt_embeds,
"audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
"timestep": video_timestep,
"audio_timestep": timestep,
"encoder_attention_mask": negative_connector_attention_mask,
"audio_encoder_attention_mask": negative_connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
if do_true_cfg
else None
)
noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
do_true_cfg=True,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
guidance_rescale=guidance_rescale,
cfg_normalize=False,
)
noise_pred_video, noise_pred_audio = self.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
if get_classifier_free_guidance_rank() == 0:
latents = self._step_video_latents_i2v(
noise_pred_video,
latents,
t,
latent_num_frames,
latent_height,
latent_width,
)
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
cfg_group = get_cfg_group()
latents = latents.contiguous()
audio_latents = audio_latents.contiguous()
cfg_group.broadcast(latents, src=0)
cfg_group.broadcast(audio_latents, src=0)
else:
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = latent_model_input.to(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(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
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=video_timestep,
audio_timestep=timestep,
encoder_attention_mask=connector_attention_mask,
audio_encoder_attention_mask=connector_attention_mask,
num_frames=latent_num_frames,
height=latent_height,
width=latent_width,
fps=frame_rate,
audio_num_frames=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()
if self.do_classifier_free_guidance:
noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
noise_pred_video = noise_pred_video_uncond + guidance_scale * (
noise_pred_video_text - noise_pred_video_uncond
)
noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
noise_pred_audio_text - noise_pred_audio_uncond
)
if guidance_rescale > 0:
noise_pred_video = rescale_noise_cfg(
noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
)
noise_pred_audio = rescale_noise_cfg(
noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
)
latents = self._step_video_latents_i2v(
noise_pred_video,
latents,
t,
latent_num_frames,
latent_height,
latent_width,
)
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
latents, audio_latents = self.scheduler_step_maybe_with_cfg(
(noise_pred_video, noise_pred_audio),
(t, t),
(latents, audio_latents),
do_true_cfg=do_true_cfg,
per_request_scheduler=i2v_scheduler,
)
latents, audio_latents = self._synchronize_cfg_parallel_step_output(
(latents, audio_latents),
do_true_cfg=do_true_cfg,
)
pbar.update()
@@ -748,10 +685,15 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
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, audio_num_frames, num_mel_bins=latent_mel_bins)
audio_latents = self._unpack_audio_latents(
audio_latents,
original_audio_num_frames,
num_mel_bins=latent_mel_bins,
)
if output_type == "latent":
video = latents