From 385df0359e25d7197b4171e352a7ae100498e98f Mon Sep 17 00:00:00 2001 From: mglyn <39453182+mglyn@users.noreply.github.com> Date: Tue, 26 May 2026 16:45:38 +0800 Subject: [PATCH] [BugFix] Fix LTX-2.3 audio latent padding for sequence parallelism (#3854) Signed-off-by: mglyn <1203789601@qq.com> --- .../models/ltx2/test_ltx2_3_pipeline.py | 99 +++++++++++++++++++ .../diffusion/models/ltx2/pipeline_ltx2_3.py | 44 ++++++++- 2 files changed, 140 insertions(+), 3 deletions(-) diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py index 665126df7..afe92698d 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py @@ -15,12 +15,29 @@ These tests verify: import json import os import tempfile +from types import SimpleNamespace import pytest +import torch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def _make_ltx23_pipeline(sequence_parallel_size: int = 1): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + + pipeline = object.__new__(LTX23Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae_temporal_compression_ratio = 4 + pipeline.audio_vae_mel_compression_ratio = 4 + pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size)) + pipeline.audio_vae = SimpleNamespace( + latents_mean=torch.tensor(0.0), + latents_std=torch.tensor(1.0), + ) + return pipeline + + class TestPipelineIndependence: """Verify LTX23Pipeline is fully independent from LTX2Pipeline.""" @@ -228,3 +245,85 @@ class TestInitExports: for name in expected_classes: assert hasattr(ltx2, name), f"{name} not exported from ltx2 package" assert name in ltx2.__all__, f"{name} not in ltx2.__all__" + + +class TestAudioLatentSPPadding: + def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(self): + pipeline = _make_ltx23_pipeline(sequence_parallel_size=2) + + latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( + batch_size=1, + num_channels_latents=8, + num_mel_bins=64, + audio_latent_length=1, + dtype=torch.float32, + device=torch.device("cpu"), + ) + + assert original_num_frames == 1 + assert padded_num_frames == 2 + assert latents.shape == (1, 2, 128) + + def test_prepare_audio_latents_pads_provided_packed_sequence_dim_for_sp(self): + pipeline = _make_ltx23_pipeline(sequence_parallel_size=4) + latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4) + + padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=8, + audio_latent_length=10, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) + + assert original_num_frames == 10 + assert padded_num_frames == 12 + assert padded.shape == (1, 12, 4) + torch.testing.assert_close(padded[:, :10], latents) + torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4)) + + def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(self): + pipeline = _make_ltx23_pipeline(sequence_parallel_size=4) + latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4) + + audio_latent_length = pipeline._resolve_audio_latent_length(10, latents) + padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=16, + audio_latent_length=audio_latent_length, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) + + assert audio_latent_length == 10 + assert original_num_frames == 10 + assert padded_num_frames == 12 + assert padded.shape == (1, 12, 8) + torch.testing.assert_close(padded, pipeline._pack_audio_latents(latents)) + + def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(self): + pipeline = _make_ltx23_pipeline(sequence_parallel_size=4) + latents = torch.zeros(1, 2, 13, 4) + + audio_latent_length = pipeline._resolve_audio_latent_length(10, latents) + + assert audio_latent_length == 13 + + def test_prepare_audio_latents_rejects_incompatible_provided_length(self): + pipeline = _make_ltx23_pipeline(sequence_parallel_size=4) + latents = torch.zeros(1, 11, 4) + + with pytest.raises(ValueError, match="incompatible audio frame count"): + pipeline.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=8, + audio_latent_length=10, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index ec3b495f5..f3bc7c928 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py @@ -445,6 +445,26 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: return latents[:, :num_frames] + @staticmethod + def _get_sp_padded_audio_latent_length(audio_latent_length: int, sp_size: int) -> int: + if sp_size > 1: + audio_latent_length += (sp_size - (audio_latent_length % sp_size)) % sp_size + return audio_latent_length + + def _resolve_audio_latent_length(self, audio_latent_length: int, audio_latents: torch.Tensor | None) -> int: + if audio_latents is None or audio_latents.ndim != 4: + return audio_latent_length + + provided_latent_length = audio_latents.shape[2] + sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1 + padded_latent_length = self._get_sp_padded_audio_latent_length(audio_latent_length, int(sp_size)) + + # Keep requested duration semantics when callers pass 4D latents that + # are already padded for SP; other 4D lengths retain shape inference. + if provided_latent_length in {audio_latent_length, padded_latent_length}: + return audio_latent_length + return provided_latent_length + # ------------------------------------------------------------------ # Latent preparation # ------------------------------------------------------------------ @@ -497,9 +517,11 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): latents: torch.Tensor | None = None, ) -> tuple[torch.Tensor, int, int]: original_latent_length = audio_latent_length - padded_latent_length = original_latent_length latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1 + padded_latent_length = self._get_sp_padded_audio_latent_length(original_latent_length, int(sp_size)) + if latents is not None: if latents.ndim == 4: latents = self._pack_audio_latents(latents) @@ -508,6 +530,23 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) latents = noise_scale * noise + (1 - noise_scale) * latents + + if latents.shape[1] not in {original_latent_length, padded_latent_length}: + raise ValueError( + "Provided `audio_latents` has incompatible audio frame count " + f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}." + ) + + if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length: + padding = torch.zeros( + latents.shape[0], + padded_latent_length - original_latent_length, + latents.shape[2], + dtype=latents.dtype, + device=latents.device, + ) + latents = torch.cat([latents, padding], dim=1) + return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins) @@ -783,8 +822,7 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) ) audio_num_frames = round(duration_s * audio_latents_per_second) - if audio_latents is not None and audio_latents.ndim == 4: - _, _, audio_num_frames, _ = audio_latents.shape + audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, audio_latents) num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64 latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio