248 lines
9.0 KiB
Python
248 lines
9.0 KiB
Python
"""Adapt the official Silero iterator to continuous PCM and async lifecycle."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import math
|
|
import struct
|
|
from collections.abc import Callable
|
|
from typing import Protocol
|
|
|
|
from voice_assistant.vad.adapter import VADAdapter
|
|
from voice_assistant.vad.models import SpeechEnded, SpeechStarted, VADEvent
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_WINDOW_SAMPLES = 512
|
|
_WINDOW_BYTES = 2 * _WINDOW_SAMPLES
|
|
|
|
|
|
class _Iterator(Protocol):
|
|
def __call__(
|
|
self, samples: list[float], *, return_seconds: bool,
|
|
) -> dict[str, int] | None: ...
|
|
|
|
def reset_states(self) -> None: ...
|
|
|
|
|
|
def _load_iterator(
|
|
threshold: float,
|
|
min_silence_duration_ms: int,
|
|
speech_pad_ms: int,
|
|
) -> _Iterator:
|
|
# The installed package contains the model. No torch.hub or URL loading.
|
|
from silero_vad import VADIterator, load_silero_vad
|
|
|
|
return VADIterator(
|
|
load_silero_vad(onnx=True),
|
|
threshold=threshold,
|
|
sampling_rate=16000,
|
|
min_silence_duration_ms=min_silence_duration_ms,
|
|
speech_pad_ms=speech_pad_ms,
|
|
)
|
|
|
|
|
|
class SileroVADAdapter(VADAdapter):
|
|
"""One independent CPU detector per input stream.
|
|
|
|
Model loading is lazy and occurs in the same worker as inference.
|
|
Cancellation preserves worker ownership and requires reset before
|
|
more audio can be accepted. Reset/close wait for outstanding work.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
threshold: float = 0.5,
|
|
min_silence_duration_ms: int = 300,
|
|
speech_pad_ms: int = 30,
|
|
worker_timeout_seconds: float = 10.0,
|
|
) -> None:
|
|
if (
|
|
isinstance(threshold, bool)
|
|
or not isinstance(threshold, (int, float))
|
|
or not math.isfinite(threshold)
|
|
or not 0 < threshold < 1
|
|
):
|
|
raise ValueError("threshold must be finite and between 0 and 1.")
|
|
if type(min_silence_duration_ms) is not int or min_silence_duration_ms <= 0:
|
|
raise ValueError("min_silence_duration_ms must be a positive integer.")
|
|
if type(speech_pad_ms) is not int or speech_pad_ms < 0:
|
|
raise ValueError("speech_pad_ms must be a nonnegative integer.")
|
|
|
|
if (
|
|
isinstance(worker_timeout_seconds, bool)
|
|
or not isinstance(worker_timeout_seconds, (int, float))
|
|
or not math.isfinite(worker_timeout_seconds)
|
|
or worker_timeout_seconds <= 0
|
|
):
|
|
raise ValueError("worker_timeout_seconds must be finite and positive.")
|
|
|
|
self._worker_timeout_seconds = worker_timeout_seconds
|
|
self._threshold = threshold
|
|
self._min_silence_duration_ms = min_silence_duration_ms
|
|
self._speech_pad_ms = speech_pad_ms
|
|
self._iterator: _Iterator | None = None
|
|
self._lock = asyncio.Lock()
|
|
self._worker: asyncio.Task | None = None
|
|
self._needs_reset = False
|
|
self._closed = False
|
|
self._buffer = bytearray()
|
|
self._received_samples = 0
|
|
self._processed_samples = 0
|
|
self._last_boundary = 0
|
|
self._speaking = False
|
|
self._finished = False
|
|
|
|
async def process(self, audio: bytes) -> tuple[VADEvent, ...]:
|
|
self.input_format.validate_chunk(audio)
|
|
return await self._run(lambda: self._process(audio))
|
|
|
|
async def flush(self) -> tuple[VADEvent, ...]:
|
|
return await self._run(self._flush)
|
|
|
|
async def reset(self) -> None:
|
|
await self._run(self._reset, recovery=True)
|
|
|
|
async def aclose(self) -> None:
|
|
await self._run(self._close, recovery=True)
|
|
|
|
async def _run[T](
|
|
self, operation: Callable[[], T], *, recovery: bool = False,
|
|
) -> T:
|
|
async with self._lock:
|
|
await self._join_worker()
|
|
if self._needs_reset and not recovery:
|
|
raise RuntimeError("VAD requires reset after failure or cancellation.")
|
|
task = asyncio.create_task(asyncio.to_thread(operation))
|
|
self._worker = task
|
|
# Observe failures even if cancellation leaves no awaiting caller.
|
|
task.add_done_callback(self._observe_worker)
|
|
try:
|
|
result = await self._wait_worker(task)
|
|
except (asyncio.CancelledError, Exception):
|
|
self._needs_reset = True
|
|
raise
|
|
else:
|
|
if recovery:
|
|
self._needs_reset = False
|
|
return result
|
|
finally:
|
|
if task.done() and self._worker is task:
|
|
self._worker = None
|
|
|
|
@staticmethod
|
|
def _observe_worker(task: asyncio.Task) -> None:
|
|
if not task.cancelled():
|
|
task.exception()
|
|
|
|
async def _wait_worker[T](self, task: asyncio.Task[T]) -> T:
|
|
# asyncio.wait does not cancel the worker on timeout or caller cancellation.
|
|
done, _ = await asyncio.wait({task}, timeout=self._worker_timeout_seconds)
|
|
if not done:
|
|
raise TimeoutError("Silero worker exceeded its time limit; ownership retained.")
|
|
return task.result()
|
|
|
|
async def _join_worker(self) -> None:
|
|
task = self._worker
|
|
if task is None:
|
|
return
|
|
try:
|
|
await self._wait_worker(task)
|
|
except Exception:
|
|
if not task.done():
|
|
raise
|
|
logger.exception("Previous Silero worker failed; recovery is required.")
|
|
finally:
|
|
if task.done() and self._worker is task:
|
|
self._worker = None
|
|
|
|
def _require_open(self) -> None:
|
|
if self._closed:
|
|
raise RuntimeError("VAD is closed.")
|
|
|
|
def _process(self, audio: bytes) -> tuple[VADEvent, ...]:
|
|
self._require_open()
|
|
if self._finished:
|
|
raise RuntimeError("VAD input is finished; reset before new audio.")
|
|
self._received_samples += len(audio) // 2
|
|
self._buffer.extend(audio)
|
|
complete_size = len(self._buffer) - len(self._buffer) % _WINDOW_BYTES
|
|
complete = bytes(self._buffer[:complete_size])
|
|
del self._buffer[:complete_size]
|
|
events: list[VADEvent] = []
|
|
for offset in range(0, complete_size, _WINDOW_BYTES):
|
|
event = self._infer(complete[offset:offset + _WINDOW_BYTES])
|
|
if event is not None:
|
|
events.append(event)
|
|
return tuple(events)
|
|
|
|
def _infer(self, window: bytes, *, final_window: bool = False) -> VADEvent | None:
|
|
if self._iterator is None:
|
|
self._iterator = _load_iterator(
|
|
self._threshold, self._min_silence_duration_ms, self._speech_pad_ms,
|
|
)
|
|
samples = [sample / 32768.0 for sample in struct.unpack("<512h", window)]
|
|
boundary = self._iterator(samples, return_seconds=False)
|
|
self._processed_samples += _WINDOW_SAMPLES
|
|
if boundary is None:
|
|
return None
|
|
limit = min(self._processed_samples, self._received_samples)
|
|
if "start" in boundary:
|
|
if self._speaking:
|
|
raise RuntimeError("Silero emitted a duplicate speech start.")
|
|
position = min(max(boundary["start"], self._last_boundary, 0), limit)
|
|
if position >= self._received_samples:
|
|
# Synthetic tail padding must not create a new speech interval.
|
|
return None
|
|
self._speaking = True
|
|
self._last_boundary = position
|
|
return SpeechStarted(position)
|
|
if "end" in boundary:
|
|
if not self._speaking:
|
|
raise RuntimeError("Silero emitted a speech end without a start.")
|
|
position = min(max(boundary["end"], self._last_boundary, 0), limit)
|
|
self._speaking = False
|
|
self._last_boundary = position
|
|
return SpeechEnded(
|
|
position, "end_of_input" if final_window else "silence",
|
|
)
|
|
raise RuntimeError("Silero emitted an unsupported boundary.")
|
|
|
|
def _flush(self) -> tuple[VADEvent, ...]:
|
|
self._require_open()
|
|
if self._finished:
|
|
return ()
|
|
events: list[VADEvent] = []
|
|
if self._buffer:
|
|
window = bytes(self._buffer).ljust(_WINDOW_BYTES, b"\x00")
|
|
self._buffer.clear()
|
|
event = self._infer(window, final_window=True)
|
|
if event is not None:
|
|
events.append(event)
|
|
if self._speaking:
|
|
events.append(SpeechEnded(self._received_samples, "end_of_input"))
|
|
self._last_boundary = self._received_samples
|
|
self._speaking = False
|
|
self._finished = True
|
|
return tuple(events)
|
|
|
|
def _clear_input(self) -> None:
|
|
self._buffer.clear()
|
|
self._received_samples = 0
|
|
self._processed_samples = 0
|
|
self._last_boundary = 0
|
|
self._speaking = False
|
|
self._finished = False
|
|
|
|
def _reset(self) -> None:
|
|
self._require_open()
|
|
if self._iterator is not None:
|
|
self._iterator.reset_states()
|
|
self._clear_input()
|
|
|
|
def _close(self) -> None:
|
|
if self._closed:
|
|
return
|
|
self._iterator = None
|
|
self._clear_input()
|
|
self._closed = True
|