diff --git a/FaceSwapNode.py b/FaceSwapNode.py new file mode 100644 index 0000000..73555b6 --- /dev/null +++ b/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 diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e730e3b --- /dev/null +++ b/__init__.py @@ -0,0 +1,8 @@ +from .FaceSwapNode import FaceSwapNode +from .install import install + +NODE_CLASS_MAPPINGS = { + "FaceSwapNode": FaceSwapNode, +} + +install() \ No newline at end of file diff --git a/install.py b/install.py new file mode 100644 index 0000000..a923e6c --- /dev/null +++ b/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) + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..3d0c6cc --- /dev/null +++ b/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