A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
Support custom implementations of VideoLoader backends. (#18091)
Chenheli Hua committed
May 14, 2025 at 22:26 UTC
4f07a640759283a09862e1ad74e390b469891918
3 files changed
+83
-2
tests/multimodal/test_video.py
new
+41
@@ -0,0 +1,41 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import numpy as np
3
+import numpy.typing as npt
4
+import pytest
5
+
6
+from vllm.multimodal.video import VIDEO_LOADER_REGISTRY, VideoLoader
7
+
8
+NUM_FRAMES = 10
9
+FAKE_OUTPUT_1 = np.random.rand(NUM_FRAMES, 1280, 720, 3)
10
+FAKE_OUTPUT_2 = np.random.rand(NUM_FRAMES, 1280, 720, 3)
11
+
12
+
13
+@VIDEO_LOADER_REGISTRY.register("test_video_loader_1")
14
+class TestVideoLoader1(VideoLoader):
15
+
16
+ @classmethod
17
+ def load_bytes(cls, data: bytes, num_frames: int = -1) -> npt.NDArray:
18
+ return FAKE_OUTPUT_1
19
+
20
+
21
+@VIDEO_LOADER_REGISTRY.register("test_video_loader_2")
22
+class TestVideoLoader2(VideoLoader):
23
+
24
+ @classmethod
25
+ def load_bytes(cls, data: bytes, num_frames: int = -1) -> npt.NDArray:
26
+ return FAKE_OUTPUT_2
27
+
28
+
29
+def test_video_loader_registry():
30
+ custom_loader_1 = VIDEO_LOADER_REGISTRY.load("test_video_loader_1")
31
+ output_1 = custom_loader_1.load_bytes(b"test")
32
+ np.testing.assert_array_equal(output_1, FAKE_OUTPUT_1)
33
+
34
+ custom_loader_2 = VIDEO_LOADER_REGISTRY.load("test_video_loader_2")
35
+ output_2 = custom_loader_2.load_bytes(b"test")
36
+ np.testing.assert_array_equal(output_2, FAKE_OUTPUT_2)
37
+
38
+
39
+def test_video_loader_type_doesnt_exist():
40
+ with pytest.raises(AssertionError):
41
+ VIDEO_LOADER_REGISTRY.load("non_existing_video_loader")
vllm/envs.py
+11
@@ -55,6 +55,7 @@ if TYPE_CHECKING:
55
VLLM_IMAGE_FETCH_TIMEOUT: int = 5
56
VLLM_VIDEO_FETCH_TIMEOUT: int = 30
57
VLLM_AUDIO_FETCH_TIMEOUT: int = 10
58
+ VLLM_VIDEO_LOADER_BACKEND: str = "opencv"
59
VLLM_MM_INPUT_CACHE_GIB: int = 8
60
VLLM_TARGET_DEVICE: str = "cuda"
61
MAX_JOBS: Optional[str] = None
@@ -446,6 +447,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
447
"VLLM_AUDIO_FETCH_TIMEOUT":
448
lambda: int(os.getenv("VLLM_AUDIO_FETCH_TIMEOUT", "10")),
449
450
+ # Backend for Video IO
451
+ # - "opencv": Default backend that uses OpenCV stream buffered backend.
452
+ #
453
+ # Custom backend implementations can be registered
454
+ # via `@VIDEO_LOADER_REGISTRY.register("my_custom_video_loader")` and
455
+ # imported at runtime.
456
+ # If a non-existing backend is used, an AssertionError will be thrown.
457
+ "VLLM_VIDEO_LOADER_BACKEND":
458
+ lambda: os.getenv("VLLM_VIDEO_LOADER_BACKEND", "opencv"),
459
+
460
# Cache size (in GiB) for multimodal input cache
461
# Default is 4 GiB
462
"VLLM_MM_INPUT_CACHE_GIB":
vllm/multimodal/video.py
+31
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import base64
4
+from abc import abstractmethod
5
from functools import partial
6
from io import BytesIO
7
from pathlib import Path
@@ -9,6 +10,8 @@ import numpy as np
10
import numpy.typing as npt
11
from PIL import Image
12
13
+from vllm import envs
14
+
15
from .base import MediaIO
16
from .image import ImageMediaIO
17
@@ -48,10 +51,35 @@ def sample_frames_from_video(frames: npt.NDArray,
51
class VideoLoader:
52
53
@classmethod
51
- def load_bytes(self, data: bytes, num_frames: int = -1) -> npt.NDArray:
54
+ @abstractmethod
55
+ def load_bytes(cls, data: bytes, num_frames: int = -1) -> npt.NDArray:
56
raise NotImplementedError
57
58
59
+class VideoLoaderRegistry:
60
+
61
+ def __init__(self) -> None:
62
+ self.name2class: dict[str, type] = {}
63
+
64
+ def register(self, name: str):
65
+
66
+ def wrap(cls_to_register):
67
+ self.name2class[name] = cls_to_register
68
+ return cls_to_register
69
+
70
+ return wrap
71
+
72
+ @staticmethod
73
+ def load(cls_name: str) -> VideoLoader:
74
+ cls = VIDEO_LOADER_REGISTRY.name2class.get(cls_name)
75
+ assert cls is not None, f"VideoLoader class {cls_name} not found"
76
+ return cls()
77
+
78
+
79
+VIDEO_LOADER_REGISTRY = VideoLoaderRegistry()
80
+
81
+
82
+@VIDEO_LOADER_REGISTRY.register("opencv")
83
class OpenCVVideoBackend(VideoLoader):
84
85
def get_cv2_video_api(self):
@@ -122,7 +150,8 @@ class VideoMediaIO(MediaIO[npt.NDArray]):
150
151
self.image_io = image_io
152
self.num_frames = num_frames
125
- self.video_loader = OpenCVVideoBackend
153
+ video_loader_backend = envs.VLLM_VIDEO_LOADER_BACKEND
154
+ self.video_loader = VIDEO_LOADER_REGISTRY.load(video_loader_backend)
155
156
def load_bytes(self, data: bytes) -> npt.NDArray:
157
return self.video_loader.load_bytes(data, self.num_frames)