diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 336be1503..e4767cabb 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -57,7 +57,8 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): return torch.zeros(1), torch.zeros(1) pipeline.diffuse = diffuse - pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + decoded_audio = torch.zeros(1) + pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), decoded_audio)) sampling = OmniDiffusionSamplingParams( quality="high", width=1344, @@ -86,7 +87,9 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): num_inference_steps=50, extra_args={"task": "t2va", "aspect_ratio": "16:9"}, ) - assert output.output == pipeline.decode.return_value + assert output.output[0].shape == (1, 1, 1, 1, 3) + assert output.output[0].dtype == torch.uint8 + assert output.output[1] is decoded_audio def test_pipeline_import_registry_and_component_discovery(): @@ -505,6 +508,26 @@ def test_joint_postprocess_is_multiprocessing_picklable(): assert result["fps"] == 24 +def test_joint_video_output_is_quantized_before_transfer(): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _minimax_h3_post_process, + _prepare_minimax_h3_video_output, + ) + + video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5) + expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1) + prepared = _prepare_minimax_h3_video_output(video) + + assert prepared.dtype == torch.uint8 + assert prepared.is_contiguous() + assert prepared.shape == (2, 2, 4, 5, 3) + torch.testing.assert_close(prepared, expected, rtol=0, atol=0) + + result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6))) + assert result["video"][0].dtype == np.uint8 + np.testing.assert_array_equal(result["video"][0], expected[0].numpy()) + + def test_cfg_parallel_is_rejected_for_distilled_checkpoint(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline @@ -627,7 +650,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition): return torch.zeros(1), torch.zeros(1) pipeline.diffuse = diffuse - pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), torch.zeros(1))) return pipeline diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index be27c4c97..03ddb8786 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -223,6 +223,19 @@ def test_channel_first_video_tensor_is_converted_to_channel_last(): np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy()) +def test_contiguous_uint8_video_reaches_mux_without_conversion(monkeypatch): + video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3) + + def fail_prepare(*args, **kwargs): + raise AssertionError("ready-to-encode uint8 video must not be normalized") + + monkeypatch.setattr(video_api_utils, "_prepare_video_frames", fail_prepare) + + frames = video_api_utils._coerce_video_to_uint8_frames(video) + + assert frames is video + + def test_channel_first_video_tensor_uses_direct_planar_mux(monkeypatch): calls = [] diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 5117dc02d..d0f63fe8e 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -361,7 +361,10 @@ def _minimax_h3_post_process(output, output_type: str = "np"): if output_type == "latent": return output if output_type == "np": - video = video.detach().float().cpu().permute(0, 2, 3, 4, 1).clamp(0, 1).numpy() + if video.dtype == torch.uint8 and video.ndim == 5 and video.shape[-1] == 3: + video = video.detach().cpu().numpy() + else: + video = video.detach().float().cpu().permute(0, 2, 3, 4, 1).clamp(0, 1).numpy() audio = audio.detach().float().cpu().numpy() video = [sample for sample in video] return { @@ -372,6 +375,16 @@ def _minimax_h3_post_process(output, output_type: str = "np"): } +def _prepare_minimax_h3_video_output(video: torch.Tensor) -> torch.Tensor: + """Quantize decoded frames in place before worker-to-engine transfer.""" + video = video.detach().float() + video.clamp_(0, 1).mul_(255).round_() + return video.permute(0, 2, 3, 4, 1).to( + dtype=torch.uint8, + memory_format=torch.contiguous_format, + ) + + def _register_dlo_component_cache(cache: BoundedAllocatorCache, *components: Any) -> None: for component in components: if component is not None: @@ -2406,10 +2419,10 @@ class MiniMaxH3Pipeline( height=context["height"], width=context["width"], ) - videos.append(video) + videos.append(_prepare_minimax_h3_video_output(video)) audios.append(audio) - video = torch.cat(videos, dim=0) - audio = torch.cat(audios, dim=0) + video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0) + audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), @@ -2711,6 +2724,7 @@ class MiniMaxH3Pipeline( height=shape["height"], width=shape["width"], ) + video = _prepare_minimax_h3_video_output(video) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index dcaf2451b..b5b64712c 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -553,6 +553,14 @@ def _coerce_prepared_video_to_uint8_frames( def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray: """Convert a video payload into contiguous uint8 frames shaped (F, H, W, 3).""" + if ( + isinstance(video, np.ndarray) + and video.dtype == np.uint8 + and video.ndim == 4 + and video.shape[-1] == 3 + and video.flags.c_contiguous + ): + return video frames, frame_shape, common_dtype = _prepare_video_frames(video) return _coerce_prepared_video_to_uint8_frames(frames, frame_shape, common_dtype)