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

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)