From 0cac8f7650b0582523d133dc45108510e801748c Mon Sep 17 00:00:00 2001
From: matt3o <matt3o@gmail.com>
Date: Fri, 23 Feb 2024 20:48:17 +0100
Subject: [PATCH] make dlib or insightface optional

---
 faceanalysis.py | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/faceanalysis.py b/faceanalysis.py
index e825b38..3d25d2b 100644
--- a/faceanalysis.py
+++ b/faceanalysis.py
@@ -1,5 +1,26 @@
-import dlib
-from insightface.app import FaceAnalysis
+IS_DLIB_INSTALLED = False
+try:
+    import dlib
+    IS_DLIB_INSTALLED = True
+except ImportError:
+    pass
+
+IS_INSIGHTFACE_INSTALLED = False
+try:
+    from insightface.app import FaceAnalysis
+    IS_INSIGHTFACE_INSTALLED = True
+except ImportError:
+    pass
+
+if not IS_DLIB_INSTALLED and not IS_INSIGHTFACE_INSTALLED:
+    raise Exception("Please install either dlib or insightface to use this node.")
+
+INSTALLED_LIBRARIES = []
+if IS_DLIB_INSTALLED:
+    INSTALLED_LIBRARIES.append("dlib")
+if IS_INSIGHTFACE_INSTALLED:
+    INSTALLED_LIBRARIES.append("insightface")
+
 import torch
 import torchvision.transforms.v2 as T
 import os
@@ -14,7 +35,7 @@ class FaceAnalysisModels:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
-            "library": (["dlib", "insightface"], ),
+            "library": (INSTALLED_LIBRARIES, ),
             "provider": (["CPU", "CUDA", "DirectML", "OpenVINO", "ROCM", "CoreML"], ),
         }}
 
@@ -69,9 +90,6 @@ class FaceEmbedDistance:
 
         self.analysis_models = analysis_models
 
-        #if reference.shape[0] > 1:
-        #    reference = torch.mean(reference, dim=0).unsqueeze(0)
-
         ref = []
         for i in reference:
             ref_emb = self.get_descriptor(np.array(T.ToPILImage()(i.permute(2, 0, 1)).convert('RGB')))