diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index 589eabf91d..84cf092c21 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -375,3 +375,85 @@ def test_vllm_config_runs_the_mm_processor_device_check(): pytest.raises(ValueError, match="also runs the language model"), ): VllmConfig._validate_mm_processor_device(vllm_config) + + +def _resolve_mm_video_decode_device( + *, + ec_role: ECRole | None, + mm_tensor_ipc: str = "torch_shm", + device: str | None = None, + video_kwargs: dict | None = None, + torchcodec_available: bool = True, +) -> dict: + """Run processor-device then video-decode-device resolution and report + the resulting video media IO kwargs.""" + mm_config = MultiModalConfig( + mm_processor_kwargs={} if device is None else {"device": device}, + mm_tensor_ipc=mm_tensor_ipc, # type: ignore[arg-type] + media_io_kwargs={} if video_kwargs is None else {"video": dict(video_kwargs)}, + ) + model_config = MagicMock(spec=ModelConfig) + model_config.multimodal_config = mm_config + vllm_config = MagicMock(spec=VllmConfig) + vllm_config.model_config = model_config + vllm_config.ec_transfer_config = ( + None + if ec_role is None + else ECTransferConfig(ec_connector="ECExampleConnector", ec_role=ec_role) + ) + + with ( + patch("vllm.platforms.current_platform.device_type", "cuda"), + patch( + "vllm.utils.import_utils.check_torchcodec_available", + side_effect=None if torchcodec_available else ImportError("torchcodec"), + ), + ): + VllmConfig._resolve_mm_processor_device(vllm_config) + VllmConfig._resolve_mm_video_decode_device(vllm_config) + return mm_config.media_io_kwargs.get("video", {}) + + +def test_auto_video_decode_uses_nvdec_on_encoder_instance(): + """The processor runs on the accelerator there, so decoded frames should + stay on-device too.""" + assert _resolve_mm_video_decode_device(ec_role="ec_producer") == { + "backend": "torchcodec", + "device": "cuda", + } + + +@pytest.mark.parametrize("ec_role", [None, "ec_consumer", "ec_both"]) +def test_auto_video_decode_stays_on_cpu_off_encoder_instance( + ec_role: ECRole | None, +): + """The processor stays on CPU off encode-only instances, so video + decoding should too.""" + assert _resolve_mm_video_decode_device(ec_role=ec_role) == {} + + +def test_auto_video_decode_follows_explicit_processor_device(): + assert _resolve_mm_video_decode_device(ec_role="ec_producer", device="cpu") == {} + assert _resolve_mm_video_decode_device(ec_role="ec_producer", device="cuda") == { + "backend": "torchcodec", + "device": "cuda", + } + + +def test_auto_video_decode_respects_explicit_media_io_kwargs(): + assert _resolve_mm_video_decode_device( + ec_role="ec_producer", video_kwargs={"backend": "opencv"} + ) == {"backend": "opencv"} + assert _resolve_mm_video_decode_device( + ec_role="ec_producer", video_kwargs={"device": "cpu"} + ) == {"device": "cpu"} + + +def test_auto_video_decode_skipped_without_torchcodec(): + assert ( + _resolve_mm_video_decode_device( + ec_role="ec_producer", + torchcodec_available=False, + ) + == {} + ) diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index a3ea0890a2..79070ac50d 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -532,6 +532,50 @@ class TestMergeKwargsGpuBackendPolicy: def test_unknown_backend_not_treated_as_gpu(self): assert not VIDEO_LOADER_REGISTRY.backend_requires_gpu("totally_unknown") + def test_strips_request_level_device(self): + """The decode device is a startup-only knob: a request must not move + decoding onto the GPU when the startup config did not opt in, nor off + it when it did.""" + result = VideoMediaIO.merge_kwargs( + default_kwargs={"backend": "torchcodec"}, + runtime_kwargs={"device": "cuda"}, + ) + assert "device" not in result + + result = VideoMediaIO.merge_kwargs( + default_kwargs={"backend": "torchcodec", "device": "cuda"}, + runtime_kwargs={"device": "cpu", "num_frames": 8}, + ) + assert result["device"] == "cuda" + assert result["num_frames"] == 8 + + +@pytest.mark.parametrize( + "default_kwargs", + [ + {"backend": "torchcodec", "device": "cuda", "seek_mode": "approximate"}, + {"backend": "pynvvideocodec", "hw_decoders": 2}, + ], +) +def test_switching_backend_drops_stale_codec_options(default_kwargs): + """Codec-specific options from the static config must not leak into a + different codec backend selected per-request, where they would fail the + new backend's option validation.""" + result = VideoMediaIO.merge_kwargs( + default_kwargs={**default_kwargs, "num_frames": 8}, + runtime_kwargs={"backend": "opencv"}, + ) + assert result == {"backend": "opencv", "num_frames": 8} + + +def test_same_backend_keeps_codec_options(): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"backend": "torchcodec", "device": "cuda"}, + runtime_kwargs={"backend": "torchcodec", "num_frames": 8}, + ) + assert result["device"] == "cuda" + assert result["num_frames"] == 8 + @pytest.mark.parametrize("layout", ["nhwc", "nchw"]) def test_pynvvc_frames_normalized_to_nhwc(layout: str): diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index 4810081ad4..d0be316790 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -137,6 +137,33 @@ def test_hash_collision_video_num_frames(): ) +def test_hash_video_tensor_frames(): + """Videos holding tensor frames (e.g. NVDEC-decoded) hash like + array-framed ones, from the original bytes without a D2H copy.""" + source = b"x" * 100 + + def item_for_hash(frames): + metadata = { + "total_num_frames": 2, + "fps": 2.0, + "duration": 1.0, + "video_backend": "torchcodec", + "frames_indices": [0, 1], + "do_sample_frames": False, + } + video = MediaWithBytes((frames, metadata), source) + items = MultiModalDataParser()._parse_video_data([video]) + return items.get_all_items_for_hash()[0] + + np_frames = np.zeros((2, 8, 8, 3), dtype=np.uint8) + torch_frames = torch.zeros((2, 8, 8, 3), dtype=torch.uint8) + + hasher = MultiModalHasher + assert hasher.hash_kwargs("blake3", video=item_for_hash(np_frames)) == ( + hasher.hash_kwargs("blake3", video=item_for_hash(torch_frames)) + ) + + def test_hash_non_contiguous_array(): arr = np.arange(24).reshape(4, 6).T assert not arr.flags.c_contiguous diff --git a/tests/multimodal/test_parse.py b/tests/multimodal/test_parse.py index b4d37b0fe8..2649f80777 100644 --- a/tests/multimodal/test_parse.py +++ b/tests/multimodal/test_parse.py @@ -56,6 +56,28 @@ def test_frame_size_hwc_chw(frame): assert items.get_frame_size(0) == (W, H) +def test_video_with_metadata_tensor_passthrough(): + """Tensor frames pass through unchanged regardless of device: HF video + processors accept tensors, and device-resident frames (e.g. NVDEC-decoded) + must not be copied back to host.""" + frames = torch.zeros((4, H, W, 3), dtype=torch.uint8) + video, metadata = MultiModalDataParser()._get_video_with_metadata(frames) + + assert video is frames + assert metadata is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_video_with_metadata_keeps_device_tensor(): + """Device-resident frames (e.g. NVDEC-decoded) pass through as tensors, + so a device-side HF processor can consume them without a D2H copy.""" + frames = torch.zeros((4, H, W, 3), dtype=torch.uint8, device="cuda") + video, metadata = MultiModalDataParser()._get_video_with_metadata(frames) + + assert video is frames + assert metadata is None + + @pytest.mark.parametrize( "frames", [ diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 9c24b5f969..0bc3b21b42 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -11,6 +11,7 @@ from pathlib import Path import numpy as np import numpy.typing as npt import pytest +import torch from transformers import AutoVideoProcessor from transformers.video_utils import VideoMetadata @@ -173,9 +174,14 @@ def test_decode_video_imports_only_selected_backend( [ ( "torchcodec", - {"min_frames": 4, "num_ffmpeg_threads": 2, "seek_mode": "approximate"}, + { + "min_frames": 4, + "num_ffmpeg_threads": 2, + "seek_mode": "approximate", + "device": "cuda", + }, {"min_frames": 4}, - {"num_ffmpeg_threads": 2, "seek_mode": "approximate"}, + {"num_ffmpeg_threads": 2, "seek_mode": "approximate", "device": "cuda"}, ), ( "deepstream", @@ -205,6 +211,11 @@ def test_video_backend_rejects_options_for_another_decoder(): ): resolve_video_backend_kwargs("opencv", {"num_ffmpeg_threads": 2}) + with pytest.raises( + ValueError, match="device is not supported by the 'opencv' backend" + ): + resolve_video_backend_kwargs("opencv", {"device": "cuda"}) + @pytest.mark.parametrize( ("backend", "error"), @@ -1205,6 +1216,30 @@ def test_torchcodec_backend_rejects_frame_recovery(dummy_video_path): ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_torchcodec_backend_cuda_decodes_on_gpu(dummy_video_path): + """With device="cuda", torchcodec decodes via NVDEC and keeps the frames + on the GPU instead of returning a host-side numpy array.""" + pytest.importorskip("torchcodec") + + with open(dummy_video_path, "rb") as f: + video_data = f.read() + + loader = VIDEO_LOADER_REGISTRY.load("opencv") + frames, metadata = loader.load_bytes( + video_data, num_frames=8, backend="torchcodec", device="cuda" + ) + + assert isinstance(frames, torch.Tensor) + assert frames.device.type == "cuda" + assert frames.dtype == torch.uint8 + assert frames.ndim == 4 + assert frames.shape[3] == 3 # RGB + assert frames.shape[0] == 8 + assert frames.shape[0] == len(metadata["frames_indices"]) + assert metadata["video_backend"] == "torchcodec" + + def test_torchcodec_backend_returns_target_frames_not_keyframes(): """Regression test: torchcodec must return the requested frames, not the GOP keyframe they seek back to. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 39c2cb4d52..05377a2735 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1736,6 +1736,7 @@ class VllmConfig: self._resolve_mm_embedding_inputs() self._resolve_mm_processor_device() + self._resolve_mm_video_decode_device() self._validate_mm_processor_device() if self.use_v2_model_runner: @@ -2623,6 +2624,61 @@ class VllmConfig: device_type, ) + def _resolve_mm_video_decode_device(self) -> None: + """Default video decoding to torchcodec GPU backend for EPD encoder-only + instance if the mm processor runs on CUDA. + + The processor consumes the decoded frames on-device in that case, so + keeping the frames on the GPU skips the host round-trip through the + CPU media path. An explicit codec/backend choice in + `--media-io-kwargs` is left alone, and the default is skipped where + torchcodec (or its FFmpeg runtime) is unavailable. + """ + if self.model_config is None or self.model_config.multimodal_config is None: + return + mm_config = self.model_config.multimodal_config + + ec_config = self.ec_transfer_config + # An EC producer that is not also a consumer runs no forward pass and + # allocates no KV cache, so frontend accelerator work has the device to + # itself. + if ec_config is None or not ec_config.is_encode_only: + return + + from vllm.platforms import current_platform + + device_type = current_platform.device_type + if ( + device_type != "cuda" + or mm_config.get_mm_processor_device_type() != device_type + ): + return + + # User set video backend or device explicitly + video_kwargs = mm_config.media_io_kwargs.setdefault("video", {}) + if "backend" in video_kwargs or "device" in video_kwargs: + return + + from vllm.utils.import_utils import check_torchcodec_available + + try: + check_torchcodec_available() + except (ImportError, RuntimeError): + # torchcodec is not installed, or is installed without a usable + # FFmpeg runtime (it raises rather than returning False). + logger.info_once( + "EPD encoder instance: keeping CPU video decoding because " + "torchcodec is not available (needs a CUDA build with FFmpeg)." + ) + return + + video_kwargs["backend"] = "torchcodec" + video_kwargs["device"] = device_type + logger.info_once( + "EPD encoder instance: decoding video with NVDEC (torchcodec device=%s).", + device_type, + ) + def _validate_mm_processor_device(self) -> None: """Hand the EC config to `MultiModalConfig`, which owns the rule.""" model_config = self.model_config diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py index 3d86b33d51..a432afdc9b 100644 --- a/vllm/multimodal/hasher.py +++ b/vllm/multimodal/hasher.py @@ -89,8 +89,11 @@ class MultiModalHasher: ) return cls.iter_item_to_bytes("image", obj.original_bytes) - if isinstance(obj, MediaWithBytes) and isinstance(obj.media, np.ndarray): + if isinstance(obj, MediaWithBytes) and isinstance( + obj.media, (np.ndarray, torch.Tensor) + ): frames = obj.media + # Both np.ndarray and torch.Tensor expose .nbytes. if frames.nbytes < len(obj.original_bytes): return cls.iter_item_to_bytes("video", frames) return cls.iter_item_to_bytes("video", obj.original_bytes) diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index e5c073f5f3..f1f5783b0e 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -19,17 +19,23 @@ from vllm.utils.sparse_utils import ( safe_to_dense, ) -from ..video import VIDEO_LOADER_REGISTRY +from ..video import VIDEO_LOADER_REGISTRY, DecodedFrames +from ..video_decoders import incompatible_backend_options from .base import MediaIO, MediaWithBytes from .image import MAGIC_NUMPY_PREFIX, ImageMediaIO logger = init_logger(__name__) -class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]): +class VideoMediaIO(MediaIO[MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]]): """Configuration values can be user-provided either by --media-io-kwargs or by the runtime API field "media_io_kwargs". Ensure proper validation and error handling. + + Use --media-io-kwargs '{"video": {"backend": "torchcodec", + "device": "cuda"}}' to decode with NVDEC and keep frames on the GPU, + which pairs with a device-side multi-modal processor + (--mm-processor-device, EPD encode-only instances). """ @classmethod @@ -43,6 +49,9 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]): runtime_kwargs = dict(runtime_kwargs) runtime_kwargs.pop("hw_decoders", None) runtime_kwargs.pop("pool_size", None) + # The decode device determines which device the decoded frames + # (and thus the processor pipeline) live on; startup-only. + runtime_kwargs.pop("device", None) # Block request-level selection of GPU video backends that # were not configured (and VRAM-reserved) at startup. @@ -70,6 +79,14 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]): merged.pop("fps", None) elif "fps" in runtime_kwargs and "num_frames" not in runtime_kwargs: merged.pop("num_frames", None) + + requested_backend = runtime_kwargs.get("backend") + static_backend = (default_kwargs or {}).get("backend") + if requested_backend and requested_backend != static_backend: + for name in incompatible_backend_options( + static_backend, requested_backend + ): + merged.pop(name, None) return merged def __init__( @@ -101,7 +118,7 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]): def load_bytes( self, data: bytes - ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: + ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]: video = self.video_loader.load_bytes( data, num_frames=self.num_frames, **self.kwargs ) @@ -109,7 +126,7 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]): def load_base64( self, media_type: str, data: str - ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: + ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]: if media_type.lower() == "video/jpeg": load_frame = partial( self.image_io.load_base64, @@ -177,7 +194,7 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]): def load_file( self, filepath: Path - ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: + ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]: with filepath.open("rb") as f: data = f.read() diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index 124a8845cc..b2f8664246 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -36,6 +36,7 @@ from .inputs import ( VideoItem, ) from .media import MediaWithBytes +from .video import DecodedFrames _T = TypeVar("_T") _I = TypeVar("_I") @@ -653,7 +654,7 @@ class MultiModalDataParser: def _get_video_with_metadata( self, video: VideoItem, - ) -> tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any] | None]: + ) -> tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any] | None]: if isinstance(video, MediaWithBytes): new_video, metadata = self._get_video_with_metadata(video.media) return MediaWithBytes(new_video, video.original_bytes), metadata @@ -661,10 +662,11 @@ class MultiModalDataParser: return video if isinstance(video, list): return np.array(video), None - if isinstance(video, np.ndarray): + if isinstance(video, (np.ndarray, torch.Tensor)): + # Tensors pass through untouched: HF video processors accept them + # directly, and device tensors (e.g. NVDEC-decoded frames) must + # stay on-device to avoid a D2H round-trip. return video, None - if isinstance(video, torch.Tensor): - return video.numpy(), None assert_never(video) @@ -769,9 +771,9 @@ class MultiModalDataParser: data_items = data # type: ignore[assignment] new_videos = list[ - np.ndarray - | MediaWithBytes[np.ndarray] - | tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any]] + DecodedFrames + | MediaWithBytes[DecodedFrames] + | tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any]] | None ]() metadata_lst: list[dict[str, Any] | None] = [] diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 854b5572b3..4142febc3f 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from abc import abstractmethod -from typing import Any, ClassVar, cast +from typing import Any, ClassVar, TypeAlias, cast import numpy as np import numpy.typing as npt @@ -28,6 +28,10 @@ except ImportError: logger = init_logger(__name__) +DecodedFrames: TypeAlias = npt.NDArray | torch.Tensor +"""Decoded video frames: a host ``np.ndarray``, or a device ``torch.Tensor`` +when a GPU decoding codec is used (e.g. torchcodec with ``device="cuda"``).""" + class VideoLoaderRegistry(ExtensionManager): def __init__(self) -> None: @@ -164,8 +168,13 @@ class VideoLoader: cls, data: bytes, **kwargs, - ) -> tuple[npt.NDArray, dict[str, Any]]: - """Load video frames from bytes and return (frames_array, metadata_dict).""" + ) -> tuple[DecodedFrames, dict[str, Any]]: + """Load video frames from bytes and return (frames, metadata_dict). + + ``frames`` is a CPU ``np.ndarray`` unless a GPU decoding codec is + used (e.g. torchcodec with ``device="cuda"``), in which case it is + a ``torch.Tensor`` living on that device. + """ raise NotImplementedError @classmethod @@ -239,7 +248,7 @@ class VideoBackend(VideoLoader): *, backend: VideoDecoderBackend = "opencv", **kwargs: Any, - ) -> tuple[npt.NDArray, dict[str, Any]]: + ) -> tuple[DecodedFrames, dict[str, Any]]: """Load sampled frames from raw video bytes. Args: @@ -265,6 +274,11 @@ class VideoBackend(VideoLoader): creation at the cost of relying on the file's metadata. See https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html for details. + - ``device`` (TorchCodec): ``"cpu"`` (default) decodes on the + host and returns a ``np.ndarray``; ``"cuda"`` decodes with + NVDEC and returns the frames as a CUDA ``torch.Tensor``, + which a device-side HF video processor can consume without + a host round-trip. - ``hw_decoders`` (PyNvVideoCodec): maximum number of concurrent decoder slots. Defaults to 2 and must be a positive integer. @@ -272,7 +286,8 @@ class VideoBackend(VideoLoader): size and pool acquisition timeout in seconds. Returns: - Tuple of ``(frames_array, metadata_dict)``. + Tuple of ``(frames, metadata_dict)``, where ``frames`` is a + CPU ``np.ndarray`` unless TorchCodec decodes on ``device="cuda"``. """ target = VideoTargetMetadata( num_frames=num_frames, fps=fps, max_duration=max_duration diff --git a/vllm/multimodal/video_decoders/__init__.py b/vllm/multimodal/video_decoders/__init__.py index 1e6106e199..9931f60f87 100644 --- a/vllm/multimodal/video_decoders/__init__.py +++ b/vllm/multimodal/video_decoders/__init__.py @@ -19,6 +19,7 @@ _BACKEND_OPTION_DEFAULTS: dict[str, dict[str, Any]] = { "torchcodec": { "num_ffmpeg_threads": 0, "seek_mode": "exact", + "device": "cpu", }, PYNVVIDEOCODEC_VIDEO_BACKEND: { "hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, @@ -40,6 +41,17 @@ def _get_backend_option_defaults(backend: str) -> dict[str, Any]: ) from None +# TODO(Isotr0py): Make kwargs structured. +def incompatible_backend_options( + old_backend: str | None, + new_backend: str, +) -> set[str]: + """Option names accepted by ``old_backend`` but not by ``new_backend``.""" + old = _BACKEND_OPTION_DEFAULTS.get(old_backend or "", {}) + new = _BACKEND_OPTION_DEFAULTS.get(new_backend, {}) + return old.keys() - new.keys() + + def resolve_video_backend_kwargs( backend: str, kwargs: dict[str, Any], @@ -103,5 +115,6 @@ __all__ = [ "VideoTargetMetadata", "check_frame_pixel_limit", "decode_video", + "incompatible_backend_options", "resolve_video_backend_kwargs", ] diff --git a/vllm/multimodal/video_decoders/torchcodec.py b/vllm/multimodal/video_decoders/torchcodec.py index f1f3290629..346f7b645b 100644 --- a/vllm/multimodal/video_decoders/torchcodec.py +++ b/vllm/multimodal/video_decoders/torchcodec.py @@ -5,7 +5,10 @@ from typing import Literal import numpy as np import numpy.typing as npt +import torch +from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available from .base import ( @@ -21,6 +24,8 @@ except (ImportError, RuntimeError): "decoders.VideoDecoder" ) +logger = init_logger(__name__) + def decode_torchcodec( loader_cls, @@ -30,12 +35,14 @@ def decode_torchcodec( *, num_ffmpeg_threads: int = 0, seek_mode: Literal["exact", "approximate"] = "exact", -) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]: + device: str = "cpu", +) -> tuple[npt.NDArray | torch.Tensor, VideoSourceMetadata, list[int], list[int]]: check_torchcodec_available() decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder( data, num_ffmpeg_threads=num_ffmpeg_threads, seek_mode=seek_mode, + device=device, ) check_frame_pixel_limit( decoder.metadata.width or 0, @@ -48,7 +55,7 @@ def decode_torchcodec( source=source, target=target, **sampling_kwargs ) frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames( - decoder, frame_idx + decoder, frame_idx, device=device ) return frames, source, frame_idx, valid @@ -67,7 +74,19 @@ class TorchCodecVideoBackendMixin: *, num_ffmpeg_threads: int = 0, seek_mode: Literal["exact", "approximate"] = "exact", + device: str = "cpu", ) -> "VideoDecoder": + torch_device = torch.device(device) + if torch_device.type == "cuda" and not current_platform.is_cuda(): + raise ValueError( + f"torchcodec video decoding on device {device!r} requires " + "a CUDA-capable platform." + ) + elif torch_device.type not in ("cpu", "cuda"): + raise ValueError( + f"torchcodec video decoding only supports 'cpu' and 'cuda' " + f"devices, got {device!r}." + ) # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest # of the pipeline expects, avoiding a transpose. return VideoDecoder( @@ -75,6 +94,7 @@ class TorchCodecVideoBackendMixin: dimension_order="NHWC", num_ffmpeg_threads=num_ffmpeg_threads, seek_mode=seek_mode, + device=device, ) @staticmethod @@ -91,10 +111,23 @@ class TorchCodecVideoBackendMixin: def decode_torchcodec_frames( decoder: "VideoDecoder", frame_indices: list[int], - ) -> tuple[npt.NDArray, list[int]]: + *, + device: str = "cpu", + ) -> tuple[npt.NDArray | torch.Tensor, list[int]]: """Decode the requested indices in one batched, index-exact call.""" if not frame_indices: return np.empty((0,), dtype=np.uint8), [] # Note: torchcodec releases the GIL for the entire call batch = decoder.get_frames_at(frame_indices) - return batch.data.numpy(), list(frame_indices) + frames = batch.data + if ( + torch.device(device).type != "cpu" + and (status := getattr(decoder, "cpu_fallback", None)) is not None + and bool(status) + ): + logger.warning_once( + "torchcodec could not use NVDEC for this video and " + "decoded on CPU instead; check codec support and " + "libnvcuvid." + ) + return frames, list(frame_indices)