From ef9573c2946c8d482d35abdaa3253af5bb1c4391 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 28 Aug 2026 07:25:30 +0000 Subject: [PATCH] feat: OpenVINO CPU MetricGAN mask; 100 ms hop Per-worker OpenVINO CPU runs the MetricGAN BLSTM mask; SpeechBrain keeps STFT/ISTFT. ENHANCE_INFER=auto with torch fallback. Hop is 0.1 s after isolated last_ms ~6 ms (model+ov+beamform). --- .env.example | 6 ++- README.md | 15 +++--- audio_cleanup.py | 80 ++++++++++++++++++++++++------ config.py | 9 +++- ov_metricgan.py | 99 +++++++++++++++++++++++++++++++++++++ run.py | 1 + run_worker.py | 4 +- tests/test_audio_cleanup.py | 19 ++++++- tests/test_config.py | 4 ++ 9 files changed, 211 insertions(+), 26 deletions(-) create mode 100644 ov_metricgan.py diff --git a/.env.example b/.env.example index 6b8337a..bd6b8ce 100644 --- a/.env.example +++ b/.env.example @@ -32,9 +32,11 @@ 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 -# model window vs publish hop: 1 s context, 200 ms emit +# model window vs publish hop: 1 s context, 100 ms emit ENHANCE_CHUNK_S=1.0 -ENHANCE_HOP_S=0.2 +ENHANCE_HOP_S=0.1 +# MetricGAN mask: auto (OpenVINO CPU, Torch fallback) | openvino | torch +ENHANCE_INFER=auto # C920 stereo Delay-and-Sum (SpeechBrain GCC-PHAT) before MetricGAN # auto (default): beamform if speechbrain loads, else average channels # force: require DelaySum_Beamformer, off: average L/R diff --git a/README.md b/README.md index c6e67ab..26bd640 100644 --- a/README.md +++ b/README.md @@ -101,22 +101,25 @@ HDMI-A-7 is replaced by the grid). ## Audio cleanup -C920 mics are 16 kHz **stereo**. Each hop (default 200 ms) is kept +C920 mics are 16 kHz **stereo**. Each hop (default 100 ms) is kept stereo through a SpeechBrain preprocessor, then published as mono: 1. Delay-and-Sum beamform (GCC-PHAT TDOA, frozen for `BEAMFORM_TDOA_EVERY=5` hops) — live path from the SpeechBrain [multi-microphone beamforming tutorial](https://speechbrain.readthedocs.io/en/latest/tutorials/preprocessing/multi-microphone-beamforming.html) 2. `speechbrain/metricgan-plus-voicebank` neural enhancement on a 1 s - rolling window (`ENHANCE_CHUNK_S=1.0`, `ENHANCE_HOP_S=0.2`) + rolling window (`ENHANCE_CHUNK_S=1.0`, `ENHANCE_HOP_S=0.1`). + The MetricGAN **mask** runs on OpenVINO CPU (`ENHANCE_INFER=auto`); + STFT/ISTFT stay in Torch. `ENHANCE_INFER=torch` forces SpeechBrain + `enhance_batch`. Do not use the Arc B580 for this net. The 1 s window is **past context**, not +1 s delay: only the latest hop is published. Algorithmic delay ≈ hop + compute. -Each worker caps Torch/BLAS to `TORCH_NUM_THREADS=1` (default) so 14–20 -processes do not oversubscribe the i7-14700. This host has an Intel Arc -B580; the venv Torch is NVIDIA CUDA and **cannot** use it. Do not set -`run_opts={"device":"cuda"}`. +Each worker caps Torch/BLAS/OpenVINO to `TORCH_NUM_THREADS=1` (default) +so 14–20 processes do not oversubscribe the i7-14700. This host has an +Intel Arc B580; the venv Torch is NVIDIA CUDA and **cannot** use it. Do +not set `run_opts={"device":"cuda"}`. One worker process per camera loads the modules once and reuses them. `enhance_batch` must be called with `lengths=torch.tensor([1.0])` diff --git a/audio_cleanup.py b/audio_cleanup.py index 95d43d3..3d6cebd 100644 --- a/audio_cleanup.py +++ b/audio_cleanup.py @@ -65,7 +65,8 @@ class AudioCleaner: hop_s: float = 0.2, beamform: str = "auto", torch_num_threads: int = 1, - tdoa_every: int = 5): + tdoa_every: int = 5, + enhance_infer: str = "auto"): self.mode = mode self.beamform_mode = beamform self.sample_rate = sample_rate @@ -87,6 +88,11 @@ class AudioCleaner: self._istft = None self._last_tdoas = None self._tdoa_age = 0 + self.enhance_infer = (enhance_infer or "auto").lower() + self.infer_backend = "torch" + self._ov_compiled = None + self._ov_request = None + self._ov_input = None self._load_torch() if beamform == "force": @@ -107,12 +113,20 @@ class AudioCleaner: elif mode == "auto": self._load_model(model_source) + if self._enhancer is not None and self.enhance_infer != "torch": + self._load_openvino() + if self.enhance_infer == "openvino" and self._ov_compiled is None: + raise RuntimeError( + "ENHANCE_INFER=openvino but OpenVINO CPU mask failed to load") + self._refresh_backend() def _refresh_backend(self) -> None: parts = [] if self._enhancer is not None: parts.append("model") + if self._ov_compiled is not None: + parts.append("ov") if self._beamformer is not None: parts.append("beamform") self.backend = "+".join(parts) if parts else "off" @@ -188,6 +202,28 @@ class AudioCleaner: "(using fallback)", source, e) self._enhancer = None + def _load_openvino(self) -> None: + """Compile MetricGAN mask on OpenVINO CPU. STFT/ISTFT stay in Torch.""" + try: + from ov_metricgan import compile_cpu, ensure_onnx + except Exception as e: # noqa: BLE001 + log.info("openvino helper unavailable (%s)", e) + return + try: + onnx_path = ensure_onnx(self._enhancer, self._torch) + compiled, request, ov_input = compile_cpu( + onnx_path, num_threads=self.torch_num_threads) + self._ov_compiled = compiled + self._ov_request = request + self._ov_input = ov_input + self.infer_backend = "openvino" + except Exception as e: # noqa: BLE001 + log.warning("OpenVINO CPU mask failed (%s); using torch enhance", e) + self._ov_compiled = None + self._ov_request = None + self._ov_input = None + self.infer_backend = "torch" + # -- DSP fallback ------------------------------------------------------ @staticmethod @@ -258,6 +294,19 @@ class AudioCleaner: out = np.pad(out, (0, n - out.shape[0])) return out + def _enhance_openvino(self, w): + """w: torch [1, T] on CPU. Returns numpy [T].""" + from ov_metricgan import infer_mask + torch = self._torch + with torch.no_grad(): + feats = self._enhancer.compute_features(w) + feats_np = np.ascontiguousarray(feats.detach().cpu().numpy()) + mask_np = infer_mask(self._ov_request, self._ov_input, feats_np) + mask = torch.from_numpy(mask_np) + enhanced = torch.mul(mask, feats) + wav = self._enhancer.hparams.resynth(torch.expm1(enhanced), w) + return wav.squeeze(0).detach().cpu().numpy() + def _to_int16(self, y: np.ndarray) -> np.ndarray: return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16) @@ -309,20 +358,23 @@ class AudioCleaner: w = torch.from_numpy(self._ctx.copy()).unsqueeze(0) if hasattr(self._enhancer, "device"): w = w.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(w.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) - enhanced = out.squeeze(0).detach().cpu().numpy() + if self._ov_compiled is not None: + enhanced = self._enhance_openvino(w) + else: + # 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(w.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) + enhanced = out.squeeze(0).detach().cpu().numpy() except Exception as e: # noqa: BLE001 - log.warning("enhance_batch failed (%s); using fallback", e) + log.warning("enhance failed (%s); using fallback", e) return self._to_int16(self._simple_clean(mono_hop)) hop_out = enhanced[-n:] if len(enhanced) >= n else np.pad( diff --git a/config.py b/config.py index 9a69a8a..cc6cbb5 100644 --- a/config.py +++ b/config.py @@ -41,7 +41,8 @@ class Config: enhance_mode: str = "auto" # auto | force | off enhance_model: str = "speechbrain/metricgan-plus-voicebank" enhance_chunk_s: float = 1.0 - enhance_hop_s: float = 0.2 + enhance_hop_s: float = 0.1 + enhance_infer: str = "auto" # auto | openvino | torch # SpeechBrain Delay-and-Sum (GCC-PHAT) on C920 stereo before enhance beamform_mode: str = "auto" # auto | force | off torch_num_threads: int = 1 @@ -118,7 +119,8 @@ def load_config(env_path: Optional[Path] = None) -> Config: "ENHANCE_MODEL", "speechbrain/metricgan-plus-voicebank" ), enhance_chunk_s=float(os.environ.get("ENHANCE_CHUNK_S", "1.0")), - enhance_hop_s=float(os.environ.get("ENHANCE_HOP_S", "0.2")), + enhance_hop_s=float(os.environ.get("ENHANCE_HOP_S", "0.1")), + enhance_infer=os.environ.get("ENHANCE_INFER", "auto").lower(), beamform_mode=os.environ.get("BEAMFORM_MODE", "auto").lower(), torch_num_threads=int(os.environ.get("TORCH_NUM_THREADS", "1")), beamform_tdoa_every=int(os.environ.get("BEAMFORM_TDOA_EVERY", "5")), @@ -146,6 +148,9 @@ def load_config(env_path: Optional[Path] = None) -> Config: if cfg.enhance_mode not in ("auto", "force", "off"): raise ValueError(f"ENHANCE_MODE must be auto|force|off, got {cfg.enhance_mode}") + if cfg.enhance_infer not in ("auto", "openvino", "torch"): + raise ValueError( + f"ENHANCE_INFER must be auto|openvino|torch, got {cfg.enhance_infer}") if cfg.beamform_mode not in ("auto", "force", "off"): raise ValueError( f"BEAMFORM_MODE must be auto|force|off, got {cfg.beamform_mode}") diff --git a/ov_metricgan.py b/ov_metricgan.py new file mode 100644 index 0000000..095b8da --- /dev/null +++ b/ov_metricgan.py @@ -0,0 +1,99 @@ +"""OpenVINO CPU compile of the MetricGAN mask net. + +STFT / ISTFT stay in SpeechBrain; only EnhancementGenerator runs here. +Device is CPU — the Arc B580 lost the GPU kill test. +""" +from __future__ import annotations + +import fcntl +import logging +from pathlib import Path + +import numpy as np + +log = logging.getLogger("audio.cleanup") + +CACHE_DIR = Path("/root/.cache/speechbrain-enhancement") +ONNX_NAME = "metricgan_enhance.onnx" +LOCK_NAME = "metricgan_enhance.lock" +SHIPPED_ONNX = Path(__file__).resolve().parent / "scripts" / "metricgan_enhance.onnx" + + +def ensure_onnx(enhancer, torch) -> Path: + """Export or reuse the BLSTM mask ONNX. One process writes; others wait.""" + CACHE_DIR.mkdir(parents=True, exist_ok=True) + onnx_path = CACHE_DIR / ONNX_NAME + lock_path = CACHE_DIR / LOCK_NAME + with open(lock_path, "w") as lockf: + fcntl.flock(lockf.fileno(), fcntl.LOCK_EX) + if onnx_path.is_file() and onnx_path.stat().st_size > 1000: + return onnx_path + if SHIPPED_ONNX.is_file() and SHIPPED_ONNX.stat().st_size > 1000: + onnx_path.write_bytes(SHIPPED_ONNX.read_bytes()) + log.info("copied shipped MetricGAN ONNX to %s", onnx_path) + return onnx_path + log.info("exporting MetricGAN enhance_model to ONNX %s", onnx_path) + import torch.nn as nn + + class MaskNet(nn.Module): + def __init__(self, net): + super().__init__() + self.net = net + + def forward(self, feats): + lengths = torch.ones(feats.shape[0], dtype=feats.dtype) + return self.net(feats, lengths=lengths) + + wrap = MaskNet(enhancer.mods.enhance_model).eval() + feats = enhancer.compute_features(torch.randn(1, 16000)) + tmp = onnx_path.with_suffix(".onnx.tmp") + with torch.no_grad(): + torch.onnx.export( + wrap, + feats, + str(tmp), + input_names=["feats"], + output_names=["mask"], + dynamic_axes={ + "feats": {0: "batch", 1: "time"}, + "mask": {0: "batch", 1: "time"}, + }, + opset_version=17, + dynamo=False, + ) + tmp.replace(onnx_path) + log.info("wrote %s (%d bytes)", onnx_path, onnx_path.stat().st_size) + return onnx_path + + +_COMPILED = None + + +def compile_cpu(onnx_path: Path, num_threads: int = 1): + """Compile ONNX for OpenVINO CPU with a 1-thread infer request.""" + global _COMPILED + if _COMPILED is not None: + compiled, _, ov_input = _COMPILED + request = compiled.create_infer_request() + return compiled, request, ov_input + import openvino as ov + + core = ov.Core() + model = ov.convert_model(str(onnx_path)) + compiled = core.compile_model( + model, + "CPU", + { + "INFERENCE_NUM_THREADS": num_threads, + "NUM_STREAMS": 1, + }, + ) + request = compiled.create_infer_request() + log.info("OpenVINO CPU MetricGAN mask ready (%s)", ov.__version__) + _COMPILED = (compiled, None, compiled.inputs[0]) + return compiled, request, compiled.inputs[0] + + +def infer_mask(request, ov_input, feats: np.ndarray) -> np.ndarray: + result = request.infer({ov_input: feats}) + return np.asarray(next(iter(result.values()))) diff --git a/run.py b/run.py index 6a3a30f..35357c1 100644 --- a/run.py +++ b/run.py @@ -46,6 +46,7 @@ def worker_args_for(cfg, cam) -> list[str]: "--enhance-model", cfg.enhance_model, "--enhance-chunk-s", str(cfg.enhance_chunk_s), "--enhance-hop-s", str(cfg.enhance_hop_s), + "--enhance-infer", cfg.enhance_infer, "--beamform-mode", cfg.beamform_mode, "--torch-num-threads", str(cfg.torch_num_threads), "--tdoa-every", str(cfg.beamform_tdoa_every), diff --git a/run_worker.py b/run_worker.py index 623ea56..3e04b9f 100644 --- a/run_worker.py +++ b/run_worker.py @@ -47,7 +47,8 @@ def main() -> int: ap.add_argument("--enhance-model", default="speechbrain/metricgan-plus-voicebank") ap.add_argument("--enhance-chunk-s", type=float, default=1.0) - ap.add_argument("--enhance-hop-s", type=float, default=0.2) + ap.add_argument("--enhance-hop-s", type=float, default=0.1) + ap.add_argument("--enhance-infer", default="auto") ap.add_argument("--beamform-mode", default="auto") ap.add_argument("--torch-num-threads", type=int, default=1) ap.add_argument("--tdoa-every", type=int, default=5) @@ -81,6 +82,7 @@ def main() -> int: beamform=args.beamform_mode, torch_num_threads=args.torch_num_threads, tdoa_every=args.tdoa_every, + enhance_infer=args.enhance_infer, ) log.info("audio backend: %s", cleaner.backend) diff --git a/tests/test_audio_cleanup.py b/tests/test_audio_cleanup.py index c377381..c57a231 100644 --- a/tests/test_audio_cleanup.py +++ b/tests/test_audio_cleanup.py @@ -94,7 +94,8 @@ def test_process_records_last_dt_ms(): def test_stereo_beamform_then_model_preserves_hop(): c = AudioCleaner( mode="force", beamform="force", chunk_s=0.5, hop_s=0.25) - assert c.backend == "model+beamform" + assert "model" in c.backend + assert "beamform" in c.backend hop = 4000 t = np.arange(hop) ch0 = (np.sin(2 * np.pi * 440 * t / 16000) * 8000).astype(np.int16) @@ -104,3 +105,19 @@ def test_stereo_beamform_then_model_preserves_hop(): out = c.process(stereo) assert out.dtype == np.int16 assert out.shape == (hop,) + + +def test_openvino_cpu_infer_loads_and_preserves_hop(): + c = AudioCleaner( + mode="force", beamform="off", enhance_infer="openvino", + chunk_s=0.5, hop_s=0.25) + assert c.infer_backend == "openvino" + assert c._ov_compiled is not None + assert "ov" in c.backend + hop = 4000 + pcm = (np.sin(2 * np.pi * 440 * np.arange(hop) / 16000) * 8000).astype( + np.int16) + for _ in range(4): + out = c.process(pcm) + assert out.dtype == np.int16 + assert out.shape == (hop,) diff --git a/tests/test_config.py b/tests/test_config.py index 0beda23..0058b62 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -13,3 +13,7 @@ def test_default_beamform_mode_auto(): def test_torch_num_threads_default_one(): assert Config().torch_num_threads == 1 + + +def test_enhance_infer_default_auto(): + assert Config().enhance_infer == "auto"