[Bugfix][Multimodal] Fix PyNvVideoCodec video backend returning NCHW instead of NHWC (#51076)
Signed-off-by: Junpu Yu <davidyu@nvidia.com>
Junpu Yu committed
Aug 7, 2026 at 21:49 UTC
d3621c1eb3c4cb8d8a97e05d8c58d7f7571e34db
2 files changed
+39
-1
tests/multimodal/media/test_video.py
+20
@@ -20,6 +20,7 @@ from vllm.multimodal.video import (
20
PYNVVIDEOCODEC_VIDEO_BACKEND,
21
VIDEO_LOADER_REGISTRY,
22
VideoLoader,
23
+ _pynvvc_frames_to_nhwc,
24
)
25
26
from ..utils import cosine_similarity, create_video_from_image, normalize_image
@@ -498,3 +499,22 @@ class TestMergeKwargsGpuBackendPolicy:
499
500
def test_unknown_backend_not_treated_as_gpu(self):
501
assert not VIDEO_LOADER_REGISTRY.backend_requires_gpu("totally_unknown")
502
+
503
+
504
+@pytest.mark.parametrize("layout", ["nhwc", "nchw"])
505
+def test_pynvvc_frames_normalized_to_nhwc(layout: str):
506
+ """PyNvVideoCodec frame batches are normalized to NHWC regardless of the
507
+ per-frame layout the decoder emits (it has varied across versions), so the
508
+ HF video processors (which materialize a PIL image per frame) receive the
509
+ same NHWC shape as every other video backend."""
510
+ torch = pytest.importorskip("torch")
511
+
512
+ n, h, w, c = 4, 5, 6, 3
513
+ nhwc = torch.arange(n * h * w * c, dtype=torch.uint8).reshape(n, h, w, c)
514
+ frames = nhwc if layout == "nhwc" else nhwc.permute(0, 3, 1, 2).contiguous()
515
+
516
+ out = _pynvvc_frames_to_nhwc(frames)
517
+
518
+ assert out.shape == (n, h, w, c)
519
+ assert out.is_contiguous()
520
+ assert torch.equal(out, nhwc) # content preserved / correctly transposed
vllm/multimodal/video.py
+19
-1
@@ -637,6 +637,24 @@ class TorchCodecVideoBackendMixin:
637
return batch.data.numpy(), list(frame_indices)
638
639
640
+def _pynvvc_frames_to_nhwc(frames):
641
+ """Return a stacked PyNvVideoCodec frame batch as contiguous NHWC.
642
+
643
+ PyNvVideoCodec's per-frame layout has varied across versions (HWC vs CHW),
644
+ so detect the channel axis rather than assuming a fixed order. NHWC is the
645
+ layout the other video backends return and the HF video processors expect.
646
+
647
+ Args:
648
+ frames: A ``(N, ?, ?, ?)`` uint8 tensor in either NHWC or NCHW order.
649
+
650
+ Returns:
651
+ The same frames as a contiguous ``(N, H, W, C)`` tensor.
652
+ """
653
+ if frames.shape[-1] != 3 and frames.shape[-3] == 3:
654
+ frames = frames.permute(0, 2, 3, 1) # NCHW -> NHWC
655
+ return frames.contiguous()
656
+
657
+
658
class PyNvVideoCodecVideoBackendMixin:
659
"""PyNvVideoCodec utilities for GPU-backed frame decode."""
660
@@ -807,7 +825,7 @@ class PyNvVideoCodecVideoBackendMixin:
825
"PyNvVideoCodec returned frames with unexpected shape "
826
f"{tuple(device_frames.shape)}"
827
)
810
- device_frames = device_frames.permute(0, 3, 1, 2).contiguous()
828
+ device_frames = _pynvvc_frames_to_nhwc(device_frames)
829
host_frames = torch.empty(
830
device_frames.shape,
831
dtype=device_frames.dtype,