1095 lines
46 KiB
Python
1095 lines
46 KiB
Python
import argparse
|
|
import json
|
|
import logging
|
|
import os
|
|
import signal
|
|
import sys
|
|
from copy import copy, deepcopy
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from queue import Queue
|
|
from sys import platform
|
|
from threading import Event
|
|
from types import FrameType
|
|
from typing import Any, Optional
|
|
|
|
import nltk
|
|
import torch
|
|
from openai.types.realtime import RealtimeSessionCreateRequest
|
|
from rich.console import Console
|
|
from transformers import HfArgumentParser
|
|
|
|
from speech_to_speech.api.openai_realtime.pipeline_unit import PipelineUnit
|
|
from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig
|
|
from speech_to_speech.arguments_classes.chat_completions_language_model_arguments import (
|
|
ChatCompletionsLanguageModelHandlerArguments,
|
|
)
|
|
from speech_to_speech.arguments_classes.chat_tts_arguments import ChatTTSHandlerArguments
|
|
from speech_to_speech.arguments_classes.facebookmms_tts_arguments import FacebookMMSTTSHandlerArguments
|
|
from speech_to_speech.arguments_classes.faster_whisper_stt_arguments import (
|
|
FasterWhisperSTTHandlerArguments,
|
|
)
|
|
from speech_to_speech.arguments_classes.kokoro_tts_arguments import KokoroTTSHandlerArguments
|
|
from speech_to_speech.arguments_classes.language_model_arguments import LanguageModelHandlerArguments
|
|
from speech_to_speech.arguments_classes.mlx_audio_whisper_arguments import (
|
|
MLXAudioWhisperSTTHandlerArguments,
|
|
)
|
|
from speech_to_speech.arguments_classes.module_arguments import ModuleArguments
|
|
from speech_to_speech.arguments_classes.paraformer_stt_arguments import ParaformerSTTHandlerArguments
|
|
from speech_to_speech.arguments_classes.parakeet_tdt_arguments import (
|
|
ParakeetTDTSTTHandlerArguments,
|
|
)
|
|
from speech_to_speech.arguments_classes.pocket_tts_arguments import PocketTTSHandlerArguments
|
|
from speech_to_speech.arguments_classes.qwen3_tts_arguments import Qwen3TTSHandlerArguments
|
|
from speech_to_speech.arguments_classes.responses_api_language_model_arguments import (
|
|
ResponsesApiLanguageModelHandlerArguments,
|
|
)
|
|
from speech_to_speech.arguments_classes.socket_receiver_arguments import SocketReceiverArguments
|
|
from speech_to_speech.arguments_classes.socket_sender_arguments import SocketSenderArguments
|
|
from speech_to_speech.arguments_classes.vad_arguments import VADHandlerArguments
|
|
from speech_to_speech.arguments_classes.websocket_streamer_arguments import WebSocketStreamerArguments
|
|
from speech_to_speech.arguments_classes.whisper_stt_arguments import WhisperSTTHandlerArguments
|
|
from speech_to_speech.baseHandler import BaseHandler
|
|
from speech_to_speech.LLM.chat import Chat
|
|
from speech_to_speech.pipeline.cancel_scope import CancelScope
|
|
from speech_to_speech.pipeline.handler_types import LLMIn, LLMOut, STTIn, STTOut, TTSIn, TTSOut
|
|
from speech_to_speech.pipeline.queue_types import (
|
|
AudioInItem,
|
|
AudioOutItem,
|
|
LMOutItem,
|
|
STTOutItem,
|
|
TextEventItem,
|
|
TextPromptItem,
|
|
TTSInItem,
|
|
VADOutItem,
|
|
)
|
|
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
|
|
from speech_to_speech.STT.transcription_notifier import TranscriptionNotifier
|
|
from speech_to_speech.utils.thread_manager import ThreadManager
|
|
from speech_to_speech.VAD.vad_handler import VADHandler
|
|
|
|
# Ensure that the necessary NLTK resources are available
|
|
try:
|
|
nltk.data.find("tokenizers/punkt_tab")
|
|
except (LookupError, OSError):
|
|
nltk.download("punkt_tab")
|
|
try:
|
|
nltk.data.find("tokenizers/averaged_perceptron_tagger_eng")
|
|
except (LookupError, OSError):
|
|
nltk.download("averaged_perceptron_tagger_eng")
|
|
|
|
# caching allows ~50% compilation time reduction
|
|
# see https://docs.google.com/document/d/1y5CRfMLdwEoF1nTk9q8qEu1mgMUuUtvhklPKJ2emLU8/edit#heading=h.o2asbxsrp1ma
|
|
CURRENT_DIR = Path(__file__).resolve().parent
|
|
os.environ["TORCHINDUCTOR_CACHE_DIR"] = os.path.join(CURRENT_DIR, "tmp")
|
|
|
|
console = Console()
|
|
logger = logging.getLogger(__name__)
|
|
logging.getLogger("numba").setLevel(logging.WARNING) # quiet down numba logs
|
|
|
|
|
|
@dataclass
|
|
class ParsedArguments:
|
|
module_kwargs: ModuleArguments
|
|
socket_receiver_kwargs: SocketReceiverArguments
|
|
socket_sender_kwargs: SocketSenderArguments
|
|
websocket_streamer_kwargs: WebSocketStreamerArguments
|
|
vad_handler_kwargs: VADHandlerArguments
|
|
whisper_stt_handler_kwargs: WhisperSTTHandlerArguments
|
|
paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments
|
|
faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments
|
|
mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments
|
|
parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments
|
|
language_model_handler_kwargs: LanguageModelHandlerArguments
|
|
responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments
|
|
chat_tts_handler_kwargs: ChatTTSHandlerArguments
|
|
facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments
|
|
pocket_tts_handler_kwargs: PocketTTSHandlerArguments
|
|
kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments
|
|
qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments
|
|
|
|
|
|
def rename_args(args: Any, prefix: str) -> None:
|
|
"""
|
|
Rename arguments by removing the prefix and prepares the gen_kwargs.
|
|
"""
|
|
gen_kwargs = {}
|
|
for key in copy(args.__dict__):
|
|
if key.startswith(prefix):
|
|
value = args.__dict__.pop(key)
|
|
new_key = key[len(prefix) + 1 :] # Remove prefix and underscore
|
|
if new_key.startswith("gen_"):
|
|
gen_kwargs[new_key[4:]] = value # Remove 'gen_' and add to dict
|
|
else:
|
|
args.__dict__[new_key] = value
|
|
|
|
args.__dict__["gen_kwargs"] = gen_kwargs
|
|
|
|
|
|
def parse_arguments() -> ParsedArguments:
|
|
# Pre-parse to determine which LM backend is selected, so only one of the
|
|
# mutually exclusive LM argument classes is registered with HfArgumentParser
|
|
# (avoids duplicate field names from the shared LanguageModelBaseArguments base).
|
|
# chat-completions reuses the responses-api connection fields via subclassing.
|
|
_backend_lm_class = {
|
|
"responses-api": ResponsesApiLanguageModelHandlerArguments,
|
|
"chat-completions": ChatCompletionsLanguageModelHandlerArguments,
|
|
}
|
|
_is_json = len(sys.argv) == 2 and sys.argv[1].endswith(".json")
|
|
if _is_json:
|
|
with open(sys.argv[1]) as _f:
|
|
_backend = json.load(_f).get("llm_backend")
|
|
else:
|
|
_pre = argparse.ArgumentParser(add_help=False)
|
|
_pre.add_argument("--llm_backend", default="responses-api")
|
|
_backend = _pre.parse_known_args()[0].llm_backend
|
|
|
|
_lm_class = _backend_lm_class.get(_backend, LanguageModelHandlerArguments)
|
|
logger.debug("LLM backend pre-parse: backend=%s, registering %s", _backend, _lm_class.__name__)
|
|
|
|
parser = HfArgumentParser(
|
|
( # type: ignore[arg-type]
|
|
ModuleArguments,
|
|
SocketReceiverArguments,
|
|
SocketSenderArguments,
|
|
WebSocketStreamerArguments,
|
|
VADHandlerArguments,
|
|
WhisperSTTHandlerArguments,
|
|
ParaformerSTTHandlerArguments,
|
|
FasterWhisperSTTHandlerArguments,
|
|
MLXAudioWhisperSTTHandlerArguments,
|
|
ParakeetTDTSTTHandlerArguments,
|
|
_lm_class,
|
|
ChatTTSHandlerArguments,
|
|
FacebookMMSTTSHandlerArguments,
|
|
PocketTTSHandlerArguments,
|
|
KokoroTTSHandlerArguments,
|
|
Qwen3TTSHandlerArguments,
|
|
)
|
|
)
|
|
|
|
if _is_json:
|
|
parsed = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]), allow_extra_keys=True)
|
|
else:
|
|
parsed = parser.parse_args_into_dataclasses()
|
|
|
|
# Build a {type: instance} lookup so field assignment is order-independent.
|
|
by_type: dict[type, Any] = {type(obj): obj for obj in parsed}
|
|
logger.debug("Parsed %d argument classes: %s", len(by_type), [t.__name__ for t in by_type])
|
|
|
|
return ParsedArguments(
|
|
module_kwargs=by_type[ModuleArguments],
|
|
socket_receiver_kwargs=by_type[SocketReceiverArguments],
|
|
socket_sender_kwargs=by_type[SocketSenderArguments],
|
|
websocket_streamer_kwargs=by_type[WebSocketStreamerArguments],
|
|
vad_handler_kwargs=by_type[VADHandlerArguments],
|
|
whisper_stt_handler_kwargs=by_type[WhisperSTTHandlerArguments],
|
|
paraformer_stt_handler_kwargs=by_type[ParaformerSTTHandlerArguments],
|
|
faster_whisper_stt_handler_kwargs=by_type[FasterWhisperSTTHandlerArguments],
|
|
mlx_audio_whisper_stt_handler_kwargs=by_type[MLXAudioWhisperSTTHandlerArguments],
|
|
parakeet_tdt_stt_handler_kwargs=by_type[ParakeetTDTSTTHandlerArguments],
|
|
language_model_handler_kwargs=by_type.get(LanguageModelHandlerArguments, LanguageModelHandlerArguments()),
|
|
# The OpenAI-compatible slot holds whichever class was registered:
|
|
# ChatCompletions... (a subclass) for chat-completions, else ResponsesApi....
|
|
responses_api_language_model_handler_kwargs=by_type.get(
|
|
ChatCompletionsLanguageModelHandlerArguments,
|
|
by_type.get(ResponsesApiLanguageModelHandlerArguments, ResponsesApiLanguageModelHandlerArguments()),
|
|
),
|
|
chat_tts_handler_kwargs=by_type[ChatTTSHandlerArguments],
|
|
facebook_mms_tts_handler_kwargs=by_type[FacebookMMSTTSHandlerArguments],
|
|
pocket_tts_handler_kwargs=by_type[PocketTTSHandlerArguments],
|
|
kokoro_tts_handler_kwargs=by_type[KokoroTTSHandlerArguments],
|
|
qwen3_tts_handler_kwargs=by_type[Qwen3TTSHandlerArguments],
|
|
)
|
|
|
|
|
|
def setup_logger(log_level: str) -> None:
|
|
global logger
|
|
from speech_to_speech.pipeline.log_context import PipelineLogFilter
|
|
|
|
logging.basicConfig(
|
|
level=log_level.upper(),
|
|
format="%(asctime)s - %(pipeline_prefix)s%(name)s - %(levelname)s - %(message)s",
|
|
)
|
|
# Attach the filter to every existing handler so each LogRecord gets a
|
|
# `pipeline_prefix` attribute (matching the format string above).
|
|
pipeline_filter = PipelineLogFilter()
|
|
for h in logging.getLogger().handlers:
|
|
h.addFilter(pipeline_filter)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# torch compile logs
|
|
if log_level == "debug":
|
|
torch._logging.set_logs(graph_breaks=True, recompiles=True, cudagraphs=True)
|
|
|
|
|
|
MLX_DEFAULT_LM_MODEL = "mlx-community/Qwen3-4B-Instruct-2507-bf16"
|
|
TRANSFORMERS_DEFAULT_LM_MODEL = "Qwen/Qwen3-4B-Instruct-2507"
|
|
|
|
|
|
def optimal_mac_settings(mac_optimal_settings: bool, *handler_kwargs: Any) -> None:
|
|
if mac_optimal_settings:
|
|
for kwargs in handler_kwargs:
|
|
if hasattr(kwargs, "device"):
|
|
kwargs.device = "mps"
|
|
if hasattr(kwargs, "mode"):
|
|
kwargs.mode = "local"
|
|
if hasattr(kwargs, "stt"):
|
|
kwargs.stt = "parakeet-tdt"
|
|
if hasattr(kwargs, "llm_backend"):
|
|
kwargs.llm_backend = "mlx-lm"
|
|
if hasattr(kwargs, "tts"):
|
|
kwargs.tts = "qwen3"
|
|
if hasattr(kwargs, "model_name"):
|
|
if kwargs.model_name == TRANSFORMERS_DEFAULT_LM_MODEL:
|
|
kwargs.model_name = MLX_DEFAULT_LM_MODEL
|
|
|
|
|
|
def check_mac_settings(module_kwargs: ModuleArguments) -> None:
|
|
if platform == "darwin":
|
|
if module_kwargs.device == "cuda":
|
|
raise ValueError("Cannot use CUDA on macOS. Please set the device to 'cpu' or 'mps'.")
|
|
if module_kwargs.llm_backend != "mlx-lm":
|
|
logger.warning(
|
|
"For macOS users, it is recommended to use mlx-lm. You can activate it by passing --llm_backend mlx-lm."
|
|
)
|
|
if module_kwargs.tts not in ("pocket", "kokoro", "qwen3"):
|
|
logger.warning(
|
|
"For macOS users, it is recommended to use qwen3 for TTS (pocket and kokoro are also valid options)."
|
|
)
|
|
|
|
|
|
def overwrite_device_argument(common_device: Optional[str], *handler_kwargs: Any) -> None:
|
|
if common_device:
|
|
for kwargs in handler_kwargs:
|
|
if hasattr(kwargs, "llm_device"):
|
|
kwargs.llm_device = common_device
|
|
if hasattr(kwargs, "tts_device"):
|
|
kwargs.tts_device = common_device
|
|
if hasattr(kwargs, "stt_device"):
|
|
kwargs.stt_device = common_device
|
|
if hasattr(kwargs, "paraformer_stt_device"):
|
|
kwargs.paraformer_stt_device = common_device
|
|
if hasattr(kwargs, "facebook_mms_device"):
|
|
kwargs.facebook_mms_device = common_device
|
|
if hasattr(kwargs, "qwen3_tts_device"):
|
|
kwargs.qwen3_tts_device = common_device
|
|
|
|
|
|
def prepare_module_args(module_kwargs: ModuleArguments, *handler_kwargs: Any) -> None:
|
|
optimal_mac_settings(module_kwargs.local_mac_optimal_settings, module_kwargs, *handler_kwargs)
|
|
if module_kwargs.tts is None:
|
|
module_kwargs.tts = "qwen3"
|
|
if platform == "darwin":
|
|
check_mac_settings(module_kwargs)
|
|
overwrite_device_argument(module_kwargs.device, *handler_kwargs)
|
|
|
|
|
|
def prepare_all_args(
|
|
module_kwargs: ModuleArguments,
|
|
whisper_stt_handler_kwargs: WhisperSTTHandlerArguments,
|
|
paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments,
|
|
faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments,
|
|
mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments,
|
|
parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments,
|
|
language_model_handler_kwargs: LanguageModelHandlerArguments,
|
|
responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments,
|
|
chat_tts_handler_kwargs: ChatTTSHandlerArguments,
|
|
facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments,
|
|
pocket_tts_handler_kwargs: PocketTTSHandlerArguments,
|
|
kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments,
|
|
qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments,
|
|
) -> None:
|
|
prepare_module_args(
|
|
module_kwargs,
|
|
whisper_stt_handler_kwargs,
|
|
faster_whisper_stt_handler_kwargs,
|
|
paraformer_stt_handler_kwargs,
|
|
mlx_audio_whisper_stt_handler_kwargs,
|
|
parakeet_tdt_stt_handler_kwargs,
|
|
language_model_handler_kwargs,
|
|
responses_api_language_model_handler_kwargs,
|
|
chat_tts_handler_kwargs,
|
|
facebook_mms_tts_handler_kwargs,
|
|
pocket_tts_handler_kwargs,
|
|
kokoro_tts_handler_kwargs,
|
|
qwen3_tts_handler_kwargs,
|
|
)
|
|
|
|
rename_args(whisper_stt_handler_kwargs, "stt")
|
|
rename_args(faster_whisper_stt_handler_kwargs, "faster_whisper_stt")
|
|
rename_args(paraformer_stt_handler_kwargs, "paraformer_stt")
|
|
rename_args(mlx_audio_whisper_stt_handler_kwargs, "mlx_audio_whisper")
|
|
rename_args(parakeet_tdt_stt_handler_kwargs, "parakeet_tdt")
|
|
rename_args(language_model_handler_kwargs, "llm")
|
|
rename_args(responses_api_language_model_handler_kwargs, "responses_api")
|
|
rename_args(chat_tts_handler_kwargs, "chat_tts")
|
|
rename_args(facebook_mms_tts_handler_kwargs, "facebook_mms")
|
|
rename_args(pocket_tts_handler_kwargs, "pocket_tts")
|
|
rename_args(kokoro_tts_handler_kwargs, "kokoro")
|
|
rename_args(qwen3_tts_handler_kwargs, "qwen3_tts")
|
|
|
|
|
|
def initialize_queues_and_events() -> dict[str, Any]:
|
|
return {
|
|
"stop_event": Event(),
|
|
"should_listen": Event(),
|
|
"response_playing": Event(),
|
|
"cancel_scope": CancelScope(),
|
|
"recv_audio_chunks_queue": Queue[AudioInItem](),
|
|
"send_audio_chunks_queue": Queue[AudioOutItem](),
|
|
"spoken_prompt_queue": Queue[VADOutItem](),
|
|
"stt_output_queue": Queue[STTOutItem](),
|
|
"text_prompt_queue": Queue[TextPromptItem](),
|
|
"lm_response_queue": Queue[LMOutItem](),
|
|
"lm_processed_queue": Queue[TTSInItem](), # NEW: LLM -> LM processor -> TTS
|
|
"text_output_queue": Queue[TextEventItem](), # NEW: for text messages to WebSocket
|
|
}
|
|
|
|
|
|
def _build_pipeline_handlers(
|
|
*,
|
|
stop_event: Event,
|
|
should_listen: Event,
|
|
recv_audio_chunks_queue: Queue[AudioInItem],
|
|
spoken_prompt_queue: Queue[VADOutItem],
|
|
stt_output_queue: Queue[STTOutItem],
|
|
text_prompt_queue: Queue[TextPromptItem],
|
|
lm_response_queue: Queue[LMOutItem],
|
|
lm_processed_queue: Queue[TTSInItem],
|
|
send_audio_chunks_queue: Queue[AudioOutItem],
|
|
text_output_queue: Optional[Queue[TextEventItem]],
|
|
transcription_notifier_setup: dict[str, Any],
|
|
module_kwargs: ModuleArguments,
|
|
vad_handler_kwargs: VADHandlerArguments,
|
|
whisper_stt_handler_kwargs: WhisperSTTHandlerArguments,
|
|
faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments,
|
|
paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments,
|
|
mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments,
|
|
parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments,
|
|
language_model_handler_kwargs: LanguageModelHandlerArguments,
|
|
responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments,
|
|
chat_tts_handler_kwargs: ChatTTSHandlerArguments,
|
|
facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments,
|
|
pocket_tts_handler_kwargs: PocketTTSHandlerArguments,
|
|
kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments,
|
|
qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments,
|
|
speculative_turns: SpeculativeTurnTracker | None = None,
|
|
) -> list[Any]:
|
|
"""Build the shared handler chain: VAD → STT → TranscriptionNotifier → LM → LMOutputProcessor → TTS.
|
|
|
|
Callers own the queues, events, and any per-mode mutations to handler kwargs (cancel_scope,
|
|
text_output_queue, live transcription flags). `transcription_notifier_setup` differs by mode:
|
|
realtime mode passes a service-driven setup; non-realtime modes inject a RuntimeConfig.
|
|
"""
|
|
from speech_to_speech.LLM.lm_output_processor import LMOutputProcessor
|
|
|
|
vad = VADHandler(
|
|
stop_event,
|
|
queue_in=recv_audio_chunks_queue,
|
|
queue_out=spoken_prompt_queue,
|
|
setup_args=(should_listen,),
|
|
setup_kwargs=vars(vad_handler_kwargs),
|
|
)
|
|
|
|
transcription_notifier = TranscriptionNotifier(
|
|
stop_event,
|
|
queue_in=stt_output_queue,
|
|
queue_out=text_prompt_queue, # type: ignore[arg-type]
|
|
setup_kwargs=transcription_notifier_setup,
|
|
)
|
|
|
|
stt = get_stt_handler(
|
|
module_kwargs,
|
|
stop_event,
|
|
spoken_prompt_queue,
|
|
stt_output_queue,
|
|
speculative_turns,
|
|
whisper_stt_handler_kwargs,
|
|
faster_whisper_stt_handler_kwargs,
|
|
paraformer_stt_handler_kwargs,
|
|
mlx_audio_whisper_stt_handler_kwargs,
|
|
parakeet_tdt_stt_handler_kwargs,
|
|
)
|
|
|
|
lm = get_llm_handler(
|
|
module_kwargs,
|
|
stop_event,
|
|
text_prompt_queue,
|
|
lm_response_queue,
|
|
language_model_handler_kwargs,
|
|
responses_api_language_model_handler_kwargs,
|
|
)
|
|
|
|
lm_processor = LMOutputProcessor(
|
|
stop_event,
|
|
queue_in=lm_response_queue,
|
|
queue_out=lm_processed_queue,
|
|
setup_kwargs={"text_output_queue": text_output_queue, "speculative_turns": speculative_turns},
|
|
)
|
|
|
|
tts = get_tts_handler(
|
|
module_kwargs,
|
|
stop_event,
|
|
lm_processed_queue,
|
|
send_audio_chunks_queue,
|
|
should_listen,
|
|
chat_tts_handler_kwargs,
|
|
facebook_mms_tts_handler_kwargs,
|
|
pocket_tts_handler_kwargs,
|
|
kokoro_tts_handler_kwargs,
|
|
qwen3_tts_handler_kwargs,
|
|
)
|
|
|
|
return [vad, stt, transcription_notifier, lm, lm_processor, tts]
|
|
|
|
|
|
def _build_realtime_pipeline_unit(
|
|
*,
|
|
index: int,
|
|
stop_event: Event,
|
|
module_kwargs: ModuleArguments,
|
|
vad_handler_kwargs: VADHandlerArguments,
|
|
whisper_stt_handler_kwargs: WhisperSTTHandlerArguments,
|
|
faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments,
|
|
paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments,
|
|
mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments,
|
|
parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments,
|
|
language_model_handler_kwargs: LanguageModelHandlerArguments,
|
|
responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments,
|
|
chat_tts_handler_kwargs: ChatTTSHandlerArguments,
|
|
facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments,
|
|
pocket_tts_handler_kwargs: PocketTTSHandlerArguments,
|
|
kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments,
|
|
qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments,
|
|
) -> "PipelineUnit":
|
|
"""Build one isolated realtime pipeline (own queues, events, service, handlers).
|
|
|
|
Returns a PipelineUnit that lives inside RealtimeServer's pool. Handler instances
|
|
are returned in `unit.handlers`; the caller threads them via ThreadManager. No uvicorn
|
|
is created here — the single server is owned by RealtimeServer.
|
|
"""
|
|
from speech_to_speech.api.openai_realtime.service import RealtimeService
|
|
from speech_to_speech.api.openai_realtime.delay_fillers import (
|
|
load_delay_filler_config,
|
|
pregenerate_delay_filler_clips,
|
|
)
|
|
|
|
# Per-unit deep copies so mutations (cancel_scope, text_output_queue) don't bleed across pipelines.
|
|
vad_kw = deepcopy(vad_handler_kwargs)
|
|
whisper_kw = deepcopy(whisper_stt_handler_kwargs)
|
|
faster_whisper_kw = deepcopy(faster_whisper_stt_handler_kwargs)
|
|
paraformer_kw = deepcopy(paraformer_stt_handler_kwargs)
|
|
mlx_audio_whisper_kw = deepcopy(mlx_audio_whisper_stt_handler_kwargs)
|
|
parakeet_kw = deepcopy(parakeet_tdt_stt_handler_kwargs)
|
|
lm_kw = deepcopy(language_model_handler_kwargs)
|
|
responses_api_kw = deepcopy(responses_api_language_model_handler_kwargs)
|
|
chat_tts_kw = deepcopy(chat_tts_handler_kwargs)
|
|
facebook_mms_kw = deepcopy(facebook_mms_tts_handler_kwargs)
|
|
pocket_tts_kw = deepcopy(pocket_tts_handler_kwargs)
|
|
kokoro_tts_kw = deepcopy(kokoro_tts_handler_kwargs)
|
|
qwen3_tts_kw = deepcopy(qwen3_tts_handler_kwargs)
|
|
|
|
should_listen = Event()
|
|
response_playing = Event()
|
|
cancel_scope = CancelScope()
|
|
speculative_turns = SpeculativeTurnTracker()
|
|
recv_audio_chunks_queue: Queue[AudioInItem] = Queue()
|
|
send_audio_chunks_queue: Queue[AudioOutItem] = Queue()
|
|
spoken_prompt_queue: Queue[VADOutItem] = Queue()
|
|
stt_output_queue: Queue[STTOutItem] = Queue()
|
|
text_prompt_queue: Queue[TextPromptItem] = Queue()
|
|
lm_response_queue: Queue[LMOutItem] = Queue()
|
|
lm_processed_queue: Queue[TTSInItem] = Queue()
|
|
text_output_queue: Queue[TextEventItem] = Queue()
|
|
|
|
vars(vad_kw)["text_output_queue"] = text_output_queue
|
|
vars(vad_kw)["speculative_turns"] = speculative_turns
|
|
for kw in (
|
|
lm_kw,
|
|
responses_api_kw,
|
|
kokoro_tts_kw,
|
|
qwen3_tts_kw,
|
|
pocket_tts_kw,
|
|
chat_tts_kw,
|
|
facebook_mms_kw,
|
|
):
|
|
vars(kw)["cancel_scope"] = cancel_scope
|
|
vars(kw)["speculative_turns"] = speculative_turns
|
|
|
|
if module_kwargs.llm_backend in ("responses-api", "chat-completions"):
|
|
chat_size = vars(responses_api_kw).get("chat_size", 10)
|
|
else:
|
|
chat_size = vars(lm_kw).get("chat_size", 10)
|
|
|
|
service = RealtimeService(
|
|
text_prompt_queue=text_prompt_queue,
|
|
should_listen=should_listen,
|
|
chat_size=chat_size,
|
|
speculative_turns=speculative_turns,
|
|
)
|
|
|
|
if module_kwargs.enable_live_transcription:
|
|
vad_kw.enable_realtime_transcription = True
|
|
vad_kw.realtime_processing_pause = module_kwargs.live_transcription_update_interval
|
|
|
|
handlers = _build_pipeline_handlers(
|
|
stop_event=stop_event,
|
|
should_listen=should_listen,
|
|
recv_audio_chunks_queue=recv_audio_chunks_queue,
|
|
spoken_prompt_queue=spoken_prompt_queue,
|
|
stt_output_queue=stt_output_queue,
|
|
text_prompt_queue=text_prompt_queue,
|
|
lm_response_queue=lm_response_queue,
|
|
lm_processed_queue=lm_processed_queue,
|
|
send_audio_chunks_queue=send_audio_chunks_queue,
|
|
text_output_queue=text_output_queue,
|
|
transcription_notifier_setup={
|
|
"text_output_queue": text_output_queue,
|
|
"should_listen": should_listen,
|
|
},
|
|
module_kwargs=module_kwargs,
|
|
vad_handler_kwargs=vad_kw,
|
|
whisper_stt_handler_kwargs=whisper_kw,
|
|
faster_whisper_stt_handler_kwargs=faster_whisper_kw,
|
|
paraformer_stt_handler_kwargs=paraformer_kw,
|
|
mlx_audio_whisper_stt_handler_kwargs=mlx_audio_whisper_kw,
|
|
parakeet_tdt_stt_handler_kwargs=parakeet_kw,
|
|
language_model_handler_kwargs=lm_kw,
|
|
responses_api_language_model_handler_kwargs=responses_api_kw,
|
|
chat_tts_handler_kwargs=chat_tts_kw,
|
|
facebook_mms_tts_handler_kwargs=facebook_mms_kw,
|
|
pocket_tts_handler_kwargs=pocket_tts_kw,
|
|
kokoro_tts_handler_kwargs=kokoro_tts_kw,
|
|
qwen3_tts_handler_kwargs=qwen3_tts_kw,
|
|
speculative_turns=speculative_turns,
|
|
)
|
|
for h in handlers:
|
|
h.pipeline_index = index
|
|
|
|
delay_filler_cfg = load_delay_filler_config()
|
|
delay_filler_short_clips: list[bytes] = []
|
|
delay_filler_full_clips: list[bytes] = []
|
|
delay_filler_threshold_s = 0.0
|
|
if delay_filler_cfg.enabled:
|
|
delay_filler_short_clips = pregenerate_delay_filler_clips(handlers[-1], delay_filler_cfg.short_phrases)
|
|
delay_filler_full_clips = pregenerate_delay_filler_clips(handlers[-1], delay_filler_cfg.full_phrases)
|
|
if delay_filler_short_clips or delay_filler_full_clips:
|
|
delay_filler_threshold_s = delay_filler_cfg.threshold_s
|
|
|
|
return PipelineUnit(
|
|
index=index,
|
|
service=service,
|
|
cancel_scope=cancel_scope,
|
|
should_listen=should_listen,
|
|
response_playing=response_playing,
|
|
input_queue=recv_audio_chunks_queue,
|
|
output_queue=send_audio_chunks_queue,
|
|
text_output_queue=text_output_queue,
|
|
text_prompt_queue=text_prompt_queue,
|
|
handlers=handlers,
|
|
delay_filler_threshold_s=delay_filler_threshold_s,
|
|
delay_filler_short_clips=delay_filler_short_clips,
|
|
delay_filler_full_clips=delay_filler_full_clips,
|
|
)
|
|
|
|
|
|
def build_pipeline(
|
|
module_kwargs: ModuleArguments,
|
|
socket_receiver_kwargs: SocketReceiverArguments,
|
|
socket_sender_kwargs: SocketSenderArguments,
|
|
websocket_streamer_kwargs: WebSocketStreamerArguments,
|
|
vad_handler_kwargs: VADHandlerArguments,
|
|
whisper_stt_handler_kwargs: WhisperSTTHandlerArguments,
|
|
faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments,
|
|
paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments,
|
|
mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments,
|
|
parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments,
|
|
language_model_handler_kwargs: LanguageModelHandlerArguments,
|
|
responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments,
|
|
chat_tts_handler_kwargs: ChatTTSHandlerArguments,
|
|
facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments,
|
|
pocket_tts_handler_kwargs: PocketTTSHandlerArguments,
|
|
kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments,
|
|
qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments,
|
|
queues_and_events: dict[str, Any],
|
|
) -> ThreadManager:
|
|
stop_event = queues_and_events["stop_event"]
|
|
should_listen = queues_and_events["should_listen"]
|
|
recv_audio_chunks_queue = queues_and_events["recv_audio_chunks_queue"]
|
|
send_audio_chunks_queue = queues_and_events["send_audio_chunks_queue"]
|
|
spoken_prompt_queue = queues_and_events["spoken_prompt_queue"]
|
|
stt_output_queue = queues_and_events["stt_output_queue"]
|
|
text_prompt_queue = queues_and_events["text_prompt_queue"]
|
|
lm_response_queue = queues_and_events["lm_response_queue"]
|
|
lm_processed_queue = queues_and_events["lm_processed_queue"]
|
|
text_output_queue = (
|
|
None # Only set for websocket/realtime modes; kept None otherwise to avoid unbounded queue growth
|
|
)
|
|
|
|
comms_handlers: list[Any] = []
|
|
if module_kwargs.mode == "local":
|
|
from speech_to_speech.connections.local_audio_streamer import LocalAudioStreamer
|
|
|
|
local_audio_streamer = LocalAudioStreamer(
|
|
input_queue=recv_audio_chunks_queue,
|
|
output_queue=send_audio_chunks_queue,
|
|
should_listen=should_listen,
|
|
)
|
|
comms_handlers = [local_audio_streamer]
|
|
should_listen.set()
|
|
elif module_kwargs.mode == "websocket":
|
|
from speech_to_speech.connections.websocket_streamer import WebSocketStreamer
|
|
|
|
text_output_queue = queues_and_events["text_output_queue"]
|
|
websocket_streamer = WebSocketStreamer(
|
|
stop_event,
|
|
input_queue=recv_audio_chunks_queue,
|
|
output_queue=send_audio_chunks_queue,
|
|
should_listen=should_listen,
|
|
text_output_queue=text_output_queue,
|
|
host=websocket_streamer_kwargs.ws_host,
|
|
port=websocket_streamer_kwargs.ws_port,
|
|
)
|
|
comms_handlers = [websocket_streamer]
|
|
elif module_kwargs.mode == "realtime":
|
|
from speech_to_speech.api.openai_realtime.server import RealtimeServer
|
|
|
|
pool_size = max(1, module_kwargs.num_pipelines)
|
|
pool = [
|
|
_build_realtime_pipeline_unit(
|
|
index=i,
|
|
stop_event=stop_event,
|
|
module_kwargs=module_kwargs,
|
|
vad_handler_kwargs=vad_handler_kwargs,
|
|
whisper_stt_handler_kwargs=whisper_stt_handler_kwargs,
|
|
faster_whisper_stt_handler_kwargs=faster_whisper_stt_handler_kwargs,
|
|
paraformer_stt_handler_kwargs=paraformer_stt_handler_kwargs,
|
|
mlx_audio_whisper_stt_handler_kwargs=mlx_audio_whisper_stt_handler_kwargs,
|
|
parakeet_tdt_stt_handler_kwargs=parakeet_tdt_stt_handler_kwargs,
|
|
language_model_handler_kwargs=language_model_handler_kwargs,
|
|
responses_api_language_model_handler_kwargs=responses_api_language_model_handler_kwargs,
|
|
chat_tts_handler_kwargs=chat_tts_handler_kwargs,
|
|
facebook_mms_tts_handler_kwargs=facebook_mms_tts_handler_kwargs,
|
|
pocket_tts_handler_kwargs=pocket_tts_handler_kwargs,
|
|
kokoro_tts_handler_kwargs=kokoro_tts_handler_kwargs,
|
|
qwen3_tts_handler_kwargs=qwen3_tts_handler_kwargs,
|
|
)
|
|
for i in range(pool_size)
|
|
]
|
|
|
|
realtime_server = RealtimeServer(
|
|
stop_event=stop_event,
|
|
pool=pool,
|
|
host=websocket_streamer_kwargs.ws_host,
|
|
port=websocket_streamer_kwargs.ws_port,
|
|
)
|
|
|
|
all_handlers: list[Any] = [realtime_server]
|
|
for unit in pool:
|
|
all_handlers.extend(unit.handlers)
|
|
return ThreadManager(all_handlers)
|
|
else:
|
|
from speech_to_speech.connections.socket_receiver import SocketReceiver
|
|
from speech_to_speech.connections.socket_sender import SocketSender
|
|
|
|
comms_handlers = [
|
|
SocketReceiver(
|
|
stop_event,
|
|
recv_audio_chunks_queue,
|
|
should_listen,
|
|
host=socket_receiver_kwargs.recv_host,
|
|
port=socket_receiver_kwargs.recv_port,
|
|
chunk_size=socket_receiver_kwargs.chunk_size,
|
|
),
|
|
SocketSender(
|
|
stop_event,
|
|
send_audio_chunks_queue,
|
|
should_listen,
|
|
host=socket_sender_kwargs.send_host,
|
|
port=socket_sender_kwargs.send_port,
|
|
),
|
|
]
|
|
|
|
# Set VAD realtime transcription parameters from module_kwargs
|
|
if module_kwargs.enable_live_transcription:
|
|
vad_handler_kwargs.enable_realtime_transcription = True
|
|
vad_handler_kwargs.realtime_processing_pause = module_kwargs.live_transcription_update_interval
|
|
|
|
if module_kwargs.llm_backend in ("responses-api", "chat-completions"):
|
|
_lm_vars = vars(responses_api_language_model_handler_kwargs)
|
|
else:
|
|
_lm_vars = vars(language_model_handler_kwargs)
|
|
transcription_notifier_setup: dict[str, Any] = {
|
|
"text_output_queue": text_output_queue,
|
|
"should_listen": should_listen,
|
|
"runtime_config": RuntimeConfig(
|
|
chat=Chat(_lm_vars.get("chat_size", 30)),
|
|
session=RealtimeSessionCreateRequest(
|
|
type="realtime",
|
|
instructions=_lm_vars.get("init_chat_prompt"),
|
|
),
|
|
),
|
|
}
|
|
|
|
pipeline_handlers = _build_pipeline_handlers(
|
|
stop_event=stop_event,
|
|
should_listen=should_listen,
|
|
recv_audio_chunks_queue=recv_audio_chunks_queue,
|
|
spoken_prompt_queue=spoken_prompt_queue,
|
|
stt_output_queue=stt_output_queue,
|
|
text_prompt_queue=text_prompt_queue,
|
|
lm_response_queue=lm_response_queue,
|
|
lm_processed_queue=lm_processed_queue,
|
|
send_audio_chunks_queue=send_audio_chunks_queue,
|
|
text_output_queue=text_output_queue,
|
|
transcription_notifier_setup=transcription_notifier_setup,
|
|
module_kwargs=module_kwargs,
|
|
vad_handler_kwargs=vad_handler_kwargs,
|
|
whisper_stt_handler_kwargs=whisper_stt_handler_kwargs,
|
|
faster_whisper_stt_handler_kwargs=faster_whisper_stt_handler_kwargs,
|
|
paraformer_stt_handler_kwargs=paraformer_stt_handler_kwargs,
|
|
mlx_audio_whisper_stt_handler_kwargs=mlx_audio_whisper_stt_handler_kwargs,
|
|
parakeet_tdt_stt_handler_kwargs=parakeet_tdt_stt_handler_kwargs,
|
|
language_model_handler_kwargs=language_model_handler_kwargs,
|
|
responses_api_language_model_handler_kwargs=responses_api_language_model_handler_kwargs,
|
|
chat_tts_handler_kwargs=chat_tts_handler_kwargs,
|
|
facebook_mms_tts_handler_kwargs=facebook_mms_tts_handler_kwargs,
|
|
pocket_tts_handler_kwargs=pocket_tts_handler_kwargs,
|
|
kokoro_tts_handler_kwargs=kokoro_tts_handler_kwargs,
|
|
qwen3_tts_handler_kwargs=qwen3_tts_handler_kwargs,
|
|
)
|
|
|
|
return ThreadManager([*comms_handlers, *pipeline_handlers])
|
|
|
|
|
|
def get_stt_handler(
|
|
module_kwargs: ModuleArguments,
|
|
stop_event: Event,
|
|
spoken_prompt_queue: Queue[VADOutItem],
|
|
text_prompt_queue: Queue[STTOutItem],
|
|
speculative_turns: SpeculativeTurnTracker | None,
|
|
whisper_stt_handler_kwargs: WhisperSTTHandlerArguments,
|
|
faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments,
|
|
paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments,
|
|
mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments,
|
|
parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments,
|
|
) -> BaseHandler[STTIn, STTOut]:
|
|
from speech_to_speech.STT.base_stt_handler import BaseSTTHandler
|
|
|
|
def with_speculative_turns(handler: BaseSTTHandler) -> BaseSTTHandler:
|
|
if speculative_turns is not None:
|
|
handler.speculative_turns = speculative_turns
|
|
return handler
|
|
|
|
if module_kwargs.stt == "whisper":
|
|
from speech_to_speech.STT.whisper_stt_handler import WhisperSTTHandler
|
|
|
|
return with_speculative_turns(
|
|
WhisperSTTHandler(
|
|
stop_event,
|
|
queue_in=spoken_prompt_queue,
|
|
queue_out=text_prompt_queue,
|
|
setup_kwargs=vars(whisper_stt_handler_kwargs),
|
|
)
|
|
)
|
|
elif module_kwargs.stt == "whisper-mlx":
|
|
from speech_to_speech.STT.lightning_whisper_mlx_handler import LightningWhisperSTTHandler
|
|
|
|
return with_speculative_turns(
|
|
LightningWhisperSTTHandler(
|
|
stop_event,
|
|
queue_in=spoken_prompt_queue,
|
|
queue_out=text_prompt_queue,
|
|
setup_kwargs=vars(whisper_stt_handler_kwargs),
|
|
)
|
|
)
|
|
elif module_kwargs.stt == "mlx-audio-whisper":
|
|
from speech_to_speech.STT.mlx_audio_whisper_handler import MLXAudioWhisperSTTHandler
|
|
|
|
# Merge MLX Audio Whisper kwargs with shared language parameter from Whisper kwargs
|
|
setup_kwargs = {**vars(mlx_audio_whisper_stt_handler_kwargs), "language": whisper_stt_handler_kwargs.language}
|
|
return with_speculative_turns(
|
|
MLXAudioWhisperSTTHandler(
|
|
stop_event,
|
|
queue_in=spoken_prompt_queue,
|
|
queue_out=text_prompt_queue,
|
|
setup_kwargs=setup_kwargs,
|
|
)
|
|
)
|
|
elif module_kwargs.stt == "paraformer":
|
|
from speech_to_speech.STT.paraformer_handler import ParaformerSTTHandler
|
|
|
|
return with_speculative_turns(
|
|
ParaformerSTTHandler(
|
|
stop_event,
|
|
queue_in=spoken_prompt_queue,
|
|
queue_out=text_prompt_queue,
|
|
setup_kwargs=vars(paraformer_stt_handler_kwargs),
|
|
)
|
|
)
|
|
elif module_kwargs.stt == "faster-whisper":
|
|
from speech_to_speech.STT.faster_whisper_handler import FasterWhisperSTTHandler
|
|
|
|
return with_speculative_turns(
|
|
FasterWhisperSTTHandler(
|
|
stop_event,
|
|
queue_in=spoken_prompt_queue,
|
|
queue_out=text_prompt_queue,
|
|
setup_kwargs=vars(faster_whisper_stt_handler_kwargs),
|
|
)
|
|
)
|
|
elif module_kwargs.stt == "parakeet-tdt":
|
|
from speech_to_speech.STT.parakeet_tdt_handler import ParakeetTDTSTTHandler
|
|
|
|
# Add live transcription parameters to setup_kwargs
|
|
setup_kwargs = {
|
|
**vars(parakeet_tdt_stt_handler_kwargs),
|
|
"enable_live_transcription": module_kwargs.enable_live_transcription,
|
|
"live_transcription_update_interval": module_kwargs.live_transcription_update_interval,
|
|
}
|
|
|
|
return with_speculative_turns(
|
|
ParakeetTDTSTTHandler(
|
|
stop_event,
|
|
queue_in=spoken_prompt_queue,
|
|
queue_out=text_prompt_queue,
|
|
setup_kwargs=setup_kwargs,
|
|
)
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
"The STT should be either whisper, whisper-mlx, mlx-audio-whisper, faster-whisper, parakeet-tdt, or paraformer."
|
|
)
|
|
|
|
|
|
def get_llm_handler(
|
|
module_kwargs: ModuleArguments,
|
|
stop_event: Event,
|
|
text_prompt_queue: Queue[TextPromptItem],
|
|
lm_response_queue: Queue[LMOutItem],
|
|
language_model_handler_kwargs: LanguageModelHandlerArguments,
|
|
responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments,
|
|
) -> BaseHandler[LLMIn, LLMOut]:
|
|
if module_kwargs.llm_backend == "responses-api":
|
|
from speech_to_speech.LLM.responses_api_language_model import ResponsesApiModelHandler
|
|
|
|
return ResponsesApiModelHandler(
|
|
stop_event,
|
|
queue_in=text_prompt_queue,
|
|
queue_out=lm_response_queue,
|
|
setup_kwargs=vars(responses_api_language_model_handler_kwargs),
|
|
)
|
|
|
|
if module_kwargs.llm_backend == "chat-completions":
|
|
from speech_to_speech.LLM.chat_completions_language_model import ChatCompletionsApiModelHandler
|
|
|
|
# Reuses the responses-api argument class (identical fields: model_name,
|
|
# base_url, api_key, stream, disable_thinking, ...).
|
|
return ChatCompletionsApiModelHandler(
|
|
stop_event,
|
|
queue_in=text_prompt_queue,
|
|
queue_out=lm_response_queue,
|
|
setup_kwargs=vars(responses_api_language_model_handler_kwargs),
|
|
)
|
|
|
|
if module_kwargs.llm_backend in ("transformers", "mlx-lm"):
|
|
lm_kwargs = vars(language_model_handler_kwargs)
|
|
is_vlm = lm_kwargs.pop("is_vlm", False)
|
|
if module_kwargs.llm_backend == "mlx-lm":
|
|
lm_kwargs["backend"] = "mlx"
|
|
|
|
if is_vlm:
|
|
from speech_to_speech.LLM.language_model import VisionLanguageModelHandler
|
|
|
|
return VisionLanguageModelHandler(
|
|
stop_event,
|
|
queue_in=text_prompt_queue,
|
|
queue_out=lm_response_queue,
|
|
setup_kwargs=lm_kwargs,
|
|
)
|
|
from speech_to_speech.LLM.language_model import LanguageModelHandler
|
|
|
|
return LanguageModelHandler(
|
|
stop_event,
|
|
queue_in=text_prompt_queue,
|
|
queue_out=lm_response_queue,
|
|
setup_kwargs=lm_kwargs,
|
|
)
|
|
|
|
raise ValueError("The LLM should be either transformers, mlx-lm, responses-api or chat-completions")
|
|
|
|
|
|
def get_tts_handler(
|
|
module_kwargs: ModuleArguments,
|
|
stop_event: Event,
|
|
lm_response_queue: Queue[TTSInItem],
|
|
send_audio_chunks_queue: Queue[AudioOutItem],
|
|
should_listen: Event,
|
|
chat_tts_handler_kwargs: ChatTTSHandlerArguments,
|
|
facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments,
|
|
pocket_tts_handler_kwargs: PocketTTSHandlerArguments,
|
|
kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments,
|
|
qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments,
|
|
) -> BaseHandler[TTSIn, TTSOut]:
|
|
if module_kwargs.tts == "chatTTS":
|
|
try:
|
|
from speech_to_speech.TTS.chatTTS_handler import ChatTTSHandler
|
|
except (ImportError, RuntimeError) as e:
|
|
logger.error('Error importing ChatTTSHandler. Install it with `pip install "speech-to-speech[chattts]"`.')
|
|
raise e
|
|
return ChatTTSHandler(
|
|
stop_event,
|
|
queue_in=lm_response_queue,
|
|
queue_out=send_audio_chunks_queue,
|
|
setup_args=(should_listen,),
|
|
setup_kwargs=vars(chat_tts_handler_kwargs),
|
|
)
|
|
elif module_kwargs.tts == "facebookMMS":
|
|
from speech_to_speech.TTS.facebookmms_handler import FacebookMMSTTSHandler
|
|
|
|
return FacebookMMSTTSHandler(
|
|
stop_event,
|
|
queue_in=lm_response_queue,
|
|
queue_out=send_audio_chunks_queue,
|
|
setup_args=(should_listen,),
|
|
setup_kwargs=vars(facebook_mms_tts_handler_kwargs),
|
|
)
|
|
elif module_kwargs.tts == "pocket":
|
|
try:
|
|
from speech_to_speech.TTS.pocket_tts_handler import PocketTTSHandler
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
'Pocket TTS is optional. Install it with `pip install "speech-to-speech[pocket]"`.'
|
|
) from e
|
|
|
|
return PocketTTSHandler(
|
|
stop_event,
|
|
queue_in=lm_response_queue,
|
|
queue_out=send_audio_chunks_queue,
|
|
setup_args=(should_listen,),
|
|
setup_kwargs=vars(pocket_tts_handler_kwargs),
|
|
)
|
|
elif module_kwargs.tts == "kokoro":
|
|
try:
|
|
from speech_to_speech.TTS.kokoro_handler import KokoroTTSHandler
|
|
except ImportError as e:
|
|
raise ImportError('Kokoro is optional. Install it with `pip install "speech-to-speech[kokoro]"`.') from e
|
|
|
|
return KokoroTTSHandler(
|
|
stop_event,
|
|
queue_in=lm_response_queue,
|
|
queue_out=send_audio_chunks_queue,
|
|
setup_args=(should_listen,),
|
|
setup_kwargs=vars(kokoro_tts_handler_kwargs),
|
|
)
|
|
elif module_kwargs.tts == "qwen3":
|
|
from speech_to_speech.TTS.qwen3_tts_handler import Qwen3TTSHandler
|
|
|
|
return Qwen3TTSHandler(
|
|
stop_event,
|
|
queue_in=lm_response_queue,
|
|
queue_out=send_audio_chunks_queue,
|
|
setup_args=(should_listen,),
|
|
setup_kwargs=vars(qwen3_tts_handler_kwargs),
|
|
)
|
|
else:
|
|
raise ValueError("The TTS should be either chatTTS, facebookMMS, pocket, kokoro, or qwen3")
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_arguments()
|
|
|
|
setup_logger(args.module_kwargs.log_level)
|
|
|
|
if args.module_kwargs.num_pipelines < 1:
|
|
raise ValueError(f"--num_pipelines must be >= 1, got {args.module_kwargs.num_pipelines}")
|
|
|
|
prepare_all_args(
|
|
args.module_kwargs,
|
|
args.whisper_stt_handler_kwargs,
|
|
args.paraformer_stt_handler_kwargs,
|
|
args.faster_whisper_stt_handler_kwargs,
|
|
args.mlx_audio_whisper_stt_handler_kwargs,
|
|
args.parakeet_tdt_stt_handler_kwargs,
|
|
args.language_model_handler_kwargs,
|
|
args.responses_api_language_model_handler_kwargs,
|
|
args.chat_tts_handler_kwargs,
|
|
args.facebook_mms_tts_handler_kwargs,
|
|
args.pocket_tts_handler_kwargs,
|
|
args.kokoro_tts_handler_kwargs,
|
|
args.qwen3_tts_handler_kwargs,
|
|
)
|
|
|
|
# Validate after prepare_all_args(): --local_mac_optimal_settings mutates
|
|
# module_kwargs.mode to "local", so checking before would let
|
|
# --local_mac_optimal_settings --num_pipelines 2 sneak past this guard.
|
|
if args.module_kwargs.num_pipelines > 1 and args.module_kwargs.mode != "realtime":
|
|
raise ValueError(
|
|
f"--num_pipelines > 1 is only supported with --mode realtime "
|
|
f"(got mode={args.module_kwargs.mode!r}, num_pipelines={args.module_kwargs.num_pipelines})"
|
|
)
|
|
|
|
# On Apple Silicon, all MLX inference serializes through a global lock (utils/mlx_lock.py).
|
|
# The progressive STT path uses a short timeout and drops work under contention, producing
|
|
# a flood of warnings without affecting final transcripts. With a pool, pre-emptively turn
|
|
# it off so logs stay readable; the final STT path is unaffected. Non-darwin platforms
|
|
# don't share this lock, so leave their live transcription alone.
|
|
if args.module_kwargs.num_pipelines > 1 and platform == "darwin" and args.module_kwargs.enable_live_transcription:
|
|
logger.info(
|
|
"MLX contention: --num_pipelines=%d > 1 on Apple Silicon → disabling live transcription "
|
|
"(progressive STT contends on the global MLX lock)",
|
|
args.module_kwargs.num_pipelines,
|
|
)
|
|
args.module_kwargs.enable_live_transcription = False
|
|
|
|
queues_and_events = initialize_queues_and_events()
|
|
|
|
pipeline_manager = build_pipeline(
|
|
args.module_kwargs,
|
|
args.socket_receiver_kwargs,
|
|
args.socket_sender_kwargs,
|
|
args.websocket_streamer_kwargs,
|
|
args.vad_handler_kwargs,
|
|
args.whisper_stt_handler_kwargs,
|
|
args.faster_whisper_stt_handler_kwargs,
|
|
args.paraformer_stt_handler_kwargs,
|
|
args.mlx_audio_whisper_stt_handler_kwargs,
|
|
args.parakeet_tdt_stt_handler_kwargs,
|
|
args.language_model_handler_kwargs,
|
|
args.responses_api_language_model_handler_kwargs,
|
|
args.chat_tts_handler_kwargs,
|
|
args.facebook_mms_tts_handler_kwargs,
|
|
args.pocket_tts_handler_kwargs,
|
|
args.kokoro_tts_handler_kwargs,
|
|
args.qwen3_tts_handler_kwargs,
|
|
queues_and_events,
|
|
)
|
|
|
|
# Set up graceful shutdown handler
|
|
shutdown_requested = [False] # Use list for nonlocal mutation
|
|
|
|
def signal_handler(_sig: int, _frame: Optional[FrameType]) -> None:
|
|
if not shutdown_requested[0]:
|
|
shutdown_requested[0] = True
|
|
console.print("\n[yellow]Shutting down gracefully...[/yellow]")
|
|
pipeline_manager.stop()
|
|
console.print("[green]✓ Pipeline stopped successfully[/green]")
|
|
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
|
|
try:
|
|
pipeline_manager.start()
|
|
pipeline_manager.wait()
|
|
except KeyboardInterrupt:
|
|
if not shutdown_requested[0]:
|
|
console.print("\n[yellow]Shutting down gracefully...[/yellow]")
|
|
pipeline_manager.stop()
|
|
console.print("[green]✓ Pipeline stopped successfully[/green]")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|