vLLM/vllm/multimodal/utils.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import bisect
import mimetypes
from collections import defaultdict
from collections.abc import Generator, Sequence
from dataclasses import replace
from itertools import groupby
from typing import TYPE_CHECKING, Any

import numpy as np
import numpy.typing as npt
from PIL import Image

from vllm.inputs import MultiModalPlaceholders
from vllm.utils.import_utils import LazyLoader

from .inputs import (
    BatchedTensorInputs,
    MultiModalFeatureSpec,
    MultiModalFieldElem,
    MultiModalKwargsItem,
    MultiModalSharedField,
    nested_tensors_equal,
)
from .media import (
    AudioMediaIO,
    ImageMediaIO,
    MediaConnector,
    MediaWithBytes,
    VideoMediaIO,
)

if TYPE_CHECKING:
    import torch.types
else:
    torch = LazyLoader("torch", globals(), "torch")


def encode_audio_base64(
    audio: np.ndarray,
    sampling_rate: int,
    *,
    format: str = "WAV",
) -> str:
    """Encode audio as base64."""
    audio_io = AudioMediaIO()
    return audio_io.encode_base64((audio, sampling_rate), audio_format=format)


def encode_audio_url(
    audio: np.ndarray,
    sampling_rate: int,
    *,
    format: str = "WAV",
) -> str:
    """Encode audio as a data URL."""
    audio_b64 = encode_audio_base64(audio, sampling_rate, format=format)
    mimetype = mimetypes.types_map.get("." + format.lower(), f"audio/{format.lower()}")
    return f"data:{mimetype};base64,{audio_b64}"


def encode_image_base64(
    image: Image.Image,
    *,
    image_mode: str | None = "RGB",
    format: str = "PNG",
) -> str:
    """
    Encode a pillow image to base64 format.

    By default, the image is converted into RGB format before being encoded.
    Pass `image_mode=None` to keep the original image mode.
    """
    image_io = ImageMediaIO(image_mode=image_mode)
    return image_io.encode_base64(image, image_format=format)


def encode_image_url(
    image: Image.Image,
    *,
    image_mode: str | None = "RGB",
    format: str = "PNG",
) -> str:
    """
    Encode a pillow image as a data URL.

    By default, the image is converted into RGB format before being encoded.
    Pass `image_mode=None` to keep the original image mode.
    """
    image_b64 = encode_image_base64(image, image_mode=image_mode, format=format)
    mimetype = mimetypes.types_map.get("." + format.lower(), f"image/{format.lower()}")
    return f"data:{mimetype};base64,{image_b64}"


def encode_video_base64(
    frames: npt.NDArray,
    *,
    format: str = "JPEG",
) -> str:
    image_io = ImageMediaIO()
    video_io = VideoMediaIO(image_io)
    return video_io.encode_base64(frames, video_format=format)


def encode_video_url(
    frames: npt.NDArray,
    *,
    format: str = "JPEG",
) -> str:
    video_b64 = encode_video_base64(frames, format=format)

    if format.lower() == "jpeg":
        mimetype = "video/jpeg"
    else:
        mimetype = mimetypes.types_map.get(
            "." + format.lower(), f"video/{format.lower()}"
        )

    return f"data:{mimetype};base64,{video_b64}"


def get_mm_features_in_window(
    mm_features: list[MultiModalFeatureSpec],
    start: int,
    end: int,
) -> tuple[int, int]:
    """Return (lo, hi) indices for features overlapping [start, end).

    Assumes mm_features are sorted by offset and non-overlapping, so
    offset + length is also sorted.
    """
    lo = bisect.bisect_left(
        mm_features,
        start + 1,
        key=lambda f: f.mm_position.offset + f.mm_position.length,
    )
    hi = bisect.bisect_left(
        mm_features,
        end,
        key=lambda f: f.mm_position.offset,
    )
    return lo, hi


def argsort_mm_positions(
    mm_positions: MultiModalPlaceholders,
) -> list[tuple[str, int]]:
    """
    Given a `MultiModalPlaceholders`, output a sequence of keys to
    sort the dictionary by `offset` (starting index in the input sequence)
    in ascending order.

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

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

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


def _can_batch_mm_items(
    left: MultiModalKwargsItem,
    right: MultiModalKwargsItem,
) -> bool:
    if left.keys() != right.keys():
        return False

    for key, left_elem in left.items():
        right_elem = right[key]
        left_field, right_field = left_elem.field, right_elem.field
        is_shared_field = isinstance(left_field, MultiModalSharedField) and isinstance(
            right_field, MultiModalSharedField
        )
        if (type(left_field) is not type(right_field)) or (
            is_shared_field
            and not nested_tensors_equal(
                left_elem.data, right_elem.data, check_dtype=True
            )
        ):
            return False

    return True


def _batch_mm_items(
    items: Sequence[MultiModalKwargsItem],
    *,
    device: torch.types.Device = None,
    pin_memory: bool = False,
):
    elems = defaultdict[str, list[MultiModalFieldElem]](list)
    for item in items:
        for key, elem in item.items():
            elems[key].append(elem)

    return {
        key: elems[0].field.reduce_data(
            elems,
            device=device,
            pin_memory=pin_memory,
        )
        for key, elems in elems.items()
    }


def strip_covered_mm_data(
    mm_features: list[MultiModalFeatureSpec],
    num_computed_tokens: int,
    uses_mrope: bool = False,
    uses_xdrope: bool = False,
) -> list[MultiModalFeatureSpec]:
    """Drop the tensor data of mm items whose placeholder span is fully inside
    a prefix-cache-covered region: no encoder run can be scheduled for them,
    so the workers never consume the payload fields. M-RoPE and XD-RoPE models
    keep CPU-side metadata fields used to compute positions. SHM address items
    are also kept so workers can balance the sender's reference count. The
    scheduler-side ``Request`` keeps the full features."""
    if not mm_features or num_computed_tokens == 0:
        return mm_features

    def maybe_strip(f: MultiModalFeatureSpec) -> MultiModalFeatureSpec:
        if (
            f.data is None
            # SHM address items must reach the worker so it can acknowledge
            # the sender's reference count before the item is evicted. This
            # mirrors how the SHM receiver identifies them ("address" in item).
            or "address" in f.data
            or (f.mm_position.offset + f.mm_position.length > num_computed_tokens)
        ):
            return f

        data = None
        if uses_mrope or uses_xdrope:
            data = MultiModalKwargsItem(
                {k: elem for k, elem in f.data.items() if elem.field.keep_on_cpu}
            )
        return replace(f, data=data)

    return [maybe_strip(f) for f in mm_features]


def group_and_batch_mm_items(
    items: Sequence[MultiModalKwargsItem],
    *,
    device: torch.types.Device = None,
    pin_memory: bool = False,
) -> Generator[tuple[int, BatchedTensorInputs]]:
    """
    Group consecutive items (possibly from different requests) into batches.

    Items must be split across groups if any of the following occurs,
    as the batch would otherwise be invalid:
    - They have different fields (e.g. mixed image and embedding inputs).
    - They have different values in `MultiModalSharedField`.

    Args:
        items: List of `MultiModalKwargsItem`.
        device: The device to place the grouped tensors on.
        pin_memory: Whether to pin memory for faster host-to-device transfer.

    Yields:
        A tuple `(num_items, grouped_kwargs)`, where:
        - `kwargs` is a dictionary of keyword arguments to pass to the model;
        - `num_items` is the corresponding number of items.
    """
    start_idx = 0
    for end_idx in range(1, len(items) + 1):
        if end_idx < len(items) and _can_batch_mm_items(
            items[end_idx - 1], items[end_idx]
        ):
            continue

        group_data = _batch_mm_items(
            items[start_idx:end_idx],
            device=device,
            pin_memory=pin_memory,
        )

        yield end_idx - start_idx, group_data
        start_idx = end_idx

    assert start_idx == len(items)


def group_and_batch_mm_kwargs(
    mm_kwargs: list[tuple[str, MultiModalKwargsItem]],
    *,
    device: torch.types.Device = None,
    pin_memory: bool = False,
) -> Generator[tuple[str, int, BatchedTensorInputs], None, None]:
    """
    Group consecutive items (possibly from different requests) into batches.

    Items must be split across groups if any of the following occurs,
    as the batch would otherwise be invalid:
    - They have different fields (e.g. mixed image and embedding inputs).
    - They have different values in `MultiModalSharedField`.

    To simplify the implementation of `embed_multimodal`, we add another
    restriction that the items in a batch must belong to the same modality.

    Args:
        mm_kwargs: List of `(modality, item)`.
        device: The device to place the grouped tensors on.
        pin_memory: Whether to pin memory for faster host-to-device transfer.

    Yields:
        A tuple `(modality, num_items, grouped_kwargs)`, where:
        - `modality` is the modality of the batch;
        - `kwargs` is a dictionary of keyword arguments to pass to the model;
        - `num_items` is the corresponding number of items.
    """
    for modality, group in groupby(mm_kwargs, key=lambda x: x[0]):
        items_lst = [item for _, item in group]

        for num_items, mm_kwargs_batch in group_and_batch_mm_items(
            items_lst,
            device=device,
            pin_memory=pin_memory,
        ):
            yield modality, num_items, mm_kwargs_batch


def fetch_audio(
    audio_url: str,
    audio_io_kwargs: dict[str, Any] | None = None,
) -> tuple[np.ndarray, int | float]:
    """
    Args:
        audio_url: URL of the audio file to fetch.
        audio_io_kwargs: Additional kwargs passed to handle audio IO.

    Warning:
        This method has direct access to local files and is only intended
        to be called by user code. Never call this from the online server!
    """
    media_io_kwargs = None if not audio_io_kwargs else {"audio": audio_io_kwargs}
    media_connector = MediaConnector(
        media_io_kwargs=media_io_kwargs,
        allowed_local_media_path="/",
    )
    return media_connector.fetch_audio(audio_url)


def fetch_image(
    image_url: str,
    image_io_kwargs: dict[str, Any] | None = None,
) -> Image.Image:
    """
    Args:
        image_url: URL of the image file to fetch.
        image_io_kwargs: Additional kwargs passed to handle image IO.

    Warning:
        This method has direct access to local files and is only intended
        to be called by user code. Never call this from the online server!
    """
    media_io_kwargs = None if not image_io_kwargs else {"image": image_io_kwargs}
    media_connector = MediaConnector(
        media_io_kwargs=media_io_kwargs,
        allowed_local_media_path="/",
    )
    return media_connector.fetch_image(image_url)


def fetch_video(
    video_url: str,
    video_io_kwargs: dict[str, Any] | None = None,
) -> MediaWithBytes[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.

    Warning:
        This method has direct access to local files and is only intended
        to be called by user code. Never call this from the online server!
    """
    media_io_kwargs = None if not video_io_kwargs else {"video": video_io_kwargs}
    media_connector = MediaConnector(
        media_io_kwargs=media_io_kwargs,
        allowed_local_media_path="/",
    )
    return media_connector.fetch_video(video_url)


def set_mm_embedding_modality(embed: "torch.Tensor", modality: str) -> "torch.Tensor":
    """Attach modality metadata to a gathered multimodal embedding tensor.

    Used by interleaved Omni merge paths that need to group embeddings by
    modality without threading a parallel modalities list through
    ``embed_input_ids``.
    """
    embed.modality = modality  # type: ignore[attr-defined]
    return embed


def copy_mm_embedding_modality(
    src: "torch.Tensor", dst: "torch.Tensor"
) -> "torch.Tensor":
    """Copy ``modality`` from ``src`` onto ``dst`` if present."""
    modality = getattr(src, "modality", None)
    if modality is not None:
        dst.modality = modality  # type: ignore[attr-defined]
    return dst


def get_mm_embedding_modalities(
    multimodal_embeddings: Sequence["torch.Tensor"],
) -> list[str]:
    """Collect per-embedding modalities previously set on the tensors."""
    modalities: list[str] = []
    for i, emb in enumerate(multimodal_embeddings):
        modality = getattr(emb, "modality", None)
        if modality is None:
            raise ValueError(
                f"Missing modality on multimodal embedding at index {i}. "
                "Encoder gather must set embed.modality before interleaved "
                "audio-in-video merge."
            )
        modalities.append(modality)
    return modalities