From 61e93eb4cc41a9d2f808823f6b3752ad5bcd93b2 Mon Sep 17 00:00:00 2001 From: ssit Date: Tue, 29 Aug 2023 21:46:37 -0400 Subject: [PATCH] Add face orderings to new node --- nodes.py | 61 +++++++++++++++++++++++++++++++++++++++++++-- scripts/faceswap.py | 8 ++++++ scripts/swapper.py | 25 ++++++++++++++++--- 3 files changed, 88 insertions(+), 6 deletions(-) diff --git a/nodes.py b/nodes.py index 70709bd..f4e6700 100644 --- a/nodes.py +++ b/nodes.py @@ -10,6 +10,9 @@ def model_names(): return {os.path.basename(x): x for x in models} +ORDERINGS = ["left to right", "up to down", "largest to smallest"] +DEFAULT_ORDERING = ORDERINGS[0] + class roop: @classmethod def INPUT_TYPES(s): @@ -28,7 +31,7 @@ class roop: RETURN_TYPES = ("IMAGE",) FUNCTION = "execute" - CATEGORY = "image/postprocessing" + CATEGORY = "roop" def execute(self, image, reference_image, swap_model, faces_index, reference_faces_index, console_logging_level): apply_logging_patch(console_logging_level) @@ -37,9 +40,61 @@ class roop: pil_images = batch_tensor_to_pil(image) source = tensor_to_pil(reference_image) p = StableDiffusionProcessingImg2Img(pil_images) + face_order = DEFAULT_ORDERING + reverse_order = False + reference_order = DEFAULT_ORDERING + reverse_reference_order = False script.process( p=p, img=source, enable=True, faces_index=faces_index, - reference_faces_index=reference_faces_index, model=swap_model, + reference_faces_index=reference_faces_index, + face_order=face_order, reverse_order=reverse_order, + reference_order=reference_order, reverse_reference_order=reverse_reference_order, + model=swap_model, + face_restorer_name=None, face_restorer_visibility=None, + upscaler_name=None, upscaler_scale=None, upscaler_visibility=None, + swap_in_source=True, swap_in_generated=True + ) + result = batched_pil_to_tensor(p.init_images) + return (result,) + + +class RoopImproved: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "reference_image": ("IMAGE",), + "swap_model": (list(model_names().keys()),), + # Comma separated face number(s) + "faces_index": ("STRING", {"default": "0"}), + "reference_faces_index": ("STRING", {"default": "0"}), + "face_order": (ORDERINGS, {"default": DEFAULT_ORDERING}), + "reverse_order": ("BOOLEAN", {"default": False}), + "reference_order": (ORDERINGS, {"default": DEFAULT_ORDERING}), + "reverse_reference_order": ("BOOLEAN", {"default": False}), + # Allow user to change the logging amount, going from minimal to verbose + "console_logging_level": ([0, 1, 2],), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "execute" + CATEGORY = "roop" + + def execute(self, image, reference_image, swap_model, faces_index, reference_faces_index, face_order, reverse_order, reference_order, reverse_reference_order, console_logging_level): + apply_logging_patch(console_logging_level) + + script = FaceSwapScript() + pil_images = batch_tensor_to_pil(image) + source = tensor_to_pil(reference_image) + p = StableDiffusionProcessingImg2Img(pil_images) + script.process( + p=p, img=source, enable=True, faces_index=faces_index, + reference_faces_index=reference_faces_index, + face_order=face_order, reverse_order=reverse_order, + reference_order=reference_order, reverse_reference_order=reverse_reference_order, + model=swap_model, face_restorer_name=None, face_restorer_visibility=None, upscaler_name=None, upscaler_scale=None, upscaler_visibility=None, swap_in_source=True, swap_in_generated=True @@ -50,8 +105,10 @@ class roop: NODE_CLASS_MAPPINGS = { "roop": roop, + "RoopImproved": RoopImproved, } NODE_DISPLAY_NAME_MAPPINGS = { "roop": "roop", + "RoopImproved": "Roop (Improved)", } diff --git a/scripts/faceswap.py b/scripts/faceswap.py index bd639a3..a153c0c 100644 --- a/scripts/faceswap.py +++ b/scripts/faceswap.py @@ -132,6 +132,10 @@ class FaceSwapScript(scripts.Script): enable, faces_index, reference_faces_index, + face_order, + reverse_order, + reference_order, + reverse_reference_order, model, face_restorer_name, face_restorer_visibility, @@ -176,6 +180,10 @@ class FaceSwapScript(scripts.Script): reference_faces_index=self.reference_faces_index, model=self.model, upscale_options=self.upscale_options, + face_order=face_order, + reverse_order=reverse_order, + reference_order=reference_order, + reverse_reference_order=reverse_reference_order, ) p.init_images[i] = result.image() else: diff --git a/scripts/swapper.py b/scripts/swapper.py index 577940c..1db162a 100644 --- a/scripts/swapper.py +++ b/scripts/swapper.py @@ -77,7 +77,14 @@ def upscale_image(image: Image, upscale_options: UpscaleOptions): return result_image -def get_face_single(img_data: np.ndarray, face_index=0, det_size=(640, 640)): +def order_largest_smallest(faces, reverse=False): + def area(face): + x1, y1, x2, y2 = face.bbox + return (x2 - x1) * (y2 - y1) + return sorted(faces, key=lambda x: area(x), reverse=not reverse) + + +def get_face_single(img_data: np.ndarray, face_index=0, det_size=(640, 640), sorter="left to right", reverse_order=False): face_analyser = insightface.app.FaceAnalysis(name="buffalo_l", providers=providers, root=insightface_path) buffalo_path = os.path.join(insightface_models_path, "buffalo_l.zip") @@ -93,7 +100,13 @@ def get_face_single(img_data: np.ndarray, face_index=0, det_size=(640, 640)): 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] + if sorter == "left to right": + result = sorted(face, key=lambda x: x.bbox[0], reverse=reverse_order) + elif sorter == "up to down": + result = sorted(face, key=lambda x: x.bbox[1], reverse=reverse_order) + elif sorter == "largest to smallest": + result = order_largest_smallest(face, reverse=reverse_order) + return result[face_index] except IndexError: return None @@ -116,6 +129,10 @@ def swap_face( faces_index: Set[int] = {0}, reference_faces_index: Set[int] = {0}, upscale_options: Union[UpscaleOptions, None] = None, + face_order: str = "left to right", + reverse_order: bool = False, + reference_order: str = "left to right", + reverse_reference_order: bool = False, ) -> ImageResult: result_image = target_img converted = convert_to_sd(target_img) @@ -138,7 +155,7 @@ def swap_face( # source_faces = [] for face_num in reference_faces_index: - source_face = get_face_single(source_img, face_index=face_num) + source_face = get_face_single(source_img, face_index=face_num, sorter=reference_order, reverse_order=reverse_reference_order) if source_face is not None: source_faces.append(source_face) else: @@ -152,7 +169,7 @@ def swap_face( face_swapper = getFaceSwapModel(model_path) for face_num in faces_index: - target_face = get_face_single(target_img, face_index=face_num) + target_face = get_face_single(target_img, face_index=face_num, sorter=face_order, reverse_order=reverse_order) if target_face is not None: source_face = source_faces[source_face_idx]