""" models.py --------- Single source of truth for every model used in the pipeline. Component Model Params Device ----------- --------------------------------------- ------- ------ STT openai/whisper-base.en 74 M GPU/CPU (fp16/fp32) LLM HuggingFaceTB/SmolLM2-1.7B-Instruct 1.7 B GPU/CPU (fp16/fp32) TTS facebook/mms-tts-eng (VITS) ~430 M CPU (fp32) VAD silero-vad 2 MB CPU All models are loaded once at startup. Inference methods are synchronous (blocking) — call them from a thread-pool executor inside async handlers. """ from __future__ import annotations import logging import time from typing import Optional import numpy as np import torch from transformers import ( AutoModelForCausalLM, AutoTokenizer, VitsModel, WhisperForConditionalGeneration, WhisperProcessor, ) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # # Model IDs (swap here for lighter / heavier variants) # # --------------------------------------------------------------------------- # STT_MODEL_ID = "openai/whisper-base.en" # 74 M — swap tiny.en for CPU-only LLM_MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct" # 1.7 B — swap 360M for CPU-only TTS_MODEL_ID = "facebook/mms-tts-eng" # VITS, 16 kHz output STT_SR = 16_000 # Whisper expects 16 kHz TTS_SR = 16_000 # MMS-TTS-ENG native output rate MAX_NEW_TOKENS = 256 # Cap LLM response length for voice interaction # --------------------------------------------------------------------------- # # ModelManager # # --------------------------------------------------------------------------- # class ModelManager: """Load and expose all model inference in one place.""" def __init__(self): self.device = "cuda" if torch.cuda.is_available() else "cpu" self.compute_dtype = ( torch.float16 if self.device == "cuda" else torch.float32 ) logger.info("Device: %s | dtype: %s", self.device, self.compute_dtype) # TTS always runs on CPU (saves GPU VRAM for STT + LLM) self.tts_device = "cpu" self._load_vad() self._load_stt() self._load_llm() self._load_tts() logger.info("✅ All models ready.") # ------------------------------------------------------------------ # # VAD # # ------------------------------------------------------------------ # def _load_vad(self): logger.info("Loading VAD (silero-vad) …") from silero_vad import VADIterator, load_silero_vad # type: ignore vad_model = load_silero_vad() # silence_ms=700 → 700 ms of quiet ends an utterance self.vad_iter = VADIterator( vad_model, sampling_rate=STT_SR, threshold=0.50, min_silence_duration_ms=700, speech_pad_ms=50, ) logger.info("VAD ready.") def vad_reset(self): """Reset VAD state between conversations.""" self.vad_iter.reset_states() # ------------------------------------------------------------------ # # STT — Whisper # # ------------------------------------------------------------------ # def _load_stt(self): logger.info("Loading STT: %s …", STT_MODEL_ID) self.stt_processor = WhisperProcessor.from_pretrained(STT_MODEL_ID) self.stt_model = WhisperForConditionalGeneration.from_pretrained( STT_MODEL_ID, torch_dtype=self.compute_dtype, ).to(self.device) self.stt_model.eval() logger.info("STT ready on %s.", self.device) def transcribe(self, audio_float32: np.ndarray) -> str: """ Transcribe a 16 kHz float32 numpy array to text. Parameters ---------- audio_float32 : np.ndarray shape (N,), dtype float32, range [-1, 1] """ t0 = time.perf_counter() inputs = self.stt_processor( audio_float32, sampling_rate=STT_SR, return_tensors="pt", ) features = inputs.input_features.to( device=self.device, dtype=self.compute_dtype ) with torch.no_grad(): ids = self.stt_model.generate(features) text = self.stt_processor.batch_decode(ids, skip_special_tokens=True)[0].strip() logger.info("STT (%.0f ms): '%s'", (time.perf_counter() - t0) * 1000, text) return text # ------------------------------------------------------------------ # # LLM — SmolLM2 # # ------------------------------------------------------------------ # def _load_llm(self): logger.info("Loading LLM: %s …", LLM_MODEL_ID) self.llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID) self.llm = AutoModelForCausalLM.from_pretrained( LLM_MODEL_ID, torch_dtype=self.compute_dtype, ).to(self.device) self.llm.eval() logger.info("LLM ready on %s.", self.device) def build_prompt(self, messages: list[dict]) -> str: """Apply the model's chat template and return a formatted string.""" return self.llm_tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) def tokenize(self, prompt: str) -> dict: """Tokenize a prompt string → dict of tensors on self.device.""" enc = self.llm_tokenizer(prompt, return_tensors="pt") return {k: v.to(self.device) for k, v in enc.items()} # ------------------------------------------------------------------ # # TTS — MMS-TTS-ENG (VITS) # # ------------------------------------------------------------------ # def _load_tts(self): logger.info("Loading TTS: %s …", TTS_MODEL_ID) self.tts_tokenizer = AutoTokenizer.from_pretrained(TTS_MODEL_ID) self.tts_model = VitsModel.from_pretrained(TTS_MODEL_ID) # fp32 on CPU self.tts_model.eval() self.tts_sample_rate: int = self.tts_model.config.sampling_rate # 16 000 logger.info("TTS ready on CPU (sample_rate=%d Hz).", self.tts_sample_rate) def synthesize(self, text: str) -> Optional[bytes]: """ Synthesize text → raw PCM bytes (int16, mono, 16 kHz). Returns None if text is empty or synthesis fails. """ text = text.strip() if not text: return None t0 = time.perf_counter() try: inputs = self.tts_tokenizer(text, return_tensors="pt") with torch.no_grad(): output = self.tts_model(**inputs) # waveform: (1, T) float32 in roughly [-1, 1] wav: np.ndarray = output.waveform.squeeze().cpu().numpy() # Convert float32 → int16 PCM wav_int16 = np.clip(wav * 32_767.0, -32_768, 32_767).astype(np.int16) pcm_bytes = wav_int16.tobytes() logger.debug( "TTS (%.0f ms): %d chars → %d bytes", (time.perf_counter() - t0) * 1000, len(text), len(pcm_bytes), ) return pcm_bytes except Exception as exc: logger.error("TTS synthesis failed for %r: %s", text[:60], exc) return None # ------------------------------------------------------------------ # # Convenience # # ------------------------------------------------------------------ # def warm_up(self): """Run a tiny inference on each model to pre-JIT all kernels.""" logger.info("Warming up models …") silence = np.zeros(STT_SR, dtype=np.float32) self.transcribe(silence) dummy_msgs = [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "hi"}, ] # Just tokenize, don't run full generation during warm-up self.tokenize(self.build_prompt(dummy_msgs)) self.synthesize("Hello.") logger.info("Warm-up complete.")