684 lines
27 KiB
Python
684 lines
27 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."""
|
|
|
|
import sys
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
|
|
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
|
|
pipe.vae_spatial_compression_ratio = 32
|
|
pipe.vae_temporal_compression_ratio = 8
|
|
return pipe
|
|
|
|
|
|
class TestLTXImageToVideoForwardStages:
|
|
def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self):
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images
|
|
|
|
pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
|
|
pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
|
|
image = Image.fromarray(pixels.numpy())
|
|
|
|
actual = _preprocess_i2v_pil_images(image, height=4, width=4)
|
|
expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
|
|
|
|
torch.testing.assert_close(actual, expected)
|
|
|
|
def test_ltx25_i2v_pil_preprocessing_preserves_aspect_ratio_after_crf(self, monkeypatch):
|
|
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
|
|
|
|
pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
|
|
pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
|
|
image = Image.fromarray(pixels.numpy())
|
|
monkeypatch.setattr(conditioning, "_apply_image_conditioning_crf", lambda image_array, _crf: image_array)
|
|
|
|
actual = conditioning._preprocess_i2v_pil_images(image, height=4, width=4, crf=18)
|
|
expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
|
|
|
|
torch.testing.assert_close(actual, expected)
|
|
|
|
def test_ltx25_i2v_applies_crf18_before_resize(self, monkeypatch):
|
|
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
captured_crfs = []
|
|
captured_prepare_kwargs = []
|
|
|
|
def fake_preprocess(image, *, height, width, crf, device, dtype):
|
|
del image, device
|
|
captured_crfs.append(crf)
|
|
return torch.zeros(1, 3, height, width, dtype=dtype)
|
|
|
|
def fake_prepare_latents(**kwargs):
|
|
captured_prepare_kwargs.append(kwargs)
|
|
return kwargs["image"], None
|
|
|
|
monkeypatch.setattr(conditioning, "_preprocess_i2v_pil_images", fake_preprocess)
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.model_version = "2.5"
|
|
object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
|
|
object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
|
|
request_inputs = SimpleNamespace(
|
|
height=4,
|
|
width=4,
|
|
num_frames=3,
|
|
num_videos_per_prompt=1,
|
|
generator=None,
|
|
latents=None,
|
|
)
|
|
prompt_context = SimpleNamespace(
|
|
batch_size=1,
|
|
positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
|
|
)
|
|
|
|
pipe._prepare_video_latents_stage(
|
|
request_inputs,
|
|
prompt_context,
|
|
device=torch.device("cpu"),
|
|
noise_scale=0.0,
|
|
image=Image.new("RGB", (8, 4)),
|
|
)
|
|
|
|
assert captured_crfs == [18]
|
|
assert captured_prepare_kwargs[0]["dtype"] is torch.bfloat16
|
|
assert captured_prepare_kwargs[0]["image"].dtype is torch.bfloat16
|
|
|
|
def test_ltx25_i2v_rejects_inputs_that_cannot_receive_crf18(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.model_version = "2.5"
|
|
|
|
with pytest.raises(ValueError, match="requires PIL images"):
|
|
pipe._prepare_video_latents_stage(
|
|
SimpleNamespace(),
|
|
SimpleNamespace(),
|
|
device=torch.device("cpu"),
|
|
noise_scale=0.0,
|
|
image=torch.zeros(3, 4, 4),
|
|
)
|
|
|
|
def test_ltx25_i2v_tensor_input_can_opt_out_of_crf18(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
captured_prepare_kwargs = []
|
|
|
|
def fake_prepare_latents(**kwargs):
|
|
captured_prepare_kwargs.append(kwargs)
|
|
return kwargs["image"], None
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.model_version = "2.5"
|
|
object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
|
|
object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
|
|
request_inputs = SimpleNamespace(
|
|
height=4,
|
|
width=4,
|
|
num_frames=3,
|
|
num_videos_per_prompt=1,
|
|
generator=None,
|
|
latents=None,
|
|
image_crf=0,
|
|
)
|
|
prompt_context = SimpleNamespace(
|
|
batch_size=1,
|
|
positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
|
|
)
|
|
|
|
pipe._prepare_video_latents_stage(
|
|
request_inputs,
|
|
prompt_context,
|
|
device=torch.device("cpu"),
|
|
noise_scale=0.0,
|
|
image=torch.zeros(3, 4, 4),
|
|
)
|
|
|
|
assert captured_prepare_kwargs[0]["image"].shape == (1, 3, 4, 4)
|
|
|
|
def test_ltx25_i2v_crf_preserves_single_pixel_dimension(self, monkeypatch):
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
|
|
|
|
image = torch.zeros(1, 4, 3, dtype=torch.uint8).numpy()
|
|
monkeypatch.setitem(sys.modules, "av", None)
|
|
|
|
assert _apply_image_conditioning_crf(image, 18) is image
|
|
|
|
def test_ltx25_i2v_reports_missing_pyav(self, monkeypatch):
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
|
|
|
|
monkeypatch.setitem(sys.modules, "av", None)
|
|
|
|
with pytest.raises(ImportError, match="PyAV with a libx264 encoder"):
|
|
_apply_image_conditioning_crf(torch.zeros(4, 4, 3, dtype=torch.uint8).numpy(), 18)
|
|
|
|
def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
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(LTX2Pipeline)
|
|
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_run_recipe(req_arg, request_inputs, **kwargs):
|
|
seen["req"] = req_arg
|
|
seen["request_inputs"] = request_inputs
|
|
seen["kwargs"] = kwargs
|
|
return ["i2v-delegated"]
|
|
|
|
object.__setattr__(pipe, "_run_recipe", fake_run_recipe)
|
|
|
|
output = pipe.forward(req)
|
|
|
|
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"]["request_sigmas"] is None
|
|
|
|
def test_unified_entry_selects_t2v_without_images_and_rejects_mixed_batches(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
request_inputs = SimpleNamespace(latents=None)
|
|
t2v_req = SimpleNamespace(prompts=["text only", {"prompt": "also text only"}])
|
|
|
|
assert pipe._resolve_request_image(t2v_req, None, request_inputs) is None
|
|
|
|
mixed_req = SimpleNamespace(
|
|
prompts=[
|
|
"text only",
|
|
{"prompt": "image conditioned", "multi_modal_data": {"image": torch.zeros(3, 8, 8)}},
|
|
]
|
|
)
|
|
with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
|
|
pipe._resolve_request_image(mixed_req, None, request_inputs)
|
|
|
|
def test_denoise_timestep_kwargs_masks_video_and_keeps_audio_per_token(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
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(guidance_parallel_ready=False),
|
|
denoise_ctx,
|
|
video_token_count=2,
|
|
audio_token_count=1,
|
|
)
|
|
|
|
torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]]))
|
|
torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None])
|
|
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 LTX2Pipeline
|
|
|
|
assert LTX2Pipeline.support_image_input is True
|
|
|
|
def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
image = object()
|
|
|
|
assert LTX2Pipeline._resolve_single_prompt_image([image]) is image
|
|
with pytest.raises(ValueError, match="exactly one image per prompt"):
|
|
LTX2Pipeline._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 LTX2Pipeline
|
|
|
|
image = torch.zeros(1, 3, 4, 4)
|
|
additional = {
|
|
"preprocessed_image": None,
|
|
"pixel_values": image,
|
|
"image": torch.ones_like(image),
|
|
}
|
|
|
|
assert LTX2Pipeline._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 LTX2Pipeline
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
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
|
|
object.__setattr__(
|
|
pipe,
|
|
"_encode_i2v_image_latents",
|
|
lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]),
|
|
)
|
|
|
|
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=torch.zeros(1, 3, 1, 1),
|
|
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([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))
|
|
|
|
@pytest.mark.parametrize(
|
|
("model_version", "sampled_shape"),
|
|
[("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
|
|
)
|
|
def test_i2v_5d_latents_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
|
|
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 LTX2Pipeline
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.model_version = model_version
|
|
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),
|
|
)
|
|
object.__setattr__(
|
|
pipe,
|
|
"_encode_i2v_image_latents",
|
|
lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]),
|
|
)
|
|
|
|
sampled_shapes = []
|
|
|
|
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
|
|
sampled_shapes.append(tuple(shape))
|
|
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=torch.zeros(1, 3, 1, 1),
|
|
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([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
|
|
assert sampled_shapes == [sampled_shape]
|
|
|
|
@pytest.mark.parametrize(
|
|
("model_version", "sampled_shape"),
|
|
[("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
|
|
)
|
|
def test_i2v_image_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
|
|
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 LTX2Pipeline
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.model_version = model_version
|
|
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(
|
|
encode=lambda image: image,
|
|
latents_mean=torch.zeros(2),
|
|
latents_std=torch.ones(2),
|
|
)
|
|
monkeypatch.setattr(
|
|
ltx2_conditioning,
|
|
"retrieve_latents",
|
|
lambda *_args, **_kwargs: torch.tensor([[[[[10.0]]], [[[11.0]]]]]),
|
|
)
|
|
sampled_shapes = []
|
|
|
|
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
|
|
sampled_shapes.append(tuple(shape))
|
|
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)
|
|
|
|
out, conditioning_mask = pipe.prepare_latents(
|
|
image=torch.zeros(1, 3, 1, 1),
|
|
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"),
|
|
)
|
|
|
|
assert sampled_shapes == [sampled_shape]
|
|
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_ltx25_re_noise_uses_official_scalar_noise_scale(self, monkeypatch):
|
|
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.model_version = "2.5"
|
|
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),
|
|
)
|
|
object.__setattr__(
|
|
pipe,
|
|
"_encode_i2v_image_latents",
|
|
lambda *_args, **_kwargs: torch.zeros(1, 2, 1, 1, 1, dtype=torch.bfloat16),
|
|
)
|
|
captured = {}
|
|
|
|
def fake_create_noised_state(latents, noise_scale, generator=None):
|
|
del generator
|
|
captured["noise_scale"] = noise_scale
|
|
return latents
|
|
|
|
monkeypatch.setattr(ltx2_latents, "create_noised_state", fake_create_noised_state)
|
|
pipe.prepare_latents(
|
|
image=torch.zeros(1, 3, 1, 1),
|
|
batch_size=1,
|
|
num_channels_latents=2,
|
|
height=1,
|
|
width=1,
|
|
num_frames=3,
|
|
noise_scale=0.909375,
|
|
dtype=torch.bfloat16,
|
|
device=torch.device("cpu"),
|
|
latents=torch.zeros(1, 2, 3, 1, 1, dtype=torch.bfloat16),
|
|
)
|
|
|
|
assert captured["noise_scale"] == 0.909375
|
|
|
|
def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.transformer_spatial_patch_size = 1
|
|
pipe.transformer_temporal_patch_size = 1
|
|
|
|
class FakeScheduler:
|
|
sigmas = torch.tensor([1.0, 0.5])
|
|
|
|
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,
|
|
0,
|
|
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:] - 0.5 * noise_pred[:, 1:])
|
|
|
|
def test_shared_step_adapter_uses_official_euler_for_t2v(self):
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
|
|
|
|
pipeline = SimpleNamespace(scheduler=SimpleNamespace(sigmas=torch.tensor([1.0, 0.5])))
|
|
adapter = LTXVideoAudioStepAdapter(
|
|
pipeline,
|
|
SimpleNamespace(sigmas=torch.tensor([1.0, 0.25])),
|
|
latent_num_frames=1,
|
|
latent_height=1,
|
|
latent_width=1,
|
|
image_conditioned=False,
|
|
)
|
|
video = torch.tensor([[[2.0]]])
|
|
audio = torch.tensor([[[4.0]]])
|
|
|
|
((video, audio),) = adapter.step(
|
|
(torch.tensor([[[2.0]]]), torch.tensor([[[4.0]]])),
|
|
(torch.tensor(1.0), torch.tensor(1.0)),
|
|
(video, audio),
|
|
)
|
|
|
|
torch.testing.assert_close(video, torch.tensor([[[1.0]]]))
|
|
torch.testing.assert_close(audio, torch.tensor([[[1.0]]]))
|
|
|
|
def test_shared_step_adapter_uses_official_euler_for_i2v_audio(self):
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
|
|
|
|
class FakePipeline:
|
|
scheduler = object()
|
|
|
|
def __init__(self):
|
|
self.step_indices = []
|
|
|
|
def _step_video_latents_i2v(
|
|
self,
|
|
noise_pred,
|
|
latents,
|
|
step_index,
|
|
latent_num_frames,
|
|
latent_height,
|
|
latent_width,
|
|
):
|
|
del noise_pred, latent_num_frames, latent_height, latent_width
|
|
self.step_indices.append(step_index)
|
|
return latents
|
|
|
|
pipeline = FakePipeline()
|
|
adapter = LTXVideoAudioStepAdapter(
|
|
pipeline,
|
|
SimpleNamespace(sigmas=torch.tensor([1.0, 0.5, 0.0])),
|
|
latent_num_frames=1,
|
|
latent_height=1,
|
|
latent_width=1,
|
|
image_conditioned=True,
|
|
)
|
|
video = torch.tensor([[[1.0]]])
|
|
audio = torch.tensor([[[2.0]]])
|
|
video_velocity = torch.zeros_like(video)
|
|
audio_velocity = torch.tensor([[[4.0]]])
|
|
|
|
((video, audio),) = adapter.step(
|
|
(video_velocity, audio_velocity),
|
|
(torch.tensor(1.0), torch.tensor(1.0)),
|
|
(video, audio),
|
|
)
|
|
torch.testing.assert_close(audio, torch.tensor([[[0.0]]]))
|
|
|
|
((video, audio),) = adapter.step(
|
|
(video_velocity, audio_velocity),
|
|
(torch.tensor(0.5), torch.tensor(0.5)),
|
|
(video, audio),
|
|
)
|
|
torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
|
|
assert pipeline.step_indices == [0, 1]
|
|
|
|
def test_ltx25_ancestral_adapter_matches_official_seeded_step_and_preserves_i2v_mask(self):
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
|
|
|
|
device = torch.device("cpu")
|
|
sigmas = torch.tensor([1.0, 0.5, 0.0])
|
|
pipeline = SimpleNamespace(device=device, scheduler=SimpleNamespace(sigmas=sigmas))
|
|
mask = torch.tensor([[1.0, 0.0]])
|
|
adapter = LTXVideoAudioStepAdapter(
|
|
pipeline,
|
|
SimpleNamespace(sigmas=sigmas.clone()),
|
|
latent_num_frames=1,
|
|
latent_height=1,
|
|
latent_width=1,
|
|
image_conditioned=True,
|
|
sampler="euler_ancestral",
|
|
generator=torch.Generator(device=device).manual_seed(42),
|
|
conditioning_mask=mask,
|
|
)
|
|
video = torch.tensor([[[3.0], [2.0]]])
|
|
audio = torch.tensor([[[4.0]]])
|
|
video_velocity = torch.tensor([[[7.0], [0.25]]])
|
|
audio_velocity = torch.tensor([[[0.5]]])
|
|
|
|
reference_generator = torch.Generator(device=device).manual_seed(10042)
|
|
video_noise = torch.randn(video.shape, generator=reference_generator)
|
|
audio_noise = torch.randn(audio.shape, generator=reference_generator)
|
|
|
|
def official_step(sample, velocity, noise):
|
|
sigma = sigmas[0]
|
|
sigma_next = sigmas[1]
|
|
denoised = sample.float() - velocity.float() * sigma
|
|
sigma_down = sigma_next * (sigma_next / sigma)
|
|
ratio = sigma_down / sigma
|
|
deterministic = ratio * sample.float() + (1.0 - ratio) * denoised
|
|
alpha_next = 1.0 - sigma_next
|
|
alpha_down = 1.0 - sigma_down
|
|
coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
|
|
return (alpha_next / alpha_down) * deterministic + noise * coeff
|
|
|
|
expected_video = official_step(video, video_velocity, video_noise)
|
|
expected_video[:, :1] = video[:, :1]
|
|
expected_audio = official_step(audio, audio_velocity, audio_noise)
|
|
|
|
((actual_video, actual_audio),) = adapter.step(
|
|
(video_velocity, audio_velocity),
|
|
(sigmas[0], sigmas[0]),
|
|
(video, audio),
|
|
)
|
|
torch.testing.assert_close(actual_video, expected_video)
|
|
torch.testing.assert_close(actual_audio, expected_audio)
|
|
|
|
terminal_video_velocity = torch.full_like(actual_video, 0.125)
|
|
terminal_audio_velocity = torch.full_like(actual_audio, 0.25)
|
|
((terminal_video, terminal_audio),) = adapter.step(
|
|
(terminal_video_velocity, terminal_audio_velocity),
|
|
(sigmas[1], sigmas[1]),
|
|
(actual_video, actual_audio),
|
|
)
|
|
expected_terminal_video = actual_video - terminal_video_velocity * sigmas[1]
|
|
expected_terminal_video[:, :1] = actual_video[:, :1]
|
|
expected_terminal_audio = actual_audio - terminal_audio_velocity * sigmas[1]
|
|
torch.testing.assert_close(terminal_video, expected_terminal_video)
|
|
torch.testing.assert_close(terminal_audio, expected_terminal_audio)
|
|
|
|
def test_ltx25_ancestral_quantizes_x0_to_model_dtype_before_fp32_step(self):
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import _ancestral_euler_step_from_velocity
|
|
|
|
sample = torch.tensor([[[3.140625, -1.234375]]], dtype=torch.bfloat16)
|
|
velocity = torch.tensor([[[0.333984375, 2.71875]]], dtype=torch.bfloat16)
|
|
sigmas = torch.tensor([0.73, 0.41, 0.0], dtype=torch.float32)
|
|
generator = torch.Generator(device="cpu").manual_seed(10042)
|
|
reference_generator = torch.Generator(device="cpu").manual_seed(10042)
|
|
|
|
sigma, sigma_next = sigmas[:2]
|
|
denoised = (sample.float() - velocity.float() * sigma).to(sample.dtype).float()
|
|
sigma_down = sigma_next * (sigma_next / sigma)
|
|
sigma_down_ratio = sigma_down / sigma
|
|
expected = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised
|
|
alpha_next = 1.0 - sigma_next
|
|
alpha_down = 1.0 - sigma_down
|
|
coefficient = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
|
|
noise = torch.randn(
|
|
sample.shape,
|
|
generator=reference_generator,
|
|
dtype=sample.dtype,
|
|
device=sample.device,
|
|
)
|
|
expected = ((alpha_next / alpha_down) * expected + noise.float() * coefficient).to(sample.dtype)
|
|
|
|
actual = _ancestral_euler_step_from_velocity(sample, velocity, sigmas, 0, generator)
|
|
|
|
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
|
|
|
def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
denoise_ctx = LTXDenoiseContext(
|
|
latents=torch.empty(1, 3, 1),
|
|
audio_latents=torch.empty(1, 1, 1),
|
|
video_coords=torch.empty(1),
|
|
audio_coords=torch.empty(1),
|
|
conditioning_mask=torch.tensor([[1.0, 0.0, 0.0]]),
|
|
)
|
|
|
|
actual = pipe._video_guidance_model_sigma(torch.tensor(0.75), denoise_ctx)
|
|
|
|
torch.testing.assert_close(actual, torch.tensor([[[0.0], [0.75], [0.75]]]))
|