From 9fb834448758fc887b10aaabdae79f6dcfb6ce46 Mon Sep 17 00:00:00 2001
From: Laurent <laurent@lagon.tech>
Date: Mon, 14 Oct 2024 14:40:15 +0000
Subject: [PATCH] check conversion deps are installed when using
 refiners.conversion

---
 src/refiners/conversion/__init__.py | 33 +++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/src/refiners/conversion/__init__.py b/src/refiners/conversion/__init__.py
index 13610db..791462f 100644
--- a/src/refiners/conversion/__init__.py
+++ b/src/refiners/conversion/__init__.py
@@ -1,3 +1,36 @@
+import sys
+from importlib import import_module
+from importlib.metadata import requires
+
+from packaging.requirements import Requirement
+
+refiners_requires = requires("refiners")
+assert refiners_requires is not None
+
+# Some dependencies have different module names than their package names
+req_to_module: dict[str, str] = {
+    "huggingface-hub": "huggingface_hub",
+    "segment-anything-py": "segment_anything",
+}
+
+for dep in refiners_requires:
+    req = Requirement(dep)
+    marker = req.marker
+    if marker is None or not marker.evaluate({"extra": "conversion"}):
+        continue
+
+    module_name = req_to_module.get(req.name, req.name)
+
+    try:
+        import_module(module_name)
+    except ImportError:
+        print(
+            f"Some dependencies are missing: {req.name}. "
+            "Please install refiners with the `conversion` extra, e.g. `pip install refiners[conversion]`",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
 from .models import (
     autoencoder_sd15,
     autoencoder_sdxl,