ONNX ROCm Linux.

installer.py

@@ -440,6 +440,14 @@ def check_torch():
             # ROCm 5.5 is oldest for PyTorch 2.1
             torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/rocm5.5')
         xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
+        if rocm_ver is not None:
+            install(os.environ.get('ONNXRUNTIME_PACKAGE', get_onnxruntime_source_for_rocm(arr)), "onnxruntime-training built with ROCm", ignore=True)
+            try:
+                import onnxruntime
+                if "ROCMExecutionProvider" not in onnxruntime.get_available_providers():
+                    log.warn('Failed to automatically install onxnruntime package for ROCm. Please manually install it if you need.')
+            except Exception:
+                pass
     elif allow_ipex and (args.use_ipex or shutil.which('sycl-ls') is not None or shutil.which('sycl-ls.exe') is not None or os.environ.get('ONEAPI_ROOT') is not None or os.path.exists('/opt/intel/oneapi') or os.path.exists("C:/Program Files (x86)/Intel/oneAPI") or os.path.exists("C:/oneAPI")):
         args.use_ipex = True # pylint: disable=attribute-defined-outside-init
         log.info('Intel OneAPI Toolkit detected')
@@ -833,6 +841,24 @@ def get_version():
     return version
 
 
+def get_onnxruntime_source_for_rocm(rocm_ver):
+    ort_version = "1.16.3"
+
+    try:
+        import onnxruntime
+        ort_version = onnxruntime.__version__
+    except ImportError:
+        pass
+
+    cp_str = f"{sys.version_info.major}{sys.version_info.minor}"
+
+    if rocm_ver is None:
+        command = subprocess.run('hipconfig --version', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        rocm_ver = command.stdout.decode(encoding="utf8", errors="ignore").split('.')
+
+    return f"https://download.onnxruntime.ai/onnxruntime_training-{ort_version}%2Brocm{rocm_ver[0]}{rocm_ver[1]}-cp{cp_str}-cp{cp_str}-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"
+
+
 # check version of the main repo and optionally upgrade it
 def check_version(offline=False, reset=True): # pylint: disable=unused-argument
     if args.skip_all:

modules/onnx.py

@@ -141,7 +141,11 @@ def initialize():
     if initialized:
         return
 
-    from modules import onnx_pipelines as pipelines
+    from modules import onnx_pipelines as pipelines, devices
+    from modules.onnx_ep import ExecutionProvider, TORCH_DEVICE_TO_EP
+
+    if devices.backend == "rocm":
+        TORCH_DEVICE_TO_EP["cuda"] = ExecutionProvider.ROCm
 
     # OnnxRuntimeModel Hijack.
     OnnxRuntimeModel.__module__ = 'diffusers'

modules/onnx_ep.py

@@ -10,6 +10,7 @@ class ExecutionProvider(str, Enum):
     DirectML = "DmlExecutionProvider"
     CUDA = "CUDAExecutionProvider"
     ROCm = "ROCMExecutionProvider"
+    MIGraphX = "MIGraphXExecutionProvider"
     OpenVINO = "OpenVINOExecutionProvider"
 
 
@@ -19,6 +20,7 @@ EP_TO_NAME = {
     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 = {
@@ -42,7 +44,7 @@ def get_default_execution_provider() -> ExecutionProvider:
         if ExecutionProvider.ROCm in available_execution_providers:
             return ExecutionProvider.ROCm
         else:
-            log.warning("Currently, there's no pypi release for onnxruntime-rocm. Please download and install .whl file from https://download.onnxruntime.ai/")
+            return ExecutionProvider.CPU
     elif devices.backend == "ipex" or devices.backend == "openvino":
         return ExecutionProvider.OpenVINO
     return ExecutionProvider.CPU
@@ -59,8 +61,6 @@ def get_execution_provider_options():
         if ExecutionProvider.ROCm in available_execution_providers:
             execution_provider_options["tunable_op_enable"] = 1
             execution_provider_options["tunable_op_tuning_enable"] = 1
-        else:
-            log.warning("Currently, there's no pypi release for onnxruntime-rocm. Please download and install .whl file from https://download.onnxruntime.ai/ The inference will be fall back to CPU.")
     elif opts.onnx_execution_provider == ExecutionProvider.OpenVINO:
         from modules.intel.openvino import get_device as get_raw_openvino_device
         raw_openvino_device = get_raw_openvino_device()
@@ -79,7 +79,7 @@ def get_provider() -> Tuple:
 
 
 def install_execution_provider(ep: ExecutionProvider):
-    from installer import pip, uninstall, installed
+    from installer import pip, uninstall, installed, get_onnxruntime_source_for_rocm
 
     if installed("onnxruntime"):
         uninstall("onnxruntime")
@@ -103,12 +103,7 @@ def install_execution_provider(ep: ExecutionProvider):
             log.warn("ROCMExecutionProvider is not supported on Windows.")
             return
 
-        try:
-            major, minor = sys.version_info
-            cp_str = f"{major}{minor}"
-            packages.append(f"https://download.onnxruntime.ai/onnxruntime_training-1.16.3%2Brocm56-cp{cp_str}-cp{cp_str}-manylinux_2_17_x86_64.manylinux2014_x86_64.whl")
-        except Exception:
-            log.warn("Failed to install onnxruntime for ROCm.")
+        packages.append(get_onnxruntime_source_for_rocm())
     elif ep == ExecutionProvider.OpenVINO:
         if installed("openvino"):
             uninstall("openvino")