utils.py 16.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, Optional, TypeVar, Union
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
from typing_extensions import deprecated
19

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

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

29
_M = TypeVar("_M")
30

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

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

44

45
class MediaConnector:
46

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

67
68
        self.media_io_kwargs: dict[str, dict[
            str, Any]] = media_io_kwargs if media_io_kwargs else {}
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
        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 "
                    f"{allowed_local_media_path_} does not exist.")
            if not allowed_local_media_path_.is_dir():
                raise ValueError(
                    "Invalid `--allowed-local-media-path`: The path "
                    f"{allowed_local_media_path_} must be a directory.")
        else:
            allowed_local_media_path_ = None

        self.allowed_local_media_path = allowed_local_media_path_
86
87
88
        if allowed_media_domains is None:
            allowed_media_domains = []
        self.allowed_media_domains = allowed_media_domains
89
90
91
92
93

    def _load_data_url(
        self,
        url_spec: ParseResult,
        media_io: MediaIO[_M],
94
    ) -> _M:  # type: ignore[type-var]
95
96
97
98
99
100
101
102
103
104
105
106
107
        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],
108
    ) -> _M:  # type: ignore[type-var]
109
110
111
112
113
        allowed_local_media_path = self.allowed_local_media_path
        if allowed_local_media_path is None:
            raise RuntimeError("Cannot load local files without "
                               "`--allowed-local-media-path`.")

114
        filepath = Path(url2pathname(url_spec.path))
115
        if allowed_local_media_path not in filepath.resolve().parents:
116
            raise ValueError(
117
118
                f"The file path {filepath} must be a subpath "
                f"of `--allowed-local-media-path` {allowed_local_media_path}.")
119

120
        return media_io.load_file(filepath)
121

122
123
124
125
126
127
128
129
    def _assert_url_in_allowed_media_domains(self, url_spec) -> None:
        if self.allowed_media_domains and url_spec.hostname not in \
            self.allowed_media_domains:
            raise ValueError(
                f"The URL must be from one of the allowed domains: "
                f"{self.allowed_media_domains}. Input URL domain: "
                f"{url_spec.hostname}")

130
131
132
133
134
135
    def load_from_url(
        self,
        url: str,
        media_io: MediaIO[_M],
        *,
        fetch_timeout: Optional[int] = None,
136
    ) -> _M:  # type: ignore[type-var]
137
        url_spec = urlparse(url)
138

139
        if url_spec.scheme.startswith("http"):
140
141
            self._assert_url_in_allowed_media_domains(url_spec)

142
            connection = self.connection
143
144
145
146
147
            data = connection.get_bytes(
                url,
                timeout=fetch_timeout,
                allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
            )
148

149
            return media_io.load_bytes(data)
150

151
152
        if url_spec.scheme == "data":
            return self._load_data_url(url_spec, media_io)
153

154
155
        if url_spec.scheme == "file":
            return self._load_file_url(url_spec, media_io)
156

157
158
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
159

160
161
162
163
164
165
166
167
    async def load_from_url_async(
        self,
        url: str,
        media_io: MediaIO[_M],
        *,
        fetch_timeout: Optional[int] = None,
    ) -> _M:
        url_spec = urlparse(url)
168
        loop = asyncio.get_running_loop()
169

170
        if url_spec.scheme.startswith("http"):
171
172
            self._assert_url_in_allowed_media_domains(url_spec)

173
            connection = self.connection
174
175
176
177
178
            data = await connection.async_get_bytes(
                url,
                timeout=fetch_timeout,
                allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
            )
179
180
181
            future = loop.run_in_executor(global_thread_pool,
                                          media_io.load_bytes, data)
            return await future
182

183
        if url_spec.scheme == "data":
184
185
186
187
            future = loop.run_in_executor(global_thread_pool,
                                          self._load_data_url, url_spec,
                                          media_io)
            return await future
188

189
        if url_spec.scheme == "file":
190
191
192
193
            future = loop.run_in_executor(global_thread_pool,
                                          self._load_file_url, url_spec,
                                          media_io)
            return await future
194
195
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
196

197
198
199
200
201
202
203
    def fetch_audio(
        self,
        audio_url: str,
    ) -> tuple[np.ndarray, Union[int, float]]:
        """
        Load audio from a URL.
        """
204
        audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
205

206
        return self.load_from_url(
207
            audio_url,
208
209
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
210
        )
211

212
213
214
215
216
217
218
    async def fetch_audio_async(
        self,
        audio_url: str,
    ) -> tuple[np.ndarray, Union[int, float]]:
        """
        Asynchronously fetch audio from a URL.
        """
219
        audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
220

221
        return await self.load_from_url_async(
222
            audio_url,
223
224
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
225
        )
226

227
228
229
230
231
232
233
    def fetch_image(
        self,
        image_url: str,
        *,
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
234
        Load a PIL image from an HTTP or base64 data URL.
235

236
237
        By default, the image is converted into RGB format.
        """
238
239
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
240

241
242
243
244
245
246
247
248
249
        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
250

251
252
    async def fetch_image_async(
        self,
253
254
        image_url: str,
        *,
255
256
257
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
258
        Asynchronously load a PIL image from an HTTP or base64 data URL.
259

260
261
        By default, the image is converted into RGB format.
        """
262
263
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
264

265
266
267
268
269
270
271
272
273
        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
274

275
276
277
278
279
    def fetch_video(
        self,
        video_url: str,
        *,
        image_mode: str = "RGB",
280
    ) -> tuple[npt.NDArray, dict[str, Any]]:
281
        """
282
        Load video from an HTTP or base64 data URL.
283
        """
284
285
286
287
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
        video_io = VideoMediaIO(image_io,
                                **self.media_io_kwargs.get("video", {}))
288
289
290
291
292
293

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

295
296
297
    async def fetch_video_async(
        self,
        video_url: str,
298
        *,
299
        image_mode: str = "RGB",
300
    ) -> tuple[npt.NDArray, dict[str, Any]]:
301
        """
302
        Asynchronously load video from an HTTP or base64 data URL.
303
304
305

        By default, the image is converted into RGB format.
        """
306
307
308
309
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
        video_io = VideoMediaIO(image_io,
                                **self.media_io_kwargs.get("video", {}))
310
311
312
313
314
315

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

317
318
319
320
321
322
323
324
325
326
327
    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)

328

329
330
def encode_audio_base64(
    audio: np.ndarray,
331
    sampling_rate: int,
332
333
) -> str:
    """Encode audio as base64."""
334
335
    audio_io = AudioMediaIO()
    return audio_io.encode_base64((audio, sampling_rate))
336
337


338
339
340
341
342
343
344
345
def encode_image_base64(
    image: Image.Image,
    *,
    image_mode: str = "RGB",
    format: str = "JPEG",
) -> str:
    """
    Encode a pillow image to base64 format.
346

347
348
    By default, the image is converted into RGB format before being encoded.
    """
349
350
    image_io = ImageMediaIO(image_mode=image_mode)
    return image_io.encode_base64(image, image_format=format)
351
352


353
def encode_video_base64(frames: npt.NDArray) -> str:
354
355
356
    image_io = ImageMediaIO()
    video_io = VideoMediaIO(image_io)
    return video_io.encode_base64(frames)
357
358


359
360
361
362
363
364
def argsort_mm_positions(
        mm_positions: MultiModalPlaceholderDict) -> list[tuple[str, int]]:
    """
    Given a `MultiModalPlaceholderDict`, output a sequence of keys to
    sort the dictionary by `offset` (starting index in the input sequence)
    in ascending order.
365
366

    Returns:
367
368
        A list of `(modality, idx)`, which can be used to access an item
        by `mm_positions[modality][idx]`.
369
    """
370
371
372
    flat_items = ((modality, idx, item)
                  for modality, items in mm_positions.items()
                  for idx, item in enumerate(items))
373

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

376
    return [(modality, idx) for modality, idx, _ in sorted_flat_items]
377
378


379
380
381
382
# Temporary back-compatibility for plugins that define model runner
@deprecated("`group_mm_inputs_by_modality` is superseded by "
            "`group_mm_kwargs_by_modality` and will be removed in v0.13. "
            "Please use `group_mm_kwargs_by_modality` instead.")
383
def group_mm_inputs_by_modality(
384
385
    mm_inputs: list[MultiModalKwargsItems]
) -> list[list[MultiModalKwargsItems]]:
386
387
388
    if not mm_inputs:
        return []

389
390
    def modality_group_func(
            mm_input: MultiModalKwargsItems) -> Union[str, int]:
391
        # If the input has multiple modalities, return an id as the unique key
392
        # for the mm_input input.
393
        if len(mm_input) > 1:
394
395
            return id(mm_input)

396
397
        elif len(mm_input) == 1:
            return next(iter(mm_input.keys()))
398

399
        raise AssertionError("This line should be unreachable.")
400
401
402
403

    return [
        list(group) for _, group in groupby(mm_inputs, key=modality_group_func)
    ]
404
405


406
407
408
409
410
def group_mm_kwargs_by_modality(
    mm_kwargs: list[MultiModalKwargsItem],
    *,
    device: torch.types.Device = None,
    pin_memory: bool = False,
411
    merge_by_field_config: bool = False,
412
413
414
415
416
) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
    """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same
    modality together into the same `MultiModalKwargs` instance.

    Args:
417
418
419
        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.
420
421
422
423

    Yields:
        A tuple `(modality, num_items, grouped_kwargs)`.
    """
424
    from vllm.multimodal.inputs import MultiModalKwargs, MultiModalKwargsItems
425
426
427
428

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

429
        # TODO: Enable `merge_by_field_config` for all models
430
431
432
        # to avoid creating an extra batch dimension (except for fields
        # that are meant to be stacked anyway).
        # We will also need to update each model to remove `flatten_bn`.
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
        if merge_by_field_config:
            mm_kwargs_group: BatchedTensorInputs = dict(
                MultiModalKwargsItems.from_seq(items_lst).get_data(
                    pin_memory=pin_memory))

            if device is not None:
                mm_kwargs_group = json_map_leaves(
                    lambda x: x.to(device=device),
                    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,
            )
454
455
456
457

        yield modality, len(items_lst), mm_kwargs_group


458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
def fetch_audio(
    audio_url: str,
    audio_io_kwargs: Optional[dict[str, Any]] = None,
) -> tuple[np.ndarray, Union[int, float]]:
    """
    Args:
        audio_url: URL of the audio file to fetch.
        audio_io_kwargs: Additional kwargs passed to handle audio IO.
    """
    media_io_kwargs = None if not audio_io_kwargs else {
        "audio": audio_io_kwargs
    }
    media_connector = MediaConnector(media_io_kwargs=media_io_kwargs)
    return media_connector.fetch_audio(audio_url)


def fetch_image(
    image_url: str,
    image_io_kwargs: Optional[dict[str, Any]] = None,
) -> Image.Image:
    """
    Args:
        image_url: URL of the image file to fetch.
        image_io_kwargs: Additional kwargs passed to handle image IO.
    """
    media_io_kwargs = None if not image_io_kwargs else {
        "image": image_io_kwargs
    }
    media_connector = MediaConnector(media_io_kwargs=media_io_kwargs)
    return media_connector.fetch_image(image_url)


def fetch_video(
    video_url: str,
    video_io_kwargs: Optional[dict[str, Any]] = None,
) -> 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.
    """
    media_io_kwargs = None if not video_io_kwargs else {
        "video": video_io_kwargs
    }
    media_connector = MediaConnector(media_io_kwargs=media_io_kwargs)
503
    return media_connector.fetch_video(video_url)