220 lines
8.4 KiB
Python
220 lines
8.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Unit tests for LTX image-to-video input and conditioning behavior."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
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
|
|
|
|
|
|
class TestLTXImageToVideoForwardStages:
|
|
def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 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
|
|
|
|
pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline)
|
|
image = torch.zeros(3, 8, 8)
|
|
req = DiffusionRequestBatch(
|
|
[
|
|
OmniDiffusionRequest(
|
|
prompt={
|
|
"prompt": "make the image move",
|
|
"negative_prompt": "jitter",
|
|
"multi_modal_data": {"image": image},
|
|
},
|
|
sampling_params=OmniDiffusionSamplingParams(
|
|
height=384,
|
|
width=512,
|
|
num_frames=25,
|
|
num_inference_steps=2,
|
|
),
|
|
request_id="ltx23-i2v-forward-stage-delegation",
|
|
)
|
|
]
|
|
)
|
|
seen = {}
|
|
|
|
def fake_forward_impl(req_arg, request_inputs, **kwargs):
|
|
seen["req"] = req_arg
|
|
seen["request_inputs"] = request_inputs
|
|
seen["kwargs"] = kwargs
|
|
return ["i2v-delegated"]
|
|
|
|
object.__setattr__(pipe, "_forward_impl", fake_forward_impl)
|
|
|
|
output = pipe.forward(req, noise_scale=0.5)
|
|
|
|
assert output == ["i2v-delegated"]
|
|
assert seen["req"] is req
|
|
assert seen["request_inputs"].prompt == ["make the image move"]
|
|
assert seen["request_inputs"].negative_prompt == ["jitter"]
|
|
assert seen["kwargs"]["image"] is image
|
|
assert seen["kwargs"]["noise_scale"] == 0.5
|
|
|
|
def test_denoise_timestep_kwargs_masks_video_only(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
|
|
|
|
pipe = object.__new__(LTX23ImageToVideoPipeline)
|
|
ts = torch.tensor([2.0, 4.0])
|
|
denoise_ctx = SimpleNamespace(
|
|
conditioning_mask=torch.tensor([[1.0, 0.0]]),
|
|
conditioning_mask_for_model=torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
|
|
)
|
|
|
|
kwargs = pipe._denoise_timestep_kwargs(
|
|
ts,
|
|
SimpleNamespace(cfg_parallel_ready=False),
|
|
denoise_ctx,
|
|
)
|
|
|
|
torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]]))
|
|
torch.testing.assert_close(kwargs["audio_timestep"], ts)
|
|
torch.testing.assert_close(kwargs["sigma"], ts)
|
|
|
|
|
|
class TestLTXImageToVideoConditioning:
|
|
def test_ltx23_i2v_supports_image_input(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
|
|
|
|
assert LTX23ImageToVideoPipeline.support_image_input is True
|
|
|
|
def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
|
|
|
|
image = object()
|
|
|
|
assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image
|
|
with pytest.raises(ValueError, match="exactly one image per prompt"):
|
|
LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()])
|
|
|
|
def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
|
|
|
|
image = torch.zeros(1, 3, 4, 4)
|
|
additional = {
|
|
"preprocessed_image": None,
|
|
"pixel_values": image,
|
|
"image": torch.ones_like(image),
|
|
}
|
|
|
|
assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image
|
|
|
|
def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch):
|
|
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
|
|
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
|
|
|
|
pipe = object.__new__(LTX23ImageToVideoPipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.vae_spatial_compression_ratio = 1
|
|
pipe.vae_temporal_compression_ratio = 1
|
|
pipe.transformer_spatial_patch_size = 1
|
|
pipe.transformer_temporal_patch_size = 1
|
|
|
|
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
|
|
raise AssertionError("packed I2V latents should not be noised")
|
|
|
|
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
|
|
monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
|
|
|
|
latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]])
|
|
|
|
out, conditioning_mask = pipe.prepare_latents(
|
|
image=None,
|
|
batch_size=1,
|
|
num_channels_latents=2,
|
|
height=1,
|
|
width=1,
|
|
num_frames=3,
|
|
noise_scale=1.0,
|
|
dtype=torch.float32,
|
|
device=torch.device("cpu"),
|
|
latents=latents,
|
|
)
|
|
|
|
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
|
|
torch.testing.assert_close(out, latents)
|
|
|
|
def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
|
|
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
|
|
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
|
|
|
|
pipe = object.__new__(LTX23ImageToVideoPipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.vae_spatial_compression_ratio = 1
|
|
pipe.vae_temporal_compression_ratio = 1
|
|
pipe.transformer_spatial_patch_size = 1
|
|
pipe.transformer_temporal_patch_size = 1
|
|
pipe.vae = SimpleNamespace(
|
|
latents_mean=torch.zeros(2),
|
|
latents_std=torch.ones(2),
|
|
config=SimpleNamespace(scaling_factor=1.0),
|
|
)
|
|
|
|
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
|
|
return torch.ones(shape, device=device, dtype=dtype)
|
|
|
|
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
|
|
monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
|
|
|
|
latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
|
|
|
|
out, conditioning_mask = pipe.prepare_latents(
|
|
image=None,
|
|
batch_size=1,
|
|
num_channels_latents=2,
|
|
height=1,
|
|
width=1,
|
|
num_frames=3,
|
|
noise_scale=1.0,
|
|
dtype=torch.float32,
|
|
device=torch.device("cpu"),
|
|
latents=latents,
|
|
)
|
|
|
|
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
|
|
torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
|
|
|
|
def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
|
|
|
|
pipe = object.__new__(LTX23ImageToVideoPipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.transformer_spatial_patch_size = 1
|
|
pipe.transformer_temporal_patch_size = 1
|
|
|
|
class FakeScheduler:
|
|
def step(self, noise_pred, t, latents, return_dict=False):
|
|
return (latents + noise_pred + t,)
|
|
|
|
pipe.scheduler = FakeScheduler()
|
|
latents = torch.tensor([[[1.0], [2.0], [3.0]]])
|
|
noise_pred = torch.full_like(latents, 10.0)
|
|
|
|
out = pipe._step_video_latents_i2v(
|
|
noise_pred,
|
|
latents,
|
|
torch.tensor(0.5),
|
|
latent_num_frames=3,
|
|
latent_height=1,
|
|
latent_width=1,
|
|
)
|
|
|
|
torch.testing.assert_close(out[:, :1], latents[:, :1])
|
|
torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5)
|