turn_sim.py

"""Offline simulation of pipecat's VAD -> Smart Turn v3 pipeline.

This replays an audio file through the *same* logic pipecat runs live, so you can
see exactly when a turn would end. It faithfully ports, from the cloned repo:

  - Silero VAD model wrapper      pipecat/audio/vad/silero.py        (SileroOnnxModel)
  - VAD state machine + params    pipecat/audio/vad/vad_analyzer.py  (VADAnalyzer / VADParams)
  - loudness gate                 pipecat/audio/utils.py             (calculate_audio_volume)
  - Smart Turn audio buffering    pipecat/audio/turn/smart_turn/base_smart_turn.py
  - Smart Turn v3 inference       reused from predict.py (mirrors local_smart_turn_v3.py)

Instead of a live transport we feed the waveform in fixed-size chunks (default
20 ms, matching pipecat's local transport). Every chunk is pushed through VAD;
when VAD reports the user stopped speaking (after `stop_secs` of silence) the
buffered segment is handed to Smart Turn, exactly like
``TurnAnalyzerUserTurnStopStrategy`` does in pipecat.

Time is simulated from the sample position (NOT wall-clock), so results are
deterministic and independent of how fast the file is processed.
"""

import os
import warnings
from dataclasses import dataclass
from enum import Enum

import numpy as np
import onnxruntime as ort

from predict import SAMPLE_RATE, SmartTurn  # Smart Turn v3 inference (faithful port)

# Silero VAD model. Prefer a local copy next to this file (used when deployed,
# e.g. on HF Spaces); otherwise fall back to the one bundled in the cloned
# pipecat repo (local development).
_HERE = os.path.dirname(os.path.abspath(__file__))
_LOCAL_SILERO = os.path.join(_HERE, "silero_vad.onnx")
SILERO_VAD_PATH = (
    _LOCAL_SILERO
    if os.path.exists(_LOCAL_SILERO)
    else os.path.join(
        _HERE, "pipecat-exploration/pipecat/src/pipecat/audio/vad/data/silero_vad.onnx"
    )
)

# pyloudnorm emits cosmetic warnings on short/quiet blocks; pipecat ignores them.
warnings.filterwarnings("ignore")
import pyloudnorm as pyln  # noqa: E402


# --------------------------------------------------------------------------- #
# Parameters (mirror pipecat's VADParams and SmartTurnParams defaults)
# --------------------------------------------------------------------------- #
@dataclass
class VADParams:
    """pipecat/audio/vad/vad_analyzer.py defaults."""

    confidence: float = 0.7
    start_secs: float = 0.2
    stop_secs: float = 0.2
    min_volume: float = 0.6


@dataclass
class SmartTurnParams:
    """pipecat/audio/turn/smart_turn/base_smart_turn.py defaults."""

    stop_secs: float = 3.0  # silence-timeout fallback (forces end-of-turn)
    pre_speech_ms: float = 500.0  # lead-in audio prepended to the segment
    max_duration_secs: float = 8.0  # cap on segment length fed to the model


# --------------------------------------------------------------------------- #
# Audio loudness gate (ported from pipecat/audio/utils.py)
# --------------------------------------------------------------------------- #
def _normalize_value(value, min_value, max_value):
    normalized = (value - min_value) / (max_value - min_value)
    return max(0.0, min(1.0, normalized))


def _calculate_audio_volume(audio_int16: np.ndarray, sample_rate: int) -> float:
    audio_float = audio_int16.astype(np.float64)
    block_size = audio_int16.size / sample_rate
    try:
        meter = pyln.Meter(sample_rate, block_size=block_size)
        loudness = meter.integrated_loudness(audio_float)
    except Exception:
        return 0.0
    if not np.isfinite(loudness):
        return 0.0
    return _normalize_value(loudness, -20, 80)


def _exp_smoothing(value, prev_value, factor):
    return prev_value + factor * (value - prev_value)


# --------------------------------------------------------------------------- #
# Silero ONNX wrapper (ported verbatim from pipecat/audio/vad/silero.py)
# --------------------------------------------------------------------------- #
class _SileroOnnxModel:
    def __init__(self, path, force_onnx_cpu=True):
        opts = ort.SessionOptions()
        opts.inter_op_num_threads = 1
        opts.intra_op_num_threads = 1
        self.session = ort.InferenceSession(
            path, providers=["CPUExecutionProvider"], sess_options=opts
        )
        self.sample_rates = [8000, 16000]
        self.reset_states()

    def reset_states(self, batch_size=1):
        self._state = np.zeros((2, batch_size, 128), dtype="float32")
        self._context = np.zeros((batch_size, 0), dtype="float32")
        self._last_sr = 0
        self._last_batch_size = 0

    def __call__(self, x, sr: int):
        if np.ndim(x) == 1:
            x = np.expand_dims(x, 0)
        num_samples = 512 if sr == 16000 else 256
        batch_size = np.shape(x)[0]
        context_size = 64 if sr == 16000 else 32

        if not self._last_batch_size:
            self.reset_states(batch_size)
        if self._last_sr and self._last_sr != sr:
            self.reset_states(batch_size)
        if self._last_batch_size and self._last_batch_size != batch_size:
            self.reset_states(batch_size)

        if not np.shape(self._context)[1]:
            self._context = np.zeros((batch_size, context_size), dtype="float32")

        x = np.concatenate((self._context, x), axis=1)
        ort_inputs = {"input": x, "state": self._state, "sr": np.array(sr, dtype="int64")}
        out, state = self.session.run(None, ort_inputs)
        self._state = state
        self._context = x[..., -context_size:]
        self._last_sr = sr
        self._last_batch_size = batch_size
        return out


# --------------------------------------------------------------------------- #
# VAD state machine (ported from pipecat/audio/vad/vad_analyzer.py)
# --------------------------------------------------------------------------- #
class VADState(Enum):
    QUIET = 1
    STARTING = 2
    SPEAKING = 3
    STOPPING = 4


class SileroVAD:
    """Silero VAD with pipecat's exact start/stop confirmation state machine."""

    def __init__(self, sample_rate: int, params: VADParams):
        self.sample_rate = sample_rate
        self._params = params
        self._model = _SileroOnnxModel(SILERO_VAD_PATH)
        self._vad_buffer = b""
        self._smoothing_factor = 0.2
        self._prev_volume = 0.0
        self._set_params()

    def _num_frames_required(self) -> int:
        return 512 if self.sample_rate == 16000 else 256

    def _set_params(self):
        self._vad_frames = self._num_frames_required()
        self._vad_frames_num_bytes = self._vad_frames * 2  # mono int16
        vad_frames_per_sec = self._vad_frames / self.sample_rate
        self._vad_start_frames = round(self._params.start_secs / vad_frames_per_sec)
        self._vad_stop_frames = round(self._params.stop_secs / vad_frames_per_sec)
        self._vad_starting_count = 0
        self._vad_stopping_count = 0
        self._vad_state = VADState.QUIET

    def _voice_confidence(self, audio_int16: np.ndarray) -> float:
        audio_float32 = audio_int16.astype(np.float32) / 32768.0
        return float(np.asarray(self._model(audio_float32, self.sample_rate)).flatten()[0])

    def analyze_audio(self, buffer: bytes) -> VADState:
        """Feed one chunk; returns the current confirmed VAD state."""
        self._vad_buffer += buffer
        n = self._vad_frames_num_bytes
        if len(self._vad_buffer) < n:
            return self._vad_state

        while len(self._vad_buffer) >= n:
            frame_bytes = self._vad_buffer[:n]
            self._vad_buffer = self._vad_buffer[n:]
            audio_int16 = np.frombuffer(frame_bytes, dtype=np.int16)

            confidence = self._voice_confidence(audio_int16)
            volume = _exp_smoothing(
                _calculate_audio_volume(audio_int16, self.sample_rate),
                self._prev_volume,
                self._smoothing_factor,
            )
            self._prev_volume = volume

            speaking = (
                confidence >= self._params.confidence and volume >= self._params.min_volume
            )

            if speaking:
                if self._vad_state == VADState.QUIET:
                    self._vad_state = VADState.STARTING
                    self._vad_starting_count = 1
                elif self._vad_state == VADState.STARTING:
                    self._vad_starting_count += 1
                elif self._vad_state == VADState.STOPPING:
                    self._vad_state = VADState.SPEAKING
                    self._vad_stopping_count = 0
            else:
                if self._vad_state == VADState.STARTING:
                    self._vad_state = VADState.QUIET
                    self._vad_starting_count = 0
                elif self._vad_state == VADState.SPEAKING:
                    self._vad_state = VADState.STOPPING
                    self._vad_stopping_count = 1
                elif self._vad_state == VADState.STOPPING:
                    self._vad_stopping_count += 1

            if (
                self._vad_state == VADState.STARTING
                and self._vad_starting_count >= self._vad_start_frames
            ):
                self._vad_state = VADState.SPEAKING
                self._vad_starting_count = 0

            if (
                self._vad_state == VADState.STOPPING
                and self._vad_stopping_count >= self._vad_stop_frames
            ):
                self._vad_state = VADState.QUIET
                self._vad_stopping_count = 0

        return self._vad_state


# --------------------------------------------------------------------------- #
# Smart Turn audio buffer (ported from base_smart_turn.py, sim-clock version)
# --------------------------------------------------------------------------- #
class EndOfTurnState(Enum):
    COMPLETE = 1
    INCOMPLETE = 2


class SmartTurnBuffer:
    """Replicates BaseSmartTurn.append_audio + _process_speech_segment.

    Uses a simulated clock (sample position) instead of time.monotonic().
    """

    def __init__(self, sample_rate: int, params: SmartTurnParams):
        self.sample_rate = sample_rate
        self._params = params
        self._stop_ms = params.stop_secs * 1000
        self._audio_buffer = []  # list of (t_seconds, int16 ndarray)
        self._speech_triggered = False
        self._silence_ms = 0.0
        self._speech_start_time = 0.0
        self._vad_start_secs = 0.0

    def update_vad_start_secs(self, v: float):
        self._vad_start_secs = v

    def append_audio(self, audio_int16: np.ndarray, is_speech: bool, t: float) -> EndOfTurnState:
        """Append one chunk at sim time `t`. Returns COMPLETE only on silence timeout."""
        self._audio_buffer.append((t, audio_int16))
        state = EndOfTurnState.INCOMPLETE

        if is_speech:
            self._silence_ms = 0.0
            self._speech_triggered = True
            if self._speech_start_time == 0:
                self._speech_start_time = t
        else:
            if self._speech_triggered:
                chunk_ms = len(audio_int16) / (self.sample_rate / 1000)
                self._silence_ms += chunk_ms
                if self._silence_ms >= self._stop_ms:
                    state = EndOfTurnState.COMPLETE
                    self._clear(state)
            else:
                max_buffer_time = (
                    (self._params.pre_speech_ms / 1000)
                    + self._params.stop_secs
                    + self._params.max_duration_secs
                )
                while self._audio_buffer and self._audio_buffer[0][0] < t - max_buffer_time:
                    self._audio_buffer.pop(0)
        return state

    def analyze_end_of_turn(self, model: SmartTurn, threshold: float):
        """Run Smart Turn on the buffered segment. Returns a result dict or None."""
        if not self._audio_buffer:
            return None

        effective_pre_speech_ms = self._params.pre_speech_ms + (self._vad_start_secs * 1000)
        start_time = self._speech_start_time - (effective_pre_speech_ms / 1000)
        start_index = 0
        for i, (t, _) in enumerate(self._audio_buffer):
            if t >= start_time:
                start_index = i
                break
        end_index = len(self._audio_buffer) - 1

        chunks = [c for _, c in self._audio_buffer[start_index : end_index + 1]]
        segment_int16 = np.concatenate(chunks)
        segment = segment_int16.astype(np.float32) / 32768.0

        max_samples = int(self._params.max_duration_secs * self.sample_rate)
        if len(segment) > max_samples:
            segment = segment[-max_samples:]

        if len(segment) == 0:
            return None

        prob = model.predict(segment)["probability"]
        complete = prob >= threshold
        result = {
            "probability": prob,
            "complete": complete,
            "segment_secs": len(segment) / self.sample_rate,
            "segment": segment,  # the exact audio fed to Smart Turn (float32, 16 kHz)
        }
        if complete:
            self._clear(EndOfTurnState.COMPLETE)
        return result

    def has_pending_speech(self) -> bool:
        """True if the buffer holds speech that was never finalized by a VAD stop."""
        return (
            self._speech_triggered
            and len(self._audio_buffer) > 0
            and self._speech_start_time > 0
        )

    def _clear(self, state: EndOfTurnState):
        self._speech_triggered = state == EndOfTurnState.INCOMPLETE
        self._audio_buffer = []
        self._speech_start_time = 0.0
        self._silence_ms = 0.0


# --------------------------------------------------------------------------- #
# Driver
# --------------------------------------------------------------------------- #
def simulate(
    wav: np.ndarray,
    model: SmartTurn,
    vad_params: VADParams,
    st_params: SmartTurnParams,
    threshold: float = 0.5,
    chunk_ms: float = 20.0,
    eval_at_end: bool = True,
    sample_rate: int = SAMPLE_RATE,
):
    """Replay `wav` (mono float32, 16 kHz) through VAD -> Smart Turn.

    Returns (events, trace) where each event is one end-of-turn the pipeline
    would have fired: VAD-stop time, segment fed to the model, P(complete),
    verdict, and whether it came from the model or the silence-timeout fallback.

    If `eval_at_end` is True, Smart Turn is *also* run when the recording ends,
    even if VAD never reported a stop (e.g. the clip cuts off mid-speech). This
    guarantees a prediction on the full audio. Such events are tagged
    ``end_of_recording``.
    """
    vad = SileroVAD(sample_rate, vad_params)
    buf = SmartTurnBuffer(sample_rate, st_params)

    chunk_samples = max(1, int(sample_rate * chunk_ms / 1000))
    vad_user_speaking = False
    prev_state = VADState.QUIET
    last_speaking_t = 0.0  # sim time VAD was last actively SPEAKING
    events = []
    trace = []  # (t, vad_state, speaking) for optional inspection

    n = len(wav)
    for i in range(0, n, chunk_samples):
        chunk = wav[i : i + chunk_samples]
        if len(chunk) == 0:
            break
        t_start = i / sample_rate
        t_end = (i + len(chunk)) / sample_rate
        chunk_int16 = np.clip(chunk * 32768.0, -32768, 32767).astype(np.int16)

        # --- VAD: detect confirmed started/stopped transitions ---
        new_state = vad.analyze_audio(chunk_int16.tobytes())
        if new_state == VADState.SPEAKING:
            last_speaking_t = t_end
        event = None
        if new_state != prev_state and new_state in (VADState.SPEAKING, VADState.QUIET):
            event = "started" if new_state == VADState.SPEAKING else "stopped"
            prev_state = new_state

        if event == "started":
            vad_user_speaking = True
            buf.update_vad_start_secs(vad_params.start_secs)

        # --- Smart Turn buffer accumulates every chunk ---
        st_state = buf.append_audio(chunk_int16, vad_user_speaking, t_end)
        if st_state == EndOfTurnState.COMPLETE:
            # Silence-timeout fallback: turn forced complete WITHOUT the model.
            silence_secs = round(st_params.stop_secs, 3)
            msg = (
                f"VAD detected silence for {silence_secs:.2f}s "
                f"(fallback timeout) and Smart Turn was skipped "
                f"-> turn forced COMPLETE"
            )
            print(f"[turn {len(events) + 1}] {msg}", flush=True)
            events.append(
                {
                    "t": round(t_end, 3),
                    "trigger": "silence_timeout",
                    "silence_secs": silence_secs,
                    "probability": None,
                    "complete": True,
                    "segment_secs": None,
                    "segment": None,
                    "message": msg,
                }
            )

        if event == "stopped":
            vad_user_speaking = False
            silence_secs = round(t_end - last_speaking_t, 3)
            res = buf.analyze_end_of_turn(model, threshold)
            if res is not None:
                verdict = "COMPLETE" if res["complete"] else "INCOMPLETE"
                msg = (
                    f"VAD detected silence for {silence_secs:.2f}s and "
                    f"Smart Turn predicted {verdict} "
                    f"(P(complete)={res['probability']:.4f})"
                )
                print(f"[turn {len(events) + 1}] {msg}", flush=True)
                events.append(
                    {
                        "t": round(t_end, 3),
                        "trigger": "vad_stop",
                        "silence_secs": silence_secs,
                        "probability": res["probability"],
                        "complete": res["complete"],
                        "segment_secs": res["segment_secs"],
                        "segment": res["segment"],
                        "message": msg,
                    }
                )

        trace.append((round(t_start, 3), new_state.name, vad_user_speaking))

    # --- End-of-recording trigger ---------------------------------------- #
    # Real transports fire on a VAD stop. When the audio simply ends (no
    # trailing silence), VAD may never report a stop, so Smart Turn would never
    # run. If requested, force one final evaluation on the full audio.
    if eval_at_end:
        t_final = n / sample_rate
        if buf.has_pending_speech():
            # Speech is still buffered (clip ended mid-utterance, or after an
            # INCOMPLETE stop) -> evaluate exactly what pipecat holds.
            res = buf.analyze_end_of_turn(model, threshold)
        elif not events:
            # Nothing buffered and nothing ever fired (e.g. VAD ignored a quiet
            # clip) -> evaluate the last 8 s of the whole recording directly.
            seg = wav[-int(st_params.max_duration_secs * sample_rate) :]
            if len(seg) > 0:
                prob = model.predict(seg)["probability"]
                res = {
                    "probability": prob,
                    "complete": prob >= threshold,
                    "segment_secs": len(seg) / sample_rate,
                    "segment": seg.astype(np.float32),
                }
            else:
                res = None
        else:
            res = None

        if res is not None:
            verdict = "COMPLETE" if res["complete"] else "INCOMPLETE"
            msg = (
                f"Recording ended (no VAD stop) and Smart Turn predicted "
                f"{verdict} (P(complete)={res['probability']:.4f})"
            )
            print(f"[turn {len(events) + 1}] {msg}", flush=True)
            events.append(
                {
                    "t": round(t_final, 3),
                    "trigger": "end_of_recording",
                    "silence_secs": None,
                    "probability": res["probability"],
                    "complete": res["complete"],
                    "segment_secs": res["segment_secs"],
                    "segment": res["segment"],
                    "message": msg,
                }
            )

    return events, trace