From 185e3d6df60276166349e69fdc1171fd0e50c79f Mon Sep 17 00:00:00 2001 From: Librium <129784133+pxljs@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:05:57 +0800 Subject: [PATCH] [Bugfix] Preserve channel-last video tensor layout (#5418) Signed-off-by: pxljs <1621352782@qq.com> --- .../openai_api/test_video_api_utils.py | 19 +++++++++++++++++++ .../entrypoints/openai/video_api_utils.py | 5 ++--- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 722819682..652adef4d 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -48,6 +48,25 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch): assert mux_calls[0]["audio"] is None +@pytest.mark.parametrize("frame_count", [3, 4]) +def test_channel_last_video_tensor_preserves_channel_sized_frame_count(frame_count): + video = torch.arange(frame_count * 2 * 5 * 3, dtype=torch.uint8).reshape(frame_count, 2, 5, 3) + + frames = video_api_utils._coerce_video_to_uint8_frames(video) + + assert frames.shape == (frame_count, 2, 5, 3) + np.testing.assert_array_equal(frames, video.numpy()) + + +def test_channel_first_video_tensor_is_converted_to_channel_last(): + video = torch.arange(3 * 5 * 2 * 6, dtype=torch.uint8).reshape(3, 5, 2, 6) + + frames = video_api_utils._coerce_video_to_uint8_frames(video) + + assert frames.shape == (5, 2, 6, 3) + np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy()) + + def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch): muxers = [] diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index aa6479189..f0d2a59f2 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -308,9 +308,6 @@ def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray: video_tensor = video_tensor.detach().cpu() if video_tensor.dim() == 5: raise ValueError("Batched video tensors are not supported for single-video encoding.") - elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4): - # [C, F, H, W] -> [F, H, W, C] - video_tensor = video_tensor.permute(1, 2, 3, 0) if video_tensor.is_floating_point(): # Cast to float32 first: bf16 (e.g. SANA-WM's refiner output) has no @@ -329,6 +326,8 @@ def _normalize_single_video_array(video_array: np.ndarray) -> np.ndarray: if video_array.ndim == 4: # Convert channel-first layouts to channel-last + # Prefer an explicit channel-last dimension for ambiguous 3/4-frame + # videos before interpreting a leading dimension as channels. if video_array.shape[0] in (3, 4) and video_array.shape[-1] not in (3, 4): video_array = np.transpose(video_array, (1, 2, 3, 0)) elif video_array.shape[1] in (3, 4) and video_array.shape[-1] not in (3, 4):