637 lines
25 KiB
Python
637 lines
25 KiB
Python
# 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,
|
|
)
|