diff --git a/__init__.py b/__init__.py index 82ac223..fa47ed4 100644 --- a/__init__.py +++ b/__init__.py @@ -60,12 +60,15 @@ NODE_CLASS_MAPPINGS = { "Deglaze Image (mtb)": DeglazeImage, "Smart Step (mtb)": SmartStep, "Styles Loader (mtb)": StylesLoader, + "Text to Image (mtb)": TextToImage, "Load Image Sequence (mtb)": LoadImageSequence, "Save Image Sequence (mtb)": SaveImageSequence, "Mask to Image (mtb)": MaskToImage, "Image Remove Background RemBG (mtb)": ImageRemoveBackgroundRembg, "Colored Image (mtb)": ColoredImage, "Image Premultiply (mtb)": ImagePremultiply, + "Face Swap [roop] (mtb)": Roop, + # "MMPose Estimation (mtb)": MMPoseEstimation, # "Load Geometry (mtb)": LoadGeometry, # "Geometry Info (mtb)": GeometryInfo, } diff --git a/nodes/roop.py b/nodes/roop.py new file mode 100644 index 0000000..8294202 --- /dev/null +++ b/nodes/roop.py @@ -0,0 +1,152 @@ +# region imports +from ifnude import detect +from logging import getLogger +from pathlib import Path +from PIL import Image +from typing import List, Set, Tuple +import cv2 +import folder_paths +import glob +import insightface +import numpy as np +import onnxruntime +import os +import tempfile +import torch + +from ..utils import pil2tensor, tensor2pil + +# endregion + +logger = getLogger(__name__) +providers = onnxruntime.get_available_providers() + +# region roop node +class Roop: + model = None + model_path = None + + def __init__(self) -> None: + pass + + @staticmethod + def get_models() -> List[Path]: + models_path = os.path.join(folder_paths.models_dir, "roop/*") + models = glob.glob(models_path) + models = [Path(x) for x in models if x.endswith(".onnx") or x.endswith(".pth")] + return models + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "reference": ("IMAGE",), + "faces_index": ("STRING", {"default": "0"}), + "roop_model": ([x.name for x in cls.get_models()], {"default": "None"}), + }, + "optional": { + "image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "swap" + CATEGORY = "image" + + def swap( + self, + image: torch.Tensor, + reference: torch.Tensor, + faces_index: str, + roop_model: str, + ): + image = tensor2pil(image) + reference = tensor2pil(reference) + faces_index = { + int(x) for x in faces_index.strip(",").split(",") if x.isnumeric() + } + + roop_model = self.getFaceSwapModel(roop_model) + swapped = swap_face(reference, image, roop_model, faces_index) + return (pil2tensor(swapped),) + + def getFaceSwapModel(self, model_path: str): + model_path = os.path.join(folder_paths.models_dir, "roop", model_path) + if self.model_path is None or self.model_path != model_path: + self.model_path = model_path + self.model = insightface.model_zoo.get_model( + model_path, providers=providers + ) + + return self.model + + +# endregion + +# region face swap utils +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 convert_to_sd(img) -> Tuple[bool, str]: + chunks = detect(img) + shapes = [chunk["score"] > 0.7 for chunk in chunks] + return [any(shapes), tempfile.NamedTemporaryFile(delete=False, suffix=".png")] + + +def swap_face( + source_img: Image.Image, + target_img: Image.Image, + face_swapper_model=None, + faces_index: Set[int] = None, +) -> Image.Image: + if faces_index is None: + faces_index = {0} + result_image = target_img + converted = convert_to_sd(target_img) + scale, fn = converted[0], converted[1] + if face_swapper_model is not None and not scale: + if isinstance(source_img, str): # source_img is a base64 string + import base64, io + + if ( + "base64," in source_img + ): # check if the base64 string has a data URL scheme + base64_data = source_img.split("base64,")[-1] + img_bytes = base64.b64decode(base64_data) + else: + # if no data URL scheme, just decode + img_bytes = base64.b64decode(source_img) + source_img = Image.open(io.BytesIO(img_bytes)) + source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR) + target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) + source_face = get_face_single(source_img, face_index=0) + if source_face is not None: + result = target_img + + for face_num in faces_index: + target_face = get_face_single(target_img, face_index=face_num) + if target_face is not None: + result = face_swapper_model.get(result, target_face, source_face) + else: + logger.info(f"No target face found for {face_num}") + + result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) + else: + logger.info("No source face found") + return result_image + + +# endregion face swap utils