diff --git a/CHANGELOG.md b/CHANGELOG.md index 5954e6baf..1f4811184 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -119,6 +119,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 case there's no need to push audio to the rest of the pipeline, but this is not a very common case. +- Added `RivaSegmentedSTTService`, which allows Riva offline/batch models, such + as to be "canary-1b-asr" used in Pipecat. + ### Deprecated - Function calls with parameters @@ -134,6 +137,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `TransportParams.vad_audio_passthrough` parameter is now deprecated, use `TransportParams.audio_in_passthrough` instead. +- `ParakeetSTTService` is now deprecated, use `RivaSTTService` instead, which uses + the model "parakeet-ctc-1.1b-asr" by default. + +- `FastPitchTTSService` is now deprecated, use `RivaTTSService` instead, which uses + the model "magpie-tts-multilingual" by default. + ### Fixed - Fixed an issue where `OpenAIRealtimeBetaLLMService` would add two assistant diff --git a/examples/foundational/07r-interruptible-riva-nim.py b/examples/foundational/07r-interruptible-riva-nim.py index 915beda51..ddb80181c 100644 --- a/examples/foundational/07r-interruptible-riva-nim.py +++ b/examples/foundational/07r-interruptible-riva-nim.py @@ -16,8 +16,12 @@ from pipecat.pipeline.runner import PipelineRunner from pipecat.pipeline.task import PipelineParams, PipelineTask from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext from pipecat.services.nim.llm import NimLLMService -from pipecat.services.riva.stt import ParakeetSTTService -from pipecat.services.riva.tts import FastPitchTTSService +from pipecat.services.riva.stt import ( + ParakeetSTTService, + RivaSegmentedSTTService, + RivaSTTService, +) +from pipecat.services.riva.tts import FastPitchTTSService, RivaTTSService from pipecat.transports.base_transport import TransportParams from pipecat.transports.network.small_webrtc import SmallWebRTCTransport from pipecat.transports.network.webrtc_connection import SmallWebRTCConnection @@ -37,11 +41,11 @@ async def run_bot(webrtc_connection: SmallWebRTCConnection, _: argparse.Namespac ), ) - stt = ParakeetSTTService(api_key=os.getenv("NVIDIA_API_KEY")) + stt = RivaSTTService(api_key=os.getenv("NVIDIA_API_KEY")) llm = NimLLMService(api_key=os.getenv("NVIDIA_API_KEY"), model="meta/llama-3.1-405b-instruct") - tts = FastPitchTTSService(api_key=os.getenv("NVIDIA_API_KEY")) + tts = RivaTTSService(api_key=os.getenv("NVIDIA_API_KEY")) messages = [ { diff --git a/pyproject.toml b/pyproject.toml index 6d1cf840c..30d72f161 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,7 +78,7 @@ perplexity = [] playht = [ "pyht~=0.1.12", "websockets~=13.1" ] qwen = [] rime = [ "websockets~=13.1" ] -riva = [ "nvidia-riva-client~=2.19.0" ] +riva = [ "nvidia-riva-client~=2.19.1" ] sentry = [ "sentry-sdk~=2.23.1" ] local-smart-turn = [ "coremltools>=8.0", "transformers", "torch==2.5.0", "torchaudio==2.5.0" ] remote-smart-turn = [] diff --git a/src/pipecat/services/riva/stt.py b/src/pipecat/services/riva/stt.py index 6328bcb65..2b263f03e 100644 --- a/src/pipecat/services/riva/stt.py +++ b/src/pipecat/services/riva/stt.py @@ -5,7 +5,7 @@ # import asyncio -from typing import AsyncGenerator, Optional +from typing import AsyncGenerator, List, Mapping, Optional from loguru import logger from pydantic import BaseModel @@ -13,12 +13,13 @@ from pydantic import BaseModel from pipecat.frames.frames import ( CancelFrame, EndFrame, + ErrorFrame, Frame, InterimTranscriptionFrame, StartFrame, TranscriptionFrame, ) -from pipecat.services.stt_service import STTService +from pipecat.services.stt_service import SegmentedSTTService, STTService from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 @@ -31,7 +32,59 @@ except ModuleNotFoundError as e: raise Exception(f"Missing module: {e}") -class ParakeetSTTService(STTService): +def language_to_riva_language(language: Language) -> Optional[str]: + """Maps Language enum to Riva ASR language codes. + + Source: + https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-riva-build-table.html?highlight=fr%20fr + + Args: + language: Language enum value. + + Returns: + Optional[str]: Riva language code or None if not supported. + """ + language_map = { + # Arabic + Language.AR: "ar-AR", + # English + Language.EN: "en-US", # Default to US + Language.EN_US: "en-US", + Language.EN_GB: "en-GB", + # French + Language.FR: "fr-FR", + Language.FR_FR: "fr-FR", + # German + Language.DE: "de-DE", + Language.DE_DE: "de-DE", + # Hindi + Language.HI: "hi-IN", + Language.HI_IN: "hi-IN", + # Italian + Language.IT: "it-IT", + Language.IT_IT: "it-IT", + # Japanese + Language.JA: "ja-JP", + Language.JA_JP: "ja-JP", + # Korean + Language.KO: "ko-KR", + Language.KO_KR: "ko-KR", + # Portuguese + Language.PT: "pt-BR", # Default to Brazilian + Language.PT_BR: "pt-BR", + # Russian + Language.RU: "ru-RU", + Language.RU_RU: "ru-RU", + # Spanish + Language.ES: "es-ES", # Default to Spain + Language.ES_ES: "es-ES", + Language.ES_US: "es-US", # US Spanish + } + + return language_map.get(language) + + +class RivaSTTService(STTService): class InputParams(BaseModel): language: Optional[Language] = Language.EN_US @@ -40,7 +93,10 @@ class ParakeetSTTService(STTService): *, api_key: str, server: str = "grpc.nvcf.nvidia.com:443", - function_id: str = "1598d209-5e27-4d3c-8079-4751568b1081", + model_function_map: Mapping[str, str] = { + "function_id": "1598d209-5e27-4d3c-8079-4751568b1081", + "model_name": "parakeet-ctc-1.1b-asr", + }, sample_rate: Optional[int] = None, params: InputParams = InputParams(), **kwargs, @@ -48,7 +104,7 @@ class ParakeetSTTService(STTService): super().__init__(sample_rate=sample_rate, **kwargs) self._api_key = api_key self._profanity_filter = False - self._automatic_punctuation = False + self._automatic_punctuation = True self._no_verbatim_transcripts = False self._language_code = params.language self._boosted_lm_words = None @@ -60,11 +116,12 @@ class ParakeetSTTService(STTService): self._stop_history_eou = -1 self._stop_threshold_eou = -1.0 self._custom_configuration = "" + self._function_id = model_function_map.get("function_id") - self.set_model_name("parakeet-ctc-1.1b-asr") + self.set_model_name(model_function_map.get("model_name")) metadata = [ - ["function-id", function_id], + ["function-id", self._function_id], ["authorization", f"Bearer {api_key}"], ] auth = riva.client.Auth(None, True, server, metadata) @@ -79,6 +136,13 @@ class ParakeetSTTService(STTService): def can_generate_metrics(self) -> bool: return False + async def set_model(self, model: str): + logger.warning(f"Cannot set model after initialization. Set model and function id like so:") + example = {"function_id": "", "model_name": ""} + logger.warning( + f"{self.__class__.__name__}(api_key=, model_function_map={example})" + ) + async def start(self, frame: StartFrame): await super().start(frame) @@ -196,3 +260,262 @@ class ParakeetSTTService(STTService): def __iter__(self): return self + + +class RivaSegmentedSTTService(SegmentedSTTService): + """Speech-to-text service using NVIDIA Riva's offline/batch models. + + By default, his service uses NVIDIA's Riva Canary ASR API to perform speech-to-text + transcription on audio segments. It inherits from SegmentedSTTService to handle + audio buffering and speech detection. + + Args: + api_key: NVIDIA API key for authentication + server: Riva server address (defaults to NVIDIA Cloud Function endpoint) + model_function_map: Mapping of model name and its corresponding NVIDIA Cloud Function ID + sample_rate: Audio sample rate in Hz. If not provided, uses the pipeline's rate + params: Additional configuration parameters for Riva + **kwargs: Additional arguments passed to SegmentedSTTService + """ + + class InputParams(BaseModel): + language: Optional[Language] = Language.EN_US + profanity_filter: bool = False + automatic_punctuation: bool = True + verbatim_transcripts: bool = False + boosted_lm_words: Optional[List[str]] = None + boosted_lm_score: float = 4.0 + + def __init__( + self, + *, + api_key: str, + server: str = "grpc.nvcf.nvidia.com:443", + model_function_map: Mapping[str, str] = { + "function_id": "ee8dc628-76de-4acc-8595-1836e7e857bd", + "model_name": "canary-1b-asr", + }, + sample_rate: Optional[int] = None, + params: InputParams = InputParams(), + **kwargs, + ): + super().__init__(sample_rate=sample_rate, **kwargs) + + # Set model name + self.set_model_name(model_function_map.get("model_name")) + + # Initialize Riva settings + self._api_key = api_key + self._server = server + self._function_id = model_function_map.get("function_id") + self._model_name = model_function_map.get("model_name") + + # Store the language as a Language enum and as a string + self._language_enum = params.language or Language.EN_US + self._language = self.language_to_service_language(self._language_enum) or "en-US" + + # Configure transcription parameters + self._profanity_filter = params.profanity_filter + self._automatic_punctuation = params.automatic_punctuation + self._verbatim_transcripts = params.verbatim_transcripts + self._boosted_lm_words = params.boosted_lm_words + self._boosted_lm_score = params.boosted_lm_score + + # Voice activity detection thresholds (use Riva defaults) + self._start_history = -1 + self._start_threshold = -1.0 + self._stop_history = -1 + self._stop_threshold = -1.0 + self._stop_history_eou = -1 + self._stop_threshold_eou = -1.0 + self._custom_configuration = "" + + # Create Riva client + self._config = None + self._asr_service = None + self._settings = {"language": self._language_enum} + + def language_to_service_language(self, language: Language) -> Optional[str]: + """Convert pipecat Language enum to Riva's language code.""" + return language_to_riva_language(language) + + def _initialize_client(self): + """Initialize the Riva ASR client with authentication metadata.""" + if self._asr_service is not None: + return + + # Set up authentication metadata for NVIDIA Cloud Functions + metadata = [ + ["function-id", self._function_id], + ["authorization", f"Bearer {self._api_key}"], + ] + + # Create authenticated client + auth = riva.client.Auth(None, True, self._server, metadata) + self._asr_service = riva.client.ASRService(auth) + + logger.info(f"Initialized RivaSegmentedSTTService with model: {self.model_name}") + + def _create_recognition_config(self): + """Create the Riva ASR recognition configuration.""" + # Create base configuration + config = riva.client.RecognitionConfig( + language_code=self._language, # Now using the string, not a tuple + max_alternatives=1, + profanity_filter=self._profanity_filter, + enable_automatic_punctuation=self._automatic_punctuation, + verbatim_transcripts=self._verbatim_transcripts, + ) + + # Add word boosting if specified + if self._boosted_lm_words: + riva.client.add_word_boosting_to_config( + config, self._boosted_lm_words, self._boosted_lm_score + ) + + # Add voice activity detection parameters + riva.client.add_endpoint_parameters_to_config( + config, + self._start_history, + self._start_threshold, + self._stop_history, + self._stop_history_eou, + self._stop_threshold, + self._stop_threshold_eou, + ) + + # Add any custom configuration + if self._custom_configuration: + riva.client.add_custom_configuration_to_config(config, self._custom_configuration) + + return config + + def can_generate_metrics(self) -> bool: + """Indicates whether this service can generate processing metrics.""" + return True + + async def set_model(self, model: str): + logger.warning(f"Cannot set model after initialization. Set model and function id like so:") + example = {"function_id": "", "model_name": ""} + logger.warning( + f"{self.__class__.__name__}(api_key=, model_function_map={example})" + ) + + async def start(self, frame: StartFrame): + """Initialize the service when the pipeline starts.""" + await super().start(frame) + self._initialize_client() + self._config = self._create_recognition_config() + + async def set_language(self, language: Language): + """Set the language for the STT service.""" + logger.info(f"Switching STT language to: [{language}]") + self._language_enum = language + self._language = self.language_to_service_language(language) or "en-US" + self._settings["language"] = language + + # Update configuration with new language + if self._config: + self._config.language_code = self._language + + async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]: + """Transcribe an audio segment. + + Args: + audio: Raw audio bytes in WAV format (already converted by base class). + + Yields: + Frame: TranscriptionFrame containing the transcribed text. + """ + try: + await self.start_processing_metrics() + await self.start_ttfb_metrics() + + # Make sure the client is initialized + if self._asr_service is None: + self._initialize_client() + + # Make sure the config is created + if self._config is None: + self._config = self._create_recognition_config() + + # Type assertion to satisfy the IDE + assert self._asr_service is not None, "ASR service not initialized" + assert self._config is not None, "Recognition config not created" + + # Process audio with Riva ASR - explicitly request non-future response + raw_response = self._asr_service.offline_recognize(audio, self._config, future=False) + + await self.stop_ttfb_metrics() + await self.stop_processing_metrics() + + # Process the response - handle different possible return types + try: + # If it's a future-like object, get the result + if hasattr(raw_response, "result"): + response = raw_response.result() + else: + response = raw_response + + # Process transcription results + transcription_found = False + + # Now we can safely check results + # Type hint for the IDE + results = getattr(response, "results", []) + + for result in results: + alternatives = getattr(result, "alternatives", []) + if alternatives: + text = alternatives[0].transcript.strip() + if text: + logger.debug(f"Transcription: [{text}]") + yield TranscriptionFrame( + text, "", time_now_iso8601(), self._language_enum + ) + transcription_found = True + + if not transcription_found: + logger.debug("No transcription results found in Riva response") + + except AttributeError as ae: + logger.error(f"Unexpected response structure from Riva: {ae}") + yield ErrorFrame(f"Unexpected Riva response format: {str(ae)}") + + except Exception as e: + logger.exception(f"Riva Canary ASR error: {e}") + yield ErrorFrame(f"Riva Canary ASR error: {str(e)}") + + +class ParakeetSTTService(RivaSTTService): + """Deprecated: Use RivaSTTService instead.""" + + def __init__( + self, + *, + api_key: str, + server: str = "grpc.nvcf.nvidia.com:443", + model_function_map: Mapping[str, str] = { + "function_id": "1598d209-5e27-4d3c-8079-4751568b1081", + "model_name": "parakeet-ctc-1.1b-asr", + }, + sample_rate: Optional[int] = None, + params: RivaSTTService.InputParams = RivaSTTService.InputParams(), # Use parent class's type + **kwargs, + ): + super().__init__( + api_key=api_key, + server=server, + model_function_map=model_function_map, + sample_rate=sample_rate, + params=params, + **kwargs, + ) + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("always") + warnings.warn( + "`ParakeetSTTService` is deprecated, use `RivaSTTService` instead.", + DeprecationWarning, + ) diff --git a/src/pipecat/services/riva/tts.py b/src/pipecat/services/riva/tts.py index 0fd0c3ce0..4fc8b8489 100644 --- a/src/pipecat/services/riva/tts.py +++ b/src/pipecat/services/riva/tts.py @@ -5,7 +5,11 @@ # import asyncio -from typing import AsyncGenerator, Optional +import os +from typing import AsyncGenerator, Mapping, Optional + +# Suppress gRPC fork warnings +os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" from loguru import logger from pydantic import BaseModel @@ -27,10 +31,10 @@ except ModuleNotFoundError as e: logger.error("In order to use NVIDIA Riva TTS, you need to `pip install pipecat-ai[riva]`.") raise Exception(f"Missing module: {e}") -FASTPITCH_TIMEOUT_SECS = 5 +RIVA_TTS_TIMEOUT_SECS = 5 -class FastPitchTTSService(TTSService): +class RivaTTSService(TTSService): class InputParams(BaseModel): language: Optional[Language] = Language.EN_US quality: Optional[int] = 20 @@ -38,11 +42,14 @@ class FastPitchTTSService(TTSService): def __init__( self, *, - api_key: str, + api_key: str = None, server: str = "grpc.nvcf.nvidia.com:443", - voice_id: str = "English-US.Female-1", + voice_id: str = "Magpie-Multilingual.EN-US.Ray", sample_rate: Optional[int] = None, - function_id: str = "0149dedb-2be8-4195-b9a0-e57e0e14f972", + model_function_map: Mapping[str, str] = { + "function_id": "877104f7-e885-42b9-8de8-f6e4c6303969", + "model_name": "magpie-tts-multilingual", + }, params: InputParams = InputParams(), **kwargs, ): @@ -51,12 +58,13 @@ class FastPitchTTSService(TTSService): self._voice_id = voice_id self._language_code = params.language self._quality = params.quality + self._function_id = model_function_map.get("function_id") - self.set_model_name("fastpitch-hifigan-tts") + self.set_model_name(model_function_map.get("model_name")) self.set_voice(voice_id) metadata = [ - ["function-id", function_id], + ["function-id", self._function_id], ["authorization", f"Bearer {api_key}"], ] auth = riva.client.Auth(None, True, server, metadata) @@ -68,6 +76,13 @@ class FastPitchTTSService(TTSService): riva.client.proto.riva_tts_pb2.RivaSynthesisConfigRequest() ) + async def set_model(self, model: str): + logger.warning(f"Cannot set model after initialization. Set model and function id like so:") + example = {"function_id": "", "model_name": ""} + logger.warning( + f"{self.__class__.__name__}(api_key=, model_function_map={example})" + ) + async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]: def read_audio_responses(queue: asyncio.Queue): def add_response(r): @@ -100,7 +115,7 @@ class FastPitchTTSService(TTSService): await asyncio.to_thread(read_audio_responses, queue) # Wait for the thread to start. - resp = await asyncio.wait_for(queue.get(), FASTPITCH_TIMEOUT_SECS) + resp = await asyncio.wait_for(queue.get(), RIVA_TTS_TIMEOUT_SECS) while resp: await self.stop_ttfb_metrics() frame = TTSAudioRawFrame( @@ -109,9 +124,46 @@ class FastPitchTTSService(TTSService): num_channels=1, ) yield frame - resp = await asyncio.wait_for(queue.get(), FASTPITCH_TIMEOUT_SECS) + resp = await asyncio.wait_for(queue.get(), RIVA_TTS_TIMEOUT_SECS) except asyncio.TimeoutError: logger.error(f"{self} timeout waiting for audio response") await self.start_tts_usage_metrics(text) yield TTSStoppedFrame() + + +class FastPitchTTSService(RivaTTSService): + class InputParams(BaseModel): + language: Optional[Language] = Language.EN_US + quality: Optional[int] = 20 + + def __init__( + self, + *, + api_key: str = None, + server: str = "grpc.nvcf.nvidia.com:443", + voice_id: str = "English-US.Female-1", + sample_rate: Optional[int] = None, + model_function_map: Mapping[str, str] = { + "function_id": "0149dedb-2be8-4195-b9a0-e57e0e14f972", + "model_name": "fastpitch-hifigan-tts", + }, + params: InputParams = InputParams(), + **kwargs, + ): + super().__init__( + api_key=api_key, + voice_id=voice_id, + sample_rate=sample_rate, + model_function_map=model_function_map, + params=params, + **kwargs, + ) + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("always") + warnings.warn( + "`FastPitchTTSService` is deprecated, use `RivaTTSService` instead.", + DeprecationWarning, + )