Add face orderings to new node

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

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:

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]