137 lines
5.0 KiB
Python
137 lines
5.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""CUDA coverage for LTX decoded-video output preparation."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from diffusers.video_processor import VideoProcessor
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _prepare_ltx2_video_output
|
|
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
|
|
|
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
|
@pytest.mark.parametrize("do_normalize", [True, False])
|
|
def test_ltx_uint8_output_runs_on_cuda_with_bounded_rounding_drift(dtype, do_normalize):
|
|
processor = VideoProcessor(vae_scale_factor=8, do_normalize=do_normalize)
|
|
source = torch.linspace(
|
|
-1.2 if do_normalize else 0.0,
|
|
1.2 if do_normalize else 1.0,
|
|
2 * 3 * 4 * 5 * 7,
|
|
device="cuda",
|
|
dtype=torch.float32,
|
|
).reshape(2, 3, 4, 5, 7)
|
|
source = source.to(dtype)
|
|
|
|
legacy = processor.postprocess_video(source.clone(), output_type="np")
|
|
if do_normalize:
|
|
expected = np.clip(legacy, 0.0, 1.0)
|
|
else:
|
|
# With normalization disabled, VideoProcessor's contract is already-[0, 1] input.
|
|
np.testing.assert_array_less(legacy, 1.0 + np.finfo(np.float32).eps)
|
|
np.testing.assert_array_less(-np.finfo(np.float32).eps, legacy)
|
|
expected = legacy
|
|
expected = np.rint(expected * 255.0).astype(np.uint8)
|
|
|
|
prepared = _prepare_ltx2_video_output(source.clone(), do_normalize=do_normalize)
|
|
|
|
assert prepared.is_cuda
|
|
assert prepared.shape == (2, 4, 5, 7, 3)
|
|
assert prepared.dtype == torch.uint8
|
|
assert prepared.is_contiguous()
|
|
delta = np.abs(prepared.cpu().numpy().astype(np.int16) - expected.astype(np.int16))
|
|
assert delta.max() <= 1
|
|
|
|
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
|
def test_ltx_uint8_output_defensively_clamps_unnormalized_values_on_cuda():
|
|
source = (
|
|
torch.tensor([-0.2, 0.0, 0.5, 1.0, 1.2], device="cuda", dtype=torch.bfloat16)
|
|
.reshape(1, 1, 1, 1, 5)
|
|
.expand(1, 3, 1, 1, 5)
|
|
.clone()
|
|
)
|
|
|
|
prepared = _prepare_ltx2_video_output(source, do_normalize=False)
|
|
|
|
expected = torch.tensor([0, 0, 128, 255, 255], device="cuda", dtype=torch.uint8).reshape(1, 1, 1, 5)
|
|
assert torch.equal(prepared[..., 0], expected)
|
|
assert torch.equal(prepared[..., 1], expected)
|
|
assert torch.equal(prepared[..., 2], expected)
|
|
|
|
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
|
@pytest.mark.parametrize("use_diffusion_decoder", [False, True], ids=["conv-vae", "diff-vae"])
|
|
def test_ltx_bf16_decoder_output_uses_cuda_transport_path(use_diffusion_decoder):
|
|
decoded_video = torch.linspace(
|
|
-1.2,
|
|
1.2,
|
|
1 * 3 * 2 * 3 * 4,
|
|
device="cuda",
|
|
dtype=torch.float32,
|
|
).reshape(1, 3, 2, 3, 4)
|
|
decoded_video = decoded_video.to(torch.bfloat16)
|
|
processor = VideoProcessor(vae_scale_factor=8)
|
|
legacy = processor.postprocess_video(decoded_video.clone(), output_type="np")
|
|
expected = np.rint(np.clip(legacy, 0.0, 1.0) * 255.0).astype(np.uint8)
|
|
decode_calls = []
|
|
|
|
class VideoDecoder:
|
|
dtype = torch.bfloat16
|
|
config = SimpleNamespace(timestep_conditioning=False)
|
|
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
def decode(self, *_args, **_kwargs):
|
|
decode_calls.append(self.name)
|
|
return (decoded_video.clone(),)
|
|
|
|
class AudioVae:
|
|
dtype = torch.bfloat16
|
|
|
|
def decode(self, audio_latents, *, return_dict):
|
|
assert return_dict is False
|
|
return (audio_latents,)
|
|
|
|
pipe = object.__new__(LTX2Pipeline)
|
|
torch.nn.Module.__init__(pipe)
|
|
pipe.reports_stage_durations = False
|
|
pipe.distributed_video_decode = False
|
|
pipe.use_diffusion_decoder = use_diffusion_decoder
|
|
pipe.vae = VideoDecoder("conv-vae")
|
|
pipe.diffusion_decoder = VideoDecoder("diff-vae")
|
|
pipe.audio_vae = AudioVae()
|
|
pipe.vocoder = torch.nn.Identity()
|
|
pipe.video_processor = processor
|
|
|
|
output = pipe._decode_output(
|
|
latents=torch.ones(1, 1, device="cuda", dtype=torch.bfloat16),
|
|
audio_latents=torch.ones(1, 1, device="cuda", dtype=torch.bfloat16),
|
|
output_type="np",
|
|
connector_prompt_embeds=torch.ones(1, 1, device="cuda", dtype=torch.bfloat16),
|
|
generator=None,
|
|
device=torch.device("cuda"),
|
|
decode_timestep=0.0,
|
|
decode_noise_scale=None,
|
|
prompt_batch_size=1,
|
|
)
|
|
|
|
prepared = output.output[0]
|
|
selected_decoder = "diff-vae" if use_diffusion_decoder else "conv-vae"
|
|
assert decode_calls == [selected_decoder]
|
|
assert prepared.is_cuda
|
|
assert prepared.shape == (1, 2, 3, 4, 3)
|
|
assert prepared.dtype == torch.uint8
|
|
assert prepared.is_contiguous()
|
|
delta = np.abs(prepared.cpu().numpy().astype(np.int16) - expected.astype(np.int16))
|
|
assert delta.max() <= 1
|