Update onnxruntime package candidates (#2666)

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

Update onnxruntime package candidates (#2666)
* update import onnxruntime package, enable onnxruntime-rocm and onnxruntime-training * add ort_nightly_gpu
c4892f18 · PeixuanZuo · GitHub · f6feb699 · c4892f18
Unverified Commit c4892f18 authored Mar 24, 2023 by PeixuanZuo Committed by GitHub Mar 24, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

src/diffusers/utils/import_utils.py src/diffusers/utils/import_utils.py +3 -0

No files found.
--- 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