File size: 3,036 Bytes
e405f21 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | #!/usr/bin/env python3
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))
|