[Multimodal] Reorganize video decoder backends (#49155)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
Isotr0py
2026-08-18 15:06:18 +00:00
committed by GitHub
parent 689be2bcd3
commit 3bb9c18f0c
12 changed files with 1416 additions and 1097 deletions
+7 -3
View File
@@ -21,8 +21,10 @@ from vllm.multimodal.media import ImageMediaIO, VideoMediaIO
from vllm.multimodal.video import (
PYNVVIDEOCODEC_VIDEO_BACKEND,
VIDEO_LOADER_REGISTRY,
PyNvVideoCodecVideoBackend,
VideoLoader,
)
from vllm.multimodal.video_decoders.pynvvideocodec import (
PyNvVideoCodecVideoBackendMixin,
_pynvvc_frames_to_nhwc,
)
@@ -381,13 +383,15 @@ def test_pynvvideocodec_unrelated_error_propagates(
raise original_error
monkeypatch.setattr(
PyNvVideoCodecVideoBackend,
PyNvVideoCodecVideoBackendMixin,
"_read_source_metadata",
classmethod(raise_unrelated_error),
)
with pytest.raises(RuntimeError) as exc_info:
PyNvVideoCodecVideoBackend.decode_frames_pynvvideocodec(b"video", None)
PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec(
None, b"video", None
)
assert exc_info.value is original_error
+7 -28
View File
@@ -6,7 +6,6 @@ import time
import pytest
import vllm.config.multimodal as multimodal_config_module
from vllm.config.multimodal import MultiModalConfig
from vllm.multimodal.gpu_ipc_memory import (
MultiModalGPUMemoryPool,
@@ -15,10 +14,10 @@ from vllm.multimodal.gpu_ipc_memory import (
reserve_mm_ipc_gpu_memory,
set_mm_gpu_ipc_pool,
)
from vllm.multimodal.video import (
from vllm.multimodal.video_decoders import PYNVVIDEOCODEC_VIDEO_BACKEND
from vllm.multimodal.video_decoders.pynvvideocodec import (
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
from vllm.utils.mem_constants import GiB_bytes
@@ -186,11 +185,7 @@ def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(
monkeypatch: pytest.MonkeyPatch,
video_backend: str | None,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
"opencv",
)
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")
mm_config = _mm_config(
mm_ipc_gpu_memory_gb=0.25,
video_backend=video_backend,
@@ -202,11 +197,7 @@ def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(
def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
"opencv",
)
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")
mm_config = _mm_config(
mm_ipc_gpu_memory_gb=0.25,
video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,
@@ -221,11 +212,7 @@ def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(
def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", PYNVVIDEOCODEC_VIDEO_BACKEND)
available_bytes = 4 * GiB_bytes
assert reserve_mm_ipc_gpu_memory(available_bytes, _mm_config()) == (
@@ -236,11 +223,7 @@ def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(
def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", PYNVVIDEOCODEC_VIDEO_BACKEND)
available_bytes = 8 * GiB_bytes
assert reserve_mm_ipc_gpu_memory(
@@ -253,11 +236,7 @@ def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers(
def test_reserve_mm_ipc_gpu_memory_uses_configured_hw_decoders(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
"opencv",
)
monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")
available_bytes = 4 * GiB_bytes
mm_config = _mm_config(
video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,
+155 -18
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import itertools
import subprocess
import sys
import threading
from contextlib import ExitStack, contextmanager
@@ -16,24 +17,26 @@ from transformers.video_utils import VideoMetadata
from vllm.assets.base import get_vllm_public_assets
from vllm.models.minimax_m3.common.mm_preprocess import MiniMaxM3VideoBackend
from vllm.multimodal.video import (
PYNVVIDEOCODEC_DECODER_CACHE_SIZE,
PYNVVIDEOCODEC_VIDEO_BACKEND,
VIDEO_LOADER_REGISTRY,
DynamicVideoBackend,
GLM46VVideoBackend,
Molmo2VideoBackend,
PyNvVideoCodecDecoderSlot,
PyNvVideoCodecVideoBackend,
PyNvVideoCodecVideoBackendMixin,
Qwen2VLVideoBackend,
Qwen3VLVideoBackend,
VideoBackend,
VideoLoader,
VideoSourceMetadata,
VideoTargetMetadata,
_pynv_decoder_pool,
get_video_loader_backend_for_processor,
)
from vllm.multimodal.video_decoders import decode_video, resolve_video_backend_kwargs
from vllm.multimodal.video_decoders.pynvvideocodec import (
PYNVVIDEOCODEC_DECODER_CACHE_SIZE,
PyNvVideoCodecDecoderSlot,
PyNvVideoCodecVideoBackendMixin,
_pynv_decoder_pool,
)
from vllm.platforms import current_platform
from vllm.transformers_utils.processor import get_video_processor_cls_name_from_config
@@ -99,6 +102,126 @@ def test_video_loader_type_doesnt_exist():
VIDEO_LOADER_REGISTRY.load("non_existing_video_loader")
def test_video_decoder_backends_are_lazy_imported():
code = """
import sys
import vllm.multimodal.video # noqa: F401
backend_modules = {
f"vllm.multimodal.video_decoders.{backend}"
for backend in ("opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream")
}
loaded = sorted(backend_modules & sys.modules.keys())
assert not loaded, loaded
"""
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
)
assert result.returncode == 0, result.stderr
@pytest.mark.parametrize(
"backend",
["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"],
)
def test_decode_video_imports_only_selected_backend(
backend: str,
monkeypatch: pytest.MonkeyPatch,
):
imports = []
decoded = object()
def fake_decoder(*args, **kwargs):
return decoded
class FakeBackendModule:
pass
setattr(FakeBackendModule, f"decode_{backend}", fake_decoder)
def fake_import_module(name: str, package: str):
imports.append((name, package))
return FakeBackendModule
monkeypatch.setattr(
"vllm.multimodal.video_decoders.import_module", fake_import_module
)
result = decode_video(
backend,
loader_cls=None,
data=b"",
target=VideoTargetMetadata(-1, -1, 300),
sampling_kwargs={},
backend_kwargs={},
frame_recovery=False,
)
assert result is decoded
assert imports == [(f".{backend}", "vllm.multimodal.video_decoders")]
@pytest.mark.parametrize(
("backend", "kwargs", "expected_sampling", "expected_backend"),
[
(
"torchcodec",
{"min_frames": 4, "num_ffmpeg_threads": 2, "seek_mode": "approximate"},
{"min_frames": 4},
{"num_ffmpeg_threads": 2, "seek_mode": "approximate"},
),
(
"deepstream",
{"max_frames": 16, "pool_size": 3, "timeout_sec": 10.0},
{"max_frames": 16},
{"pool_size": 3, "timeout_sec": 10.0},
),
],
)
def test_video_backend_kwargs_are_separated_from_sampling_kwargs(
backend: str,
kwargs: dict,
expected_sampling: dict,
expected_backend: dict,
):
original_kwargs = dict(kwargs)
sampling_kwargs, backend_kwargs = resolve_video_backend_kwargs(backend, kwargs)
assert sampling_kwargs == expected_sampling
assert backend_kwargs == expected_backend
assert kwargs == original_kwargs
def test_video_backend_rejects_options_for_another_decoder():
with pytest.raises(
ValueError, match="num_ffmpeg_threads is not supported by the 'pyav' backend"
):
resolve_video_backend_kwargs("pyav", {"num_ffmpeg_threads": 2})
@pytest.mark.parametrize(
("backend", "error"),
[
("pyav", AssertionError),
(PYNVVIDEOCODEC_VIDEO_BACKEND, ValueError),
],
)
def test_video_decoder_spec_validates_frame_recovery(
backend: str, error: type[Exception]
):
with pytest.raises(error, match="frame_recovery is not supported"):
decode_video(
backend,
loader_cls=None,
data=b"",
target=VideoTargetMetadata(-1, -1, 300),
sampling_kwargs={},
backend_kwargs={},
frame_recovery=True,
)
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
def test_pynvvideocodec_backend_accounts_raw_decoded_frames(
monkeypatch: pytest.MonkeyPatch,
@@ -145,7 +268,9 @@ def test_pynvvideocodec_backend_accounts_raw_decoded_frames(
"vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool
)
monkeypatch.setattr(
PyNvVideoCodecVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode)
PyNvVideoCodecVideoBackendMixin,
"_decode_to_pinned_host",
classmethod(fake_decode),
)
loader = VIDEO_LOADER_REGISTRY.load(PYNVVIDEOCODEC_VIDEO_BACKEND)
@@ -205,7 +330,9 @@ def test_pynvvideocodec_codec_uses_dynamic_sampling_strategy(
"vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool
)
monkeypatch.setattr(
DynamicVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode)
PyNvVideoCodecVideoBackendMixin,
"_decode_to_pinned_host",
classmethod(fake_decode),
)
loader = VIDEO_LOADER_REGISTRY.load("opencv_dynamic")
@@ -273,7 +400,7 @@ def test_pynvvideocodec_decoder_slots_are_bounded(
create_count = 0
with _fresh_decoder_pool():
PyNvVideoCodecVideoBackend._configure_decoder_slots(hw_decoders)
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(hw_decoders)
def fake_create_slot(cls):
nonlocal create_count
@@ -281,7 +408,7 @@ def test_pynvvideocodec_decoder_slots_are_bounded(
return FakeSlot()
monkeypatch.setattr(
PyNvVideoCodecVideoBackend,
PyNvVideoCodecVideoBackendMixin,
"_create_decoder_slot",
classmethod(fake_create_slot),
)
@@ -291,12 +418,16 @@ def test_pynvvideocodec_decoder_slots_are_bounded(
with ExitStack() as stack:
retained_slots = [
stack.enter_context(PyNvVideoCodecVideoBackend._borrow_decoder_slot())
stack.enter_context(
PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot()
)
for _ in range(hw_decoders)
]
def borrow_extra_slot():
with PyNvVideoCodecVideoBackend._borrow_decoder_slot() as extra_slot:
with (
PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot()
) as extra_slot:
seen_slots.append(extra_slot)
borrowed.set()
@@ -317,11 +448,11 @@ def test_pynvvideocodec_decoder_slots_are_configured_once(
):
monkeypatch.setattr(_pynv_decoder_pool, "max_slots", None)
PyNvVideoCodecVideoBackend._configure_decoder_slots(2)
PyNvVideoCodecVideoBackend._configure_decoder_slots(2)
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(2)
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(2)
with pytest.raises(RuntimeError, match="already configured as 2, got 3"):
PyNvVideoCodecVideoBackend._configure_decoder_slots(3)
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(3)
def test_pynvvideocodec_failed_rebuild_invalidates_decoder_slot():
@@ -367,7 +498,7 @@ def test_pynvvideocodec_failed_rebuild_invalidates_decoder_slot():
with (
pytest.raises(RuntimeError, match="construct failed"),
PyNvVideoCodecVideoBackend._borrow_decoder_slot() as borrowed,
PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot() as borrowed,
):
assert borrowed is slot
borrowed.get_decoder(
@@ -456,6 +587,12 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():
counters via ClassVar shadowing.
"""
class MixinSubclassA(PyNvVideoCodecVideoBackendMixin):
pass
class MixinSubclassB(PyNvVideoCodecVideoBackendMixin):
pass
class FakeSlot:
pass
@@ -475,8 +612,8 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():
)
try:
with ExitStack() as stack:
stack.enter_context(VideoBackend._borrow_decoder_slot())
stack.enter_context(Qwen3VLVideoBackend._borrow_decoder_slot())
stack.enter_context(MixinSubclassA._borrow_decoder_slot())
stack.enter_context(MixinSubclassB._borrow_decoder_slot())
assert pool.active == 2
blocked = threading.Event()
@@ -484,7 +621,7 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():
def try_borrow():
blocked.set()
with Qwen2VLVideoBackend._borrow_decoder_slot():
with MixinSubclassB._borrow_decoder_slot():
acquired.set()
t = threading.Thread(target=try_borrow)
+5 -3
View File
@@ -188,11 +188,13 @@ def reserve_mm_ipc_gpu_memory(
return available_kv_cache_memory_bytes
from vllm import envs
from vllm.multimodal.video import (
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
from vllm.multimodal.video_decoders import (
PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
from vllm.multimodal.video_decoders.pynvvideocodec import (
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
validate_pynvvideocodec_hw_decoders,
)
+66 -1045
View File
File diff suppressed because it is too large Load Diff
+110
View File
@@ -0,0 +1,110 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from importlib import import_module
from typing import Any, Literal
from .base import (
PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
PYNVVIDEOCODEC_VIDEO_BACKEND,
VideoSourceMetadata,
VideoTargetMetadata,
check_frame_pixel_limit,
)
VideoDecoderBackend = Literal[
"opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"
]
_BACKEND_OPTION_DEFAULTS: dict[str, dict[str, Any]] = {
"opencv": {},
"pyav": {},
"torchcodec": {
"num_ffmpeg_threads": 0,
"seek_mode": "exact",
},
PYNVVIDEOCODEC_VIDEO_BACKEND: {
"hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
},
"deepstream": {
"pool_size": None,
"timeout_sec": 120.0,
},
}
def _get_backend_option_defaults(backend: str) -> dict[str, Any]:
try:
return _BACKEND_OPTION_DEFAULTS[backend]
except KeyError:
valid = ", ".join(repr(name) for name in _BACKEND_OPTION_DEFAULTS)
raise ValueError(
f"Unknown video codec backend {backend!r}; valid options: {valid}."
) from None
def resolve_video_backend_kwargs(
backend: str,
kwargs: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, Any]]:
"""Split frame-sampling kwargs from options owned by a decoder backend."""
defaults = _get_backend_option_defaults(backend)
sampling_kwargs = dict(kwargs)
backend_kwargs = dict(defaults)
backend_option_names = {
name for options in _BACKEND_OPTION_DEFAULTS.values() for name in options
}
misplaced = (sampling_kwargs.keys() & backend_option_names) - defaults.keys()
if misplaced:
names = ", ".join(sorted(misplaced))
raise ValueError(f"{names} is not supported by the {backend!r} backend")
for name in defaults:
if name in sampling_kwargs:
backend_kwargs[name] = sampling_kwargs.pop(name)
return sampling_kwargs, backend_kwargs
def decode_video(
backend: str,
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict[str, Any],
backend_kwargs: dict[str, Any],
*,
frame_recovery: bool,
):
"""Decode a sampled video, importing only the selected backend."""
_get_backend_option_defaults(backend)
if frame_recovery and backend != "opencv":
error = (
ValueError if backend == PYNVVIDEOCODEC_VIDEO_BACKEND else AssertionError
)
raise error(f"frame_recovery is not supported by the {backend!r} backend")
decoder_kwargs = dict(backend_kwargs)
if backend == "opencv":
decoder_kwargs["frame_recovery"] = frame_recovery
module = import_module(f".{backend}", __name__)
decoder = getattr(module, f"decode_{backend}")
return decoder(
loader_cls,
data,
target,
sampling_kwargs,
**decoder_kwargs,
)
__all__ = [
"PYNVVIDEOCODEC_DEFAULT_HW_DECODERS",
"PYNVVIDEOCODEC_VIDEO_BACKEND",
"VideoDecoderBackend",
"VideoSourceMetadata",
"VideoTargetMetadata",
"check_frame_pixel_limit",
"decode_video",
"resolve_video_backend_kwargs",
]
+37
View File
@@ -0,0 +1,37 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Literal, NamedTuple
from vllm import envs
PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec"
PYNVVIDEOCODEC_DEFAULT_HW_DECODERS = 2
class VideoTargetMetadata(NamedTuple):
"""Metadata describing the requested video sample."""
num_frames: int
fps: float
max_duration: float
class VideoSourceMetadata(NamedTuple):
"""Metadata describing the encoded video source."""
total_frames_num: int
original_fps: float
duration: float
def check_frame_pixel_limit(width: int, height: int) -> None:
"""Reject video frames exceeding ``VLLM_MAX_IMAGE_PIXELS``."""
max_pixels = envs.VLLM_MAX_IMAGE_PIXELS
if max_pixels > 0 and width * height > max_pixels:
raise ValueError(
f"Video frame dimensions {width}x{height} "
f"({width * height} pixels) exceed the maximum of "
f"{max_pixels} pixels. Set VLLM_MAX_IMAGE_PIXELS to "
f"increase this limit."
)
@@ -0,0 +1,151 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
from typing import Any, ClassVar
import numpy.typing as npt
from vllm.logger import init_logger
from .base import VideoSourceMetadata, VideoTargetMetadata, check_frame_pixel_limit
logger = init_logger(__name__)
def decode_deepstream(
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict,
*,
pool_size: int | None = None,
timeout_sec: float = 120.0,
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
from nvidia.deepstream_videodecode import probe_metadata
total_frames, original_fps, duration, _width, _height, codec = probe_metadata(data)
check_frame_pixel_limit(_width, _height)
source = loader_cls._prepare_source(
VideoSourceMetadata(total_frames, original_fps, duration)
)
frame_idx = loader_cls.compute_frames_index_to_sample(
source=source, target=target, **sampling_kwargs
)
frames, valid = DeepStreamVideoBackendMixin.decode_indices(
data,
frame_idx,
source,
codec=codec,
pool_size=pool_size,
timeout_sec=timeout_sec,
)
return frames, source, frame_idx, valid
class DeepStreamVideoBackendMixin:
"""NVIDIA DeepStream (NVDEC) GPU-decode codec utilities.
Decoding runs on a shared pool of daemon threads inside one CUDA
context (see the ``nvidia-deepstream-videodecode-cu13`` package). The
container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no
local file path is required — HTTP and base64 sources decode identically
to local files.
Like the OpenCV/PyAV mixins, this provides only the codec layer.
Frame *selection* lives in the loader's
``compute_frames_index_to_sample`` and arrives here as an explicit
list of frame indices.
"""
# Process-wide lazy decode pool, shared across all DeepStream backends.
_pool: ClassVar[Any] = None
_pool_lock: ClassVar[Any] = None
@classmethod
def _get_pool(cls, pool_size: int | None = None):
"""Lazy-initialize the shared decode pool on first use.
``pool_size`` (number of decode worker threads) comes from
``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it
defaults to the existing ``VLLM_MEDIA_LOADING_THREAD_COUNT`` so no
DeepStream-specific env var is needed. The pool is a process-wide
singleton, so the first decode's value wins.
"""
if cls._pool is not None:
return cls._pool
if cls._pool_lock is None:
cls._pool_lock = threading.Lock()
with cls._pool_lock:
if cls._pool is not None:
return cls._pool
import os
from nvidia.deepstream_videodecode import DecodePool
if pool_size is None:
pool_size = int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", 8))
pool_size = max(1, min(int(pool_size), 16))
logger.info(
"[DeepStream] initializing decode pool with %d workers",
pool_size,
)
cls._pool = DecodePool(num_workers=pool_size)
return cls._pool
@classmethod
def decode_indices(
cls,
data: bytes,
frame_indices: list[int],
source: VideoSourceMetadata,
codec: str = "",
pool_size: int | None = None,
timeout_sec: float = 120.0,
) -> tuple[npt.NDArray, list[int]]:
"""Decode the requested frame indices from raw container bytes.
The whole stream is decoded; the pool keeps exactly the frames whose
decode-order index is in ``frame_indices`` (1:1, frame-exact) and
sends EOS once the last one is matched.
``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC
session warm across same-codec streams and rebuild only on a codec
change. Frames are returned as a CPU NHWC uint8 array so the
upstream multimodal parser sees the same shape as the other
backends.
"""
if not frame_indices:
raise ValueError("DeepStream backend received no frame indices")
result = cls._get_pool(pool_size).decode(
data,
target_indices=frame_indices,
codec=codec,
max_frames=len(frame_indices),
timeout_sec=timeout_sec,
)
if result.error:
raise ValueError(f"DeepStream decode failed: {result.error}")
if result.frames is None or result.n_kept == 0:
raise ValueError("DeepStream decode produced no frames")
valid = frame_indices[: result.n_kept]
# GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps
# the array shape identical to the OpenCV/PyAV backends. Copy into
# PINNED host memory (reused across calls by PyTorch's pinned caching
# allocator) so the D2H runs at full PCIe bandwidth (~13 GB/s) rather
# than the ~1 GB/s pageable path that plain ``.cpu()`` takes — ~12x
# faster for a 1080p x8 frame batch (~46ms -> ~4ms). ``numpy()`` keeps
# the pinned tensor alive via the array's base.
import torch
gpu = result.frames
if gpu.is_cuda:
host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True)
host.copy_(gpu, non_blocking=True)
torch.cuda.current_stream().synchronize()
arr = host.numpy()
else:
arr = gpu.numpy()
return arr, valid
+295
View File
@@ -0,0 +1,295 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from io import BytesIO
import numpy as np
import numpy.typing as npt
from vllm.logger import init_logger
from vllm.utils.import_utils import PlaceholderModule
from .base import (
VideoSourceMetadata,
VideoTargetMetadata,
check_frame_pixel_limit,
)
try:
import cv2
import cv2.videoio_registry as vr
except ImportError:
cv2 = PlaceholderModule("cv2")
vr = PlaceholderModule("cv2").placeholder_attr("videoio_registry")
logger = init_logger(__name__)
def decode_opencv(
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict,
*,
frame_recovery: bool = False,
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
cap = OpenCVVideoBackendMixin.open_video_capture(data)
check_frame_pixel_limit(
int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
)
source = loader_cls._prepare_source(OpenCVVideoBackendMixin.get_video_metadata(cap))
frame_idx = loader_cls.compute_frames_index_to_sample(
source=source, target=target, **sampling_kwargs
)
frames, valid = OpenCVVideoBackendMixin.read_frames(
cap,
frame_idx,
total_frames_num=source.total_frames_num,
frame_recovery=frame_recovery,
)
return frames, source, frame_idx, valid
class OpenCVVideoBackendMixin:
@staticmethod
def get_cv2_video_api():
api_pref = None
for backend in vr.getStreamBufferedBackends():
if not vr.hasBackend(backend):
continue
if not vr.isBackendBuiltIn(backend):
_, abi, api = vr.getStreamBufferedBackendPluginVersion(backend)
if abi < 1 or (abi == 1 and api < 2):
continue
api_pref = backend
break
return api_pref
@classmethod
def open_video_capture(cls, data: bytes) -> "cv2.VideoCapture":
backend = cls.get_cv2_video_api()
cap = cv2.VideoCapture(BytesIO(data), backend, [])
if not cap.isOpened():
raise ValueError("Could not open video stream")
return cap
@staticmethod
def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata:
total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
original_fps = cap.get(cv2.CAP_PROP_FPS)
duration = total_frames_num / original_fps if original_fps > 0 else 0
return VideoSourceMetadata(
total_frames_num=total_frames_num,
original_fps=original_fps,
duration=duration,
)
@classmethod
def _can_use_for_recovery(
cls,
idx: int,
failed_frames: list[int],
next_target_map: dict[int, int],
total_frames: int,
) -> bool:
"""Check if current frame can recover the oldest failed frame."""
if not failed_frames:
return False
oldest_failed = failed_frames[0]
limit = next_target_map.get(oldest_failed, total_frames)
return idx < limit
@classmethod
def _read_frames_with_recovery(
cls,
cap: "cv2.VideoCapture",
frame_indices: list[int],
total_frames: int,
) -> tuple[npt.NDArray, list[int], dict[int, int]]:
"""
Read frames with dynamic window forward-scan recovery.
When a target frame fails to load, the next successfully grabbed
frame (before the next target frame) will be used to recover it.
Args:
cap: OpenCV VideoCapture object
frame_indices: Sorted list of target frame indices to load
total_frames: Total number of frames in the video
Returns:
Tuple of (frames_array, valid_frame_indices, recovered_map)
- frames_array: Array of loaded frames
- valid_frame_indices: List of frame indices that were loaded
- recovered_map: Dict mapping recovered_idx -> source_idx
"""
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
assert width > 0 and height > 0, (
f"Invalid video frame size: width={width}, height={height}"
)
frame_idx_set = set(frame_indices)
max_frame_idx = frame_indices[-1] if frame_indices else 0
# Build map: target_idx -> next_target_idx (for recovery window)
next_target_map: dict[int, int] = {}
for k in range(len(frame_indices) - 1):
next_target_map[frame_indices[k]] = frame_indices[k + 1]
next_target_map[frame_indices[-1]] = total_frames
frames_list: list[npt.NDArray] = []
valid_frame_indices: list[int] = []
failed_frames_idx: list[int] = []
recovered_map: dict[int, int] = {}
i = 0
for idx in range(max_frame_idx + 1):
is_target_frame = idx in frame_idx_set
# Attempt to grab the current frame
ok = cap.grab()
if not ok:
if is_target_frame:
logger.debug(
"Failed to grab frame %d during video loading.",
idx,
)
failed_frames_idx.append(idx)
continue
# Check if we should retrieve: target frame OR can recover a failed one
can_recover = cls._can_use_for_recovery(
idx, failed_frames_idx, next_target_map, total_frames
)
if is_target_frame or can_recover:
ret, frame = cap.retrieve()
if ret and frame is not None and frame.size > 0:
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frames_list.append(rgb_frame)
valid_frame_indices.append(idx)
i += 1
if can_recover:
recovered_idx = failed_frames_idx.pop(0)
recovered_map[recovered_idx] = idx
logger.info(
"Recovered frame %d using frame %d (delay: %d)",
recovered_idx,
idx,
idx - recovered_idx,
)
elif is_target_frame:
logger.debug(
"Failed to retrieve frame %d during video loading.",
idx,
)
failed_frames_idx.append(idx)
# Log any remaining failed frames
for failed_idx in failed_frames_idx:
logger.debug(
"Frame %d could not be recovered (end of video).",
failed_idx,
)
# Stack frames
if frames_list:
frames = np.stack(frames_list)
else:
frames = np.empty((0, height, width, 3), dtype=np.uint8)
return frames, valid_frame_indices, recovered_map
@classmethod
def _read_frames_no_recovery(
cls,
cap,
frame_indices: set[int],
max_frame_idx: int,
) -> tuple[npt.NDArray, list[int]]:
num_expected_frames = len(frame_indices)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
frames = np.empty((num_expected_frames, height, width, 3), dtype=np.uint8)
i = 0
valid_frame_indices = []
for idx in range(max_frame_idx + 1):
ok = cap.grab()
if not ok:
# Frame is broken/unreadable, skip it
if idx in frame_indices:
logger.debug(
"Failed to grab frame %d during video loading. "
"This frame will be skipped.",
idx,
)
continue
if idx in frame_indices:
ret, frame = cap.retrieve()
if ret:
frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
valid_frame_indices.append(idx)
i += 1
else:
# retrieve() failed even though grab() succeeded
logger.debug(
"Failed to retrieve frame %d during video loading. "
"This frame will be skipped.",
idx,
)
valid_num_frames = len(valid_frame_indices)
if valid_num_frames < num_expected_frames:
logger.warning(
"Video loading completed with %d broken/unreadable frames. "
"Expected %d frames but only loaded %d frames.",
num_expected_frames - valid_num_frames,
num_expected_frames,
valid_num_frames,
)
return frames[:valid_num_frames], valid_frame_indices
@classmethod
def read_frames(
cls,
cap: "cv2.VideoCapture",
frame_idx: list[int],
total_frames_num: int,
*,
frame_recovery: bool = False,
) -> tuple[npt.NDArray, list[int]]:
if frame_recovery:
num_frames_to_sample = len(frame_idx)
frames, valid_frame_indices, recovered_map = cls._read_frames_with_recovery(
cap, frame_idx, total_frames_num
)
if recovered_map:
logger.info(
"Frame recovery: %d frames recovered using forward scan.",
len(recovered_map),
)
else:
frame_idx_set = set(frame_idx)
num_frames_to_sample = len(frame_idx_set)
frames, valid_frame_indices = cls._read_frames_no_recovery(
cap, frame_idx_set, max(frame_idx)
)
valid_num_frames = len(valid_frame_indices)
if valid_num_frames < num_frames_to_sample:
logger.warning(
"Video loading completed with %d broken/unreadable frames. "
"Expected to sample %d frames but only loaded %d frames.",
num_frames_to_sample - valid_num_frames,
num_frames_to_sample,
valid_num_frames,
)
return frames, valid_frame_indices
+111
View File
@@ -0,0 +1,111 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from io import BytesIO
import numpy as np
import numpy.typing as npt
from vllm.utils.import_utils import PlaceholderModule
from .base import (
VideoSourceMetadata,
VideoTargetMetadata,
check_frame_pixel_limit,
)
try:
import av
except ImportError:
av = PlaceholderModule("av") # type: ignore[assignment]
def decode_pyav(
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict,
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
with av.open(BytesIO(data)) as container:
stream = container.streams.video[0]
check_frame_pixel_limit(stream.width, stream.height)
source = loader_cls._prepare_source(
PyAVVideoBackendMixin.get_metadata(container)
)
frame_idx = loader_cls.compute_frames_index_to_sample(
source=source, target=target, **sampling_kwargs
)
frames, valid = PyAVVideoBackendMixin.decode_frames(
container, frame_idx, source.original_fps, source.duration
)
return frames, source, frame_idx, valid
class PyAVVideoBackendMixin:
"""PyAV (in-process FFmpeg bindings) codec utilities.
Reads stream metadata and decodes target frames via per-frame
``container.seek()``. The seek releases the GIL between frames and
scales with the number of sampled frames rather than the video
length, enabling concurrent decoding under serving load.
"""
@staticmethod
def get_metadata(
container: "av.container.InputContainer",
) -> VideoSourceMetadata:
if not container.streams.video:
raise ValueError("No video streams found in container")
stream = container.streams.video[0]
total_frames = stream.frames or 0
fps = float(stream.average_rate) if stream.average_rate else 0.0
duration = float(stream.duration * stream.time_base) if stream.duration else 0.0
if total_frames == 0 and duration > 0 and fps > 0:
total_frames = int(duration * fps)
return VideoSourceMetadata(total_frames, fps, duration)
@staticmethod
def decode_frames(
container: "av.container.InputContainer",
frame_indices: list[int],
fps: float,
duration: float,
) -> tuple[npt.NDArray, list[int]]:
"""Decode target frames via per-frame seek + forward decode to PTS."""
stream = container.streams.video[0]
# SLICE parallelizes within a single frame without the
# one-frame-per-thread latency penalty of FRAME threading.
stream.thread_type = "SLICE"
time_base = stream.time_base
frames_list: list[npt.NDArray] = []
valid_indices: list[int] = []
frame_interval = 1.0 / fps if fps > 0 else 0.1
max_ts = max(0.0, duration - frame_interval) if duration > 0 else float("inf")
decoder = None
last_pts = None
for idx in frame_indices:
ts = min(idx / fps, max_ts) if fps > 0 else 0.0
pts = int(ts / time_base)
# seek() snaps backward to a keyframe; reuse the running decoder
# while targets advance monotonically to avoid re-decoding the
# GOP prefix once per requested frame.
if decoder is None or last_pts is None or pts <= last_pts:
container.seek(pts, stream=stream)
decoder = container.decode(video=0)
chosen = None
for frame in decoder:
if frame.pts is not None and frame.pts >= pts:
chosen = frame
last_pts = frame.pts
break
if chosen is not None:
frames_list.append(chosen.to_ndarray(format="rgb24"))
valid_indices.append(idx)
else:
decoder = None
if not frames_list:
return np.empty((0,), dtype=np.uint8), valid_indices
return np.stack(frames_list), valid_indices
@@ -0,0 +1,372 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
import tempfile
import threading
from contextlib import contextmanager, suppress
from typing import ClassVar, NamedTuple
import numpy as np
import numpy.typing as npt
from vllm.logger import init_logger
from vllm.utils.mem_constants import MiB_bytes
from .base import (
PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
VideoSourceMetadata,
VideoTargetMetadata,
check_frame_pixel_limit,
)
logger = init_logger(__name__)
def decode_pynvvideocodec(
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict,
*,
hw_decoders: int = PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(hw_decoders)
return PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec(
loader_cls,
data,
target,
**sampling_kwargs,
)
class PyNvVideoCodecSourceMetadata(NamedTuple):
"""Metadata needed before GPU video decode."""
source: VideoSourceMetadata
width: int
height: int
# Per-decoder upper bound reserved for persistent PyNvVideoCodec surfaces.
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes
PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2
# Per-API-server CUDA context and driver allocation, measured with
# PyNvVideoCodec 2.0.4 on H100.
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES = int(1.8 * 1024 * MiB_bytes)
def validate_pynvvideocodec_hw_decoders(hw_decoders: object) -> int:
if (
isinstance(hw_decoders, bool)
or not isinstance(hw_decoders, int)
or hw_decoders < 1
):
raise ValueError("hw_decoders must be a positive integer")
return hw_decoders
def _pynvvideocodec_exception_types(nvc) -> tuple[type[Exception], ...]:
return tuple(
exception_type
for name in dir(nvc)
if name.startswith("PyNvVCException")
and isinstance((exception_type := getattr(nvc, name)), type)
and issubclass(exception_type, Exception)
)
def _pynvvc_frames_to_nhwc(frames):
"""Return a stacked PyNvVideoCodec frame batch as contiguous NHWC."""
if frames.shape[-1] != 3 and frames.shape[-3] == 3:
frames = frames.permute(0, 2, 3, 1)
return frames.contiguous()
class PyNvVideoCodecDecoderSlot:
"""A retained PyNv decoder slot and its CUDA stream.
The decoder is reused across requests: ``reconfigure_decoder`` repoints the
existing decoder at each new source instead of paying a fresh
``SimpleDecoder`` construction per request. Construction (CUVID parser +
decoder + surface-pool allocation) is the dominant per-request cost, so
reconfiguring is far cheaper. A single decoder serves both metadata
(``len``/``get_stream_metadata``) and frame decode -- no separate
metadata decoder.
"""
def __init__(self, stream) -> None:
self.stream = stream
self.decoder = None
self.source_path: str | None = None
def invalidate(self) -> None:
self.decoder = None
self.source_path = None
def _construct(self, file_path: str, nvc, device_index: int) -> None:
self.invalidate()
decoder = nvc.SimpleDecoder(
file_path,
output_color_type=nvc.OutputColorType.RGB,
use_device_memory=True,
need_scanned_stream_metadata=True,
gpu_id=device_index,
cuda_stream=self.stream.cuda_stream,
decoder_cache_size=PYNVVIDEOCODEC_DECODER_CACHE_SIZE,
)
self.decoder = decoder
self.source_path = file_path
def get_decoder(self, file_path: str, nvc, device_index: int):
if self.decoder is None:
self._construct(file_path, nvc, device_index)
elif self.source_path != file_path:
try:
self.decoder.reconfigure_decoder(file_path)
self.source_path = file_path
except Exception:
# reconfigure unsupported/unsafe for this source -> rebuild.
self._construct(file_path, nvc, device_index)
return self.decoder
class _PyNvDecoderPool:
"""Process-wide singleton managing PyNvVideoCodec decoder slot state.
Prevents subclass counter shadowing (GHSA-j682-9xp5-rrf3) by storing
all mutable pool state in a single module-level instance rather than
in ClassVar attributes that get shadowed by Python's augmented
assignment semantics on subclasses.
"""
def __init__(self) -> None:
self.slots: list[PyNvVideoCodecDecoderSlot] = []
self.active: int = 0
self.cond: threading.Condition = threading.Condition()
self.max_slots: int | None = None
def configure(self, hw_decoders: int) -> None:
with self.cond:
if self.max_slots is None:
self.max_slots = hw_decoders
elif self.max_slots != hw_decoders:
raise RuntimeError(
"PyNvVideoCodec decoder count is already configured as "
f"{self.max_slots}, got {hw_decoders}"
)
_pynv_decoder_pool = _PyNvDecoderPool()
class PyNvVideoCodecVideoBackendMixin:
"""PyNvVideoCodec utilities for GPU-backed frame decode."""
_DEVICE_INDEX: ClassVar[int] = 0
@classmethod
def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot:
import torch
return PyNvVideoCodecDecoderSlot(torch.cuda.Stream(device=cls._DEVICE_INDEX))
@classmethod
def _configure_decoder_slots(cls, hw_decoders: object) -> None:
hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders)
_pynv_decoder_pool.configure(hw_decoders)
@staticmethod
@contextmanager
def _torch_stream_context(stream):
import torch
torch.accelerator.set_device_index(stream.device.index)
previous_stream = torch.accelerator.current_stream()
torch.accelerator.set_stream(stream)
try:
yield
finally:
torch.accelerator.set_stream(previous_stream)
@classmethod
@contextmanager
def _borrow_decoder_slot(cls):
pool = _pynv_decoder_pool
create_slot = False
with pool.cond:
if pool.max_slots is None:
raise RuntimeError("PyNvVideoCodec decoder slots are not configured")
while True:
if pool.slots:
slot = pool.slots.pop()
break
if pool.active < pool.max_slots:
pool.active += 1
create_slot = True
break
pool.cond.wait()
if create_slot:
try:
slot = cls._create_decoder_slot()
except Exception:
with pool.cond:
pool.active -= 1
pool.cond.notify()
raise
borrow_succeeded = False
try:
yield slot
borrow_succeeded = True
finally:
if not borrow_succeeded:
slot.invalidate()
with pool.cond:
pool.slots.append(slot)
pool.cond.notify()
@staticmethod
def _metadata_value(metadata, *names: str, default=None):
for name in names:
value = getattr(metadata, name, None)
if value is not None:
return value
return default
@classmethod
def _read_source_metadata(
cls,
file_path: str,
nvc,
) -> PyNvVideoCodecSourceMetadata:
with cls._borrow_decoder_slot() as decoder_slot:
with cls._torch_stream_context(decoder_slot.stream):
decoder = decoder_slot.get_decoder(
file_path, nvc, device_index=cls._DEVICE_INDEX
)
metadata = decoder.get_stream_metadata()
total_frames_num = len(decoder)
width = int(cls._metadata_value(metadata, "width", default=0))
height = int(cls._metadata_value(metadata, "height", default=0))
original_fps = float(
cls._metadata_value(
metadata,
"average_fps",
"avg_frame_rate",
"frame_rate",
"frameRate",
default=0.0,
)
)
duration = float(
cls._metadata_value(metadata, "duration", default=0.0)
or (total_frames_num / original_fps if original_fps > 0 else 0.0)
)
if total_frames_num <= 0:
raise ValueError("Could not determine video frame count")
if width <= 0 or height <= 0:
raise ValueError("Could not determine video dimensions")
return PyNvVideoCodecSourceMetadata(
source=VideoSourceMetadata(total_frames_num, original_fps, duration),
width=width,
height=height,
)
@classmethod
def _decode_to_pinned_host(
cls,
file_path: str,
frame_idx: list[int],
nvc,
) -> npt.NDArray:
import torch
if not frame_idx:
return np.empty((0,), dtype=np.uint8)
with cls._borrow_decoder_slot() as decoder_slot:
stream = decoder_slot.stream
with cls._torch_stream_context(stream):
try:
decoder = decoder_slot.get_decoder(
file_path, nvc, device_index=cls._DEVICE_INDEX
)
decoded_frames = decoder.get_batch_frames_by_index(frame_idx)
except Exception as exc:
if not isinstance(
exc,
_pynvvideocodec_exception_types(nvc) + (IndexError,),
):
raise
raise ValueError("Invalid or unsupported video file.") from exc
if len(decoded_frames) < len(frame_idx):
logger.warning(
"pynvvideocodec video loading: expected %d frames but got %d.",
len(frame_idx),
len(decoded_frames),
)
torch_frames = [torch.from_dlpack(frame) for frame in decoded_frames]
if not torch_frames:
return np.empty((0,), dtype=np.uint8)
device_frames = torch.stack(torch_frames)
if device_frames.ndim != 4:
raise ValueError(
"PyNvVideoCodec returned frames with unexpected shape "
f"{tuple(device_frames.shape)}"
)
device_frames = _pynvvc_frames_to_nhwc(device_frames)
host_frames = torch.empty(
device_frames.shape,
dtype=device_frames.dtype,
device="cpu",
pin_memory=True,
)
host_frames.copy_(device_frames, non_blocking=True)
stream.synchronize()
host_array = host_frames.numpy()
del decoded_frames, torch_frames, device_frames
return host_array
@classmethod
def decode_frames_pynvvideocodec(
cls,
loader_cls,
data: bytes,
target: VideoTargetMetadata,
**kwargs,
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
import PyNvVideoCodec as nvc
from vllm.multimodal.gpu_ipc_memory import get_mm_gpu_ipc_pool
temp_fd, temp_path = tempfile.mkstemp(suffix=".mp4")
try:
with os.fdopen(temp_fd, "wb") as temp_file:
temp_file.write(data)
try:
gpu_source = cls._read_source_metadata(temp_path, nvc)
except Exception as exc:
if not isinstance(exc, _pynvvideocodec_exception_types(nvc)):
raise
raise ValueError("Invalid or unsupported video file.") from exc
check_frame_pixel_limit(gpu_source.width, gpu_source.height)
source = loader_cls._prepare_source(gpu_source.source)
frame_idx = loader_cls.compute_frames_index_to_sample(
source=source, target=target, **kwargs
)
raw_frame_bytes = len(frame_idx) * gpu_source.height * gpu_source.width * 3
pool = get_mm_gpu_ipc_pool()
if pool is None or raw_frame_bytes == 0:
frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc)
else:
with pool.acquire(raw_frame_bytes):
frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc)
finally:
with suppress(FileNotFoundError):
os.unlink(temp_path)
valid_frame_indices = frame_idx[: int(frames.shape[0])]
return frames, source, frame_idx, valid_frame_indices
@@ -0,0 +1,100 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Literal
import numpy as np
import numpy.typing as npt
from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available
from .base import (
VideoSourceMetadata,
VideoTargetMetadata,
check_frame_pixel_limit,
)
try:
from torchcodec.decoders import VideoDecoder
except (ImportError, RuntimeError):
VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment]
"decoders.VideoDecoder"
)
def decode_torchcodec(
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict,
*,
num_ffmpeg_threads: int = 0,
seek_mode: Literal["exact", "approximate"] = "exact",
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
check_torchcodec_available()
decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder(
data,
num_ffmpeg_threads=num_ffmpeg_threads,
seek_mode=seek_mode,
)
check_frame_pixel_limit(
decoder.metadata.width or 0,
decoder.metadata.height or 0,
)
source = loader_cls._prepare_source(
TorchCodecVideoBackendMixin.get_torchcodec_metadata(decoder)
)
frame_idx = loader_cls.compute_frames_index_to_sample(
source=source, target=target, **sampling_kwargs
)
frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames(
decoder, frame_idx
)
return frames, source, frame_idx, valid
class TorchCodecVideoBackendMixin:
"""TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities.
Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory
bytes and extracts the sampled indices with a single batched
``get_frames_at`` call, while releasing the GIL during decode.
"""
@staticmethod
def make_torchcodec_decoder(
data: bytes,
*,
num_ffmpeg_threads: int = 0,
seek_mode: Literal["exact", "approximate"] = "exact",
) -> "VideoDecoder":
# NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest
# of the pipeline expects, avoiding a transpose.
return VideoDecoder(
data,
dimension_order="NHWC",
num_ffmpeg_threads=num_ffmpeg_threads,
seek_mode=seek_mode,
)
@staticmethod
def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata:
md = decoder.metadata
total_frames = md.num_frames or 0
fps = float(md.average_fps) if md.average_fps else 0.0
duration = float(md.duration_seconds) if md.duration_seconds else 0.0
if total_frames == 0 and duration > 0 and fps > 0:
total_frames = int(duration * fps)
return VideoSourceMetadata(total_frames, fps, duration)
@staticmethod
def decode_torchcodec_frames(
decoder: "VideoDecoder",
frame_indices: list[int],
) -> tuple[npt.NDArray, list[int]]:
"""Decode the requested indices in one batched, index-exact call."""
if not frame_indices:
return np.empty((0,), dtype=np.uint8), []
# Note: torchcodec releases the GIL for the entire call
batch = decoder.get_frames_at(frame_indices)
return batch.data.numpy(), list(frame_indices)