Add widget for indexing faces in the reference image

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

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,
                         )

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)