| """RNNoise AX620E 推理会话（NPU 专用发布版：无 onnxruntime/torch 回退）。""" |
| import numpy as np |
|
|
| from . import dsp |
|
|
| DEFAULT_PROVIDER = "AxEngineExecutionProvider" |
|
|
| INPUT_NAMES = ["features", "conv1_mem", "conv2_mem", |
| "gru1_s", "gru2_s", "gru3_s"] |
| INPUT_SHAPES = {"features": (1, 65), "conv1_mem": (1, 130), |
| "conv2_mem": (1, 256), "gru1_s": (1, 384), |
| "gru2_s": (1, 384), "gru3_s": (1, 384)} |
| OUTPUT_NAMES = ["gains", "vad", "conv1_mem_new", "conv2_mem_new", |
| "gru1_s_new", "gru2_s_new", "gru3_s_new"] |
| _STATE_INPUTS = ["conv1_mem", "conv2_mem", "gru1_s", "gru2_s", "gru3_s"] |
|
|
|
|
| class RNNoiseDenoiser: |
| """48k 单声道实时降噪器（AX 芯片端到端，无 CPU 回退）。""" |
|
|
| def __init__(self, model_path, providers=None): |
| try: |
| import axengine as axe |
| except ImportError as exc: |
| raise RuntimeError( |
| "SDK 为 NPU 专用发布版，仅支持在 AX 芯片上运行；请先安装 " |
| "requirements.txt 并在板端执行（无 onnxruntime/torch 回退）" |
| ) from exc |
| self.session = axe.InferenceSession( |
| model_path, providers=providers or [DEFAULT_PROVIDER]) |
| self.backend = "axengine" |
| self.input_names = [i.name for i in self.session.get_inputs()] |
| self.output_names = [o.name for o in self.session.get_outputs()] |
| self.reset() |
|
|
| def reset(self): |
| self.st = dsp.RNNoiseState() |
| self._states = {k: np.zeros(INPUT_SHAPES[k], dtype=np.float32) |
| for k in _STATE_INPUTS} |
|
|
| def process_frame(self, frame): |
| frame = np.asarray(frame, dtype=np.float32).reshape(-1) |
| if frame.size != dsp.FRAME_SIZE: |
| raise ValueError(f"帧长必须为 {dsp.FRAME_SIZE}，实际 {frame.size}") |
| ana = dsp.analyze_frame(self.st, frame) |
| if ana["silence"]: |
| out, _ = dsp.synthesize_frame(self.st, ana, None, 0.0) |
| return out, 0.0 |
| feeds = { |
| "features": np.ascontiguousarray( |
| ana["features"][None, :].astype(np.float32)), |
| } |
| feeds.update({k: np.ascontiguousarray(v) |
| for k, v in self._states.items()}) |
| outs = self.session.run(None, feeds) |
| out_idx = {n: i for i, n in enumerate(self.output_names)} |
| gains = outs[out_idx["gains"]] |
| vad = outs[out_idx["vad"]] |
| for k in _STATE_INPUTS: |
| self._states[k] = np.asarray( |
| outs[out_idx[k + "_new"]], dtype=np.float32) |
| out, vad = dsp.synthesize_frame(self.st, ana, gains, vad) |
| return out, float(vad) |
|
|
| def process(self, pcm): |
| pcm = np.asarray(pcm, dtype=np.float32) |
| if pcm.ndim == 1: |
| n = pcm.size // dsp.FRAME_SIZE |
| frames = pcm[:n * dsp.FRAME_SIZE].reshape(n, dsp.FRAME_SIZE) |
| else: |
| frames = pcm.reshape(-1, dsp.FRAME_SIZE) |
| outs, vads = [], [] |
| for fr in frames: |
| o, v = self.process_frame(fr) |
| outs.append(o) |
| vads.append(v) |
| return np.concatenate(outs), np.asarray(vads, dtype=np.float32) |
|
|
