Просмотр исходного кода

feat(vad): replace energy-based VAD with ONNX Silero VAD (xiaozhi-style)

- ONNX Silero VAD model loaded via onnxruntime (per-participant state)
- Dual-threshold hysteresis (0.5/0.2) + sliding window (5 frames, 50%)
- Automatic fallback to energy-based VAD if ONNX model not available
- Shared ONNX session across all participants (thread-safe)
- Model path, thresholds configurable via env vars
- State reset on each utterance boundary
wenhongquan 3 недель назад
Родитель
Сommit
7a424e632b
1 измененных файлов с 120 добавлено и 55 удалено
  1. 120 55
      asr_agent/conversation_worker.py

+ 120 - 55
asr_agent/conversation_worker.py

@@ -1,6 +1,6 @@
 """
 Full Pipeline Worker: VAD -> ASR -> LLM -> TTS
-  VAD: dual-threshold energy-based + sliding window (xiaozhi-inspired)
+  VAD: Silero VAD (ONNX) with dual-threshold + sliding window
   VP:  AEC/ANS/AGC handled by LiveKit WebRTC
   ASR: Qwen3-ASR (local model, batch on speech-end)
   LLM: streaming via OpenAI-compatible API (vLLM), <think> tag filtering
@@ -20,9 +20,15 @@ import time
 from collections import deque
 from dataclasses import dataclass, field
 
+import numpy as np
+try:
+    import onnxruntime
+    _HAS_ONNX = True
+except ImportError:
+    _HAS_ONNX = False
+
 import aiohttp
 import jwt
-import numpy as np
 from livekit import rtc
 
 _root = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "whisper-qt-client", "whisper_asr")
@@ -41,14 +47,6 @@ VLLM_URL = os.environ.get("VLLM_URL", "http://127.0.0.1:8000/v1")
 LLM_MODEL = os.environ.get("LLM_MODEL", "qwen3.6-35b-awq")
 MIMO_KEY = os.environ.get("MIMO_KEY", "")
 
-# ── VAD ──
-VAD_THRESHOLD_HIGH = 0.015
-VAD_THRESHOLD_LOW  = 0.005
-VAD_WINDOW_SIZE    = 8
-VAD_VOICE_RATIO    = 0.5
-MIN_SPEECH_S       = 0.3
-MIN_SILENCE_S      = 0.8
-MAX_SPEECH_S       = 8.0
 
 # ── LLM ──
 LLM_MAX_TOKENS  = 128
@@ -198,56 +196,123 @@ class _TTSSegmenter:
         return remaining
 
 
+
+# ── VAD ──
+VAD_MODEL_PATH = os.environ.get("VAD_MODEL_PATH", "/data/silero_vad/silero_vad.onnx")
+VAD_THRESHOLD_HIGH = float(os.environ.get("VAD_THRESHOLD_HIGH", "0.5"))
+VAD_THRESHOLD_LOW  = float(os.environ.get("VAD_THRESHOLD_LOW", "0.2"))
+VAD_WINDOW_SIZE    = int(os.environ.get("VAD_WINDOW_SIZE", "5"))
+VAD_VOICE_RATIO    = float(os.environ.get("VAD_VOICE_RATIO", "0.5"))
+MIN_SPEECH_S       = float(os.environ.get("MIN_SPEECH_S", "0.3"))
+MIN_SILENCE_S      = float(os.environ.get("MIN_SILENCE_S", "0.8"))
+MAX_SPEECH_S       = float(os.environ.get("MAX_SPEECH_S", "8.0"))
+
+# Global ONNX session (shared across participants)
+_vad_session = None
+
+
+def _get_vad_session():
+    global _vad_session
+    if _vad_session is not None:
+        return _vad_session
+    if not _HAS_ONNX or not os.path.exists(VAD_MODEL_PATH):
+        return None
+    opts = onnxruntime.SessionOptions()
+    opts.inter_op_num_threads = 1
+    opts.intra_op_num_threads = 1
+    _vad_session = onnxruntime.InferenceSession(
+        VAD_MODEL_PATH, providers=["CPUExecutionProvider"], sess_options=opts,
+    )
+    logger.info("Silero VAD loaded (%s)", VAD_MODEL_PATH)
+    return _vad_session
+
+
 # ═══════════════════════════════════════════════════════════════
-#  VAD (dual-threshold + sliding window)
+#  Silero VAD (ONNX, per-participant state)
 # ═══════════════════════════════════════════════════════════════
 
-@dataclass
-class _VADState:
-    window: deque = field(default_factory=lambda: deque(maxlen=VAD_WINDOW_SIZE))
-    in_speech: bool = False
-    speech_start_sample: int = 0
-    silence_counter: int = 0
-    total_samples: int = 0
+class _SileroVAD:
+    """Xiaozhi-style Silero VAD with dual-threshold + sliding window."""
 
-    def reset(self):
-        self.window.clear()
+    def __init__(self):
+        self._sess = _get_vad_session()
+        self._state = np.zeros((2, 1, 128), dtype=np.float32)
+        self._context = np.zeros((1, 64), dtype=np.float32)
+        self._window: deque = deque(maxlen=max(VAD_WINDOW_SIZE, 1))
         self.in_speech = False
         self.speech_start_sample = 0
         self.silence_counter = 0
+        self.total_samples = 0
+        self._sr = np.array(16000, dtype=np.int64)
+
+    def process(self, frame: np.ndarray) -> bool:
+        """Process 160-sample (10ms) frame. Call every frame. Returns in_speech."""
+        n = len(frame)
+
+        if self._sess is not None:
+            # ── Silero ONNX inference ──
+            # Silero processes 512-sample windows. Accumulate frames.
+            if not hasattr(self, '_buf'):
+                self._buf = np.array([], dtype=np.float32)
+            self._buf = np.concatenate([self._buf, frame]).astype(np.float32)
+
+            is_voice = False
+            while len(self._buf) >= 512:
+                chunk = self._buf[:512]
+                self._buf = self._buf[512:]
+                audio_in = chunk.reshape(1, -1)
+                inp = np.concatenate([self._context, audio_in], axis=1).astype(np.float32)
+                out, self._state = self._sess.run(
+                    None, {"input": inp, "state": self._state, "sr": self._sr},
+                )
+                self._context = inp[:, -64:]
+                prob = out.item()
+                is_voice = prob >= VAD_THRESHOLD_HIGH or (
+                    prob > VAD_THRESHOLD_LOW and (
+                        self._window[-1] if self._window else False
+                    )
+                )
+                self._window.append(is_voice)
+        else:
+            # ── Fallback: energy-based ──
+            energy = float(np.sqrt(np.mean(frame ** 2)))
+            if energy > 0.015:
+                is_voice = True
+            elif energy < 0.005:
+                is_voice = False
+            else:
+                is_voice = self._window[-1] if self._window else False
+            self._window.append(is_voice)
+
+        have_voice = sum(self._window) / max(len(self._window), 1) >= VAD_VOICE_RATIO
+
+        if have_voice and not self.in_speech:
+            logger.info("VAD START")
+            self.in_speech = True
+            self.speech_start_sample = self.total_samples
+            self.silence_counter = 0
+        elif not have_voice and self.in_speech:
+            self.silence_counter += n
+        elif have_voice:
+            self.silence_counter = 0
+        self.total_samples += n
+        return self.in_speech
+
+    def should_end(self) -> bool:
+        return self.in_speech and self.silence_counter >= int(MIN_SILENCE_S * 16000)
+
+    def min_speech_met(self) -> bool:
+        return (self.total_samples - self.speech_start_sample) / 16000 >= MIN_SPEECH_S
 
-
-def _vad_process(state: _VADState, frame: np.ndarray, sample_rate: int) -> bool:
-    n = len(frame)
-    energy = float(np.sqrt(np.mean(frame ** 2)))
-    if energy > VAD_THRESHOLD_HIGH:
-        is_voice = True
-    elif energy < VAD_THRESHOLD_LOW:
-        is_voice = False
-    else:
-        is_voice = state.window[-1] if state.window else False
-    state.window.append(is_voice)
-    have_voice = sum(state.window) / max(len(state.window), 1) >= VAD_VOICE_RATIO
-
-    if have_voice and not state.in_speech:
-        logger.info("VAD START")
-        state.in_speech = True
-        state.speech_start_sample = state.total_samples
-        state.silence_counter = 0
-    elif not have_voice and state.in_speech:
-        state.silence_counter += n
-    elif have_voice:
-        state.silence_counter = 0
-    state.total_samples += n
-    return state.in_speech
-
-
-def _vad_should_end(state: _VADState, sample_rate: int) -> bool:
-    return state.in_speech and state.silence_counter >= int(MIN_SILENCE_S * sample_rate)
-
-
-def _vad_min_speech_met(state: _VADState, sample_rate: int) -> bool:
-    return (state.total_samples - state.speech_start_sample) / sample_rate >= MIN_SPEECH_S
+    def reset(self):
+        self._window.clear()
+        self.in_speech = False
+        self.speech_start_sample = 0
+        self.silence_counter = 0
+        self._state = np.zeros((2, 1, 128), dtype=np.float32)
+        self._context = np.zeros((1, 64), dtype=np.float32)
+        if hasattr(self, '_buf'):
+            self._buf = np.array([], dtype=np.float32)
 
 
 # ═══════════════════════════════════════════════════════════════
@@ -298,7 +363,7 @@ class Worker:
         logger.info("_run %s", sid)
         buf = AudioBuffer(max_duration=MAX_SPEECH_S + 4, sample_rate=16000)
         sr = 16000
-        vad = _VADState()
+        vad = _SileroVAD()
         busy = False
         cur_tts = None
 
@@ -314,14 +379,14 @@ class Worker:
                 buf.append(arr)
 
                 was_speech = vad.in_speech
-                _vad_process(vad, arr, sr)
+                vad.process(arr)
 
                 if fc % 100 == 0:
                     logger.info("[%s] af#%d e=%.4f vad=%s blen=%d",
                                 sid, fc, float(np.sqrt(np.mean(arr ** 2))),
                                 vad.in_speech, len(buf.buffer))
 
-                if was_speech and _vad_should_end(vad, sr) and _vad_min_speech_met(vad, sr):
+                if was_speech and vad.should_end() and vad.min_speech_met():
                     logger.info("[%s] VAD END blen=%d", sid, len(buf.buffer))
                     if cur_tts:
                         cur_tts = None