# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for LTX-2.3 pipeline integration. These tests verify: - Pipeline is properly registered in the diffusion registry - Post-process function is registered - Cache-DiT enablers are registered - Pipeline does NOT inherit from LTX2Pipeline - Vocoder sample rate detection logic - Re-export module works correctly """ 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.""" def test_ltx23_pipeline_does_not_inherit_from_ltx2(self): """LTX23Pipeline must NOT inherit from LTX2Pipeline.""" from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline assert not issubclass(LTX23Pipeline, LTX2Pipeline), ( "LTX23Pipeline should be fully independent and not inherit from LTX2Pipeline" ) def test_ltx23_pipeline_is_nn_module(self): """LTX23Pipeline must be an nn.Module.""" import torch.nn as nn from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline assert issubclass(LTX23Pipeline, nn.Module) def test_ltx23_pipeline_has_progress_bar(self): """LTX23Pipeline must mix in ProgressBarMixin.""" from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin 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", ), request_id="ltx23-cfg-parallel-forward-test", ) 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.""" def test_pipeline_models_registered(self): """LTX-2.3 pipeline variants must be in _DIFFUSION_MODELS.""" from vllm_omni.diffusion.registry import _DIFFUSION_MODELS expected = [ "LTX23Pipeline", "LTX23ImageToVideoPipeline", ] for name in expected: assert name in _DIFFUSION_MODELS, f"{name} not found in _DIFFUSION_MODELS" def test_pipeline_module_paths(self): """Registry entries must point to the correct modules.""" from vllm_omni.diffusion.registry import _DIFFUSION_MODELS # T2V -> pipeline_ltx2_3 assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2_3", "LTX23Pipeline") # I2V -> pipeline_ltx2_3_image2video assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == ( "ltx2", "pipeline_ltx2_3_image2video", "LTX23ImageToVideoPipeline", ) def test_post_process_funcs_registered(self): """Pipeline variants must map to get_ltx2_post_process_func.""" from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS expected = [ "LTX23Pipeline", "LTX23ImageToVideoPipeline", ] for name in expected: assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS" assert _DIFFUSION_POST_PROCESS_FUNCS[name] == "get_ltx2_post_process_func" def test_cache_dit_enablers_registered(self): """Pipeline variants must be registered in CUSTOM_DIT_ENABLERS.""" from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS expected = [ "LTX23Pipeline", "LTX23ImageToVideoPipeline", ] for name in expected: assert name in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS" class TestVocoderSampleRateDetection: """Test _detect_vocoder_output_sample_rate logic.""" def test_detects_48khz_from_config(self): """Should detect output_sampling_rate=48000 from vocoder/config.json.""" from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate with tempfile.TemporaryDirectory() as tmpdir: vocoder_dir = os.path.join(tmpdir, "vocoder") os.makedirs(vocoder_dir) with open(os.path.join(vocoder_dir, "config.json"), "w") as f: json.dump({"output_sampling_rate": 48000, "input_sampling_rate": 16000}, f) result = _detect_vocoder_output_sample_rate(tmpdir) assert result == 48000 def test_returns_none_for_no_output_sr(self): """Should return None if vocoder config has no output_sampling_rate.""" from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate with tempfile.TemporaryDirectory() as tmpdir: vocoder_dir = os.path.join(tmpdir, "vocoder") os.makedirs(vocoder_dir) with open(os.path.join(vocoder_dir, "config.json"), "w") as f: json.dump({"sampling_rate": 16000}, f) result = _detect_vocoder_output_sample_rate(tmpdir) assert result is None def test_returns_none_for_missing_directory(self): """Should return None if vocoder directory doesn't exist.""" from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate result = _detect_vocoder_output_sample_rate("/nonexistent/path") assert result is None class TestPostProcessFunction: """Test the post-process function factory.""" def test_post_process_includes_audio_sample_rate(self): """Post-process func should include audio_sample_rate when detected.""" import torch from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func with tempfile.TemporaryDirectory() as tmpdir: vocoder_dir = os.path.join(tmpdir, "vocoder") os.makedirs(vocoder_dir) with open(os.path.join(vocoder_dir, "config.json"), "w") as f: json.dump({"output_sampling_rate": 48000}, f) # Create a minimal od_config mock class MockConfig: model = tmpdir func = get_ltx2_post_process_func(MockConfig()) video = torch.zeros(1, 3, 4, 64, 64) audio = torch.zeros(1, 1, 48000) result = func((video, audio)) assert isinstance(result, dict) assert "video" in result assert "audio" in result assert result["audio_sample_rate"] == 48000 def test_post_process_without_vocoder_config(self): """Post-process func should work without vocoder config (no audio_sample_rate key).""" import torch from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func class MockConfig: model = "/nonexistent/path" func = get_ltx2_post_process_func(MockConfig()) video = torch.zeros(1, 3, 4, 64, 64) audio = torch.zeros(1, 1, 16000) result = func((video, audio)) assert isinstance(result, dict) assert "video" in result assert "audio" in result assert "audio_sample_rate" not in result class TestReExportModule: """Test that pipeline_ltx2_3_image2video.py correctly re-exports.""" def test_i2v_classes_importable(self): """I2V classes must be importable from the re-export module.""" from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline assert LTX23ImageToVideoPipeline is not None def test_post_process_func_importable(self): """get_ltx2_post_process_func must be importable from re-export module.""" from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func assert callable(get_ltx2_post_process_func) def test_i2v_classes_are_same_as_direct_import(self): """Re-exported classes must be the same objects as direct imports.""" from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23ImageToVideoPipeline as Direct from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import ( LTX23ImageToVideoPipeline as ReExported, ) assert Direct is ReExported class TestInitExports: """Test that __init__.py exports all LTX-2.3 classes.""" def test_all_ltx23_classes_exported(self): """All LTX-2.3 pipeline classes must be in the ltx2 package __all__.""" from vllm_omni.diffusion.models import ltx2 expected_classes = [ "LTX23Pipeline", "LTX23ImageToVideoPipeline", ] 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, )