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).
2026-08-28 07:25:30 +00:00
parent 42443420be
commit ef9573c294
9 changed files with 211 additions and 26 deletions
@@ -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
@@ -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])`
@@ -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,6 +358,9 @@ class AudioCleaner:
            w = torch.from_numpy(self._ctx.copy()).unsqueeze(0)
            if hasattr(self._enhancer, "device"):
                w = w.to(self._enhancer.device)
+            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".
@@ -322,7 +374,7 @@ class AudioCleaner:
                        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(
@@ -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}")
@@ -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())))
@@ -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),
@@ -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)

@@ -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,)
@@ -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"