[BugFix] Fix LTX-2.3 audio latent padding for sequence parallelism (#3854)

Signed-off-by: mglyn <1203789601@qq.com>

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,
+            )

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