[BugFix] token2wav code out of range (#655)

Signed-off-by: dengyunyang <584797741@qq.com>

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -101,6 +101,12 @@ class Qwen2_5OmniForConditionalGeneration(
             self.talker.init_multi_modal(thinker_config)
             self.model = self.talker
             self.token2wav = None
+            # set suppress start id according to token2wav
+            t2w_token_end_id = getattr(
+                getattr(getattr(config, "token2wav_config", None), "dit_config", None), "num_embeds", None
+            )
+            if t2w_token_end_id:
+                self.model.set_suppress_start_id(t2w_token_end_id + 1)
 
         elif self.model_stage == "code2wav":
             self.thinker = None

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py

@@ -84,6 +84,9 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         )
         self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
 
+        # suppress start id
+        self.suppress_start_id = None
+
     def init_multi_modal(self, thinker_config):
         self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
         self.visual = Qwen2_5_VisionTransformer(
@@ -150,6 +153,10 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         return hidden_states
 
     def bad_word_processor(self, logits: torch.Tensor) -> torch.Tensor:
+        # suppress token IDs unsupported by token2wav
+        if self.suppress_start_id and self.suppress_start_id < logits.size(-1):
+            logits[..., self.suppress_start_id : logits.size(-1)] = -1e9
+
         if hasattr(self.config, "tts_codec_start_token_id"):
             bos_id = int(getattr(self.config, "tts_codec_start_token_id"))
             logits[..., bos_id] = -1e9
@@ -234,3 +241,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
                 audio_embeddings = self._process_audio_input(multimodal_input)
                 multimodal_embeddings += audio_embeddings
         return multimodal_embeddings
+
+    def set_suppress_start_id(self, start_id: int):
+        self.suppress_start_id = start_id
+        self.logger.debug(f"Set suppress start id to {self.suppress_start_id}")