From 8e0a338d969462ce610d302e6876067f7610c709 Mon Sep 17 00:00:00 2001 From: filipi87 Date: Tue, 12 May 2026 11:15:42 -0300 Subject: [PATCH] Nvidia Sagemaker Magpie TTS service --- .../services/nvidia/sagemaker/__init__.py | 0 src/pipecat/services/nvidia/sagemaker/tts.py | 452 ++++++++++++++++++ src/pipecat/utils/text/tts_text_sanitizer.py | 132 +++++ 3 files changed, 584 insertions(+) create mode 100644 src/pipecat/services/nvidia/sagemaker/__init__.py create mode 100644 src/pipecat/services/nvidia/sagemaker/tts.py create mode 100644 src/pipecat/utils/text/tts_text_sanitizer.py diff --git a/src/pipecat/services/nvidia/sagemaker/__init__.py b/src/pipecat/services/nvidia/sagemaker/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/pipecat/services/nvidia/sagemaker/tts.py b/src/pipecat/services/nvidia/sagemaker/tts.py new file mode 100644 index 000000000..e38494b77 --- /dev/null +++ b/src/pipecat/services/nvidia/sagemaker/tts.py @@ -0,0 +1,452 @@ +# +# Copyright (c) 2024-2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""NVIDIA Magpie TTS service backed by an AWS SageMaker endpoint.""" + +import asyncio +import base64 +import json +import os +from collections.abc import AsyncGenerator +from dataclasses import dataclass +from typing import Optional + +import aioboto3 +from loguru import logger + +from pipecat.frames.frames import ( + CancelFrame, + EndFrame, + ErrorFrame, + Frame, + InterruptionFrame, + StartFrame, + TTSAudioRawFrame, +) +from pipecat.processors.frame_processor import FrameDirection +from pipecat.services.aws.sagemaker.bidi_client import SageMakerBidiClient +from pipecat.services.settings import TTSSettings +from pipecat.services.tts_service import InterruptibleTTSService, TTSService +from pipecat.utils.text.tts_text_sanitizer import sanitize_text_for_tts +from pipecat.utils.tracing.service_decorators import traced_tts + + +@dataclass +class NvidiaSageMakerTTSSettings(TTSSettings): + """Settings for NvidiaSageMakerHTTPTTSService. + + Parameters: + voice: NIM voice name (e.g. ``Magpie-Multilingual.EN-US.Aria``). + language: BCP-47 language code passed to NIM (e.g. ``en-US``). + """ + + voice: str = "Magpie-Multilingual.EN-US.Aria" + language: str = "en-US" + + +class NvidiaSageMakerHTTPTTSService(TTSService): + """NVIDIA Magpie TTS service that calls a SageMaker HTTP endpoint. + + Sends each text segment to the wrapper's ``POST /invocations`` endpoint + as a JSON body and streams the raw PCM audio response back to bot + as :class:`TTSAudioRawFrame` frames. + + Example:: + + tts = NvidiaSageMakerHTTPTTSService( + endpoint_name=os.getenv("SAGEMAKER_MAGPIE_ENDPOINT_NAME"), + region=os.getenv("AWS_REGION", "us-west-2"), + settings=NvidiaSageMakerHTTPTTSService.Settings( + voice="Magpie-Multilingual.EN-US.Aria", + language="en-US", + ), + ) + """ + + Settings = NvidiaSageMakerTTSSettings + + def __init__( + self, + *, + endpoint_name: str, + region: str = "us-west-2", + sample_rate: int | None = None, + settings: NvidiaSageMakerTTSSettings | None = None, + **kwargs, + ): + """Initialize the SageMaker HTTP TTS service. + + Args: + endpoint_name: Name of the deployed SageMaker endpoint. + region: AWS region where the endpoint lives. + sample_rate: Output sample rate in Hz. Defaults to bot's pipeline rate. + params: Deprecated — use ``settings`` instead. + settings: Runtime-updatable settings (voice, language). + **kwargs: Forwarded to :class:`TTSService`. + """ + default_settings = self.Settings(model="magpie") + + if settings is not None: + default_settings.apply_update(settings) + + super().__init__( + sample_rate=sample_rate, + push_start_frame=True, + push_stop_frames=True, + settings=default_settings, + **kwargs, + ) + + self._endpoint_name = endpoint_name + self._region = region + self._client = None + self._client_ctx = None + + def can_generate_metrics(self) -> bool: + return True + + # ── Lifecycle ───────────────────────────────────────────────────────────── + + async def start(self, frame: StartFrame): + await super().start(frame) + session = aioboto3.Session( + aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"), + aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"), + region_name=self._region, + ) + self._client_ctx = session.client("sagemaker-runtime") + self._client = await self._client_ctx.__aenter__() + logger.debug(f"{self}: connected to SageMaker endpoint '{self._endpoint_name}'") + + async def _close_client(self): + if self._client_ctx is not None: + try: + await self._client_ctx.__aexit__(None, None, None) + except Exception as e: + logger.warning(f"{self}: error closing SageMaker client: {e}") + self._client_ctx = None + self._client = None + + async def stop(self, frame: EndFrame): + await super().stop(frame) + await self._close_client() + + async def cancel(self, frame: CancelFrame): + await super().cancel(frame) + await self._close_client() + + # ── Synthesis ───────────────────────────────────────────────────────────── + + @traced_tts + async def run_tts(self, text: str, context_id: str) -> AsyncGenerator[Frame, None]: + """Synthesize text via SageMaker and yield a single PCM audio frame. + + Args: + text: The text to synthesize. + context_id: Pipecat audio context identifier. + + Yields: + :class:`TTSAudioRawFrame` chunks of signed 16-bit mono PCM. + """ + logger.debug(f"{self}: Generating TTS [{text}]") + + text = sanitize_text_for_tts(text) + if not text: + return + + try: + body = json.dumps( + { + "text": text, + "voice_name": self._settings.voice, + "language_code": self._settings.language, + "sample_rate_hz": self.sample_rate, + } + ) + + response = await self._client.invoke_endpoint( + EndpointName=self._endpoint_name, + ContentType="application/json", + Accept="application/octet-stream", + Body=body, + ) + + if "Body" not in response: + yield ErrorFrame(error="SageMaker TTS returned no audio stream") + return + + async for chunk in response["Body"].iter_chunks(chunk_size=self.chunk_size): + if chunk: + yield TTSAudioRawFrame( + audio=chunk, + sample_rate=self.sample_rate, + num_channels=1, + context_id=context_id, + ) + except Exception as e: + logger.error(f"{self}: SageMaker TTS error: {e}") + yield ErrorFrame(error=f"SageMaker TTS error: {e}") + finally: + await self.stop_ttfb_metrics() + + +@dataclass +class NvidiaSageMakerWSTTSSettings(TTSSettings): + """Settings for NvidiaSageMakerWebsocketTTSService. + + Parameters: + voice: NIM voice name (e.g. ``Magpie-Multilingual.EN-US.Aria``). + language: BCP-47 language code passed to NIM (e.g. ``en-US``). + """ + + voice: str = "Magpie-Multilingual.EN-US.Aria" + language: str = "en-US" + + +class NvidiaSageMakerWebsocketTTSService(InterruptibleTTSService): + """NVIDIA Magpie TTS service using SageMaker bidirectional streaming. + + Maintains a persistent HTTP/2 bidi-stream session to the SageMaker endpoint + for the lifetime of the pipeline. Each text segment is sent as NIM realtime + events; audio chunks arrive asynchronously and are pushed as + :class:`TTSAudioRawFrame` frames. + + Example:: + + tts = NvidiaSageMakerWebsocketTTSService( + endpoint_name=os.getenv("SAGEMAKER_MAGPIE_ENDPOINT_NAME"), + region=os.getenv("AWS_REGION", "us-west-2"), + settings=NvidiaSageMakerWebsocketTTSService.Settings( + voice="Magpie-Multilingual.EN-US.Aria", + language="en-US", + ), + ) + """ + + Settings = NvidiaSageMakerWSTTSSettings + + def __init__( + self, + *, + endpoint_name: str, + region: str = "us-west-2", + sample_rate: int | None = None, + settings: NvidiaSageMakerWSTTSSettings | None = None, + **kwargs, + ): + default_settings = self.Settings(model="magpie") + + if settings is not None: + default_settings.apply_update(settings) + + super().__init__( + sample_rate=sample_rate, + push_start_frame=True, + push_stop_frames=True, + pause_frame_processing=True, + append_trailing_space=True, + settings=default_settings, + **kwargs, + ) + + self._endpoint_name = endpoint_name + self._region = region + self._client: SageMakerBidiClient | None = None + self._receive_task = None + self._speech_completed_event = asyncio.Event() + + def can_generate_metrics(self) -> bool: + return True + + # ── Lifecycle ───────────────────────────────────────────────────────────── + + async def start(self, frame: StartFrame): + await super().start(frame) + await self._connect() + + async def stop(self, frame: EndFrame): + await super().stop(frame) + await self._disconnect() + + async def cancel(self, frame: CancelFrame): + await super().cancel(frame) + await self._disconnect() + + # ── Connection management (WebsocketService abstract interface) ──────────── + + async def _connect(self): + await super()._connect() + await self._connect_websocket() + if self._client and self._client.is_active and not self._receive_task: + self._receive_task = self.create_task(self._receive_task_handler(self._report_error)) + + async def _disconnect(self): + await super()._disconnect() + if self._receive_task: + await self.cancel_task(self._receive_task) + self._receive_task = None + await self._disconnect_websocket() + + async def _connect_websocket(self): + if self._client and self._client.is_active: + return + + logger.debug( + f"{self}: connecting to SageMaker bidi-stream endpoint '{self._endpoint_name}'" + ) + try: + self._client = SageMakerBidiClient( + endpoint_name=self._endpoint_name, + region=self._region, + model_query_string=None, + model_invocation_path=None, + ) + await self._client.start_session() + await self._send_session_config() + logger.debug(f"{self}: connected") + await self._call_event_handler("on_connected") + except Exception as e: + logger.error(f"{self}: connection error: {e}") + self._client = None + await self._call_event_handler("on_connection_error", f"{e}") + + async def _disconnect_websocket(self): + try: + if self._client and self._client.is_active: + logger.debug(f"{self}: disconnecting") + try: + await self._client.send_json({"type": "session.end"}) + except Exception as e: + logger.warning(f"{self}: error sending session.end: {e}") + await self._client.close_session() + logger.debug(f"{self}: disconnected") + except Exception as e: + logger.warning(f"{self}: error during disconnect: {e}") + finally: + self._client = None + await self._call_event_handler("on_disconnected") + + async def _verify_connection(self): + active = self._client and self._client.is_active + logger.info(f"{self}: verifying if websocket connection is active {active}") + return active + + async def _handle_interruption(self, frame: InterruptionFrame, direction: FrameDirection): + if self._bot_speaking: + logger.debug( + f"{self}: interruption detected, sending input_text.done and waiting for speech.completed" + ) + self._disconnecting = True + self._speech_completed_event.clear() + try: + await self._client.send_json({"type": "input_text.done"}) + await asyncio.wait_for(self._speech_completed_event.wait(), timeout=5.0) + except TimeoutError: + logger.warning(f"{self}: timed out waiting for conversation.item.speech.completed") + await super()._handle_interruption(frame, direction) + + async def _receive_messages(self): + """Receive NIM JSON events and push audio frames.""" + while self._client and self._client.is_active and not self._disconnecting: + result = await self._client.receive_response() + + if self._disconnecting: + self._speech_completed_event.set() + + if result is None: + break + + if not (hasattr(result, "value") and hasattr(result.value, "bytes_")): + continue + + payload = result.value.bytes_ + if not payload: + continue + + context_id = self.get_active_audio_context_id() + + try: + msg = json.loads(payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + # Unexpected binary frame — treat as raw PCM + await self.push_frame( + TTSAudioRawFrame( + audio=payload, + sample_rate=self.sample_rate, + num_channels=1, + context_id=context_id, + ) + ) + continue + + event_type = msg.get("type", "") + + if event_type != "conversation.item.speech.data": + logger.debug(f"{self}: received event: {event_type}") + + if event_type == "conversation.item.speech.data": + chunk_b64 = msg.get("audio", "") + if chunk_b64: + await self.stop_ttfb_metrics() + await self.push_frame( + TTSAudioRawFrame( + audio=base64.b64decode(chunk_b64), + sample_rate=self.sample_rate, + num_channels=1, + context_id=context_id, + ) + ) + elif event_type == "error": + await self.push_error(error_msg=f"NIM error: {msg.get('message', msg)}") + # In case of error we need to reconnect, otherwise we are not going to receive audio from the TTS service anymore + break + elif event_type == "conversation.item.speech.completed": + # Need to reconnect to reset the synthesis state and be able to synthesize new text + break + + # synthesize_session.updated, input_text.committed, etc. are ignored. + + async def _send_session_config(self): + """Send synthesize_session.update to configure voice and audio params.""" + logger.debug(f"{self}: sending session config, sample_rate={self.sample_rate}") + await self._client.send_json( + { + "type": "synthesize_session.update", + "session": { + "input_text_synthesis": { + "voice_name": self._settings.voice, + "language_code": self._settings.language, + }, + "output_audio_params": { + "sample_rate_hz": self.sample_rate, + }, + }, + } + ) + + # ── Synthesis ───────────────────────────────────────────────────────────── + + @traced_tts + async def run_tts(self, text: str, context_id: str) -> AsyncGenerator[Frame, None]: + """Send text to NIM; audio arrives asynchronously via _receive_messages.""" + logger.debug(f"{self}: Generating TTS [{text}]") + + text = sanitize_text_for_tts(text) + + logger.debug(f"{self}: sanitized text: {text}") + if not text: + return + + try: + if not self._client or not self._client.is_active: + await self._connect() + + await self._client.send_json({"type": "input_text.append", "text": text}) + await self._client.send_json({"type": "input_text.commit"}) + yield None + except Exception as e: + logger.error(f"{self}: TTS error: {e}") + yield ErrorFrame(error=f"NvidiaSageMakerWebsocketTTSService error: {e}") diff --git a/src/pipecat/utils/text/tts_text_sanitizer.py b/src/pipecat/utils/text/tts_text_sanitizer.py new file mode 100644 index 000000000..b5e78e5ce --- /dev/null +++ b/src/pipecat/utils/text/tts_text_sanitizer.py @@ -0,0 +1,132 @@ +# +# Copyright (c) 2024-2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""Utility for stripping non-speakable characters and markdown formatting from +text before it is sent to a TTS service. + +Both NvidiaSageMakerHTTPTTSService and NvidiaSageMakerWebsocketTTSService +use :func:`sanitize_text_for_tts` so the logic lives in one place. +""" + +import re + +# --------------------------------------------------------------------------- +# Emoji / symbol ranges +# --------------------------------------------------------------------------- + +_EMOJI_PATTERN = re.compile( + "[" + "\U0001f600-\U0001f64f" # Emoticons + "\U0001f300-\U0001f5ff" # Misc Symbols and Pictographs + "\U0001f680-\U0001f6ff" # Transport and Map + "\U0001f700-\U0001f77f" # Alchemical Symbols + "\U0001f780-\U0001f7ff" # Geometric Shapes Extended + "\U0001f800-\U0001f8ff" # Supplemental Arrows-C + "\U0001f900-\U0001f9ff" # Supplemental Symbols and Pictographs + "\U0001fa00-\U0001fa6f" # Chess Symbols + "\U0001fa70-\U0001faff" # Symbols and Pictographs Extended-A + "\U00002702-\U000027b0" # Dingbats + "\U0001f1e0-\U0001f1ff" # Flags (iOS) + "]+", + flags=re.UNICODE, +) + + +def sanitize_text_for_tts(text: str) -> str: + """Remove emojis and markdown formatting that should not be spoken aloud. + + Transformations applied (in order): + 1. Fenced code blocks (``` ... ```) → removed entirely + 2. Markdown headers (# / ## / …) → header text kept, # stripped + 3. Horizontal rules (--- / *** / ___) → removed + 4. Table separator rows (|---|---|) → removed + 5. Table data rows (| a | b |) → cells joined with commas + 6. Bold / italic markers (**x**, *x*, __x__, _x_) → text kept, markers stripped + 7. Blockquote markers (> …) → marker stripped, text kept + 8. Inline code backticks (`x`) → backticks stripped, text kept + 9. Emojis → removed + 10. Curly quotes → straight quotes + Em/en dashes → comma (natural pause, not a spoken symbol) + Separator hyphens ( - ) → comma + Unordered list bullets (^- , ^* ) → removed + Remaining bare * and _ → removed + Other non-speakable symbols → removed + 11. Collapse extra whitespace + + Args: + text: Raw text, potentially containing markdown and/or emoji. + + Returns: + Plain text suitable for speech synthesis. + """ + # 1. Fenced code blocks + text = re.sub(r"```[\s\S]*?```", "", text) + + # 2. Markdown headers (# Heading → Heading) + text = re.sub(r"^#{1,6}\s+", "", text, flags=re.MULTILINE) + + # 3. Horizontal rules (---, ***, ___ on their own line) + text = re.sub(r"^\s*[-*_]{3,}\s*$", "", text, flags=re.MULTILINE) + + # 4. Table separator rows |---|:---:|---| + text = re.sub(r"^\s*\|[\s\-|:]+\|\s*$", "", text, flags=re.MULTILINE) + + # 5. Table data rows | cell | cell | → cell, cell + def _table_row_to_csv(m: re.Match) -> str: + cells = [c.strip() for c in m.group(1).split("|")] + return ", ".join(c for c in cells if c) + + text = re.sub(r"^\s*\|(.+)\|\s*$", _table_row_to_csv, text, flags=re.MULTILINE) + + # 6. Bold / italic (**x**, *x*, __x__, _x_) + # Handle triple before double before single to avoid partial matches. + text = re.sub(r"\*{3}([^*]+)\*{3}", r"\1", text) + text = re.sub(r"\*{2}([^*]+)\*{2}", r"\1", text) + text = re.sub(r"\*([^*\s][^*]*[^*\s]|\S)\*", r"\1", text) + text = re.sub(r"_{3}([^_]+)_{3}", r"\1", text) + text = re.sub(r"_{2}([^_]+)_{2}", r"\1", text) + text = re.sub(r"_([^_\s][^_]*[^_\s]|\S)_", r"\1", text) + + # 7. Blockquote markers + text = re.sub(r"^\s*>\s*", "", text, flags=re.MULTILINE) + + # 8. Inline code backticks + text = re.sub(r"`([^`]*)`", r"\1", text) + + # 9. Emojis + text = _EMOJI_PATTERN.sub("", text) + + # 10. Typographic and non-speakable characters + + # Curly quotes → straight equivalents (speakable) + text = text.replace("\u2018", "'") # LEFT SINGLE QUOTATION MARK + text = text.replace("\u2019", "'") # RIGHT SINGLE QUOTATION MARK + text = text.replace("\u201c", '"') # LEFT DOUBLE QUOTATION MARK + text = text.replace("\u201d", '"') # RIGHT DOUBLE QUOTATION MARK + + # Em/en dashes → comma (they mark a pause, not a spoken symbol) + text = text.replace("\u2014", ", ") # EM DASH + text = text.replace("\u2013", ", ") # EN DASH + + # Hyphen used as a separator ( - ) → comma; keep word-hyphens (e.g. "well-known") + text = re.sub(r"(?<= )- | -(?= )", ", ", text) + + # Unordered list bullets at line start (not caught by step 3) + text = re.sub(r"^[\s]*[-*]\s+", "", text, flags=re.MULTILINE) + + # Remaining bare * and _ (e.g. orphaned markers) + text = text.replace("*", "") + text = text.replace("_", " ") + + # Other symbols that TTS engines typically misread or glitch on + text = re.sub(r"[\\|<>{}[\]~^=+#@]", "", text) + + # 11. Collapse whitespace + text = re.sub(r"\n{3,}", "\n\n", text) # no more than one blank line + text = re.sub(r"[ \t]+", " ", text) # collapse spaces/tabs + # text = text.strip() + + return text