From d8452c0ffe69f208d63b97fc2a4dfe874cdec19d Mon Sep 17 00:00:00 2001 From: Hermes Date: Fri, 28 Aug 2026 05:17:13 +0000 Subject: [PATCH] feat: overlap-window speechbrain enhance, hop 200ms / window 1s --- .env.example | 4 ++- audio_cleanup.py | 51 ++++++++++++++++++++++--------------- config.py | 6 +++-- run.py | 1 + run_worker.py | 9 ++++--- scripts/bench_enhance.py | 23 +++++++++++++++++ tests/test_audio_cleanup.py | 20 +++++++++++++++ 7 files changed, 88 insertions(+), 26 deletions(-) create mode 100644 scripts/bench_enhance.py diff --git a/.env.example b/.env.example index b55e5f3..720f7f2 100644 --- a/.env.example +++ b/.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 diff --git a/audio_cleanup.py b/audio_cleanup.py index a712ab4..e24e4c4 100644 --- a/audio_cleanup.py +++ b/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) diff --git a/config.py b/config.py index 4a63691..c982541 100644 --- a/config.py +++ b/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")), diff --git a/run.py b/run.py index 663ad08..ac73312 100644 --- a/run.py +++ b/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), ] diff --git a/run_worker.py b/run_worker.py index 3623e31..3d7f67f 100644 --- a/run_worker.py +++ b/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 diff --git a/scripts/bench_enhance.py b/scripts/bench_enhance.py new file mode 100644 index 0000000..424bd51 --- /dev/null +++ b/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") diff --git a/tests/test_audio_cleanup.py b/tests/test_audio_cleanup.py index a1cb746..a8643b9 100644 --- a/tests/test_audio_cleanup.py +++ b/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