feat: overlap-window speechbrain enhance, hop 200ms / window 1s

.env.example

@@ -25,7 +25,9 @@ AUDIO_CHANNELS=1
 ENHANCE_MODE=auto
 # speechbrain model repo (must be 16 kHz; metricGAN-plus runs ~20-40ms/frame on CPU)
 ENHANCE_MODEL=speechbrain/metricgan-plus-voicebank
-ENHANCE_CHUNK_S=0.2
+# model window vs publish hop: 1 s context, 200 ms emit
+ENHANCE_CHUNK_S=1.0
+ENHANCE_HOP_S=0.2
 
 # Misc
 PUBLISH_TIMEOUT_S=60

audio_cleanup.py

@@ -31,9 +31,14 @@ class AudioCleaner:
 
     def __init__(self, mode: str = "auto",
                  model_source: str = "speechbrain/metricgan-plus-voicebank",
-                 sample_rate: int = SAMPLE_RATE):
+                 sample_rate: int = SAMPLE_RATE,
+                 chunk_s: float = 1.0,
+                 hop_s: float = 0.2):
         self.mode = mode
         self.sample_rate = sample_rate
+        self.chunk = max(1, int(chunk_s * sample_rate))
+        self.hop = max(1, int(hop_s * sample_rate))
+        self._ctx = np.zeros(0, dtype=np.float32)
         self._enhancer = None
         self._torch = None
         self.backend = "off"
@@ -102,20 +107,29 @@ class AudioCleaner:
     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.
+        Accumulates a `chunk_s` window and runs the model on that window,
+        emitting only the latest hop (the input length). Until a full window
+        is buffered, the DSP fallback is used so the stream is not silent.
         """
-        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)
-
+        n = len(pcm)
+        if n == 0:
+            return pcm
         x = pcm.astype(np.float32) / 32768.0
+
+        if self._enhancer is None:
+            y = self._simple_clean(x)
+            return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16)
+
+        self._ctx = np.concatenate([self._ctx, x])
+        if len(self._ctx) < self.chunk:
+            y = self._simple_clean(x)
+            return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16)
+        self._ctx = self._ctx[-self.chunk:]
+
         torch = self._torch
         try:
-            w = torch.from_numpy(x).unsqueeze(0).to(self._enhancer.device)
+            w = torch.from_numpy(self._ctx.copy()).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".
@@ -127,16 +141,13 @@ class AudioCleaner:
             except TypeError:
                 with torch.no_grad():
                     out = self._enhancer.enhance_batch(w)
-            out = out.squeeze(0).cpu().numpy()
+            enhanced = 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)
+            return (np.clip(y, -1.0, 1.0) * 32767.0).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)
+        hop_out = enhanced[-n:] if len(enhanced) >= n else np.pad(
+            enhanced, (0, n - len(enhanced)))
+        hop_out = hop_out[:n]
+        return (np.clip(hop_out, -1.0, 1.0) * 32767.0).astype(np.int16)

config.py

@@ -36,7 +36,8 @@ class Config:
     # speechbrain cleanup
     enhance_mode: str = "auto"  # auto | force | off
     enhance_model: str = "speechbrain/metricgan-plus-voicebank"
-    enhance_chunk_s: float = 0.2
+    enhance_chunk_s: float = 1.0
+    enhance_hop_s: float = 0.2
     # misc
     publish_timeout_s: float = 60.0
     log_level: str = "INFO"
@@ -80,7 +81,8 @@ def load_config(env_path: Optional[Path] = None) -> Config:
         enhance_model=os.environ.get(
             "ENHANCE_MODEL", "speechbrain/metricgan-plus-voicebank"
         ),
-        enhance_chunk_s=float(os.environ.get("ENHANCE_CHUNK_S", "0.2")),
+        enhance_chunk_s=float(os.environ.get("ENHANCE_CHUNK_S", "1.0")),
+        enhance_hop_s=float(os.environ.get("ENHANCE_HOP_S", "0.2")),
         publish_timeout_s=float(os.environ.get("PUBLISH_TIMEOUT_S", "60")),
         log_level=os.environ.get("LOG_LEVEL", "INFO").upper(),
         cameras=_parse_cameras(os.environ.get("CAMERAS", "all")),

run.py

@@ -41,6 +41,7 @@ def worker_args_for(cfg, cam) -> list[str]:
         "--enhance-mode", cfg.enhance_mode,
         "--enhance-model", cfg.enhance_model,
         "--enhance-chunk-s", str(cfg.enhance_chunk_s),
+        "--enhance-hop-s", str(cfg.enhance_hop_s),
     ]
 
 

run_worker.py

@@ -35,7 +35,8 @@ def main() -> int:
     ap.add_argument("--enhance-mode", default="auto")
     ap.add_argument("--enhance-model",
                     default="speechbrain/metricgan-plus-voicebank")
-    ap.add_argument("--enhance-chunk-s", type=float, default=0.2)
+    ap.add_argument("--enhance-chunk-s", type=float, default=1.0)
+    ap.add_argument("--enhance-hop-s", type=float, default=0.2)
     args = ap.parse_args()
 
     cam = json.loads(args.camera_json)
@@ -57,6 +58,8 @@ def main() -> int:
             mode=args.enhance_mode,
             model_source=args.enhance_model,
             sample_rate=args.audio_rate,
+            chunk_s=args.enhance_chunk_s,
+            hop_s=args.enhance_hop_s,
         )
         log.info("audio backend: %s", cleaner.backend)
 
@@ -125,8 +128,8 @@ def main() -> int:
             atrack = rtc.LocalAudioTrack.create_audio_track(
                 f"{identity}-audio", audio_source)
 
-            chunk_s = args.enhance_chunk_s
-            chunk = int(chunk_s * args.audio_rate)  # samples per channel
+            hop_s = args.enhance_hop_s
+            chunk = int(hop_s * args.audio_rate)  # samples per hop
 
             async def audio_loop():
                 assert aproc is not None and atrack is not None

scripts/bench_enhance.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+"""RTF bench for hop/window speechbrain enhance."""
+import time
+
+import numpy as np
+
+from audio_cleanup import AudioCleaner
+
+c = AudioCleaner(mode="force", chunk_s=1.0, hop_s=0.2)
+hop = 3200
+pcm = (np.random.randn(hop) * 800).astype(np.int16)
+for _ in range(8):
+    c.process(pcm)
+n = 50
+t0 = time.perf_counter()
+for _ in range(n):
+    c.process(pcm)
+dt = time.perf_counter() - t0
+audio_s = n * 0.2
+rtf = dt / audio_s
+print(f"hops={n} wall={dt:.3f}s audio={audio_s:.1f}s rtf={rtf:.3f}")
+if rtf > 0.8:
+    raise SystemExit("RTF too high for 20-way realtime")

tests/test_audio_cleanup.py

@@ -28,3 +28,23 @@ def test_process_preserves_length_and_int16():
     out = c.process(pcm)
     assert out.dtype == np.int16
     assert out.shape == pcm.shape
+
+
+def test_overlap_process_emits_hop_length():
+    c = AudioCleaner(mode="force", chunk_s=1.0, hop_s=0.2)
+    hop = 3200
+    pcm = (np.random.randn(hop) * 500).astype(np.int16)
+    out = c.process(pcm)
+    assert out.shape == (hop,)
+
+
+def test_overlap_runs_model_after_enough_context():
+    c = AudioCleaner(mode="force", chunk_s=0.5, hop_s=0.25)
+    hop = 4000
+    for _ in range(4):
+        pcm = (np.sin(2 * np.pi * 440 * np.arange(hop) / 16000) * 8000).astype(
+            np.int16
+        )
+        out = c.process(pcm)
+        assert out.shape == (hop,)
+        assert out.dtype == np.int16