NVIDIATTSService: process incoming audio frame right away

Process audio as soon as we receive it from the generator. Previously, we were
reading from the generator and adding elements into a queue until there was no
more data, then we would process the queue.
This commit is contained in:
Aleix Conchillo Flaqué
2026-01-20 13:44:52 -08:00
parent 14495c425a
commit a787fd9cd8
2 changed files with 27 additions and 28 deletions

View File

@@ -0,0 +1 @@
- Optimized `NVIDIATTSService` to process incoming audio frames immediately.

View File

@@ -12,7 +12,7 @@ gRPC API for high-quality speech synthesis.
import asyncio import asyncio
import os import os
from typing import AsyncGenerator, Mapping, Optional from typing import AsyncGenerator, AsyncIterable, Generator, Mapping, Optional
from pipecat.utils.tracing.service_decorators import traced_tts from pipecat.utils.tracing.service_decorators import traced_tts
@@ -35,14 +35,12 @@ from pipecat.transcriptions.language import Language
try: try:
import riva.client import riva.client
import riva.client.proto.riva_tts_pb2 as rtts
except ModuleNotFoundError as e: except ModuleNotFoundError as e:
logger.error(f"Exception: {e}") logger.error(f"Exception: {e}")
logger.error("In order to use NVIDIA Riva TTS, you need to `pip install pipecat-ai[nvidia]`.") logger.error("In order to use NVIDIA Riva TTS, you need to `pip install pipecat-ai[nvidia]`.")
raise Exception(f"Missing module: {e}") raise Exception(f"Missing module: {e}")
NVIDIA_TTS_TIMEOUT_SECS = 5
class NvidiaTTSService(TTSService): class NvidiaTTSService(TTSService):
"""NVIDIA Riva text-to-speech service. """NVIDIA Riva text-to-speech service.
@@ -165,11 +163,7 @@ class NvidiaTTSService(TTSService):
Frame: Audio frames containing the synthesized speech data. Frame: Audio frames containing the synthesized speech data.
""" """
def read_audio_responses(queue: asyncio.Queue): def read_audio_responses() -> Generator[rtts.SynthesizeSpeechResponse, None, None]:
def add_response(r):
asyncio.run_coroutine_threadsafe(queue.put(r), self.get_event_loop())
try:
responses = self._service.synthesize_online( responses = self._service.synthesize_online(
text, text,
self._voice_id, self._voice_id,
@@ -179,12 +173,20 @@ class NvidiaTTSService(TTSService):
zero_shot_quality=self._quality, zero_shot_quality=self._quality,
custom_dictionary={}, custom_dictionary={},
) )
for r in responses: return responses
add_response(r)
add_response(None) def async_next(it):
except Exception as e: try:
logger.error(f"{self} exception: {e}") return next(it)
add_response(None) except StopIteration:
return None
async def async_iterator(iterator) -> AsyncIterable[rtts.SynthesizeSpeechResponse]:
while True:
item = await asyncio.to_thread(async_next, iterator)
if item is None:
return
yield item
try: try:
assert self._service is not None, "TTS service not initialized" assert self._service is not None, "TTS service not initialized"
@@ -195,12 +197,9 @@ class NvidiaTTSService(TTSService):
logger.debug(f"{self}: Generating TTS [{text}]") logger.debug(f"{self}: Generating TTS [{text}]")
queue = asyncio.Queue() responses = await asyncio.to_thread(read_audio_responses)
await asyncio.to_thread(read_audio_responses, queue)
# Wait for the thread to start. async for resp in async_iterator(responses):
resp = await asyncio.wait_for(queue.get(), timeout=NVIDIA_TTS_TIMEOUT_SECS)
while resp:
await self.stop_ttfb_metrics() await self.stop_ttfb_metrics()
frame = TTSAudioRawFrame( frame = TTSAudioRawFrame(
audio=resp.audio, audio=resp.audio,
@@ -208,7 +207,6 @@ class NvidiaTTSService(TTSService):
num_channels=1, num_channels=1,
) )
yield frame yield frame
resp = await asyncio.wait_for(queue.get(), timeout=NVIDIA_TTS_TIMEOUT_SECS)
await self.start_tts_usage_metrics(text) await self.start_tts_usage_metrics(text)
yield TTSStoppedFrame() yield TTSStoppedFrame()