Mirrors
/
diffusers
mirror of https://github.com/ShivamShrirao/diffusers.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix is_onnx_available (#440) 

* Fix is_onnx_available

Fix: If user install onnxruntime-gpu, is_onnx_available() will return False.

* add more onnxruntime candidates

* Run `make style`

Co-authored-by: anton-l <anton@huggingface.co>
This commit is contained in:
SkyTNT 2022-09-16 18:13:22 +08:00 committed by GitHub
parent 214520c66a
commit 37c9d789aa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
src/diffusers/utils/import_utils.py
View File

@ -137,11 +137,19 @@ except importlib_metadata.PackageNotFoundError:
_onnx_available = importlib.util.find_spec("onnxruntime") is not None
if _onnx_available:
candidates = ("onnxruntime", "onnxruntime-gpu", "onnxruntime-directml", "onnxruntime-openvino")
_onnxruntime_version = None
# For the metadata, we have to look for both onnxruntime and onnxruntime-gpu
for pkg in candidates:
try:
_onnxruntime_version = importlib_metadata.version("onnxruntime")
logger.debug(f"Successfully imported onnxruntime version {_onnxruntime_version}")
_onnxruntime_version = importlib_metadata.version(pkg)
break
except importlib_metadata.PackageNotFoundError:
_onnx_available = False
pass
_onnx_available = _onnxruntime_version is not None
if _onnx_available:
logger.debug(f"Successfully imported onnxruntime version {_onnxruntime_version}")
_scipy_available = importlib.util.find_spec("scipy") is not None