From e28de773115d68c97e9e2e124bf935596bf2c474 Mon Sep 17 00:00:00 2001
From: toto <tototianhao@gmail.com>
Date: Mon, 23 Oct 2023 16:20:59 +0800
Subject: [PATCH] update project struct

---
 __init__.py                                   | 11 ++++---
 requirements.txt                              |  1 +
 {utils/protrait => scripts}/__init__.py       |  0
 config.py => scripts/config.py                |  0
 node.py => scripts/nodes.py                   | 32 ++++++++++++-------
 scripts/utils/__init__.py                     |  0
 .../utils/face_process_utils.py               |  0
 .../protrait => scripts/utils}/img_utils.py   |  0
 8 files changed, 27 insertions(+), 17 deletions(-)
 rename {utils/protrait => scripts}/__init__.py (100%)
 rename config.py => scripts/config.py (100%)
 rename node.py => scripts/nodes.py (90%)
 create mode 100644 scripts/utils/__init__.py
 rename face_process_utils.py => scripts/utils/face_process_utils.py (100%)
 rename {utils/protrait => scripts/utils}/img_utils.py (100%)

diff --git a/__init__.py b/__init__.py
index 9543dba..593ec65 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,15 +1,17 @@
-import subprocess
 import sys
+from scripts.config import *
+sys.path.append(root_path)
+
+
+import subprocess
 import threading
 
 import requests
 from tqdm import tqdm
-from .config import *
-
+from scripts.nodes import *
 # import pydevd_pycharm
 # pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
 
-sys.path.append(utils_path)
 
 def handle_stream(stream, prefix):
     for line in stream:
@@ -34,7 +36,6 @@ print("##  installing dependencies")
 requirements_path = os.path.join(root_path, "requirements.txt")
 run_script([sys.executable, '-s', '-m', 'pip', 'install', '-q', '-r', requirements_path])
 
-from .node import *
 
 def urldownload_progressbar(url, file_path):
     response = requests.get(url, stream=True)
diff --git a/requirements.txt b/requirements.txt
index ec4e96f..179a4ce 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,4 +6,5 @@ onnxruntime
 modelscope
 scikit-image
 matplotlib
+insightface
 diffusers==0.18.2
\ No newline at end of file
diff --git a/utils/protrait/__init__.py b/scripts/__init__.py
similarity index 100%
rename from utils/protrait/__init__.py
rename to scripts/__init__.py
diff --git a/config.py b/scripts/config.py
similarity index 100%
rename from config.py
rename to scripts/config.py
diff --git a/node.py b/scripts/nodes.py
similarity index 90%
rename from node.py
rename to scripts/nodes.py
index f592d35..af61a2c 100644
--- a/node.py
+++ b/scripts/nodes.py
@@ -6,9 +6,11 @@ from PIL import Image
 from modelscope.outputs import OutputKeys
 from modelscope.pipelines import pipeline
 from modelscope.utils.constant import Tasks
-from .face_process_utils import call_face_crop, color_transfer, Face_Skin
-from protrait.img_utils import img_to_tensor, tensor_to_img, tensor_to_np, np_to_tensor, np_to_mask, img_to_mask
-from .config import models_path
+from utils.face_process_utils import call_face_crop, color_transfer, Face_Skin
+from utils.img_utils import img_to_tensor, tensor_to_img, tensor_to_np, np_to_tensor, np_to_mask, img_to_mask
+import insightface
+from insightface.app import FaceAnalysis
+from insightface.data import get_image as ins_get_image
 
 # import pydevd_pycharm
 # pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
@@ -42,8 +44,9 @@ class FaceFusionPM:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": {"image": ("IMAGE",),
-                             "user_image": ("IMAGE",),
+        return {"required": {"source_image": ("IMAGE",),
+                             "swap_image": ("IMAGE",),
+                             "mode": (["ali", "roop"],),
                              }}
 
     RETURN_TYPES = ("IMAGE",)
@@ -51,13 +54,18 @@ class FaceFusionPM:
 
     CATEGORY = "protrait/model"
 
-    def img_face_fusion(self, image, user_image):
-        image = tensor_to_img(image)
-        user_image = tensor_to_img(user_image)
-        fusion_image = self.image_face_fusion(dict(template=image, user=user_image))[
-            OutputKeys.OUTPUT_IMG]
-        # swap_face(target_img=output_image, source_img=roop_image, model="inswapper_128.onnx", upscale_options=UpscaleOptions())
-        fusion_image = Image.fromarray(cv2.cvtColor(fusion_image, cv2.COLOR_BGR2RGB))
+    def img_face_fusion(self, source_image, swap_image, mode):
+        result_image = None
+        if mode == "ali":
+            source_image = tensor_to_img(source_image)
+            swap_image = tensor_to_img(swap_image)
+            fusion_image = self.image_face_fusion(dict(template=source_image, user=source_image))[
+                OutputKeys.OUTPUT_IMG]
+            # swap_face(target_img=output_image, source_img=roop_image, model="inswapper_128.onnx", upscale_options=UpscaleOptions())
+            result_image = Image.fromarray(cv2.cvtColor(fusion_image, cv2.COLOR_BGR2RGB))
+        else:
+            app = FaceAnalysis(name='buffalo_l')
+
         return (img_to_tensor(fusion_image),)
 
 class RatioMerge2Image:
diff --git a/scripts/utils/__init__.py b/scripts/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/face_process_utils.py b/scripts/utils/face_process_utils.py
similarity index 100%
rename from face_process_utils.py
rename to scripts/utils/face_process_utils.py
diff --git a/utils/protrait/img_utils.py b/scripts/utils/img_utils.py
similarity index 100%
rename from utils/protrait/img_utils.py
rename to scripts/utils/img_utils.py