Reduce listener recomputation with exact embedding reuse

README.md

@@ -145,6 +145,10 @@ if scores["hey_livekit"] > 0.5:
     print("Wake word detected!")
 ```
 
+`WakeWordModel.predict()` is stateless. For a sequence of overlapping full
+windows, `model.create_stream()` reuses only speech embeddings whose mel inputs
+match exactly. `WakeWordListener` enables this optimization automatically.
+
 **Async listener with microphone:**
 
 ```python

docs/export-and-inference.md

@@ -107,8 +107,29 @@ Feature extraction models (`melspectrogram.onnx`, `embedding_model.onnx`) are bu
 | Method | Returns | Description |
 |--------|---------|-------------|
 | `predict(audio_chunk)` | `dict[str, float]` | Scores for each loaded model (0-1) |
+| `create_stream()` | `WakeWordStream` | Isolated cache for overlapping full audio windows |
 | `load_model(path, name)` | `None` | Load additional wake word model |
 
+#### Streaming Prediction
+
+`WakeWordModel` remains stateless. For one continuous audio source, create a
+separate stream and pass it the same complete ~2-second windows, advanced by
+80 ms (1,280 samples), that you would otherwise pass to `model.predict()`:
+
+```python
+stream = model.create_stream()
+
+scores = stream.predict(first_audio_window)
+scores = stream.predict(next_overlapping_window)
+stream.reset()  # Clear cached mel windows and speech embeddings.
+```
+
+The stream still computes the full mel spectrogram because mel output can depend
+on the complete audio window. It reuses a speech embedding only when its 76-frame
+mel input exactly matches the shifted input from the previous window, then batches
+all changed windows into one embedding call. Use one stream per audio source and
+call it serially. `WakeWordListener` creates and resets its own stream automatically.
+
 #### Audio Input
 
 - **Format:** 16kHz mono, int16 or float32
@@ -118,6 +139,8 @@ Feature extraction models (`melspectrogram.onnx`, `embedding_model.onnx`) are bu
 ### WakeWordListener
 
 The `WakeWordListener` class provides async microphone detection with debouncing.
+It also uses `WakeWordModel.create_stream()` to avoid recomputing unchanged speech
+embeddings between overlapping microphone windows.
 
 ```python
 import asyncio

src/livekit/wakeword/__init__.py

@@ -1,7 +1,7 @@
 """livekit-wakeword — Wake word detection for voice-enabled applications."""
 
 from .inference.listener import Detection, WakeWordListener
-from .inference.model import WakeWordModel
+from .inference.model import WakeWordModel, WakeWordStream
 
 __version__ = "0.1.0"
 
@@ -35,6 +35,7 @@ def __getattr__(name: str) -> object:
     "ExportFormat",
     "WakeWordListener",
     "WakeWordModel",
+    "WakeWordStream",
     "Detection",
     "load_config",
     "run_augment",

src/livekit/wakeword/inference/__init__.py

@@ -1,6 +1,6 @@
 """Real-time inference engine."""
 
 from .listener import Detection, WakeWordListener
-from .model import WakeWordModel
+from .model import WakeWordModel, WakeWordStream
 
-__all__ = ["Detection", "WakeWordListener", "WakeWordModel"]
+__all__ = ["Detection", "WakeWordListener", "WakeWordModel", "WakeWordStream"]

src/livekit/wakeword/inference/listener.py

@@ -8,6 +8,7 @@
 import time
 from collections import deque
 from dataclasses import dataclass
+from typing import Protocol
 
 import numpy as np
 
@@ -18,10 +19,14 @@
 SAMPLE_RATE = 16000
 FRAME_SAMPLES = 1280  # 80ms per frame
 CHUNK_SECONDS = 2.0
-# Number of frames that fill a ~2-second chunk (25 × 80ms = 2000ms)
+# Number of frames that fill a ~2-second chunk (25 * 80ms = 2000ms)
 CHUNK_FRAMES = int(CHUNK_SECONDS * SAMPLE_RATE / FRAME_SAMPLES)
 
 
+class _Predictor(Protocol):
+    def predict(self, audio_chunk: np.ndarray) -> dict[str, float]: ...
+
+
 @dataclass
 class Detection:
     """Wake word detection result."""
@@ -35,9 +40,10 @@ class WakeWordListener:
     """Async wake word listener that handles audio capture.
 
     The listener owns the audio buffer and passes fixed ~2-second chunks
-    to the stateless ``WakeWordModel.predict()``.  After a detection the
-    loop pauses automatically and resumes when the consumer calls
-    ``wait_for_detection()`` again.
+    through an isolated stream predictor. The predictor reuses unchanged
+    speech embeddings while ``WakeWordModel`` remains stateless. After a
+    detection the loop pauses automatically and resumes when the consumer
+    calls ``wait_for_detection()`` again.
 
     Example:
         from livekit.wakeword import WakeWordModel, WakeWordListener
@@ -64,6 +70,7 @@ def __init__(
             debounce: Minimum seconds between detections.
         """
         self._model = model
+        self._predictor: _Predictor = model
         self._threshold = threshold
         self._debounce = debounce
 
@@ -95,6 +102,8 @@ async def __aenter__(self) -> WakeWordListener:
         """Start audio capture."""
         import pyaudio
 
+        create_stream = getattr(self._model, "create_stream", None)
+        self._predictor = create_stream() if callable(create_stream) else self._model
         self._pa = pyaudio.PyAudio()
         self._stream = self._pa.open(
             format=pyaudio.paInt16,
@@ -124,7 +133,7 @@ async def __aexit__(self, *_: object) -> None:
         if self._task:
             try:
                 await asyncio.wait_for(self._task, timeout=2.0)
-            except (asyncio.TimeoutError, asyncio.CancelledError):
+            except (TimeoutError, asyncio.CancelledError):
                 self._task.cancel()
                 try:
                     await self._task
@@ -175,7 +184,7 @@ async def _audio_loop(self) -> None:
                 chunk = np.concatenate(list(self._frame_buffer))
                 scores = await loop.run_in_executor(
                     self._executor,
-                    self._model.predict,
+                    self._predictor.predict,
                     chunk,
                 )
                 if not self._running:
@@ -193,11 +202,12 @@ async def _audio_loop(self) -> None:
                             # the detection.
                             self._listening.clear()
                             self._frame_buffer.clear()
+                            reset = getattr(self._predictor, "reset", None)
+                            if callable(reset):
+                                await loop.run_in_executor(self._executor, reset)
 
                             await self._detection_queue.put(
-                                Detection(
-                                    name=name, confidence=score, timestamp=now
-                                )
+                                Detection(name=name, confidence=score, timestamp=now)
                             )
                             break  # one detection per iteration
         except asyncio.CancelledError:
@@ -243,8 +253,6 @@ async def wait_for_detection(self) -> Detection:
             return self._detection_queue.get_nowait()
 
         if self._error is not None:
-            raise RuntimeError(
-                f"Audio loop crashed: {self._error}"
-            ) from self._error
+            raise RuntimeError(f"Audio loop crashed: {self._error}") from self._error
 
         raise RuntimeError("Audio loop ended unexpectedly")

src/livekit/wakeword/inference/model.py

@@ -1,4 +1,4 @@
-"""Stateless wake word detection model."""
+"""Wake word detection model with optional per-stream caching."""
 
 from __future__ import annotations
 
@@ -61,7 +61,10 @@ def __init__(
             )
 
         self._mel_frontend = MelSpectrogramFrontend(onnx_path=mel_path, sess_options=sess_options)
-        self._speech_embedding = SpeechEmbedding(onnx_path=embedding_path, sess_options=sess_options)
+        self._speech_embedding = SpeechEmbedding(
+            onnx_path=embedding_path,
+            sess_options=sess_options,
+        )
 
         # name -> (onnx_session, input_name)
         self._classifiers: dict[str, tuple] = {}
@@ -70,7 +73,12 @@ def __init__(
             for model_path in models:
                 self.load_model(model_path, sess_options=sess_options)
 
-    def load_model(self, model_path: str | Path, model_name: str | None = None, sess_options: SessionOptions = None ) -> None:
+    def load_model(
+        self,
+        model_path: str | Path,
+        model_name: str | None = None,
+        sess_options: SessionOptions | None = None,
+    ) -> None:
         """Load a wake word classifier model.
 
         Args:
@@ -89,12 +97,20 @@ def load_model(self, model_path: str | Path, model_name: str | None = None, sess
         session = ort.InferenceSession(
             str(model_path),
             providers=["CPUExecutionProvider"],
-            sess_options=sess_options
+            sess_options=sess_options,
         )
         input_name = session.get_inputs()[0].name
         self._classifiers[model_name] = (session, input_name)
         logger.info(f"Loaded wake word model '{model_name}' from {model_path}")
 
+    def create_stream(self) -> WakeWordStream:
+        """Create an isolated predictor for overlapping audio windows.
+
+        The returned stream caches speech embeddings. The model itself remains
+        stateless, and each stream owns independent cache state.
+        """
+        return WakeWordStream(self)
+
     def predict(self, audio_chunk: np.ndarray) -> dict[str, float]:
         """Get wake word predictions for an audio chunk.
 
@@ -112,38 +128,104 @@ def predict(self, audio_chunk: np.ndarray) -> dict[str, float]:
         if not self._classifiers:
             return {}
 
-        # Convert int16 to float32 if needed
+        windows = self._mel_windows(audio_chunk)
+        if windows is None:
+            return self._zero_scores()
+
+        embeddings = np.stack(
+            [self._speech_embedding(window[np.newaxis, :, :])[0] for window in windows],
+            axis=0,
+        )
+        return self._classify(embeddings)
+
+    def _mel_windows(self, audio_chunk: np.ndarray) -> np.ndarray | None:
         if audio_chunk.dtype == np.int16:
             audio_chunk = audio_chunk.astype(np.float32) / 32768.0
 
-        audio_chunk = audio_chunk.flatten()
-
-        # Mel spectrogram over the full chunk
-        all_mel = self._mel_frontend(audio_chunk)
+        all_mel = self._mel_frontend(audio_chunk.flatten())
         if all_mel.ndim == 3:
             all_mel = all_mel[0]
 
-        if all_mel.shape[0] < EMBEDDING_WINDOW:
-            return {name: 0.0 for name in self._classifiers}
+        starts = range(
+            0,
+            all_mel.shape[0] - EMBEDDING_WINDOW + 1,
+            EMBEDDING_STRIDE,
+        )
+        windows = [all_mel[start : start + EMBEDDING_WINDOW] for start in starts]
+        if len(windows) < MIN_EMBEDDINGS:
+            return None
+        return np.stack(windows[-MIN_EMBEDDINGS:], axis=0)
+
+    def _classify(self, embeddings: np.ndarray) -> dict[str, float]:
+        emb_input = embeddings[np.newaxis, :, :].astype(np.float32)
+        predictions: dict[str, float] = {}
+        for name, (session, input_name) in self._classifiers.items():
+            outputs = session.run(None, {input_name: emb_input})
+            predictions[name] = float(outputs[0][0, 0])
+        return predictions
+
+    def _zero_scores(self) -> dict[str, float]:
+        return {name: 0.0 for name in self._classifiers}
 
-        # Extract embeddings: 76-frame windows, stride 8
-        embeddings = []
-        for start in range(0, all_mel.shape[0] - EMBEDDING_WINDOW + 1, EMBEDDING_STRIDE):
-            window = all_mel[start : start + EMBEDDING_WINDOW]
-            emb = self._speech_embedding(window[np.newaxis, :, :])
-            embeddings.append(emb[0])
 
-        if len(embeddings) < MIN_EMBEDDINGS:
-            return {name: 0.0 for name in self._classifiers}
+class WakeWordStream:
+    """Prediction state for a sequence of overlapping full audio windows.
 
-        # Use last 16 embeddings
-        emb_sequence = np.stack(embeddings[-MIN_EMBEDDINGS:], axis=0)
-        emb_input = emb_sequence[np.newaxis, :, :].astype(np.float32)
+    Each call still computes the full mel spectrogram because that model's
+    output can depend on the complete audio window. Speech embeddings are reused
+    only when their 76-frame mel inputs exactly match the shifted prior inputs.
+    """
 
-        predictions = {}
-        for name, (session, input_name) in self._classifiers.items():
-            outputs = session.run(None, {input_name: emb_input})
-            score = float(outputs[0][0, 0])
-            predictions[name] = score
+    def __init__(self, model: WakeWordModel):
+        self._model = model
+        self._mel_windows: np.ndarray | None = None
+        self._embeddings: np.ndarray | None = None
 
-        return predictions
+    def reset(self) -> None:
+        """Clear cached windows and embeddings."""
+        self._mel_windows = None
+        self._embeddings = None
+
+    def predict(self, audio_chunk: np.ndarray) -> dict[str, float]:
+        """Predict from the next complete audio window in the stream."""
+        if not self._model._classifiers:
+            return {}
+
+        windows = self._model._mel_windows(audio_chunk)
+        if windows is None:
+            self.reset()
+            return self._model._zero_scores()
+
+        embeddings = self._embedding_sequence(windows)
+        self._mel_windows = windows
+        self._embeddings = embeddings
+        return self._model._classify(embeddings)
+
+    def _embedding_sequence(self, windows: np.ndarray) -> np.ndarray:
+        changed_indexes: list[int] = []
+        changed_windows: list[np.ndarray] = []
+        reused: dict[int, np.ndarray] = {}
+
+        for index, window in enumerate(windows):
+            previous_index = index + 1
+            if (
+                self._mel_windows is not None
+                and self._embeddings is not None
+                and previous_index < len(self._mel_windows)
+                and np.array_equal(window, self._mel_windows[previous_index])
+            ):
+                reused[index] = self._embeddings[previous_index]
+            else:
+                changed_indexes.append(index)
+                changed_windows.append(window)
+
+        fresh = self._model._speech_embedding(np.stack(changed_windows, axis=0))
+        embeddings = np.empty(
+            (len(windows), fresh.shape[1]),
+            dtype=fresh.dtype,
+        )
+        for index, embedding in reused.items():
+            embeddings[index] = embedding
+        for index, embedding in zip(changed_indexes, fresh, strict=True):
+            embeddings[index] = embedding
+        return embeddings