utils.py 16.5 KB
Newer Older
1
from functools import lru_cache
2
from itertools import groupby
3
from pathlib import Path
4
from typing import TYPE_CHECKING, Optional, TypeVar, Union
5
from urllib.parse import ParseResult, urlparse
6

7
import numpy as np
8
import numpy.typing as npt
9
10
from PIL import Image

11
import vllm.envs as envs
12
from vllm.connections import HTTPConnection, global_http_connection
13
14
15
from vllm.logger import init_logger
from vllm.transformers_utils.tokenizer import AnyTokenizer, get_tokenizer

16
17
18
19
20
from .audio import AudioMediaIO
from .base import MediaIO
from .image import ImageMediaIO
from .inputs import PlaceholderRange
from .video import VideoMediaIO
21

22
23
24
logger = init_logger(__name__)

cached_get_tokenizer = lru_cache(get_tokenizer)
25

26
_M = TypeVar("_M")
27

28
29
if TYPE_CHECKING:
    from .hasher import MultiModalHashDict
30
    from .inputs import MultiModalKwargs, MultiModalPlaceholderDict
31

32

33
class MediaConnector:
34

35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
    def __init__(
        self,
        connection: HTTPConnection = global_http_connection,
        *,
        allowed_local_media_path: str = "",
    ) -> None:
        super().__init__()

        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],
    ) -> _M:
        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],
    ) -> _M:
        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`.")

        filepath = Path(url_spec.path)
        if allowed_local_media_path not in filepath.resolve().parents:
87
            raise ValueError(
88
89
                f"The file path {filepath} must be a subpath "
                f"of `--allowed-local-media-path` {allowed_local_media_path}.")
90

91
        return media_io.load_file(filepath)
92

93
94
95
96
97
98
99
100
    def load_from_url(
        self,
        url: str,
        media_io: MediaIO[_M],
        *,
        fetch_timeout: Optional[int] = None,
    ) -> _M:
        url_spec = urlparse(url)
101

102
103
104
        if url_spec.scheme.startswith("http"):
            connection = self.connection
            data = connection.get_bytes(url, timeout=fetch_timeout)
105

106
            return media_io.load_bytes(data)
107

108
109
        if url_spec.scheme == "data":
            return self._load_data_url(url_spec, media_io)
110

111
112
        if url_spec.scheme == "file":
            return self._load_file_url(url_spec, media_io)
113

114
115
        msg = "The URL must be either a HTTP, data or file URL."
        raise ValueError(msg)
116

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

126
127
128
        if url_spec.scheme.startswith("http"):
            connection = self.connection
            data = await connection.async_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
    def fetch_audio(
        self,
        audio_url: str,
    ) -> tuple[np.ndarray, Union[int, float]]:
        """
        Load audio from a URL.
        """
        audio_io = AudioMediaIO()
149

150
        return self.load_from_url(
151
            audio_url,
152
153
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
154
        )
155

156
157
158
159
160
161
162
163
    async def fetch_audio_async(
        self,
        audio_url: str,
    ) -> tuple[np.ndarray, Union[int, float]]:
        """
        Asynchronously fetch audio from a URL.
        """
        audio_io = AudioMediaIO()
164

165
        return await self.load_from_url_async(
166
            audio_url,
167
168
            audio_io,
            fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
169
        )
170

171
172
173
174
175
176
177
178
    def fetch_image(
        self,
        image_url: str,
        *,
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
        Load a PIL image from a HTTP or base64 data URL.
179

180
181
182
        By default, the image is converted into RGB format.
        """
        image_io = ImageMediaIO(image_mode=image_mode)
183

184
185
186
187
188
        return self.load_from_url(
            image_url,
            image_io,
            fetch_timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
        )
189

190
191
    async def fetch_image_async(
        self,
192
193
        image_url: str,
        *,
194
195
196
197
        image_mode: str = "RGB",
    ) -> Image.Image:
        """
        Asynchronously load a PIL image from a HTTP or base64 data URL.
198

199
200
201
        By default, the image is converted into RGB format.
        """
        image_io = ImageMediaIO(image_mode=image_mode)
202

203
204
205
206
207
        return await self.load_from_url_async(
            image_url,
            image_io,
            fetch_timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
        )
208

209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
    def fetch_video(
        self,
        video_url: str,
        *,
        image_mode: str = "RGB",
        num_frames: int = 32,
    ) -> npt.NDArray:
        """
        Load video from a HTTP or base64 data URL.
        """
        image_io = ImageMediaIO(image_mode=image_mode)
        video_io = VideoMediaIO(image_io, num_frames=num_frames)

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

228
229
230
    async def fetch_video_async(
        self,
        video_url: str,
231
        *,
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
        image_mode: str = "RGB",
        num_frames: int = 32,
    ) -> npt.NDArray:
        """
        Asynchronously load video from a HTTP or base64 data URL.

        By default, the image is converted into RGB format.
        """
        image_io = ImageMediaIO(image_mode=image_mode)
        video_io = VideoMediaIO(image_io, num_frames=num_frames)

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


250
251
252
253
254
255
global_media_connector = MediaConnector()
"""The global :class:`MediaConnector` instance used by vLLM."""

fetch_audio = global_media_connector.fetch_audio
fetch_image = global_media_connector.fetch_image
fetch_video = global_media_connector.fetch_video
256
257


258
259
260
261
262
def encode_audio_base64(
    audio: np.ndarray,
    sampling_rate: int,
) -> str:
    """Encode audio as base64."""
263
264
    audio_io = AudioMediaIO()
    return audio_io.encode_base64((audio, sampling_rate))
265
266


267
268
269
270
271
272
273
274
def encode_image_base64(
    image: Image.Image,
    *,
    image_mode: str = "RGB",
    format: str = "JPEG",
) -> str:
    """
    Encode a pillow image to base64 format.
275

276
277
    By default, the image is converted into RGB format before being encoded.
    """
278
279
    image_io = ImageMediaIO(image_mode=image_mode)
    return image_io.encode_base64(image, image_format=format)
280
281


282
def encode_video_base64(frames: npt.NDArray) -> str:
283
284
285
    image_io = ImageMediaIO()
    video_io = VideoMediaIO(image_io)
    return video_io.encode_base64(frames)
286
287


288
289
290
291
292
293
294
295
296
297
# Utilities for input processors
_T = TypeVar("_T", str, int)


def repeat_and_pad_token(
    token: _T,
    *,
    repeat_count: int = 1,
    pad_token_left: Optional[_T] = None,
    pad_token_right: Optional[_T] = None,
298
) -> list[_T]:
299
300
301
302
303
304
305
306
307
308
309
310
    replacement = [token] * repeat_count
    if pad_token_left is not None:
        replacement = [pad_token_left] + replacement
    if pad_token_right is not None:
        replacement = replacement + [pad_token_right]

    return replacement


def repeat_and_pad_placeholder_tokens(
    tokenizer: AnyTokenizer,
    prompt: Optional[str],
311
    prompt_token_ids: list[int],
312
313
    *,
    placeholder_token_id: int,
314
    repeat_count: Union[int, list[int]],
315
316
    pad_token_left: Optional[int] = None,
    pad_token_right: Optional[int] = None,
317
) -> tuple[Optional[str], list[int], list[PlaceholderRange]]:
318
319
320
    if isinstance(repeat_count, int):
        repeat_count = [repeat_count]

321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
    if prompt is None:
        new_prompt = None
    else:
        placeholder_token_str = tokenizer.decode(placeholder_token_id)
        pad_token_str_left = (None if pad_token_left is None else
                              tokenizer.decode(pad_token_left))
        pad_token_str_right = (None if pad_token_right is None else
                               tokenizer.decode(pad_token_right))

        placeholder_token_count = prompt.count(placeholder_token_str)
        # This is an arbitrary number to distinguish between the two cases
        if placeholder_token_count > 16:
            logger.warning(
                "Please follow the prompt format that is "
                "documented on HuggingFace which does not involve "
                "repeating %s tokens.", placeholder_token_str)
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
        if placeholder_token_count < len(repeat_count):
            logger.warning(
                "The number of multi-modal placeholder tokens in the prompt "
                "is less than the number of multi-modal inputs. Extra "
                "placeholder tokens will be treated as plain text")
            repeat_count = repeat_count[:placeholder_token_count]

        prompt_parts = prompt.split(placeholder_token_str,
                                    maxsplit=len(repeat_count))
        new_prompt = ""
        for i, repeat_count_item in enumerate(repeat_count):
            replacement_str = "".join(
                repeat_and_pad_token(
                    placeholder_token_str,
                    repeat_count=repeat_count_item,
                    pad_token_left=pad_token_str_left,
                    pad_token_right=pad_token_str_right,
                ))
            # The image tokens are removed to be consistent with HuggingFace
            new_prompt += prompt_parts[i] + replacement_str
        new_prompt += prompt_parts[-1]
358

359
360
    new_token_ids = list[int]()
    placeholder_ranges = list[PlaceholderRange]()
361
    placeholder_token_idx = 0
362
363
    for i, token in enumerate(prompt_token_ids):
        if token == placeholder_token_id:
364
            curr_repeat_count = repeat_count[placeholder_token_idx]
365
366
            replacement_ids = repeat_and_pad_token(
                placeholder_token_id,
367
                repeat_count=curr_repeat_count,
368
369
370
                pad_token_left=pad_token_left,
                pad_token_right=pad_token_right,
            )
371
372
373
            offset = len(new_token_ids)
            if pad_token_left is not None:
                offset += 1
374
            placeholder_ranges.append({
375
376
                "offset": offset,
                "length": curr_repeat_count,
377
            })
378
            new_token_ids.extend(replacement_ids)
379
            placeholder_token_idx += 1
380

381
382
383
384
            # No need to further scan the list since we replaced all tokens
            if placeholder_token_idx >= len(repeat_count):
                new_token_ids.extend(prompt_token_ids[i + 1:])
                break
385
386
387
        else:
            new_token_ids.append(token)

388
389
390
    return new_prompt, new_token_ids, placeholder_ranges


391
392
393
def consecutive_placeholder_ranges(
        num_items: int,
        item_size: int,
394
        initial_offset: int = 0) -> list[PlaceholderRange]:
395
396
397
    """Returns a list of consecutive PlaceholderRanges of a fixed size"""

    return [
398
399
        PlaceholderRange(offset=initial_offset + i * item_size,
                         length=item_size) for i in range(num_items)
400
    ]
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
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
475
476
477
478
479
480


def merge_and_sort_multimodal_metadata(
    mm_positions: "MultiModalPlaceholderDict",
    mm_hashes: Optional["MultiModalHashDict"],
) -> tuple[list[str], list[PlaceholderRange], Optional[list[str]]]:
    """Given a MultiModalPlaceholderDict, merge all PlaceholderRange
    objects from all available modalities into a single list of 
    PlaceholderRange, sorted by their offset (starting index in the input 
    sequence) in the ascending order.

    Optionally if a MultiModalHashDict is given, same operation will be 
    applied to the object and the sorted list of hashes will be returned.

    Raises:
        ValueError: If the input prompt has interleaved placeholders from
            different modalities (e.g, "<image><audio><image> Describe the 
            content.")
    
    Returns:
        list[str]: Sorted list of involved modalities.
        list[PlaceholderRange]: Sorted list of all PlaceholdeRanges from 
            mm_positions.
        Optional[list[str]]: Sorted list of all hashes from mm_hashes if 
            given, None otherwise.
    """

    modalities = list(mm_positions.keys())

    assert len(modalities) > 0, "No modalities found in the mm_positions."

    # For single modality, placeholder ranges and hashes are already sorted
    # so we can return the list directly.
    if len(modalities) == 1:
        if mm_hashes is None:
            return modalities, list(mm_positions[modalities[0]]), None
        else:
            return modalities, list(mm_positions[modalities[0]]), list(
                mm_hashes[modalities[0]])

    placeholder_lists_with_modality = [(modality, mm_positions[modality])
                                       for modality in modalities]

    if mm_hashes is None:
        sorted_placeholder_lists = sorted(placeholder_lists_with_modality,
                                          key=lambda x: x[1][0]['offset'])
        sorted_hash_lists = None
    else:
        hashes_lists = [
            mm_hashes[modality] for modality in modalities
            if modality in mm_hashes
        ]
        sorted_pairs = sorted(zip(placeholder_lists_with_modality,
                                  hashes_lists),
                              key=lambda x: x[0][1][0]['offset'])
        sorted_placeholder_tuple, sorted_hash_tuple = zip(*sorted_pairs)
        sorted_placeholder_lists = list(sorted_placeholder_tuple)
        sorted_hash_lists = list(sorted_hash_tuple)

    sorted_modalities = [modality for modality, _ in sorted_placeholder_lists]

    # Flatten sorted list of lists to a single list and verify there is no
    # interleaving of placeholders from different modalities.
    merged_placeholders: list[PlaceholderRange] = []
    for modality, placeholder_list in sorted_placeholder_lists:
        if merged_placeholders and placeholder_list[0][
                'offset'] < merged_placeholders[-1]['offset']:
            raise ValueError(
                "Interleaved mixed-modality inference is currently not "
                "supported.")
        merged_placeholders.extend(placeholder_list)

    if sorted_hash_lists is not None:
        merged_hashes = []
        for hash_list in sorted_hash_lists:
            merged_hashes.extend(hash_list)
    else:
        merged_hashes = None

    return sorted_modalities, merged_placeholders, merged_hashes
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511


def group_mm_inputs_by_modality(
        mm_inputs: list["MultiModalKwargs"]) -> list[list["MultiModalKwargs"]]:
    """Group consecutive MultiModalKwargs from mm_inputs with the same modality 
    together into the same list for batching purpose. For MultiModalKwargs with 
    multiple modalities, put them into their own list.

    Args:
        mm_inputs: List of MultiModalKwargs.

    Returns:
        list[list[MultiModalKwargs]]: List of list of MultiModalKwargs, each 
        inner list contains consecutive MultiModalKwargs with same modality, or
        one with multimodal modalities.
    """
    if not mm_inputs:
        return []

    def modality_group_func(mm_input: "MultiModalKwargs") -> Union[str, int]:
        # If the input has multiple modalities, return a id as the unique key
        # for the mm_input input.
        if len(mm_input.modalities) > 1:
            return id(mm_input)

        # Otherwise return the modality string
        return list(mm_input.modalities)[0]

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