647 lines
26 KiB
Python
647 lines
26 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
|
|
|
|
|
|
def _make_ltx23_request_pipe(cls):
|
|
pipe = object.__new__(cls)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.device = torch.device("cpu")
|
|
pipe.tokenizer_max_length = 99
|
|
return pipe
|
|
|
|
|
|
def _resolve_request_inputs_for_test(pipe, req):
|
|
return pipe._resolve_request_inputs(
|
|
req,
|
|
prompt=None,
|
|
negative_prompt=None,
|
|
height=None,
|
|
width=None,
|
|
num_frames=None,
|
|
frame_rate=None,
|
|
num_inference_steps=None,
|
|
timesteps=None,
|
|
guidance_scale=4.0,
|
|
num_videos_per_prompt=1,
|
|
generator=None,
|
|
latents=None,
|
|
audio_latents=None,
|
|
prompt_embeds=None,
|
|
negative_prompt_embeds=None,
|
|
prompt_attention_mask=None,
|
|
negative_prompt_attention_mask=None,
|
|
decode_timestep=0.0,
|
|
decode_noise_scale=None,
|
|
output_type="np",
|
|
max_sequence_length=None,
|
|
)
|
|
|
|
|
|
class TestLTX23RequestParsing:
|
|
def test_t2v_and_i2v_share_request_input_resolution(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
prompt_embeds = torch.tensor([[1.0, 2.0]])
|
|
negative_prompt_embeds = torch.tensor([[3.0, 4.0]])
|
|
prompt_attention_mask = torch.tensor([True, False])
|
|
negative_attention_mask = torch.tensor([False, True])
|
|
video_latents = torch.ones(1, 2, 3)
|
|
audio_latents = torch.zeros(1, 2, 3)
|
|
generator = torch.Generator().manual_seed(123)
|
|
|
|
req = DiffusionRequestBatch(
|
|
[
|
|
OmniDiffusionRequest(
|
|
prompt={
|
|
"prompt": "shared prompt",
|
|
"negative_prompt": "shared negative",
|
|
"additional_information": {
|
|
"prompt_embeds": prompt_embeds,
|
|
"negative_prompt_embeds": negative_prompt_embeds,
|
|
"attention_mask": prompt_attention_mask,
|
|
"negative_attention_mask": negative_attention_mask,
|
|
},
|
|
},
|
|
sampling_params=OmniDiffusionSamplingParams(
|
|
height=384,
|
|
width=512,
|
|
num_frames=25,
|
|
frame_rate=12.5,
|
|
num_inference_steps=1,
|
|
num_outputs_per_prompt=2,
|
|
guidance_scale=6.0,
|
|
generator=generator,
|
|
latents=video_latents,
|
|
extra_args={"audio_latents": audio_latents},
|
|
decode_timestep=[0.1],
|
|
decode_noise_scale=[0.2],
|
|
output_type="latent",
|
|
max_sequence_length=17,
|
|
),
|
|
request_id="ltx23-shared-request-inputs",
|
|
)
|
|
]
|
|
)
|
|
|
|
resolved_t2v = _resolve_request_inputs_for_test(
|
|
_make_ltx23_request_pipe(LTX23Pipeline),
|
|
req,
|
|
)
|
|
resolved_i2v = _resolve_request_inputs_for_test(
|
|
_make_ltx23_request_pipe(LTX23ImageToVideoPipeline),
|
|
req,
|
|
)
|
|
|
|
assert resolved_i2v.prompt == resolved_t2v.prompt is None
|
|
assert resolved_i2v.negative_prompt == resolved_t2v.negative_prompt is None
|
|
assert resolved_i2v.height == resolved_t2v.height == 384
|
|
assert resolved_i2v.width == resolved_t2v.width == 512
|
|
assert resolved_i2v.num_frames == resolved_t2v.num_frames == 25
|
|
assert resolved_i2v.frame_rate == resolved_t2v.frame_rate == 12.5
|
|
assert resolved_i2v.num_inference_steps == resolved_t2v.num_inference_steps == 2
|
|
assert resolved_i2v.guidance_scale == resolved_t2v.guidance_scale == 6.0
|
|
assert resolved_i2v.num_videos_per_prompt == resolved_t2v.num_videos_per_prompt == 2
|
|
assert isinstance(resolved_i2v.generator, list)
|
|
assert isinstance(resolved_t2v.generator, list)
|
|
assert [gen.initial_seed() for gen in resolved_i2v.generator] == [123, 123]
|
|
assert [gen.initial_seed() for gen in resolved_t2v.generator] == [123, 123]
|
|
assert resolved_i2v.decode_timestep == resolved_t2v.decode_timestep == [0.1]
|
|
assert resolved_i2v.decode_noise_scale == resolved_t2v.decode_noise_scale == [0.2]
|
|
assert resolved_i2v.output_type == resolved_t2v.output_type == "latent"
|
|
assert resolved_i2v.max_sequence_length == resolved_t2v.max_sequence_length == 17
|
|
torch.testing.assert_close(resolved_i2v.latents, video_latents)
|
|
torch.testing.assert_close(resolved_t2v.latents, video_latents)
|
|
torch.testing.assert_close(resolved_i2v.audio_latents, audio_latents)
|
|
torch.testing.assert_close(resolved_t2v.audio_latents, audio_latents)
|
|
torch.testing.assert_close(resolved_i2v.prompt_embeds, torch.stack([prompt_embeds]))
|
|
torch.testing.assert_close(resolved_t2v.prompt_embeds, torch.stack([prompt_embeds]))
|
|
torch.testing.assert_close(resolved_i2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds]))
|
|
torch.testing.assert_close(resolved_t2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds]))
|
|
torch.testing.assert_close(resolved_i2v.prompt_attention_mask, torch.stack([prompt_attention_mask]))
|
|
torch.testing.assert_close(resolved_t2v.prompt_attention_mask, torch.stack([prompt_attention_mask]))
|
|
torch.testing.assert_close(
|
|
resolved_i2v.negative_prompt_attention_mask,
|
|
torch.stack([negative_attention_mask]),
|
|
)
|
|
torch.testing.assert_close(
|
|
resolved_t2v.negative_prompt_attention_mask,
|
|
torch.stack([negative_attention_mask]),
|
|
)
|
|
|
|
def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
pipe = _make_ltx23_request_pipe(LTX23Pipeline)
|
|
req = DiffusionRequestBatch(
|
|
[
|
|
OmniDiffusionRequest(
|
|
prompt={
|
|
"prompt": "with embeds",
|
|
"additional_information": {"prompt_embeds": torch.tensor([[1.0]])},
|
|
},
|
|
sampling_params=OmniDiffusionSamplingParams(
|
|
height=384,
|
|
width=512,
|
|
num_frames=25,
|
|
num_inference_steps=2,
|
|
),
|
|
request_id="ltx23-mixed-precomputed-fields-0",
|
|
),
|
|
OmniDiffusionRequest(
|
|
prompt={"prompt": "without embeds"},
|
|
sampling_params=OmniDiffusionSamplingParams(
|
|
height=384,
|
|
width=512,
|
|
num_frames=25,
|
|
num_inference_steps=2,
|
|
),
|
|
request_id="ltx23-mixed-precomputed-fields-1",
|
|
),
|
|
]
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="mix of provided and missing prompt_embeds"):
|
|
_resolve_request_inputs_for_test(pipe, req)
|
|
|
|
|
|
class TestLTX23ForwardStages:
|
|
def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
|
|
|
pipe = _make_ltx23_request_pipe(LTX23Pipeline)
|
|
prompt_embeds = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]])
|
|
negative_prompt_embeds = torch.tensor([[[-1.0], [-2.0]], [[-3.0], [-4.0]]])
|
|
prompt_attention_mask = torch.tensor([[True, False], [False, True]])
|
|
negative_prompt_attention_mask = torch.tensor([[False, False], [True, True]])
|
|
|
|
(
|
|
repeated_prompt_embeds,
|
|
repeated_prompt_attention_mask,
|
|
repeated_negative_prompt_embeds,
|
|
repeated_negative_prompt_attention_mask,
|
|
) = pipe.encode_prompt(
|
|
prompt=None,
|
|
do_classifier_free_guidance=True,
|
|
num_videos_per_prompt=2,
|
|
prompt_embeds=prompt_embeds,
|
|
negative_prompt_embeds=negative_prompt_embeds,
|
|
prompt_attention_mask=prompt_attention_mask,
|
|
negative_prompt_attention_mask=negative_prompt_attention_mask,
|
|
device=torch.device("cpu"),
|
|
)
|
|
|
|
torch.testing.assert_close(
|
|
repeated_prompt_embeds,
|
|
torch.tensor([[[1.0], [2.0]], [[1.0], [2.0]], [[3.0], [4.0]], [[3.0], [4.0]]]),
|
|
)
|
|
torch.testing.assert_close(
|
|
repeated_negative_prompt_embeds,
|
|
torch.tensor([[[-1.0], [-2.0]], [[-1.0], [-2.0]], [[-3.0], [-4.0]], [[-3.0], [-4.0]]]),
|
|
)
|
|
torch.testing.assert_close(
|
|
repeated_prompt_attention_mask,
|
|
torch.tensor([[True, False], [True, False], [False, True], [False, True]]),
|
|
)
|
|
torch.testing.assert_close(
|
|
repeated_negative_prompt_attention_mask,
|
|
torch.tensor([[False, False], [False, False], [True, True], [True, True]]),
|
|
)
|
|
|
|
def test_t2v_forward_delegates_to_shared_forward_impl(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
pipe = _make_ltx23_request_pipe(LTX23Pipeline)
|
|
req = DiffusionRequestBatch(
|
|
[
|
|
OmniDiffusionRequest(
|
|
prompt="stage refactor prompt",
|
|
sampling_params=OmniDiffusionSamplingParams(
|
|
height=384,
|
|
width=512,
|
|
num_frames=25,
|
|
frame_rate=16,
|
|
num_inference_steps=3,
|
|
guidance_scale=4.5,
|
|
output_type="latent",
|
|
),
|
|
request_id="ltx23-t2v-forward-stage-delegation",
|
|
)
|
|
]
|
|
)
|
|
sigmas = [1.0, 0.5]
|
|
timesteps = [1000, 500]
|
|
attention_kwargs = {"scale": 1.0}
|
|
seen = {}
|
|
|
|
def fake_forward_impl(req_arg, request_inputs, **kwargs):
|
|
seen["req"] = req_arg
|
|
seen["request_inputs"] = request_inputs
|
|
seen["kwargs"] = kwargs
|
|
return ["delegated"]
|
|
|
|
object.__setattr__(pipe, "_forward_impl", fake_forward_impl)
|
|
|
|
output = pipe.forward(
|
|
req,
|
|
sigmas=sigmas,
|
|
timesteps=timesteps,
|
|
noise_scale=0.25,
|
|
attention_kwargs=attention_kwargs,
|
|
)
|
|
|
|
assert output == ["delegated"]
|
|
assert seen["req"] is req
|
|
assert seen["request_inputs"].height == 384
|
|
assert seen["request_inputs"].width == 512
|
|
assert seen["request_inputs"].num_frames == 25
|
|
assert seen["request_inputs"].frame_rate == 16.0
|
|
assert seen["request_inputs"].guidance_scale == 4.5
|
|
assert seen["request_inputs"].output_type == "latent"
|
|
assert seen["kwargs"] == {
|
|
"noise_scale": 0.25,
|
|
"sigmas": sigmas,
|
|
"timesteps": timesteps,
|
|
"attention_kwargs": attention_kwargs,
|
|
}
|
|
|
|
|
|
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_declares_offload_components(self):
|
|
"""LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery."""
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
|
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
|
|
|
|
pipe = object.__new__(LTX23Pipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.transformer = torch.nn.Linear(1, 1)
|
|
pipe.text_encoder = torch.nn.Linear(1, 1)
|
|
pipe.connectors = torch.nn.Linear(1, 1)
|
|
pipe.vae = torch.nn.Linear(1, 1)
|
|
pipe.audio_vae = torch.nn.Linear(1, 1)
|
|
pipe.vocoder = torch.nn.Linear(1, 1)
|
|
|
|
modules = ModuleDiscovery.discover(pipe)
|
|
|
|
assert LTX23Pipeline._dit_modules == ["transformer"]
|
|
assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"]
|
|
assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"]
|
|
assert LTX23Pipeline._resident_modules == ["vocoder"]
|
|
assert modules.dit_names == ["transformer"]
|
|
assert modules.encoder_names == ["text_encoder", "connectors"]
|
|
assert modules.resident_names == ["vocoder"]
|
|
assert len(modules.vaes) == 2
|
|
|
|
def test_ltx23_pipeline_has_diffusion_pipeline_profiler_mixin(self):
|
|
"""LTX23Pipeline must support lightweight stage timing."""
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
|
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
|
|
|
assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin)
|
|
|
|
|
|
class TestLTX23DecodeConditioning:
|
|
def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value
|
|
|
|
assert _expand_per_prompt_decode_value(
|
|
[0.1, 0.2],
|
|
prompt_batch_size=2,
|
|
effective_batch_size=4,
|
|
field_name="decode_timestep",
|
|
) == [0.1, 0.1, 0.2, 0.2]
|
|
assert _expand_per_prompt_decode_value(
|
|
[0.3],
|
|
prompt_batch_size=2,
|
|
effective_batch_size=4,
|
|
field_name="decode_timestep",
|
|
) == [0.3, 0.3, 0.3, 0.3]
|
|
assert _expand_per_prompt_decode_value(
|
|
[0.1, 0.2, 0.3, 0.4],
|
|
prompt_batch_size=2,
|
|
effective_batch_size=4,
|
|
field_name="decode_timestep",
|
|
) == [0.1, 0.2, 0.3, 0.4]
|
|
|
|
def test_decode_conditioning_rejects_ambiguous_lengths(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value
|
|
|
|
with pytest.raises(ValueError, match="decode_timestep"):
|
|
_expand_per_prompt_decode_value(
|
|
[0.1, 0.2, 0.3],
|
|
prompt_batch_size=2,
|
|
effective_batch_size=4,
|
|
field_name="decode_timestep",
|
|
)
|
|
|
|
|
|
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_for_ltx2_does_not_have_custom_enablers_registered(self):
|
|
"""Pipeline variants are *not* registered in CUSTOM_DIT_ENABLERS."""
|
|
from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
|
|
|
|
# NOTE: We used to have custom enablers for this model, but refactored to handle
|
|
# it more generically. Now we only need to ensure it has git cache adapter config.
|
|
expected = [
|
|
"LTX23Pipeline",
|
|
"LTX23ImageToVideoPipeline",
|
|
]
|
|
for name in expected:
|
|
assert name not 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 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,
|
|
)
|