Spaces:
Running on Zero
Running on Zero
File size: 1,744 Bytes
37aeb5b 24f5025 37aeb5b | 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 | import onnxruntime
import torch
# ZeroGPU containers ship CUDA 13; the onnxruntime-gpu wheel links against
# libcublasLt.so.12 (CUDA 12) → CUDAExecutionProvider fails to load. Use CPU.
providers = ['CPUExecutionProvider']
def load_onnx(file_path: str):
assert file_path.endswith(".onnx")
sess_opt = onnxruntime.SessionOptions()
ort_session = onnxruntime.InferenceSession(file_path, sess_opt=sess_opt, providers=providers)
return ort_session
def load_onnx_caller(file_path: str, single_output=False):
ort_session = load_onnx(file_path)
def caller(*args):
torch_input = isinstance(args[0], torch.Tensor)
if torch_input:
torch_input_dtype = args[0].dtype
torch_input_device = args[0].device
# check all are torch.Tensor and have same dtype and device
assert all([isinstance(arg, torch.Tensor) for arg in args]), "All inputs should be torch.Tensor, if first input is torch.Tensor"
assert all([arg.dtype == torch_input_dtype for arg in args]), "All inputs should have same dtype, if first input is torch.Tensor"
assert all([arg.device == torch_input_device for arg in args]), "All inputs should have same device, if first input is torch.Tensor"
args = [arg.cpu().float().numpy() for arg in args]
ort_inputs = {ort_session.get_inputs()[idx].name: args[idx] for idx in range(len(args))}
ort_outs = ort_session.run(None, ort_inputs)
if torch_input:
ort_outs = [torch.tensor(ort_out, dtype=torch_input_dtype, device=torch_input_device) for ort_out in ort_outs]
if single_output:
return ort_outs[0]
return ort_outs
return caller
|