Add SmartTurn model selector

app.py

@@ -35,6 +35,22 @@ DEFAULT_PVAD_ONNX = MODELS / "pvad_core.onnx"
 DEFAULT_PVAD_H256_ONNX = MODELS / "pvad_core_h256.onnx"
 DEFAULT_SILERO_JIT = MODELS / "silero_vad.jit"
 DEFAULT_SMARTTURN_ONNX = MODELS / "smartturn-v3.1.onnx"
+CHECKPOINT50_SMARTTURN_ONNX = MODELS / "official_vi_plus_synthinc2000_headonly_lr1e5_checkpoint50_staticcalib8.onnx"
+LOCAL_CHECKPOINT50_SMARTTURN_ONNX = Path(
+    "/Utilisateurs/tnguye28/smartturn-vn/outputs/lumi_turn/onnx_exports/"
+    "official_vi_plus_synthinc2000_headonly_lr1e5_checkpoint50_staticcalib8/"
+    "model_int8_static_calib8.onnx"
+)
+EXAMPLE_SMARTTURN_ONNX = (
+    ROOT
+    / ".."
+    / "smartturn-vn"
+    / "outputs"
+    / "lumi_turn"
+    / "onnx_exports"
+    / "official_vi_plus_synthinc2000_headonly_lr1e5_checkpoint25_staticcalib8"
+    / "model_int8_static_calib8.onnx"
+).resolve()
 DUALTURN_MODEL_ID = "anyreach-ai/dualturn-qwen2.5-mimi-0.5B"
 DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -69,6 +85,10 @@ def default_audio_value(path: Path) -> str | None:
     return str(path) if path.exists() else None
 
 
+def checkpoint50_smartturn_path() -> Path:
+    return LOCAL_CHECKPOINT50_SMARTTURN_ONNX if LOCAL_CHECKPOINT50_SMARTTURN_ONNX.exists() else CHECKPOINT50_SMARTTURN_ONNX
+
+
 def write_wav(path: Path, audio: np.ndarray, sample_rate: int = 16000) -> str:
     path.parent.mkdir(parents=True, exist_ok=True)
     sf.write(str(path), np.asarray(audio, dtype=np.float32), sample_rate)
@@ -76,12 +96,13 @@ def write_wav(path: Path, audio: np.ndarray, sample_rate: int = 16000) -> str:
 
 
 def make_ort_session(path: str | Path) -> ort.InferenceSession:
+    path = str(path).strip()
     opts = ort.SessionOptions()
     opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
     opts.intra_op_num_threads = 1
     opts.inter_op_num_threads = 1
     opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-    return ort.InferenceSession(str(path), sess_options=opts, providers=["CPUExecutionProvider"])
+    return ort.InferenceSession(path, sess_options=opts, providers=["CPUExecutionProvider"])
 
 
 def frame_audio(audio: np.ndarray, frame_size: int = 512) -> np.ndarray:
@@ -327,6 +348,7 @@ def run_pipeline(
     smartturn_threshold: float,
     pvad_threshold: float,
     pvad_model_path: str,
+    smartturn_model_path: str,
     min_active_target_ms: float,
     silence_fallback_ms: float,
     asr_cut_silence_ms: float,
@@ -342,7 +364,7 @@ def run_pipeline(
         assistant = assistant[: int(round(assistant_max_playback_sec * 16000))]
 
     pvad = PvadOnnx(pvad_model_path, DEFAULT_SILERO_JIT)
-    smartturn = SmartTurnOnnx(DEFAULT_SMARTTURN_ONNX)
+    smartturn = SmartTurnOnnx(smartturn_model_path)
     dualturn = DualTurnHF(DUALTURN_MODEL_ID, device)
 
     pvad_probs = pvad.predict_target_probs(mic, enroll)
@@ -645,6 +667,7 @@ def run_gradio(
     smartturn_threshold: float,
     pvad_threshold: float,
     pvad_model_path: str,
+    smartturn_model_path: str,
     min_active_target_ms: float,
     silence_fallback_ms: float,
     asr_cut_silence_ms: float,
@@ -665,6 +688,7 @@ def run_gradio(
         smartturn_threshold=float(smartturn_threshold),
         pvad_threshold=float(pvad_threshold),
         pvad_model_path=pvad_model_path,
+        smartturn_model_path=smartturn_model_path,
         min_active_target_ms=float(min_active_target_ms),
         silence_fallback_ms=float(silence_fallback_ms),
         asr_cut_silence_ms=float(asr_cut_silence_ms),
@@ -722,7 +746,7 @@ def run_gradio(
         "models": {
             "pvad": str(pvad_model_path),
             "silero": str(DEFAULT_SILERO_JIT),
-            "smartturn": str(DEFAULT_SMARTTURN_ONNX),
+            "smartturn": str(smartturn_model_path),
             "dualturn": DUALTURN_MODEL_ID,
         },
         "state_counts": {state: sum(1 for r in result.rows if r["state"] == state) for state in STATE_COLORS},
@@ -744,6 +768,12 @@ def select_cut_audio(cut_path: str | None) -> str | None:
 
 
 def build_app() -> gr.Blocks:
+    default_smartturn = EXAMPLE_SMARTTURN_ONNX if EXAMPLE_SMARTTURN_ONNX.exists() else DEFAULT_SMARTTURN_ONNX
+    checkpoint50_smartturn = checkpoint50_smartturn_path()
+    smartturn_choices = [
+        ("original", str(default_smartturn)),
+        ("finetuned", str(checkpoint50_smartturn)),
+    ]
     with gr.Blocks(title="Half Duuplex Demo") as demo:
         gr.Markdown("## Half Duuplex Demo")
         with gr.Row():
@@ -766,6 +796,11 @@ def build_app() -> gr.Blocks:
                 label="PVAD ONNX model",
             )
             check_ms = gr.Slider(80, 1000, value=240, step=80, label="Model check interval ms")
+        smartturn_model = gr.Dropdown(
+            choices=smartturn_choices,
+            value=str(default_smartturn),
+            label="SmartTurn ONNX model",
+        )
         with gr.Row():
             min_active = gr.Slider(0, 2000, value=300, step=50, label="Min active target ms")
             silence_fallback = gr.Slider(0, 3000, value=800, step=100, label="Silence fallback END ms")
@@ -791,6 +826,7 @@ def build_app() -> gr.Blocks:
                 smart_threshold,
                 pvad_threshold,
                 pvad_model,
+                smartturn_model,
                 min_active,
                 silence_fallback,
                 asr_cut_silence,