feat: support image_embeds in OpenAI API

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

Author: ixlmar

docs/source/commands/trtllm-serve/trtllm-serve.rst

@@ -152,6 +152,24 @@ TRT-LLM multimodal supports the following modalities and data types (depending o
    `load_base64_image utility <https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/utils/load_base64_image.py>`__
    for implementation details.
 
+**Image embeddings**
+
+It is also possible to directly provide the image embeddings to use by the multimodal
+model.
+
+* Using "image_embeds" with base64-encoded data:
+
+  .. code-block:: json
+
+     {"role": "user", "content": [
+         {"type": "text", "text": "What's in this image?"},
+         {"type": "image_embeds", "image_embeds": "{image_embeddings_base64}"}}
+     ]}
+
+.. note::
+   The contents of `image_embeddings_base64` can be generated by base64-encoding
+   the result of serializing a tensor via `torch.save`.
+
 **Video**
 
 * Using "video_url":

tensorrt_llm/inputs/__init__.py

@@ -16,7 +16,8 @@
                     async_load_audio, async_load_image, async_load_video,
                     convert_image_mode, default_multimodal_input_loader,
                     encode_base64_content_from_url, encode_base64_image,
-                    get_cache_salt_id, load_image, load_video)
+                    get_cache_salt_id, load_base64_image_embeds, load_image,
+                    load_video)
 
 __all__ = [
     "ALL_SUPPORTED_MULTIMODAL_MODELS",
@@ -57,4 +58,5 @@
     "get_cache_salt_id",
     "compute_retained_tokens_count",
     "compute_retention_mask",
+    "load_base64_image_embeds",
 ]

tensorrt_llm/inputs/utils.py

@@ -113,6 +113,15 @@ def load_base64_image(parsed_url: str) -> Image.Image:
     return image
 
 
+def load_base64_image_embeds(str_content: str) -> torch.Tensor:
+    content_bytes = base64.b64decode(str_content)
+    with BytesIO(content_bytes) as buf:
+        image_data: torch.Tensor = torch.load(buf,
+                                              weights_only=True,
+                                              map_location="cpu")
+    return image_data
+
+
 def load_image(image: Union[str, Image.Image],
                format: str = "pt",
                device: str = "cpu") -> Union[Image.Image, torch.Tensor]:
@@ -465,10 +474,15 @@ def retrieve_all_sync(self) -> Optional[Dict[str, List[Any]]]:
 
         return {modality: items for modality, items in self._data.items()}
 
-    def add_data(self, media_type: str, data: Union[Coroutine, Any]):
+    def add_data(self,
+                 media_type: str,
+                 data: Union[Coroutine, Any],
+                 *,
+                 modality: Optional[str] = None):
+        modality = modality or media_type
         current_count = len(self._data[media_type]) + 1
         placeholder = retrieve_multimodal_placeholder(self._model_type,
-                                                      media_type, current_count)
+                                                      modality, current_count)
         self._data[media_type].append(data)
         if placeholder:
             self._placeholder_counts[placeholder] += 1

tensorrt_llm/serve/chat_utils.py

@@ -17,7 +17,8 @@
 from tensorrt_llm.inputs import (ConversationMessage, MultimodalData,
                                  MultimodalDataTracker,
                                  add_multimodal_placeholders, async_load_audio,
-                                 async_load_image, async_load_video)
+                                 async_load_image, async_load_video,
+                                 load_base64_image_embeds)
 from tensorrt_llm.inputs.multimodal import MultimodalServerConfig
 from tensorrt_llm.logger import logger
 
@@ -33,24 +34,38 @@ class ChatCompletionContentPartVideoParam(TypedDict, total=False):
     type: Required[Literal["video_url"]]
 
 
+class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):
+    """Type definition for image embeddings passed in base64-encoded PyTorch tensor format."""
+    image_embeds: Required[str]
+    type: Required[Literal["image_embeds"]]
+
+
 # Type Aliases and Constants
 ChatCompletionContentPartParam: TypeAlias = Union[
-    OpenAIChatCompletionContentPartParam, ChatCompletionContentPartVideoParam,
-    str]
+    OpenAIChatCompletionContentPartParam,
+    ChatCompletionContentPartVideoParam,
+    ChatCompletionContentPartImageEmbedsParam,
+    str,
+]
 
 # TODO: Add "input_audio" to support byte_encoded audio input.
 VALID_MESSAGE_CONTENT_MM_PART_TYPES = [
-    "text", "image_url", "video_url", "audio_url"
+    "text",
+    "image_url",
+    "video_url",
+    "audio_url",
+    "image_embeds",
 ]
 
 # Parser Functions
 _TextParser = partial(cast, ChatCompletionContentPartTextParam)
 _ImageParser = partial(cast, ChatCompletionContentPartImageParam)
+_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
 _VideoParser = partial(cast, ChatCompletionContentPartVideoParam)
 _AudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
 
 MM_PARSER_MAP: dict[str, Callable[[ChatCompletionContentPartParam], Union[
-    str, dict[str, str]]]] = {
+    str, dict[str, str], None]]] = {
         "text":
         lambda part: _TextParser(part).get("text", None),
         "image_url":
@@ -59,12 +74,20 @@ class ChatCompletionContentPartVideoParam(TypedDict, total=False):
         lambda part: _VideoParser(part).get("video_url", {}).get("url", None),
         "audio_url":
         lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),
+        "image_embeds":
+        lambda part: _ImageEmbedsParser(part).get("image_embeds", None),
     }
 
+# Map from content part tags used to directly provide embeddings
+# to the corresponding data modality.
+MM_EMBEDDING_MAP: dict[str, str] = {
+    "image_embeds": "image",
+}
+
 
 def _parse_chat_message_content_mm_part(
     part: ChatCompletionContentPartParam
-) -> tuple[str, Union[str, dict[str, str]]]:
+) -> tuple[str, Union[str, dict[str, str], None]]:
     """Parse a single multimodal part of a chat message."""
     assert isinstance(part, dict)
     part_type = part.get("type", None)
@@ -78,7 +101,7 @@ def _parse_chat_message_content_mm_part(
 
 
 def parse_chat_message_content_part(
-    part: ChatCompletionMessageParam,
+    part: ChatCompletionContentPartParam,
     mm_data_tracker: MultimodalDataTracker,
 ) -> Optional[Any]:
     """Parse a single part of a chat message."""
@@ -112,6 +135,19 @@ async def load_image_async():
 
         return MultimodalData(modality="image", data=load_image_async())
 
+    if part_type == "image_embeds":
+        str_content = cast(str, content)
+
+        async def decode_image_embeds_async():
+            try:
+                return load_base64_image_embeds(str_content)
+            except Exception as e:
+                logger.error(f"Failed to decode image data: {str(e)}")
+                return None
+
+        return MultimodalData(modality="image_embeds",
+                              data=decode_image_embeds_async())
+
     if part_type == "video_url":
         str_content = cast(str, content)
 
@@ -147,7 +183,7 @@ async def load_audio_async():
 
 def parse_chat_message_content_parts(
     role: str,
-    parts: Iterable[ChatCompletionMessageParam],
+    parts: Iterable[ChatCompletionContentPartParam],
     mm_data_tracker: MultimodalDataTracker,
 ) -> ConversationMessage:
     """Parse multiple parts of a chat message."""
@@ -237,7 +273,10 @@ def parse_chat_messages_coroutines(
         conversation.append(parsed_msg)
         if parsed_msg["media"]:
             for mdata in parsed_msg["media"]:
-                mm_data_tracker.add_data(mdata["modality"], mdata["data"])
+                mm_data_tracker.add_data(mdata["modality"],
+                                         mdata["data"],
+                                         modality=MM_EMBEDDING_MAP.get(
+                                             mdata["modality"], None))
         mm_placeholder_count = mm_data_tracker.placeholder_counts()
         if mm_placeholder_count:
             parsed_msg["content"] = add_multimodal_placeholders(

tensorrt_llm/serve/openai_server.py

@@ -34,7 +34,7 @@
 from tensorrt_llm.llmapi.llm import RequestOutput
 from tensorrt_llm.logger import logger
 from tensorrt_llm.metrics.collector import MetricsCollector
-from tensorrt_llm.serve.chat_utils import (load_chat_template,
+from tensorrt_llm.serve.chat_utils import (MM_EMBEDDING_MAP, load_chat_template,
                                            parse_chat_messages_coroutines)
 from tensorrt_llm.serve.cluster_storage import create_cluster_storage_client
 from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterWorker
@@ -526,7 +526,18 @@ async def create_chat_response(
 
             mm_data = await mm_coroutines
             if mm_data is not None:
-                prompt["multi_modal_data"] = mm_data
+                # single out directly provided embeddings
+                mm_embeds = {}
+                for tag in list(mm_data.keys()):
+                    if (modality := MM_EMBEDDING_MAP.get(tag, None)) is not None:
+                        mm_embeds[modality] = mm_data.pop(tag)
+
+                if mm_data:
+                    prompt["multi_modal_data"] = mm_data
+                if mm_embeds:
+                    prompt["multi_modal_embeddings"] = mm_embeds
+                if mm_data and mm_embeds:
+                    raise ValueError("Passing 'multi_modal_data' and 'multi_modal_embeddings' at the same time is not supported.")
 
             postproc_args.reasoning_parser = self.llm.args.reasoning_parser
             postproc_args.tool_parser = self.tool_parser

tests/unittest/llmapi/apps/_attach_multimodal_embeddings_patch/__init__.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# used by tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
+
+import pickle
+import tempfile
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from tensorrt_llm._torch.models.modeling_qwen2vl import Qwen2VLInputProcessorBase
+from tensorrt_llm.inputs import ExtraProcessedInputs, TextPrompt
+from tensorrt_llm.sampling_params import SamplingParams
+
+
+# signature taken from tensorrt_llm/inputs/registry.py
+def _attach_multimodal_embeddings(
+    self,
+    inputs: TextPrompt,
+    multimodal_embedding: dict[str, list[torch.Tensor]],
+    sampling_params: SamplingParams,
+) -> tuple[list[int], Optional[ExtraProcessedInputs]]:
+    tempdir = tempfile.gettempdir()
+    file_path = Path(tempdir) / "forwarded_embeddings.pickle"
+    with open(file_path, "wb") as f:
+        pickle.dump(multimodal_embedding, f)
+    raise ValueError(file_path)
+
+
+setattr(Qwen2VLInputProcessorBase, "attach_multimodal_embeddings", _attach_multimodal_embeddings)