From 486200afc525056c68ce80e35bacf12a2d7241c1 Mon Sep 17 00:00:00 2001 From: ssit Date: Fri, 25 Aug 2023 21:30:34 -0400 Subject: [PATCH] Add widget for indexing faces in the reference image --- nodes.py | 6 ++++-- scripts/faceswap.py | 13 +++++++++++-- scripts/swapper.py | 24 ++++++++++++++++++++++-- 3 files changed, 37 insertions(+), 6 deletions(-) diff --git a/nodes.py b/nodes.py index 6bc1454..70709bd 100644 --- a/nodes.py +++ b/nodes.py @@ -20,6 +20,7 @@ class roop: "swap_model": (list(model_names().keys()),), # Comma separated face number(s) "faces_index": ("STRING", {"default": "0"}), + "reference_faces_index": ("STRING", {"default": "0"}), # Allow user to change the logging amount, going from minimal to verbose "console_logging_level": ([0, 1, 2],), } @@ -29,7 +30,7 @@ class roop: FUNCTION = "execute" CATEGORY = "image/postprocessing" - def execute(self, image, reference_image, swap_model, faces_index, console_logging_level): + def execute(self, image, reference_image, swap_model, faces_index, reference_faces_index, console_logging_level): apply_logging_patch(console_logging_level) script = FaceSwapScript() @@ -37,7 +38,8 @@ class roop: source = tensor_to_pil(reference_image) p = StableDiffusionProcessingImg2Img(pil_images) script.process( - p=p, img=source, enable=True, faces_index=faces_index, model=swap_model, + p=p, img=source, enable=True, faces_index=faces_index, + reference_faces_index=reference_faces_index, 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 diff --git a/scripts/faceswap.py b/scripts/faceswap.py index 9f44feb..bd639a3 100644 --- a/scripts/faceswap.py +++ b/scripts/faceswap.py @@ -131,6 +131,7 @@ class FaceSwapScript(scripts.Script): img, enable, faces_index, + reference_faces_index, model, face_restorer_name, face_restorer_visibility, @@ -150,14 +151,21 @@ class FaceSwapScript(scripts.Script): self.swap_in_generated = swap_in_generated self.model = model self.faces_index = { - int(x) for x in faces_index.strip(",").split(",") if x.isnumeric() + int(y) for x in faces_index.split(",") if (y := x.strip()).isnumeric() + } + self.reference_faces_index = { + int(y) for x in reference_faces_index.split(",") if (y := x.strip()).isnumeric() } if len(self.faces_index) == 0: self.faces_index = {0} + if len(self.reference_faces_index) == 0: + self.reference_faces_index = {0} + logger.info(f"Faces index: {self.faces_index}") + logger.info(f"Reference faces index: {self.reference_faces_index}") if self.enable: if self.source is not None: if isinstance(p, StableDiffusionProcessingImg2Img) and swap_in_source: - logger.info(f"roop enabled, face index %s", self.faces_index) + logger.info(f"roop enabled, face index %s, reference face index %s", self.faces_index, self.reference_faces_index) for i in range(len(p.init_images)): logger.info(f"Swap in source %s", i) @@ -165,6 +173,7 @@ class FaceSwapScript(scripts.Script): self.source, p.init_images[i], faces_index=self.faces_index, + reference_faces_index=self.reference_faces_index, model=self.model, upscale_options=self.upscale_options, ) diff --git a/scripts/swapper.py b/scripts/swapper.py index a06ab0b..577940c 100644 --- a/scripts/swapper.py +++ b/scripts/swapper.py @@ -114,6 +114,7 @@ def swap_face( target_img: Image.Image, model: Union[str, None] = None, faces_index: Set[int] = {0}, + reference_faces_index: Set[int] = {0}, upscale_options: Union[UpscaleOptions, None] = None, ) -> ImageResult: result_image = target_img @@ -131,8 +132,21 @@ def swap_face( 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: + + # + # Get source faces + # + source_faces = [] + for face_num in reference_faces_index: + source_face = get_face_single(source_img, face_index=face_num) + if source_face is not None: + source_faces.append(source_face) + else: + logger.info(f"No source face found for {face_num}") + logger.info(f"Found {len(source_faces)} source faces") + source_face_idx = 0 + + if len(source_faces) > 0: result = target_img model_path = os.path.join(swapper_path, model) face_swapper = getFaceSwapModel(model_path) @@ -140,10 +154,16 @@ def swap_face( for face_num in faces_index: target_face = get_face_single(target_img, face_index=face_num) if target_face is not None: + + source_face = source_faces[source_face_idx] + logger.info(f"Swapping source face {source_face_idx} onto target face {face_num}") + result = face_swapper.get(result, target_face, source_face) else: logger.info(f"No target face found for {face_num}") + source_face_idx = (source_face_idx + 1) % len(source_faces) + result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) if upscale_options is not None: result_image = upscale_image(result_image, upscale_options)