| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
| from typing import Dict, List |
|
|
| import numpy as np |
|
|
| from contextlib import contextmanager |
|
|
|
|
| @contextmanager |
| def _suppress_native_stdio(): |
| if os.environ.get("ZIPVOICE_AXERA_VERBOSE"): |
| yield |
| return |
|
|
| stdout_fd = os.dup(1) |
| stderr_fd = os.dup(2) |
| try: |
| with open(os.devnull, "w") as devnull: |
| os.dup2(devnull.fileno(), 1) |
| os.dup2(devnull.fileno(), 2) |
| yield |
| finally: |
| os.dup2(stdout_fd, 1) |
| os.dup2(stderr_fd, 2) |
| os.close(stdout_fd) |
| os.close(stderr_fd) |
|
|
| try: |
| with _suppress_native_stdio(): |
| import axengine |
| HAS_AXENGINE = True |
| except ImportError: |
| HAS_AXENGINE = False |
|
|
|
|
| class AxeSession: |
| """Thin wrapper around axengine.InferenceSession for one axmodel.""" |
|
|
| def __init__(self, model_path: str | Path): |
| if not HAS_AXENGINE: |
| raise RuntimeError("axengine not available; run this script on the AX650 board") |
| self.path = Path(model_path) |
| if not self.path.exists(): |
| raise FileNotFoundError(f"Model not found: {self.path}") |
| with _suppress_native_stdio(): |
| self._session = axengine.InferenceSession(str(self.path)) |
| 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 run(self, feed_dict: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: |
| outputs = self._session.run(None, feed_dict) |
| if isinstance(outputs, dict): |
| return outputs |
| if isinstance(outputs, (list, tuple)): |
| return {name: value for name, value in zip(self.output_names, outputs)} |
| if len(self.output_names) == 1: |
| return {self.output_names[0]: outputs} |
| raise RuntimeError(f"Unexpected AXEngine output type from {self.path.name}: {type(outputs)}") |
|
|
