app.py

"""Gradio demo: live VAD -> Smart Turn v3.2 turn-detection pipeline, run offline.

This mirrors what a live voice pipeline does in a real call:

    audio chunks --> Silero VAD --> (after `stop_secs` of silence) --> Smart Turn

Record or upload a clip; it is replayed in 20 ms chunks through the *same* Silero
VAD state machine and Smart Turn audio-buffering logic used live (see
turn_sim.py). Every time VAD reports the user stopped speaking, the buffered
segment is handed to Smart Turn.

All three models run independently and in parallel on the same input audio, so
you can compare their end-of-turn decisions side by side.

All VAD and Smart Turn parameters are editable in the UI (collapsed by default).

Run:  ./venv/bin/python app.py   then open http://127.0.0.1:7860
"""

import librosa
import numpy as np
import gradio as gr

from predict import SmartTurn, SAMPLE_RATE
from turn_sim import VADParams, SmartTurnParams, simulate

# (dropdown label -> model filename in models/)
MODELS = {
    "Trained — int8 (cpu, 8.8 MB)": "trained-int8.onnx",
    "Trained — fp32 (gpu, 32 MB)": "trained-fp32.onnx",
    "Baseline — smart-turn-v3.2-cpu": "smart-turn-v3.2-cpu.onnx",
}
MAX_TURNS = 20  # audio-player slots per model; every run explicitly fills/hides all of them
_cache = {}


def get_model(name):
    if name not in _cache:
        _cache[name] = SmartTurn(model_path=f"models/{name}")
    return _cache[name]


def run(
    audio_path,
    # VAD params
    vad_confidence,
    vad_start_secs,
    vad_stop_secs,
    vad_min_volume,
    # Smart Turn params
    st_pre_speech_ms,
    st_max_duration_secs,
    st_fallback_stop_secs,
    threshold,
    # sim params
    chunk_ms,
    trailing_ms,
    eval_at_end,
):
    hidden = [gr.update(visible=False, value=None) for _ in range(MAX_TURNS * len(MODELS))]
    if not audio_path:
        return hidden

    wav, _ = librosa.load(audio_path, sr=SAMPLE_RATE, mono=True)
    if wav.size == 0:
        return hidden
    if trailing_ms > 0:
        wav = np.concatenate(
            [wav, np.zeros(int(trailing_ms / 1000 * SAMPLE_RATE), dtype=np.float32)]
        )

    vad_params = VADParams(
        confidence=vad_confidence,
        start_secs=vad_start_secs,
        stop_secs=vad_stop_secs,
        min_volume=vad_min_volume,
    )
    st_params = SmartTurnParams(
        stop_secs=st_fallback_stop_secs,
        pre_speech_ms=st_pre_speech_ms,
        max_duration_secs=st_max_duration_secs,
    )

    # Run each model through its own independent simulation (own VAD/buffer
    # state), so all three models' end-of-turn decisions are shown in parallel.
    all_updates = []
    for model_label, model_file in MODELS.items():
        model = get_model(model_file)
        events, _ = simulate(
            wav,
            model,
            vad_params,
            st_params,
            threshold=threshold,
            chunk_ms=chunk_ms,
            eval_at_end=eval_at_end,
        )

        updates = []
        for i, e in enumerate(events, 1):
            if len(updates) >= MAX_TURNS or e.get("segment") is None:
                continue
            seg_int16 = np.clip(e["segment"] * 32768.0, -32768, 32767).astype(np.int16)
            verdict = "✅ COMPLETE" if e["complete"] else "🔴 INCOMPLETE"
            updates.append(
                gr.update(
                    visible=True,
                    value=(SAMPLE_RATE, seg_int16),
                    label=(
                        f"Turn {i} @ {e['t']:.2f}s — {verdict} "
                        f"(P={e['probability']:.4f}, seg={e['segment_secs']:.2f}s)"
                    ),
                )
            )
        while len(updates) < MAX_TURNS:
            updates.append(gr.update(visible=False, value=None))
        all_updates.extend(updates)

    return all_updates


with gr.Blocks(title="VAD + Smart Turn simulator") as demo:
    gr.Markdown(
        "# Turn-detection simulator — VAD ➜ Smart Turn\n"
        "Replays your audio through the **exact** Silero VAD state machine and "
        "Smart Turn buffering used in a live voice pipeline, independently for "
        "all three models. When VAD detects the user stopped speaking, the "
        "buffered segment is sent to Smart Turn and each model's decision is "
        "shown below, alongside the audio it was given."
    )
    with gr.Row():
        with gr.Column(scale=1):
            audio_in = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                label="Audio (upload or record)",
            )

            with gr.Accordion("Parameters", open=False):
                gr.Markdown("**VAD (Silero)**")
                vad_conf = gr.Slider(0.0, 1.0, value=0.75, step=0.01, label="confidence")
                vad_start = gr.Slider(
                    0.0, 1.0, value=0.3, step=0.05, label="start_secs (speech-start confirm)"
                )
                vad_stop = gr.Slider(
                    0.0, 1.0, value=0.2, step=0.05,
                    label="stop_secs (silence before VAD stop → triggers Smart Turn)",
                )
                vad_vol = gr.Slider(0.0, 1.0, value=0.6, step=0.01, label="min_volume")

                gr.Markdown("**Smart Turn**")
                st_pre = gr.Slider(
                    0, 2000, value=500, step=50, label="pre_speech_ms (lead-in)"
                )
                st_max = gr.Slider(
                    1.0, 16.0, value=8.0, step=0.5, label="max_duration_secs (model window cap)"
                )
                st_fallback = gr.Slider(
                    0.5, 10.0, value=3.0, step=0.5,
                    label="stop_secs fallback (force end-of-turn on long silence)",
                )
                thr = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="decision threshold")

                gr.Markdown("**Simulation**")
                chunk = gr.Slider(
                    10, 100, value=20, step=10, label="chunk_ms (transport frame size)"
                )
                sil = gr.Slider(
                    0, 2000, value=0, step=50, label="append trailing silence (ms)"
                )
                eval_end = gr.Checkbox(
                    value=True,
                    label="Also run at end of recording (force a prediction on the "
                    "full audio even if VAD never reports a stop)",
                )

            btn = gr.Button("Run pipeline", variant="primary")

        with gr.Column(scale=2):
            turn_players = []
            with gr.Row():
                for model_label in MODELS:
                    with gr.Column():
                        gr.Markdown(f"**{model_label}**")
                        for _ in range(MAX_TURNS):
                            turn_players.append(
                                gr.Audio(visible=False, interactive=False, type="numpy")
                            )

    inputs = [
        audio_in,
        vad_conf, vad_start, vad_stop, vad_vol,
        st_pre, st_max, st_fallback, thr,
        chunk, sil, eval_end,
    ]
    outputs = turn_players
    btn.click(run, inputs, outputs)
    audio_in.stop_recording(run, inputs, outputs)


if __name__ == "__main__":
    demo.launch()