""" speech.py — pluggable STT / TTS backends ------------------------------------------ FastRTC ships English-only built-ins out of the box: - get_stt_model("moonshine/base") -> STT - get_tts_model("kokoro") -> TTS Neither supports Malayalam, so this module adds two small classes that satisfy FastRTC's STTModel / TTSModel protocols (a `.stt(audio)` method and `.tts(text, options)` / `.stream_tts_sync(text, options)` methods) and can be dropped into `ReplyOnPause` exactly like the built-ins. MalayalamSTT -> a Whisper checkpoint fine-tuned on Malayalam speech MalayalamTTS -> Meta's MMS VITS Malayalam checkpoint Swap either model id via env vars without touching any other code. """ from __future__ import annotations import os import re from dataclasses import dataclass import numpy as np import torch from transformers import ( AutoModelForSpeechSeq2Seq, AutoProcessor, VitsModel, AutoTokenizer, ) from fastrtc import KokoroTTSOptions, get_stt_model, get_tts_model DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 # Fine-tuned Whisper checkpoint for Malayalam ASR. Swap for whichever # checkpoint gives you the best WER on your audience's accent/domain -- # a few public options: vrclc/Whisper-medium-Malayalam, # kavyamanohar/whisper-small-malayalam, Jithjacob123/whisper-small-Malayalam. MALAYALAM_STT_MODEL_ID = os.environ.get( "MEDGUIDE_ML_STT_MODEL_ID", "vrclc/Whisper-medium-Malayalam" ) # Meta MMS single-speaker VITS checkpoint for Malayalam TTS. For higher # quality/expressiveness (at the cost of a heavier model + prompt # engineering) swap to ai4bharat/indic-parler-tts, which also covers # Malayalam plus 20 other Indic languages. MALAYALAM_TTS_MODEL_ID = os.environ.get( "MEDGUIDE_ML_TTS_MODEL_ID", "facebook/mms-tts-mal" ) TARGET_SR = 16000 def _to_mono_float32(audio: tuple[int, np.ndarray]) -> tuple[int, np.ndarray]: """Normalize a FastRTC (sample_rate, array) pair to mono float32 in [-1, 1].""" sr, arr = audio arr = np.asarray(arr) if arr.ndim == 2: arr = arr.mean(axis=0) if arr.shape[0] < arr.shape[1] else arr.mean(axis=1) if np.issubdtype(arr.dtype, np.integer): arr = arr.astype(np.float32) / 32768.0 else: arr = arr.astype(np.float32) return sr, arr def _resample(arr: np.ndarray, sr: int, target_sr: int = TARGET_SR) -> np.ndarray: if sr == target_sr: return arr from scipy import signal n_target = int(round(len(arr) * target_sr / sr)) return signal.resample(arr, max(n_target, 1)).astype(np.float32) class MalayalamSTT: """Implements FastRTC's STTModel protocol: `.stt(audio) -> str`.""" def __init__(self, model_id: str = MALAYALAM_STT_MODEL_ID): self.processor = AutoProcessor.from_pretrained(model_id) self.model = AutoModelForSpeechSeq2Seq.from_pretrained( model_id, dtype=DTYPE ).to(DEVICE) self.model.eval() # Force Malayalam decoding regardless of Whisper's language auto-detect. try: self.forced_ids = self.processor.get_decoder_prompt_ids( language="malayalam", task="transcribe" ) except Exception: self.forced_ids = None def stt(self, audio: tuple[int, np.ndarray]) -> str: sr, arr = _to_mono_float32(audio) arr = _resample(arr, sr, TARGET_SR) inputs = self.processor( arr, sampling_rate=TARGET_SR, return_tensors="pt" ).to(DEVICE) with torch.no_grad(): ids = self.model.generate( **inputs, forced_decoder_ids=self.forced_ids, max_new_tokens=200, ) return self.processor.batch_decode(ids, skip_special_tokens=True)[0].strip() @dataclass class SimpleTTSOptions: speed: float = 1.0 _SENTENCE_SPLIT_RE = re.compile(r"(?<=[.!?।])\s+") class MalayalamTTS: """Implements FastRTC's TTSModel protocol: `.tts()` and `.stream_tts_sync()`. MMS-VITS is not autoregressive/streamable at the token level, so "streaming" here means synthesizing sentence-by-sentence and yielding each clip as soon as it's ready -- still cuts first-audio latency noticeably versus waiting for the whole reply.""" def __init__(self, model_id: str = MALAYALAM_TTS_MODEL_ID): self.tokenizer = AutoTokenizer.from_pretrained(model_id) self.model = VitsModel.from_pretrained(model_id).to(DEVICE) self.model.eval() self.sample_rate = self.model.config.sampling_rate def _synth_one(self, text: str) -> tuple[int, np.ndarray]: inputs = self.tokenizer(text, return_tensors="pt").to(DEVICE) with torch.no_grad(): waveform = self.model(**inputs).waveform arr = waveform.squeeze().cpu().numpy().astype(np.float32) return self.sample_rate, arr.reshape(1, -1) def tts(self, text: str, options: SimpleTTSOptions | None = None): return self._synth_one(text) def stream_tts_sync(self, text: str, options: SimpleTTSOptions | None = None): chunks = [c.strip() for c in _SENTENCE_SPLIT_RE.split(text) if c.strip()] if not chunks: return for chunk in chunks: yield self._synth_one(chunk) # --------------------------------------------------------------------- # Per-language registry the UI/handler pick from # --------------------------------------------------------------------- _english_stt = None _english_tts = None _malayalam_stt = None _malayalam_tts = None def get_stt(language: str): global _english_stt, _malayalam_stt if language == "Malayalam": if _malayalam_stt is None: _malayalam_stt = MalayalamSTT() return _malayalam_stt if _english_stt is None: _english_stt = get_stt_model(model="moonshine/base") return _english_stt def get_tts(language: str): global _english_tts, _malayalam_tts if language == "Malayalam": if _malayalam_tts is None: _malayalam_tts = MalayalamTTS() return _malayalam_tts if _english_tts is None: _english_tts = get_tts_model(model="kokoro") return _english_tts ENGLISH_TTS_OPTIONS = KokoroTTSOptions(voice="af_heart", speed=1.0, lang="en-us")