From 4b55c73fbefb509c451aaf0f1b94d6570c611dc4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Mon, 9 Dec 2024 13:03:19 -0800
Subject: [PATCH] services(riva): make FastpitchTTSService asyncio

---
 src/pipecat/services/riva.py | 52 +++++++++++++++++++-----------------
 1 file changed, 27 insertions(+), 25 deletions(-)

diff --git a/src/pipecat/services/riva.py b/src/pipecat/services/riva.py
index d6c5f85f5..38920927e 100644
--- a/src/pipecat/services/riva.py
+++ b/src/pipecat/services/riva.py
@@ -60,8 +60,6 @@ class FastpitchTTSService(TTSService):
         self.voice_id = voice_id
         self.sample_rate_hz = sample_rate_hz
         self.language_code = params.language
-        self.nchannels = 1
-        self.sampwidth = 2
         self.quality = None
 
         metadata = [
@@ -73,35 +71,39 @@ class FastpitchTTSService(TTSService):
         self.service = riva.client.SpeechSynthesisService(auth)
 
     async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
+        def read_audio_responses():
+            try:
+                custom_dictionary_input = {}
+                responses = self.service.synthesize_online(
+                    text,
+                    self.voice_id,
+                    self.language_code,
+                    sample_rate_hz=self.sample_rate_hz,
+                    audio_prompt_file=None,
+                    quality=20 if self.quality is None else self.quality,
+                    custom_dictionary=custom_dictionary_input,
+                )
+                return responses
+            except Exception as e:
+                logger.error(f"{self} exception: {e}")
+                return []
+
         logger.debug(f"Generating TTS: [{text}]")
 
         await self.start_ttfb_metrics()
         yield TTSStartedFrame()
 
-        try:
-            custom_dictionary_input = {}
-            responses = self.service.synthesize_online(
-                text,
-                self.voice_id,
-                self.language_code,
-                sample_rate_hz=self.sample_rate_hz,
-                audio_prompt_file=None,
-                quality=20 if self.quality is None else self.quality,
-                custom_dictionary=custom_dictionary_input,
+        responses = await asyncio.to_thread(read_audio_responses)
+
+        for resp in responses:
+            await self.stop_ttfb_metrics()
+
+            frame = TTSAudioRawFrame(
+                audio=resp.audio,
+                sample_rate=self.sample_rate_hz,
+                num_channels=1,
             )
-
-            for resp in responses:
-                await self.stop_ttfb_metrics()
-
-                frame = TTSAudioRawFrame(
-                    audio=resp.audio,
-                    sample_rate=self.sample_rate_hz,
-                    num_channels=self.nchannels,
-                )
-                yield frame
-
-        except Exception as e:
-            logger.error(f"{self} exception: {e}")
+            yield frame
 
         await self.start_tts_usage_metrics(text)
         yield TTSStoppedFrame()