[infer] vllm video/audio inference (#7566)

Author: hiyouga

scripts/vllm_infer.py

@@ -92,8 +92,20 @@ def vllm_infer(
             multi_modal_data = {
                 "image": template_obj.mm_plugin._regularize_images(
                     sample["images"], image_max_pixels=image_max_pixels, image_min_pixels=image_min_pixels
-                )
+                )["images"]
             }
+        elif sample["videos"]:
+            multi_modal_data = {
+                "video": template_obj.mm_plugin._regularize_videos(
+                    sample["videos"], image_max_pixels=image_max_pixels, image_min_pixels=image_min_pixels
+                )["videos"]
+            }
+        elif sample["audios"]:
+            audio_data = template_obj.mm_plugin._regularize_audios(
+                sample["audios"],
+                sampling_rate=16000,
+            )
+            multi_modal_data = {"audio": zip(audio_data["audios"], audio_data["sampling_rates"])}
         else:
             multi_modal_data = None
 
@@ -131,7 +143,7 @@ def vllm_infer(
         "enable_lora": model_args.adapter_name_or_path is not None,
     }
     if template_obj.mm_plugin.__class__.__name__ != "BasePlugin":
-        engine_args["limit_mm_per_prompt"] = {"image": 4, "video": 2}
+        engine_args["limit_mm_per_prompt"] = {"image": 4, "video": 2, "audio": 2}
 
     if isinstance(model_args.vllm_config, dict):
         engine_args.update(model_args.vllm_config)

src/llamafactory/api/chat.py

@@ -23,7 +23,7 @@
 
 from ..data import Role as DataRole
 from ..extras import logging
-from ..extras.constants import IMAGE_PLACEHOLDER
+from ..extras.constants import AUDIO_PLACEHOLDER, IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER
 from ..extras.misc import is_env_enabled
 from ..extras.packages import is_fastapi_available, is_pillow_available, is_requests_available
 from .common import dictify, jsonify
@@ -56,7 +56,7 @@
 
 if TYPE_CHECKING:
     from ..chat import ChatModel
-    from ..data.mm_plugin import ImageInput
+    from ..data.mm_plugin import AudioInput, ImageInput, VideoInput
     from .protocol import ChatCompletionRequest, ScoreEvaluationRequest
 
 
@@ -72,7 +72,14 @@
 
 def _process_request(
     request: "ChatCompletionRequest",
-) -> tuple[list[dict[str, str]], Optional[str], Optional[str], Optional[list["ImageInput"]]]:
+) -> tuple[
+    list[dict[str, str]],
+    Optional[str],
+    Optional[str],
+    Optional[list["ImageInput"]],
+    Optional[list["VideoInput"]],
+    Optional[list["AudioInput"]],
+]:
     if is_env_enabled("API_VERBOSE", "1"):
         logger.info_rank0(f"==== request ====\n{json.dumps(dictify(request), indent=2, ensure_ascii=False)}")
 
@@ -88,7 +95,7 @@ def _process_request(
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Only supports u/a/u/a/u...")
 
     input_messages = []
-    images = []
+    images, videos, audios = [], [], []
     for i, message in enumerate(request.messages):
         if i % 2 == 0 and message.role not in [Role.USER, Role.TOOL]:
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid role")
@@ -107,7 +114,7 @@ def _process_request(
             for input_item in message.content:
                 if input_item.type == "text":
                     text_content += input_item.text
-                else:
+                elif input_item.type == "image_url":
                     text_content += IMAGE_PLACEHOLDER
                     image_url = input_item.image_url.url
                     if re.match(r"^data:image\/(png|jpg|jpeg|gif|bmp);base64,(.+)$", image_url):  # base64 image
@@ -118,6 +125,28 @@ def _process_request(
                         image_stream = requests.get(image_url, stream=True).raw
 
                     images.append(Image.open(image_stream).convert("RGB"))
+                elif input_item.type == "video_url":
+                    text_content += VIDEO_PLACEHOLDER
+                    video_url = input_item.video_url.url
+                    if os.path.isfile(video_url):  # local file
+                        video_stream = open(video_url, "rb")
+                    else:  # web uri
+                        video_stream = requests.get(video_url, stream=True).raw
+
+                    videos.append(video_stream)
+                elif input_item.type == "audio_url":
+                    text_content += AUDIO_PLACEHOLDER
+                    audio_url = input_item.audio_url.url
+                    if os.path.isfile(audio_url):  # local file
+                        audio_stream = open(audio_url, "rb")
+                    else:  # web uri
+                        audio_stream = requests.get(audio_url, stream=True).raw
+
+                    audios.append(audio_stream)
+                else:
+                    raise HTTPException(
+                        status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid input type {input_item.type}."
+                    )
 
             input_messages.append({"role": ROLE_MAPPING[message.role], "content": text_content})
         else:
@@ -132,7 +161,7 @@ def _process_request(
     else:
         tools = None
 
-    return input_messages, system, tools, images or None
+    return input_messages, system, tools, images or None, videos or None, audios or None
 
 
 def _create_stream_chat_completion_chunk(
@@ -151,12 +180,14 @@ async def create_chat_completion_response(
     request: "ChatCompletionRequest", chat_model: "ChatModel"
 ) -> "ChatCompletionResponse":
     completion_id = f"chatcmpl-{uuid.uuid4().hex}"
-    input_messages, system, tools, images = _process_request(request)
+    input_messages, system, tools, images, videos, audios = _process_request(request)
     responses = await chat_model.achat(
         input_messages,
         system,
         tools,
         images,
+        videos,
+        audios,
         do_sample=request.do_sample,
         temperature=request.temperature,
         top_p=request.top_p,
@@ -202,7 +233,7 @@ async def create_stream_chat_completion_response(
     request: "ChatCompletionRequest", chat_model: "ChatModel"
 ) -> AsyncGenerator[str, None]:
     completion_id = f"chatcmpl-{uuid.uuid4().hex}"
-    input_messages, system, tools, images = _process_request(request)
+    input_messages, system, tools, images, videos, audios = _process_request(request)
     if tools:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Cannot stream function calls.")
 
@@ -217,6 +248,8 @@ async def create_stream_chat_completion_response(
         system,
         tools,
         images,
+        videos,
+        audios,
         do_sample=request.do_sample,
         temperature=request.temperature,
         top_p=request.top_p,

src/llamafactory/api/protocol.py

@@ -70,14 +70,17 @@ class FunctionCall(BaseModel):
     function: Function
 
 
-class ImageURL(BaseModel):
+class URL(BaseModel):
     url: str
+    detail: Literal["auto", "low", "high"] = "auto"
 
 
 class MultimodalInputItem(BaseModel):
-    type: Literal["text", "image_url"]
+    type: Literal["text", "image_url", "video_url", "audio_url"]
     text: Optional[str] = None
-    image_url: Optional[ImageURL] = None
+    image_url: Optional[URL] = None
+    video_url: Optional[URL] = None
+    audio_url: Optional[URL] = None
 
 
 class ChatMessage(BaseModel):

src/llamafactory/chat/sglang_engine.py

@@ -33,7 +33,7 @@
 
 
 if is_sglang_available():
-    from sglang.utils import launch_server_cmd, terminate_process, wait_for_server
+    from sglang.utils import launch_server_cmd, terminate_process, wait_for_server  # type: ignore
 
 
 if TYPE_CHECKING:
@@ -134,24 +134,17 @@ async def _generate(
         audios: Optional[list["AudioInput"]] = None,
         **input_kwargs,
     ) -> AsyncIterator[dict[str, Any]]:
-        mm_input_dict = {"images": [], "videos": [], "audios": [], "imglens": [0], "vidlens": [0], "audlens": [0]}
-        if images is not None:
-            mm_input_dict.update({"images": images, "imglens": [len(images)]})
-            if not any(IMAGE_PLACEHOLDER in message["content"] for message in messages):
-                messages[0]["content"] = IMAGE_PLACEHOLDER * len(images) + messages[0]["content"]
-
-        if videos is not None:
-            mm_input_dict.update({"videos": videos, "vidlens": [len(videos)]})
-            if not any(VIDEO_PLACEHOLDER in message["content"] for message in messages):
-                messages[0]["content"] = VIDEO_PLACEHOLDER * len(videos) + messages[0]["content"]
-
-        if audios is not None:
-            mm_input_dict.update({"audios": audios, "audlens": [len(audios)]})
-            if not any(AUDIO_PLACEHOLDER in message["content"] for message in messages):
-                messages[0]["content"] = AUDIO_PLACEHOLDER * len(audios) + messages[0]["content"]
+        if images is not None and not any(IMAGE_PLACEHOLDER in message["content"] for message in messages):
+            messages[0]["content"] = IMAGE_PLACEHOLDER * len(images) + messages[0]["content"]
+
+        if videos is not None and not any(VIDEO_PLACEHOLDER in message["content"] for message in messages):
+            messages[0]["content"] = VIDEO_PLACEHOLDER * len(videos) + messages[0]["content"]
+
+        if audios is not None and not any(AUDIO_PLACEHOLDER in message["content"] for message in messages):
+            messages[0]["content"] = AUDIO_PLACEHOLDER * len(audios) + messages[0]["content"]
 
         messages = self.template.mm_plugin.process_messages(
-            messages, mm_input_dict["images"], mm_input_dict["videos"], mm_input_dict["audios"], self.processor
+            messages, images or [], videos or [], audios or [], self.processor
         )
         paired_messages = messages + [{"role": "assistant", "content": ""}]
         system = system or self.generating_args["default_system"]

src/llamafactory/chat/vllm_engine.py

@@ -83,7 +83,7 @@ def __init__(
             "max_lora_rank": model_args.vllm_max_lora_rank,
         }
         if self.template.mm_plugin.__class__.__name__ != "BasePlugin":
-            engine_args["limit_mm_per_prompt"] = {"image": 4, "video": 2}
+            engine_args["limit_mm_per_prompt"] = {"image": 4, "video": 2, "audio": 2}
 
         if isinstance(model_args.vllm_config, dict):
             engine_args.update(model_args.vllm_config)
@@ -111,24 +111,17 @@ async def _generate(
         **input_kwargs,
     ) -> AsyncIterator["RequestOutput"]:
         request_id = f"chatcmpl-{uuid.uuid4().hex}"
-        mm_input_dict = {"images": [], "videos": [], "audios": [], "imglens": [0], "vidlens": [0], "audlens": [0]}
-        if images is not None:
-            mm_input_dict.update({"images": images, "imglens": [len(images)]})
-            if not any(IMAGE_PLACEHOLDER in message["content"] for message in messages):
-                messages[0]["content"] = IMAGE_PLACEHOLDER * len(images) + messages[0]["content"]
-
-        if videos is not None:
-            mm_input_dict.update({"videos": videos, "vidlens": [len(videos)]})
-            if not any(VIDEO_PLACEHOLDER in message["content"] for message in messages):
-                messages[0]["content"] = VIDEO_PLACEHOLDER * len(videos) + messages[0]["content"]
-
-        if audios is not None:
-            mm_input_dict.update({"audios": audios, "audlens": [len(audios)]})
-            if not any(AUDIO_PLACEHOLDER in message["content"] for message in messages):
-                messages[0]["content"] = AUDIO_PLACEHOLDER * len(audios) + messages[0]["content"]
+        if images is not None and not any(IMAGE_PLACEHOLDER in message["content"] for message in messages):
+            messages[0]["content"] = IMAGE_PLACEHOLDER * len(images) + messages[0]["content"]
+
+        if videos is not None and not any(VIDEO_PLACEHOLDER in message["content"] for message in messages):
+            messages[0]["content"] = VIDEO_PLACEHOLDER * len(videos) + messages[0]["content"]
+
+        if audios is not None and not any(AUDIO_PLACEHOLDER in message["content"] for message in messages):
+            messages[0]["content"] = AUDIO_PLACEHOLDER * len(audios) + messages[0]["content"]
 
         messages = self.template.mm_plugin.process_messages(
-            messages, mm_input_dict["images"], mm_input_dict["videos"], mm_input_dict["audios"], self.processor
+            messages, images or [], videos or [], audios or [], self.processor
         )
         paired_messages = messages + [{"role": "assistant", "content": ""}]
         system = system or self.generating_args["default_system"]
@@ -186,8 +179,24 @@ async def _generate(
                     images,
                     image_max_pixels=self.model_args.image_max_pixels,
                     image_min_pixels=self.model_args.image_min_pixels,
-                )
+                )["images"]
             }
+        elif videos is not None:
+            multi_modal_data = {
+                "video": self.template.mm_plugin._regularize_videos(
+                    videos,
+                    image_max_pixels=self.model_args.video_max_pixels,
+                    image_min_pixels=self.model_args.video_min_pixels,
+                    video_fps=self.model_args.video_fps,
+                    video_maxlen=self.model_args.video_maxlen,
+                )["videos"]
+            }
+        elif audios is not None:
+            audio_data = self.template.mm_plugin._regularize_audios(
+                audios,
+                sampling_rate=self.model_args.audio_sampling_rate,
+            )
+            multi_modal_data = {"audio": zip(audio_data["audios"], audio_data["sampling_rates"])}
         else:
             multi_modal_data = None
 

src/llamafactory/data/converter.py

@@ -26,8 +26,12 @@
     from transformers import Seq2SeqTrainingArguments
 
     from ..hparams import DataArguments
+    from .mm_plugin import AudioInput, ImageInput, VideoInput
     from .parser import DatasetAttr
 
+    MediaType = Union[ImageInput, VideoInput, AudioInput]
+
+
 logger = logging.get_logger(__name__)
 
 
@@ -36,10 +40,12 @@ class DatasetConverter:
     dataset_attr: "DatasetAttr"
     data_args: "DataArguments"
 
-    def _find_medias(self, medias: Union[Any, list[Any]]) -> Optional[list[Any]]:
+    def _find_medias(self, medias: Union["MediaType", list["MediaType"], None]) -> Optional[list["MediaType"]]:
         r"""Optionally concatenate media path to media dir when loading from local disk."""
-        if not isinstance(medias, list):
-            medias = [medias] if medias is not None else []
+        if medias is None:
+            return None
+        elif not isinstance(medias, list):
+            medias = [medias]
         elif len(medias) == 0:
             return None
         else: