From c4892f1855097a68703ca2e949aca15829526958 Mon Sep 17 00:00:00 2001
From: PeixuanZuo <94887879+PeixuanZuo@users.noreply.github.com>
Date: Fri, 24 Mar 2023 19:23:05 +0800
Subject: [PATCH] Update onnxruntime package candidates (#2666)

* update import onnxruntime package, enable onnxruntime-rocm and onnxruntime-training

* add ort_nightly_gpu
---
 src/diffusers/utils/import_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py
index 7cb72525c9..5757ded65d 100644
--- a/src/diffusers/utils/import_utils.py
+++ b/src/diffusers/utils/import_utils.py
@@ -153,9 +153,12 @@ if _onnx_available:
     candidates = (
         "onnxruntime",
         "onnxruntime-gpu",
+        "ort_nightly_gpu",
         "onnxruntime-directml",
         "onnxruntime-openvino",
         "ort_nightly_directml",
+        "onnxruntime-rocm",
+        "onnxruntime-training",
     )
     _onnxruntime_version = None
     # For the metadata, we have to look for both onnxruntime and onnxruntime-gpu