vad.py

"""
Voice Activity Detection
=========================
Frame-level speech/silence classification, used for:

  - Endpointing     : knowing when the caller has FINISHED speaking
                      (so we don't wait for a push-to-talk button)
  - Preroll         : keeping ~300ms of audio BEFORE speech onset so the
                      first phoneme is never clipped
  - Barge-in        : detecting that the caller started speaking while the
                      agent's TTS is still playing, so we can stop it
  - Silence culling : never sending silence to Whisper (Whisper hallucinates
                      badly on silence — "Thank you." / "Subtitles by…")

Two backends:
  1. SileroVAD  — torch.hub snakers4/silero-vad. Small (~1.8MB), fast on CPU,
                  robust to background noise. Requires EXACTLY 512 samples
                  per frame at 16kHz.
  2. EnergyVAD  — adaptive-threshold RMS + zero-crossing rate. No download,
                  no torch.hub dependency. Degraded in noise but never fails.

Both expose the same interface:
    vad.speech_prob(frame: np.ndarray) -> float in [0, 1]
    vad.reset()
"""

import logging
import numpy as np
from typing import Optional

logger = logging.getLogger(__name__)

SAMPLE_RATE   = 16_000
FRAME_SAMPLES = 512                       # Silero requirement at 16kHz
FRAME_MS      = FRAME_SAMPLES / SAMPLE_RATE * 1000     # 32.0 ms


# ──────────────────────────────────────────────────────────────────────────────
#  Silero
# ──────────────────────────────────────────────────────────────────────────────

class SileroVAD:

    def __init__(self):
        import torch
        self.torch = torch
        logger.info("Loading Silero VAD …")
        self.model, _ = torch.hub.load(
            repo_or_dir="snakers4/silero-vad",
            model="silero_vad",
            force_reload=False,
            onnx=False,
            trust_repo=True,
        )
        self.model.eval()
        self.name = "silero"

    def speech_prob(self, frame: np.ndarray) -> float:
        """frame: float32 mono, exactly FRAME_SAMPLES long, range [-1, 1]."""
        if len(frame) != FRAME_SAMPLES:
            frame = _fit(frame, FRAME_SAMPLES)
        with self.torch.no_grad():
            t = self.torch.from_numpy(frame.astype(np.float32))
            return float(self.model(t, SAMPLE_RATE).item())

    def reset(self):
        if hasattr(self.model, "reset_states"):
            self.model.reset_states()


# ──────────────────────────────────────────────────────────────────────────────
#  Energy fallback
# ──────────────────────────────────────────────────────────────────────────────

class EnergyVAD:
    """
    Adaptive-noise-floor RMS gate with a zero-crossing-rate sanity check.

    The noise floor tracks the quietest recent frames, so it adapts to a
    caller on a noisy street or a quiet office without reconfiguration.
    """

    SEED_FRAMES  = 8         # frames used to establish the initial floor
    ABS_GATE_LO  = 0.015     # rms clearly above any plausible room noise
    ABS_GATE_HI  = 0.060     # rms unambiguously speech at normal mic gain

    def __init__(self, sensitivity: float = 3.0):
        self.sensitivity = sensitivity      # how many x above floor = speech
        self.reset()
        self.name = "energy"

    def reset(self):
        self._floor      = None
        self._floor_init = False
        self._frames     = 0

    def speech_prob(self, frame: np.ndarray) -> float:
        frame = frame.astype(np.float32)
        rms = float(np.sqrt(np.mean(frame ** 2)) + 1e-9)

        # Zero-crossing rate — speech sits in a middling band; pure hiss is high,
        # DC/rumble is very low.
        zcr = float(np.mean(np.abs(np.diff(np.sign(frame)))) / 2.0)

        self._frames += 1

        # ── Seeding ──────────────────────────────────────────────────────────
        # Seed from the MINIMUM, never a running average. If the caller starts
        # talking part-way through the seed window (or is already talking when
        # the stream opens), an averaged seed is dragged up to speech level and
        # the detector goes deaf for the rest of the session.
        if not self._floor_init:
            self._floor = rms if self._floor is None else min(self._floor, rms)
            if self._frames >= self.SEED_FRAMES:
                self._floor_init = True
            # Still answer using the absolute gate, so speech during the seed
            # window is not silently swallowed.
            return self._absolute_prob(rms) * self._zcr_penalty(zcr)

        # ── Asymmetric adaptation ────────────────────────────────────────────
        # Down fast (room got quiet → recover in ~7 frames).
        # Up very slowly (~60s time constant), so a long unbroken utterance
        # cannot drag the floor up to its own level and mute itself.
        self._floor = max(
            0.85 * self._floor + 0.15 * rms if rms < self._floor
            else 0.9995 * self._floor + 0.0005 * rms,
            1e-6,
        )

        ratio    = rms / (self._floor * self.sensitivity)
        rel_prob = float(np.clip((ratio - 0.6) / 1.4, 0.0, 1.0))

        # Absolute gate is a floor on confidence, not a cap: if the signal is
        # loud in absolute terms it is speech regardless of what the adaptive
        # estimate believes.
        prob = max(rel_prob, self._absolute_prob(rms))
        return prob * self._zcr_penalty(zcr)

    # ── Helpers ───────────────────────────────────────────────────────────────

    def _absolute_prob(self, rms: float) -> float:
        return float(np.clip(
            (rms - self.ABS_GATE_LO) / (self.ABS_GATE_HI - self.ABS_GATE_LO),
            0.0, 1.0))

    @staticmethod
    def _zcr_penalty(zcr: float) -> float:
        """Penalise implausible zero-crossing rates (hiss, rumble, DC)."""
        return 0.4 if (zcr > 0.35 or zcr < 0.005) else 1.0


# ──────────────────────────────────────────────────────────────────────────────

def _fit(frame: np.ndarray, n: int) -> np.ndarray:
    """Pad or truncate a frame to exactly n samples."""
    if len(frame) >= n:
        return frame[:n]
    return np.pad(frame, (0, n - len(frame)))


def load_vad(prefer: str = "auto"):
    """
    prefer: 'auto' | 'silero' | 'energy'
    Never raises — falls back to EnergyVAD if Silero cannot be loaded
    (no network in the Space, torch.hub blocked, etc).
    """
    if prefer == "energy":
        return EnergyVAD()
    try:
        return SileroVAD()
    except Exception as e:
        logger.warning(f"Silero VAD unavailable ({e}); using energy VAD.")
        return EnergyVAD()