mirror of https://github.com/vladmandic/automatic
117 lines
4.5 KiB
Python
117 lines
4.5 KiB
Python
import sys
|
|
from enum import Enum
|
|
from typing import Tuple, List
|
|
import onnxruntime as ort
|
|
from installer import log
|
|
from modules import devices
|
|
|
|
|
|
class ExecutionProvider(str, Enum):
|
|
CPU = "CPUExecutionProvider"
|
|
DirectML = "DmlExecutionProvider"
|
|
CUDA = "CUDAExecutionProvider"
|
|
ROCm = "ROCMExecutionProvider"
|
|
MIGraphX = "MIGraphXExecutionProvider"
|
|
OpenVINO = "OpenVINOExecutionProvider"
|
|
|
|
|
|
available_execution_providers: List[ExecutionProvider] = ort.get_available_providers()
|
|
EP_TO_NAME = {
|
|
ExecutionProvider.CPU: "gpu-cpu", # ???
|
|
ExecutionProvider.DirectML: "gpu-dml",
|
|
ExecutionProvider.CUDA: "gpu-cuda", # test required
|
|
ExecutionProvider.ROCm: "gpu-rocm", # test required
|
|
ExecutionProvider.MIGraphX: "gpu-migraphx", # test required
|
|
ExecutionProvider.OpenVINO: "gpu-openvino", # test required
|
|
}
|
|
TORCH_DEVICE_TO_EP = {
|
|
"cpu": ExecutionProvider.CPU if devices.backend != "openvino" else ExecutionProvider.OpenVINO,
|
|
"cuda": ExecutionProvider.CUDA,
|
|
"xpu": ExecutionProvider.OpenVINO,
|
|
"privateuseone": ExecutionProvider.DirectML,
|
|
"meta": None,
|
|
}
|
|
|
|
|
|
def get_default_execution_provider() -> ExecutionProvider:
|
|
if devices.backend == "cpu":
|
|
return ExecutionProvider.CPU
|
|
elif devices.backend == "directml":
|
|
return ExecutionProvider.DirectML
|
|
elif devices.backend == "cuda":
|
|
return ExecutionProvider.CUDA
|
|
elif devices.backend == "rocm":
|
|
return ExecutionProvider.ROCm
|
|
elif devices.backend == "ipex" or devices.backend == "openvino":
|
|
return ExecutionProvider.OpenVINO
|
|
return ExecutionProvider.CPU
|
|
|
|
|
|
def get_execution_provider_options():
|
|
from modules.shared import cmd_opts, opts
|
|
execution_provider_options = { "device_id": int(cmd_opts.device_id or 0) }
|
|
if opts.onnx_execution_provider == ExecutionProvider.ROCm:
|
|
if ExecutionProvider.ROCm in available_execution_providers:
|
|
execution_provider_options["tunable_op_enable"] = True
|
|
execution_provider_options["tunable_op_tuning_enable"] = True
|
|
elif opts.onnx_execution_provider == ExecutionProvider.OpenVINO:
|
|
from modules.intel.openvino import get_device as get_raw_openvino_device
|
|
device = get_raw_openvino_device()
|
|
if "HETERO:" not in device:
|
|
if opts.olive_float16:
|
|
device = f"{device}_FP16"
|
|
else:
|
|
device = f"{device}_FP32"
|
|
else:
|
|
device = ""
|
|
available_devices = opts.openvino_devices.copy()
|
|
available_devices.remove("CPU")
|
|
for hetero_device in available_devices:
|
|
if opts.olive_float16:
|
|
device = f"{device},{hetero_device}_FP16"
|
|
else:
|
|
device = f"{device},{hetero_device}_FP32"
|
|
if "CPU" in opts.openvino_devices:
|
|
if opts.olive_float16:
|
|
device = f"{device},CPU_FP16"
|
|
else:
|
|
device = f"{device},CPU_FP32"
|
|
device = f"HETERO:{device[1:]}"
|
|
|
|
execution_provider_options["device_type"] = device
|
|
del execution_provider_options["device_id"]
|
|
return execution_provider_options
|
|
|
|
|
|
def get_provider() -> Tuple:
|
|
from modules.shared import opts
|
|
return (opts.onnx_execution_provider, get_execution_provider_options(),)
|
|
|
|
|
|
def install_execution_provider(ep: ExecutionProvider):
|
|
import imp # pylint: disable=deprecated-module
|
|
from installer import installed, install, uninstall
|
|
res = "<br><pre>"
|
|
res += uninstall(["onnxruntime", "onnxruntime-directml", "onnxruntime-gpu", "onnxruntime-training", "onnxruntime-openvino"], quiet=True)
|
|
installed("onnxruntime", reload=True)
|
|
packages = ["onnxruntime"] # Failed to load olive: cannot import name '__version__' from 'onnxruntime'
|
|
if ep == ExecutionProvider.DirectML:
|
|
packages.append("onnxruntime-directml")
|
|
elif ep == ExecutionProvider.CUDA:
|
|
packages.append("onnxruntime-gpu")
|
|
elif ep == ExecutionProvider.ROCm:
|
|
if "linux" not in sys.platform:
|
|
log.warning("ROCMExecutionProvider is not supported on Windows.")
|
|
return
|
|
packages.append("--pre onnxruntime-training --index-url https://pypi.lsh.sh/60 --extra-index-url https://pypi.org/simple")
|
|
elif ep == ExecutionProvider.OpenVINO:
|
|
packages.append("openvino")
|
|
packages.append("onnxruntime-openvino")
|
|
for package in packages:
|
|
res += install(package)
|
|
res += '</pre><br>'
|
|
res += 'Server restart required'
|
|
log.info("Server restart required")
|
|
imp.reload(ort)
|
|
return res
|