From cd634a8fbba38883d421248e910c5d5bac219549 Mon Sep 17 00:00:00 2001 From: wfng92 <43742196+wfng92@users.noreply.github.com> Date: Fri, 31 Mar 2023 22:00:59 +0800 Subject: [PATCH] Check for all different packages of opencv (#2901) Co-authored-by: Patrick von Platen --- src/diffusers/utils/import_utils.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 5757ded65d..fd7538b1b5 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -175,9 +175,22 @@ if _onnx_available: # (sayakpaul): importlib.util.find_spec("opencv-python") returns None even when it's installed. # _opencv_available = importlib.util.find_spec("opencv-python") is not None try: - _opencv_version = importlib_metadata.version("opencv-python") - _opencv_available = True - logger.debug(f"Successfully imported cv2 version {_opencv_version}") + candidates = ( + "opencv-python", + "opencv-contrib-python", + "opencv-python-headless", + "opencv-contrib-python-headless", + ) + _opencv_version = None + for pkg in candidates: + try: + _opencv_version = importlib_metadata.version(pkg) + break + except importlib_metadata.PackageNotFoundError: + pass + _opencv_available = _opencv_version is not None + if _opencv_available: + logger.debug(f"Successfully imported cv2 version {_opencv_version}") except importlib_metadata.PackageNotFoundError: _opencv_available = False