@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3bb9c18f0c

[Multimodal] Reorganize video decoder backends (#49155)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

Isotr0py committed Aug 18, 2026 at 23:06 UTC 3bb9c18f0c845170679498f6aa39c6a92f2f68cb

12 files changed +1416 -1097

tests/multimodal/media/test_video.py

+7 -3

@@ -21,8 +21,10 @@ from vllm.multimodal.media import ImageMediaIO, VideoMediaIO

21 from vllm.multimodal.video import (

22 PYNVVIDEOCODEC_VIDEO_BACKEND,

23 VIDEO_LOADER_REGISTRY,

24 - PyNvVideoCodecVideoBackend,

24 VideoLoader,

25 +)

26 +from vllm.multimodal.video_decoders.pynvvideocodec import (

27 + PyNvVideoCodecVideoBackendMixin,

28 _pynvvc_frames_to_nhwc,

29 )

@@ -381,13 +383,15 @@ def test_pynvvideocodec_unrelated_error_propagates(

383 raise original_error

384

385 monkeypatch.setattr(

384 - PyNvVideoCodecVideoBackend,

386 + PyNvVideoCodecVideoBackendMixin,

387 "_read_source_metadata",

388 classmethod(raise_unrelated_error),

389 )

390

391 with pytest.raises(RuntimeError) as exc_info:

390 - PyNvVideoCodecVideoBackend.decode_frames_pynvvideocodec(b"video", None)

392 + PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec(

393 + None, b"video", None

394 + )

395

396 assert exc_info.value is original_error

397

tests/multimodal/test_gpu_ipc_memory.py

+7 -28

@@ -6,7 +6,6 @@ import time

7 import pytest

9 -import vllm.config.multimodal as multimodal_config_module

9 from vllm.config.multimodal import MultiModalConfig

10 from vllm.multimodal.gpu_ipc_memory import (

11 MultiModalGPUMemoryPool,

@@ -15,10 +14,10 @@ from vllm.multimodal.gpu_ipc_memory import (

14 reserve_mm_ipc_gpu_memory,

15 set_mm_gpu_ipc_pool,

16 )

18 -from vllm.multimodal.video import (

17 +from vllm.multimodal.video_decoders import PYNVVIDEOCODEC_VIDEO_BACKEND

18 +from vllm.multimodal.video_decoders.pynvvideocodec import (

19 PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,

20 PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,

21 - PYNVVIDEOCODEC_VIDEO_BACKEND,

21 )

22 from vllm.utils.mem_constants import GiB_bytes

@@ -186,11 +185,7 @@ def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(

185 monkeypatch: pytest.MonkeyPatch,

186 video_backend: str | None,

187 ):

189 - monkeypatch.setattr(

190 - multimodal_config_module.envs,

191 - "VLLM_VIDEO_LOADER_BACKEND",

192 - "opencv",

193 - )

188 + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")

189 mm_config = _mm_config(

190 mm_ipc_gpu_memory_gb=0.25,

191 video_backend=video_backend,

@@ -202,11 +197,7 @@ def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(

197 def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(

198 monkeypatch: pytest.MonkeyPatch,

199 ):

205 - monkeypatch.setattr(

206 - multimodal_config_module.envs,

207 - "VLLM_VIDEO_LOADER_BACKEND",

208 - "opencv",

209 - )

200 + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")

201 mm_config = _mm_config(

202 mm_ipc_gpu_memory_gb=0.25,

203 video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,

@@ -221,11 +212,7 @@ def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(

212 def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(

213 monkeypatch: pytest.MonkeyPatch,

214 ):

224 - monkeypatch.setattr(

225 - multimodal_config_module.envs,

226 - "VLLM_VIDEO_LOADER_BACKEND",

227 - PYNVVIDEOCODEC_VIDEO_BACKEND,

228 - )

215 + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", PYNVVIDEOCODEC_VIDEO_BACKEND)

216 available_bytes = 4 * GiB_bytes

217

218 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(

223 def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers(

224 monkeypatch: pytest.MonkeyPatch,

225 ):

239 - monkeypatch.setattr(

240 - multimodal_config_module.envs,

241 - "VLLM_VIDEO_LOADER_BACKEND",

242 - PYNVVIDEOCODEC_VIDEO_BACKEND,

243 - )

226 + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", PYNVVIDEOCODEC_VIDEO_BACKEND)

227 available_bytes = 8 * GiB_bytes

228

229 assert reserve_mm_ipc_gpu_memory(

@@ -253,11 +236,7 @@ def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers(

236 def test_reserve_mm_ipc_gpu_memory_uses_configured_hw_decoders(

237 monkeypatch: pytest.MonkeyPatch,

238 ):

256 - monkeypatch.setattr(

257 - multimodal_config_module.envs,

258 - "VLLM_VIDEO_LOADER_BACKEND",

259 - "opencv",

260 - )

239 + monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv")

240 available_bytes = 4 * GiB_bytes

241 mm_config = _mm_config(

242 video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,

tests/multimodal/test_video.py

+155 -18

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import itertools

5 +import subprocess

6 import sys

7 import threading

8 from contextlib import ExitStack, contextmanager

@@ -16,24 +17,26 @@ from transformers.video_utils import VideoMetadata

17 from vllm.assets.base import get_vllm_public_assets

18 from vllm.models.minimax_m3.common.mm_preprocess import MiniMaxM3VideoBackend

19 from vllm.multimodal.video import (

19 - PYNVVIDEOCODEC_DECODER_CACHE_SIZE,

20 PYNVVIDEOCODEC_VIDEO_BACKEND,

21 VIDEO_LOADER_REGISTRY,

22 DynamicVideoBackend,

23 GLM46VVideoBackend,

24 Molmo2VideoBackend,

25 - PyNvVideoCodecDecoderSlot,

26 - PyNvVideoCodecVideoBackend,

27 - PyNvVideoCodecVideoBackendMixin,

25 Qwen2VLVideoBackend,

26 Qwen3VLVideoBackend,

27 VideoBackend,

28 VideoLoader,

29 VideoSourceMetadata,

30 VideoTargetMetadata,

34 - _pynv_decoder_pool,

31 get_video_loader_backend_for_processor,

32 )

33 +from vllm.multimodal.video_decoders import decode_video, resolve_video_backend_kwargs

34 +from vllm.multimodal.video_decoders.pynvvideocodec import (

35 + PYNVVIDEOCODEC_DECODER_CACHE_SIZE,

36 + PyNvVideoCodecDecoderSlot,

37 + PyNvVideoCodecVideoBackendMixin,

38 + _pynv_decoder_pool,

39 +)

40 from vllm.platforms import current_platform

41 from vllm.transformers_utils.processor import get_video_processor_cls_name_from_config

@@ -99,6 +102,126 @@ def test_video_loader_type_doesnt_exist():

102 VIDEO_LOADER_REGISTRY.load("non_existing_video_loader")

103

104

105 +def test_video_decoder_backends_are_lazy_imported():

106 + code = """

107 +import sys

108 +import vllm.multimodal.video # noqa: F401

109 +

110 +backend_modules = {

111 + f"vllm.multimodal.video_decoders.{backend}"

112 + for backend in ("opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream")

113 +}

114 +loaded = sorted(backend_modules & sys.modules.keys())

115 +assert not loaded, loaded

116 +"""

117 + result = subprocess.run(

118 + [sys.executable, "-c", code],

119 + capture_output=True,

120 + text=True,

121 + )

122 + assert result.returncode == 0, result.stderr

123 +

124 +

125 +@pytest.mark.parametrize(

126 + "backend",

127 + ["opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"],

128 +)

129 +def test_decode_video_imports_only_selected_backend(

130 + backend: str,

131 + monkeypatch: pytest.MonkeyPatch,

132 +):

133 + imports = []

134 + decoded = object()

135 +

136 + def fake_decoder(*args, **kwargs):

137 + return decoded

138 +

139 + class FakeBackendModule:

140 + pass

141 +

142 + setattr(FakeBackendModule, f"decode_{backend}", fake_decoder)

143 +

144 + def fake_import_module(name: str, package: str):

145 + imports.append((name, package))

146 + return FakeBackendModule

147 +

148 + monkeypatch.setattr(

149 + "vllm.multimodal.video_decoders.import_module", fake_import_module

150 + )

151 + result = decode_video(

152 + backend,

153 + loader_cls=None,

154 + data=b"",

155 + target=VideoTargetMetadata(-1, -1, 300),

156 + sampling_kwargs={},

157 + backend_kwargs={},

158 + frame_recovery=False,

159 + )

160 +

161 + assert result is decoded

162 + assert imports == [(f".{backend}", "vllm.multimodal.video_decoders")]

163 +

164 +

165 +@pytest.mark.parametrize(

166 + ("backend", "kwargs", "expected_sampling", "expected_backend"),

167 + [

168 + (

169 + "torchcodec",

170 + {"min_frames": 4, "num_ffmpeg_threads": 2, "seek_mode": "approximate"},

171 + {"min_frames": 4},

172 + {"num_ffmpeg_threads": 2, "seek_mode": "approximate"},

173 + ),

174 + (

175 + "deepstream",

176 + {"max_frames": 16, "pool_size": 3, "timeout_sec": 10.0},

177 + {"max_frames": 16},

178 + {"pool_size": 3, "timeout_sec": 10.0},

179 + ),

180 + ],

181 +)

182 +def test_video_backend_kwargs_are_separated_from_sampling_kwargs(

183 + backend: str,

184 + kwargs: dict,

185 + expected_sampling: dict,

186 + expected_backend: dict,

187 +):

188 + original_kwargs = dict(kwargs)

189 + sampling_kwargs, backend_kwargs = resolve_video_backend_kwargs(backend, kwargs)

190 +

191 + assert sampling_kwargs == expected_sampling

192 + assert backend_kwargs == expected_backend

193 + assert kwargs == original_kwargs

194 +

195 +

196 +def test_video_backend_rejects_options_for_another_decoder():

197 + with pytest.raises(

198 + ValueError, match="num_ffmpeg_threads is not supported by the 'pyav' backend"

199 + ):

200 + resolve_video_backend_kwargs("pyav", {"num_ffmpeg_threads": 2})

201 +

202 +

203 +@pytest.mark.parametrize(

204 + ("backend", "error"),

205 + [

206 + ("pyav", AssertionError),

207 + (PYNVVIDEOCODEC_VIDEO_BACKEND, ValueError),

208 + ],

209 +)

210 +def test_video_decoder_spec_validates_frame_recovery(

211 + backend: str, error: type[Exception]

212 +):

213 + with pytest.raises(error, match="frame_recovery is not supported"):

214 + decode_video(

215 + backend,

216 + loader_cls=None,

217 + data=b"",

218 + target=VideoTargetMetadata(-1, -1, 300),

219 + sampling_kwargs={},

220 + backend_kwargs={},

221 + frame_recovery=True,

222 + )

223 +

224 +

225 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")

226 def test_pynvvideocodec_backend_accounts_raw_decoded_frames(

227 monkeypatch: pytest.MonkeyPatch,

@@ -145,7 +268,9 @@ def test_pynvvideocodec_backend_accounts_raw_decoded_frames(

268 "vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool

269 )

270 monkeypatch.setattr(

148 - PyNvVideoCodecVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode)

271 + PyNvVideoCodecVideoBackendMixin,

272 + "_decode_to_pinned_host",

273 + classmethod(fake_decode),

274 )

275

276 loader = VIDEO_LOADER_REGISTRY.load(PYNVVIDEOCODEC_VIDEO_BACKEND)

@@ -205,7 +330,9 @@ def test_pynvvideocodec_codec_uses_dynamic_sampling_strategy(

330 "vllm.multimodal.gpu_ipc_memory.get_mm_gpu_ipc_pool", lambda: pool

331 )

332 monkeypatch.setattr(

208 - DynamicVideoBackend, "_decode_to_pinned_host", classmethod(fake_decode)

333 + PyNvVideoCodecVideoBackendMixin,

334 + "_decode_to_pinned_host",

335 + classmethod(fake_decode),

336 )

337

338 loader = VIDEO_LOADER_REGISTRY.load("opencv_dynamic")

@@ -273,7 +400,7 @@ def test_pynvvideocodec_decoder_slots_are_bounded(

400

401 create_count = 0

402 with _fresh_decoder_pool():

276 - PyNvVideoCodecVideoBackend._configure_decoder_slots(hw_decoders)

403 + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(hw_decoders)

404

405 def fake_create_slot(cls):

406 nonlocal create_count

@@ -281,7 +408,7 @@ def test_pynvvideocodec_decoder_slots_are_bounded(

408 return FakeSlot()

409

410 monkeypatch.setattr(

284 - PyNvVideoCodecVideoBackend,

411 + PyNvVideoCodecVideoBackendMixin,

412 "_create_decoder_slot",

413 classmethod(fake_create_slot),

414 )

@@ -291,12 +418,16 @@ def test_pynvvideocodec_decoder_slots_are_bounded(

418

419 with ExitStack() as stack:

420 retained_slots = [

294 - stack.enter_context(PyNvVideoCodecVideoBackend._borrow_decoder_slot())

421 + stack.enter_context(

422 + PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot()

423 + )

424 for _ in range(hw_decoders)

425 ]

426

427 def borrow_extra_slot():

299 - with PyNvVideoCodecVideoBackend._borrow_decoder_slot() as extra_slot:

428 + with (

429 + PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot()

430 + ) as extra_slot:

431 seen_slots.append(extra_slot)

432 borrowed.set()

433

@@ -317,11 +448,11 @@ def test_pynvvideocodec_decoder_slots_are_configured_once(

448 ):

449 monkeypatch.setattr(_pynv_decoder_pool, "max_slots", None)

450

320 - PyNvVideoCodecVideoBackend._configure_decoder_slots(2)

321 - PyNvVideoCodecVideoBackend._configure_decoder_slots(2)

451 + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(2)

452 + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(2)

453

454 with pytest.raises(RuntimeError, match="already configured as 2, got 3"):

324 - PyNvVideoCodecVideoBackend._configure_decoder_slots(3)

455 + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(3)

456

457

458 def test_pynvvideocodec_failed_rebuild_invalidates_decoder_slot():

@@ -367,7 +498,7 @@ def test_pynvvideocodec_failed_rebuild_invalidates_decoder_slot():

498

499 with (

500 pytest.raises(RuntimeError, match="construct failed"),

370 - PyNvVideoCodecVideoBackend._borrow_decoder_slot() as borrowed,

501 + PyNvVideoCodecVideoBackendMixin._borrow_decoder_slot() as borrowed,

502 ):

503 assert borrowed is slot

504 borrowed.get_decoder(

@@ -456,6 +587,12 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():

587 counters via ClassVar shadowing.

588 """

589

590 + class MixinSubclassA(PyNvVideoCodecVideoBackendMixin):

591 + pass

592 +

593 + class MixinSubclassB(PyNvVideoCodecVideoBackendMixin):

594 + pass

595 +

596 class FakeSlot:

597 pass

598

@@ -475,8 +612,8 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():

612 )

613 try:

614 with ExitStack() as stack:

478 - stack.enter_context(VideoBackend._borrow_decoder_slot())

479 - stack.enter_context(Qwen3VLVideoBackend._borrow_decoder_slot())

615 + stack.enter_context(MixinSubclassA._borrow_decoder_slot())

616 + stack.enter_context(MixinSubclassB._borrow_decoder_slot())

617 assert pool.active == 2

618

619 blocked = threading.Event()

@@ -484,7 +621,7 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():

621

622 def try_borrow():

623 blocked.set()

487 - with Qwen2VLVideoBackend._borrow_decoder_slot():

624 + with MixinSubclassB._borrow_decoder_slot():

625 acquired.set()

626

627 t = threading.Thread(target=try_borrow)

vllm/multimodal/gpu_ipc_memory.py

+5 -3

@@ -188,11 +188,13 @@ def reserve_mm_ipc_gpu_memory(

188 return available_kv_cache_memory_bytes

189

190 from vllm import envs

191 - from vllm.multimodal.video import (

192 - PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,

193 - PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,

191 + from vllm.multimodal.video_decoders import (

192 PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,

193 PYNVVIDEOCODEC_VIDEO_BACKEND,

194 + )

195 + from vllm.multimodal.video_decoders.pynvvideocodec import (

196 + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,

197 + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,

198 validate_pynvvideocodec_hw_decoders,

199 )

200

vllm/multimodal/video.py

+66 -1045

@@ -1,42 +1,29 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 -import os

5 -import tempfile

6 -import threading

4 from abc import abstractmethod

8 -from contextlib import contextmanager, suppress

9 -from io import BytesIO

10 -from typing import Any, ClassVar, Literal, NamedTuple, cast

5 +from typing import Any, ClassVar, Literal, cast

7 import numpy as np

8 import numpy.typing as npt

9 import torch

16 -from vllm import envs

11 from vllm.logger import init_logger

18 -from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available

19 -from vllm.utils.mem_constants import MiB_bytes

12 +from vllm.multimodal.video_decoders import (

13 + PYNVVIDEOCODEC_VIDEO_BACKEND,

14 + VideoDecoderBackend,

15 + VideoSourceMetadata,

16 + VideoTargetMetadata,

17 + decode_video,

18 + resolve_video_backend_kwargs,

19 +)

20 +from vllm.utils.import_utils import PlaceholderModule

21 from vllm.utils.registry import ExtensionManager

23 try:

24 import cv2

24 - import cv2.videoio_registry as vr

25 except ImportError:

26 cv2 = PlaceholderModule("cv2")

27 - vr = PlaceholderModule("cv2").placeholder_attr("videoio_registry")

28 -

29 -try:

30 - import av

31 -except ImportError:

32 - av = PlaceholderModule("av") # type: ignore[assignment]

33 -

34 -try:

35 - from torchcodec.decoders import VideoDecoder

36 -except (ImportError, RuntimeError):

37 - VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment]

38 - "decoders.VideoDecoder"

39 - )

29 logger = init_logger(__name__)

@@ -106,18 +93,6 @@ def get_video_loader_backend_for_processor(

93 return VIDEO_LOADER_REGISTRY.get_backend_for_video_processor(video_processor)

109 -def _check_frame_pixel_limit(width: int, height: int) -> None:

110 - """Reject video frames exceeding VLLM_MAX_IMAGE_PIXELS before decoding."""

111 - max_pixels = envs.VLLM_MAX_IMAGE_PIXELS

112 - if max_pixels > 0 and width * height > max_pixels:

113 - raise ValueError(

114 - f"Video frame dimensions {width}x{height} "

115 - f"({width * height} pixels) exceed the maximum of "

116 - f"{max_pixels} pixels. Set VLLM_MAX_IMAGE_PIXELS to "

117 - f"increase this limit."

118 - )

119 -

120 -

96 def resize_video(frames: npt.NDArray, size: tuple[int, int]) -> npt.NDArray:

97 num_frames, _, _, channels = frames.shape

98 new_height, new_width = size

@@ -149,30 +124,6 @@ def sample_frames_from_video(frames: npt.NDArray, num_frames: int) -> npt.NDArra

124 return sampled_frames

125

126

152 -class VideoTargetMetadata(NamedTuple):

153 - """Metadata represents target video."""

154 -

155 - num_frames: int

156 - fps: float

157 - max_duration: float

158 -

159 -

160 -class VideoSourceMetadata(NamedTuple):

161 - """Metadata represents source video."""

162 -

163 - total_frames_num: int

164 - original_fps: float

165 - duration: float

166 -

167 -

168 -class PyNvVideoCodecSourceMetadata(NamedTuple):

169 - """Metadata needed before GPU video decode."""

170 -

171 - source: VideoSourceMetadata

172 - width: int

173 - height: int

174 -

175 -

127 class VideoLoader:

128 @classmethod

129 def compute_frames_index_to_sample(

@@ -184,6 +135,11 @@ class VideoLoader:

135 """Return the list of frame indices to sample from the video."""

136 raise NotImplementedError

137

138 + @classmethod

139 + def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata:

140 + """Sampling-algorithm-specific metadata adjustment hook."""

141 + return source

142 +

143 @classmethod

144 @abstractmethod

145 def load_bytes(

@@ -214,835 +170,9 @@ class VideoLoader:

170 VIDEO_LOADER_REGISTRY = VideoLoaderRegistry()

171 VIDEO_LOADER_REGISTRY.register_gpu_codec("deepstream")

172

217 -PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec"

218 -# Per-decoder upper bound reserved for persistent PyNvVideoCodec surfaces.

219 -PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes

220 -PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2

221 -PYNVVIDEOCODEC_DEFAULT_HW_DECODERS = 2

222 -# Per-API-server CUDA context and driver allocation, measured with

223 -# PyNvVideoCodec 2.0.4 on H100.

224 -PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES = int(1.8 * 1024 * MiB_bytes)

225 -

226 -

227 -def validate_pynvvideocodec_hw_decoders(hw_decoders: object) -> int:

228 - if (

229 - isinstance(hw_decoders, bool)

230 - or not isinstance(hw_decoders, int)

231 - or hw_decoders < 1

232 - ):

233 - raise ValueError("hw_decoders must be a positive integer")

234 - return hw_decoders

235 -

236 -

237 -def _pynvvideocodec_exception_types(nvc) -> tuple[type[Exception], ...]:

238 - return tuple(

239 - exception_type

240 - for name in dir(nvc)

241 - if name.startswith("PyNvVCException")

242 - and isinstance((exception_type := getattr(nvc, name)), type)

243 - and issubclass(exception_type, Exception)

244 - )

245 -

246 -

247 -class PyNvVideoCodecDecoderSlot:

248 - """A retained PyNv decoder slot and its CUDA stream.

249 -

250 - The decoder is reused across requests: ``reconfigure_decoder`` repoints the

251 - existing decoder at each new source instead of paying a fresh

252 - ``SimpleDecoder`` construction per request. Construction (CUVID parser +

253 - decoder + surface-pool allocation) is the dominant per-request cost, so

254 - reconfiguring is far cheaper. A single decoder serves both metadata

255 - (``len``/``get_stream_metadata``) and frame decode -- no separate

256 - metadata decoder.

257 - """

258 -

259 - def __init__(self, stream) -> None:

260 - self.stream = stream

261 - self.decoder = None

262 - self.source_path: str | None = None

263 -

264 - def invalidate(self) -> None:

265 - self.decoder = None

266 - self.source_path = None

267 -

268 - def _construct(self, file_path: str, nvc, device_index: int) -> None:

269 - self.invalidate()

270 - decoder = nvc.SimpleDecoder(

271 - file_path,

272 - output_color_type=nvc.OutputColorType.RGB,

273 - use_device_memory=True,

274 - need_scanned_stream_metadata=True,

275 - gpu_id=device_index,

276 - cuda_stream=self.stream.cuda_stream,

277 - decoder_cache_size=PYNVVIDEOCODEC_DECODER_CACHE_SIZE,

278 - )

279 - self.decoder = decoder

280 - self.source_path = file_path

281 -

282 - def get_decoder(self, file_path: str, nvc, device_index: int):

283 - if self.decoder is None:

284 - self._construct(file_path, nvc, device_index)

285 - elif self.source_path != file_path:

286 - try:

287 - self.decoder.reconfigure_decoder(file_path)

288 - self.source_path = file_path

289 - except Exception:

290 - # reconfigure unsupported/unsafe for this source -> rebuild.

291 - self._construct(file_path, nvc, device_index)

292 - return self.decoder

293 -

294 -

295 -class OpenCVVideoBackendMixin:

296 - @staticmethod

297 - def get_cv2_video_api():

298 - api_pref = None

299 - for backend in vr.getStreamBufferedBackends():

300 - if not vr.hasBackend(backend):

301 - continue

302 - if not vr.isBackendBuiltIn(backend):

303 - _, abi, api = vr.getStreamBufferedBackendPluginVersion(backend)

304 - if abi < 1 or (abi == 1 and api < 2):

305 - continue

306 - api_pref = backend

307 - break

308 - return api_pref

309 -

310 - @classmethod

311 - def open_video_capture(cls, data: bytes) -> "cv2.VideoCapture":

312 - backend = cls.get_cv2_video_api()

313 - cap = cv2.VideoCapture(BytesIO(data), backend, [])

314 - if not cap.isOpened():

315 - raise ValueError("Could not open video stream")

316 - return cap

317 -

318 - @staticmethod

319 - def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata:

320 - total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

321 - original_fps = cap.get(cv2.CAP_PROP_FPS)

322 - duration = total_frames_num / original_fps if original_fps > 0 else 0

323 - return VideoSourceMetadata(

324 - total_frames_num=total_frames_num,

325 - original_fps=original_fps,

326 - duration=duration,

327 - )

328 -

329 - @classmethod

330 - def _can_use_for_recovery(

331 - cls,

332 - idx: int,

333 - failed_frames: list[int],

334 - next_target_map: dict[int, int],

335 - total_frames: int,

336 - ) -> bool:

337 - """Check if current frame can recover the oldest failed frame."""

338 - if not failed_frames:

339 - return False

340 - oldest_failed = failed_frames[0]

341 - limit = next_target_map.get(oldest_failed, total_frames)

342 - return idx < limit

343 -

344 - @classmethod

345 - def _read_frames_with_recovery(

346 - cls,

347 - cap: "cv2.VideoCapture",

348 - frame_indices: list[int],

349 - total_frames: int,

350 - ) -> tuple[npt.NDArray, list[int], dict[int, int]]:

351 - """

352 - Read frames with dynamic window forward-scan recovery.

353 -

354 - When a target frame fails to load, the next successfully grabbed

355 - frame (before the next target frame) will be used to recover it.

356 -

357 - Args:

358 - cap: OpenCV VideoCapture object

359 - frame_indices: Sorted list of target frame indices to load

360 - total_frames: Total number of frames in the video

361 -

362 - Returns:

363 - Tuple of (frames_array, valid_frame_indices, recovered_map)

364 - - frames_array: Array of loaded frames

365 - - valid_frame_indices: List of frame indices that were loaded

366 - - recovered_map: Dict mapping recovered_idx -> source_idx

367 - """

368 - width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))

369 - height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

370 -

371 - assert width > 0 and height > 0, (

372 - f"Invalid video frame size: width={width}, height={height}"

373 - )

374 -

375 - frame_idx_set = set(frame_indices)

376 - max_frame_idx = frame_indices[-1] if frame_indices else 0

377 -

378 - # Build map: target_idx -> next_target_idx (for recovery window)

379 - next_target_map: dict[int, int] = {}

380 - for k in range(len(frame_indices) - 1):

381 - next_target_map[frame_indices[k]] = frame_indices[k + 1]

382 - next_target_map[frame_indices[-1]] = total_frames

383 -

384 - frames_list: list[npt.NDArray] = []

385 - valid_frame_indices: list[int] = []

386 - failed_frames_idx: list[int] = []

387 - recovered_map: dict[int, int] = {}

388 -

389 - i = 0

390 - for idx in range(max_frame_idx + 1):

391 - is_target_frame = idx in frame_idx_set

392 -

393 - # Attempt to grab the current frame

394 - ok = cap.grab()

395 -

396 - if not ok:

397 - if is_target_frame:

398 - logger.debug(

399 - "Failed to grab frame %d during video loading.",

400 - idx,

401 - )

402 - failed_frames_idx.append(idx)

403 - continue

404 -

405 - # Check if we should retrieve: target frame OR can recover a failed one

406 - can_recover = cls._can_use_for_recovery(

407 - idx, failed_frames_idx, next_target_map, total_frames

408 - )

409 -

410 - if is_target_frame or can_recover:

411 - ret, frame = cap.retrieve()

412 -

413 - if ret and frame is not None and frame.size > 0:

414 - rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

415 - frames_list.append(rgb_frame)

416 - valid_frame_indices.append(idx)

417 - i += 1

418 -

419 - if can_recover:

420 - recovered_idx = failed_frames_idx.pop(0)

421 - recovered_map[recovered_idx] = idx

422 - logger.info(

423 - "Recovered frame %d using frame %d (delay: %d)",

424 - recovered_idx,

425 - idx,

426 - idx - recovered_idx,

427 - )

428 - elif is_target_frame:

429 - logger.debug(

430 - "Failed to retrieve frame %d during video loading.",

431 - idx,

432 - )

433 - failed_frames_idx.append(idx)

434 -

435 - # Log any remaining failed frames

436 - for failed_idx in failed_frames_idx:

437 - logger.debug(

438 - "Frame %d could not be recovered (end of video).",

439 - failed_idx,

440 - )

441 -

442 - # Stack frames

443 - if frames_list:

444 - frames = np.stack(frames_list)

445 - else:

446 - frames = np.empty((0, height, width, 3), dtype=np.uint8)

447 -

448 - return frames, valid_frame_indices, recovered_map

449 -

450 - @classmethod

451 - def _read_frames_no_recovery(

452 - cls,

453 - cap,

454 - frame_indices: set[int],

455 - max_frame_idx: int,

456 - ) -> tuple[npt.NDArray, list[int]]:

457 - num_expected_frames = len(frame_indices)

458 - width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))

459 - height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

460 - frames = np.empty((num_expected_frames, height, width, 3), dtype=np.uint8)

461 -

462 - i = 0

463 - valid_frame_indices = []

464 - for idx in range(max_frame_idx + 1):

465 - ok = cap.grab()

466 - if not ok:

467 - # Frame is broken/unreadable, skip it

468 - if idx in frame_indices:

469 - logger.debug(

470 - "Failed to grab frame %d during video loading. "

471 - "This frame will be skipped.",

472 - idx,

473 - )

474 - continue

475 - if idx in frame_indices:

476 - ret, frame = cap.retrieve()

477 - if ret:

478 - frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

479 - valid_frame_indices.append(idx)

480 - i += 1

481 - else:

482 - # retrieve() failed even though grab() succeeded

483 - logger.debug(

484 - "Failed to retrieve frame %d during video loading. "

485 - "This frame will be skipped.",

486 - idx,

487 - )

488 -

489 - valid_num_frames = len(valid_frame_indices)

490 - if valid_num_frames < num_expected_frames:

491 - logger.warning(

492 - "Video loading completed with %d broken/unreadable frames. "

493 - "Expected %d frames but only loaded %d frames.",

494 - num_expected_frames - valid_num_frames,

495 - num_expected_frames,

496 - valid_num_frames,

497 - )

498 -

499 - return frames[:valid_num_frames], valid_frame_indices

500 -

501 - @classmethod

502 - def read_frames(

503 - cls,

504 - cap: "cv2.VideoCapture",

505 - frame_idx: list[int],

506 - total_frames_num: int,

507 - *,

508 - frame_recovery: bool = False,

509 - ) -> tuple[npt.NDArray, list[int]]:

510 - if frame_recovery:

511 - num_frames_to_sample = len(frame_idx)

512 - frames, valid_frame_indices, recovered_map = cls._read_frames_with_recovery(

513 - cap, frame_idx, total_frames_num

514 - )

515 -

516 - if recovered_map:

517 - logger.info(

518 - "Frame recovery: %d frames recovered using forward scan.",

519 - len(recovered_map),

520 - )

521 - else:

522 - frame_idx_set = set(frame_idx)

523 - num_frames_to_sample = len(frame_idx_set)

524 - frames, valid_frame_indices = cls._read_frames_no_recovery(

525 - cap, frame_idx_set, max(frame_idx)

526 - )

527 - valid_num_frames = len(valid_frame_indices)

528 - if valid_num_frames < num_frames_to_sample:

529 - logger.warning(

530 - "Video loading completed with %d broken/unreadable frames. "

531 - "Expected to sample %d frames but only loaded %d frames.",

532 - num_frames_to_sample - valid_num_frames,

533 - num_frames_to_sample,

534 - valid_num_frames,

535 - )

536 - return frames, valid_frame_indices

537 -

538 -

539 -class PyAVVideoBackendMixin:

540 - """PyAV (in-process FFmpeg bindings) codec utilities.

541 -

542 - Reads stream metadata and decodes target frames via per-frame

543 - ``container.seek()``. The seek releases the GIL between frames and

544 - scales with the number of sampled frames rather than the video

545 - length, enabling concurrent decoding under serving load.

546 - """

547 -

548 - @staticmethod

549 - def get_metadata(

550 - container: "av.container.InputContainer",

551 - ) -> VideoSourceMetadata:

552 - if not container.streams.video:

553 - raise ValueError("No video streams found in container")

554 - stream = container.streams.video[0]

555 - total_frames = stream.frames or 0

556 - fps = float(stream.average_rate) if stream.average_rate else 0.0

557 - duration = float(stream.duration * stream.time_base) if stream.duration else 0.0

558 - if total_frames == 0 and duration > 0 and fps > 0:

559 - total_frames = int(duration * fps)

560 - return VideoSourceMetadata(total_frames, fps, duration)

561 -

562 - @staticmethod

563 - def decode_frames(

564 - container: "av.container.InputContainer",

565 - frame_indices: list[int],

566 - fps: float,

567 - duration: float,

568 - ) -> tuple[npt.NDArray, list[int]]:

569 - """Decode target frames via per-frame seek + forward decode to PTS."""

570 - stream = container.streams.video[0]

571 - # SLICE parallelizes within a single frame without the

572 - # one-frame-per-thread latency penalty of FRAME threading.

573 - stream.thread_type = "SLICE"

574 - time_base = stream.time_base

575 -

576 - frames_list: list[npt.NDArray] = []

577 - valid_indices: list[int] = []

578 - frame_interval = 1.0 / fps if fps > 0 else 0.1

579 - max_ts = max(0.0, duration - frame_interval) if duration > 0 else float("inf")

580 -

581 - decoder = None

582 - last_pts = None

583 - for idx in frame_indices:

584 - ts = min(idx / fps, max_ts) if fps > 0 else 0.0

585 - pts = int(ts / time_base)

586 - # seek() snaps backward to a keyframe; reuse the running decoder

587 - # while targets advance monotonically to avoid re-decoding the

588 - # GOP prefix once per requested frame.

589 - if decoder is None or last_pts is None or pts <= last_pts:

590 - container.seek(pts, stream=stream)

591 - decoder = container.decode(video=0)

592 - chosen = None

593 - for frame in decoder:

594 - if frame.pts is not None and frame.pts >= pts:

595 - chosen = frame

596 - last_pts = frame.pts

597 - break

598 - if chosen is not None:

599 - frames_list.append(chosen.to_ndarray(format="rgb24"))

600 - valid_indices.append(idx)

601 - else:

602 - decoder = None

603 -

604 - if not frames_list:

605 - return np.empty((0,), dtype=np.uint8), valid_indices

606 - return np.stack(frames_list), valid_indices

607 -

608 -

609 -class TorchCodecVideoBackendMixin:

610 - """TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities.

611 -

612 - Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory

613 - bytes and extracts the sampled indices with a single batched

614 - ``get_frames_at`` call, while releasing the GIL during decode.

615 - """

616 -

617 - @staticmethod

618 - def make_torchcodec_decoder(

619 - data: bytes,

620 - *,

621 - num_ffmpeg_threads: int = 0,

622 - seek_mode: Literal["exact", "approximate"] = "exact",

623 - ) -> "VideoDecoder":

624 - # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest

625 - # of the pipeline expects, avoiding a transpose.

626 - return VideoDecoder(

627 - data,

628 - dimension_order="NHWC",

629 - num_ffmpeg_threads=num_ffmpeg_threads,

630 - seek_mode=seek_mode,

631 - )

632 -

633 - @staticmethod

634 - def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata:

635 - md = decoder.metadata

636 - total_frames = md.num_frames or 0

637 - fps = float(md.average_fps) if md.average_fps else 0.0

638 - duration = float(md.duration_seconds) if md.duration_seconds else 0.0

639 - if total_frames == 0 and duration > 0 and fps > 0:

640 - total_frames = int(duration * fps)

641 - return VideoSourceMetadata(total_frames, fps, duration)

642 -

643 - @staticmethod

644 - def decode_torchcodec_frames(

645 - decoder: "VideoDecoder",

646 - frame_indices: list[int],

647 - ) -> tuple[npt.NDArray, list[int]]:

648 - """Decode the requested indices in one batched, index-exact call."""

649 - if not frame_indices:

650 - return np.empty((0,), dtype=np.uint8), []

651 - # Note: torchcodec releases the GIL for the entire call

652 - batch = decoder.get_frames_at(frame_indices)

653 - return batch.data.numpy(), list(frame_indices)

654 -

655 -

656 -def _pynvvc_frames_to_nhwc(frames: torch.Tensor) -> torch.Tensor:

657 - """Return a stacked PyNvVideoCodec frame batch as contiguous NHWC.

658 -

659 - PyNvVideoCodec's per-frame layout has varied across versions (HWC vs CHW),

660 - so detect the channel axis rather than assuming a fixed order. NHWC is the

661 - layout the other video backends return and the HF video processors expect.

662 -

663 - Args:

664 - frames: A ``(N, ?, ?, ?)`` uint8 tensor in either NHWC or NCHW order.

665 -

666 - Returns:

667 - The same frames as a contiguous ``(N, H, W, C)`` tensor.

668 - """

669 - if frames.shape[-1] != 3 and frames.shape[-3] == 3:

670 - frames = frames.permute(0, 2, 3, 1) # NCHW -> NHWC

671 - return frames.contiguous()

672 -

673 -

674 -class _PyNvDecoderPool:

675 - """Process-wide singleton managing PyNvVideoCodec decoder slot state.

676 -

677 - Prevents subclass counter shadowing (GHSA-j682-9xp5-rrf3) by storing

678 - all mutable pool state in a single module-level instance rather than

679 - in ClassVar attributes that get shadowed by Python's augmented

680 - assignment semantics on subclasses.

681 - """

682 -

683 - def __init__(self) -> None:

684 - self.slots: list[PyNvVideoCodecDecoderSlot] = []

685 - self.active: int = 0

686 - self.cond: threading.Condition = threading.Condition()

687 - self.max_slots: int | None = None

688 -

689 - def configure(self, hw_decoders: int) -> None:

690 - with self.cond:

691 - if self.max_slots is None:

692 - self.max_slots = hw_decoders

693 - elif self.max_slots != hw_decoders:

694 - raise RuntimeError(

695 - "PyNvVideoCodec decoder count is already configured as "

696 - f"{self.max_slots}, got {hw_decoders}"

697 - )

698 -

699 -

700 -_pynv_decoder_pool = _PyNvDecoderPool()

701 -

702 -

703 -class PyNvVideoCodecVideoBackendMixin:

704 - """PyNvVideoCodec utilities for GPU-backed frame decode."""

705 -

706 - _DEVICE_INDEX: ClassVar[int] = 0

707 -

708 - @classmethod

709 - @abstractmethod

710 - def compute_frames_index_to_sample(

711 - cls,

712 - source: VideoSourceMetadata,

713 - target: VideoTargetMetadata,

714 - **kwargs,

715 - ) -> list[int]:

716 - raise NotImplementedError

717 -

718 - @classmethod

719 - @abstractmethod

720 - def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata:

721 - raise NotImplementedError

722 -

723 - @classmethod

724 - def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot:

725 - import torch

726 -

727 - return PyNvVideoCodecDecoderSlot(torch.cuda.Stream(device=cls._DEVICE_INDEX))

728 -

729 - @classmethod

730 - def _configure_decoder_slots(cls, hw_decoders: object) -> None:

731 - hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders)

732 - _pynv_decoder_pool.configure(hw_decoders)

733 -

734 - @staticmethod

735 - @contextmanager

736 - def _torch_stream_context(stream):

737 - import torch

738 -

739 - torch.accelerator.set_device_index(stream.device.index)

740 - previous_stream = torch.accelerator.current_stream()

741 - torch.accelerator.set_stream(stream)

742 - try:

743 - yield

744 - finally:

745 - torch.accelerator.set_stream(previous_stream)

746 -

747 - @classmethod

748 - @contextmanager

749 - def _borrow_decoder_slot(cls):

750 - pool = _pynv_decoder_pool

751 - create_slot = False

752 - with pool.cond:

753 - if pool.max_slots is None:

754 - raise RuntimeError("PyNvVideoCodec decoder slots are not configured")

755 - while True:

756 - if pool.slots:

757 - slot = pool.slots.pop()

758 - break

759 - if pool.active < pool.max_slots:

760 - pool.active += 1

761 - create_slot = True

762 - break

763 - pool.cond.wait()

764 -

765 - if create_slot:

766 - try:

767 - slot = cls._create_decoder_slot()

768 - except Exception:

769 - with pool.cond:

770 - pool.active -= 1

771 - pool.cond.notify()

772 - raise

773 -

774 - borrow_succeeded = False

775 - try:

776 - yield slot

777 - borrow_succeeded = True

778 - finally:

779 - if not borrow_succeeded:

780 - slot.invalidate()

781 - with pool.cond:

782 - pool.slots.append(slot)

783 - pool.cond.notify()

784 -

785 - @staticmethod

786 - def _metadata_value(metadata, *names: str, default=None):

787 - for name in names:

788 - value = getattr(metadata, name, None)

789 - if value is not None:

790 - return value

791 - return default

792 -

793 - @classmethod

794 - def _read_source_metadata(

795 - cls,

796 - file_path: str,

797 - nvc,

798 - ) -> PyNvVideoCodecSourceMetadata:

799 - with cls._borrow_decoder_slot() as decoder_slot:

800 - with cls._torch_stream_context(decoder_slot.stream):

801 - decoder = decoder_slot.get_decoder(

802 - file_path, nvc, device_index=cls._DEVICE_INDEX

803 - )

804 - metadata = decoder.get_stream_metadata()

805 - total_frames_num = len(decoder)

806 - width = int(cls._metadata_value(metadata, "width", default=0))

807 - height = int(cls._metadata_value(metadata, "height", default=0))

808 - original_fps = float(

809 - cls._metadata_value(

810 - metadata,

811 - "average_fps",

812 - "avg_frame_rate",

813 - "frame_rate",

814 - "frameRate",

815 - default=0.0,

816 - )

817 - )

818 - duration = float(

819 - cls._metadata_value(metadata, "duration", default=0.0)

820 - or (total_frames_num / original_fps if original_fps > 0 else 0.0)

821 - )

822 - if total_frames_num <= 0:

823 - raise ValueError("Could not determine video frame count")

824 - if width <= 0 or height <= 0:

825 - raise ValueError("Could not determine video dimensions")

826 - return PyNvVideoCodecSourceMetadata(

827 - source=VideoSourceMetadata(total_frames_num, original_fps, duration),

828 - width=width,

829 - height=height,

830 - )

831 -

832 - @classmethod

833 - def _decode_to_pinned_host(

834 - cls,

835 - file_path: str,

836 - frame_idx: list[int],

837 - nvc,

838 - ) -> npt.NDArray:

839 - import torch

840 -

841 - if not frame_idx:

842 - return np.empty((0,), dtype=np.uint8)

843 -

844 - with cls._borrow_decoder_slot() as decoder_slot:

845 - stream = decoder_slot.stream

846 - with cls._torch_stream_context(stream):

847 - try:

848 - decoder = decoder_slot.get_decoder(

849 - file_path, nvc, device_index=cls._DEVICE_INDEX

850 - )

851 - decoded_frames = decoder.get_batch_frames_by_index(frame_idx)

852 - except Exception as exc:

853 - if not isinstance(

854 - exc,

855 - _pynvvideocodec_exception_types(nvc) + (IndexError,),

856 - ):

857 - raise

858 - raise ValueError("Invalid or unsupported video file.") from exc

859 - if len(decoded_frames) < len(frame_idx):

860 - logger.warning(

861 - "pynvvideocodec video loading: expected %d frames but got %d.",

862 - len(frame_idx),

863 - len(decoded_frames),

864 - )

865 - torch_frames = [torch.from_dlpack(frame) for frame in decoded_frames]

866 - if not torch_frames:

867 - return np.empty((0,), dtype=np.uint8)

868 - device_frames = torch.stack(torch_frames)

869 - if device_frames.ndim != 4:

870 - raise ValueError(

871 - "PyNvVideoCodec returned frames with unexpected shape "

872 - f"{tuple(device_frames.shape)}"

873 - )

874 - device_frames = _pynvvc_frames_to_nhwc(device_frames)

875 - host_frames = torch.empty(

876 - device_frames.shape,

877 - dtype=device_frames.dtype,

878 - device="cpu",

879 - pin_memory=True,

880 - )

881 - host_frames.copy_(device_frames, non_blocking=True)

882 - stream.synchronize()

883 - host_array = host_frames.numpy()

884 - del decoded_frames, torch_frames, device_frames

885 - return host_array

886 -

887 - @classmethod

888 - def decode_frames_pynvvideocodec(

889 - cls,

890 - data: bytes,

891 - target: VideoTargetMetadata,

892 - **kwargs,

893 - ) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:

894 - import PyNvVideoCodec as nvc

895 -

896 - from vllm.multimodal.gpu_ipc_memory import get_mm_gpu_ipc_pool

897 -

898 - temp_fd, temp_path = tempfile.mkstemp(suffix=".mp4")

899 - try:

900 - with os.fdopen(temp_fd, "wb") as temp_file:

901 - temp_file.write(data)

902 -

903 - try:

904 - gpu_source = cls._read_source_metadata(temp_path, nvc)

905 - except Exception as exc:

906 - if not isinstance(exc, _pynvvideocodec_exception_types(nvc)):

907 - raise

908 - raise ValueError("Invalid or unsupported video file.") from exc

909 - _check_frame_pixel_limit(gpu_source.width, gpu_source.height)

910 - source = cls._prepare_source(gpu_source.source)

911 - frame_idx = cls.compute_frames_index_to_sample(

912 - source=source, target=target, **kwargs

913 - )

914 - raw_frame_bytes = len(frame_idx) * gpu_source.height * gpu_source.width * 3

915 - pool = get_mm_gpu_ipc_pool()

916 - if pool is None or raw_frame_bytes == 0:

917 - frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc)

918 - else:

919 - with pool.acquire(raw_frame_bytes):

920 - frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc)

921 - finally:

922 - with suppress(FileNotFoundError):

923 - os.unlink(temp_path)

924 -

925 - valid_frame_indices = frame_idx[: int(frames.shape[0])]

926 - return frames, source, frame_idx, valid_frame_indices

927 -

928 -

929 -class DeepStreamVideoBackendMixin:

930 - """NVIDIA DeepStream (NVDEC) GPU-decode codec utilities.

931 -

932 - Decoding runs on a shared pool of daemon threads inside one CUDA

933 - context (see the ``nvidia-deepstream-videodecode-cu13`` package). The

934 - container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no

935 - local file path is required — HTTP and base64 sources decode identically

936 - to local files.

937 -

938 - Like the OpenCV/PyAV mixins, this provides only the codec layer.

939 - Frame *selection* lives in the loader's

940 - ``compute_frames_index_to_sample`` and arrives here as an explicit

941 - list of frame indices.

942 - """

943 -

944 - # Process-wide lazy decode pool, shared across all DeepStream backends.

945 - _pool: ClassVar[Any] = None

946 - _pool_lock: ClassVar[Any] = None

947 -

948 - @classmethod

949 - def _get_pool(cls, pool_size: int | None = None):

950 - """Lazy-initialize the shared decode pool on first use.

951 -

952 - ``pool_size`` (number of decode worker threads) comes from

953 - ``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it

954 - defaults to the existing ``VLLM_MEDIA_LOADING_THREAD_COUNT`` so no

955 - DeepStream-specific env var is needed. The pool is a process-wide

956 - singleton, so the first decode's value wins.

957 - """

958 - if cls._pool is not None:

959 - return cls._pool

960 - if cls._pool_lock is None:

961 - cls._pool_lock = threading.Lock()

962 - with cls._pool_lock:

963 - if cls._pool is not None:

964 - return cls._pool

965 - import os

966 -

967 - from nvidia.deepstream_videodecode import DecodePool

968 -

969 - if pool_size is None:

970 - pool_size = int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", 8))

971 - pool_size = max(1, min(int(pool_size), 16))

972 - logger.info(

973 - "[DeepStream] initializing decode pool with %d workers",

974 - pool_size,

975 - )

976 - cls._pool = DecodePool(num_workers=pool_size)

977 - return cls._pool

978 -

979 - @classmethod

980 - def decode_indices(

981 - cls,

982 - data: bytes,

983 - frame_indices: list[int],

984 - source: VideoSourceMetadata,

985 - codec: str = "",

986 - pool_size: int | None = None,

987 - timeout_sec: float = 120.0,

988 - ) -> tuple[npt.NDArray, list[int]]:

989 - """Decode the requested frame indices from raw container bytes.

990 -

991 - The whole stream is decoded; the pool keeps exactly the frames whose

992 - decode-order index is in ``frame_indices`` (1:1, frame-exact) and

993 - sends EOS once the last one is matched.

994 -

995 - ``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC

996 - session warm across same-codec streams and rebuild only on a codec

997 - change. Frames are returned as a CPU NHWC uint8 array so the

998 - upstream multimodal parser sees the same shape as the other

999 - backends.

1000 - """

1001 - if not frame_indices:

1002 - raise ValueError("DeepStream backend received no frame indices")

1003 -

1004 - result = cls._get_pool(pool_size).decode(

1005 - data,

1006 - target_indices=frame_indices,

1007 - codec=codec,

1008 - max_frames=len(frame_indices),

1009 - timeout_sec=timeout_sec,

1010 - )

1011 - if result.error:

1012 - raise ValueError(f"DeepStream decode failed: {result.error}")

1013 - if result.frames is None or result.n_kept == 0:

1014 - raise ValueError("DeepStream decode produced no frames")

1015 -

1016 - valid = frame_indices[: result.n_kept]

1017 - # GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps

1018 - # the array shape identical to the OpenCV/PyAV backends. Copy into

1019 - # PINNED host memory (reused across calls by PyTorch's pinned caching

1020 - # allocator) so the D2H runs at full PCIe bandwidth (~13 GB/s) rather

1021 - # than the ~1 GB/s pageable path that plain ``.cpu()`` takes — ~12x

1022 - # faster for a 1080p x8 frame batch (~46ms -> ~4ms). ``numpy()`` keeps

1023 - # the pinned tensor alive via the array's base.

1024 - import torch

1025 -

1026 - gpu = result.frames

1027 - if gpu.is_cuda:

1028 - host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True)

1029 - host.copy_(gpu, non_blocking=True)

1030 - torch.cuda.current_stream().synchronize()

1031 - arr = host.numpy()

1032 - else:

1033 - arr = gpu.numpy()

1034 - return arr, valid

1035 -

173

174 @VIDEO_LOADER_REGISTRY.register("opencv")

1038 -class VideoBackend(

1039 - VideoLoader,

1040 - OpenCVVideoBackendMixin,

1041 - PyAVVideoBackendMixin,

1042 - TorchCodecVideoBackendMixin,

1043 - PyNvVideoCodecVideoBackendMixin,

1044 - DeepStreamVideoBackendMixin,

1045 -):

175 +class VideoBackend(VideoLoader):

176 """Uniform-sampling video backend.

177

178 Samples ``num_frames`` uniformly across the video (or one frame every

@@ -1080,11 +210,6 @@ class VideoBackend(

210 0, total_frames_num - 1, num_frames_to_sample, dtype=int

211 ).tolist()

212

1083 - @classmethod

1084 - def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata:

1085 - """Sampling-algorithm-specific metadata adjustment hook."""

1086 - return source

1087 -

213 @classmethod

214 def load_bytes(

215 cls,

@@ -1094,12 +219,7 @@ class VideoBackend(

219 max_duration: int = 300,

220 frame_recovery: bool = False,

221 *,

1097 - backend: Literal[

1098 - "opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"

1099 - ] = "opencv",

1100 - num_ffmpeg_threads: int = 0,

1101 - seek_mode: Literal["exact", "approximate"] = "exact",

1102 - hw_decoders: int = PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,

222 + backend: VideoDecoderBackend = "opencv",

223 **kwargs,

224 ) -> tuple[npt.NDArray, dict[str, Any]]:

225 """Load sampled frames from raw video bytes.

@@ -1114,20 +234,25 @@ class VideoBackend(

234 Only honored by the OpenCV codec.

235 backend: Decoding codec — ``"opencv"``, ``"pyav"``,

236 ``"torchcodec"``, ``"pynvvideocodec"`` or ``"deepstream"``.

1117 - num_ffmpeg_threads: Number of FFmpeg decoding threads, only used by

1118 - TorchCodec: ``0`` (default) relies on the FFmpeg default value

1119 - which is ``min(cpu_count + 1, 16)``.

1120 - OpenCV will always use ``min(cpu_count, 16)`` while pyav will

1121 - always use ``min(cpu_count, (height + 15) / 16)``.

1122 - seek_mode: Seek mode for the TorchCodec decoder, only used by

1123 - TorchCodec: ``"exact"`` (default) guarantees frame-accurate

1124 - sampling by scanning the file on creation, while

1125 - ``"approximate"`` skips that scan for faster decoder creation

1126 - at the cost of relying on the file's metadata. See

1127 - https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html

1128 - for details.

1129 - hw_decoders: Maximum number of concurrent PyNvVideoCodec decoder

1130 - slots. Defaults to 2 and must be a positive integer.

237 + kwargs: Codec-specific options, validated against and forwarded to

238 + ``backend``:

239 +

240 + - ``num_ffmpeg_threads`` (TorchCodec): number of FFmpeg

241 + decoding threads; ``0`` (default) relies on the FFmpeg

242 + default value which is ``min(cpu_count + 1, 16)``.

243 + OpenCV will always use ``min(cpu_count, 16)`` while pyav

244 + will always use ``min(cpu_count, (height + 15) / 16)``.

245 + - ``seek_mode`` (TorchCodec): ``"exact"`` (default) guarantees

246 + frame-accurate sampling by scanning the file on creation,

247 + while ``"approximate"`` skips that scan for faster decoder

248 + creation at the cost of relying on the file's metadata. See

249 + https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html

250 + for details.

251 + - ``hw_decoders`` (PyNvVideoCodec): maximum number of

252 + concurrent decoder slots. Defaults to 2 and must be a

253 + positive integer.

254 + - ``pool_size`` / ``timeout_sec`` (DeepStream): decoder pool

255 + size and pool acquisition timeout in seconds.

256

257 Returns:

258 Tuple of ``(frames_array, metadata_dict)``.

@@ -1135,101 +260,16 @@ class VideoBackend(

260 target = VideoTargetMetadata(

261 num_frames=num_frames, fps=fps, max_duration=max_duration

262 )

1138 -

1139 - if backend == "opencv":

1140 - cap = cls.open_video_capture(data)

1141 - _check_frame_pixel_limit(

1142 - int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),

1143 - int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),

1144 - )

1145 - source = cls._prepare_source(cls.get_video_metadata(cap))

1146 - frame_idx = cls.compute_frames_index_to_sample(

1147 - source=source, target=target, **kwargs

1148 - )

1149 - frames, valid = cls.read_frames(

1150 - cap,

1151 - frame_idx,

1152 - total_frames_num=source.total_frames_num,

1153 - frame_recovery=frame_recovery,

1154 - )

1155 - elif backend == "pyav":

1156 - assert not frame_recovery, (

1157 - "frame_recovery is only available for `opencv` backend"

1158 - )

1159 - with av.open(BytesIO(data)) as container:

1160 - stream = container.streams.video[0]

1161 - _check_frame_pixel_limit(stream.width, stream.height)

1162 - source = cls._prepare_source(cls.get_metadata(container))

1163 - frame_idx = cls.compute_frames_index_to_sample(

1164 - source=source, target=target, **kwargs

1165 - )

1166 - frames, valid = cls.decode_frames(

1167 - container, frame_idx, source.original_fps, source.duration

1168 - )

1169 - elif backend == "torchcodec":

1170 - assert not frame_recovery, (

1171 - "frame_recovery is only available for `opencv` backend"

1172 - )

1173 - check_torchcodec_available()

1174 - decoder = cls.make_torchcodec_decoder(

1175 - data,

1176 - num_ffmpeg_threads=num_ffmpeg_threads,

1177 - seek_mode=seek_mode,

1178 - )

1179 - _check_frame_pixel_limit(

1180 - decoder.metadata.width or 0,

1181 - decoder.metadata.height or 0,

1182 - )

1183 - source = cls._prepare_source(cls.get_torchcodec_metadata(decoder))

1184 - frame_idx = cls.compute_frames_index_to_sample(

1185 - source=source, target=target, **kwargs

1186 - )

1187 - frames, valid = cls.decode_torchcodec_frames(decoder, frame_idx)

1188 - elif backend == PYNVVIDEOCODEC_VIDEO_BACKEND:

1189 - if frame_recovery:

1190 - raise ValueError(

1191 - "frame_recovery is not supported for "

1192 - f"`{PYNVVIDEOCODEC_VIDEO_BACKEND}` backend"

1193 - )

1194 - cls._configure_decoder_slots(hw_decoders)

1195 - frames, source, frame_idx, valid = cls.decode_frames_pynvvideocodec(

1196 - data,

1197 - target,

1198 - **kwargs,

1199 - )

1200 - elif backend == "deepstream":

1201 - assert not frame_recovery, (

1202 - "frame_recovery is only available for `opencv` backend"

1203 - )

1204 - # Decode-pool size comes from media-io-kwargs (no env var); the

1205 - # pool is a process-wide singleton so the first decode's value

1206 - # wins. Pop it so it isn't forwarded to the frame sampler.

1207 - pool_size = kwargs.pop("pool_size", None)

1208 - # Probe container metadata from the bytes via GStreamer (in

1209 - # the deepstream video-decode wheel) — no PyAV/pymediainfo, no path.

1210 - from nvidia.deepstream_videodecode import probe_metadata

1211 -

1212 - total_frames, original_fps, duration, _w, _h, codec = probe_metadata(data)

1213 - _check_frame_pixel_limit(_w, _h)

1214 - source = cls._prepare_source(

1215 - VideoSourceMetadata(

1216 - total_frames_num=total_frames,

1217 - original_fps=original_fps,

1218 - duration=duration,

1219 - )

1220 - )

1221 - frame_idx = cls.compute_frames_index_to_sample(

1222 - source=source, target=target, **kwargs

1223 - )

1224 - frames, valid = cls.decode_indices(

1225 - data, frame_idx, source, codec=codec, pool_size=pool_size

1226 - )

1227 - else:

1228 - raise ValueError(

1229 - f"Unknown video codec backend {backend!r}; "

1230 - "valid options: 'opencv', 'pyav', 'torchcodec', "

1231 - "'pynvvideocodec' and 'deepstream'."

1232 - )

263 + sampling_kwargs, backend_kwargs = resolve_video_backend_kwargs(backend, kwargs)

264 + frames, source, frame_idx, valid = decode_video(

265 + backend,

266 + cls,

267 + data,

268 + target,

269 + sampling_kwargs,

270 + backend_kwargs,

271 + frame_recovery=frame_recovery,

272 + )

273

274 if len(valid) < len(frame_idx):

275 logger.warning(

@@ -1742,7 +782,7 @@ class GLMGAVideoBackend(VideoBackend):

782 "molmo2",

783 video_processor="Molmo2VideoProcessor",

784 )

1745 -class Molmo2VideoBackend(VideoLoader, OpenCVVideoBackendMixin):

785 +class Molmo2VideoBackend(VideoLoader):

786 @classmethod

787 def get_candidate_target_fps(

788 cls,

@@ -1975,30 +1015,21 @@ class Molmo2VideoBackend(VideoLoader, OpenCVVideoBackendMixin):

1015 frame_recovery: bool = False,

1016 **kwargs,

1017 ) -> tuple[npt.NDArray, dict[str, Any]]:

1978 - cap = cls.open_video_capture(data)

1979 - _check_frame_pixel_limit(

1980 - int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),

1981 - int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),

1982 - )

1983 -

1984 - source = OpenCVVideoBackendMixin.get_video_metadata(cap)

1018 target = VideoTargetMetadata(

1019 num_frames=num_frames,

1020 fps=sampling_fps,

1988 - max_duration=source.duration,

1021 + max_duration=-1,

1022 )

1990 -

1991 - frame_idx = cls.compute_frames_index_to_sample(

1992 - source=source,

1993 - target=target,

1994 - frame_sample_mode=frame_sample_mode,

1995 - max_fps=max_fps,

1996 - )

1997 -

1998 - frames, valid_frame_indices = cls.read_frames(

1999 - cap,

2000 - frame_idx,

2001 - total_frames_num=source.total_frames_num,

1023 + frames, source, _, valid_frame_indices = decode_video(

1024 + "opencv",

1025 + cls,

1026 + data,

1027 + target,

1028 + {

1029 + "frame_sample_mode": frame_sample_mode,

1030 + "max_fps": max_fps,

1031 + },

1032 + {},

1033 frame_recovery=frame_recovery,

1034 )

1035

@@ -2064,7 +1095,7 @@ class NemotronVLVideoBackend(VideoBackend):

1095

1096

1097 @VIDEO_LOADER_REGISTRY.register("openpangu")

2067 -class OpenCVDynamicOpenPanguVideoBackend(VideoLoader, OpenCVVideoBackendMixin):

1098 +class OpenCVDynamicOpenPanguVideoBackend(VideoLoader):

1099 @classmethod

1100 def compute_frames_index_to_sample(

1101 cls,

@@ -2130,14 +1161,6 @@ class OpenCVDynamicOpenPanguVideoBackend(VideoLoader, OpenCVVideoBackendMixin):

1161 Returns:

1162 Tuple of (frames_array, metadata_dict)

1163 """

2133 - cap = cls.open_video_capture(data)

2134 - _check_frame_pixel_limit(

2135 - int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),

2136 - int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),

2137 - )

2138 -

2139 - source = OpenCVVideoBackendMixin.get_video_metadata(cap)

2140 -

1164 # recompute source metadata with adjusted duration to ensure correct

1165 # sampling indices computation

1166 target = VideoTargetMetadata(

@@ -2146,15 +1169,13 @@ class OpenCVDynamicOpenPanguVideoBackend(VideoLoader, OpenCVVideoBackendMixin):

1169 max_duration=max_duration,

1170 )

1171

2149 - frame_indices_list = cls.compute_frames_index_to_sample(

2150 - source=source,

2151 - target=target,

2152 - )

2153 -

2154 - frames, valid_frame_indices = cls.read_frames(

2155 - cap,

2156 - frame_indices_list,

2157 - total_frames_num=source.total_frames_num,

1172 + frames, source, _, valid_frame_indices = decode_video(

1173 + "opencv",

1174 + cls,

1175 + data,

1176 + target,

1177 + {},

1178 + {},

1179 frame_recovery=frame_recovery,

1180 )

1181

vllm/multimodal/video_decoders/__init__.py new

+110

@@ -0,0 +1,110 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from importlib import import_module

5 +from typing import Any, Literal

6 +

7 +from .base import (

8 + PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,

9 + PYNVVIDEOCODEC_VIDEO_BACKEND,

10 + VideoSourceMetadata,

11 + VideoTargetMetadata,

12 + check_frame_pixel_limit,

13 +)

14 +

15 +VideoDecoderBackend = Literal[

16 + "opencv", "pyav", "torchcodec", "pynvvideocodec", "deepstream"

17 +]

18 +

19 +_BACKEND_OPTION_DEFAULTS: dict[str, dict[str, Any]] = {

20 + "opencv": {},

21 + "pyav": {},

22 + "torchcodec": {

23 + "num_ffmpeg_threads": 0,

24 + "seek_mode": "exact",

25 + },

26 + PYNVVIDEOCODEC_VIDEO_BACKEND: {

27 + "hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,

28 + },

29 + "deepstream": {

30 + "pool_size": None,

31 + "timeout_sec": 120.0,

32 + },

33 +}

34 +

35 +

36 +def _get_backend_option_defaults(backend: str) -> dict[str, Any]:

37 + try:

38 + return _BACKEND_OPTION_DEFAULTS[backend]

39 + except KeyError:

40 + valid = ", ".join(repr(name) for name in _BACKEND_OPTION_DEFAULTS)

41 + raise ValueError(

42 + f"Unknown video codec backend {backend!r}; valid options: {valid}."

43 + ) from None

44 +

45 +

46 +def resolve_video_backend_kwargs(

47 + backend: str,

48 + kwargs: dict[str, Any],

49 +) -> tuple[dict[str, Any], dict[str, Any]]:

50 + """Split frame-sampling kwargs from options owned by a decoder backend."""

51 + defaults = _get_backend_option_defaults(backend)

52 + sampling_kwargs = dict(kwargs)

53 + backend_kwargs = dict(defaults)

54 + backend_option_names = {

55 + name for options in _BACKEND_OPTION_DEFAULTS.values() for name in options

56 + }

57 + misplaced = (sampling_kwargs.keys() & backend_option_names) - defaults.keys()

58 + if misplaced:

59 + names = ", ".join(sorted(misplaced))

60 + raise ValueError(f"{names} is not supported by the {backend!r} backend")

61 +

62 + for name in defaults:

63 + if name in sampling_kwargs:

64 + backend_kwargs[name] = sampling_kwargs.pop(name)

65 +

66 + return sampling_kwargs, backend_kwargs

67 +

68 +

69 +def decode_video(

70 + backend: str,

71 + loader_cls,

72 + data: bytes,

73 + target: VideoTargetMetadata,

74 + sampling_kwargs: dict[str, Any],

75 + backend_kwargs: dict[str, Any],

76 + *,

77 + frame_recovery: bool,

78 +):

79 + """Decode a sampled video, importing only the selected backend."""

80 + _get_backend_option_defaults(backend)

81 + if frame_recovery and backend != "opencv":

82 + error = (

83 + ValueError if backend == PYNVVIDEOCODEC_VIDEO_BACKEND else AssertionError

84 + )

85 + raise error(f"frame_recovery is not supported by the {backend!r} backend")

86 +

87 + decoder_kwargs = dict(backend_kwargs)

88 + if backend == "opencv":

89 + decoder_kwargs["frame_recovery"] = frame_recovery

90 + module = import_module(f".{backend}", __name__)

91 + decoder = getattr(module, f"decode_{backend}")

92 + return decoder(

93 + loader_cls,

94 + data,

95 + target,

96 + sampling_kwargs,

97 + **decoder_kwargs,

98 + )

99 +

100 +

101 +__all__ = [

102 + "PYNVVIDEOCODEC_DEFAULT_HW_DECODERS",

103 + "PYNVVIDEOCODEC_VIDEO_BACKEND",

104 + "VideoDecoderBackend",

105 + "VideoSourceMetadata",

106 + "VideoTargetMetadata",

107 + "check_frame_pixel_limit",

108 + "decode_video",

109 + "resolve_video_backend_kwargs",

110 +]

vllm/multimodal/video_decoders/base.py new

+37

@@ -0,0 +1,37 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from typing import Literal, NamedTuple

5 +

6 +from vllm import envs

7 +

8 +PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec"

9 +PYNVVIDEOCODEC_DEFAULT_HW_DECODERS = 2

10 +

11 +

12 +class VideoTargetMetadata(NamedTuple):

13 + """Metadata describing the requested video sample."""

14 +

15 + num_frames: int

16 + fps: float

17 + max_duration: float

18 +

19 +

20 +class VideoSourceMetadata(NamedTuple):

21 + """Metadata describing the encoded video source."""

22 +

23 + total_frames_num: int

24 + original_fps: float

25 + duration: float

26 +

27 +

28 +def check_frame_pixel_limit(width: int, height: int) -> None:

29 + """Reject video frames exceeding ``VLLM_MAX_IMAGE_PIXELS``."""

30 + max_pixels = envs.VLLM_MAX_IMAGE_PIXELS

31 + if max_pixels > 0 and width * height > max_pixels:

32 + raise ValueError(

33 + f"Video frame dimensions {width}x{height} "

34 + f"({width * height} pixels) exceed the maximum of "

35 + f"{max_pixels} pixels. Set VLLM_MAX_IMAGE_PIXELS to "

36 + f"increase this limit."

37 + )

vllm/multimodal/video_decoders/deepstream.py new

+151

@@ -0,0 +1,151 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import threading

5 +from typing import Any, ClassVar

6 +

7 +import numpy.typing as npt

8 +

9 +from vllm.logger import init_logger

10 +

11 +from .base import VideoSourceMetadata, VideoTargetMetadata, check_frame_pixel_limit

12 +

13 +logger = init_logger(__name__)

14 +

15 +

16 +def decode_deepstream(

17 + loader_cls,

18 + data: bytes,

19 + target: VideoTargetMetadata,

20 + sampling_kwargs: dict,

21 + *,

22 + pool_size: int | None = None,

23 + timeout_sec: float = 120.0,

24 +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:

25 + from nvidia.deepstream_videodecode import probe_metadata

26 +

27 + total_frames, original_fps, duration, _width, _height, codec = probe_metadata(data)

28 + check_frame_pixel_limit(_width, _height)

29 + source = loader_cls._prepare_source(

30 + VideoSourceMetadata(total_frames, original_fps, duration)

31 + )

32 + frame_idx = loader_cls.compute_frames_index_to_sample(

33 + source=source, target=target, **sampling_kwargs

34 + )

35 + frames, valid = DeepStreamVideoBackendMixin.decode_indices(

36 + data,

37 + frame_idx,

38 + source,

39 + codec=codec,

40 + pool_size=pool_size,

41 + timeout_sec=timeout_sec,

42 + )

43 + return frames, source, frame_idx, valid

44 +

45 +

46 +class DeepStreamVideoBackendMixin:

47 + """NVIDIA DeepStream (NVDEC) GPU-decode codec utilities.

48 +

49 + Decoding runs on a shared pool of daemon threads inside one CUDA

50 + context (see the ``nvidia-deepstream-videodecode-cu13`` package). The

51 + container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no

52 + local file path is required — HTTP and base64 sources decode identically

53 + to local files.

54 +

55 + Like the OpenCV/PyAV mixins, this provides only the codec layer.

56 + Frame *selection* lives in the loader's

57 + ``compute_frames_index_to_sample`` and arrives here as an explicit

58 + list of frame indices.

59 + """

60 +

61 + # Process-wide lazy decode pool, shared across all DeepStream backends.

62 + _pool: ClassVar[Any] = None

63 + _pool_lock: ClassVar[Any] = None

64 +

65 + @classmethod

66 + def _get_pool(cls, pool_size: int | None = None):

67 + """Lazy-initialize the shared decode pool on first use.

68 +

69 + ``pool_size`` (number of decode worker threads) comes from

70 + ``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it

71 + defaults to the existing ``VLLM_MEDIA_LOADING_THREAD_COUNT`` so no

72 + DeepStream-specific env var is needed. The pool is a process-wide

73 + singleton, so the first decode's value wins.

74 + """

75 + if cls._pool is not None:

76 + return cls._pool

77 + if cls._pool_lock is None:

78 + cls._pool_lock = threading.Lock()

79 + with cls._pool_lock:

80 + if cls._pool is not None:

81 + return cls._pool

82 + import os

83 +

84 + from nvidia.deepstream_videodecode import DecodePool

85 +

86 + if pool_size is None:

87 + pool_size = int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", 8))

88 + pool_size = max(1, min(int(pool_size), 16))

89 + logger.info(

90 + "[DeepStream] initializing decode pool with %d workers",

91 + pool_size,

92 + )

93 + cls._pool = DecodePool(num_workers=pool_size)

94 + return cls._pool

95 +

96 + @classmethod

97 + def decode_indices(

98 + cls,

99 + data: bytes,

100 + frame_indices: list[int],

101 + source: VideoSourceMetadata,

102 + codec: str = "",

103 + pool_size: int | None = None,

104 + timeout_sec: float = 120.0,

105 + ) -> tuple[npt.NDArray, list[int]]:

106 + """Decode the requested frame indices from raw container bytes.

107 +

108 + The whole stream is decoded; the pool keeps exactly the frames whose

109 + decode-order index is in ``frame_indices`` (1:1, frame-exact) and

110 + sends EOS once the last one is matched.

111 +

112 + ``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC

113 + session warm across same-codec streams and rebuild only on a codec

114 + change. Frames are returned as a CPU NHWC uint8 array so the

115 + upstream multimodal parser sees the same shape as the other

116 + backends.

117 + """

118 + if not frame_indices:

119 + raise ValueError("DeepStream backend received no frame indices")

120 +

121 + result = cls._get_pool(pool_size).decode(

122 + data,

123 + target_indices=frame_indices,

124 + codec=codec,

125 + max_frames=len(frame_indices),

126 + timeout_sec=timeout_sec,

127 + )

128 + if result.error:

129 + raise ValueError(f"DeepStream decode failed: {result.error}")

130 + if result.frames is None or result.n_kept == 0:

131 + raise ValueError("DeepStream decode produced no frames")

132 +

133 + valid = frame_indices[: result.n_kept]

134 + # GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps

135 + # the array shape identical to the OpenCV/PyAV backends. Copy into

136 + # PINNED host memory (reused across calls by PyTorch's pinned caching

137 + # allocator) so the D2H runs at full PCIe bandwidth (~13 GB/s) rather

138 + # than the ~1 GB/s pageable path that plain ``.cpu()`` takes — ~12x

139 + # faster for a 1080p x8 frame batch (~46ms -> ~4ms). ``numpy()`` keeps

140 + # the pinned tensor alive via the array's base.

141 + import torch

142 +

143 + gpu = result.frames

144 + if gpu.is_cuda:

145 + host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True)

146 + host.copy_(gpu, non_blocking=True)

147 + torch.cuda.current_stream().synchronize()

148 + arr = host.numpy()

149 + else:

150 + arr = gpu.numpy()

151 + return arr, valid

vllm/multimodal/video_decoders/opencv.py new

+295

@@ -0,0 +1,295 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from io import BytesIO

5 +

6 +import numpy as np

7 +import numpy.typing as npt

8 +

9 +from vllm.logger import init_logger

10 +from vllm.utils.import_utils import PlaceholderModule

11 +

12 +from .base import (

13 + VideoSourceMetadata,

14 + VideoTargetMetadata,

15 + check_frame_pixel_limit,

16 +)

17 +

18 +try:

19 + import cv2

20 + import cv2.videoio_registry as vr

21 +except ImportError:

22 + cv2 = PlaceholderModule("cv2")

23 + vr = PlaceholderModule("cv2").placeholder_attr("videoio_registry")

24 +

25 +logger = init_logger(__name__)

26 +

27 +

28 +def decode_opencv(

29 + loader_cls,

30 + data: bytes,

31 + target: VideoTargetMetadata,

32 + sampling_kwargs: dict,

33 + *,

34 + frame_recovery: bool = False,

35 +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:

36 + cap = OpenCVVideoBackendMixin.open_video_capture(data)

37 + check_frame_pixel_limit(

38 + int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),

39 + int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),

40 + )

41 + source = loader_cls._prepare_source(OpenCVVideoBackendMixin.get_video_metadata(cap))

42 + frame_idx = loader_cls.compute_frames_index_to_sample(

43 + source=source, target=target, **sampling_kwargs

44 + )

45 + frames, valid = OpenCVVideoBackendMixin.read_frames(

46 + cap,

47 + frame_idx,

48 + total_frames_num=source.total_frames_num,

49 + frame_recovery=frame_recovery,

50 + )

51 + return frames, source, frame_idx, valid

52 +

53 +

54 +class OpenCVVideoBackendMixin:

55 + @staticmethod

56 + def get_cv2_video_api():

57 + api_pref = None

58 + for backend in vr.getStreamBufferedBackends():

59 + if not vr.hasBackend(backend):

60 + continue

61 + if not vr.isBackendBuiltIn(backend):

62 + _, abi, api = vr.getStreamBufferedBackendPluginVersion(backend)

63 + if abi < 1 or (abi == 1 and api < 2):

64 + continue

65 + api_pref = backend

66 + break

67 + return api_pref

68 +

69 + @classmethod

70 + def open_video_capture(cls, data: bytes) -> "cv2.VideoCapture":

71 + backend = cls.get_cv2_video_api()

72 + cap = cv2.VideoCapture(BytesIO(data), backend, [])

73 + if not cap.isOpened():

74 + raise ValueError("Could not open video stream")

75 + return cap

76 +

77 + @staticmethod

78 + def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata:

79 + total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

80 + original_fps = cap.get(cv2.CAP_PROP_FPS)

81 + duration = total_frames_num / original_fps if original_fps > 0 else 0

82 + return VideoSourceMetadata(

83 + total_frames_num=total_frames_num,

84 + original_fps=original_fps,

85 + duration=duration,

86 + )

87 +

88 + @classmethod

89 + def _can_use_for_recovery(

90 + cls,

91 + idx: int,

92 + failed_frames: list[int],

93 + next_target_map: dict[int, int],

94 + total_frames: int,

95 + ) -> bool:

96 + """Check if current frame can recover the oldest failed frame."""

97 + if not failed_frames:

98 + return False

99 + oldest_failed = failed_frames[0]

100 + limit = next_target_map.get(oldest_failed, total_frames)

101 + return idx < limit

102 +

103 + @classmethod

104 + def _read_frames_with_recovery(

105 + cls,

106 + cap: "cv2.VideoCapture",

107 + frame_indices: list[int],

108 + total_frames: int,

109 + ) -> tuple[npt.NDArray, list[int], dict[int, int]]:

110 + """

111 + Read frames with dynamic window forward-scan recovery.

112 +

113 + When a target frame fails to load, the next successfully grabbed

114 + frame (before the next target frame) will be used to recover it.

115 +

116 + Args:

117 + cap: OpenCV VideoCapture object

118 + frame_indices: Sorted list of target frame indices to load

119 + total_frames: Total number of frames in the video

120 +

121 + Returns:

122 + Tuple of (frames_array, valid_frame_indices, recovered_map)

123 + - frames_array: Array of loaded frames

124 + - valid_frame_indices: List of frame indices that were loaded

125 + - recovered_map: Dict mapping recovered_idx -> source_idx

126 + """

127 + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))

128 + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

129 +

130 + assert width > 0 and height > 0, (

131 + f"Invalid video frame size: width={width}, height={height}"

132 + )

133 +

134 + frame_idx_set = set(frame_indices)

135 + max_frame_idx = frame_indices[-1] if frame_indices else 0

136 +

137 + # Build map: target_idx -> next_target_idx (for recovery window)

138 + next_target_map: dict[int, int] = {}

139 + for k in range(len(frame_indices) - 1):

140 + next_target_map[frame_indices[k]] = frame_indices[k + 1]

141 + next_target_map[frame_indices[-1]] = total_frames

142 +

143 + frames_list: list[npt.NDArray] = []

144 + valid_frame_indices: list[int] = []

145 + failed_frames_idx: list[int] = []

146 + recovered_map: dict[int, int] = {}

147 +

148 + i = 0

149 + for idx in range(max_frame_idx + 1):

150 + is_target_frame = idx in frame_idx_set

151 +

152 + # Attempt to grab the current frame

153 + ok = cap.grab()

154 +

155 + if not ok:

156 + if is_target_frame:

157 + logger.debug(

158 + "Failed to grab frame %d during video loading.",

159 + idx,

160 + )

161 + failed_frames_idx.append(idx)

162 + continue

163 +

164 + # Check if we should retrieve: target frame OR can recover a failed one

165 + can_recover = cls._can_use_for_recovery(

166 + idx, failed_frames_idx, next_target_map, total_frames

167 + )

168 +

169 + if is_target_frame or can_recover:

170 + ret, frame = cap.retrieve()

171 +

172 + if ret and frame is not None and frame.size > 0:

173 + rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

174 + frames_list.append(rgb_frame)

175 + valid_frame_indices.append(idx)

176 + i += 1

177 +

178 + if can_recover:

179 + recovered_idx = failed_frames_idx.pop(0)

180 + recovered_map[recovered_idx] = idx

181 + logger.info(

182 + "Recovered frame %d using frame %d (delay: %d)",

183 + recovered_idx,

184 + idx,

185 + idx - recovered_idx,

186 + )

187 + elif is_target_frame:

188 + logger.debug(

189 + "Failed to retrieve frame %d during video loading.",

190 + idx,

191 + )

192 + failed_frames_idx.append(idx)

193 +

194 + # Log any remaining failed frames

195 + for failed_idx in failed_frames_idx:

196 + logger.debug(

197 + "Frame %d could not be recovered (end of video).",

198 + failed_idx,

199 + )

200 +

201 + # Stack frames

202 + if frames_list:

203 + frames = np.stack(frames_list)

204 + else:

205 + frames = np.empty((0, height, width, 3), dtype=np.uint8)

206 +

207 + return frames, valid_frame_indices, recovered_map

208 +

209 + @classmethod

210 + def _read_frames_no_recovery(

211 + cls,

212 + cap,

213 + frame_indices: set[int],

214 + max_frame_idx: int,

215 + ) -> tuple[npt.NDArray, list[int]]:

216 + num_expected_frames = len(frame_indices)

217 + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))

218 + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

219 + frames = np.empty((num_expected_frames, height, width, 3), dtype=np.uint8)

220 +

221 + i = 0

222 + valid_frame_indices = []

223 + for idx in range(max_frame_idx + 1):

224 + ok = cap.grab()

225 + if not ok:

226 + # Frame is broken/unreadable, skip it

227 + if idx in frame_indices:

228 + logger.debug(

229 + "Failed to grab frame %d during video loading. "

230 + "This frame will be skipped.",

231 + idx,

232 + )

233 + continue

234 + if idx in frame_indices:

235 + ret, frame = cap.retrieve()

236 + if ret:

237 + frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

238 + valid_frame_indices.append(idx)

239 + i += 1

240 + else:

241 + # retrieve() failed even though grab() succeeded

242 + logger.debug(

243 + "Failed to retrieve frame %d during video loading. "

244 + "This frame will be skipped.",

245 + idx,

246 + )

247 +

248 + valid_num_frames = len(valid_frame_indices)

249 + if valid_num_frames < num_expected_frames:

250 + logger.warning(

251 + "Video loading completed with %d broken/unreadable frames. "

252 + "Expected %d frames but only loaded %d frames.",

253 + num_expected_frames - valid_num_frames,

254 + num_expected_frames,

255 + valid_num_frames,

256 + )

257 +

258 + return frames[:valid_num_frames], valid_frame_indices

259 +

260 + @classmethod

261 + def read_frames(

262 + cls,

263 + cap: "cv2.VideoCapture",

264 + frame_idx: list[int],

265 + total_frames_num: int,

266 + *,

267 + frame_recovery: bool = False,

268 + ) -> tuple[npt.NDArray, list[int]]:

269 + if frame_recovery:

270 + num_frames_to_sample = len(frame_idx)

271 + frames, valid_frame_indices, recovered_map = cls._read_frames_with_recovery(

272 + cap, frame_idx, total_frames_num

273 + )

274 +

275 + if recovered_map:

276 + logger.info(

277 + "Frame recovery: %d frames recovered using forward scan.",

278 + len(recovered_map),

279 + )

280 + else:

281 + frame_idx_set = set(frame_idx)

282 + num_frames_to_sample = len(frame_idx_set)

283 + frames, valid_frame_indices = cls._read_frames_no_recovery(

284 + cap, frame_idx_set, max(frame_idx)

285 + )

286 + valid_num_frames = len(valid_frame_indices)

287 + if valid_num_frames < num_frames_to_sample:

288 + logger.warning(

289 + "Video loading completed with %d broken/unreadable frames. "

290 + "Expected to sample %d frames but only loaded %d frames.",

291 + num_frames_to_sample - valid_num_frames,

292 + num_frames_to_sample,

293 + valid_num_frames,

294 + )

295 + return frames, valid_frame_indices

vllm/multimodal/video_decoders/pyav.py new

+111

@@ -0,0 +1,111 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from io import BytesIO

5 +

6 +import numpy as np

7 +import numpy.typing as npt

8 +

9 +from vllm.utils.import_utils import PlaceholderModule

10 +

11 +from .base import (

12 + VideoSourceMetadata,

13 + VideoTargetMetadata,

14 + check_frame_pixel_limit,

15 +)

16 +

17 +try:

18 + import av

19 +except ImportError:

20 + av = PlaceholderModule("av") # type: ignore[assignment]

21 +

22 +

23 +def decode_pyav(

24 + loader_cls,

25 + data: bytes,

26 + target: VideoTargetMetadata,

27 + sampling_kwargs: dict,

28 +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:

29 + with av.open(BytesIO(data)) as container:

30 + stream = container.streams.video[0]

31 + check_frame_pixel_limit(stream.width, stream.height)

32 + source = loader_cls._prepare_source(

33 + PyAVVideoBackendMixin.get_metadata(container)

34 + )

35 + frame_idx = loader_cls.compute_frames_index_to_sample(

36 + source=source, target=target, **sampling_kwargs

37 + )

38 + frames, valid = PyAVVideoBackendMixin.decode_frames(

39 + container, frame_idx, source.original_fps, source.duration

40 + )

41 + return frames, source, frame_idx, valid

42 +

43 +

44 +class PyAVVideoBackendMixin:

45 + """PyAV (in-process FFmpeg bindings) codec utilities.

46 +

47 + Reads stream metadata and decodes target frames via per-frame

48 + ``container.seek()``. The seek releases the GIL between frames and

49 + scales with the number of sampled frames rather than the video

50 + length, enabling concurrent decoding under serving load.

51 + """

52 +

53 + @staticmethod

54 + def get_metadata(

55 + container: "av.container.InputContainer",

56 + ) -> VideoSourceMetadata:

57 + if not container.streams.video:

58 + raise ValueError("No video streams found in container")

59 + stream = container.streams.video[0]

60 + total_frames = stream.frames or 0

61 + fps = float(stream.average_rate) if stream.average_rate else 0.0

62 + duration = float(stream.duration * stream.time_base) if stream.duration else 0.0

63 + if total_frames == 0 and duration > 0 and fps > 0:

64 + total_frames = int(duration * fps)

65 + return VideoSourceMetadata(total_frames, fps, duration)

66 +

67 + @staticmethod

68 + def decode_frames(

69 + container: "av.container.InputContainer",

70 + frame_indices: list[int],

71 + fps: float,

72 + duration: float,

73 + ) -> tuple[npt.NDArray, list[int]]:

74 + """Decode target frames via per-frame seek + forward decode to PTS."""

75 + stream = container.streams.video[0]

76 + # SLICE parallelizes within a single frame without the

77 + # one-frame-per-thread latency penalty of FRAME threading.

78 + stream.thread_type = "SLICE"

79 + time_base = stream.time_base

80 +

81 + frames_list: list[npt.NDArray] = []

82 + valid_indices: list[int] = []

83 + frame_interval = 1.0 / fps if fps > 0 else 0.1

84 + max_ts = max(0.0, duration - frame_interval) if duration > 0 else float("inf")

85 +

86 + decoder = None

87 + last_pts = None

88 + for idx in frame_indices:

89 + ts = min(idx / fps, max_ts) if fps > 0 else 0.0

90 + pts = int(ts / time_base)

91 + # seek() snaps backward to a keyframe; reuse the running decoder

92 + # while targets advance monotonically to avoid re-decoding the

93 + # GOP prefix once per requested frame.

94 + if decoder is None or last_pts is None or pts <= last_pts:

95 + container.seek(pts, stream=stream)

96 + decoder = container.decode(video=0)

97 + chosen = None

98 + for frame in decoder:

99 + if frame.pts is not None and frame.pts >= pts:

100 + chosen = frame

101 + last_pts = frame.pts

102 + break

103 + if chosen is not None:

104 + frames_list.append(chosen.to_ndarray(format="rgb24"))

105 + valid_indices.append(idx)

106 + else:

107 + decoder = None

108 +

109 + if not frames_list:

110 + return np.empty((0,), dtype=np.uint8), valid_indices

111 + return np.stack(frames_list), valid_indices

vllm/multimodal/video_decoders/pynvvideocodec.py new

+372

@@ -0,0 +1,372 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import os

5 +import tempfile

6 +import threading

7 +from contextlib import contextmanager, suppress

8 +from typing import ClassVar, NamedTuple

9 +

10 +import numpy as np

11 +import numpy.typing as npt

12 +

13 +from vllm.logger import init_logger

14 +from vllm.utils.mem_constants import MiB_bytes

15 +

16 +from .base import (

17 + PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,

18 + VideoSourceMetadata,

19 + VideoTargetMetadata,

20 + check_frame_pixel_limit,

21 +)

22 +

23 +logger = init_logger(__name__)

24 +

25 +

26 +def decode_pynvvideocodec(

27 + loader_cls,

28 + data: bytes,

29 + target: VideoTargetMetadata,

30 + sampling_kwargs: dict,

31 + *,

32 + hw_decoders: int = PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,

33 +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:

34 + PyNvVideoCodecVideoBackendMixin._configure_decoder_slots(hw_decoders)

35 + return PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec(

36 + loader_cls,

37 + data,

38 + target,

39 + **sampling_kwargs,

40 + )

41 +

42 +

43 +class PyNvVideoCodecSourceMetadata(NamedTuple):

44 + """Metadata needed before GPU video decode."""

45 +

46 + source: VideoSourceMetadata

47 + width: int

48 + height: int

49 +

50 +

51 +# Per-decoder upper bound reserved for persistent PyNvVideoCodec surfaces.

52 +PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes

53 +PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2

54 +# Per-API-server CUDA context and driver allocation, measured with

55 +# PyNvVideoCodec 2.0.4 on H100.

56 +PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES = int(1.8 * 1024 * MiB_bytes)

57 +

58 +

59 +def validate_pynvvideocodec_hw_decoders(hw_decoders: object) -> int:

60 + if (

61 + isinstance(hw_decoders, bool)

62 + or not isinstance(hw_decoders, int)

63 + or hw_decoders < 1

64 + ):

65 + raise ValueError("hw_decoders must be a positive integer")

66 + return hw_decoders

67 +

68 +

69 +def _pynvvideocodec_exception_types(nvc) -> tuple[type[Exception], ...]:

70 + return tuple(

71 + exception_type

72 + for name in dir(nvc)

73 + if name.startswith("PyNvVCException")

74 + and isinstance((exception_type := getattr(nvc, name)), type)

75 + and issubclass(exception_type, Exception)

76 + )

77 +

78 +

79 +def _pynvvc_frames_to_nhwc(frames):

80 + """Return a stacked PyNvVideoCodec frame batch as contiguous NHWC."""

81 + if frames.shape[-1] != 3 and frames.shape[-3] == 3:

82 + frames = frames.permute(0, 2, 3, 1)

83 + return frames.contiguous()

84 +

85 +

86 +class PyNvVideoCodecDecoderSlot:

87 + """A retained PyNv decoder slot and its CUDA stream.

88 +

89 + The decoder is reused across requests: ``reconfigure_decoder`` repoints the

90 + existing decoder at each new source instead of paying a fresh

91 + ``SimpleDecoder`` construction per request. Construction (CUVID parser +

92 + decoder + surface-pool allocation) is the dominant per-request cost, so

93 + reconfiguring is far cheaper. A single decoder serves both metadata

94 + (``len``/``get_stream_metadata``) and frame decode -- no separate

95 + metadata decoder.

96 + """

97 +

98 + def __init__(self, stream) -> None:

99 + self.stream = stream

100 + self.decoder = None

101 + self.source_path: str | None = None

102 +

103 + def invalidate(self) -> None:

104 + self.decoder = None

105 + self.source_path = None

106 +

107 + def _construct(self, file_path: str, nvc, device_index: int) -> None:

108 + self.invalidate()

109 + decoder = nvc.SimpleDecoder(

110 + file_path,

111 + output_color_type=nvc.OutputColorType.RGB,

112 + use_device_memory=True,

113 + need_scanned_stream_metadata=True,

114 + gpu_id=device_index,

115 + cuda_stream=self.stream.cuda_stream,

116 + decoder_cache_size=PYNVVIDEOCODEC_DECODER_CACHE_SIZE,

117 + )

118 + self.decoder = decoder

119 + self.source_path = file_path

120 +

121 + def get_decoder(self, file_path: str, nvc, device_index: int):

122 + if self.decoder is None:

123 + self._construct(file_path, nvc, device_index)

124 + elif self.source_path != file_path:

125 + try:

126 + self.decoder.reconfigure_decoder(file_path)

127 + self.source_path = file_path

128 + except Exception:

129 + # reconfigure unsupported/unsafe for this source -> rebuild.

130 + self._construct(file_path, nvc, device_index)

131 + return self.decoder

132 +

133 +

134 +class _PyNvDecoderPool:

135 + """Process-wide singleton managing PyNvVideoCodec decoder slot state.

136 +

137 + Prevents subclass counter shadowing (GHSA-j682-9xp5-rrf3) by storing

138 + all mutable pool state in a single module-level instance rather than

139 + in ClassVar attributes that get shadowed by Python's augmented

140 + assignment semantics on subclasses.

141 + """

142 +

143 + def __init__(self) -> None:

144 + self.slots: list[PyNvVideoCodecDecoderSlot] = []

145 + self.active: int = 0

146 + self.cond: threading.Condition = threading.Condition()

147 + self.max_slots: int | None = None

148 +

149 + def configure(self, hw_decoders: int) -> None:

150 + with self.cond:

151 + if self.max_slots is None:

152 + self.max_slots = hw_decoders

153 + elif self.max_slots != hw_decoders:

154 + raise RuntimeError(

155 + "PyNvVideoCodec decoder count is already configured as "

156 + f"{self.max_slots}, got {hw_decoders}"

157 + )

158 +

159 +

160 +_pynv_decoder_pool = _PyNvDecoderPool()

161 +

162 +

163 +class PyNvVideoCodecVideoBackendMixin:

164 + """PyNvVideoCodec utilities for GPU-backed frame decode."""

165 +

166 + _DEVICE_INDEX: ClassVar[int] = 0

167 +

168 + @classmethod

169 + def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot:

170 + import torch

171 +

172 + return PyNvVideoCodecDecoderSlot(torch.cuda.Stream(device=cls._DEVICE_INDEX))

173 +

174 + @classmethod

175 + def _configure_decoder_slots(cls, hw_decoders: object) -> None:

176 + hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders)

177 + _pynv_decoder_pool.configure(hw_decoders)

178 +

179 + @staticmethod

180 + @contextmanager

181 + def _torch_stream_context(stream):

182 + import torch

183 +

184 + torch.accelerator.set_device_index(stream.device.index)

185 + previous_stream = torch.accelerator.current_stream()

186 + torch.accelerator.set_stream(stream)

187 + try:

188 + yield

189 + finally:

190 + torch.accelerator.set_stream(previous_stream)

191 +

192 + @classmethod

193 + @contextmanager

194 + def _borrow_decoder_slot(cls):

195 + pool = _pynv_decoder_pool

196 + create_slot = False

197 + with pool.cond:

198 + if pool.max_slots is None:

199 + raise RuntimeError("PyNvVideoCodec decoder slots are not configured")

200 + while True:

201 + if pool.slots:

202 + slot = pool.slots.pop()

203 + break

204 + if pool.active < pool.max_slots:

205 + pool.active += 1

206 + create_slot = True

207 + break

208 + pool.cond.wait()

209 +

210 + if create_slot:

211 + try:

212 + slot = cls._create_decoder_slot()

213 + except Exception:

214 + with pool.cond:

215 + pool.active -= 1

216 + pool.cond.notify()

217 + raise

218 +

219 + borrow_succeeded = False

220 + try:

221 + yield slot

222 + borrow_succeeded = True

223 + finally:

224 + if not borrow_succeeded:

225 + slot.invalidate()

226 + with pool.cond:

227 + pool.slots.append(slot)

228 + pool.cond.notify()

229 +

230 + @staticmethod

231 + def _metadata_value(metadata, *names: str, default=None):

232 + for name in names:

233 + value = getattr(metadata, name, None)

234 + if value is not None:

235 + return value

236 + return default

237 +

238 + @classmethod

239 + def _read_source_metadata(

240 + cls,

241 + file_path: str,

242 + nvc,

243 + ) -> PyNvVideoCodecSourceMetadata:

244 + with cls._borrow_decoder_slot() as decoder_slot:

245 + with cls._torch_stream_context(decoder_slot.stream):

246 + decoder = decoder_slot.get_decoder(

247 + file_path, nvc, device_index=cls._DEVICE_INDEX

248 + )

249 + metadata = decoder.get_stream_metadata()

250 + total_frames_num = len(decoder)

251 + width = int(cls._metadata_value(metadata, "width", default=0))

252 + height = int(cls._metadata_value(metadata, "height", default=0))

253 + original_fps = float(

254 + cls._metadata_value(

255 + metadata,

256 + "average_fps",

257 + "avg_frame_rate",

258 + "frame_rate",

259 + "frameRate",

260 + default=0.0,

261 + )

262 + )

263 + duration = float(

264 + cls._metadata_value(metadata, "duration", default=0.0)

265 + or (total_frames_num / original_fps if original_fps > 0 else 0.0)

266 + )

267 + if total_frames_num <= 0:

268 + raise ValueError("Could not determine video frame count")

269 + if width <= 0 or height <= 0:

270 + raise ValueError("Could not determine video dimensions")

271 + return PyNvVideoCodecSourceMetadata(

272 + source=VideoSourceMetadata(total_frames_num, original_fps, duration),

273 + width=width,

274 + height=height,

275 + )

276 +

277 + @classmethod

278 + def _decode_to_pinned_host(

279 + cls,

280 + file_path: str,

281 + frame_idx: list[int],

282 + nvc,

283 + ) -> npt.NDArray:

284 + import torch

285 +

286 + if not frame_idx:

287 + return np.empty((0,), dtype=np.uint8)

288 +

289 + with cls._borrow_decoder_slot() as decoder_slot:

290 + stream = decoder_slot.stream

291 + with cls._torch_stream_context(stream):

292 + try:

293 + decoder = decoder_slot.get_decoder(

294 + file_path, nvc, device_index=cls._DEVICE_INDEX

295 + )

296 + decoded_frames = decoder.get_batch_frames_by_index(frame_idx)

297 + except Exception as exc:

298 + if not isinstance(

299 + exc,

300 + _pynvvideocodec_exception_types(nvc) + (IndexError,),

301 + ):

302 + raise

303 + raise ValueError("Invalid or unsupported video file.") from exc

304 + if len(decoded_frames) < len(frame_idx):

305 + logger.warning(

306 + "pynvvideocodec video loading: expected %d frames but got %d.",

307 + len(frame_idx),

308 + len(decoded_frames),

309 + )

310 + torch_frames = [torch.from_dlpack(frame) for frame in decoded_frames]

311 + if not torch_frames:

312 + return np.empty((0,), dtype=np.uint8)

313 + device_frames = torch.stack(torch_frames)

314 + if device_frames.ndim != 4:

315 + raise ValueError(

316 + "PyNvVideoCodec returned frames with unexpected shape "

317 + f"{tuple(device_frames.shape)}"

318 + )

319 + device_frames = _pynvvc_frames_to_nhwc(device_frames)

320 + host_frames = torch.empty(

321 + device_frames.shape,

322 + dtype=device_frames.dtype,

323 + device="cpu",

324 + pin_memory=True,

325 + )

326 + host_frames.copy_(device_frames, non_blocking=True)

327 + stream.synchronize()

328 + host_array = host_frames.numpy()

329 + del decoded_frames, torch_frames, device_frames

330 + return host_array

331 +

332 + @classmethod

333 + def decode_frames_pynvvideocodec(

334 + cls,

335 + loader_cls,

336 + data: bytes,

337 + target: VideoTargetMetadata,

338 + **kwargs,

339 + ) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:

340 + import PyNvVideoCodec as nvc

341 +

342 + from vllm.multimodal.gpu_ipc_memory import get_mm_gpu_ipc_pool

343 +

344 + temp_fd, temp_path = tempfile.mkstemp(suffix=".mp4")

345 + try:

346 + with os.fdopen(temp_fd, "wb") as temp_file:

347 + temp_file.write(data)

348 +

349 + try:

350 + gpu_source = cls._read_source_metadata(temp_path, nvc)

351 + except Exception as exc:

352 + if not isinstance(exc, _pynvvideocodec_exception_types(nvc)):

353 + raise

354 + raise ValueError("Invalid or unsupported video file.") from exc

355 + check_frame_pixel_limit(gpu_source.width, gpu_source.height)

356 + source = loader_cls._prepare_source(gpu_source.source)

357 + frame_idx = loader_cls.compute_frames_index_to_sample(

358 + source=source, target=target, **kwargs

359 + )

360 + raw_frame_bytes = len(frame_idx) * gpu_source.height * gpu_source.width * 3

361 + pool = get_mm_gpu_ipc_pool()

362 + if pool is None or raw_frame_bytes == 0:

363 + frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc)

364 + else:

365 + with pool.acquire(raw_frame_bytes):

366 + frames = cls._decode_to_pinned_host(temp_path, frame_idx, nvc)

367 + finally:

368 + with suppress(FileNotFoundError):

369 + os.unlink(temp_path)

370 +

371 + valid_frame_indices = frame_idx[: int(frames.shape[0])]

372 + return frames, source, frame_idx, valid_frame_indices

vllm/multimodal/video_decoders/torchcodec.py new

+100

@@ -0,0 +1,100 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from typing import Literal

5 +

6 +import numpy as np

7 +import numpy.typing as npt

8 +

9 +from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available

10 +

11 +from .base import (

12 + VideoSourceMetadata,

13 + VideoTargetMetadata,

14 + check_frame_pixel_limit,

15 +)

16 +

17 +try:

18 + from torchcodec.decoders import VideoDecoder

19 +except (ImportError, RuntimeError):

20 + VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment]

21 + "decoders.VideoDecoder"

22 + )

23 +

24 +

25 +def decode_torchcodec(

26 + loader_cls,

27 + data: bytes,

28 + target: VideoTargetMetadata,

29 + sampling_kwargs: dict,

30 + *,

31 + num_ffmpeg_threads: int = 0,

32 + seek_mode: Literal["exact", "approximate"] = "exact",

33 +) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:

34 + check_torchcodec_available()

35 + decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder(

36 + data,

37 + num_ffmpeg_threads=num_ffmpeg_threads,

38 + seek_mode=seek_mode,

39 + )

40 + check_frame_pixel_limit(

41 + decoder.metadata.width or 0,

42 + decoder.metadata.height or 0,

43 + )

44 + source = loader_cls._prepare_source(

45 + TorchCodecVideoBackendMixin.get_torchcodec_metadata(decoder)

46 + )

47 + frame_idx = loader_cls.compute_frames_index_to_sample(

48 + source=source, target=target, **sampling_kwargs

49 + )

50 + frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames(

51 + decoder, frame_idx

52 + )

53 + return frames, source, frame_idx, valid

54 +

55 +

56 +class TorchCodecVideoBackendMixin:

57 + """TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities.

58 +

59 + Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory

60 + bytes and extracts the sampled indices with a single batched

61 + ``get_frames_at`` call, while releasing the GIL during decode.

62 + """

63 +

64 + @staticmethod

65 + def make_torchcodec_decoder(

66 + data: bytes,

67 + *,

68 + num_ffmpeg_threads: int = 0,

69 + seek_mode: Literal["exact", "approximate"] = "exact",

70 + ) -> "VideoDecoder":

71 + # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest

72 + # of the pipeline expects, avoiding a transpose.

73 + return VideoDecoder(

74 + data,

75 + dimension_order="NHWC",

76 + num_ffmpeg_threads=num_ffmpeg_threads,

77 + seek_mode=seek_mode,

78 + )

79 +

80 + @staticmethod

81 + def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata:

82 + md = decoder.metadata

83 + total_frames = md.num_frames or 0

84 + fps = float(md.average_fps) if md.average_fps else 0.0

85 + duration = float(md.duration_seconds) if md.duration_seconds else 0.0

86 + if total_frames == 0 and duration > 0 and fps > 0:

87 + total_frames = int(duration * fps)

88 + return VideoSourceMetadata(total_frames, fps, duration)

89 +

90 + @staticmethod

91 + def decode_torchcodec_frames(

92 + decoder: "VideoDecoder",

93 + frame_indices: list[int],

94 + ) -> tuple[npt.NDArray, list[int]]:

95 + """Decode the requested indices in one batched, index-exact call."""

96 + if not frame_indices:

97 + return np.empty((0,), dtype=np.uint8), []

98 + # Note: torchcodec releases the GIL for the entire call

99 + batch = decoder.get_frames_at(frame_indices)

100 + return batch.data.numpy(), list(frame_indices)