Merge pull request #222 from jiveabillion/main

UPDATE: Mask Update, Multiple Source Files, Rearranged UI

scripts/inferencers/bisenet_mask_generator.py

@@ -4,16 +4,19 @@ import cv2
 import modules.shared as shared
 import numpy as np
 import torch
+from scripts.reactor_logger import logger
 from facexlib.parsing import init_parsing_model
 from facexlib.utils.misc import img2tensor
 from torchvision.transforms.functional import normalize
 from PIL import Image
+from scripts.inferencers.vignette_mask_generator import VignetteMaskGenerator
 from scripts.inferencers.mask_generator import MaskGenerator
-from scripts.reactor_logger import logger
+
 
 class BiSeNetMaskGenerator(MaskGenerator):
     def __init__(self) -> None:
         self.mask_model = init_parsing_model(device=shared.device)
+        self.fallback_mask_generator = VignetteMaskGenerator()
 
     def name(self):
         return "BiSeNet"
@@ -25,7 +28,7 @@ class BiSeNetMaskGenerator(MaskGenerator):
         affected_areas: List[str],
         mask_size: int,
         use_minimal_area: bool,
-        fallback_ratio: float = 0.25,
+        fallback_ratio: float = 0.10,
         **kwargs,
     ) -> np.ndarray:
         original_face_image = face_image
@@ -59,11 +62,11 @@ class BiSeNetMaskGenerator(MaskGenerator):
         if w != 512 or h != 512:
             mask = cv2.resize(mask, dsize=(w, h))
         
-        """if MaskGenerator.calculate_mask_coverage(mask) < fallback_ratio:
+        if MaskGenerator.calculate_mask_coverage(mask) < fallback_ratio:
-            logger.info("Use fallback mask generator")
+            logger.status(F"Mask coverage less than fallback ratio of {fallback_ratio}. Using vignette mask generator.")
             mask = self.fallback_mask_generator.generate_mask(
                 original_face_image, face_area_on_image, use_minimal_area=True
-            )"""
+            )
        
         return mask
 

scripts/inferencers/vignette_mask_generator.py

@@ -0,0 +1,50 @@
+from typing import Tuple
+
+import cv2
+import numpy as np
+
+from scripts.inferencers.mask_generator import MaskGenerator
+
+
+class VignetteMaskGenerator(MaskGenerator):
+    def name(self):
+        return "Vignette"
+
+    def generate_mask(
+        self,
+        face_image: np.ndarray,
+        face_area_on_image: Tuple[int, int, int, int],
+        use_minimal_area: bool,
+        sigma: float = -1,
+        keep_safe_area: bool = False,
+        **kwargs,
+    ) -> np.ndarray:
+        (left, top, right, bottom) = face_area_on_image
+        w, h = right - left, bottom - top
+        mask = np.zeros((face_image.shape[0], face_image.shape[1]), dtype=np.uint8)
+        if use_minimal_area:
+            sigma = 120 if sigma == -1 else sigma
+            mask[top : top + h, left : left + w] = 255
+        else:
+            sigma = 180 if sigma == -1 else sigma
+            h, w = face_image.shape[0], face_image.shape[1]
+            mask[:, :] = 255
+
+        Y = np.linspace(0, h, h, endpoint=False)
+        X = np.linspace(0, w, w, endpoint=False)
+        Y, X = np.meshgrid(Y, X)
+        Y -= h / 2
+        X -= w / 2
+
+        gaussian = np.exp(-(X**2 + Y**2) / (2 * sigma**2))
+        gaussian_mask = np.uint8(255 * gaussian.T)
+        if use_minimal_area:
+            mask[top : top + h, left : left + w] = gaussian_mask
+        else:
+            mask[:, :] = gaussian_mask
+
+        if keep_safe_area:
+            mask = cv2.ellipse(mask, ((left + right) // 2, (top + bottom) // 2), (w // 2, h // 2), 0, 0, 360, 255, -1)
+
+        mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
+        return mask

scripts/reactor_faceswap.py

@@ -1,5 +1,6 @@
 import os, glob
 import gradio as gr
+import tempfile
 from PIL import Image
 try:
     import torch.cuda as cuda
@@ -8,7 +9,7 @@ except:
     EP_is_visible = False
 
 from typing import List
-
+from PIL import Image
 import modules.scripts as scripts
 from modules.upscaler import Upscaler, UpscalerData
 from modules import scripts, shared, images, scripts_postprocessing
@@ -28,13 +29,14 @@ except:
         model_path = os.path.abspath("models")
 
 from scripts.reactor_logger import logger
-from scripts.reactor_swapper import EnhancementOptions, swap_face, check_process_halt, reset_messaged
+from scripts.reactor_swapper import EnhancementOptions,MaskOptions,MaskOption, swap_face, check_process_halt, reset_messaged
 from scripts.reactor_version import version_flag, app_title
 from scripts.console_log_patch import apply_logging_patch
 from scripts.reactor_helpers import make_grid, get_image_path, set_Device
 from scripts.reactor_globals import DEVICE, DEVICE_LIST
 
 
+
 MODELS_PATH = None
 
 def get_models():
@@ -61,12 +63,17 @@ class FaceSwapScript(scripts.Script):
 
     def ui(self, is_img2img):
         with gr.Accordion(f"{app_title}", open=False):
+            enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}")
+            gr.Markdown("<br>")
             with gr.Tab("Main"):
+
                 with gr.Column():
-                    img = gr.Image(type="pil")
+                    with gr.Tab("Single Source Image"):
-                    enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}")
+                        img = gr.Image(type="pil")
+                    with gr.Tab("Multiple Source Images"):
+                        face_files = gr.File(label="Multiple Source Face Files",file_count="multiple",file_types=["image"],info="Upload multiple face files and each file will be processed in post processing")
                     save_original = gr.Checkbox(False, label="Save Original", info="Save the original image(s) made before swapping; If you use \"img2img\" - this option will affect with \"Swap in generated\" only")
-                    mask_face = gr.Checkbox(False, label="Mask Faces", info="Attempt to mask only the faces and eliminate pixelation of the image around the contours.")
+                    mask_face = gr.Checkbox(False, label="Mask Faces", info="Attempt to mask only the faces and eliminate pixelation of the image around the contours. Additional settings in the Masking tab.")
                     
                     gr.Markdown("<br>")
                     gr.Markdown("Source Image (above):")
@@ -140,6 +147,26 @@ class FaceSwapScript(scripts.Script):
                     upscaler_visibility = gr.Slider(
                         0, 1, 1, step=0.1, label="Upscaler Visibility (if scale = 1)"
                     )
+            with gr.Tab("Masking"):
+                save_face_mask = gr.Checkbox(False, label="Save Face Mask", info="Save the face mask as a separate image with alpha transparency.")
+                use_minimal_area = gr.Checkbox(MaskOption.DEFAULT_USE_MINIMAL_AREA, label="Use Minimal Area", info="Use the least amount of area for the mask as possible. This is good for multiple faces that are close together or for preserving the most of the surrounding image.")
+                
+                mask_areas = gr.CheckboxGroup(
+                    label="Mask areas", choices=["Face", "Hair", "Hat", "Neck"], type="value", value= MaskOption.DEFAULT_FACE_AREAS
+                )
+                face_size = gr.Radio(
+                    label = "Face Size", choices = [512,256,128],value=MaskOption.DEFAULT_FACE_SIZE,type="value", info="Size of the masked area. Use larger numbers if the face is expected to be large, smaller if small. Default is 512."
+                )
+                mask_blur = gr.Slider(label="Mask blur", minimum=0, maximum=64, step=1, value=12,info="The number of pixels from the outer edge of the mask to blur.")
+
+                mask_vignette_fallback_threshold = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.01,
+                    value=MaskOption.DEFAULT_VIGNETTE_THRESHOLD,
+                    label="Vignette fallback threshold",
+                    info="Switch to a rectangular vignette mask when masked area is only this specified percentage of Face Size."
+                )
             with gr.Tab("Settings"):
                 models = get_models()
                 with gr.Row(visible=EP_is_visible):
@@ -213,7 +240,14 @@ class FaceSwapScript(scripts.Script):
             source_hash_check,
             target_hash_check,
             device,
-            mask_face
+            mask_face,
+            save_face_mask,
+            mask_areas,
+            mask_blur,
+            use_minimal_area,
+            face_size,
+            mask_vignette_fallback_threshold,
+            face_files
         ]
 
 
@@ -242,6 +276,16 @@ class FaceSwapScript(scripts.Script):
             restorer_visibility=self.face_restorer_visibility,
             codeformer_weight=self.codeformer_weight,
         )
+    @property
+    def mask_options(self) -> MaskOptions:
+        return MaskOptions(
+            mask_areas = self.mask_areas,
+            save_face_mask = self.save_face_mask,
+            mask_blur = self.mask_blur,
+            face_size = self.mask_face_size,
+            vignette_fallback_threshold = self.mask_vignette_fallback_threshold,
+            use_minimal_area = self.mask_use_minimal_area,
+        )
     
     def process(
         self,
@@ -267,7 +311,14 @@ class FaceSwapScript(scripts.Script):
         source_hash_check,
         target_hash_check,
         device,
-        mask_face
+        mask_face,
+        save_face_mask:bool,
+        mask_areas,
+        mask_blur:int,
+        mask_use_minimal_area,
+        mask_face_size,
+        mask_vignette_fallback_threshold,
+        face_files
     ):
         self.enable = enable
         if self.enable:
@@ -277,6 +328,7 @@ class FaceSwapScript(scripts.Script):
                 return
             
             global MODELS_PATH
+            
             self.source = img
             self.face_restorer_name = face_restorer_name
             self.upscaler_scale = upscaler_scale
@@ -296,6 +348,13 @@ class FaceSwapScript(scripts.Script):
             self.target_hash_check = target_hash_check
             self.device = device
             self.mask_face = mask_face
+            self.save_face_mask = save_face_mask
+            self.mask_blur = mask_blur
+            self.mask_areas = mask_areas
+            self.mask_face_size = mask_face_size
+            self.mask_vignette_fallback_threshold = mask_vignette_fallback_threshold
+            self.mask_use_minimal_area = mask_use_minimal_area
+            self.face_files = face_files
             if self.gender_source is None or self.gender_source == "No":
                 self.gender_source = 0
             if self.gender_target is None or self.gender_target == "No":
@@ -318,9 +377,10 @@ class FaceSwapScript(scripts.Script):
                 self.target_hash_check = False
 
             set_Device(self.device)
-            
+            apply_logging_patch(console_logging_level)
             if self.source is not None:
-                apply_logging_patch(console_logging_level)
+                
+               
                 if isinstance(p, StableDiffusionProcessingImg2Img) and self.swap_in_source:
                     logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
 
@@ -339,7 +399,8 @@ class FaceSwapScript(scripts.Script):
                             source_hash_check=self.source_hash_check,
                             target_hash_check=self.target_hash_check,
                             device=self.device,
-                            mask_face=mask_face
+                            mask_face=mask_face,
+                            mask_options=self.mask_options
                         )
                         p.init_images[i] = result
                         # result_path = get_image_path(p.init_images[i], p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], "txt", p=p, suffix="-swapped")
@@ -350,7 +411,7 @@ class FaceSwapScript(scripts.Script):
                         if shared.state.interrupted or shared.state.skipped:
                             return
             
-            else:
+            elif self.face_files is None or len(self.face_files) == 0:
                 logger.error("Please provide a source face")
 
     def postprocess(self, p: StableDiffusionProcessing, processed: Processed, *args):
@@ -359,27 +420,29 @@ class FaceSwapScript(scripts.Script):
             reset_messaged()
             if check_process_halt():
                 return
+            postprocess_run: bool = True
+
+            orig_images : List[Image.Image] = processed.images[processed.index_of_first_image:]
+            orig_infotexts : List[str] = processed.infotexts[processed.index_of_first_image:]
+            result_images:List[Image.Image] = []
+            
             
             if self.save_original:
 
-                postprocess_run: bool = True
-
-                orig_images : List[Image.Image] = processed.images[processed.index_of_first_image:]
-                orig_infotexts : List[str] = processed.infotexts[processed.index_of_first_image:]
-
                 result_images: List = processed.images
                 # result_info: List = processed.infotexts
 
                 if self.swap_in_generated:
                     logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
                     if self.source is not None:
+                       
                         for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)):
                             if check_process_halt():
                                 postprocess_run = False
                                 break
                             if len(orig_images) > 1:
                                 logger.status("Swap in %s", i)
-                            result, output, swapped = swap_face(
+                            result, output, swapped, masked_faces = swap_face(
                                 self.source,
                                 img,
                                 source_faces_index=self.source_faces_index,
@@ -391,7 +454,8 @@ class FaceSwapScript(scripts.Script):
                                 source_hash_check=self.source_hash_check,
                                 target_hash_check=self.target_hash_check,
                                 device=self.device,
-                                mask_face=self.mask_face
+                                mask_face=self.mask_face,
+                                mask_options=self.mask_options
                             )
                             if result is not None and swapped > 0:
                                 result_images.append(result)
@@ -400,6 +464,14 @@ class FaceSwapScript(scripts.Script):
                                     img_path = save_image(result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
                                 except:
                                     logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)")
+                                if self.mask_face and self.save_face_mask and masked_faces is not None:
+                                    result_images.append(masked_faces)
+                                    suffix = "-mask"
+                                    try:
+                                        img_path = save_image(masked_faces, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
+                                    except:
+                                        logger.error("Cannot save a Masked Face image - please, check SD WebUI Settings (Saving and Paths)")
+                                  
                             elif result is None:
                                 logger.error("Cannot create a result image")
                             
@@ -408,7 +480,48 @@ class FaceSwapScript(scripts.Script):
                             #     fullfn = split_fullfn[0] + ".txt"
                             #     with open(fullfn, 'w', encoding="utf8") as f:
                             #         f.writelines(output)
+                    if self.face_files is not None and len(self.face_files) > 0:
                        
+                        for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)):
+                            for j,f_img in enumerate(self.face_files):
+                                
+                                if check_process_halt():
+                                    postprocess_run = False
+                                    break
+                                if len(self.face_files) > 1:
+                                    logger.status("Swap in face file #%s", j+1)
+                                result, output, swapped, masked_faces = swap_face(
+                                    Image.open(os.path.abspath(f_img.name)),
+                                    img,
+                                    source_faces_index=self.source_faces_index,
+                                    faces_index=self.faces_index,
+                                    model=self.model,
+                                    enhancement_options=self.enhancement_options,
+                                    gender_source=self.gender_source,
+                                    gender_target=self.gender_target,
+                                    source_hash_check=self.source_hash_check,
+                                    target_hash_check=self.target_hash_check,
+                                    device=self.device,
+                                    mask_face=self.mask_face,
+                                    mask_options=self.mask_options
+                                )
+                                if result is not None and swapped > 0:
+                                    result_images.append(result)
+                                    suffix = f"-swapped-ff-{j+1}"
+                                    try:
+                                        img_path = save_image(result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
+                                    except:
+                                        logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)")
+                                    if self.mask_face and self.save_face_mask and masked_faces is not None:
+                                        result_images.append(masked_faces)
+                                        suffix = f"-mask-ff-{j+1}"
+                                        try:
+                                            img_path = save_image(masked_faces, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
+                                        except:
+                                            logger.error("Cannot save a Masked Face image - please, check SD WebUI Settings (Saving and Paths)")
+                                    
+                                elif result is None:
+                                    logger.error("Cannot create a result image")
                 if shared.opts.return_grid and len(result_images) > 2 and postprocess_run:
                     grid = make_grid(result_images)
                     result_images.insert(0, grid)
@@ -419,13 +532,62 @@ class FaceSwapScript(scripts.Script):
                 
                 processed.images = result_images
                 # processed.infotexts = result_info
+            elif self.face_files is not None and len(self.face_files) > 0:
+                for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)):
+                    for j,f_img in enumerate(self.face_files):
                         
+                        if check_process_halt():
+                            postprocess_run = False
+                            break
+                        if len(self.face_files) > 1:
+                            logger.status("Swap in face file #%s", j+1)
+                        result, output, swapped, masked_faces = swap_face(
+                            Image.open(os.path.abspath(f_img.name)),
+                            img,
+                            source_faces_index=self.source_faces_index,
+                            faces_index=self.faces_index,
+                            model=self.model,
+                            enhancement_options=self.enhancement_options,
+                            gender_source=self.gender_source,
+                            gender_target=self.gender_target,
+                            source_hash_check=self.source_hash_check,
+                            target_hash_check=self.target_hash_check,
+                            device=self.device,
+                            mask_face=self.mask_face,
+                            mask_options=self.mask_options
+                        )
+                        if result is not None and swapped > 0:
+                            result_images.append(result)
+                            suffix = f"-swapped-ff-{j+1}"
+                            try:
+                                img_path = save_image(result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
+                            except:
+                                logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)")
+                            if self.mask_face and self.save_face_mask and masked_faces is not None:
+                                result_images.append(masked_faces)
+                                suffix = f"-mask-ff-{j+1}"
+                                try:
+                                    img_path = save_image(masked_faces, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
+                                except:
+                                    logger.error("Cannot save a Masked Face image - please, check SD WebUI Settings (Saving and Paths)")
+                            
+                        elif result is None:
+                            logger.error("Cannot create a result image")
+                if shared.opts.return_grid and len(result_images) > 2 and postprocess_run:
+                    grid = make_grid(result_images)
+                    result_images.insert(0, grid)
+                    try:
+                        save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], shared.opts.grid_format, info=info, short_filename=not shared.opts.grid_extended_filename, p=p, grid=True)
+                    except:
+                        logger.error("Cannot save a grid - please, check SD WebUI Settings (Saving and Paths)")
+                
+                processed.images = result_images
     def postprocess_batch(self, p, *args, **kwargs):
         if self.enable and not self.save_original:
             images = kwargs["images"]
 
     def postprocess_image(self, p, script_pp: scripts.PostprocessImageArgs, *args):
-        if self.enable and self.swap_in_generated and not self.save_original:
+        if self.enable and self.swap_in_generated and not ( self.save_original or ( self.face_files is not None and len(self.face_files) > 0)):
 
             current_job_number = shared.state.job_no + 1
             job_count = shared.state.job_count
@@ -437,7 +599,7 @@ class FaceSwapScript(scripts.Script):
             if self.source is not None:
                 logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
                 image: Image.Image = script_pp.image
-                result, output, swapped = swap_face(
+                result, output, swapped, masked_faces = swap_face(
                     self.source,
                     image,
                     source_faces_index=self.source_faces_index,
@@ -449,7 +611,8 @@ class FaceSwapScript(scripts.Script):
                     source_hash_check=self.source_hash_check,
                     target_hash_check=self.target_hash_check,
                     device=self.device,
-                    mask_face=self.mask_face
+                    mask_face=self.mask_face,
+                    mask_options=self.mask_options
                 )
                 try:
                     pp = scripts_postprocessing.PostprocessedImage(result)
@@ -475,9 +638,9 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
             with gr.Tab("Main"):
                 with gr.Column():
                     img = gr.Image(type="pil")
+                    #face_files = gr.File(file_count="multiple",file_types=["image"],info="Upload multiple face files and each file will be processed in post processing")
                     enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}")
-                    mask_face = gr.Checkbox(False, label="Mask Faces", info="Attempt to mask only the faces and eliminate pixelation of the image around the contours.")
+                    mask_face = gr.Checkbox(False, label="Mask Faces", info="Attempt to mask only the faces and eliminate pixelation of the image around the contours. Additional settings in the Masking tab.")
-                    
                     gr.Markdown("Source Image (above):")
                     with gr.Row():
                         source_faces_index = gr.Textbox(
@@ -536,6 +699,27 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
                     upscaler_visibility = gr.Slider(
                         0, 1, 1, step=0.1, label="Upscaler Visibility (if scale = 1)"
                     )
+            with gr.Tab("Masking"):
+                #save_face_mask = gr.Checkbox(False, label="Save Face Mask", info="Save the face mask as a separate image with alpha transparency.")
+                use_minimal_area = gr.Checkbox(MaskOption.DEFAULT_USE_MINIMAL_AREA, label="Use Minimal Area", info="Use the least amount of area for the mask as possible. This is good for multiple faces that are close together or for preserving the most of the surrounding image.")
+                
+                mask_areas = gr.CheckboxGroup(
+                    label="Mask areas", choices=["Face", "Hair", "Hat", "Neck"], type="value", value= MaskOption.DEFAULT_FACE_AREAS
+                )
+                face_size = gr.Radio(
+                    label = "Face Size", choices = [512,256,128],value=MaskOption.DEFAULT_FACE_SIZE,type="value", info="Size of the masked area. Use larger numbers if the face is expected to be large, smaller if small. Default is 512."
+                )
+                mask_blur = gr.Slider(label="Mask blur", minimum=0, maximum=64, step=1, value=12,info="The number of pixels from the outer edge of the mask to blur.")
+
+                mask_vignette_fallback_threshold = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.01,
+                    value=MaskOption.DEFAULT_VIGNETTE_THRESHOLD,
+                    label="Vignette fallback threshold",
+                    info="Switch to a rectangular vignette mask when masked area is only this specified percentage of Face Size."
+                )
+                
             with gr.Tab("Settings"):
                 models = get_models()
                 with gr.Row(visible=EP_is_visible):
@@ -592,7 +776,14 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
             'gender_target': gender_target,
             'codeformer_weight': codeformer_weight,
             'device': device,
-            'mask_face':mask_face
+            'mask_face':mask_face,
+           
+            'mask_areas':mask_areas,
+            'mask_blur':mask_blur,
+            'mask_vignette_fallback_threshold':mask_vignette_fallback_threshold,
+            'face_size':face_size,
+            'use_minimal_area':use_minimal_area,
+           
         }
         return args
 
@@ -621,14 +812,23 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
             restorer_visibility=self.face_restorer_visibility,
             codeformer_weight=self.codeformer_weight,
         )
-
+    @property
+    def mask_options(self) -> MaskOptions:
+        return MaskOptions(
+            mask_areas = self.mask_areas,
+            save_face_mask = self.save_face_mask,
+            mask_blur = self.mask_blur,
+            face_size = self.face_size,
+            vignette_fallback_threshold = self.mask_vignette_fallback_threshold,
+            use_minimal_area = self.use_minimal_area,
+        )
     def process(self, pp: scripts_postprocessing.PostprocessedImage, **args):
         if args['enable']:
             reset_messaged()
             if check_process_halt():
                 return  
-
             global MODELS_PATH
+            
             self.source = args['img']
             self.face_restorer_name = args['face_restorer_name']
             self.upscaler_scale = args['upscaler_scale']
@@ -643,6 +843,13 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
             self.codeformer_weight = args['codeformer_weight']
             self.device = args['device']
             self.mask_face = args['mask_face']
+            self.save_face_mask = None
+            self.mask_areas= args['mask_areas']
+            self.mask_blur= args['mask_blur']
+            self.mask_vignette_fallback_threshold= args['mask_vignette_fallback_threshold']
+            self.face_size= args['face_size']
+            self.use_minimal_area= args['use_minimal_area']
+            self.face_files = None
             if self.gender_source is None or self.gender_source == "No":
                 self.gender_source = 0
             if self.gender_target is None or self.gender_target == "No":
@@ -664,12 +871,12 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
                 reset_messaged()
 
             set_Device(self.device)
-            
+            apply_logging_patch(self.console_logging_level)
             if self.source is not None:
-                apply_logging_patch(self.console_logging_level)
+                
                 logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
                 image: Image.Image = pp.image
-                result, output, swapped = swap_face(
+                result, output, swapped,masked_faces = swap_face(
                     self.source,
                     image,
                     source_faces_index=self.source_faces_index,
@@ -681,7 +888,8 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
                     source_hash_check=True,
                     target_hash_check=True,
                     device=self.device,
-                    mask_face=self.mask_face
+                    mask_face=self.mask_face,
+                    mask_options=self.mask_options
                 )
                 try:
                     pp.info["ReActor"] = True

scripts/reactor_swapper.py

@@ -1,7 +1,7 @@
 import copy
 import os
 from dataclasses import dataclass
-from typing import List, Union
+from typing import List, Tuple, Union
 
 import cv2
 import numpy as np
@@ -41,7 +41,12 @@ if DEVICE == "CUDA":
     PROVIDERS = ["CUDAExecutionProvider"]
 else:
     PROVIDERS = ["CPUExecutionProvider"]
-
+class MaskOption:
+    DEFAULT_FACE_AREAS = ["Face"]
+    DEFAULT_FACE_SIZE:int = 512
+    DEFAULT_VIGNETTE_THRESHOLD:float = 0.1
+    DEFAULT_MASK_BLUR:int = 12,
+    DEFAULT_USE_MINIMAL_AREA:bool = True
 
 @dataclass
 class EnhancementOptions:
@@ -54,6 +59,15 @@ class EnhancementOptions:
     codeformer_weight: float = 0.5
     
 
+@dataclass
+class MaskOptions:
+    mask_areas:List[str]
+    save_face_mask: bool = False
+    mask_blur:int = 12   
+    face_size:int  = 512
+    vignette_fallback_threshold:float =0.10
+    use_minimal_area:bool = True
+
 MESSAGED_STOPPED = False
 MESSAGED_SKIPPED = False
 
@@ -175,7 +189,7 @@ def enhance_image(image: Image, enhancement_options: EnhancementOptions):
         result_image = restore_face(result_image, enhancement_options)
 
     return result_image
-def enhance_image_and_mask(image: Image.Image, enhancement_options: EnhancementOptions,target_img_orig:Image.Image,entire_mask_image:Image.Image)->Image.Image:
+def enhance_image_and_mask(image: Image.Image, enhancement_options: EnhancementOptions,target_img_orig:Image.Image,entire_mask_image:Image.Image)->Tuple[Image.Image,Image.Image]:
     result_image = image
     
     if check_process_halt(msgforced=True):
@@ -183,7 +197,11 @@ def enhance_image_and_mask(image: Image.Image, enhancement_options: EnhancementO
     
     if enhancement_options.do_restore_first:
         
+      
         result_image = restore_face(result_image, enhancement_options)
+
+        transparent = Image.new("RGBA",result_image.size)
+        masked_faces = Image.composite(result_image.convert("RGBA"),transparent,entire_mask_image)
         result_image = Image.composite(result_image,target_img_orig,entire_mask_image)
         result_image = upscale_image(result_image, enhancement_options)
 
@@ -191,10 +209,13 @@ def enhance_image_and_mask(image: Image.Image, enhancement_options: EnhancementO
 
         result_image = upscale_image(result_image, enhancement_options)
         entire_mask_image = Image.fromarray(cv2.resize(np.array(entire_mask_image),result_image.size, interpolation=cv2.INTER_AREA)).convert("L")
+       
         result_image = Image.composite(result_image,target_img_orig,entire_mask_image)
         result_image = restore_face(result_image, enhancement_options)
+        transparent = Image.new("RGBA",result_image.size)
+        masked_faces = Image.composite(result_image.convert("RGBA"),transparent,entire_mask_image)
+    return result_image, masked_faces
 
-    return result_image
 
     
 def get_gender(face, face_index):
@@ -310,15 +331,18 @@ def swap_face(
     source_hash_check: bool = True,
     target_hash_check: bool = False,
     device: str = "CPU",
-    mask_face:bool = False
+    mask_face:bool = False,
+    mask_options:Union[MaskOptions, None]= None
 ):
     global SOURCE_FACES, SOURCE_IMAGE_HASH, TARGET_FACES, TARGET_IMAGE_HASH, PROVIDERS
     result_image = target_img
-
+    masked_faces = None
     PROVIDERS = ["CUDAExecutionProvider"] if device == "CUDA" else ["CPUExecutionProvider"]
     
     if check_process_halt():
         return result_image, [], 0
+    if mask_options is None:
+        mask_options = MaskOptions()
         
     if model is not None:
 
@@ -444,7 +468,7 @@ def swap_face(
                             swapped_image = face_swapper.get(result, target_face, source_face)
                                                     
                             if mask_face:
-                                result = apply_face_mask(swapped_image=swapped_image,target_image=result,target_face=target_face,entire_mask_image=entire_mask_image)
+                                result = apply_face_mask(swapped_image=swapped_image,target_image=result,target_face=target_face,entire_mask_image=entire_mask_image,mask_options=mask_options)
                             else:
                                 result = swapped_image
                             swapped += 1
@@ -480,8 +504,8 @@ def swap_face(
                 result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
                 
                 if enhancement_options is not None and swapped > 0:
-                    if mask_face and entire_mask_image is not None:
+                    if  mask_face and entire_mask_image is not None:
-                        result_image = enhance_image_and_mask(result_image, enhancement_options,Image.fromarray(target_img_orig),Image.fromarray(entire_mask_image).convert("L"))    
+                       result_image, masked_faces = enhance_image_and_mask(result_image, enhancement_options,Image.fromarray(target_img_orig),Image.fromarray(entire_mask_image).convert("L"))    
                     else:    
                         result_image = enhance_image(result_image, enhancement_options)
                 elif mask_face and entire_mask_image is not None and swapped > 0:
@@ -492,20 +516,19 @@ def swap_face(
         else:
             logger.status("No source face(s) found")
     
-    return result_image, output, swapped
+    return result_image, output, swapped,masked_faces
 
 
 
-def apply_face_mask(swapped_image:np.ndarray,target_image:np.ndarray,target_face,entire_mask_image:np.array)->np.ndarray:
+def apply_face_mask(swapped_image:np.ndarray,target_image:np.ndarray,target_face,entire_mask_image:np.array,mask_options:Union[MaskOptions,None] = None)->np.ndarray:
     logger.status("Masking Face") 
     mask_generator =  BiSeNetMaskGenerator()
-    face = Face(target_image,Rect.from_ndarray(np.array(target_face.bbox)),1.6,512,"")
+    face = Face(target_image,Rect.from_ndarray(np.array(target_face.bbox)),1.6,mask_options.face_size,"")
     face_image = np.array(face.image)
-    process_face_image(face)
     face_area_on_image = face.face_area_on_image
-    mask = mask_generator.generate_mask(face_image,face_area_on_image=face_area_on_image,affected_areas=["Face"],mask_size=0,use_minimal_area=True)
+   
-    mask = cv2.blur(mask, (12, 12))
+    mask = mask_generator.generate_mask(face_image,face_area_on_image=face_area_on_image,affected_areas=mask_options.mask_areas,mask_size=0,use_minimal_area=mask_options.use_minimal_area)
-    """entire_mask_image = np.zeros_like(target_image)"""
+    mask = cv2.blur(mask, (mask_options.mask_blur, mask_options.mask_blur))
     larger_mask = cv2.resize(mask, dsize=(face.width, face.height))
     entire_mask_image[
         face.top : face.bottom,
@@ -554,22 +577,7 @@ def color_generator(colors):
 
 
 color_iter = color_generator(colors)
-def process_face_image(
-        face: Face,
-        **kwargs,
-    ) -> Image:
-        image = np.array(face.image)
-        overlay = image.copy()
-        cv2.rectangle(overlay, (0, 0), (image.shape[1], image.shape[0]), next(color_iter), -1)
-        l, t, r, b = face.face_area_on_image
-        cv2.rectangle(overlay, (l, t), (r, b), (0, 0, 0), 10)
-        if face.landmarks_on_image is not None:
-            for landmark in face.landmarks_on_image:
-                cv2.circle(overlay, (int(landmark.x), int(landmark.y)), 6, (0, 0, 0), 10)
-        alpha = 0.3
-        output = cv2.addWeighted(image, 1 - alpha, overlay, alpha, 0)
 
-        return Image.fromarray(output)
 def dilate_erode(img: Image.Image, value: int) -> Image.Image:
     """
     The dilate_erode function takes an image and a value.