models.py

"""
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.")