From 2b999429c4d408a890be8504996be2f9b4f6204a Mon Sep 17 00:00:00 2001 From: mglyn <39453182+mglyn@users.noreply.github.com> Date: Thu, 28 May 2026 12:36:33 +0800 Subject: [PATCH] [Feature] Add cfg-parallel for LTX-2.3 (#3905) Signed-off-by: mglyn <1203789601@qq.com> --- docs/user_guide/diffusion_features.md | 2 +- recipes/LTX/LTX-2.3.md | 1 - .../models/ltx2/test_ltx2_3_pipeline.py | 306 ++++++++++++++++++ .../diffusion/models/ltx2/pipeline_ltx2_3.py | 304 ++++++++++++++--- 4 files changed, 560 insertions(+), 53 deletions(-) diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 52c24ab9e..9400c828e 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -145,7 +145,7 @@ The following tables show which models support each feature: | **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ | | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | -| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/recipes/LTX/LTX-2.3.md b/recipes/LTX/LTX-2.3.md index 36a9eac09..269ea6425 100644 --- a/recipes/LTX/LTX-2.3.md +++ b/recipes/LTX/LTX-2.3.md @@ -86,7 +86,6 @@ curl -X POST http://localhost:8000/v1/videos \ - Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1) - Known limitations: - No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder) - - No CFG-parallel support (single-GPU only) - Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git) ## Hardware Support diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py index afe92698d..5b254d50b 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py @@ -65,6 +65,312 @@ class TestPipelineIndependence: assert issubclass(LTX23Pipeline, ProgressBarMixin) + def test_ltx23_pipeline_has_cfg_parallel_mixin(self): + """LTX23Pipeline must use the shared CFG parallel implementation.""" + from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + + assert issubclass(LTX23Pipeline, CFGParallelMixin) + + +class TestCFGParallelHelpers: + """Test LTX-2.3 CFG helper math without loading model weights.""" + + def test_combine_cfg_noise_matches_x0_space_formula(self): + import torch + + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + + pipe = object.__new__(LTX23Pipeline) + video_sample = torch.tensor([[[1.0, -2.0]]]) + audio_sample = torch.tensor([[[0.5, 3.0]]]) + video_pos = torch.tensor([[[0.2, -0.3]]]) + video_neg = torch.tensor([[[-0.4, 0.1]]]) + audio_pos = torch.tensor([[[0.7, -0.2]]]) + audio_neg = torch.tensor([[[0.1, 0.4]]]) + video_sigma = torch.tensor(0.25) + audio_sigma = torch.tensor(0.5) + scale = 4.0 + + video_combined, audio_combined = pipe.combine_cfg_noise( + (video_pos, audio_pos), + (video_neg, audio_neg), + scale, + video_latents=video_sample, + audio_latents=audio_sample, + video_sigma=video_sigma, + audio_sigma=audio_sigma, + ) + + x0_video_cond = video_sample - video_pos * video_sigma + x0_video_uncond = video_sample - video_neg * video_sigma + x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond) + expected_video = (video_sample - x0_video_guided) / video_sigma + + x0_audio_cond = audio_sample - audio_pos * audio_sigma + x0_audio_uncond = audio_sample - audio_neg * audio_sigma + x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond) + expected_audio = (audio_sample - x0_audio_guided) / audio_sigma + assert torch.allclose(video_combined, expected_video) + assert torch.allclose(audio_combined, expected_audio) + + def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 + + pipe = object.__new__(ltx23.LTX23Pipeline) + video_sample = torch.tensor([[[1.0, -2.0]]]) + audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]]) + video_pos = torch.tensor([[[0.2, -0.3]]]) + video_neg = torch.tensor([[[-0.4, 0.1]]]) + audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]]) + audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]]) + video_sigma = torch.tensor(0.25) + audio_sigma = torch.tensor(0.5) + scale = 4.0 + + class FakeCfgGroup: + def all_gather(self, tensor, separate_tensors=True): + assert separate_tensors + if tensor.shape == video_pos.shape: + return [video_pos, video_neg] + return [audio_pos, audio_neg] + + monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2) + monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup()) + + expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg( + video_sample, + video_pos, + video_neg, + video_sigma, + scale, + ) + expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg( + audio_sample, + audio_pos, + audio_neg, + audio_sigma, + scale, + ) + + for rank, expected_branch in ((0, "positive"), (1, "negative")): + calls = [] + monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank) + + def fake_predict_noise(**kwargs): + calls.append(kwargs["branch"]) + if kwargs["branch"] == "positive": + return video_pos, audio_pos + return video_neg, audio_neg + + object.__setattr__(pipe, "predict_noise", fake_predict_noise) + video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg( + true_cfg_scale=scale, + positive_kwargs={"branch": "positive"}, + negative_kwargs={"branch": "negative"}, + cfg_normalize=False, + video_latents=video_sample, + audio_latents=audio_sample, + video_sigma=video_sigma, + audio_sigma=audio_sigma, + ) + + assert calls == [expected_branch] + torch.testing.assert_close(video_combined, expected_video) + torch.testing.assert_close(audio_combined, expected_audio) + + assert "_cfg_video_latents" not in pipe.__dict__ + assert "_cfg_audio_latents" not in pipe.__dict__ + + +class TestCFGParallelForwardPath: + """Test the LTX-2.3 CFG-parallel denoising path without loading model weights.""" + + @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)]) + def test_forward_cfg_parallel_steps_video_and_audio_scheduler( + self, + monkeypatch, + cfg_rank, + expected_prompt_value, + ): + from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = object.__new__(ltx23.LTX23Pipeline) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 1 + pipe.vae_spatial_compression_ratio = 32 + pipe.vae_temporal_compression_ratio = 1 + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + pipe.audio_sampling_rate = 1 + pipe.audio_hop_length = 1 + pipe.audio_vae_temporal_compression_ratio = 1 + pipe.audio_vae_mel_compression_ratio = 1 + pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)) + pipe.tokenizer = SimpleNamespace(padding_side="left") + pipe.vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(scaling_factor=1.0), + ) + pipe.audio_vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(mel_bins=2, latent_channels=1), + ) + + video_pos = torch.tensor([[[0.2, -0.3]]]) + video_neg = torch.tensor([[[-0.4, 0.1]]]) + audio_pos = torch.tensor([[[0.7, -0.2]]]) + audio_neg = torch.tensor([[[0.1, 0.4]]]) + + class FakeCfgGroup: + def all_gather(self, tensor, separate_tensors=True): + assert separate_tensors + if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg): + return [video_pos, video_neg] + if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg): + return [audio_pos, audio_neg] + raise AssertionError(f"Unexpected gathered tensor: {tensor}") + + monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2) + monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank) + monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup()) + + def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None): + scheduler.sigmas = torch.tensor([0.25, 0.25], device=device) + return torch.tensor([1.0, 0.5], device=device), 2 + + monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps) + + class FakeScheduler: + def __init__(self, name="video", calls=None): + self.name = name + self.calls = [] if calls is None else calls + self.config = { + "max_image_seq_len": 4096, + "base_image_seq_len": 1024, + "base_shift": 0.95, + "max_shift": 2.05, + } + self.sigmas = torch.tensor([0.25, 0.25]) + + def __deepcopy__(self, memo): + return FakeScheduler("audio", self.calls) + + def step(self, noise_pred, t, latents, return_dict=False, generator=None): + self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone())) + return (latents - noise_pred,) + + class FakeConnectors: + def to(self, device): + return self + + def __call__(self, prompt_embeds, prompt_attention_mask, padding_side): + assert padding_side == "left" + assert prompt_embeds.shape[0] == 2 + return prompt_embeds, prompt_embeds, prompt_attention_mask + + class FakeRope: + def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps): + return torch.zeros(batch_size, num_frames * height * width, 3, device=device) + + def prepare_audio_coords(self, batch_size, num_frames, device): + return torch.zeros(batch_size, num_frames, 1, device=device) + + class FakeTransformer: + def __init__(self): + self.config = SimpleNamespace(in_channels=2) + self.rope = FakeRope() + self.audio_rope = FakeRope() + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + expected_prompt = torch.full((1, 1, 1), expected_prompt_value) + torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt) + torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt) + assert kwargs["hidden_states"].shape == (1, 1, 2) + assert kwargs["audio_hidden_states"].shape == (1, 1, 2) + if cfg_rank == 0: + return video_pos, audio_pos + return video_neg, audio_neg + + class DummyProgress: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def update(self): + pass + + pipe.scheduler = FakeScheduler() + pipe.connectors = FakeConnectors() + pipe.transformer = FakeTransformer() + object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress()) + + def fake_encode_prompt(**kwargs): + return ( + torch.ones(1, 1, 1), + torch.ones(1, 1, dtype=torch.bool), + torch.zeros(1, 1, 1), + torch.ones(1, 1, dtype=torch.bool), + ) + + object.__setattr__(pipe, "encode_prompt", fake_encode_prompt) + + video_latents = torch.tensor([[[1.0, -2.0]]]) + audio_latents = torch.tensor([[[0.5, 3.0]]]) + req = OmniDiffusionRequest( + prompts=[{"prompt": "prompt", "negative_prompt": "negative"}], + sampling_params=OmniDiffusionSamplingParams( + height=32, + width=32, + num_frames=1, + frame_rate=1.0, + num_inference_steps=2, + guidance_scale=4.0, + latents=video_latents, + audio_latents=audio_latents, + output_type="latent", + ), + ) + + output = pipe.forward(req) + + expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg( + video_latents, + video_pos, + video_neg, + pipe.scheduler.sigmas[0], + 4.0, + ) + expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg( + audio_latents, + audio_pos, + audio_neg, + pipe.scheduler.sigmas[0], + 4.0, + ) + scheduler_call_names = [call[0] for call in pipe.scheduler.calls] + assert scheduler_call_names == ["video", "audio", "video", "audio"] + assert len(pipe.transformer.calls) == 2 + torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise) + torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise) + torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise) + torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise) + torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise) + torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise) + + video_out, audio_out = output.output + torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1)) + torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2)) + class TestRegistryIntegration: """Verify all LTX-2.3 pipeline variants are registered.""" diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index c16f74eef..3c18e70dc 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py @@ -37,6 +37,12 @@ from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, +) from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin @@ -44,6 +50,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from .pipeline_ltx2 import ( _get_prompt_field, + _VideoAudioScheduler, calculate_shift, create_transformer_from_config, load_transformer_config, @@ -104,7 +111,7 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig): return post_process_func -class LTX23Pipeline(nn.Module, ProgressBarMixin): +class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): """Fully independent LTX-2.3 pipeline. Key differences from LTX2Pipeline: @@ -553,6 +560,12 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) latents = self._pack_audio_latents(latents) return latents, original_latent_length, padded_latent_length @@ -635,6 +648,121 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): return cache_context(context_name) return nullcontext() + # ------------------------------------------------------------------ + # CFG helpers + # ------------------------------------------------------------------ + + @staticmethod + def _combine_x0_space_cfg( + sample: torch.Tensor, + positive_noise_pred: torch.Tensor, + negative_noise_pred: torch.Tensor, + sigma: torch.Tensor, + guidance_scale: float, + ) -> torch.Tensor: + x0_cond = sample - positive_noise_pred * sigma + x0_uncond = sample - negative_noise_pred * sigma + x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond) + return (sample - x0_guided) / sigma + + def predict_noise(self, **kwargs): + with self._transformer_cache_context("cond_uncond"): + noise_pred_video, noise_pred_audio = self.transformer(**kwargs) + return noise_pred_video.float(), noise_pred_audio.float() + + def combine_cfg_noise( + self, + positive_noise_pred, + negative_noise_pred, + true_cfg_scale, + cfg_normalize=False, + *, + video_latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + video_sigma: torch.Tensor | None = None, + audio_sigma: torch.Tensor | None = None, + ): + if video_latents is None or audio_latents is None or video_sigma is None or audio_sigma is None: + raise ValueError("LTX23Pipeline applies CFG in x0-space and requires video/audio latents and sigmas.") + + video_pos, audio_pos = positive_noise_pred + video_neg, audio_neg = negative_noise_pred + video_combined = self._combine_x0_space_cfg( + video_latents, + video_pos, + video_neg, + video_sigma, + true_cfg_scale, + ) + audio_combined = self._combine_x0_space_cfg( + audio_latents, + audio_pos, + audio_neg, + audio_sigma, + true_cfg_scale, + ) + if cfg_normalize: + video_combined = self.cfg_normalize_function(video_pos, video_combined) + audio_combined = self.cfg_normalize_function(audio_pos, audio_combined) + return video_combined, audio_combined + + def predict_noise_with_parallel_cfg( + self, + true_cfg_scale: float, + positive_kwargs: dict[str, Any], + negative_kwargs: dict[str, Any], + cfg_normalize: bool = True, + output_slice: int | None = None, + *, + video_latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + video_sigma: torch.Tensor | None = None, + audio_sigma: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + def maybe_slice(pred: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + if output_slice is None: + return pred + return pred[0][:, :output_slice], pred[1][:, :output_slice] + + cfg_world_size = get_classifier_free_guidance_world_size() + if cfg_world_size != 2: + raise ValueError(f"LTX23Pipeline parallel CFG requires cfg_parallel_size 2, but got {cfg_world_size}.") + + cfg_group = get_cfg_group() + cfg_rank = get_classifier_free_guidance_rank() + branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs + local_video_pred, local_audio_pred = maybe_slice(self.predict_noise(**branch_kwargs)) + + gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True) + gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True) + positive_noise_pred = (gathered_video[0], gathered_audio[0]) + negative_noise_pred = (gathered_video[1], gathered_audio[1]) + + return self.combine_cfg_noise( + positive_noise_pred, + negative_noise_pred, + true_cfg_scale, + cfg_normalize, + video_latents=video_latents, + audio_latents=audio_latents, + video_sigma=video_sigma, + audio_sigma=audio_sigma, + ) + + def _synchronize_cfg_parallel_step_output( + self, + latents: tuple[torch.Tensor, torch.Tensor], + do_true_cfg: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1): + return latents + + latents = tuple(tensor.contiguous() for tensor in latents) + device = next((tensor.device for tensor in latents if tensor.is_cuda), None) + if device is not None: + torch.cuda.current_stream(device).synchronize() + return latents + # ------------------------------------------------------------------ # Forward pass # ------------------------------------------------------------------ @@ -753,6 +881,12 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): self._attention_kwargs = attention_kwargs self._interrupt = False self._current_timestep = None + cfg_world_size = get_classifier_free_guidance_world_size() + if self.do_classifier_free_guidance and cfg_world_size not in (1, 2): + raise ValueError( + f"LTX23Pipeline supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}." + ) + cfg_parallel_ready = self.do_classifier_free_guidance and cfg_world_size > 1 if prompt is not None and isinstance(prompt, str): batch_size = 1 @@ -799,6 +933,21 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): if torch.cuda.is_available(): torch.accelerator.empty_cache() + positive_connector_prompt_embeds = connector_prompt_embeds + positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds + positive_connector_attention_mask = connector_attention_mask + negative_connector_prompt_embeds = None + negative_connector_audio_prompt_embeds = None + negative_connector_attention_mask = None + if self.do_classifier_free_guidance: + split_batch = batch_size * num_videos_per_prompt + negative_connector_prompt_embeds = connector_prompt_embeds[:split_batch] + positive_connector_prompt_embeds = connector_prompt_embeds[split_batch:] + negative_connector_audio_prompt_embeds = connector_audio_prompt_embeds[:split_batch] + positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds[split_batch:] + negative_connector_attention_mask = connector_attention_mask[:split_batch] + positive_connector_attention_mask = connector_attention_mask[split_batch:] + # ---- Prepare latents ---- latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 latent_height = height // self.vae_spatial_compression_ratio @@ -854,6 +1003,7 @@ class LTX23Pipeline(nn.Module, 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, device, timesteps, sigmas=sigmas, mu=mu) timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, @@ -880,9 +1030,9 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): audio_latents.device, ) - # ---- CFG: duplicate coords for batch=2 ---- + # ---- CFG: duplicate coords for single-rank batch=2 CFG ---- # Connector outputs are already batch=2 (neg+pos concatenated before connector call) - if self.do_classifier_free_guidance: + if self.do_classifier_free_guidance and not cfg_parallel_ready: video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1)) @@ -897,60 +1047,112 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): self._current_timestep = t - # Duplicate latents for CFG (uncond + cond) - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = latent_model_input.to(connector_prompt_embeds.dtype) - audio_latent_model_input = ( - torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents - ) - audio_latent_model_input = audio_latent_model_input.to(connector_prompt_embeds.dtype) - ts = t.expand(latent_model_input.shape[0]) - - with self._transformer_cache_context("cond_uncond"): - noise_pred_video, noise_pred_audio = self.transformer( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=connector_prompt_embeds, - audio_encoder_hidden_states=connector_audio_prompt_embeds, - timestep=ts, - sigma=ts, # LTX-2.3: sigma for prompt_adaln - encoder_attention_mask=connector_attention_mask, - audio_encoder_attention_mask=connector_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=frame_rate, - audio_num_frames=padded_audio_num_frames, - video_coords=video_coords, - audio_coords=audio_coords, - attention_kwargs=attention_kwargs, - return_dict=False, + if cfg_parallel_ready: + latent_model_input = latents.to(positive_connector_prompt_embeds.dtype) + audio_latent_model_input = audio_latents.to(positive_connector_prompt_embeds.dtype) + ts = t.expand(latent_model_input.shape[0]) + positive_kwargs = { + "hidden_states": latent_model_input, + "audio_hidden_states": audio_latent_model_input, + "encoder_hidden_states": positive_connector_prompt_embeds, + "audio_encoder_hidden_states": positive_connector_audio_prompt_embeds, + "timestep": ts, + "sigma": ts, + "encoder_attention_mask": positive_connector_attention_mask, + "audio_encoder_attention_mask": positive_connector_attention_mask, + "num_frames": latent_num_frames, + "height": latent_height, + "width": latent_width, + "fps": frame_rate, + "audio_num_frames": padded_audio_num_frames, + "video_coords": video_coords, + "audio_coords": audio_coords, + "attention_kwargs": attention_kwargs, + "return_dict": False, + } + negative_kwargs = { + **positive_kwargs, + "encoder_hidden_states": negative_connector_prompt_embeds, + "audio_encoder_hidden_states": negative_connector_audio_prompt_embeds, + "encoder_attention_mask": negative_connector_attention_mask, + "audio_encoder_attention_mask": negative_connector_attention_mask, + } + noise_pred_video, noise_pred_audio = self.predict_noise_with_parallel_cfg( + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + video_latents=latents, + audio_latents=audio_latents, + video_sigma=self.scheduler.sigmas[i], + audio_sigma=audio_scheduler.sigmas[i], ) - noise_pred_video = noise_pred_video.float() - noise_pred_audio = noise_pred_audio.float() + latents, audio_latents = self.scheduler_step_maybe_with_cfg( + (noise_pred_video, noise_pred_audio), + (t, t), + (latents, audio_latents), + do_true_cfg=self.do_classifier_free_guidance, + per_request_scheduler=video_audio_scheduler, + ) + latents, audio_latents = self._synchronize_cfg_parallel_step_output( + (latents, audio_latents), + do_true_cfg=self.do_classifier_free_guidance, + ) + else: + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(connector_prompt_embeds.dtype) + audio_latent_model_input = ( + torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents + ) + audio_latent_model_input = audio_latent_model_input.to(connector_prompt_embeds.dtype) + ts = t.expand(latent_model_input.shape[0]) - # CFG in x0-space (delta formulation matching diffusers) - if self.do_classifier_free_guidance: - noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2) - # Convert velocity to x0: x0 = sample - velocity * sigma - x0_video_cond = latents - noise_pred_video_cond * self.scheduler.sigmas[i] - x0_video_uncond = latents - noise_pred_video_uncond * self.scheduler.sigmas[i] - video_cfg_delta = (guidance_scale - 1) * (x0_video_cond - x0_video_uncond) - x0_video_guided = x0_video_cond + video_cfg_delta + with self._transformer_cache_context("cond_uncond"): + noise_pred_video, noise_pred_audio = self.transformer( + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=connector_prompt_embeds, + audio_encoder_hidden_states=connector_audio_prompt_embeds, + timestep=ts, + sigma=ts, # LTX-2.3: sigma for prompt_adaln + encoder_attention_mask=connector_attention_mask, + audio_encoder_attention_mask=connector_attention_mask, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=frame_rate, + audio_num_frames=padded_audio_num_frames, + video_coords=video_coords, + audio_coords=audio_coords, + attention_kwargs=attention_kwargs, + return_dict=False, + ) - noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2) - x0_audio_cond = audio_latents - noise_pred_audio_cond * audio_scheduler.sigmas[i] - x0_audio_uncond = audio_latents - noise_pred_audio_uncond * audio_scheduler.sigmas[i] - audio_cfg_delta = (guidance_scale - 1) * (x0_audio_cond - x0_audio_uncond) - x0_audio_guided = x0_audio_cond + audio_cfg_delta + noise_pred_video = noise_pred_video.float() + noise_pred_audio = noise_pred_audio.float() - # Convert x0 back to velocity: v = (sample - x0) / sigma - noise_pred_video = (latents - x0_video_guided) / self.scheduler.sigmas[i] - noise_pred_audio = (audio_latents - x0_audio_guided) / audio_scheduler.sigmas[i] + if self.do_classifier_free_guidance: + noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2) + noise_pred_video = self._combine_x0_space_cfg( + latents, + noise_pred_video_cond, + noise_pred_video_uncond, + self.scheduler.sigmas[i], + guidance_scale, + ) - 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] + noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2) + noise_pred_audio = self._combine_x0_space_cfg( + audio_latents, + noise_pred_audio_cond, + noise_pred_audio_uncond, + audio_scheduler.sigmas[i], + guidance_scale, + ) + + latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0] + audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0] pbar.update()