From 1cd52104404ab87db3057d0a5bf96646da53db9f Mon Sep 17 00:00:00 2001 From: Yangshen Deng Date: Wed, 8 Apr 2026 04:44:04 +0100 Subject: [PATCH] [Diffusion] Refactor LTX2 to use unified CFG parallel framework (#2160) Signed-off-by: Yangshen Deng Co-authored-by: Claude Opus 4.6 (1M context) --- tests/dfx/perf/tests/test_ltx2_vllm_omni.json | 217 +++++++++ .../ltx2/test_ltx2_cfg_parallel_adaptation.py | 58 +++ .../test_ltx2_cfg_parallel_parity.py | 243 ++++++++++ .../diffusion/models/ltx2/pipeline_ltx2.py | 416 +++++++----------- .../models/ltx2/pipeline_ltx2_image2video.py | 254 +++++------ 5 files changed, 764 insertions(+), 424 deletions(-) create mode 100644 tests/dfx/perf/tests/test_ltx2_vllm_omni.json create mode 100644 tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py create mode 100644 tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py diff --git a/tests/dfx/perf/tests/test_ltx2_vllm_omni.json b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json new file mode 100644 index 000000000..4a6f9e350 --- /dev/null +++ b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json @@ -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 + } + ] + } +] diff --git a/tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py b/tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py new file mode 100644 index 000000000..bbfe63dfa --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py @@ -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) diff --git a/tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py b/tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py new file mode 100644 index 000000000..659040929 --- /dev/null +++ b/tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py @@ -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") diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index efc342e93..c60b192f0 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -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 diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py index 11091518b..65e7454b7 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py @@ -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