Update onnxruntime package candidates (#2666)

* update import onnxruntime package, enable onnxruntime-rocm and onnxruntime-training * add ort_nightly_gpu
2026-01-27 17:22:53 +03:00 · 2023-03-24 19:23:05 +08:00
parent f6feb69991
commit c4892f1855
1 changed files with 3 additions and 0 deletions
--- 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