utils.py 16.8 KB
Newer Older
1
import base64
2
import os
3
from functools import lru_cache
4
from io import BytesIO
5
from typing import Any, List, Optional, Tuple, TypeVar, Union
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 global_http_connection
13
14
15
from vllm.logger import init_logger
from vllm.transformers_utils.tokenizer import AnyTokenizer, get_tokenizer

16
17
from .inputs import MultiModalDataDict, PlaceholderRange

18
19
20
logger = init_logger(__name__)

cached_get_tokenizer = lru_cache(get_tokenizer)
21
22


23
def _load_image_from_bytes(b: bytes) -> Image.Image:
24
25
26
27
28
    image = Image.open(BytesIO(b))
    image.load()
    return image


29
30
31
32
33
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
def _is_subpath(image_path: str, allowed_local_media_path: str) -> bool:
    # Get the common path
    common_path = os.path.commonpath([
        os.path.abspath(image_path),
        os.path.abspath(allowed_local_media_path)
    ])
    # Check if the common path is the same as allowed_local_media_path
    return common_path == os.path.abspath(allowed_local_media_path)


def _load_image_from_file(image_url: str,
                          allowed_local_media_path: str) -> Image.Image:
    if not allowed_local_media_path:
        raise ValueError("Invalid 'image_url': Cannot load local files without"
                         "'--allowed-local-media-path'.")
    if allowed_local_media_path:
        if not os.path.exists(allowed_local_media_path):
            raise ValueError(
                "Invalid '--allowed-local-media-path': "
                f"The path {allowed_local_media_path} does not exist.")
        if not os.path.isdir(allowed_local_media_path):
            raise ValueError(
                "Invalid '--allowed-local-media-path': "
                f"The path {allowed_local_media_path} must be a directory.")

    # Only split once and assume the second part is the image path
    _, image_path = image_url.split("file://", 1)
    if not _is_subpath(image_path, allowed_local_media_path):
        raise ValueError(
            f"Invalid 'image_url': The file path {image_path} must"
            " be a subpath of '--allowed-local-media-path'"
            f" '{allowed_local_media_path}'.")

    image = Image.open(image_path)
    image.load()
    return image


def _load_image_from_data_url(image_url: str) -> Image.Image:
68
69
70
71
72
    # Only split once and assume the second part is the base64 encoded image
    _, image_base64 = image_url.split(",", 1)
    return load_image_from_base64(image_base64)


73
74
75
76
def fetch_image(image_url: str,
                *,
                image_mode: str = "RGB",
                allowed_local_media_path: str = "") -> Image.Image:
77
78
79
80
81
    """
    Load a PIL image from a HTTP or base64 data URL.

    By default, the image is converted into RGB format.
    """
82
    if image_url.startswith('http'):
83
        image_raw = global_http_connection.get_bytes(
84
85
86
            image_url,
            timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
        )
87
88
89
90
        image = _load_image_from_bytes(image_raw)

    elif image_url.startswith('data:image'):
        image = _load_image_from_data_url(image_url)
91
92
    elif image_url.startswith('file://'):
        image = _load_image_from_file(image_url, allowed_local_media_path)
93
94
    else:
        raise ValueError("Invalid 'image_url': A valid 'image_url' must start "
95
                         "with either 'data:image', 'file://' or 'http'.")
96

97
    return image.convert(image_mode)
98
99


100
101
async def async_fetch_image(image_url: str,
                            *,
102
103
                            image_mode: str = "RGB",
                            allowed_local_media_path: str = "") -> Image.Image:
104
105
    """
    Asynchronously load a PIL image from a HTTP or base64 data URL.
106

107
108
109
110
    By default, the image is converted into RGB format.
    """
    if image_url.startswith('http'):
        image_raw = await global_http_connection.async_get_bytes(
111
112
113
            image_url,
            timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
        )
114
        image = _load_image_from_bytes(image_raw)
115

116
117
    elif image_url.startswith('data:image'):
        image = _load_image_from_data_url(image_url)
118
119
    elif image_url.startswith('file://'):
        image = _load_image_from_file(image_url, allowed_local_media_path)
120
121
    else:
        raise ValueError("Invalid 'image_url': A valid 'image_url' must start "
122
                         "with either 'data:image', 'file://' or 'http'.")
123

124
    return image.convert(image_mode)
125
126


127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def _load_video_frames_from_bytes(b: bytes):
    frame = Image.open(BytesIO(b))
    return np.array(frame)


def load_video_frames_from_base64(frame: Union[bytes, str]):
    """Load frame from base64 format."""
    return _load_video_frames_from_bytes(base64.b64decode(frame))


def _load_video_from_bytes(b: bytes, num_frames: int = 32):
    _, decord = try_import_video_packages()

    video_path = BytesIO(b)
    vr = decord.VideoReader(video_path, num_threads=1)
    total_frame_num = len(vr)

    if total_frame_num > num_frames:
        uniform_sampled_frames = np.linspace(0,
                                             total_frame_num - 1,
                                             num_frames,
                                             dtype=int)
        frame_idx = uniform_sampled_frames.tolist()
    else:
        frame_idx = [i for i in range(0, total_frame_num)]
    frames = vr.get_batch(frame_idx).asnumpy()

    return frames


def _load_video_from_data_url(video_url: str):
    # Only split once and assume the second part is the base64 encoded image
    frames_base64 = video_url.split(",")[1:]
    return np.stack([
        load_video_frames_from_base64(frame_base64)
        for frame_base64 in frames_base64
    ])


def fetch_video(video_url: str, *, num_frames: int = 32) -> npt.NDArray:
    """
    Load video from a HTTP or base64 data URL.
    """
    if video_url.startswith('http') or video_url.startswith('https'):
        video_raw = global_http_connection.get_bytes(
            video_url,
            timeout=envs.VLLM_VIDEO_FETCH_TIMEOUT,
        )
        video = _load_video_from_bytes(video_raw, num_frames)
    elif video_url.startswith('data:video'):
        video = _load_video_from_data_url(video_url)
    else:
        raise ValueError("Invalid 'video_url': A valid 'video_url' must start "
                         "with either 'data:video' or 'http'.")
    return video


async def async_fetch_video(video_url: str,
                            *,
                            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.
    """
    if video_url.startswith('http') or video_url.startswith('https'):
        video_raw = await global_http_connection.async_get_bytes(
            video_url,
            timeout=envs.VLLM_VIDEO_FETCH_TIMEOUT,
        )
        video = _load_video_from_bytes(video_raw, num_frames)
    elif video_url.startswith('data:video'):
        video = _load_video_from_data_url(video_url)
    else:
        raise ValueError("Invalid 'video_url': A valid 'video_url' must start "
                         "with either 'data:video' or 'http'.")
    return video


206
207
208
209
def try_import_audio_packages() -> Tuple[Any, Any]:
    try:
        import librosa
        import soundfile
210
    except ImportError as exc:
211
        raise ImportError(
212
            "Please install vllm[audio] for audio support.") from exc
213
214
215
    return librosa, soundfile


216
217
218
219
def fetch_audio(audio_url: str) -> Tuple[np.ndarray, Union[int, float]]:
    """
    Load audio from a URL.
    """
220
221
    librosa, _ = try_import_audio_packages()

222
223
    if audio_url.startswith("http"):
        audio_bytes = global_http_connection.get_bytes(
224
225
226
            audio_url,
            timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
        )
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    elif audio_url.startswith("data:audio"):
        _, audio_base64 = audio_url.split(",", 1)
        audio_bytes = base64.b64decode(audio_base64)
    else:
        raise ValueError("Invalid 'audio_url': A valid 'audio_url' must start "
                         "with either 'data:audio' or 'http'.")

    return librosa.load(BytesIO(audio_bytes), sr=None)


async def async_fetch_audio(
        audio_url: str) -> Tuple[np.ndarray, Union[int, float]]:
    """
    Asynchronously fetch audio from a URL.
    """
242
243
    librosa, _ = try_import_audio_packages()

244
245
    if audio_url.startswith("http"):
        audio_bytes = await global_http_connection.async_get_bytes(
246
247
248
            audio_url,
            timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
        )
249
250
251
252
253
254
255
256
257
258
    elif audio_url.startswith("data:audio"):
        _, audio_base64 = audio_url.split(",", 1)
        audio_bytes = base64.b64decode(audio_base64)
    else:
        raise ValueError("Invalid 'audio_url': A valid 'audio_url' must start "
                         "with either 'data:audio' or 'http'.")

    return librosa.load(BytesIO(audio_bytes), sr=None)


259
260
261
262
263
def get_and_parse_audio(audio_url: str) -> MultiModalDataDict:
    audio, sr = fetch_audio(audio_url)
    return {"audio": (audio, sr)}


264
265
266
267
268
269
def get_and_parse_image(
        image_url: str,
        *,
        allowed_local_media_path: str = "") -> MultiModalDataDict:
    image = fetch_image(image_url,
                        allowed_local_media_path=allowed_local_media_path)
270
271
272
    return {"image": image}


273
274
275
276
277
def get_and_parse_video(video_url: str) -> MultiModalDataDict:
    video = fetch_video(video_url)
    return {"video": video}


278
279
280
281
282
async def async_get_and_parse_audio(audio_url: str) -> MultiModalDataDict:
    audio, sr = await async_fetch_audio(audio_url)
    return {"audio": (audio, sr)}


283
284
285
286
287
288
async def async_get_and_parse_image(
        image_url: str,
        *,
        allowed_local_media_path: str = "") -> MultiModalDataDict:
    image = await async_fetch_image(
        image_url, allowed_local_media_path=allowed_local_media_path)
289
290
291
    return {"image": image}


292
293
294
295
296
async def async_get_and_parse_video(video_url: str) -> MultiModalDataDict:
    video = await async_fetch_video(video_url)
    return {"video": video}


297
298
299
300
301
def encode_audio_base64(
    audio: np.ndarray,
    sampling_rate: int,
) -> str:
    """Encode audio as base64."""
302
303
    _, soundfile = try_import_audio_packages()

304
305
306
307
308
309
    buffered = BytesIO()
    soundfile.write(buffered, audio, sampling_rate, format="WAV")

    return base64.b64encode(buffered.getvalue()).decode('utf-8')


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

319
320
    By default, the image is converted into RGB format before being encoded.
    """
321
    buffered = BytesIO()
322
    image = image.convert(image_mode)
323
324
325
326
327
328
    image.save(buffered, format)
    return base64.b64encode(buffered.getvalue()).decode('utf-8')


def load_image_from_base64(image: Union[bytes, str]) -> Image.Image:
    """Load image from base64 format."""
329
    return _load_image_from_bytes(base64.b64decode(image))
330
331


332
333
334
def rescale_image_size(image: Image.Image,
                       size_factor: float,
                       transpose: int = -1) -> Image.Image:
335
336
337
    """Rescale the dimensions of an image by a constant factor."""
    new_width = int(image.width * size_factor)
    new_height = int(image.height * size_factor)
338
339
340
341
    image = image.resize((new_width, new_height))
    if transpose >= 0:
        image = image.transpose(Image.Transpose(transpose))
    return image
342
343


344
345
346
def try_import_video_packages() -> Any:
    try:
        import cv2
347
        import decord
348
    except ImportError as exc:
349
        raise ImportError(
350
            "Please install vllm[video] for video support.") from exc
351
    return cv2, decord
352
353
354


def resize_video(frames: npt.NDArray, size: Tuple[int, int]) -> npt.NDArray:
355
    cv2, _ = try_import_video_packages()
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385

    num_frames, _, _, channels = frames.shape
    new_height, new_width = size
    resized_frames = np.empty((num_frames, new_height, new_width, channels),
                              dtype=frames.dtype)
    for i, frame in enumerate(frames):
        resized_frame = cv2.resize(frame, (new_width, new_height))
        resized_frames[i] = resized_frame
    return resized_frames


def rescale_video_size(frames: npt.NDArray, size_factor: float) -> npt.NDArray:
    _, height, width, _ = frames.shape
    new_height = int(height * size_factor)
    new_width = int(width * size_factor)

    return resize_video(frames, (new_height, new_width))


def sample_frames_from_video(frames: npt.NDArray,
                             num_frames: int) -> npt.NDArray:
    total_frames = frames.shape[0]
    if num_frames == -1:
        return frames
    else:
        frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
        sampled_frames = frames[frame_indices, ...]
        return sampled_frames


386
387
388
389
390
391
392
393
394
def encode_video_base64(frames: npt.NDArray):
    base64_frames = []
    frames_list = [frames[i] for i in range(frames.shape[0])]
    for frame in frames_list:
        img_base64 = encode_image_base64(Image.fromarray(frame))
        base64_frames.append(img_base64)
    return ",".join(base64_frames)


395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
# 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,
) -> List[_T]:
    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],
    prompt_token_ids: List[int],
    *,
    placeholder_token_id: int,
421
    repeat_count: Union[int, List[int]],
422
423
    pad_token_left: Optional[int] = None,
    pad_token_right: Optional[int] = None,
424
) -> Tuple[Optional[str], List[int], List[PlaceholderRange]]:
425
426
427
    if isinstance(repeat_count, int):
        repeat_count = [repeat_count]

428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
    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)
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
        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]
465
466

    new_token_ids: List[int] = []
467
    placeholder_ranges: List[PlaceholderRange] = []
468
    placeholder_token_idx = 0
469
470
471
472
    for i, token in enumerate(prompt_token_ids):
        if token == placeholder_token_id:
            replacement_ids = repeat_and_pad_token(
                placeholder_token_id,
473
                repeat_count=repeat_count[placeholder_token_idx],
474
475
476
                pad_token_left=pad_token_left,
                pad_token_right=pad_token_right,
            )
477
478
479
480
            placeholder_ranges.append({
                "offset": len(new_token_ids),
                "length": len(replacement_ids)
            })
481
            new_token_ids.extend(replacement_ids)
482
            placeholder_token_idx += 1
483

484
485
486
487
            # 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
488
489
490
        else:
            new_token_ids.append(token)

491
492
493
494
495
496
497
498
499
500
501
    return new_prompt, new_token_ids, placeholder_ranges


def consecutive_placeholder_ranges(num_items: int,
                                   item_size: int) -> List[PlaceholderRange]:
    """Returns a list of consecutive PlaceholderRanges of a fixed size"""

    return [
        PlaceholderRange(offset=i * item_size, length=item_size)
        for i in range(num_items)
    ]