utils.py 15.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
import asyncio
import atexit
6
from collections.abc import Iterable
7
from concurrent.futures import ThreadPoolExecutor
8
from itertools import groupby
9
from pathlib import Path
10
from typing import TYPE_CHECKING, Any, TypeVar
11
from urllib.parse import ParseResult, urlparse
12
from urllib.request import url2pathname
13

14
import numpy as np
15
import numpy.typing as npt
16
import torch
17
from PIL import Image, UnidentifiedImageError
18

19
import vllm.envs as envs
20
from vllm.connections import HTTPConnection, global_http_connection
21
from vllm.utils.jsontree import json_map_leaves
22

23
24
from .audio import AudioMediaIO
from .base import MediaIO
25
from .image import ImageEmbeddingMediaIO, ImageMediaIO
26
from .video import VideoMediaIO
27

28
_M = TypeVar("_M")
29

30
if TYPE_CHECKING:
31
32
33
34
35
    from .inputs import (
        BatchedTensorInputs,
        MultiModalKwargsItem,
        MultiModalPlaceholderDict,
    )
36
else:
37
38
    BatchedTensorInputs = Any
    MultiModalKwargsItem = Any
39
    MultiModalPlaceholderDict = Any
40

41
global_thread_pool = ThreadPoolExecutor(
42
43
    max_workers=envs.VLLM_MEDIA_LOADING_THREAD_COUNT
)
44
45
atexit.register(global_thread_pool.shutdown)

46

47
48
49
class MediaConnector:
    def __init__(
        self,
50
        media_io_kwargs: dict[str, dict[str, Any]] | None = None,
51
52
53
        connection: HTTPConnection = global_http_connection,
        *,
        allowed_local_media_path: str = "",
54
        allowed_media_domains: list[str] | None = None,
55
    ) -> None:
56
57
        """
        Args:
58
59
60
            media_io_kwargs: Additional args passed to process media
                             inputs, keyed by modalities. For example,
                             to set num_frames for video, set
61
                             `--media-io-kwargs '{"video":{"num_frames":40}}'`
62
            connection: HTTP connection client to download media contents.
63
64
            allowed_local_media_path: A local directory to load media files
                                      from.
65
        """
66
67
        super().__init__()

68
69
70
        self.media_io_kwargs: dict[str, dict[str, Any]] = (
            media_io_kwargs if media_io_kwargs else {}
        )
71
72
73
74
75
76
77
78
        self.connection = connection

        if allowed_local_media_path:
            allowed_local_media_path_ = Path(allowed_local_media_path)

            if not allowed_local_media_path_.exists():
                raise ValueError(
                    "Invalid `--allowed-local-media-path`: The path "
79
80
                    f"{allowed_local_media_path_} does not exist."
                )
81
82
83
            if not allowed_local_media_path_.is_dir():
                raise ValueError(
                    "Invalid `--allowed-local-media-path`: The path "
84
85
                    f"{allowed_local_media_path_} must be a directory."
                )
86
87
88
89
        else:
            allowed_local_media_path_ = None

        self.allowed_local_media_path = allowed_local_media_path_
90
91
92
        if allowed_media_domains is None:
            allowed_media_domains = []
        self.allowed_media_domains = allowed_media_domains
93
94
95
96
97

    def _load_data_url(
        self,
        url_spec: ParseResult,
        media_io: MediaIO[_M],
98
    ) -> _M:  # type: ignore[type-var]
99
100
101
102
103
104
105
106
107
108
109
110
111
        data_spec, data = url_spec.path.split(",", 1)
        media_type, data_type = data_spec.split(";", 1)

        if data_type != "base64":
            msg = "Only base64 data URLs are supported for now."
            raise NotImplementedError(msg)

        return media_io.load_base64(media_type, data)

    def _load_file_url(
        self,
        url_spec: ParseResult,
        media_io: MediaIO[_M],
112
    ) -> _M:  # type: ignore[type-var]
113
114
        allowed_local_media_path = self.allowed_local_media_path
        if allowed_local_media_path is None:
115
116
117
            raise RuntimeError(
                "Cannot load local files without `--allowed-local-media-path`."
            )
118

119
        filepath = Path(url2pathname(url_spec.path))
120
        if allowed_local_media_path not in filepath.resolve().parents:
121
            raise ValueError(
122
                f"The file path {filepath} must be a subpath "
123
124
                f"of `--allowed-local-media-path` {allowed_local_media_path}."
            )
125

126
        return media_io.load_file(filepath)
127

128
    def _assert_url_in_allowed_media_domains(self, url_spec) -> None:
129
130
131
132
        if (
            self.allowed_media_domains
            and url_spec.hostname not in self.allowed_media_domains
        ):
133
134
135
            raise ValueError(
                f"The URL must be from one of the allowed domains: "
                f"{self.allowed_media_domains}. Input URL domain: "
136
137
                f"{url_spec.hostname}"
            )
138

139
140
141
142
143
    def load_from_url(
        self,
        url: str,
        media_io: MediaIO[_M],
        *,
144
        fetch_timeout: int | None = None,
145
    ) -> _M:  # type: ignore[type-var]
146
        url_spec = urlparse(url)
147

148
        if url_spec.scheme.startswith("http"):
149
150
            self._assert_url_in_allowed_media_domains(url_spec)

151
            connection = self.connection
152
153
154
155
156
            data = connection.get_bytes(
                url,
                timeout=fetch_timeout,
                allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
            )
157

158
            return media_io.load_bytes(data)
159

160
161
        if url_spec.scheme == "data":
            return self._load_data_url(url_spec, media_io)
162

163
164
        if url_spec.scheme == "file":
            return self._load_file_url(url_spec, media_io)
165

166
167
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
168

169
170
171
172
173
    async def load_from_url_async(
        self,
        url: str,
        media_io: MediaIO[_M],
        *,
174
        fetch_timeout: int | None = None,
175
176
    ) -> _M:
        url_spec = urlparse(url)
177
        loop = asyncio.get_running_loop()
178

179
        if url_spec.scheme.startswith("http"):
180
181
            self._assert_url_in_allowed_media_domains(url_spec)

182
            connection = self.connection
183
184
185
186
187
            data = await connection.async_get_bytes(
                url,
                timeout=fetch_timeout,
                allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
            )
188
            future = loop.run_in_executor(global_thread_pool, media_io.load_bytes, data)
189
            return await future
190

191
        if url_spec.scheme == "data":
192
193
194
            future = loop.run_in_executor(
                global_thread_pool, self._load_data_url, url_spec, media_io
            )
195
            return await future
196

197
        if url_spec.scheme == "file":
198
199
200
            future = loop.run_in_executor(
                global_thread_pool, self._load_file_url, url_spec, media_io
            )
201
            return await future
202
203
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
204

205
206
207
    def fetch_audio(
        self,
        audio_url: str,
208
    ) -> tuple[np.ndarray, int | float]:
209
210
211
        """
        Load audio from a URL.
        """
212
        audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
213

214
        return self.load_from_url(
215
            audio_url,
216
217
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
218
        )
219

220
221
222
    async def fetch_audio_async(
        self,
        audio_url: str,
223
    ) -> tuple[np.ndarray, int | float]:
224
225
226
        """
        Asynchronously fetch audio from a URL.
        """
227
        audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
228

229
        return await self.load_from_url_async(
230
            audio_url,
231
232
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
233
        )
234

235
236
237
238
239
240
241
    def fetch_image(
        self,
        image_url: str,
        *,
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
242
        Load a PIL image from an HTTP or base64 data URL.
243

244
245
        By default, the image is converted into RGB format.
        """
246
247
248
        image_io = ImageMediaIO(
            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
        )
249

250
251
252
253
254
255
256
257
258
        try:
            return self.load_from_url(
                image_url,
                image_io,
                fetch_timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
            )
        except UnidentifiedImageError as e:
            # convert to ValueError to be properly caught upstream
            raise ValueError(str(e)) from e
259

260
261
    async def fetch_image_async(
        self,
262
263
        image_url: str,
        *,
264
265
266
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
267
        Asynchronously load a PIL image from an HTTP or base64 data URL.
268

269
270
        By default, the image is converted into RGB format.
        """
271
272
273
        image_io = ImageMediaIO(
            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
        )
274

275
276
277
278
279
280
281
282
283
        try:
            return await self.load_from_url_async(
                image_url,
                image_io,
                fetch_timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
            )
        except UnidentifiedImageError as e:
            # convert to ValueError to be properly caught upstream
            raise ValueError(str(e)) from e
284

285
286
287
288
289
    def fetch_video(
        self,
        video_url: str,
        *,
        image_mode: str = "RGB",
290
    ) -> tuple[npt.NDArray, dict[str, Any]]:
291
        """
292
        Load video from an HTTP or base64 data URL.
293
        """
294
295
296
297
        image_io = ImageMediaIO(
            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
        )
        video_io = VideoMediaIO(image_io, **self.media_io_kwargs.get("video", {}))
298
299
300
301
302
303

        return self.load_from_url(
            video_url,
            video_io,
            fetch_timeout=envs.VLLM_VIDEO_FETCH_TIMEOUT,
        )
304

305
306
307
    async def fetch_video_async(
        self,
        video_url: str,
308
        *,
309
        image_mode: str = "RGB",
310
    ) -> tuple[npt.NDArray, dict[str, Any]]:
311
        """
312
        Asynchronously load video from an HTTP or base64 data URL.
313
314
315

        By default, the image is converted into RGB format.
        """
316
317
318
319
        image_io = ImageMediaIO(
            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
        )
        video_io = VideoMediaIO(image_io, **self.media_io_kwargs.get("video", {}))
320
321
322
323
324
325

        return await self.load_from_url_async(
            video_url,
            video_io,
            fetch_timeout=envs.VLLM_VIDEO_FETCH_TIMEOUT,
        )
326

327
328
329
330
331
332
333
334
335
336
337
    def fetch_image_embedding(
        self,
        data: str,
    ) -> torch.Tensor:
        """
        Load image embedding from a URL.
        """
        image_embedding_io = ImageEmbeddingMediaIO()

        return image_embedding_io.load_base64("", data)

338

339
340
def encode_audio_base64(
    audio: np.ndarray,
341
    sampling_rate: int,
342
343
) -> str:
    """Encode audio as base64."""
344
345
    audio_io = AudioMediaIO()
    return audio_io.encode_base64((audio, sampling_rate))
346
347


348
349
350
351
352
353
354
355
def encode_image_base64(
    image: Image.Image,
    *,
    image_mode: str = "RGB",
    format: str = "JPEG",
) -> str:
    """
    Encode a pillow image to base64 format.
356

357
358
    By default, the image is converted into RGB format before being encoded.
    """
359
360
    image_io = ImageMediaIO(image_mode=image_mode)
    return image_io.encode_base64(image, image_format=format)
361
362


363
def encode_video_base64(frames: npt.NDArray) -> str:
364
365
366
    image_io = ImageMediaIO()
    video_io = VideoMediaIO(image_io)
    return video_io.encode_base64(frames)
367
368


369
def argsort_mm_positions(
370
371
    mm_positions: MultiModalPlaceholderDict,
) -> list[tuple[str, int]]:
372
373
374
375
    """
    Given a `MultiModalPlaceholderDict`, output a sequence of keys to
    sort the dictionary by `offset` (starting index in the input sequence)
    in ascending order.
376
377

    Returns:
378
379
        A list of `(modality, idx)`, which can be used to access an item
        by `mm_positions[modality][idx]`.
380
    """
381
382
383
384
385
    flat_items = (
        (modality, idx, item)
        for modality, items in mm_positions.items()
        for idx, item in enumerate(items)
    )
386

387
    sorted_flat_items = sorted(flat_items, key=lambda x: x[2].offset)
388

389
    return [(modality, idx) for modality, idx, _ in sorted_flat_items]
390
391


392
393
394
395
396
def group_mm_kwargs_by_modality(
    mm_kwargs: list[MultiModalKwargsItem],
    *,
    device: torch.types.Device = None,
    pin_memory: bool = False,
397
    merge_by_field_config: bool | None = None,
398
399
400
401
402
) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
    """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same
    modality together into the same `MultiModalKwargs` instance.

    Args:
403
404
405
        mm_kwargs: List of `MultiModalKwargsItem`.
        device: The device to place the grouped tensors on.
        pin_memory: Whether to pin memory for faster host-to-device transfer.
406
407
408
409

    Yields:
        A tuple `(modality, num_items, grouped_kwargs)`.
    """
410
411
412
413
    if merge_by_field_config is None:
        raise RuntimeError(
            "`group_mm_kwargs_by_modality` now requires "
            "`merge_by_field_config` arg, please update your model runner "
414
415
            "according to https://github.com/vllm-project/vllm/pull/25676."
        )
416

417
    from vllm.multimodal.inputs import MultiModalKwargs, MultiModalKwargsItems
418
419
420
421

    for modality, items in groupby(mm_kwargs, key=lambda item: item.modality):
        items_lst = list(items)

422
423
        # TODO: Deprecate `merge_by_field_config` once
        # we have migrated all in-tree models
424
425
426
        if merge_by_field_config:
            mm_kwargs_group: BatchedTensorInputs = dict(
                MultiModalKwargsItems.from_seq(items_lst).get_data(
427
428
429
                    pin_memory=pin_memory
                )
            )
430
431
432

            if device is not None:
                mm_kwargs_group = json_map_leaves(
433
                    lambda x: x.to(device=device) if isinstance(x, torch.Tensor) else x,
434
435
436
437
438
439
440
441
442
443
444
445
446
                    mm_kwargs_group,
                )
        else:
            mm_kwargs_group = MultiModalKwargs.as_kwargs(
                MultiModalKwargs.batch(
                    [
                        MultiModalKwargsItems.from_seq([item]).get_data()
                        for item in items_lst
                    ],
                    pin_memory=pin_memory,
                ),
                device=device,
            )
447
448
449
450

        yield modality, len(items_lst), mm_kwargs_group


451
452
def fetch_audio(
    audio_url: str,
453
454
    audio_io_kwargs: dict[str, Any] | None = None,
) -> tuple[np.ndarray, int | float]:
455
456
457
458
459
    """
    Args:
        audio_url: URL of the audio file to fetch.
        audio_io_kwargs: Additional kwargs passed to handle audio IO.
    """
460
    media_io_kwargs = None if not audio_io_kwargs else {"audio": audio_io_kwargs}
461
462
463
464
465
466
    media_connector = MediaConnector(media_io_kwargs=media_io_kwargs)
    return media_connector.fetch_audio(audio_url)


def fetch_image(
    image_url: str,
467
    image_io_kwargs: dict[str, Any] | None = None,
468
469
470
471
472
473
) -> Image.Image:
    """
    Args:
        image_url: URL of the image file to fetch.
        image_io_kwargs: Additional kwargs passed to handle image IO.
    """
474
    media_io_kwargs = None if not image_io_kwargs else {"image": image_io_kwargs}
475
476
477
478
479
480
    media_connector = MediaConnector(media_io_kwargs=media_io_kwargs)
    return media_connector.fetch_image(image_url)


def fetch_video(
    video_url: str,
481
    video_io_kwargs: dict[str, Any] | None = None,
482
483
484
485
486
487
) -> tuple[npt.NDArray, dict[str, Any]]:
    """
    Args:
        video_url: URL of the video file to fetch.
        video_io_kwargs: Additional kwargs passed to handle video IO.
    """
488
    media_io_kwargs = None if not video_io_kwargs else {"video": video_io_kwargs}
489
    media_connector = MediaConnector(media_io_kwargs=media_io_kwargs)
490
    return media_connector.fetch_video(video_url)