File size: 3,304 Bytes
3fef26f | 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 | """Official FuXi-S2S ONNX adapter.
The public FuXi-S2S release contains an ONNX model and inference code. This
module exposes the official inference contract through ONNX Runtime.
"""
from pathlib import Path
import numpy as np
class FuXiS2SModel:
def __init__(self, model_path, device="cpu", providers=None):
self.model_path = Path(model_path)
self.device = device
self.providers = list(providers) if providers else self._default_providers(device)
try:
import onnxruntime as ort
except ImportError as exc:
raise RuntimeError("onnxruntime is required for official FuXi-S2S inference") from exc
if not self.model_path.is_file():
raise FileNotFoundError(f"FuXi-S2S ONNX model not found: {self.model_path}")
external_data_path = self.model_path.parent / "fuxi_s2s"
if not external_data_path.is_file():
raise FileNotFoundError(
f"FuXi-S2S ONNX external data not found: {external_data_path}. "
"The official archive requires both fuxi_s2s.onnx and fuxi_s2s."
)
available_providers = ort.get_available_providers()
missing_providers = [provider for provider in self.providers if provider not in available_providers]
if missing_providers:
raise RuntimeError(
f"ONNX Runtime provider(s) unavailable: {missing_providers}. "
f"Available providers: {available_providers}. "
"Install the matching ONNX Runtime build and update model.providers."
)
session_options = ort.SessionOptions()
if device == "dcu":
# CPU constant folding can warn on kernels that are intended for the
# accelerator. Let the configured DCU provider execute the graph.
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
self.session = ort.InferenceSession(
str(self.model_path),
sess_options=session_options,
providers=self.providers,
)
@staticmethod
def _default_providers(device):
if device == "cpu":
return ["CPUExecutionProvider"]
if device == "cuda":
return ["CUDAExecutionProvider", "CPUExecutionProvider"]
if device == "dcu":
raise ValueError(
"DCU inference requires an explicit ONNX Runtime provider in "
"model.providers, for example ['DCUExecutionProvider', "
"'CPUExecutionProvider']"
)
raise ValueError("device must be 'cpu', 'cuda', or 'dcu'")
@property
def input_names(self):
return [item.name for item in self.session.get_inputs()]
@property
def output_names(self):
return [item.name for item in self.session.get_outputs()]
def step(self, inputs):
payload = {}
for name in self.input_names:
if name not in inputs:
raise KeyError(f"Missing FuXi-S2S ONNX input: {name}")
payload[name] = np.asarray(inputs[name], dtype=np.float32)
outputs = self.session.run(None, payload)
return dict(zip(self.output_names, outputs))
def __call__(self, inputs):
return self.step(inputs)
|