From 0a08098511e80abf432b0494d361667279b4a0dd Mon Sep 17 00:00:00 2001
From: wfng92 <43742196+wfng92@users.noreply.github.com>
Date: Thu, 30 Mar 2023 14:26:36 +0800
Subject: [PATCH] Check for all different packages of opencv

---
 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 5757ded65dac..fd7538b1b5e9 100644
--- a/src/diffusers/utils/import_utils.py
+++ b/src/diffusers/utils/import_utils.py
@@ -175,9 +175,22 @@
 # (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