| 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 | import torch |
| 9 | |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available |
| 13 | |
| 14 | from .base import ( |
| 15 | VideoSourceMetadata, |
| 16 | VideoTargetMetadata, |
| 17 | check_frame_pixel_limit, |
| 18 | ) |
| 19 | |
| 20 | try: |
| 21 | from torchcodec.decoders import VideoDecoder |
| 22 | except (ImportError, RuntimeError): |
| 23 | VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment] |
| 24 | "decoders.VideoDecoder" |
| 25 | ) |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | |
| 30 | def decode_torchcodec( |
| 31 | loader_cls, |
| 32 | data: bytes, |
| 33 | target: VideoTargetMetadata, |
| 34 | sampling_kwargs: dict, |
| 35 | *, |
| 36 | num_ffmpeg_threads: int = 0, |
| 37 | seek_mode: Literal["exact", "approximate"] = "exact", |
| 38 | device: str = "cpu", |
| 39 | ) -> tuple[npt.NDArray | torch.Tensor, VideoSourceMetadata, list[int], list[int]]: |
| 40 | check_torchcodec_available() |
| 41 | decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder( |
| 42 | data, |
| 43 | num_ffmpeg_threads=num_ffmpeg_threads, |
| 44 | seek_mode=seek_mode, |
| 45 | device=device, |
| 46 | ) |
| 47 | check_frame_pixel_limit( |
| 48 | decoder.metadata.width or 0, |
| 49 | decoder.metadata.height or 0, |
| 50 | ) |
| 51 | source = loader_cls._prepare_source( |
| 52 | TorchCodecVideoBackendMixin.get_torchcodec_metadata(decoder) |
| 53 | ) |
| 54 | frame_idx = loader_cls.compute_frames_index_to_sample( |
| 55 | source=source, target=target, **sampling_kwargs |
| 56 | ) |
| 57 | frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames( |
| 58 | decoder, frame_idx, device=device |
| 59 | ) |
| 60 | return frames, source, frame_idx, valid |
| 61 | |
| 62 | |
| 63 | class TorchCodecVideoBackendMixin: |
| 64 | """TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities. |
| 65 | |
| 66 | Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory |
| 67 | bytes and extracts the sampled indices with a single batched |
| 68 | ``get_frames_at`` call, while releasing the GIL during decode. |
| 69 | """ |
| 70 | |
| 71 | @staticmethod |
| 72 | def make_torchcodec_decoder( |
| 73 | data: bytes, |
| 74 | *, |
| 75 | num_ffmpeg_threads: int = 0, |
| 76 | seek_mode: Literal["exact", "approximate"] = "exact", |
| 77 | device: str = "cpu", |
| 78 | ) -> "VideoDecoder": |
| 79 | torch_device = torch.device(device) |
| 80 | if torch_device.type == "cuda" and not current_platform.is_cuda(): |
| 81 | raise ValueError( |
| 82 | f"torchcodec video decoding on device {device!r} requires " |
| 83 | "a CUDA-capable platform." |
| 84 | ) |
| 85 | elif torch_device.type not in ("cpu", "cuda"): |
| 86 | raise ValueError( |
| 87 | f"torchcodec video decoding only supports 'cpu' and 'cuda' " |
| 88 | f"devices, got {device!r}." |
| 89 | ) |
| 90 | # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest |
| 91 | # of the pipeline expects, avoiding a transpose. |
| 92 | return VideoDecoder( |
| 93 | data, |
| 94 | dimension_order="NHWC", |
| 95 | num_ffmpeg_threads=num_ffmpeg_threads, |
| 96 | seek_mode=seek_mode, |
| 97 | device=device, |
| 98 | ) |
| 99 | |
| 100 | @staticmethod |
| 101 | def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata: |
| 102 | md = decoder.metadata |
| 103 | total_frames = md.num_frames or 0 |
| 104 | fps = float(md.average_fps) if md.average_fps else 0.0 |
| 105 | duration = float(md.duration_seconds) if md.duration_seconds else 0.0 |
| 106 | if total_frames == 0 and duration > 0 and fps > 0: |
| 107 | total_frames = int(duration * fps) |
| 108 | return VideoSourceMetadata(total_frames, fps, duration) |
| 109 | |
| 110 | @staticmethod |
| 111 | def decode_torchcodec_frames( |
| 112 | decoder: "VideoDecoder", |
| 113 | frame_indices: list[int], |
| 114 | *, |
| 115 | device: str = "cpu", |
| 116 | ) -> tuple[npt.NDArray | torch.Tensor, list[int]]: |
| 117 | """Decode the requested indices in one batched, index-exact call.""" |
| 118 | if not frame_indices: |
| 119 | return np.empty((0,), dtype=np.uint8), [] |
| 120 | # Note: torchcodec releases the GIL for the entire call |
| 121 | batch = decoder.get_frames_at(frame_indices) |
| 122 | frames = batch.data |
| 123 | if ( |
| 124 | torch.device(device).type != "cpu" |
| 125 | and (status := getattr(decoder, "cpu_fallback", None)) is not None |
| 126 | and bool(status) |
| 127 | ): |
| 128 | logger.warning_once( |
| 129 | "torchcodec could not use NVDEC for this video and " |
| 130 | "decoded on CPU instead; check codec support and " |
| 131 | "libnvcuvid." |
| 132 | ) |
| 133 | return frames, list(frame_indices) |