[V1] Extend beyond image modality and support mixed-modality inference with Llava-OneVision

vllm/model_executor/models/llava_onevision.py

@@ -38,8 +38,8 @@
 from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
                     maybe_prefix, merge_multimodal_embeddings)
 
-# Result in the max possible feature size (2x2 grid of 336x336px tiles)
-MAX_IMAGE_FEATURE_SIZE_HEIGHT = MAX_IMAGE_FEATURE_SIZE_WIDTH = 448
+# Ref: https://github.com/LLaVA-VL/LLaVA-NeXT/blob/main/docs/LLaVA_OneVision.md?plain=1#L14
+MAX_IMAGE_FEATURE_SIZE_HEIGHT = MAX_IMAGE_FEATURE_SIZE_WIDTH = 2304
 
 # For profile run
 _MAX_FRAMES_PER_VIDEO = 16
@@ -366,9 +366,11 @@ def input_processor_for_llava_onevision(ctx: InputContext,
                                     and "image" not in multi_modal_data):
         return inputs
     if "image" in multi_modal_data:
-        return input_processor_when_multimodal_input_image(ctx, inputs)
+        inputs = input_processor_when_multimodal_input_image(ctx, inputs)
     if "video" in multi_modal_data:
         return input_processor_when_multimodal_input_video(ctx, inputs)
+    else:
+        return inputs
 
     msg = "Unsupported multi data type"
     raise NotImplementedError(msg)
@@ -832,21 +834,18 @@ def get_multimodal_embeddings(
         if not modalities:
             return None
 
-        # We make a tuple of each embedding with its modality string. This is a
-        # temporary workaround for models to handle mixed modalities when
-        # get_multimodal_embeddings and get_input_embeddings are called
-        # separately.
-        # TODO(ywang96): Add support for mixed-modality inference for v1.
-        multimodal_embeddings: List[Tuple[NestedTensors, str]] = []
+        # The result multimoal_embeddings is tuple of tensors, with each
+        # tensor correspoending to a multimodal data item (image or video).
+        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
 
         if "images" in modalities:
             image_input = modalities["images"]
             vision_embeddings = self._process_image_input(image_input)
-            multimodal_embeddings.append((vision_embeddings, "image"))
+            multimodal_embeddings += tuple(vision_embeddings)
         if "videos" in modalities:
             video_input = modalities["videos"]
             video_embeddings = self._process_video_pixels(video_input)
-            multimodal_embeddings.append((video_embeddings, "video"))
+            multimodal_embeddings += tuple(video_embeddings)
 
         return multimodal_embeddings
 
@@ -858,15 +857,9 @@ def get_input_embeddings(
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.get_input_embeddings(input_ids)
         if multimodal_embeddings is not None:
-            for embeddings, modality in multimodal_embeddings:
-                if modality == "image":
-                    inputs_embeds = merge_multimodal_embeddings(
-                        input_ids, inputs_embeds, embeddings,
-                        self.config.image_token_index)
-                if modality == "video":
-                    inputs_embeds = merge_multimodal_embeddings(
-                        input_ids, inputs_embeds, embeddings,
-                        self.config.video_token_index)
+            inputs_embeds = merge_multimodal_embeddings(
+                input_ids, inputs_embeds, multimodal_embeddings,
+                [self.config.image_token_index, self.config.video_token_index])
         return inputs_embeds
 
     def forward(

vllm/multimodal/utils.py

@@ -1,6 +1,6 @@
 from functools import lru_cache
 from pathlib import Path
-from typing import Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Optional, TypeVar, Union
 from urllib.parse import ParseResult, urlparse
 
 import numpy as np
@@ -25,6 +25,9 @@
 
 _M = TypeVar("_M")
 
+if TYPE_CHECKING:
+    from ..multimodal import MultiModalPlaceholderDict
+
 
 class MediaConnector:
 
@@ -437,3 +440,48 @@ def consecutive_placeholder_ranges(
         PlaceholderRange(offset=initial_offset + i * item_size,
                          length=item_size) for i in range(num_items)
     ]
+
+
+def merge_and_sort_placeholders_from_modalities(
+    mm_positions: "MultiModalPlaceholderDict"
+) -> tuple[list[str], list[PlaceholderRange]]:
+    """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.
+
+    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.
+    """
+
+    modalities = list(mm_positions.keys())
+
+    # For single modality, its placeholder ranges are already sorted.
+    if len(modalities) == 1:
+        return modalities, list(mm_positions[modalities[0]])
+
+    placeholder_lists_with_modality = [(modality, mm_positions[modality])
+                                       for modality in modalities
+                                       if modality in mm_positions]
+
+    sorted_lists_with_modality = sorted(placeholder_lists_with_modality,
+                                        key=lambda x: x[1][0]['offset'])
+
+    # Verify if the sorted order avoids interleaving
+    merged: list[PlaceholderRange] = []
+    for modality, placeholder_list in sorted_lists_with_modality:
+        if merged and placeholder_list[0]['offset'] < merged[-1]['offset']:
+            raise ValueError(
+                "Interleaved mixed-modality inference is currently not "
+                "supported.")
+        merged.extend(placeholder_list)
+
+    # Return the order of modalities and the merged placeholder ranges
+    return [modality for modality, _ in sorted_lists_with_modality], merged

vllm/v1/request.py

@@ -4,6 +4,7 @@
 from vllm.inputs import DecoderOnlyInputs, SingletonInputsAdapter, token_inputs
 from vllm.lora.request import LoRARequest
 from vllm.multimodal import MultiModalKwargs
+from vllm.multimodal.utils import merge_and_sort_placeholders_from_modalities
 from vllm.sampling_params import SamplingParams
 from vllm.sequence import RequestMetrics
 from vllm.v1.engine import EngineCoreRequest
@@ -51,15 +52,37 @@ def __init__(
         # Multi-modal input metadata.
         mm_positions = self.inputs.multi_modal_placeholders
         if mm_positions:
-            # FIXME(woosuk): Support other modalities.
-            self.mm_positions = mm_positions.get("image", [])
+            sorted_modalities, sorted_mm_positions = merge_and_sort_placeholders_from_modalities(  # noqa: E501
+                mm_positions)
+            self.mm_positions = sorted_mm_positions
         else:
+            sorted_modalities = []
             self.mm_positions = []
+
         # Output of the mm input mapper (e.g., image tensors).
         self.mm_inputs: List[MultiModalKwargs] = []
         if self.inputs.multi_modal_inputs:
+            # NOTE: We only need to sort multimodal kwargs when there
+            # are multiple modalities involved.
+            if len(sorted_modalities) > 1:
+                modality_order_dict = {
+                    modality: order
+                    for order, modality in enumerate(sorted_modalities)
+                }
+
+                # Sanity check to make sure each multimodal input
+                # has only one modality key.
+                for mm_input in self.inputs.multi_modal_inputs:
+                    assert len(mm_input.modalities) == 1
+
+                # Sort MultiModalKwags to match sorted_mm_positions
+                self.inputs.multi_modal_inputs.sort(
+                    key=lambda mm_input: modality_order_dict[list(
+                        mm_input.modalities)[0]])
+
             self.mm_inputs = self.inputs.multi_modal_inputs
 
+        assert len(self.mm_inputs) == len(self.mm_positions)
         self.mm_hashes: List[str] = self.inputs.multi_modal_hashes
 
         # Cache the computed kv block hashes of the request to avoid