diff --git a/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py b/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py index f3e0b1ce3..45218db5a 100644 --- a/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py +++ b/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py @@ -148,7 +148,8 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested(): assert array.shape == (2, 4, 4, 3) video = post({"video": frames}, SimpleNamespace(output_type="pt")) assert set(video) == {"video"} - assert video["video"] is frames + assert isinstance(video["video"], np.ndarray) + np.testing.assert_array_equal(video["video"], frames.numpy()) image = post( {"image": torch.ones(4, 4, 3)}, SimpleNamespace(output_type="pt"), @@ -162,6 +163,40 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested(): assert latent_output["image"] is latent +@pytest.mark.parametrize("output_type", [None, "pt", "np"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_postprocess_preserves_video_range_for_serving(output_type: str | None, dtype: torch.dtype) -> None: + from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func + from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_frames + + # Include black and midtones: all-white frames cannot expose double normalization. + frames = torch.tensor([0.0, 0.25, 0.5, 0.75, 1.0], dtype=dtype) + frames = frames.reshape(5, 1, 1, 1).expand(5, 2, 4, 3).clone() + expected = frames.float().numpy().copy() + sampling = None if output_type is None else OmniDiffusionSamplingParams(output_type=output_type) + post = get_lingbot_video_post_process_func(None) + + processed = post({"video": frames}, sampling) + encoded_frames = np.stack(_coerce_video_to_frames(processed["video"])) + + assert set(processed) == {"video"} + assert encoded_frames.dtype == np.float32 + np.testing.assert_array_equal(encoded_frames, expected) + np.testing.assert_array_equal(frames.float().numpy(), expected) + + +@pytest.mark.parametrize("output_key", ["image", "video"]) +def test_postprocess_preserves_latent_payload(output_key: str) -> None: + from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func + + latents = torch.tensor([-2.0, 0.0, 2.0]).reshape(1, 3, 1, 1, 1) + sampling = OmniDiffusionSamplingParams(output_type="latent") + processed = get_lingbot_video_post_process_func(None)({output_key: latents}, sampling) + + assert set(processed) == {output_key} + assert processed[output_key] is latents + + def test_postprocess_preserves_image_contract_for_serving(): from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func @@ -181,7 +216,9 @@ def test_postprocess_preserves_image_contract_for_serving(): sampling_params=sampling, request_id="lingbot-image-contract-test", ) - raw_output = {"image": torch.ones(8, 12, 3)} + image = torch.zeros(8, 12, 3) + image[4:] = 1.0 + raw_output = {"image": image} processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling) normalized = normalize_diffusion_postprocess_output(processed) @@ -199,6 +236,8 @@ def test_postprocess_preserves_image_contract_for_serving(): endpoint_images = _extract_images_from_result(result) assert len(endpoint_images) == 1 assert endpoint_images[0].size == (12, 8) + assert endpoint_images[0].getpixel((0, 0)) == (0, 0, 0) + assert endpoint_images[0].getpixel((0, 7)) == (255, 255, 255) def test_forward_resolves_t2v_sampling_and_flow_shift_alias(): diff --git a/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py index 51687a241..50fb33a33 100644 --- a/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py +++ b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py @@ -288,13 +288,13 @@ def get_lingbot_video_post_process_func(od_config: OmniDiffusionConfig): output_key = output_keys[0] frames = frames[output_key] output_type = getattr(sampling_params, "output_type", None) or "pt" - # The image serving path currently accepts PIL images or NumPy arrays, - # while LingBot decodes T2I outputs to tensors. Keep that compatibility - # conversion without discarding the model's image/video payload key. + # Decoded LingBot pixels are already in [0, 1]. The video API treats + # floating tensors as [-1, 1], so use its NumPy path to avoid normalizing + # twice. Images also need NumPy for serving; keep latent tensors intact. if ( isinstance(frames, torch.Tensor) and output_type != "latent" - and (output_type == "np" or output_key == "image") + and (output_type == "np" or output_key in {"image", "video"}) ): frames = frames.float().cpu().numpy() return {output_key: frames} if output_key is not None else frames