832 lines
33 KiB
Python
832 lines
33 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Unit tests for the shared LTX pipeline runtime and public contracts."""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
|
|
LTX2_COMPONENT_PROFILE,
|
|
LTX23_COMPONENT_PROFILE,
|
|
)
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
|
|
LTXDenoiseExecutor,
|
|
LTXPhaseResult,
|
|
prepare_scheduler_stage,
|
|
)
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
|
|
LTX_LEGACY_VELOCITY_GUIDANCE,
|
|
LTX_OFFICIAL_X0_GUIDANCE,
|
|
)
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import LTXPipelineRuntime
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
|
|
LTX2ImageToVideoPipeline,
|
|
LTX2Pipeline,
|
|
LTX23ImageToVideoPipeline,
|
|
LTX23Pipeline,
|
|
LTXOneStagePipeline,
|
|
)
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
|
|
LTX2ImageToVideoTwoStagesPipeline,
|
|
LTX2TwoStagesPipeline,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _make_ltx_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,
|
|
)
|
|
|
|
|
|
def test_ltx_versions_share_runtime_without_cross_version_inheritance():
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
|
|
DistributedAutoencoderKLLTX2Video,
|
|
)
|
|
|
|
assert issubclass(LTX2Pipeline, LTXPipelineRuntime)
|
|
assert issubclass(LTX23Pipeline, LTXPipelineRuntime)
|
|
assert issubclass(LTX2TwoStagesPipeline, LTXPipelineRuntime)
|
|
assert not issubclass(LTX23Pipeline, LTX2Pipeline)
|
|
assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
|
|
assert LTX23Pipeline.component_profile is LTX23_COMPONENT_PROFILE
|
|
assert LTX2Pipeline.one_stage_recipe is LTX2_ONE_STAGE_RECIPE
|
|
assert LTX23Pipeline.one_stage_recipe is LTX23_ONE_STAGE_RECIPE
|
|
assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
|
|
assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
|
|
|
|
|
|
def test_ltx_one_stage_variants_expose_shared_request_batch_and_distributed_decode_contracts():
|
|
one_stage_variants = (
|
|
LTX2Pipeline,
|
|
LTX2ImageToVideoPipeline,
|
|
LTX23Pipeline,
|
|
LTX23ImageToVideoPipeline,
|
|
)
|
|
assert all(pipeline_cls.supports_request_batch for pipeline_cls in one_stage_variants)
|
|
assert all(pipeline_cls.distributed_video_decode for pipeline_cls in one_stage_variants)
|
|
assert not LTX2TwoStagesPipeline.supports_request_batch
|
|
assert not LTX2ImageToVideoTwoStagesPipeline.supports_request_batch
|
|
|
|
|
|
@pytest.mark.parametrize("pipeline_cls", [LTX2Pipeline, LTX23Pipeline])
|
|
def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cls):
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
requests = [
|
|
OmniDiffusionRequest(
|
|
prompt=f"prompt-{index}",
|
|
sampling_params=OmniDiffusionSamplingParams(),
|
|
request_id=f"request-{index}",
|
|
)
|
|
for index in range(2)
|
|
]
|
|
request_batch = DiffusionRequestBatch(requests)
|
|
video = torch.tensor([[1.0], [2.0]])
|
|
audio = torch.tensor([[3.0], [4.0]])
|
|
pipe = object.__new__(pipeline_cls)
|
|
torch.nn.Module.__init__(pipe)
|
|
object.__setattr__(pipe, "_decode_output", lambda **_kwargs: DiffusionOutput(output=(video, audio)))
|
|
forward_ctx = SimpleNamespace(
|
|
req=request_batch,
|
|
request_inputs=SimpleNamespace(
|
|
output_type="np",
|
|
generator=None,
|
|
decode_timestep=0.0,
|
|
decode_noise_scale=None,
|
|
),
|
|
prompt_context=SimpleNamespace(connector_prompt_embeds=torch.empty(2, 1)),
|
|
device=torch.device("cpu"),
|
|
batch_size=2,
|
|
num_videos_per_prompt=1,
|
|
)
|
|
|
|
outputs = pipe.decode_phase(
|
|
LTXPhaseResult(
|
|
forward_context=forward_ctx,
|
|
video=video,
|
|
audio=audio,
|
|
)
|
|
)
|
|
|
|
assert len(outputs) == 2
|
|
torch.testing.assert_close(outputs[0].output[0], video[:1])
|
|
torch.testing.assert_close(outputs[0].output[1], audio[:1])
|
|
torch.testing.assert_close(outputs[1].output[0], video[1:])
|
|
torch.testing.assert_close(outputs[1].output[1], audio[1:])
|
|
|
|
|
|
def test_ltx_one_stage_variants_share_forward_template():
|
|
assert LTX2Pipeline._forward_impl is LTXOneStagePipeline._forward_impl
|
|
assert LTX2ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl
|
|
assert LTX23Pipeline._forward_impl is LTXOneStagePipeline._forward_impl
|
|
assert LTX23ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl
|
|
assert LTX2ImageToVideoPipeline.forward is LTX23ImageToVideoPipeline.forward
|
|
|
|
|
|
def test_ltx_versions_share_request_prompt_and_step_templates():
|
|
shared_methods = (
|
|
"_get_gemma_prompt_embeds",
|
|
"encode_prompt",
|
|
"check_inputs",
|
|
"_resolve_request_inputs",
|
|
"_prepare_prompt_context",
|
|
"_denoise_step",
|
|
)
|
|
for method_name in shared_methods:
|
|
base_method = getattr(LTXPipelineRuntime, method_name)
|
|
assert getattr(LTX2Pipeline, method_name) is base_method
|
|
assert getattr(LTX23Pipeline, method_name) is base_method
|
|
|
|
|
|
def test_ltx_versions_select_guidance_without_overriding_control_flow():
|
|
assert LTX2Pipeline.guidance_strategy is LTX_LEGACY_VELOCITY_GUIDANCE
|
|
assert LTX23Pipeline.guidance_strategy is LTX_OFFICIAL_X0_GUIDANCE
|
|
assert LTX2Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step
|
|
assert LTX23Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step
|
|
assert LTX2Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise
|
|
assert LTX23Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise
|
|
|
|
|
|
def test_ltx2_two_stage_variants_share_stage_orchestration():
|
|
assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTX2TwoStagesPipeline)
|
|
assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTXI2VConditioningMixin)
|
|
assert LTX2ImageToVideoTwoStagesPipeline._run_two_stage is LTX2TwoStagesPipeline._run_two_stage
|
|
|
|
|
|
def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
|
|
assert "_denoise_loop" not in LTX2Pipeline.__dict__
|
|
assert "_denoise_loop" not in LTX23Pipeline.__dict__
|
|
assert issubclass(LTX2ImageToVideoPipeline, LTXI2VConditioningMixin)
|
|
assert issubclass(LTX23ImageToVideoPipeline, LTXI2VConditioningMixin)
|
|
assert LTX2ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents
|
|
assert LTX23ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents
|
|
assert LTX2Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents
|
|
assert LTX23Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents
|
|
assert LTX2Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents
|
|
assert LTX23Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents
|
|
assert LTX2Pipeline.decode_phase is LTXPipelineRuntime.decode_phase
|
|
assert LTX23Pipeline.decode_phase is LTXPipelineRuntime.decode_phase
|
|
|
|
|
|
def test_denoise_executor_owns_progress_and_interrupt():
|
|
updates = []
|
|
|
|
class Progress:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, traceback):
|
|
return False
|
|
|
|
def update(self):
|
|
updates.append(True)
|
|
|
|
pipeline = SimpleNamespace(
|
|
interrupt=False,
|
|
progress_bar=lambda total: Progress(),
|
|
)
|
|
seen: list[tuple[int, float]] = []
|
|
timesteps = torch.tensor([3.0, 2.0, 1.0])
|
|
initial_state = LTXAVState(video=torch.tensor(0.0), audio=torch.tensor(10.0))
|
|
|
|
def step(index, timestep, state):
|
|
seen.append((index, timestep.item()))
|
|
pipeline.interrupt = True
|
|
return LTXAVState(video=state.video + 1, audio=state.audio + 1)
|
|
|
|
state = LTXDenoiseExecutor.run(pipeline, initial_state, timesteps, step)
|
|
|
|
assert seen == [(0, 3.0)]
|
|
assert updates == [True]
|
|
torch.testing.assert_close(state.video, torch.tensor(1.0))
|
|
torch.testing.assert_close(state.audio, torch.tensor(11.0))
|
|
|
|
|
|
def test_ltx2_two_stage_reuses_prompt_context_between_phases():
|
|
request_inputs = LTXRequestInputs(
|
|
prompt="prompt",
|
|
negative_prompt="negative",
|
|
height=32,
|
|
width=32,
|
|
num_frames=1,
|
|
frame_rate=24.0,
|
|
num_inference_steps=4,
|
|
guidance_scale=4.0,
|
|
guidance_rescale=0.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.25,
|
|
decode_noise_scale=0.5,
|
|
output_type="np",
|
|
max_sequence_length=16,
|
|
)
|
|
prompt_context = object()
|
|
phase_calls = []
|
|
|
|
def resolve_request_inputs(req, **kwargs):
|
|
return request_inputs
|
|
|
|
def run_phase(req, inputs, *, prompt_context=None, **kwargs):
|
|
phase_calls.append((inputs, prompt_context))
|
|
if len(phase_calls) == 1:
|
|
assert prompt_context is None
|
|
assert inputs.decode_timestep == 0.25
|
|
assert inputs.decode_noise_scale == 0.5
|
|
context = prompt_context_sentinel
|
|
video = torch.tensor([1.0])
|
|
audio = torch.tensor([2.0])
|
|
else:
|
|
assert prompt_context is prompt_context_sentinel
|
|
torch.testing.assert_close(inputs.latents, torch.tensor([11.0]))
|
|
torch.testing.assert_close(inputs.audio_latents, torch.tensor([2.0]))
|
|
assert inputs.guidance_scale == 1.0
|
|
assert inputs.num_inference_steps == 3
|
|
assert inputs.decode_timestep == 0.0
|
|
assert inputs.decode_noise_scale is None
|
|
context = prompt_context
|
|
video = torch.tensor([3.0])
|
|
audio = torch.tensor([4.0])
|
|
return LTXPhaseResult(
|
|
forward_context=SimpleNamespace(prompt_context=context),
|
|
video=video,
|
|
audio=audio,
|
|
)
|
|
|
|
def decode_phase(phase):
|
|
return DiffusionOutput(output=(phase.video, phase.audio))
|
|
|
|
class FakeUpsampler(torch.nn.Module):
|
|
def forward(self, *, latents, output_type, return_dict):
|
|
assert output_type == "latent"
|
|
assert not return_dict
|
|
return (latents + 10,)
|
|
|
|
prompt_context_sentinel = prompt_context
|
|
pipeline = object.__new__(LTX2TwoStagesPipeline)
|
|
torch.nn.Module.__init__(pipeline)
|
|
pipeline.distilled = True
|
|
pipeline.upsample_pipe = FakeUpsampler()
|
|
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
|
|
object.__setattr__(pipeline, "run_phase", run_phase)
|
|
object.__setattr__(pipeline, "decode_phase", decode_phase)
|
|
|
|
output = pipeline.forward(SimpleNamespace())
|
|
|
|
assert len(phase_calls) == 2
|
|
assert phase_calls[1][1] is prompt_context_sentinel
|
|
torch.testing.assert_close(output.output[0], torch.tensor([3.0]))
|
|
torch.testing.assert_close(output.output[1], torch.tensor([4.0]))
|
|
|
|
|
|
def test_ltx_custom_sigmas_bypass_scheduler_shifting():
|
|
from diffusers import FlowMatchEulerDiscreteScheduler
|
|
|
|
scheduler = FlowMatchEulerDiscreteScheduler(
|
|
use_dynamic_shifting=True,
|
|
base_shift=0.95,
|
|
max_shift=2.05,
|
|
shift_terminal=0.1,
|
|
)
|
|
pipeline = SimpleNamespace(
|
|
scheduler=scheduler,
|
|
_make_video_audio_scheduler=lambda *args, **kwargs: object(),
|
|
)
|
|
sigmas = [1.0, 0.75, 0.25, 0.0]
|
|
|
|
audio_scheduler, _, timesteps = prepare_scheduler_stage(
|
|
pipeline,
|
|
SimpleNamespace(num_inference_steps=3),
|
|
device=torch.device("cpu"),
|
|
sigmas=sigmas,
|
|
timesteps=None,
|
|
latent_num_frames=1,
|
|
latent_height=1,
|
|
latent_width=1,
|
|
)
|
|
|
|
expected = torch.tensor(sigmas, dtype=torch.float32)
|
|
torch.testing.assert_close(pipeline.scheduler.sigmas, expected)
|
|
torch.testing.assert_close(audio_scheduler.sigmas, expected)
|
|
torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps)
|
|
|
|
|
|
class TestLTXRequestParsing:
|
|
@pytest.mark.parametrize(
|
|
("t2v_cls", "i2v_cls"),
|
|
[
|
|
(LTX2Pipeline, LTX2ImageToVideoPipeline),
|
|
(LTX23Pipeline, LTX23ImageToVideoPipeline),
|
|
],
|
|
)
|
|
def test_t2v_and_i2v_share_request_input_resolution(self, t2v_cls, i2v_cls):
|
|
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_ltx_request_pipe(t2v_cls),
|
|
req,
|
|
)
|
|
resolved_i2v = _resolve_request_inputs_for_test(
|
|
_make_ltx_request_pipe(i2v_cls),
|
|
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 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_ltx_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 import LTX23Pipeline
|
|
|
|
pipe = _make_ltx_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 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_ltx_request_pipe(LTX23Pipeline)
|
|
request_sigmas = [1.0, 0.4]
|
|
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",
|
|
sigmas=request_sigmas,
|
|
),
|
|
request_id="ltx23-t2v-forward-stage-delegation",
|
|
)
|
|
]
|
|
)
|
|
fallback_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=fallback_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": request_sigmas,
|
|
"timesteps": timesteps,
|
|
"attention_kwargs": attention_kwargs,
|
|
}
|
|
|
|
|
|
class TestPipelineComponents:
|
|
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 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
|
|
|
|
|
|
class TestLTX23DecodeConditioning:
|
|
def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self):
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime 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.ltx2_pipeline_runtime 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_module_paths(self):
|
|
"""Registry entries must point to the correct modules."""
|
|
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
|
|
|
|
assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX23Pipeline")
|
|
|
|
assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == (
|
|
"ltx2",
|
|
"pipeline_ltx2",
|
|
"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 = [
|
|
"LTX2Pipeline",
|
|
"LTX2ImageToVideoPipeline",
|
|
"LTX2TwoStagesPipeline",
|
|
"LTX2ImageToVideoTwoStagesPipeline",
|
|
"LTX2T2VDMD2Pipeline",
|
|
"LTX2I2VDMD2Pipeline",
|
|
"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"
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_class_name",
|
|
[
|
|
"LTX2Pipeline",
|
|
"LTX2ImageToVideoPipeline",
|
|
"LTX2TwoStagesPipeline",
|
|
"LTX2ImageToVideoTwoStagesPipeline",
|
|
"LTX2T2VDMD2Pipeline",
|
|
"LTX2I2VDMD2Pipeline",
|
|
"LTX23Pipeline",
|
|
"LTX23ImageToVideoPipeline",
|
|
],
|
|
)
|
|
def test_post_process_func_resolves_from_every_entry_module(self, model_class_name):
|
|
from vllm_omni.diffusion.registry import get_diffusion_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 config_file:
|
|
json.dump({"output_sampling_rate": 48000}, config_file)
|
|
|
|
post_process = get_diffusion_post_process_func(
|
|
SimpleNamespace(model_class_name=model_class_name, model=tmpdir)
|
|
)
|
|
|
|
result = post_process((torch.zeros(1), torch.zeros(1)))
|
|
assert result["audio_sample_rate"] == 48000
|
|
|
|
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.ltx2_components 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.ltx2_components 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.ltx2_components 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 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 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
|