initial version

FaceSwapNode.py

@@ -0,0 +1,115 @@
+import insightface
+import onnxruntime
+import torch
+import glob
+import tempfile
+import numpy as np
+import cv2
+import os
+from PIL import Image
+from typing import List, Union, Dict, Set, Tuple
+import folder_paths
+import torchvision.transforms as T
+from comfy import model_management
+
+providers = ["CPUExecutionProvider"]
+model_path = folder_paths.models_dir
+onnx_path = os.path.join(model_path, "roop")
+FS_MODEL = None
+CURRENT_FS_MODEL_PATH = None
+device = model_management.get_torch_device()
+
+
+def get_models():
+    models_path = os.path.join(onnx_path + os.path.sep + "*")
+    models = glob.glob(models_path)
+    models = [x for x in models if x.endswith(".onnx") or x.endswith(".pth")]
+    return models
+
+
+def convert_to_sd(img):
+    return [False, tempfile.NamedTemporaryFile(delete=False, suffix=".png")]
+
+
+class FaceSwapNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"face": ("IMAGE",),
+                             "image": ("IMAGE",),
+                             "source_face_index": ("INT", {"default": 0, "min": 0, "step": 1}),
+                             "target_face_indices": ("STRING", {"multiline": False}),
+                             }}
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "swap"
+
+    CATEGORY = "image/faceswap"
+
+    def swap(self, face: torch.Tensor, image: torch.Tensor, source_face_index=0, target_face_indices="0"):
+        models = get_models()
+
+        target_faces = {int(x) for x in target_face_indices.strip(",").split(",") if x.isnumeric()}
+        result = swap_face(face, image, models[0], source_face_index, target_faces)
+
+        result_tensor = np.array(result).astype(np.float32) / 255.0
+        result_tensor = torch.from_numpy(result_tensor)[None,]
+
+        return (result_tensor,)
+
+
+def getFaceSwapModel(model_path: str):
+    global FS_MODEL
+    global CURRENT_FS_MODEL_PATH
+    if CURRENT_FS_MODEL_PATH is None or CURRENT_FS_MODEL_PATH != model_path:
+        CURRENT_FS_MODEL_PATH = model_path
+        FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers)
+    return FS_MODEL
+
+
+def get_face_single(img_data: np.ndarray, face_index=0, det_size=(640, 640)):
+    face_analyser = insightface.app.FaceAnalysis(name="buffalo_l", providers=providers)
+    face_analyser.prepare(ctx_id=0, det_size=det_size)
+
+    face = face_analyser.get(img_data)
+
+    if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320:
+        det_size_half = (det_size[0] // 2, det_size[1] // 2)
+        return get_face_single(img_data, face_index=face_index, det_size=det_size_half)
+
+    try:
+        return sorted(face, key=lambda x: x.bbox[0])[face_index]
+    except IndexError:
+        return None
+
+
+def swap_face(
+        source_img: torch.Tensor,
+        target_img: torch.Tensor,
+        model: Union[str, None] = None,
+        source_face: [int] = 0,
+        target_face_list: Set[int] = {0},
+) -> Image.Image:
+    result_image = target_img
+    converted = convert_to_sd(target_img)
+    scale, fn = converted[0], converted[1]
+    if model is not None and not scale:
+
+        source_img = (source_img[0].detach().numpy() * 255).astype(np.uint8)
+        target_img = (target_img[0].detach().numpy() * 255).astype(np.uint8)
+
+        source_img = cv2.cvtColor(source_img, cv2.COLOR_RGB2BGR)
+        target_img = cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR)
+        source_face = get_face_single(source_img, face_index=source_face)
+
+        if source_face is not None:
+            result = target_img
+            model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model)
+            face_swapper = getFaceSwapModel(model_path)
+
+            for face_num in target_face_list:
+                target_face = get_face_single(target_img, face_index=face_num)
+                if target_face is not None:
+                    result = face_swapper.get(result, target_face, source_face)
+                    result_image = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
+
+    return result_image

__init__.py

@@ -0,0 +1,8 @@
+from .FaceSwapNode import FaceSwapNode
+from .install import install
+
+NODE_CLASS_MAPPINGS = {
+    "FaceSwapNode": FaceSwapNode,
+}
+
+install()

install.py

@@ -0,0 +1,42 @@
+import os
+import sys
+import subprocess
+
+comfy_path = '../..'
+if sys.argv[0] == 'install.py':
+    sys.path.append('.')   # for portable version
+
+impact_path = os.path.join(os.path.dirname(__file__), "modules")
+
+sys.path.append(impact_path)
+sys.path.append(comfy_path)
+
+import platform
+import folder_paths
+from torchvision.datasets.utils import download_url
+
+print("### ComfyUI-FaceSwapper: Check dependencies")
+
+if "python_embeded" in sys.executable or "python_embedded" in sys.executable:
+    pip_install = [sys.executable, '-s', '-m', 'pip', 'install']
+else:
+    pip_install = [sys.executable, '-m', 'pip', 'install']
+
+def ensure_pip_packages():
+    try:
+        import cython
+    except Exception:
+        my_path = os.path.dirname(__file__)
+        requirements_path = os.path.join(my_path, "requirements.txt")
+        subprocess.check_call(pip_install + ['-r', requirements_path])
+
+def install():
+    ensure_pip_packages()
+    # Download model
+    print("### ComfyUI-Impact-Pack: Check basic models")
+    model_path = folder_paths.models_dir
+    onnx_path = os.path.join(model_path, "roop")
+
+    if not os.path.exists(onnx_path):
+        download_url("https://huggingface.co/henryruhs/roop/resolve/main/inswapper_128.onnx", onnx_path)
+

requirements.txt

@@ -0,0 +1,5 @@
+insightface==0.7.3
+onnx==1.14.0
+onnxruntime==1.15.0
+opencv-python==4.7.0.72
+cython