utils.py 16.1 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
143
            connection = self.connection
            data = connection.get_bytes(url, timeout=fetch_timeout)
144

145
            return media_io.load_bytes(data)
146

147
148
        if url_spec.scheme == "data":
            return self._load_data_url(url_spec, media_io)
149

150
151
        if url_spec.scheme == "file":
            return self._load_file_url(url_spec, media_io)
152

153
154
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
155

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

166
        if url_spec.scheme.startswith("http"):
167
168
            self._assert_url_in_allowed_media_domains(url_spec)

169
170
            connection = self.connection
            data = await connection.async_get_bytes(url, timeout=fetch_timeout)
171
172
173
            future = loop.run_in_executor(global_thread_pool,
                                          media_io.load_bytes, data)
            return await future
174

175
        if url_spec.scheme == "data":
176
177
178
179
            future = loop.run_in_executor(global_thread_pool,
                                          self._load_data_url, url_spec,
                                          media_io)
            return await future
180

181
        if url_spec.scheme == "file":
182
183
184
185
            future = loop.run_in_executor(global_thread_pool,
                                          self._load_file_url, url_spec,
                                          media_io)
            return await future
186
187
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
188

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

198
        return self.load_from_url(
199
            audio_url,
200
201
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
202
        )
203

204
205
206
207
208
209
210
    async def fetch_audio_async(
        self,
        audio_url: str,
    ) -> tuple[np.ndarray, Union[int, float]]:
        """
        Asynchronously fetch audio from a URL.
        """
211
        audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
212

213
        return await self.load_from_url_async(
214
            audio_url,
215
216
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
217
        )
218

219
220
221
222
223
224
225
    def fetch_image(
        self,
        image_url: str,
        *,
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
226
        Load a PIL image from an HTTP or base64 data URL.
227

228
229
        By default, the image is converted into RGB format.
        """
230
231
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
232

233
234
235
236
237
238
239
240
241
        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
242

243
244
    async def fetch_image_async(
        self,
245
246
        image_url: str,
        *,
247
248
249
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
250
        Asynchronously load a PIL image from an HTTP or base64 data URL.
251

252
253
        By default, the image is converted into RGB format.
        """
254
255
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
256

257
258
259
260
261
262
263
264
265
        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
266

267
268
269
270
271
    def fetch_video(
        self,
        video_url: str,
        *,
        image_mode: str = "RGB",
272
    ) -> tuple[npt.NDArray, dict[str, Any]]:
273
        """
274
        Load video from an HTTP or base64 data URL.
275
        """
276
277
278
279
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
        video_io = VideoMediaIO(image_io,
                                **self.media_io_kwargs.get("video", {}))
280
281
282
283
284
285

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

287
288
289
    async def fetch_video_async(
        self,
        video_url: str,
290
        *,
291
        image_mode: str = "RGB",
292
    ) -> tuple[npt.NDArray, dict[str, Any]]:
293
        """
294
        Asynchronously load video from an HTTP or base64 data URL.
295
296
297

        By default, the image is converted into RGB format.
        """
298
299
300
301
        image_io = ImageMediaIO(image_mode=image_mode,
                                **self.media_io_kwargs.get("image", {}))
        video_io = VideoMediaIO(image_io,
                                **self.media_io_kwargs.get("video", {}))
302
303
304
305
306
307

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

309
310
311
312
313
314
315
316
317
318
319
    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)

320

321
322
def encode_audio_base64(
    audio: np.ndarray,
323
    sampling_rate: int,
324
325
) -> str:
    """Encode audio as base64."""
326
327
    audio_io = AudioMediaIO()
    return audio_io.encode_base64((audio, sampling_rate))
328
329


330
331
332
333
334
335
336
337
def encode_image_base64(
    image: Image.Image,
    *,
    image_mode: str = "RGB",
    format: str = "JPEG",
) -> str:
    """
    Encode a pillow image to base64 format.
338

339
340
    By default, the image is converted into RGB format before being encoded.
    """
341
342
    image_io = ImageMediaIO(image_mode=image_mode)
    return image_io.encode_base64(image, image_format=format)
343
344


345
def encode_video_base64(frames: npt.NDArray) -> str:
346
347
348
    image_io = ImageMediaIO()
    video_io = VideoMediaIO(image_io)
    return video_io.encode_base64(frames)
349
350


351
352
353
354
355
356
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.
357
358

    Returns:
359
360
        A list of `(modality, idx)`, which can be used to access an item
        by `mm_positions[modality][idx]`.
361
    """
362
363
364
    flat_items = ((modality, idx, item)
                  for modality, items in mm_positions.items()
                  for idx, item in enumerate(items))
365

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

368
    return [(modality, idx) for modality, idx, _ in sorted_flat_items]
369
370


371
372
373
374
# 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.")
375
def group_mm_inputs_by_modality(
376
377
    mm_inputs: list[MultiModalKwargsItems]
) -> list[list[MultiModalKwargsItems]]:
378
379
380
    if not mm_inputs:
        return []

381
382
    def modality_group_func(
            mm_input: MultiModalKwargsItems) -> Union[str, int]:
383
        # If the input has multiple modalities, return an id as the unique key
384
        # for the mm_input input.
385
        if len(mm_input) > 1:
386
387
            return id(mm_input)

388
389
        elif len(mm_input) == 1:
            return next(iter(mm_input.keys()))
390

391
        raise AssertionError("This line should be unreachable.")
392
393
394
395

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


398
399
400
401
402
def group_mm_kwargs_by_modality(
    mm_kwargs: list[MultiModalKwargsItem],
    *,
    device: torch.types.Device = None,
    pin_memory: bool = False,
403
    merge_by_field_config: bool = False,
404
405
406
407
408
) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
    """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same
    modality together into the same `MultiModalKwargs` instance.

    Args:
409
410
411
        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.
412
413
414
415

    Yields:
        A tuple `(modality, num_items, grouped_kwargs)`.
    """
416
    from vllm.multimodal.inputs import MultiModalKwargs, MultiModalKwargsItems
417
418
419
420

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

421
        # TODO: Enable `merge_by_field_config` for all models
422
423
424
        # 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`.
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
        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,
            )
446
447
448
449

        yield modality, len(items_lst), mm_kwargs_group


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
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
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)
495
    return media_connector.fetch_video(video_url)