File size: 3,507 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 89 90 91 92 93 94 95 | #!/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 small AxeSession interface used by the board 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)
# The exported part0 ONNX keeps t/guidance_scale as scalar inputs ([]),
# while the axmodel path feeds them as shape [1]. Normalize only for ONNX.
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 Decoder4ZipVoiceBoardRuntimePart0Onnx(Decoder4ZipVoiceBoardRuntime):
"""Runs decoder part0 with ONNX Runtime, all other split models with axmodel."""
def _load_models(self) -> None:
self.sessions = {}
encoder_name = self.encoder_info["name"]
encoder_path = self.models_dir / self.encoder_info["file"]
logging.debug("Loading %s from %s", encoder_name, encoder_path)
self.sessions[encoder_name] = AxeSession(encoder_path)
for index, info in enumerate(self.decoder_parts):
name = info["name"]
if index == 0:
path = self.models_dir / "fm_decoder_part0.onnx"
logging.info("part0 使用 ONNX Runtime: %s", path)
self.sessions[name] = _OrtSession(path)
else:
path = self.models_dir / info["file"]
logging.debug("Loading %s from %s", name, path)
self.sessions[name] = AxeSession(path)
self.decoder_label = "decoder4(part0_onnx+part1-3_axmodel)"
logging.debug("Loaded encoder axmodel + part0 ONNX + %d decoder axmodels", len(self.decoder_parts) - 1)
|