utils.py 15.4 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
53
        connection: HTTPConnection = global_http_connection,
        *,
        allowed_local_media_path: str = "",
    ) -> None:
54
55
56
57
58
        """
        Args:
            media_io_kwargs: Additional args passed to process media 
                             inputs, keyed by modalities. For example, 
                             to set num_frames for video, set 
59
                             `--media-io-kwargs '{"video":{"num_frames":40}}'`
60
            connection: HTTP connection client to download media contents.
61
62
            allowed_local_media_path: A local directory to load media files
                                      from.
63
        """
64
65
        super().__init__()

66
67
        self.media_io_kwargs: dict[str, dict[
            str, Any]] = media_io_kwargs if media_io_kwargs else {}
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
        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_

    def _load_data_url(
        self,
        url_spec: ParseResult,
        media_io: MediaIO[_M],
90
    ) -> _M:  # type: ignore[type-var]
91
92
93
94
95
96
97
98
99
100
101
102
103
        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],
104
    ) -> _M:  # type: ignore[type-var]
105
106
107
108
109
        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`.")

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

116
        return media_io.load_file(filepath)
117

118
119
120
121
122
123
    def load_from_url(
        self,
        url: str,
        media_io: MediaIO[_M],
        *,
        fetch_timeout: Optional[int] = None,
124
    ) -> _M:  # type: ignore[type-var]
125
        url_spec = urlparse(url)
126

127
128
129
        if url_spec.scheme.startswith("http"):
            connection = self.connection
            data = connection.get_bytes(url, timeout=fetch_timeout)
130

131
            return media_io.load_bytes(data)
132

133
134
        if url_spec.scheme == "data":
            return self._load_data_url(url_spec, media_io)
135

136
137
        if url_spec.scheme == "file":
            return self._load_file_url(url_spec, media_io)
138

139
140
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
141

142
143
144
145
146
147
148
149
    async def load_from_url_async(
        self,
        url: str,
        media_io: MediaIO[_M],
        *,
        fetch_timeout: Optional[int] = None,
    ) -> _M:
        url_spec = urlparse(url)
150
        loop = asyncio.get_running_loop()
151

152
153
154
        if url_spec.scheme.startswith("http"):
            connection = self.connection
            data = await connection.async_get_bytes(url, timeout=fetch_timeout)
155
156
157
            future = loop.run_in_executor(global_thread_pool,
                                          media_io.load_bytes, data)
            return await future
158

159
        if url_spec.scheme == "data":
160
161
162
163
            future = loop.run_in_executor(global_thread_pool,
                                          self._load_data_url, url_spec,
                                          media_io)
            return await future
164

165
        if url_spec.scheme == "file":
166
167
168
169
            future = loop.run_in_executor(global_thread_pool,
                                          self._load_file_url, url_spec,
                                          media_io)
            return await future
170
171
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
172

173
174
175
176
177
178
179
    def fetch_audio(
        self,
        audio_url: str,
    ) -> tuple[np.ndarray, Union[int, float]]:
        """
        Load audio from a URL.
        """
180
        audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
181

182
        return self.load_from_url(
183
            audio_url,
184
185
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
186
        )
187

188
189
190
191
192
193
194
    async def fetch_audio_async(
        self,
        audio_url: str,
    ) -> tuple[np.ndarray, Union[int, float]]:
        """
        Asynchronously fetch audio from a URL.
        """
195
        audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
196

197
        return await self.load_from_url_async(
198
            audio_url,
199
200
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
201
        )
202

203
204
205
206
207
208
209
    def fetch_image(
        self,
        image_url: str,
        *,
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
210
        Load a PIL image from an HTTP or base64 data URL.
211

212
213
        By default, the image is converted into RGB format.
        """
214
215
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
216

217
218
219
220
221
222
223
224
225
        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
226

227
228
    async def fetch_image_async(
        self,
229
230
        image_url: str,
        *,
231
232
233
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
234
        Asynchronously 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 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
250

251
252
253
254
255
    def fetch_video(
        self,
        video_url: str,
        *,
        image_mode: str = "RGB",
256
    ) -> tuple[npt.NDArray, dict[str, Any]]:
257
        """
258
        Load video from an HTTP or base64 data URL.
259
        """
260
261
262
263
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
        video_io = VideoMediaIO(image_io,
                                **self.media_io_kwargs.get("video", {}))
264
265
266
267
268
269

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

271
272
273
    async def fetch_video_async(
        self,
        video_url: str,
274
        *,
275
        image_mode: str = "RGB",
276
    ) -> tuple[npt.NDArray, dict[str, Any]]:
277
        """
278
        Asynchronously load video from an HTTP or base64 data URL.
279
280
281

        By default, the image is converted into RGB format.
        """
282
283
284
285
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
        video_io = VideoMediaIO(image_io,
                                **self.media_io_kwargs.get("video", {}))
286
287
288
289
290
291

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

293
294
295
296
297
298
299
300
301
302
303
    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)

304

305
306
def encode_audio_base64(
    audio: np.ndarray,
307
    sampling_rate: int,
308
309
) -> str:
    """Encode audio as base64."""
310
311
    audio_io = AudioMediaIO()
    return audio_io.encode_base64((audio, sampling_rate))
312
313


314
315
316
317
318
319
320
321
def encode_image_base64(
    image: Image.Image,
    *,
    image_mode: str = "RGB",
    format: str = "JPEG",
) -> str:
    """
    Encode a pillow image to base64 format.
322

323
324
    By default, the image is converted into RGB format before being encoded.
    """
325
326
    image_io = ImageMediaIO(image_mode=image_mode)
    return image_io.encode_base64(image, image_format=format)
327
328


329
def encode_video_base64(frames: npt.NDArray) -> str:
330
331
332
    image_io = ImageMediaIO()
    video_io = VideoMediaIO(image_io)
    return video_io.encode_base64(frames)
333
334


335
336
337
338
339
340
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.
341
342

    Returns:
343
344
        A list of `(modality, idx)`, which can be used to access an item
        by `mm_positions[modality][idx]`.
345
    """
346
347
348
    flat_items = ((modality, idx, item)
                  for modality, items in mm_positions.items()
                  for idx, item in enumerate(items))
349

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

352
    return [(modality, idx) for modality, idx, _ in sorted_flat_items]
353
354


355
356
357
358
# 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.")
359
def group_mm_inputs_by_modality(
360
361
    mm_inputs: list[MultiModalKwargsItems]
) -> list[list[MultiModalKwargsItems]]:
362
363
364
    if not mm_inputs:
        return []

365
366
    def modality_group_func(
            mm_input: MultiModalKwargsItems) -> Union[str, int]:
367
        # If the input has multiple modalities, return an id as the unique key
368
        # for the mm_input input.
369
        if len(mm_input) > 1:
370
371
            return id(mm_input)

372
373
        elif len(mm_input) == 1:
            return next(iter(mm_input.keys()))
374

375
        raise AssertionError("This line should be unreachable.")
376
377
378
379

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


382
383
384
385
386
def group_mm_kwargs_by_modality(
    mm_kwargs: list[MultiModalKwargsItem],
    *,
    device: torch.types.Device = None,
    pin_memory: bool = False,
387
    merge_by_field_config: bool = False,
388
389
390
391
392
) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
    """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same
    modality together into the same `MultiModalKwargs` instance.

    Args:
393
394
395
        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.
396
397
398
399

    Yields:
        A tuple `(modality, num_items, grouped_kwargs)`.
    """
400
    from vllm.multimodal.inputs import MultiModalKwargs, MultiModalKwargsItems
401
402
403
404

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

405
        # TODO: Enable `merge_by_field_config` for all models
406
407
408
        # 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`.
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
        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,
            )
430
431
432
433

        yield modality, len(items_lst), mm_kwargs_group


434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
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)
479
    return media_connector.fetch_video(video_url)