thread simplification + handling interuption

This commit is contained in:
sathwika
2026-04-10 17:18:22 +05:30
parent 8abda808ca
commit 746fadc2b5
3 changed files with 204 additions and 113 deletions

View File

@@ -1,9 +1,6 @@
- Added cross-sentence stitching support to NvidiaTTSService. Multiple sentences within an LLM turn are fed into a single `SynthesizeOnline` gRPC stream, enabling Magpie's seamless audio across sentence boundaries (requires Magpie TTS model v1.7.0+).
- Added new parameters to NvidiaTTSService:
- `custom_dictionary` for IPA-based custom pronunciation
- `encoding` for output audio encoding format
- `zero_shot_audio_prompt_file` for zero-shot voice cloning with file caching
- `audio_prompt_encoding` for zero-shot audio prompt format
- Added `can_generate_metrics` returning True for NvidiaTTSService.
- Added `stop_all_metrics()` call on audio context interruption.
- Added gRPC error handling for synthesis config retrieval.
- Added enhancements to `NvidiaTTSService`:
- Cross-sentence stitching: multiple sentences within an LLM turn are fed into a single `SynthesizeOnline` gRPC stream for seamless audio across sentence boundaries (requires Magpie TTS model v1.7.0+).
- `custom_dictionary` and `encoding` parameters for IPA-based custom pronunciation and output audio encoding.
- Metrics generation (`can_generate_metrics` returns true) and `stop_all_metrics()` when an audio context is interrupted.
- gRPC error handling around synthesis config retrieval (`GetRivaSynthesisConfig`).

View File

@@ -1,5 +1,8 @@
- Updated NvidiaTTSService:
- Made `api_key` optional for local NIM deployments
- Fixed `_update_settings` to correctly indicate that runtime settings updates (voice, language, quality) take effect on the next turn
- Replaced per-sentence synchronous `synthesize_online` calls with async queue-backed gRPC streaming
- Renamed Riva references to Nemotron Speech
- Updated `NvidiaTTSService`:
- Made `api_key` optional for local NIM deployments.
- Voice, language, and quality can be updated without reconnecting the gRPC client; new values take effect on the next synthesis turn, not for the current turn's in-flight requests.
- Replaced per-sentence synchronous `synthesize_online` calls with async queue-backed gRPC streaming.
- Streaming now uses asyncio tasks with explicit gRPC cancellation on interruption and stale-response filtering when a stream is aborted or replaced.
- Renamed Riva references to Nemotron Speech in docs and messages.
- Disabled automatic TTS start frames at the service level (`push_start_frame=False`) and emit `TTSStartedFrame` when a stitched synthesis stream is started for a context.

View File

@@ -13,11 +13,6 @@ gRPC API for high-quality speech synthesis.
Refer to the NVIDIA TTS NIM documentation for usage, customization,
and local deployment steps:
https://docs.nvidia.com/nim/speech/latest/tts/
For zero-shot voice cloning, request access to the Magpie TTS Zero-Shot model:
https://developer.nvidia.com/riva-tts-zeroshot-models
Local or private cloud deployment is recommended for best zero-shot performance.
"""
import asyncio
@@ -26,7 +21,6 @@ import queue
import textwrap
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, AsyncGenerator, Mapping, Optional
from pipecat.utils.tracing.service_decorators import traced_tts
@@ -74,6 +68,19 @@ class NvidiaTTSSettings(TTSSettings):
quality: int | _NotGiven = field(default_factory=lambda: NOT_GIVEN)
@dataclass
class _SynthesisStreamState:
"""Runtime state for one active synthesis stream."""
context_id: str
text_queue: queue.Queue
response_queue: asyncio.Queue
stop_event: threading.Event
rpc_call: Any = None
synth_task: Optional[asyncio.Task] = None
response_task: Optional[asyncio.Task] = None
class NvidiaTTSService(TTSService):
"""NVIDIA Nemotron Speech text-to-speech service.
@@ -116,8 +123,6 @@ class NvidiaTTSService(TTSService):
use_ssl: bool = True,
custom_dictionary: Optional[dict] = None,
encoding: Optional[AudioEncoding] = AudioEncoding.LINEAR_PCM,
zero_shot_audio_prompt_file: Optional[Path] = None,
audio_prompt_encoding: Optional[AudioEncoding] = AudioEncoding.ENCODING_UNSPECIFIED,
**kwargs,
):
"""Initialize the NVIDIA Nemotron Speech TTS service.
@@ -149,13 +154,6 @@ class NvidiaTTSService(TTSService):
https://docs.nvidia.com/nim/speech/latest/tts/phoneme-support.html
for the list of supported IPA phonemes.
encoding: Output audio encoding format. Defaults to ``AudioEncoding.LINEAR_PCM``.
zero_shot_audio_prompt_file: Path to audio prompt file for zero-shot voice
cloning. Audio length should be between 3-10 seconds. The file
is read once at init time and cached in memory. Requires the
Magpie TTS Zero-Shot model. See
https://docs.nvidia.com/nim/speech/latest/tts/voice-cloning.html
audio_prompt_encoding: Encoding of the zero-shot audio prompt file.
Defaults to ``AudioEncoding.ENCODING_UNSPECIFIED``.
**kwargs: Additional arguments passed to parent TTSService.
"""
# 1. Initialize default_settings with hardcoded defaults
@@ -202,31 +200,12 @@ class NvidiaTTSService(TTSService):
entries = [f"{k} {v}" for k, v in custom_dictionary.items()]
self._custom_dictionary = ",".join(entries)
self._encoding = encoding
self._audio_prompt_encoding = audio_prompt_encoding
self._zero_shot_audio_prompt_file = zero_shot_audio_prompt_file
self._zero_shot_audio_prompt: Optional[bytes] = None
if self._zero_shot_audio_prompt_file is not None:
if not self._zero_shot_audio_prompt_file.exists():
raise FileNotFoundError(
f"Zero-shot audio prompt file not found: {self._zero_shot_audio_prompt_file}"
)
with self._zero_shot_audio_prompt_file.open("rb") as f:
self._zero_shot_audio_prompt = f.read()
logger.debug(
f"Loaded zero-shot audio prompt from {self._zero_shot_audio_prompt_file} "
f"({len(self._zero_shot_audio_prompt)} bytes)"
)
self._service = None
self._config = None
# Persistent gRPC stream state for cross-sentence stitching
self._text_queue: Optional[queue.Queue] = None
self._synth_thread: Optional[threading.Thread] = None
self._response_task: Optional[asyncio.Task] = None
self._response_queue: asyncio.Queue = asyncio.Queue()
self._active_context_id: Optional[str] = None
# Runtime state for the active streaming turn.
self._stream_state: Optional[_SynthesisStreamState] = None
def can_generate_metrics(self) -> bool:
"""Check if this service can generate metrics.
@@ -245,7 +224,6 @@ class NvidiaTTSService(TTSService):
Example::
NvidiaTTSService(
api_key=...,
model_function_map={"function_id": "<UUID>", "model_name": "<model_name>"},
@@ -320,19 +298,19 @@ class NvidiaTTSService(TTSService):
logger.debug(f"Initialized NvidiaTTSService with model: {self._settings.model}")
async def stop(self, frame: EndFrame):
"""Stop the NVIDIA Nemotron Speech TTS service and clean up resources.
"""Stop the NVIDIA Nemotron Speech TTS service.
Args:
frame: EndFrame indicating pipeline stop.
frame: The end frame.
"""
await super().stop(frame)
await self._close_synthesis_stream()
async def cancel(self, frame: CancelFrame):
"""Cancel the NVIDIA Nemotron Speech TTS service operation.
"""Cancel the NVIDIA Nemotron Speech TTS service.
Args:
frame: CancelFrame indicating operation cancellation.
frame: The cancel frame.
"""
await super().cancel(frame)
await self._close_synthesis_stream()
@@ -344,18 +322,20 @@ class NvidiaTTSService(TTSService):
``synthesize_online``. The gRPC stream stays open until a ``None``
sentinel is pushed into the queue.
"""
self._text_queue = queue.Queue()
self._active_context_id = context_id
self._response_queue = asyncio.Queue()
self._synth_thread = threading.Thread(
target=self._synthesis_thread_handler,
daemon=True,
name="nvidia-tts-synth",
state = _SynthesisStreamState(
context_id=context_id,
text_queue=queue.Queue(),
response_queue=asyncio.Queue(),
stop_event=threading.Event(),
)
self._synth_thread.start()
self._response_task = self.create_task(
self._response_consumer(), name="nvidia-tts-response"
self._stream_state = state
logger.debug(f"{self}: starting synthesis stream")
state.synth_task = self.create_task(
self._synth_task_handler(state), name="nvidia-tts-synth"
)
state.response_task = self.create_task(
self._process_responses(state), name="nvidia-tts-response"
)
def _build_base_request(self) -> rtts.SynthesizeSpeechRequest:
@@ -369,105 +349,212 @@ class NvidiaTTSService(TTSService):
voice = self._settings.voice
if voice:
req.voice_name = voice
if self._zero_shot_audio_prompt is not None:
req.zero_shot_data.audio_prompt = self._zero_shot_audio_prompt
req.zero_shot_data.encoding = self._audio_prompt_encoding
req.zero_shot_data.quality = self._settings.quality
if self._custom_dictionary:
req.custom_dictionary = self._custom_dictionary
return req
def _synthesis_thread_handler(self):
"""Run ``SynthesizeOnline`` gRPC stream in a background thread.
def _synthesis_handler(self, state: _SynthesisStreamState):
"""Run ``SynthesizeOnline`` gRPC stream in a blocking call.
Uses a queue-backed generator to feed text chunks into a single
``SynthesizeOnline`` call, enabling Magpie's cross-sentence stitching.
Audio responses are forwarded to the async response queue.
"""
event_loop = self.get_event_loop()
base_req = self._build_base_request()
def request_generator():
while True:
text = self._text_queue.get()
if text is None:
if state.stop_event.is_set():
break
text = state.text_queue.get()
if text is None or state.stop_event.is_set():
break
base_req.text = text
yield base_req
try:
responses = self._service.stub.SynthesizeOnline(
call = self._service.stub.SynthesizeOnline(
request_generator(),
metadata=self._service.auth.get_auth_metadata(),
)
for resp in responses:
asyncio.run_coroutine_threadsafe(
self._response_queue.put(resp), self.get_event_loop()
)
except Exception as e:
logger.error(f"{self} gRPC synthesis stream error: {e}")
asyncio.run_coroutine_threadsafe(self._response_queue.put(e), self.get_event_loop())
finally:
asyncio.run_coroutine_threadsafe(self._response_queue.put(None), self.get_event_loop())
state.rpc_call = call
async def _response_consumer(self):
for resp in call:
if state.stop_event.is_set():
break
asyncio.run_coroutine_threadsafe(state.response_queue.put(resp), event_loop)
except Exception as e:
if not state.stop_event.is_set():
logger.error(f"{self} gRPC synthesis stream error: {e}")
asyncio.run_coroutine_threadsafe(state.response_queue.put(e), event_loop)
finally:
state.rpc_call = None
asyncio.run_coroutine_threadsafe(state.response_queue.put(None), event_loop)
async def _synth_task_handler(self, state: _SynthesisStreamState):
"""Wrap ``_synthesis_handler`` as an asyncio-managed task."""
await asyncio.to_thread(self._synthesis_handler, state)
def _cancel_stream_call(self, state: _SynthesisStreamState):
"""Best-effort cancellation of in-flight gRPC call."""
call = state.rpc_call
if call is not None and hasattr(call, "cancel"):
try:
call.cancel()
except Exception as e:
logger.debug(f"{self}: failed to cancel gRPC stream call: {e}")
async def _process_responses(self, state: _SynthesisStreamState):
"""Consume gRPC responses and append audio to the active audio context."""
while True:
item = await self._response_queue.get()
item = await state.response_queue.get()
if item is None:
break
if isinstance(item, Exception):
await self.push_error(f"{self} synthesis error: {item}")
# Ignore stale exceptions from interrupted streams.
if self._stream_state is state:
await self.push_error(f"{self} synthesis error: {item}")
break
# Stale stream responses must never leak into a newer stream context.
if self._stream_state is not state:
continue
await self.stop_ttfb_metrics()
frame = TTSAudioRawFrame(
audio=item.audio,
sample_rate=self.sample_rate,
num_channels=1,
context_id=self._active_context_id,
context_id=state.context_id,
)
await self.append_to_audio_context(self._active_context_id, frame)
await self.append_to_audio_context(state.context_id, frame)
# Finalize ownership once the stream drains naturally.
if self._stream_state is state and not state.stop_event.is_set():
self._stream_state = None
def _signal_synthesis_close(self, state: _SynthesisStreamState):
"""Signal the active synthesis request generator to close."""
state.text_queue.put(None)
async def _close_synthesis_stream(self):
"""Close the active gRPC synthesis stream.
"""Close the active gRPC synthesis stream gracefully.
Sends a sentinel to end the request generator, waits for the gRPC
thread to finish producing all remaining audio, then lets the
response consumer drain naturally before cleaning up.
Sends a sentinel to end the request generator, waits for the
synthesis task to finish producing all remaining audio, then lets
the response task drain naturally before cleaning up.
"""
if self._text_queue is not None:
self._text_queue.put(None)
state = self._stream_state
if state is None:
return
if self._synth_thread is not None:
await asyncio.to_thread(self._synth_thread.join)
self._synth_thread = None
self._signal_synthesis_close(state)
self._text_queue = None
if self._response_task is not None:
if state.synth_task is not None:
try:
await self._response_task
await state.synth_task
except asyncio.CancelledError:
pass
self._response_task = None
state.synth_task = None
self._active_context_id = None
if state.response_task is not None:
try:
await state.response_task
except asyncio.CancelledError:
pass
state.response_task = None
if self._stream_state is state:
self._stream_state = None
async def _wait_for_synthesis_close_interruptibly(self, state: _SynthesisStreamState):
"""Wait for synthesis close unless interruption preempts this stream."""
while True:
if self._stream_state is not state or state.stop_event.is_set():
# Interruption took ownership of stream shutdown.
return
synth_done = state.synth_task is None or state.synth_task.done()
response_done = state.response_task is None or state.response_task.done()
if synth_done and response_done:
break
# Poll in short intervals to keep this wait interruptible.
await asyncio.sleep(0.05)
if state.synth_task is not None:
try:
await state.synth_task
except asyncio.CancelledError:
pass
state.synth_task = None
if state.response_task is not None:
try:
await state.response_task
except asyncio.CancelledError:
pass
state.response_task = None
if self._stream_state is state:
self._stream_state = None
async def _abort_synthesis_stream(self):
"""Abort the active gRPC synthesis stream immediately.
Cancels the response task first to stop delivering audio, then
drains the text queue and signals the synthesis handler to stop.
Unlike ``_close_synthesis_stream``, pending audio is discarded.
"""
state = self._stream_state
if state is None:
return
state.stop_event.set()
if state.response_task is not None:
await self.cancel_task(state.response_task)
state.response_task = None
while not state.text_queue.empty():
try:
state.text_queue.get_nowait()
except queue.Empty:
break
self._signal_synthesis_close(state)
self._cancel_stream_call(state)
if state.synth_task is not None:
await self.cancel_task(state.synth_task)
state.synth_task = None
if self._stream_state is state:
self._stream_state = None
async def flush_audio(self, context_id: Optional[str] = None):
"""Flush the synthesis stream at the end of an LLM turn.
Sends a sentinel to the gRPC stream, waits for remaining audio,
then delegates to the base class for audio context cleanup.
"""Flush any pending audio and finalize the current context.
Args:
context_id: The audio context to flush.
context_id: The specific context to flush. If None, falls back to the
currently active context.
"""
await self._close_synthesis_stream()
state = self._stream_state
if state is not None:
self._signal_synthesis_close(state)
await self._wait_for_synthesis_close_interruptibly(state)
await super().flush_audio(context_id)
async def on_audio_context_interrupted(self, context_id: str):
"""Handle interruption by closing the active synthesis stream."""
"""Cancel the active gRPC synthesis stream when the bot is interrupted.
Args:
context_id: The ID of the audio context that was interrupted.
"""
await self.stop_all_metrics()
await self._close_synthesis_stream()
await self._abort_synthesis_stream()
await super().on_audio_context_interrupted(context_id)
@staticmethod
def _split_text_into_chunks(text: str) -> list[str]:
@@ -530,9 +617,13 @@ class NvidiaTTSService(TTSService):
logger.debug(f"{self}: Generating TTS [{text}]")
state = self._stream_state
if state is None:
raise RuntimeError("Synthesis stream not started")
for chunk in self._split_text_into_chunks(text):
if any(c.isalnum() for c in chunk):
self._text_queue.put(chunk)
state.text_queue.put(chunk)
await self.start_tts_usage_metrics(text)
yield None