WIP: support video and audio, refactor

Author: CUHKSZzxy

lmdeploy/pytorch/models/qwen3_omni_moe_thinker.py

@@ -899,7 +899,7 @@ def _make_image_mm_data(self, input_mm: Dict[str, Any]) -> MultiModalData:
         offset = input_mm['offset']
         start = offset
         image_token_id = input_mm['image_token_id']
-        num_pad = input_mm['image_tokens']
+        num_pad = input_mm['mm_token_num']
         if isinstance(num_pad, torch.Tensor):
             num_pad = num_pad.item()
 
@@ -917,7 +917,7 @@ def _make_video_mm_data(self, input_mm: Dict[str, Any]) -> MultiModalData:
         offset = input_mm['offset']
         start = offset
         video_token_id = input_mm['video_token_id']
-        num_pad = input_mm['video_tokens']
+        num_pad = input_mm['mm_token_num']
         if isinstance(num_pad, torch.Tensor):
             num_pad = num_pad.item()
 

lmdeploy/pytorch/models/qwen3_vl.py

@@ -15,7 +15,6 @@ class MultiModalData:
     start: int
     end: int = None
     meta: Dict[str, Any] = None
-
     modality: Modality = Modality.IMAGE
 
     def __post_init__(self):

lmdeploy/pytorch/multimodal/data_type.py

@@ -8,6 +8,7 @@
 from lmdeploy.tokenizer import Tokenizer
 from lmdeploy.utils import get_logger
 from lmdeploy.vl.constants import Modality
+from lmdeploy.vl.media.audio import AudioMediaIO
 from lmdeploy.vl.media.connection import load_from_url
 from lmdeploy.vl.media.image import ImageMediaIO
 from lmdeploy.vl.media.time_series import TimeSeriesMediaIO
@@ -124,6 +125,10 @@ def _parse_multimodal_item(i: int, in_messages: List[Dict], out_messages: List[D
                 vid_io = VideoMediaIO(image_io=ImageMediaIO(), **media_io_kwargs.get('video', {}))
                 data, metadata = load_from_url(data_src, vid_io)
                 item_params['video_metadata'] = metadata
+            elif item_type == 'audio_url':
+                modality = Modality.AUDIO
+                audio_io = AudioMediaIO(**media_io_kwargs.get('audio', {}))
+                data = load_from_url(data_src, audio_io)
             elif item_type == 'time_series_url':
                 modality = Modality.TIME_SERIES
                 ts_io = TimeSeriesMediaIO(**media_io_kwargs.get('time_series', {}))
@@ -304,7 +309,7 @@ def _re_format_prompt_images_pair(prompt: Tuple) -> Dict:
 
     def _has_multimodal_input(self, messages: List[Dict]) -> bool:
         """Check if messages contain multimodal input (images)."""
-        multimodal_types = ['image_url', 'image_data', 'video_url', 'time_series_url']
+        multimodal_types = ['image_url', 'image_data', 'video_url', 'audio_url', 'time_series_url']
         return any(
             isinstance(message.get('content'), list) and any(
                 item.get('type') in multimodal_types for item in message['content']) for message in messages)

lmdeploy/serve/processors/multimodal.py

@@ -305,6 +305,10 @@ def _get_and_verify_max_len(
     for key in llm_keys:
         hf_config = getattr(hf_config, key, hf_config)
 
+    # for qwen3-omni thinker
+    if hasattr(hf_config, 'thinker_config'):
+        hf_config = hf_config.thinker_config.text_config
+
     logger = get_logger('lmdeploy')
     derived_max_model_len = float('inf')
     possible_keys = [

lmdeploy/utils.py

@@ -0,0 +1,57 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+# adapted from https://github.com/vllm-project/vllm/blob/main/vllm/multimodal/media/audio.py
+
+import base64
+from io import BytesIO
+from pathlib import Path
+
+import numpy.typing as npt
+
+from .base import MediaIO
+
+try:
+    import librosa
+except ImportError:
+    raise ImportError('Please install librosa via `pip install librosa`.')
+
+try:
+    import soundfile
+except ImportError:
+    raise ImportError('Please install soundfile via `pip install soundfile`.')
+
+
+class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]):
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+
+        # for potential custom arguments from --media-io-kwargs
+        self.kwargs = kwargs
+
+    def load_bytes(self, data: bytes) -> tuple[npt.NDArray, float]:
+        # sr=None preserves the original sampling rate of the audio file
+        return librosa.load(BytesIO(data), sr=None)
+
+    def load_base64(
+        self,
+        media_type: str,
+        data: str,
+    ) -> tuple[npt.NDArray, float]:
+        return self.load_bytes(base64.b64decode(data))
+
+    def load_file(self, filepath: Path) -> tuple[npt.NDArray, float]:
+        return librosa.load(filepath, sr=None)
+
+    def encode_base64(
+        self,
+        media: tuple[npt.NDArray, int],
+        *,
+        audio_format: str = 'WAV',
+    ) -> str:
+        audio, sr = media
+
+        with BytesIO() as buffer:
+            soundfile.write(buffer, audio, sr, format=audio_format)
+            data = buffer.getvalue()
+
+        return base64.b64encode(data).decode('utf-8')

lmdeploy/vl/media/audio.py

@@ -252,40 +252,32 @@ def to_pytorch_with_input_ids(self, messages):
 
         return dict(prompt=None, input_ids=input_ids, multimodal=preps)
 
-    def to_pytorch_aux(self, messages, prompt, IMAGE_TOKEN, tokenizer, sequence_start):
+    def to_pytorch_aux(self, messages, prompt, mm_placeholder, tokenizer, sequence_start):
         """Auxiliary function to pack the preprocessing results in a format
-        compatible with what is required by pytorch engine.
-
-        Args:
-            messages(List[Dict]): the output of `preprocess`
-            prompt(str): the prompt after applying chat template
-            IMAGE_TOKEN(str): a placeholder where image tokens will be
-                inserted
-            tokenzer: the tokenizer model
-            sequence_start: starting flag of a sequence
-        """
-        # collect all preprocessing result from messages
-        preps = [x['content'] for x in messages if x['role'] == 'preprocess']
-        assert len(preps) == 1
-        preps = preps[0]
+        compatible with what is required by pytorch engine."""
+        # collect all multi-modal preprocessing result from messages, keyed by 'preprocess'
+        mm_items = [x['content'] for x in messages if x['role'] == 'preprocess']
+        assert len(mm_items) == 1
+        mm_items = mm_items[0]
 
         # split prompt into segments and validate data
-        segs = prompt.split(IMAGE_TOKEN)
-        assert len(segs) == len(preps) + 1, (f'the number of {IMAGE_TOKEN} is not equal '
-                                             f'to input images, {len(segs) - 1} vs {len(preps)}')
+        prompt_segments = prompt.split(mm_placeholder)
+        assert len(prompt_segments) == len(mm_items) + 1, (
+            f'the number of {mm_placeholder} is not equal '
+            f'to input multi modal items, {len(mm_items) - 1} vs {len(prompt_segments)}')
 
-        # calculate the image token offset for each image
+        # calculate the token offset for each multi modal item
         input_ids = []
-        for i, seg in enumerate(segs):
-            if i > 0 and i <= len(preps):
-                preps[i - 1].update(offset=len(input_ids))
-                image_tokens = preps[i - 1]['image_tokens']
-                assert self.image_token_id == preps[i - 1]['image_token_id']
-                input_ids.extend([self.image_token_id] * image_tokens)
+        mm_placeholder_id = tokenizer.encode(mm_placeholder, add_special_tokens=False)[-1]
+        for i, seg in enumerate(prompt_segments):
+            if i > 0 and i <= len(mm_items):
+                mm_items[i - 1].update(offset=len(input_ids))
+                mm_token_num = mm_items[i - 1]['mm_token_num']
+                input_ids.extend([mm_placeholder_id] * mm_token_num)
             token_ids = tokenizer.encode(seg, add_bos=((i == 0) and sequence_start))
             input_ids.extend(token_ids)
 
-        return dict(prompt=prompt, input_ids=input_ids, multimodal=preps)
+        return dict(prompt=prompt, input_ids=input_ids, multimodal=mm_items)
 
     def to_turbomind_aux(self, messages, prompt, IMAGE_TOKEN, tokenizer, sequence_start):
         """Auxiliary function to pack the forwarding results in a format

lmdeploy/vl/model/base.py

@@ -93,7 +93,7 @@ def _preprocess_image(self,
                                                 return_tensors='pt')
         merge_length = self.processor.image_processor.merge_size**2
         image_tokens = result['image_grid_thw'].prod(dim=1) // merge_length
-        result.update(dict(image_size=image.size, image_tokens=image_tokens, image_token_id=self.image_token_id))
+        result.update(dict(image_size=image.size, mm_token_num=image_tokens, image_token_id=self.image_token_id))
         return result
 
     def _preprocess_video(self,
@@ -206,7 +206,7 @@ def to_pytorch_aux_video(self, messages, prompt, VIDEO_TOKEN, tokenizer, sequenc
                 video_token_ids = tokenizer.encode(video_placeholder)
                 input_ids.extend(video_token_ids)
 
-                preps[i - 1].update(video_tokens=len(video_token_ids))
+                preps[i - 1].update(mm_token_num=len(video_token_ids))
 
             token_ids = tokenizer.encode(seg, add_bos=((i == 0) and sequence_start))
             input_ids.extend(token_ids)