diff --git a/tests/diffusion/models/ltx2/test_ltx2_output_cuda.py b/tests/diffusion/models/ltx2/test_ltx2_output_cuda.py new file mode 100644 index 000000000..94c350627 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_output_cuda.py @@ -0,0 +1,136 @@ +# 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 diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index 9cafa22b6..c2fe7cbfa 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -10,8 +10,10 @@ from dataclasses import replace from types import SimpleNamespace from typing import Any +import numpy as np import pytest import torch +from diffusers.video_processor import VideoProcessor from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.models.ltx2 import ( @@ -65,7 +67,10 @@ from vllm_omni.diffusion.models.ltx2.ltx2_request import ( validate_ltx_checkpoint, validate_pipeline_request, ) -from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime +from vllm_omni.diffusion.models.ltx2.ltx2_runtime import ( + LTXRuntime, + _prepare_ltx2_video_output, +) from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( LTX2DistilledOneStagePipeline, LTX2I2VDMD2Pipeline, @@ -123,6 +128,103 @@ def _resolve_request_inputs_for_test( ) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("do_normalize", [True, False]) +def test_ltx_device_uint8_output_stays_within_one_level_of_legacy_rounding(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, + 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.shape == (2, 4, 5, 7, 3) + assert prepared.dtype == torch.uint8 + assert prepared.is_contiguous() + delta = np.abs(prepared.numpy().astype(np.int16) - expected.astype(np.int16)) + assert delta.max() <= (0 if dtype == torch.float32 else 1) + + +@pytest.mark.parametrize( + ("video", "error"), + [ + (torch.zeros(1, 3, 4, 5), "BCTHW"), + (torch.zeros(1, 4, 2, 4, 5), "RGB"), + (torch.zeros(1, 3, 2, 4, 5, dtype=torch.uint8), "floating-point"), + ], +) +def test_ltx_device_uint8_output_rejects_invalid_contract(video, error): + with pytest.raises(ValueError, match=error): + _prepare_ltx2_video_output(video, do_normalize=True) + + +def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]})) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + profiler_setup: dict[str, Any] = {} + + def capture_profiler_setup(_pipe, **kwargs): + profiler_setup.update(kwargs) + _pipe.enable_diffusion_pipeline_profiler = False + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", capture_profiler_setup) + + LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False)) + + assert "video_processor.postprocess_video" in profiler_setup["profiler_targets"] + assert "_prepare_video_output_for_transport" in profiler_setup["profiler_targets"] + + +def test_ltx_i2v_rewraps_replaced_video_processor_for_profiler(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_conditioning, ltx2_runtime + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]})) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + def enable_profiler(pipe, **_kwargs): + pipe.enable_diffusion_pipeline_profiler = True + + wrapped: list[tuple[Any, list[str]]] = [] + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", enable_profiler) + monkeypatch.setattr( + ltx2_conditioning, + "wrap_methods_by_paths", + lambda pipe, paths: wrapped.append((pipe.video_processor, paths)), + ) + + pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=True)) + + assert wrapped == [(pipe.video_processor, ["video_processor.postprocess_video"])] + + def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries(): from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( DistributedAutoencoderKLLTX2Video, @@ -2140,6 +2242,12 @@ class TestPostProcessFunction: assert "audio" in result assert result["audio_sample_rate"] == 48000 + prepared_video = torch.zeros(1, 4, 64, 64, 3, dtype=torch.uint8) + prepared_result = func((prepared_video, audio)) + assert isinstance(prepared_result["video"], np.ndarray) + assert prepared_result["video"].shape == (1, 4, 64, 64, 3) + assert prepared_result["video"].dtype == np.uint8 + def test_post_process_without_vocoder_config(self): """Post-process func should work without vocoder config (no audio_sample_rate key).""" import torch diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 6325df870..03e03c8f7 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -550,6 +550,13 @@ def get_ltx2_post_process_func(od_config: Any): if not (isinstance(output, tuple) and len(output) == 2): return output video, audio = output + if isinstance(video, torch.Tensor) and video.dtype == torch.uint8: + if video.ndim != 5 or video.shape[-1] != 3 or not video.is_contiguous(): + raise ValueError( + "LTX uint8 video output must be contiguous BTHWC RGB, " + f"got shape={tuple(video.shape)} stride={video.stride()}" + ) + video = video.detach().cpu().numpy() if isinstance(audio, torch.Tensor): audio = audio.detach().cpu() result: dict[str, Any] = {"video": video, "audio": audio} diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py index 8b4491e71..a0fe5f011 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py @@ -18,6 +18,8 @@ from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img impo from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import wrap_methods_by_paths + from . import ltx2_latents as latent_ops from .ltx2_guidance import euler_step_from_velocity @@ -392,6 +394,8 @@ class LTXI2VConditioningMixin: vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear", ) + if getattr(self, "enable_diffusion_pipeline_profiler", False): + wrap_methods_by_paths(self, ["video_processor.postprocess_video"]) @staticmethod def _resolve_single_prompt_image(raw_image: Any) -> Any: diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index fc589fcb8..0ffa09dcb 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -84,6 +84,45 @@ def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.T return vocoder(generated_mel.float()).to(input_dtype) +def _prepare_ltx2_video_output( + video: torch.Tensor, + *, + do_normalize: bool, +) -> torch.Tensor: + """Consume decoded BCTHW video and return contiguous uint8 BTHWC frames. + + Keep denormalization, scaling, and rounding in the decoder output dtype, + then convert directly to uint8. ``do_normalize=False`` means the decoder + output is expected to already be in ``[0, 1]``; the clamp remains a + defensive bound before byte conversion. Performing the reduction before + worker IPC cuts the video D2H payload from FP32 to uint8 without + materializing a full-size FP32 CUDA tensor. + """ + if video.ndim != 5: + raise ValueError(f"Expected decoded BCTHW video, got shape {tuple(video.shape)}") + if video.shape[1] != 3: + raise ValueError(f"Expected RGB BCTHW video, got shape {tuple(video.shape)}") + if not video.is_floating_point(): + raise ValueError(f"Expected floating-point decoded video, got {video.dtype}") + + video = video.detach() + if do_normalize: + # Match VaeImageProcessor.denormalize in the source dtype. These + # in-place operations have the same dtype-rounding boundaries while + # avoiding another full-size decoded-video allocation. + video.mul_(0.5).add_(0.5).clamp_(0.0, 1.0) + else: + video.clamp_(0.0, 1.0) + + # Source-dtype quantization can differ from the legacy FP32 presentation + # path by at most one uint8 level, while avoiding a full-size FP32 tensor. + video.mul_(255.0).round_() + return video.permute(0, 2, 3, 4, 1).to( + dtype=torch.uint8, + memory_format=torch.contiguous_format, + ) + + def _expand_per_prompt_decode_value( value: float | list[float], *, @@ -191,7 +230,16 @@ class LTXRuntime( initialize_pipeline_components(self, od_config) self._phase_adapter: LTXPhaseAdapterRuntime | None = build_ltx_phase_adapter(self) self.setup_diffusion_pipeline_profiler( - enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + profiler_targets=[ + "vae.encode", + "vae.decode", + "text_encoder.forward", + "video_processor.postprocess_video", + "_prepare_video_output_for_transport", + "audio_vae.decode", + "vocoder.forward", + ], + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler, ) def _forward_request( @@ -534,6 +582,12 @@ class LTXRuntime( ) return DiffusionOutput(output=output) + def _prepare_video_output_for_transport(self, video: torch.Tensor) -> torch.Tensor: + return _prepare_ltx2_video_output( + video, + do_normalize=bool(self.video_processor.config.do_normalize), + ) + def _decode_output( self, *, @@ -599,7 +653,10 @@ class LTXRuntime( ) if video.numel() > 0: - video = self.video_processor.postprocess_video(video, output_type=output_type) + if output_type == "np": + video = self._prepare_video_output_for_transport(video) + else: + video = self.video_processor.postprocess_video(video, output_type=output_type) generated_mel = self.audio_vae.decode(audio_latents.to(self.audio_vae.dtype), return_dict=False)[0] audio = _run_ltx_vocoder(self.vocoder, generated_mel) return self._make_output((video, audio))