add-audio

Author: sangchengmeng

lightllm/models/qwen3_omni_moe_thinker/model.py

@@ -33,53 +33,43 @@ def _get_feat_extract_output_lengths(input_lengths):
     return output_lengths
 
 
-# <|audio_start|><|audio_pad|><|audio_end|>
-AUDIO_START_TOKEN = "<|audio_start|>"
-AUDIO_END_TOKEN = "<|audio_end|>"
-AUDIO_TOKEN_TOKEN = "<|audio_pad|>"
 MIN_AUDIO_LEN = 480
 
 
 class QWen3OmniTokenizer(QWen3VLTokenizer):
-    def __init__(self, tokenizer=None, image_processor=None, **kwargs):
+    def __init__(self, tokenizer=None, processor=None, **kwargs):
         self.tokenizer = tokenizer
-        self.image_processor = image_processor
+        # image
+        self.image_processor = processor.image_processor
         self.min_pixel = self.image_processor.min_pixels
         self.max_pixel = self.image_processor.max_pixels
         self.patch_size = self.image_processor.patch_size
         self.merge_size = self.image_processor.merge_size
+
+        # audio
+        self.audio_processor = processor.feature_extractor
+        self.sampling_rate = self.audio_processor.sampling_rate
+        self.n_samples = self.audio_processor.n_samples
+        self.hop_length = self.audio_processor.hop_length
+
         self.image_start_id = kwargs["model_cfg"]["vision_start_token_id"]
         self.image_end_id = kwargs["model_cfg"]["vision_end_token_id"]
         self.image_token_id = kwargs["model_cfg"]["image_token_id"]
 
-        self.audio_start_tag = AUDIO_START_TOKEN
-        self.audio_start_id = tokenizer.convert_tokens_to_ids(self.audio_start_tag)
-
-        self.audio_end_tag = AUDIO_END_TOKEN
-        self.audio_end_id = tokenizer.convert_tokens_to_ids(self.audio_end_tag)
-
-        self.audio_token_tag = AUDIO_TOKEN_TOKEN
-        self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token_tag)
-
-        # 这些太hard了, 后面改一下,可以直接从audio_processor里取?
-        self.sampling_rate = 16000
-        self.chunk_length = 30
-        self.n_samples = self.chunk_length * self.sampling_rate
-        self.hop_length = 160
+        self.audio_start_id = kwargs["model_cfg"]["audio_start_token_id"]
+        self.audio_end_id = kwargs["model_cfg"]["audio_end_token_id"]
+        self.audio_token_id = kwargs["model_cfg"]["audio_token_id"]
 
     def init_audioitem_extral_params(
         self, audio: AudioItem, multi_params: MultimodalParams, sampling_params: SamplingParams
     ):
         return
 
     def get_audio_token_length(self, audio: AudioItem):
-        # audio_bytes = audio._preload_data
-        # audio_values, _ = librosa.load(BytesIO(audio_bytes), sr=self.sampling_rate)
-        # length = max(int(audio_values.shape[0]), int(MIN_AUDIO_LEN)) #这个最短还有必要吗?稍等再检查一下
-        # L_eff = min(length, int(self.n_samples))
-        # num_frames = L_eff // int(self.hop_length)
-
-        return 290
+        length = min(audio.audio_length, int(self.n_samples))
+        token_num = length // int(self.hop_length)
+        print(f"token_num is {token_num}")
+        return token_num
 
     def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs):
         origin_ids = self.tokenizer.encode(prompt)

lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py

@@ -340,33 +340,37 @@ def forward(
     def encode(self, audio_items: List[AudioItem], cpu_embed_cache_client: CpuEmbedCacheClient):
         uuids = []
         items: List[AudioItem] = []
+        per_audio_features: List[torch.Tensor] = []
         for i, item in enumerate(audio_items):
             if isinstance(item, AudioItem):
                 uuids.append(item.uuid)
                 items.append(item)
                 audio_data = read_shm(get_shm_name_data(item.uuid))
                 audio = BytesIO(audio_data)
-                audio, _ = librosa.load(audio, sr=16000)
+                audio, _ = librosa.load(audio, sr=self.processor.sampling_rate)
             else:
                 raise ValueError(f"cannot read audio which type is {type(item)}!")
 
-        input_features, feature_attention_mask = self.processor._preprocess(audio, return_attention_mask=True)
-        print(f"input_features is {input_features}, input_features.shape is {input_features.shape}")
-        print(f"feature_attention_mask is {feature_attention_mask}, shape is {feature_attention_mask.shape}")
-        if feature_attention_mask is not None:
-            audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
-            input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
-        else:
-            audio_feature_lengths = None
-        print(f"input_features is {input_features}, input_features.shape is {input_features.shape}")
-
-        feature_lens = audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1)
-        print(f"feature_lens is {feature_lens}")
-        audio_features = self.forward(
-            input_features,
-            feature_lens=feature_lens,
-        )
-        print(f"audio_features is {audio_features}, shape is {audio_features.shape}")
+            input_features, feature_attention_mask = self.processor._preprocess(audio, return_attention_mask=True)
+            print(f"input_features is {input_features}, input_features.shape is {input_features.shape}")
+            print(f"feature_attention_mask is {feature_attention_mask}, shape is {feature_attention_mask.shape}")
+            if feature_attention_mask is not None:
+                audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
+                input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
+            else:
+                audio_feature_lengths = None
+            print(f"input_features is {input_features}, input_features.shape is {input_features.shape}")
+
+            feature_lens = (
+                audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1)
+            )
+            print(f"feature_lens is {feature_lens}")
+            audio_features = self.forward(
+                input_features,
+                feature_lens=feature_lens,
+            )
+            per_audio_features.append(audio_features)
+            print(f"audio_features is {audio_features}, shape is {audio_features.shape}")
 
         ready_audio = obtain(self.cache_client.root.get_items_embed(uuids))
         ids_to_set = []
@@ -377,8 +381,9 @@ def encode(self, audio_items: List[AudioItem], cpu_embed_cache_client: CpuEmbedC
             uid = uuids[i]
             item = items[i]
 
+            cur_embed = per_audio_features[i]
             cpu_embed_cache_client.copy_to_cache(
-                embed_tensor=audio_features, start_index_in_cache=item.start_index_in_embed_cache
+                embed_tensor=cur_embed, start_index_in_cache=item.start_index_in_embed_cache
             )
             ids_to_set.append(uid)
 

lightllm/server/tokenizer.py

@@ -106,7 +106,7 @@ def get_tokenizer(
 
         model_cfg = model_cfg["thinker_config"]
         processor = AutoProcessor.from_pretrained(tokenizer_name)
-        tokenizer = QWen3OmniTokenizer(tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg)
+        tokenizer = QWen3OmniTokenizer(tokenizer, processor=processor, model_cfg=model_cfg)
     elif model_type == "internvl_chat":
         tokenizer = InternvlTokenizer(tokenizer, model_cfg, weight_dir=tokenizer_name)
     elif model_type == "gemma3":