| 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["opencv", "torchcodec", "pynvvideocodec", "deepstream"] |
| 16 | |
| 17 | _BACKEND_OPTION_DEFAULTS: dict[str, dict[str, Any]] = { |
| 18 | "opencv": {}, |
| 19 | "torchcodec": { |
| 20 | "num_ffmpeg_threads": 0, |
| 21 | "seek_mode": "exact", |
| 22 | "device": "cpu", |
| 23 | }, |
| 24 | PYNVVIDEOCODEC_VIDEO_BACKEND: { |
| 25 | "hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, |
| 26 | }, |
| 27 | "deepstream": { |
| 28 | "pool_size": None, |
| 29 | "timeout_sec": 120.0, |
| 30 | }, |
| 31 | } |
| 32 | |
| 33 | |
| 34 | def _get_backend_option_defaults(backend: str) -> dict[str, Any]: |
| 35 | try: |
| 36 | return _BACKEND_OPTION_DEFAULTS[backend] |
| 37 | except KeyError: |
| 38 | valid = ", ".join(repr(name) for name in _BACKEND_OPTION_DEFAULTS) |
| 39 | raise ValueError( |
| 40 | f"Unknown video codec backend {backend!r}; valid options: {valid}." |
| 41 | ) from None |
| 42 | |
| 43 | |
| 44 | # TODO(Isotr0py): Make kwargs structured. |
| 45 | def incompatible_backend_options( |
| 46 | old_backend: str | None, |
| 47 | new_backend: str, |
| 48 | ) -> set[str]: |
| 49 | """Option names accepted by ``old_backend`` but not by ``new_backend``.""" |
| 50 | old = _BACKEND_OPTION_DEFAULTS.get(old_backend or "", {}) |
| 51 | new = _BACKEND_OPTION_DEFAULTS.get(new_backend, {}) |
| 52 | return old.keys() - new.keys() |
| 53 | |
| 54 | |
| 55 | def resolve_video_backend_kwargs( |
| 56 | backend: str, |
| 57 | kwargs: dict[str, Any], |
| 58 | ) -> tuple[dict[str, Any], dict[str, Any]]: |
| 59 | """Split frame-sampling kwargs from options owned by a decoder backend.""" |
| 60 | defaults = _get_backend_option_defaults(backend) |
| 61 | sampling_kwargs = dict(kwargs) |
| 62 | backend_kwargs = dict(defaults) |
| 63 | backend_option_names = { |
| 64 | name for options in _BACKEND_OPTION_DEFAULTS.values() for name in options |
| 65 | } |
| 66 | misplaced = (sampling_kwargs.keys() & backend_option_names) - defaults.keys() |
| 67 | if misplaced: |
| 68 | names = ", ".join(sorted(misplaced)) |
| 69 | raise ValueError(f"{names} is not supported by the {backend!r} backend") |
| 70 | |
| 71 | for name in defaults: |
| 72 | if name in sampling_kwargs: |
| 73 | backend_kwargs[name] = sampling_kwargs.pop(name) |
| 74 | |
| 75 | return sampling_kwargs, backend_kwargs |
| 76 | |
| 77 | |
| 78 | def decode_video( |
| 79 | backend: str, |
| 80 | loader_cls, |
| 81 | data: bytes, |
| 82 | target: VideoTargetMetadata, |
| 83 | sampling_kwargs: dict[str, Any], |
| 84 | backend_kwargs: dict[str, Any], |
| 85 | *, |
| 86 | frame_recovery: bool, |
| 87 | ): |
| 88 | """Decode a sampled video, importing only the selected backend.""" |
| 89 | _get_backend_option_defaults(backend) |
| 90 | if frame_recovery and backend != "opencv": |
| 91 | error = ( |
| 92 | ValueError if backend == PYNVVIDEOCODEC_VIDEO_BACKEND else AssertionError |
| 93 | ) |
| 94 | raise error(f"frame_recovery is not supported by the {backend!r} backend") |
| 95 | |
| 96 | decoder_kwargs = dict(backend_kwargs) |
| 97 | if backend == "opencv": |
| 98 | decoder_kwargs["frame_recovery"] = frame_recovery |
| 99 | module = import_module(f".{backend}", __name__) |
| 100 | decoder = getattr(module, f"decode_{backend}") |
| 101 | return decoder( |
| 102 | loader_cls, |
| 103 | data, |
| 104 | target, |
| 105 | sampling_kwargs, |
| 106 | **decoder_kwargs, |
| 107 | ) |
| 108 | |
| 109 | |
| 110 | __all__ = [ |
| 111 | "PYNVVIDEOCODEC_DEFAULT_HW_DECODERS", |
| 112 | "PYNVVIDEOCODEC_VIDEO_BACKEND", |
| 113 | "VideoDecoderBackend", |
| 114 | "VideoSourceMetadata", |
| 115 | "VideoTargetMetadata", |
| 116 | "check_frame_pixel_limit", |
| 117 | "decode_video", |
| 118 | "incompatible_backend_options", |
| 119 | "resolve_video_backend_kwargs", |
| 120 | ] |