utils.py 15.2 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

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
    from .inputs import (BatchedTensorInputs, MultiModalKwargsItem,
                         MultiModalKwargsItems, MultiModalPlaceholderDict)
33
else:
34
35
    BatchedTensorInputs = Any
    MultiModalKwargsItem = Any
36
    MultiModalKwargsItems = Any
37
    MultiModalPlaceholderDict = Any
38

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

43

44
class MediaConnector:
45

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

65
66
        self.media_io_kwargs: dict[str, dict[
            str, Any]] = media_io_kwargs if media_io_kwargs else {}
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
        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],
89
    ) -> _M:  # type: ignore[type-var]
90
91
92
93
94
95
96
97
98
99
100
101
102
        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],
103
    ) -> _M:  # type: ignore[type-var]
104
105
106
107
108
        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`.")

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

115
        return media_io.load_file(filepath)
116

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

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

130
            return media_io.load_bytes(data)
131

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

303

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


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

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


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


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

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

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

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


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

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

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

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

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


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

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

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

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

403
        # mm_kwargs_group = MultiModalKwargsItems.from_items(items_lst) \
404
        #    .get_data(pin_memory=pin_memory)
405
406

        # if device is not None:
407
408
409
410
        #     mm_kwargs_group = json_map_leaves(
        #         lambda x: x.to(device=device),
        #         mm_kwargs_group,
        #     )
411
412
413
414
415
416
417

        # TODO: Once V0 is removed, we can use the merging logic above
        # 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`.
        mm_kwargs_group = MultiModalKwargs.as_kwargs(
            MultiModalKwargs.batch(
418
419
420
421
                [
                    MultiModalKwargsItems.from_seq([item]).get_data()
                    for item in items_lst
                ],
422
423
424
425
426
427
428
429
                pin_memory=pin_memory,
            ),
            device=device,
        )

        yield modality, len(items_lst), mm_kwargs_group


430
431
432
433
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
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)
475
    return media_connector.fetch_video(video_url)