utils.py 18.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
import torch
10
11
from PIL import Image

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

17
18
from .inputs import MultiModalDataDict, PlaceholderRange

19
20
21
logger = init_logger(__name__)

cached_get_tokenizer = lru_cache(get_tokenizer)
22
23


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


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
68
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:
69
70
71
72
73
    # 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)


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

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

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

98
    return image.convert(image_mode)
99
100


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

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

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

125
    return image.convert(image_mode)
126
127


128
def _load_video_from_bytes(b: bytes, num_frames: int = 32) -> npt.NDArray:
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
    _, 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


148
149
150
151
152
153
154
155
156
157
158
def _load_video_from_data_url(video_url: str) -> npt.NDArray:
    # Only split once and assume the second part is the base64 encoded video
    _, video_base64 = video_url.split(",", 1)

    if video_url.startswith("data:video/jpeg;"):
        return np.stack([
            np.array(load_image_from_base64(frame_base64))
            for frame_base64 in video_base64.split(",")
        ])

    return load_video_from_base64(video_base64)
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


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


201
202
203
204
def try_import_audio_packages() -> Tuple[Any, Any]:
    try:
        import librosa
        import soundfile
205
    except ImportError as exc:
206
        raise ImportError(
207
            "Please install vllm[audio] for audio support.") from exc
208
209
210
    return librosa, soundfile


211
212
213
214
def fetch_audio(audio_url: str) -> Tuple[np.ndarray, Union[int, float]]:
    """
    Load audio from a URL.
    """
215
216
    librosa, _ = try_import_audio_packages()

217
218
    if audio_url.startswith("http"):
        audio_bytes = global_http_connection.get_bytes(
219
220
221
            audio_url,
            timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
        )
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
    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.
    """
237
238
    librosa, _ = try_import_audio_packages()

239
240
    if audio_url.startswith("http"):
        audio_bytes = await global_http_connection.async_get_bytes(
241
242
243
            audio_url,
            timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
        )
244
245
246
247
248
249
250
251
252
253
    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)


254
255
256
257
258
def get_and_parse_audio(audio_url: str) -> MultiModalDataDict:
    audio, sr = fetch_audio(audio_url)
    return {"audio": (audio, sr)}


259
260
261
262
263
264
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)
265
266
267
    return {"image": image}


268
269
270
271
272
def get_and_parse_video(video_url: str) -> MultiModalDataDict:
    video = fetch_video(video_url)
    return {"video": video}


273
274
275
276
277
async def async_get_and_parse_audio(audio_url: str) -> MultiModalDataDict:
    audio, sr = await async_fetch_audio(audio_url)
    return {"audio": (audio, sr)}


278
279
280
281
282
283
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)
284
285
286
    return {"image": image}


287
288
289
290
291
async def async_get_and_parse_video(video_url: str) -> MultiModalDataDict:
    video = await async_fetch_video(video_url)
    return {"video": video}


292
293
294
295
296
def encode_audio_base64(
    audio: np.ndarray,
    sampling_rate: int,
) -> str:
    """Encode audio as base64."""
297
298
    _, soundfile = try_import_audio_packages()

299
300
301
302
303
304
    buffered = BytesIO()
    soundfile.write(buffered, audio, sampling_rate, format="WAV")

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


305
306
307
308
309
310
311
312
def encode_image_base64(
    image: Image.Image,
    *,
    image_mode: str = "RGB",
    format: str = "JPEG",
) -> str:
    """
    Encode a pillow image to base64 format.
313

314
315
    By default, the image is converted into RGB format before being encoded.
    """
316
    buffered = BytesIO()
317
    image = image.convert(image_mode)
318
319
320
321
322
323
    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."""
324
    return _load_image_from_bytes(base64.b64decode(image))
325
326


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


339
def try_import_video_packages():
340
341
    try:
        import cv2
342
        import decord
343
    except ImportError as exc:
344
        raise ImportError(
345
            "Please install vllm[video] for video support.") from exc
346
    return cv2, decord
347
348
349


def resize_video(frames: npt.NDArray, size: Tuple[int, int]) -> npt.NDArray:
350
    cv2, _ = try_import_video_packages()
351
352
353
354
355
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

    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


381
def encode_video_base64(frames: npt.NDArray) -> str:
382
383
384
385
386
387
388
389
    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)


390
391
392
393
394
def load_video_from_base64(video: Union[bytes, str]) -> npt.NDArray:
    """Load video from base64 format."""
    return _load_video_from_bytes(base64.b64decode(video))


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
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
def resolve_visual_encoder_outputs(
    encoder_outputs: Union[torch.Tensor, list[torch.Tensor]],
    feature_sample_layers: Optional[list[int]],
    post_layer_norm: Optional[torch.nn.LayerNorm],
    max_possible_layers: int,
) -> torch.Tensor:
    """Given the outputs a visual encoder module that may correspond to the
    output of the last layer, or a list of hidden states to be stacked,
    handle post normalization and resolve it into a single output tensor.

    Args:
        encoder_outputs: Output of encoder's last layer or all hidden states.
        feature_sample_layers: Optional layer indices to grab from the encoder
            outputs; if provided, encoder outputs must be a list.
        post_layer_norm: Post norm to apply to the output of the encoder.
        max_possible_layers: Total layers in the fully loaded visual encoder.

    """
    if feature_sample_layers is None:
        if post_layer_norm is not None:
            return post_layer_norm(encoder_outputs)
        return encoder_outputs

    # Get the hidden states corresponding to the layer indices.
    # Negative values are relative to the full visual encoder,
    # so offset them depending on how many layers were loaded.
    # NOTE: this assumes that encoder_outputs contains a list
    # of hidden states in the same order as the encoder layers
    # that produced them.
    offset = max_possible_layers - len(encoder_outputs)
    hs_pool = [
        encoder_outputs[layer_idx]
        if layer_idx >= 0 else encoder_outputs[layer_idx + offset]
        for layer_idx in feature_sample_layers
    ]

    # Apply post-norm on the final hidden state if we are using it
    uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1)
    if post_layer_norm is not None and uses_last_layer:
        hs_pool[-1] = post_layer_norm(encoder_outputs)
    return torch.cat(hs_pool, dim=-1)


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
# 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,
464
    repeat_count: Union[int, List[int]],
465
466
    pad_token_left: Optional[int] = None,
    pad_token_right: Optional[int] = None,
467
) -> Tuple[Optional[str], List[int], List[PlaceholderRange]]:
468
469
470
    if isinstance(repeat_count, int):
        repeat_count = [repeat_count]

471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
    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)
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
        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]
508
509

    new_token_ids: List[int] = []
510
    placeholder_ranges: List[PlaceholderRange] = []
511
    placeholder_token_idx = 0
512
513
514
515
    for i, token in enumerate(prompt_token_ids):
        if token == placeholder_token_id:
            replacement_ids = repeat_and_pad_token(
                placeholder_token_id,
516
                repeat_count=repeat_count[placeholder_token_idx],
517
518
519
                pad_token_left=pad_token_left,
                pad_token_right=pad_token_right,
            )
520
521
522
523
            placeholder_ranges.append({
                "offset": len(new_token_ids),
                "length": len(replacement_ids)
            })
524
            new_token_ids.extend(replacement_ids)
525
            placeholder_token_idx += 1
526

527
528
529
530
            # 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
531
532
533
        else:
            new_token_ids.append(token)

534
535
536
    return new_prompt, new_token_ids, placeholder_ranges


537
538
539
540
def consecutive_placeholder_ranges(
        num_items: int,
        item_size: int,
        initial_offset: int = 0) -> List[PlaceholderRange]:
541
542
543
    """Returns a list of consecutive PlaceholderRanges of a fixed size"""

    return [
544
545
        PlaceholderRange(offset=initial_offset + i * item_size,
                         length=item_size) for i in range(num_items)
546
    ]