| |
|
|
| from __future__ import annotations |
|
|
| import logging |
| from pathlib import Path |
| from typing import Dict, List |
|
|
| import numpy as np |
| import onnxruntime as ort |
|
|
| from scripts.zipvoice_decoder4_runtime import Decoder4ZipVoiceBoardRuntime |
| from scripts.zipvoice_runtime import AxeSession |
|
|
|
|
| class _OrtSession: |
| """ONNX Runtime wrapper with the AxeSession interface used by split runtime.""" |
|
|
| _TYPE_TO_DTYPE = { |
| "tensor(float)": np.float32, |
| "tensor(float32)": np.float32, |
| "tensor(double)": np.float64, |
| "tensor(int32)": np.int32, |
| "tensor(int64)": np.int64, |
| "tensor(uint8)": np.uint8, |
| "tensor(bool)": np.bool_, |
| } |
|
|
| def __init__(self, model_path: str | Path): |
| self.path = Path(model_path) |
| if not self.path.exists(): |
| raise FileNotFoundError(f"ONNX model not found: {self.path}") |
| self._session = ort.InferenceSession( |
| str(self.path), |
| providers=["CPUExecutionProvider"], |
| ) |
| self._inputs = self._session.get_inputs() |
| self._outputs = self._session.get_outputs() |
|
|
| @property |
| def input_names(self) -> List[str]: |
| return [item.name for item in self._inputs] |
|
|
| @property |
| def output_names(self) -> List[str]: |
| return [item.name for item in self._outputs] |
|
|
| def _coerce_one(self, name: str, value: np.ndarray) -> np.ndarray: |
| info = next((item for item in self._inputs if item.name == name), None) |
| array = np.asarray(value) |
| if info is None: |
| return np.ascontiguousarray(array) |
|
|
| dtype = self._TYPE_TO_DTYPE.get(str(info.type).lower()) |
| if dtype is not None: |
| array = array.astype(dtype, copy=False) |
|
|
| if list(info.shape) == [] and array.shape == (1,): |
| array = array.reshape(()) |
|
|
| return np.ascontiguousarray(array) |
|
|
| def run(self, feed_dict: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: |
| feed = {name: self._coerce_one(name, value) for name, value in feed_dict.items()} |
| outputs = self._session.run(None, feed) |
| return {name: value for name, value in zip(self.output_names, outputs)} |
|
|
|
|
| class Decoder4ZipVoiceBoardRuntimeEncoderOnnx(Decoder4ZipVoiceBoardRuntime): |
| """Runs encoder with ONNX Runtime, all decoder parts with axmodel.""" |
|
|
| def _load_models(self) -> None: |
| self.sessions = {} |
|
|
| encoder_name = self.encoder_info["name"] |
| encoder_path = self.models_dir / "encoder_core.onnx" |
| logging.info("encoder 使用 ONNX Runtime: %s", encoder_path) |
| self.sessions[encoder_name] = _OrtSession(encoder_path) |
|
|
| for info in self.decoder_parts: |
| name = info["name"] |
| path = self.models_dir / info["file"] |
| logging.debug("Loading %s from %s", name, path) |
| self.sessions[name] = AxeSession(path) |
|
|
| self.decoder_label = "decoder4(encoder_onnx+part0-3_axmodel)" |
| logging.debug("Loaded encoder ONNX + %d decoder axmodels", len(self.decoder_parts)) |
|
|
