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
@@ -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())))