root
ef9573c294
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).
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
"""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())))
|