Update vLLM wheel version in Dockerfile and enhance input processing for omni models. Added functionality to reinject omni-specific fields into processed prompts and improved handling of multimodal audio extraction in serving_chat.py. Adjusted SharedMemoryConnector logic for payload size checks and refined input preprocessing methods for better compatibility with multimodal data.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/inputs/preprocess.py

@@ -4,7 +4,7 @@ from typing_extensions import assert_never
 from vllm.inputs.data import EmbedsInputs, SingletonInputs
 from vllm.inputs.preprocess import InputPreprocessor
 from vllm.logger import init_logger
-from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict
+from vllm.multimodal.inputs import MultiModalInputs
 from vllm.renderers.inputs import SingletonDictPrompt
 
 from vllm_omni.inputs.data import (
@@ -30,8 +30,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         parsed_content: OmniTextPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> OmniTokenInputs | MultiModalInputs:
         prompt_text = parsed_content["prompt"]
 
@@ -42,14 +40,8 @@ class OmniInputPreprocessor(InputPreprocessor):
                 multi_modal_data,
                 parsed_content.get("mm_processor_kwargs") or {},
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
-            prompt_embeds = parsed_content.get("prompt_embeds")
-            if prompt_embeds is not None:
-                inputs["prompt_embeds"] = prompt_embeds
-            additional_information = parsed_content.get("additional_information")
-            if additional_information is not None:
-                inputs["additional_information"] = additional_information
+
         else:
             prompt_token_ids = self._tokenize_prompt(
                 prompt_text,
@@ -60,7 +52,12 @@ class OmniInputPreprocessor(InputPreprocessor):
                 prompt_embeds=parsed_content.get("prompt_embeds"),
                 additional_information=parsed_content.get("additional_information"),
             )
-
+        prompt_embeds = parsed_content.get("prompt_embeds")
+        if prompt_embeds is not None:
+            inputs["prompt_embeds"] = prompt_embeds
+        additional_information = parsed_content.get("additional_information")
+        if additional_information is not None:
+            inputs["additional_information"] = additional_information
         if cache_salt := parsed_content.get("cache_salt"):
             inputs["cache_salt"] = cache_salt
 
@@ -70,8 +67,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         parsed_content: OmniTokensPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> OmniTokenInputs | MultiModalInputs:
         prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs)
         prompt_embeds = parsed_content.get("prompt_embeds")
@@ -82,21 +77,21 @@ class OmniInputPreprocessor(InputPreprocessor):
             inputs = self._process_multimodal(
                 prompt_token_ids,
                 multi_modal_data,
-                parsed_content.get("mm_processor_kwargs") or {},
+                parsed_content.get("mm_processor_kwargs"),
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
+                mm_uuids=parsed_content.get("multi_modal_uuids"),
             )
-            if prompt_embeds is not None:
-                inputs["prompt_embeds"] = prompt_embeds
-            if additional_information is not None:
-                inputs["additional_information"] = additional_information
+
         else:
             inputs = token_inputs_omni(
                 prompt_token_ids=prompt_token_ids,
                 prompt_embeds=prompt_embeds,
                 additional_information=additional_information,
             )
-
+        if prompt_embeds is not None:
+            inputs["prompt_embeds"] = prompt_embeds
+        if additional_information is not None:
+            inputs["additional_information"] = additional_information
         if cache_salt := parsed_content.get("cache_salt"):
             inputs["cache_salt"] = cache_salt
 
@@ -125,8 +120,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         self,
         prompt: SingletonDictPrompt,
         tokenization_kwargs: dict[str, Any] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> SingletonInputs:
         """
         Extract the singleton inputs from a prompt.
@@ -142,7 +135,6 @@ class OmniInputPreprocessor(InputPreprocessor):
         if "prompt_token_ids" in prompt:
             return self._process_tokens(
                 prompt,  # type: ignore[arg-type]
-                mm_uuids=mm_uuids,
             )
 
         if "prompt_embeds" in prompt:
@@ -152,7 +144,6 @@ class OmniInputPreprocessor(InputPreprocessor):
             return self._process_text(
                 prompt,  # type: ignore[arg-type]
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
 
         assert_never(prompt)  # type: ignore[arg-type]