feat: use token2wav in non-streaming tts

modeling_minicpmo.py

@@ -252,33 +252,19 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
         assert os.path.exists(model_dir), f"Asset directory not found: {model_dir}"
         return model_dir
 
-    def init_tts(self, streaming=False, model_dir=None, enable_float16=False, n_timesteps=10):
-        if streaming:
-            if self.config.tts_config.audio_tokenizer_type != "s3tokenizer_step_audio":
-                logger.warning("audio tokenizer type is set to s3tokenizer_step_audio")
-                self.tts.config.audio_tokenizer_type = "s3tokenizer_step_audio"
-
-            try:
-                from stepaudio2 import Token2wav
-            except ImportError:
-                raise ImportError("Please install Token2wav via: pip install minicpmo-utils[all]")
-
-            model_dir = self._ensure_asset_dir("assets/token2wav", model_dir)
-            self.tts.audio_tokenizer = Token2wav(model_dir, float16=enable_float16, n_timesteps=n_timesteps)
-            return self.tts.audio_tokenizer
-        else:
-            if self.config.tts_config.audio_tokenizer_type != "s3tokenizer":
-                logger.warning("audio tokenizer type is set to s3tokenizer")
-                self.tts.config.audio_tokenizer_type = "s3tokenizer"
+    def init_tts(self, model_dir=None, enable_float16=False, n_timesteps=10, **kwargs):
+        if self.config.tts_config.audio_tokenizer_type != "s3tokenizer_step_audio":
+            logger.warning("audio tokenizer type is set to s3tokenizer_step_audio")
+            self.tts.config.audio_tokenizer_type = "s3tokenizer_step_audio"
 
-            try:
-                from cosyvoice.cli.cosyvoice import CosyVoice2
-            except ImportError:
-                raise ImportError("Please install CosyVoice via: pip install minicpmo-utils[all]")
+        try:
+            from stepaudio2 import Token2wav
+        except ImportError:
+            raise ImportError("Please install Token2wav via: pip install minicpmo-utils[all]")
 
-            model_dir = self._ensure_asset_dir("assets/CosyVoice2-0.5B", model_dir)
-            self.tts.audio_tokenizer = CosyVoice2(model_dir=model_dir, load_jit=False, load_trt=False, fp16=False)
-            return self.tts.audio_tokenizer
+        model_dir = self._ensure_asset_dir("assets/token2wav", model_dir)
+        self.tts.audio_tokenizer = Token2wav(model_dir, float16=enable_float16, n_timesteps=n_timesteps)
+        return self.tts.audio_tokenizer
 
     def get_input_embeddings(self):
         return self.llm.get_input_embeddings()
@@ -1336,27 +1322,25 @@ class MiniCPMO(MiniCPMOPreTrainedModel):
             ),
         )
 
-        if self.tts.config.audio_tokenizer_type == "s3tokenizer":
-            generated_tokens = outputs.new_ids.squeeze(-1)
-            reference_audio = audio_prompt
-            prompt_speech_16k = None
-            if reference_audio is not None:
-                logger.debug("use reference audio in data to generate waveform")
-                prompt_speech_16k = torch.tensor(reference_audio).unsqueeze(0)
-
-            for i, j in enumerate(
-                self.tts.audio_tokenizer.token2wav(
-                    speech_token=generated_tokens,
-                    speech_token_len=torch.tensor([generated_tokens.shape[1]], device=generated_tokens.device),
-                    prompt_speech_16k=prompt_speech_16k,
-                    stream=False,
-                )
-            ):
-                waveform_pred = j["tts_speech"]
-                waveform_sample_rate = self.tts.audio_tokenizer.sample_rate  # 24000 here, not 16000 input.
-            return waveform_pred[0]
-        else:
-            raise NotImplementedError
+        import io
+
+        import soundfile as sf
+
+        generated_tokens = outputs.new_ids.squeeze(-1)
+        reference_audio = audio_prompt
+        prompt_wav_path = None
+        if reference_audio is not None:
+            logger.debug("use reference audio in data to generate waveform")
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_wav:
+                prompt_wav_path = tmp_wav.name
+                sf.write(prompt_wav_path, reference_audio, 16000)
+        wav_bytes = self.tts.audio_tokenizer(
+            generated_tokens.squeeze(0).tolist(),
+            prompt_wav_path,
+        )
+        # convert wav bytes back to tensor for caller compatibility
+        waveform, sr = sf.read(io.BytesIO(wav_bytes))
+        return torch.tensor(waveform, dtype=torch.float32)
 
     @torch.inference_mode()
     def init_token2wav_cache(self, prompt_speech_16k):
@@ -2511,11 +2495,7 @@ class MiniCPMODuplex:
         # Initialize TTS (same as __init__)
         enable_float16 = get_param("enable_float16")
         n_timesteps = get_param("n_timesteps")
-        instance.model.init_tts(
-            streaming=True,
-            enable_float16=enable_float16,
-            n_timesteps=n_timesteps,
-        )
+        instance.model.init_tts(enable_float16=enable_float16, n_timesteps=n_timesteps)
 
         instance.break_event = threading.Event()
         instance.session_stop_event = threading.Event()