116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
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
|