| """3 个 axmodel 的 axengine 会话封装 (enc/erb_dec/df_dec 流水线)。""" |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
|
|
| T = 99 |
|
|
|
|
| @dataclass(frozen=True) |
| class TensorInfo: |
| name: str |
| shape: tuple[int, ...] |
| dtype: np.dtype |
|
|
|
|
| def _numpy_dtype(value: Any) -> np.dtype: |
| text = str(value).lower() |
| mapping = ( |
| (("tensor(float)", "float32", "fp32", "f32"), np.float32), |
| (("tensor(float16)", "float16", "fp16", "f16"), np.float16), |
| (("tensor(int64)", "int64", "s64"), np.int64), |
| (("tensor(int32)", "int32", "s32"), np.int32), |
| (("tensor(uint16)", "uint16", "u16"), np.uint16), |
| (("tensor(uint8)", "uint8", "u8"), np.uint8), |
| ) |
| for aliases, dtype in mapping: |
| if any(alias in text for alias in aliases): |
| return np.dtype(dtype) |
| raise ValueError(f"unsupported runtime tensor dtype: {value}") |
|
|
|
|
| def _tensor_info(value: Any) -> TensorInfo: |
| shape = getattr(value, "shape", None) |
| if shape is None: |
| shape = getattr(value, "dims", None) |
| if shape is None or any(dim is None for dim in shape): |
| raise ValueError(f"dynamic or missing tensor shape for {value.name}: {shape}") |
| dtype = getattr(value, "dtype", None) |
| if dtype is None: |
| dtype = getattr(value, "type", None) |
| return TensorInfo( |
| name=value.name, |
| shape=tuple(int(dim) for dim in shape), |
| dtype=_numpy_dtype(dtype), |
| ) |
|
|
|
|
| class _AxSession: |
| """单个 axmodel 的 axengine 会话 (参考 AXERA 官方推理工程封装)。""" |
|
|
| def __init__(self, model_path: str | Path): |
| import axengine |
| self.path = Path(model_path) |
| if not self.path.is_file(): |
| raise FileNotFoundError(self.path) |
| self._session = axengine.InferenceSession(str(self.path)) |
| self.inputs = [_tensor_info(v) for v in self._session.get_inputs()] |
| self.outputs = [_tensor_info(v) for v in self._session.get_outputs()] |
| self.input_by_name = {v.name: v for v in self.inputs} |
|
|
| def run(self, feed: dict[str, np.ndarray]) -> dict[str, np.ndarray]: |
| missing = [v.name for v in self.inputs if v.name not in feed] |
| if missing: |
| raise KeyError(f"missing inputs for {self.path.name}: {missing}") |
| prepared = { |
| name: np.ascontiguousarray(np.asarray(feed[name], dtype=meta.dtype)) |
| for name, meta in self.input_by_name.items() |
| } |
| values = self._session.run(None, prepared) |
| if isinstance(values, dict): |
| return {name: np.asarray(v) for name, v in values.items()} |
| if not isinstance(values, (list, tuple)): |
| values = [values] |
| if len(values) != len(self.outputs): |
| raise RuntimeError( |
| f"unexpected output count from {self.path.name}: " |
| f"{len(values)} != {len(self.outputs)}") |
| return {meta.name: np.asarray(v) for meta, v in zip(self.outputs, values)} |
|
|
|
|
| class DfNet3Ax: |
| """DeepFilterNet3 三子图 NPU 流水线 (与 dsp.pipeline 配合)。 |
| |
| enc: feat_erb(1,1,T,32)+feat_spec(1,2,T,96) -> e0..e3, emb, c0, lsnr |
| erb_dec: emb,e3,e2,e1,e0 -> m(1,1,T,32) |
| df_dec: emb,c0 -> coefs(1,T,96,10) |
| """ |
|
|
| def __init__(self, model_dir: str | Path): |
| model_dir = Path(model_dir) |
| self.enc = _AxSession(model_dir / "enc.axmodel") |
| self.erb_dec = _AxSession(model_dir / "erb_dec.axmodel") |
| self.df_dec = _AxSession(model_dir / "df_dec.axmodel") |
| |
| self.stateful = all("h0" in s.input_by_name for s in (self.enc, self.df_dec)) |
|
|
| def run_enc(self, feat_erb: np.ndarray, feat_spec: np.ndarray, |
| h: np.ndarray | None = None) -> tuple[dict[str, np.ndarray], np.ndarray | None]: |
| feed = {"feat_erb": feat_erb, "feat_spec": feat_spec} |
| if self.stateful: |
| if h is None: |
| h = np.zeros((1, 1, 256), np.float32) |
| feed["h0"] = h |
| out = self.enc.run(feed) |
| return out, np.asarray(out["h_out"]) |
| return self.enc.run(feed), None |
|
|
| def run_erb_dec(self, emb, e3, e2, e1, e0) -> np.ndarray: |
| return self.erb_dec.run({"emb": emb, "e3": e3, "e2": e2, "e1": e1, "e0": e0})["m"] |
|
|
| def run_df_dec(self, emb, c0, h: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray | None]: |
| feed = {"emb": emb, "c0": c0} |
| if self.stateful: |
| if h is None: |
| h = np.zeros((2, 1, 256), np.float32) |
| feed["h0"] = h |
| out = self.df_dec.run(feed) |
| return np.asarray(out["coefs"]), np.asarray(out["h_out"]) |
| return self.df_dec.run(feed)["coefs"], None |
|
|
