From 4dd532185223a1fa5978446e7bb75d32d77ebdb5 Mon Sep 17 00:00:00 2001
From: melMass <mel@melmassadian.com>
Date: Mon, 7 Aug 2023 23:53:34 +0200
Subject: [PATCH] =?UTF-8?q?fix:=20=E2=AC=87=EF=B8=8F=20download=5Fantelope?=
 =?UTF-8?q?v2?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

the url used in insightface returns 404.
fixes #55
---
 nodes/faceswap.py |  5 ++++-
 utils.py          | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/nodes/faceswap.py b/nodes/faceswap.py
index faf433d..b21396b 100644
--- a/nodes/faceswap.py
+++ b/nodes/faceswap.py
@@ -12,7 +12,7 @@ import os
 import tempfile
 import torch
 from insightface.model_zoo.inswapper import INSwapper
-from ..utils import pil2tensor, tensor2pil
+from ..utils import pil2tensor, tensor2pil, download_antelopev2
 from ..log import mklog, NullWriter
 import sys
 import comfy.model_management as model_management
@@ -53,6 +53,9 @@ class LoadFaceAnalysisModel:
     CATEGORY = "mtb/facetools"
 
     def load_model(self, faceswap_model: str):
+        if faceswap_model == "antelopev2":
+            download_antelopev2()
+
         face_analyser = insightface.app.FaceAnalysis(
             name=faceswap_model,
             root=os.path.join(folder_paths.models_dir, "insightface"),
diff --git a/utils.py b/utils.py
index 1152504..4068c6b 100644
--- a/utils.py
+++ b/utils.py
@@ -4,6 +4,7 @@ import torch
 from pathlib import Path
 import sys
 from typing import List
+from .log import log
 
 
 # region MISC Utilities
@@ -120,3 +121,43 @@ def tensor2np(tensor: torch.Tensor) -> List[np.ndarray]:
 
 
 # endregion
+
+# region MODEL Utilities
+def download_antelopev2():
+    antelopev2_url = "https://drive.google.com/uc?id=18wEUfMNohBJ4K3Ly5wpTejPfDzp-8fI8"
+
+    try:
+        import gdown
+
+        import folder_paths
+
+        log.debug("Loading antelopev2 model")
+
+        dest = Path(folder_paths.models_dir) / "insightface"
+        archive = dest / "antelopev2.zip"
+        final_path = dest / "models" / "antelopev2"
+        if not final_path.exists():
+            log.info(f"antelopev2 not found, downloading to {dest}")
+            gdown.download(
+                antelopev2_url,
+                archive.as_posix(),
+                resume=True,
+            )
+
+            log.info(f"Unzipping antelopev2 to {final_path}")
+
+            if archive.exists():
+                # we unzip it
+                import zipfile
+
+                with zipfile.ZipFile(archive.as_posix(), "r") as zip_ref:
+                    zip_ref.extractall(final_path.parent.as_posix())
+
+    except Exception as e:
+        log.error(
+            f"Could not load or download antelopev2 model, download it manually from {antelopev2_url}"
+        )
+        raise e
+
+
+# endregion