chore: snapshot livekit-cameras after serial pin + speechbrain lengths fix

2026-08-28 05:15:00 +00:00
commit dc46cbef04
14 changed files with 1369 additions and 0 deletions
@@ -0,0 +1,142 @@
+"""Audio cleanup for camera mics, using speechbrain.
+
+Pipeline per audio frame (16 kHz mono int16, ~100 ms chunks from the mic):
+  int16 -> float32 [-1, 1] -> speechbrain enhancement -> int16
+
+Two implementations, selected by availability and the ENHANCE_MODE config:
+  * "model"  - speechbrain SpectralMaskEnhancement (metricGAN-plus). Real
+               ML-based noise suppression. Runs on CPU; ~20-40 ms per 100 ms
+               chunk on a modern core, so latency stays well under a frame
+               budget if workers are spread across cores.
+  * "gain"   - lightweight DSP fallback (soft-clip + high-pass) used when the
+               model can't be loaded (ENHANCE_MODE=off / auto + no model).
+
+The enhancer is created once per worker process and reused across frames.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import numpy as np
+
+log = logging.getLogger("audio.cleanup")
+
+SAMPLE_RATE = 16000
+
+
+class AudioCleaner:
+    """Stateful audio cleaner. Call .process(int16 mono) -> int16 mono."""
+
+    def __init__(self, mode: str = "auto",
+                 model_source: str = "speechbrain/metricgan-plus-voicebank",
+                 sample_rate: int = SAMPLE_RATE):
+        self.mode = mode
+        self.sample_rate = sample_rate
+        self._enhancer = None
+        self._torch = None
+        self.backend = "off"
+
+        if mode == "force":
+            self._load_model(model_source)
+            if self.backend == "off":
+                raise RuntimeError(
+                    f"ENHANCE_MODE=force but failed to load model "
+                    f"'{model_source}'")
+        elif mode == "auto":
+            self._load_model(model_source)
+        # mode == "off": passthrough + basic gain shaping
+
+    def _load_model(self, source: str) -> None:
+        try:
+            import torch
+            from speechbrain.inference.enhancement import (
+                SpectralMaskEnhancement,
+            )
+        except Exception as e:  # noqa: BLE001
+            log.info("speechbrain unavailable (%s); using fallback", e)
+            self._enhancer = None
+            return
+        try:
+            log.info("loading speechbrain enhancement model %s ...", source)
+            self._torch = torch
+            model = SpectralMaskEnhancement.from_hparams(
+                source=source,
+                savedir="/root/.cache/speechbrain-enhancement",
+            )
+            model.mods.enhance_model.eval()
+            self._enhancer = model
+            self.backend = "model"
+            log.info("speechbrain enhancement ready (16 kHz in/out)")
+        except Exception as e:  # noqa: BLE001
+            log.warning("failed to load enhancement model %s: %s "
+                        "(using fallback)", source, e)
+            self._enhancer = None
+
+    # -- DSP fallback ------------------------------------------------------
+
+    @staticmethod
+    def _simple_clean(x: np.ndarray) -> np.ndarray:
+        """Lightweight cleanup: high-pass + soft knee + gentle gain.
+
+        Good enough to remove DC offset/low rumble and taming clipping
+        without any ML model.
+        """
+        # first-order high-pass at ~80 Hz
+        if len(x) > 1:
+            alpha = 0.998
+            y = np.empty_like(x)
+            y[0] = x[0]
+            for i in range(1, len(x)):
+                y[i] = alpha * (y[i - 1] + x[i] - x[i - 1])
+            x = y
+        # soft-clip to avoid harsh distortion
+        x = np.tanh(x * 0.9) / np.tanh(0.9)
+        # gentle makeup gain
+        x = x * 1.1
+        return np.clip(x, -1.0, 1.0)
+
+    # -- public API ---------------------------------------------------------
+
+    def process(self, pcm: np.ndarray) -> np.ndarray:
+        """pcm: int16 mono numpy array. Returns int16 mono, same length.
+
+        If the model changes the length, we trim/pad back to the original
+        length so downstream chunking stays aligned.
+        """
+        if self._enhancer is None:
+            if len(pcm) == 0:
+                return pcm
+            x = pcm.astype(np.float32) / 32768.0
+            y = self._simple_clean(x)
+            return (y * 32767.0).astype(np.float32).astype(np.int16)
+
+        x = pcm.astype(np.float32) / 32768.0
+        torch = self._torch
+        try:
+            w = torch.from_numpy(x).unsqueeze(0).to(self._enhancer.device)
+            # speechbrain >= 1.1 (e.g. metricGAN-plus) requires `lengths`
+            # on enhance_batch. The value is RELATIVE (multiplied by the
+            # number of frames internally), so 1.0 = "the whole sequence".
+            # Older versions accept the call without it (TypeError fallback).
+            lengths = torch.tensor([1.0]).to(self._enhancer.device)
+            try:
+                with torch.no_grad():
+                    out = self._enhancer.enhance_batch(w, lengths=lengths)
+            except TypeError:
+                with torch.no_grad():
+                    out = self._enhancer.enhance_batch(w)
+            out = out.squeeze(0).cpu().numpy()
+        except Exception as e:  # noqa: BLE001
+            log.warning("enhance_batch failed (%s); using fallback", e)
+            y = self._simple_clean(x)
+            return (y * 32767.0).astype(np.float32).astype(np.int16)
+
+        # align length
+        n = len(x)
+        if len(out) >= n:
+            out = out[:n]
+        else:
+            out = np.pad(out, (0, n - len(out)))
+        return (out * 32767.0).astype(np.float32).astype(np.int16)