riva: make sure we don't block on fastpitch

This commit is contained in:
Aleix Conchillo Flaqué
2024-12-13 07:22:05 -08:00
parent 8f24ca4e58
commit aac907aadb

View File

@@ -35,6 +35,8 @@ except ModuleNotFoundError as e:
) )
raise Exception(f"Missing module: {e}") raise Exception(f"Missing module: {e}")
FASTPITCH_TIMEOUT_SECS = 5
class FastPitchTTSService(TTSService): class FastPitchTTSService(TTSService):
class InputParams(BaseModel): class InputParams(BaseModel):
@@ -102,20 +104,23 @@ class FastPitchTTSService(TTSService):
logger.debug(f"Generating TTS: [{text}]") logger.debug(f"Generating TTS: [{text}]")
queue = asyncio.Queue() try:
await asyncio.to_thread(read_audio_responses, queue) queue = asyncio.Queue()
await asyncio.to_thread(read_audio_responses, queue)
# Wait for the thread to start. # Wait for the thread to start.
resp = await queue.get() resp = await asyncio.wait_for(queue.get(), FASTPITCH_TIMEOUT_SECS)
while resp: while resp:
await self.stop_ttfb_metrics() await self.stop_ttfb_metrics()
frame = TTSAudioRawFrame( frame = TTSAudioRawFrame(
audio=resp.audio, audio=resp.audio,
sample_rate=self._sample_rate, sample_rate=self._sample_rate,
num_channels=1, num_channels=1,
) )
yield frame yield frame
resp = await queue.get() resp = await asyncio.wait_for(queue.get(), FASTPITCH_TIMEOUT_SECS)
except asyncio.TimeoutError:
logger.error(f"{self} timeout waiting for audio response")
await self.start_tts_usage_metrics(text) await self.start_tts_usage_metrics(text)
yield TTSStoppedFrame() yield TTSStoppedFrame()