add pp&mask options for each faces. Improve API. Requires more testing

scripts/faceswaplab.py

@@ -72,14 +72,6 @@ class FaceSwapScript(scripts.Script):
     def units_count(self) -> int:
         return opts.data.get("faceswaplab_units_count", 3)
 
-    @property
-    def upscaled_swapper_in_generated(self) -> bool:
-        return opts.data.get("faceswaplab_upscaled_swapper", False)
-
-    @property
-    def upscaled_swapper_in_source(self) -> bool:
-        return opts.data.get("faceswaplab_upscaled_swapper_in_source", False)
-
     @property
     def enabled(self) -> bool:
         """Return True if any unit is enabled and the state is not interupted"""
@@ -152,7 +144,6 @@ class FaceSwapScript(scripts.Script):
                         get_current_model(),
                         self.swap_in_source_units,
                         images=init_images,
-                        upscaled_swapper=self.upscaled_swapper_in_source,
                         force_blend=True,
                     )
                     logger.info(f"processed init images: {len(init_images)}")
@@ -187,7 +178,6 @@ class FaceSwapScript(scripts.Script):
                             get_current_model(),
                             self.swap_in_generated_units,
                             images=[(img, info)],
-                            upscaled_swapper=self.upscaled_swapper_in_generated,
                         )
                         if swapped_images is None:
                             continue

scripts/faceswaplab_api/faceswaplab_api.py

@@ -91,6 +91,8 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None:
         if src_image is not None:
             if request.postprocessing:
                 pp_options = PostProcessingOptions.from_api_dto(request.postprocessing)
+            else:
+                pp_options = None
             units = get_faceswap_units_settings(request.units)
 
             swapped_images = swapper.batch_process(

scripts/faceswaplab_settings/faceswaplab_settings.py

@@ -54,7 +54,7 @@ def on_ui_settings() -> None:
         "faceswaplab_pp_default_face_restorer",
         shared.OptionInfo(
             None,
-            "UI Default post processing face restorer (requires restart)",
+            "UI Default global post processing face restorer (requires restart)",
             gr.Dropdown,
             {
                 "interactive": True,
@@ -67,7 +67,7 @@ def on_ui_settings() -> None:
         "faceswaplab_pp_default_face_restorer_visibility",
         shared.OptionInfo(
             1,
-            "UI Default post processing face restorer visibility (requires restart)",
+            "UI Default global  post processing face restorer visibility (requires restart)",
             gr.Slider,
             {"minimum": 0, "maximum": 1, "step": 0.001},
             section=section,
@@ -77,7 +77,7 @@ def on_ui_settings() -> None:
         "faceswaplab_pp_default_face_restorer_weight",
         shared.OptionInfo(
             1,
-            "UI Default post processing face restorer weight (requires restart)",
+            "UI Default global post processing face restorer weight (requires restart)",
             gr.Slider,
             {"minimum": 0, "maximum": 1, "step": 0.001},
             section=section,
@@ -87,7 +87,7 @@ def on_ui_settings() -> None:
         "faceswaplab_pp_default_upscaler",
         shared.OptionInfo(
             None,
-            "UI Default post processing upscaler (requires restart)",
+            "UI Default global post processing upscaler (requires restart)",
             gr.Dropdown,
             {
                 "interactive": True,
@@ -100,13 +100,15 @@ def on_ui_settings() -> None:
         "faceswaplab_pp_default_upscaler_visibility",
         shared.OptionInfo(
             1,
-            "UI Default post processing upscaler visibility(requires restart)",
+            "UI Default global post processing upscaler visibility(requires restart)",
             gr.Slider,
             {"minimum": 0, "maximum": 1, "step": 0.001},
             section=section,
         ),
     )
 
+    # Inpainting
+
     shared.opts.add_option(
         "faceswaplab_pp_default_inpainting_prompt",
         shared.OptionInfo(
@@ -132,20 +134,10 @@ def on_ui_settings() -> None:
     # UPSCALED SWAPPER
 
     shared.opts.add_option(
-        "faceswaplab_upscaled_swapper",
-        shared.OptionInfo(
-            False,
-            "Upscaled swapper. Applied only to the swapped faces. Apply transformations before merging with the original image.",
-            gr.Checkbox,
-            {"interactive": True},
-            section=section,
-        ),
-    )
-    shared.opts.add_option(
-        "faceswaplab_upscaled_swapper_upscaler",
+        "faceswaplab_default_upscaled_swapper_upscaler",
         shared.OptionInfo(
             None,
-            "Upscaled swapper upscaler (Recommanded : LDSR but slow)",
+            "Default Upscaled swapper upscaler (Recommanded : LDSR but slow) (requires restart)",
             gr.Dropdown,
             {
                 "interactive": True,
@@ -155,40 +147,40 @@ def on_ui_settings() -> None:
         ),
     )
     shared.opts.add_option(
-        "faceswaplab_upscaled_swapper_sharpen",
+        "faceswaplab_default_upscaled_swapper_sharpen",
         shared.OptionInfo(
             False,
-            "Upscaled swapper sharpen",
+            "Default Upscaled swapper sharpen",
             gr.Checkbox,
             {"interactive": True},
             section=section,
         ),
     )
     shared.opts.add_option(
-        "faceswaplab_upscaled_swapper_fixcolor",
+        "faceswaplab_default_upscaled_swapper_fixcolor",
         shared.OptionInfo(
             False,
-            "Upscaled swapper color correction",
+            "Default Upscaled swapper color corrections (requires restart)",
             gr.Checkbox,
             {"interactive": True},
             section=section,
         ),
     )
     shared.opts.add_option(
-        "faceswaplab_upscaled_improved_mask",
+        "faceswaplab_default_upscaled_swapper_improved_mask",
         shared.OptionInfo(
             True,
-            "Use improved segmented mask (use pastenet to mask only the face)",
+            "Default Use improved segmented mask (use pastenet to mask only the face) (requires restart)",
             gr.Checkbox,
             {"interactive": True},
             section=section,
         ),
     )
     shared.opts.add_option(
-        "faceswaplab_upscaled_swapper_face_restorer",
+        "faceswaplab_default_upscaled_swapper_face_restorer",
         shared.OptionInfo(
             None,
-            "Upscaled swapper face restorer",
+            "Default Upscaled swapper face restorer (requires restart)",
             gr.Dropdown,
             {
                 "interactive": True,
@@ -198,40 +190,30 @@ def on_ui_settings() -> None:
         ),
     )
     shared.opts.add_option(
-        "faceswaplab_upscaled_swapper_face_restorer_visibility",
+        "faceswaplab_default_upscaled_swapper_face_restorer_visibility",
         shared.OptionInfo(
             1,
-            "Upscaled swapper face restorer visibility",
+            "Default Upscaled swapper face restorer visibility (requires restart)",
             gr.Slider,
             {"minimum": 0, "maximum": 1, "step": 0.001},
             section=section,
         ),
     )
     shared.opts.add_option(
-        "faceswaplab_upscaled_swapper_face_restorer_weight",
+        "faceswaplab_default_upscaled_swapper_face_restorer_weight",
         shared.OptionInfo(
             1,
-            "Upscaled swapper face restorer weight (codeformer)",
+            "Default Upscaled swapper face restorer weight (codeformer) (requires restart)",
             gr.Slider,
             {"minimum": 0, "maximum": 1, "step": 0.001},
             section=section,
         ),
     )
     shared.opts.add_option(
-        "faceswaplab_upscaled_swapper_fthresh",
-        shared.OptionInfo(
-            10,
-            "Upscaled swapper fthresh (diff sensitivity) 10 = default behaviour. Low impact.",
-            gr.Slider,
-            {"minimum": 5, "maximum": 250, "step": 1},
-            section=section,
-        ),
-    )
-    shared.opts.add_option(
-        "faceswaplab_upscaled_swapper_erosion",
+        "faceswaplab_default_upscaled_swapper_erosion",
         shared.OptionInfo(
             1,
-            "Upscaled swapper mask erosion factor, 1 = default behaviour. The larger it is, the more blur is applied around the face. Too large and the facial change is no longer visible.",
+            "Default Upscaled swapper mask erosion factor, 1 = default behaviour. The larger it is, the more blur is applied around the face. Too large and the facial change is no longer visible. (requires restart)",
             gr.Slider,
             {"minimum": 0, "maximum": 10, "step": 0.001},
             section=section,

scripts/faceswaplab_swapping/swapper.py

@@ -1,8 +1,14 @@
 import copy
 import os
 from dataclasses import dataclass
-from typing import Any, Dict, List, Set, Tuple, Optional
+from pprint import pformat
+from typing import Any, Dict, Generator, List, Set, Tuple, Optional
 import tempfile
+from tqdm import tqdm
+import sys
+from io import StringIO
+from contextlib import contextmanager
+import hashlib
 
 import cv2
 import insightface
@@ -13,6 +19,7 @@ from PIL import Image
 from sklearn.metrics.pairwise import cosine_similarity
 
 from scripts.faceswaplab_swapping import upscaled_inswapper
+from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions
 from scripts.faceswaplab_utils.imgutils import (
     pil_to_cv2,
     check_against_nsfw,
@@ -117,19 +124,16 @@ def batch_process(
             for src_image in src_images:
                 current_images = []
                 swapped_images = process_images_units(
-                    get_current_model(),
-                    images=[(src_image, None)],
-                    units=units,
-                    upscaled_swapper=opts.data.get(
-                        "faceswaplab_upscaled_swapper", False
-                    ),
+                    get_current_model(), images=[(src_image, None)], units=units
                 )
                 if len(swapped_images) > 0:
                     current_images += [img for img, _ in swapped_images]
 
                 logger.info("%s images generated", len(current_images))
-                for i, img in enumerate(current_images):
-                    current_images[i] = enhance_image(img, postprocess_options)
+
+                if postprocess_options:
+                    for i, img in enumerate(current_images):
+                        current_images[i] = enhance_image(img, postprocess_options)
 
                 if save_path:
                     for img in current_images:
@@ -225,6 +229,33 @@ class FaceModelException(Exception):
         super().__init__(self.message)
 
 
+@contextmanager
+def capture_stdout() -> Generator[StringIO, None, None]:
+    """
+    Capture and yield the printed messages to stdout.
+
+    This context manager temporarily replaces sys.stdout with a StringIO object,
+    capturing all printed output. After the context block is exited, sys.stdout
+    is restored to its original value.
+
+    Example usage:
+    with capture_stdout() as captured:
+        print("Hello, World!")
+        output = captured.getvalue()
+        # output now contains "Hello, World!\n"
+
+    Returns:
+        A StringIO object containing the captured output.
+    """
+    original_stdout = sys.stdout  # Type: ignore
+    captured_stdout = StringIO()
+    sys.stdout = captured_stdout  # Type: ignore
+    try:
+        yield captured_stdout
+    finally:
+        sys.stdout = original_stdout  # Type: ignore
+
+
 @lru_cache(maxsize=1)
 def getAnalysisModel() -> insightface.app.FaceAnalysis:
     """
@@ -237,11 +268,20 @@ def getAnalysisModel() -> insightface.app.FaceAnalysis:
         if not os.path.exists(faceswaplab_globals.ANALYZER_DIR):
             os.makedirs(faceswaplab_globals.ANALYZER_DIR)
 
-        logger.info("Load analysis model, will take some time.")
+        logger.info("Load analysis model, will take some time. (> 30s)")
         # Initialize the analysis model with the specified name and providers
-        return insightface.app.FaceAnalysis(
-            name="buffalo_l", providers=providers, root=faceswaplab_globals.ANALYZER_DIR
-        )
+
+        with tqdm(total=1, desc="Loading analysis model", unit="model") as pbar:
+            with capture_stdout() as captured:
+                model = insightface.app.FaceAnalysis(
+                    name="buffalo_l",
+                    providers=providers,
+                    root=faceswaplab_globals.ANALYZER_DIR,
+                )
+            pbar.update(1)
+        logger.info("%s", pformat(captured.getvalue()))
+
+        return model
     except Exception as e:
         logger.error(
             "Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)"
@@ -249,11 +289,8 @@ def getAnalysisModel() -> insightface.app.FaceAnalysis:
         raise FaceModelException("Loading of analysis model failed")
 
 
-import hashlib
-
-
 def is_sha1_matching(file_path: str, expected_sha1: str) -> bool:
-    sha1_hash = hashlib.sha1()
+    sha1_hash = hashlib.sha1(usedforsecurity=False)
 
     with open(file_path, "rb") as file:
         for byte_block in iter(lambda: file.read(4096), b""):
@@ -284,10 +321,15 @@ def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper:
                 expected_sha1,
             )
 
-        # Initializes the face swap model using the specified model path.
-        return upscaled_inswapper.UpscaledINSwapper(
-            insightface.model_zoo.get_model(model_path, providers=providers)
-        )
+        with tqdm(total=1, desc="Loading swap model", unit="model") as pbar:
+            with capture_stdout() as captured:
+                model = upscaled_inswapper.UpscaledINSwapper(
+                    insightface.model_zoo.get_model(model_path, providers=providers)
+                )
+            pbar.update(1)
+        logger.info("%s", pformat(captured.getvalue()))
+        return model
+
     except Exception as e:
         logger.error(
             "Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)"
@@ -498,7 +540,7 @@ def swap_face(
     target_img: PILImage,
     target_faces: List[Face],
     model: str,
-    upscaled_swapper: bool = False,
+    swapping_options: Optional[InswappperOptions],
     compute_similarity: bool = True,
 ) -> ImageResult:
     """
@@ -532,7 +574,7 @@ def swap_face(
                     img=result,
                     target_face=swapped_face,
                     source_face=source_face,
-                    upscale=upscaled_swapper,
+                    options=swapping_options,
                 )  # type: ignore
 
             result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
@@ -581,7 +623,6 @@ def process_image_unit(
     unit: FaceSwapUnitSettings,
     image: PILImage,
     info: str = None,
-    upscaled_swapper: bool = False,
     force_blend: bool = False,
 ) -> List[Tuple[PILImage, str]]:
     """Process one image and return a List of (image, info) (one if blended, many if not).
@@ -639,7 +680,7 @@ def process_image_unit(
                 target_img=current_image,
                 target_faces=target_faces,
                 model=model,
-                upscaled_swapper=upscaled_swapper,
+                swapping_options=unit.swapping_options,
                 compute_similarity=unit.compute_similarity,
             )
             # Apply post-inpainting to image
@@ -694,7 +735,6 @@ def process_images_units(
     model: str,
     units: List[FaceSwapUnitSettings],
     images: List[Tuple[Optional[PILImage], Optional[str]]],
-    upscaled_swapper: bool = False,
     force_blend: bool = False,
 ) -> Optional[List[Tuple[PILImage, str]]]:
     """
@@ -725,13 +765,9 @@ def process_images_units(
     processed_images = []
     for i, (image, info) in enumerate(images):
         logger.debug("Processing image %s", i)
-        swapped = process_image_unit(
-            model, units[0], image, info, upscaled_swapper, force_blend
-        )
+        swapped = process_image_unit(model, units[0], image, info, force_blend)
         logger.debug("Image %s -> %s images", i, len(swapped))
-        nexts = process_images_units(
-            model, units[1:], swapped, upscaled_swapper, force_blend
-        )
+        nexts = process_images_units(model, units[1:], swapped, force_blend)
         if nexts:
             processed_images.extend(nexts)
         else:

scripts/faceswaplab_swapping/upcaled_inswapper_options.py

@@ -0,0 +1,38 @@
+from dataclasses import *
+from client_api import api_utils
+
+
+@dataclass
+class InswappperOptions:
+    face_restorer_name: str = None
+    restorer_visibility: float = 1
+    codeformer_weight: float = 1
+    upscaler_name: str = None
+    improved_mask: bool = False
+    color_corrections: bool = False
+    sharpen: bool = False
+    erosion_factor: float = 1.0
+
+    @staticmethod
+    def from_api_dto(dto: api_utils.InswappperOptions) -> "InswappperOptions":
+        """
+        Converts a InpaintingOptions object from an API DTO (Data Transfer Object).
+
+        :param options: An object of api_utils.InpaintingOptions representing the
+                        post-processing options as received from the API.
+        :return: A InpaintingOptions instance containing the translated values
+                from the API DTO.
+        """
+        if dto is None:
+            return InswappperOptions()
+
+        return InswappperOptions(
+            face_restorer_name=dto.face_restorer_name,
+            restorer_visibility=dto.restorer_visibility,
+            codeformer_weight=dto.codeformer_weight,
+            upscaler_name=dto.upscaler_name,
+            improved_mask=dto.improved_mask,
+            color_corrections=dto.color_corrections,
+            sharpen=dto.sharpen,
+            erosion_factor=dto.erosion_factor,
+        )

scripts/faceswaplab_swapping/upscaled_inswapper.py

@@ -12,8 +12,10 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
     PostProcessingOptions,
 )
 from scripts.faceswaplab_swapping.facemask import generate_face_mask
+from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions
 from scripts.faceswaplab_utils.imgutils import cv2_to_pil, pil_to_cv2
 from scripts.faceswaplab_utils.typing import CV2ImgU8, Face
+from scripts.faceswaplab_utils.faceswaplab_logging import logger
 
 
 def get_upscaler() -> UpscalerData:
@@ -127,26 +129,25 @@ class UpscaledINSwapper(INSwapper):
     def __init__(self, inswapper: INSwapper):
         self.__dict__.update(inswapper.__dict__)
 
-    def super_resolution(self, img: CV2ImgU8, k: int = 2) -> CV2ImgU8:
+    def upscale_and_restore(
+        self, img: CV2ImgU8, k: int = 2, inswapper_options: InswappperOptions = None
+    ) -> CV2ImgU8:
         pil_img = cv2_to_pil(img)
-        options = PostProcessingOptions(
-            upscaler_name=opts.data.get(
-                "faceswaplab_upscaled_swapper_upscaler", "LDSR"
-            ),
+        pp_options = PostProcessingOptions(
+            upscaler_name=inswapper_options.upscaler_name,
             upscale_visibility=1,
             scale=k,
-            face_restorer_name=opts.data.get(
-                "faceswaplab_upscaled_swapper_face_restorer", ""
-            ),
-            codeformer_weight=opts.data.get(
-                "faceswaplab_upscaled_swapper_face_restorer_weight", 1
-            ),
-            restorer_visibility=opts.data.get(
-                "faceswaplab_upscaled_swapper_face_restorer_visibility", 1
-            ),
+            face_restorer_name=inswapper_options.face_restorer_name,
+            codeformer_weight=inswapper_options.codeformer_weight,
+            restorer_visibility=inswapper_options.restorer_visibility,
         )
-        upscaled = upscaling.upscale_img(pil_img, options)
-        upscaled = upscaling.restore_face(upscaled, options)
+
+        upscaled = pil_img
+        if pp_options.upscaler_name:
+            upscaled = upscaling.upscale_img(pil_img, pp_options)
+        if pp_options.face_restorer_name:
+            upscaled = upscaling.restore_face(upscaled, pp_options)
+
         return pil_to_cv2(upscaled)
 
     def get(
@@ -155,7 +156,7 @@ class UpscaledINSwapper(INSwapper):
         target_face: Face,
         source_face: Face,
         paste_back: bool = True,
-        upscale: bool = True,
+        options: InswappperOptions = None,
     ) -> Union[CV2ImgU8, Tuple[CV2ImgU8, Any]]:
         aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0])
         blob = cv2.dnn.blobFromImage(
@@ -190,43 +191,48 @@ class UpscaledINSwapper(INSwapper):
                     fake_diff[:, -2:] = 0
                     return fake_diff
 
-                if upscale:
-                    print("*" * 80)
-                    print(
-                        f"Upscaled inswapper using {opts.data.get('faceswaplab_upscaled_swapper_upscaler', 'LDSR')}"
-                    )
-                    print("*" * 80)
+                if options:
+                    logger.info("*" * 80)
+                    logger.info(f"Upscaled inswapper")
 
-                    k = 4
-                    aimg, M = face_align.norm_crop2(
-                        img, target_face.kps, self.input_size[0] * k
-                    )
+                    if options.upscaler_name:
+                        # Upscale original image
+                        k = 4
+                        aimg, M = face_align.norm_crop2(
+                            img, target_face.kps, self.input_size[0] * k
+                        )
+                    else:
+                        k = 1
 
                     # upscale and restore face :
-                    bgr_fake = self.super_resolution(bgr_fake, k)
+                    bgr_fake = self.upscale_and_restore(
+                        bgr_fake, inswapper_options=options, k=k
+                    )
 
-                    if opts.data.get("faceswaplab_upscaled_improved_mask", True):
+                    if options.improved_mask:
                         mask = get_face_mask(aimg, bgr_fake)
                         bgr_fake = merge_images_with_mask(aimg, bgr_fake, mask)
 
                     # compute fake_diff before sharpen and color correction (better result)
                     fake_diff = compute_diff(bgr_fake, aimg)
 
-                    if opts.data.get("faceswaplab_upscaled_swapper_sharpen", True):
-                        print("sharpen")
+                    if options.sharpen:
+                        logger.info("sharpen")
                         # Add sharpness
                         blurred = cv2.GaussianBlur(bgr_fake, (0, 0), 3)
                         bgr_fake = cv2.addWeighted(bgr_fake, 1.5, blurred, -0.5, 0)
 
                     # Apply color corrections
-                    if opts.data.get("faceswaplab_upscaled_swapper_fixcolor", True):
-                        print("color correction")
+                    if options.color_corrections:
+                        logger.info("color correction")
                         correction = processing.setup_color_correction(cv2_to_pil(aimg))
                         bgr_fake_pil = processing.apply_color_correction(
                             correction, cv2_to_pil(bgr_fake)
                         )
                         bgr_fake = pil_to_cv2(bgr_fake_pil)
 
+                    logger.info("*" * 80)
+
                 else:
                     fake_diff = compute_diff(bgr_fake, aimg)
 
@@ -254,7 +260,7 @@ class UpscaledINSwapper(INSwapper):
                     borderValue=0.0,
                 )
                 img_white[img_white > 20] = 255
-                fthresh = opts.data.get("faceswaplab_upscaled_swapper_fthresh", 10)
+                fthresh = 10
                 print("fthresh", fthresh)
                 fake_diff[fake_diff < fthresh] = 0
                 fake_diff[fake_diff >= fthresh] = 255
@@ -263,9 +269,8 @@ class UpscaledINSwapper(INSwapper):
                 mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
                 mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
                 mask_size = int(np.sqrt(mask_h * mask_w))
-                erosion_factor = opts.data.get(
-                    "faceswaplab_upscaled_swapper_erosion", 1
-                )
+                erosion_factor = options.erosion_factor
+
                 k = max(int(mask_size // 10 * erosion_factor), int(10 * erosion_factor))
 
                 kernel = np.ones((k, k), np.uint8)

scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py

@@ -17,7 +17,7 @@ def postprocessing_ui() -> List[gr.components.Component]:
                 choices=["None"] + [x.name() for x in shared.face_restorers],
                 value=lambda: opts.data.get(
                     "faceswaplab_pp_default_face_restorer",
-                    shared.face_restorers[0].name(),
+                    "None",
                 ),
                 type="value",
                 elem_id="faceswaplab_pp_face_restorer",

scripts/faceswaplab_ui/faceswaplab_tab.py

@@ -249,6 +249,8 @@ def tools_ui() -> None:
                 preview = gr.components.Image(
                     type="pil",
                     label="Preview",
+                    width=512,
+                    height=512,
                     interactive=False,
                     elem_id="faceswaplab_build_preview_face",
                 )

scripts/faceswaplab_ui/faceswaplab_unit_settings.py

@@ -6,6 +6,7 @@ from typing import List, Optional, Set, Union
 import gradio as gr
 from insightface.app.common import Face
 from PIL import Image
+from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions
 from scripts.faceswaplab_utils.imgutils import pil_to_cv2
 from scripts.faceswaplab_utils.faceswaplab_logging import logger
 from scripts.faceswaplab_utils import face_checkpoints_utils
@@ -51,6 +52,8 @@ class FaceSwapUnitSettings:
     swap_in_generated: bool
     # Pre inpainting configuration (Don't use optional for this or gradio parsing will fail) :
     pre_inpainting: InpaintingOptions
+    # Configure swapping options
+    swapping_options: InswappperOptions
     # Post inpainting configuration (Don't use optional for this or gradio parsing will fail) :
     post_inpainting: InpaintingOptions
 
@@ -81,6 +84,7 @@ class FaceSwapUnitSettings:
             swap_in_generated=True,
             swap_in_source=False,
             pre_inpainting=InpaintingOptions.from_api_dto(dto.pre_inpainting),
+            swapping_options=InswappperOptions.from_api_dto(dto.swapping_options),
             post_inpainting=InpaintingOptions.from_api_dto(dto.post_inpainting),
         )
 

scripts/faceswaplab_ui/faceswaplab_unit_ui.py

@@ -2,6 +2,100 @@ from typing import List
 from scripts.faceswaplab_ui.faceswaplab_inpainting_ui import face_inpainting_ui
 from scripts.faceswaplab_utils.face_checkpoints_utils import get_face_checkpoints
 import gradio as gr
+from modules.shared import opts
+from modules import shared
+
+
+def faceswap_unit_advanced_options(
+    is_img2img: bool, unit_num: int = 1, id_prefix: str = "faceswaplab_"
+) -> List[gr.components.Component]:
+    with gr.Accordion(f"Post-Processing & Advanced Mask Options", open=False):
+        gr.Markdown("""Post-processing and mask settings for unit faces""")
+        with gr.Row():
+            face_restorer_name = gr.Radio(
+                label="Restore Face",
+                choices=["None"] + [x.name() for x in shared.face_restorers],
+                value=lambda: opts.data.get(
+                    "faceswaplab_default_upscaled_swapper_face_restorer",
+                    "None",
+                ),
+                type="value",
+                elem_id=f"{id_prefix}_face{unit_num}_face_restorer",
+            )
+            with gr.Column():
+                face_restorer_visibility = gr.Slider(
+                    0,
+                    1,
+                    value=lambda: opts.data.get(
+                        "faceswaplab_default_upscaled_swapper_face_restorer_visibility",
+                        1.0,
+                    ),
+                    step=0.001,
+                    label="Restore visibility",
+                    elem_id=f"{id_prefix}_face{unit_num}_face_restorer_visibility",
+                )
+                codeformer_weight = gr.Slider(
+                    0,
+                    1,
+                    value=lambda: opts.data.get(
+                        "faceswaplab_default_upscaled_swapper_face_restorer_weight", 1.0
+                    ),
+                    step=0.001,
+                    label="codeformer weight",
+                    elem_id=f"{id_prefix}_face{unit_num}_face_restorer_weight",
+                )
+        upscaler_name = gr.Dropdown(
+            choices=[upscaler.name for upscaler in shared.sd_upscalers],
+            value=lambda: opts.data.get(
+                "faceswaplab_default_upscaled_swapper_upscaler", ""
+            ),
+            label="Upscaler",
+            elem_id=f"{id_prefix}_face{unit_num}_upscaler",
+        )
+
+        improved_mask = gr.Checkbox(
+            lambda: opts.data.get(
+                "faceswaplab_default_upscaled_swapper_improved_mask", False
+            ),
+            interactive=True,
+            label="Use improved segmented mask (use pastenet to mask only the face)",
+            elem_id=f"{id_prefix}_face{unit_num}_improved_mask",
+        )
+        color_corrections = gr.Checkbox(
+            lambda: opts.data.get(
+                "faceswaplab_default_upscaled_swapper_fixcolor", False
+            ),
+            interactive=True,
+            label="Use color corrections",
+            elem_id=f"{id_prefix}_face{unit_num}_color_corrections",
+        )
+        sharpen_face = gr.Checkbox(
+            lambda: opts.data.get(
+                "faceswaplab_default_upscaled_swapper_sharpen", False
+            ),
+            interactive=True,
+            label="sharpen face",
+            elem_id=f"{id_prefix}_face{unit_num}_sharpen_face",
+        )
+        erosion_factor = gr.Slider(
+            0.0,
+            10.0,
+            lambda: opts.data.get("faceswaplab_default_upscaled_swapper_erosion", 1.0),
+            step=0.01,
+            label="Upscaled swapper mask erosion factor, 1 = default behaviour.",
+            elem_id=f"{id_prefix}_face{unit_num}_erosion_factor",
+        )
+
+    return [
+        face_restorer_name,
+        face_restorer_visibility,
+        codeformer_weight,
+        upscaler_name,
+        improved_mask,
+        color_corrections,
+        sharpen_face,
+        erosion_factor,
+    ]
 
 
 def faceswap_unit_ui(
@@ -62,35 +156,6 @@ def faceswap_unit_ui(
                     elem_id=f"{id_prefix}_face{unit_num}_blend_faces",
                     interactive=True,
                 )
-            gr.Markdown("""Discard images with low similarity or no faces :""")
-            with gr.Row():
-                check_similarity = gr.Checkbox(
-                    False,
-                    placeholder="discard",
-                    label="Check similarity",
-                    elem_id=f"{id_prefix}_face{unit_num}_check_similarity",
-                )
-                compute_similarity = gr.Checkbox(
-                    False,
-                    label="Compute similarity",
-                    elem_id=f"{id_prefix}_face{unit_num}_compute_similarity",
-                )
-            min_sim = gr.Slider(
-                0,
-                1,
-                0,
-                step=0.01,
-                label="Min similarity",
-                elem_id=f"{id_prefix}_face{unit_num}_min_similarity",
-            )
-            min_ref_sim = gr.Slider(
-                0,
-                1,
-                0,
-                step=0.01,
-                label="Min reference similarity",
-                elem_id=f"{id_prefix}_face{unit_num}_min_ref_similarity",
-            )
 
             gr.Markdown(
                 """Select the face to be swapped, you can sort by size or use the same gender as the desired face:"""
@@ -143,11 +208,46 @@ def faceswap_unit_ui(
                 visible=is_img2img,
                 elem_id=f"{id_prefix}_face{unit_num}_swap_in_generated",
             )
+
+        with gr.Accordion("Similarity", open=False):
+            gr.Markdown("""Discard images with low similarity or no faces :""")
+            with gr.Row():
+                check_similarity = gr.Checkbox(
+                    False,
+                    placeholder="discard",
+                    label="Check similarity",
+                    elem_id=f"{id_prefix}_face{unit_num}_check_similarity",
+                )
+                compute_similarity = gr.Checkbox(
+                    False,
+                    label="Compute similarity",
+                    elem_id=f"{id_prefix}_face{unit_num}_compute_similarity",
+                )
+            min_sim = gr.Slider(
+                0,
+                1,
+                0,
+                step=0.01,
+                label="Min similarity",
+                elem_id=f"{id_prefix}_face{unit_num}_min_similarity",
+            )
+            min_ref_sim = gr.Slider(
+                0,
+                1,
+                0,
+                step=0.01,
+                label="Min reference similarity",
+                elem_id=f"{id_prefix}_face{unit_num}_min_ref_similarity",
+            )
+
         pre_inpainting = face_inpainting_ui(
             name="Pre-Inpainting (Before swapping)",
             id_prefix=f"{id_prefix}_face{unit_num}_preinpainting",
             description="Pre-inpainting sends face to inpainting before swapping",
         )
+
+        options = faceswap_unit_advanced_options(is_img2img, unit_num, id_prefix)
+
         post_inpainting = face_inpainting_ui(
             name="Post-Inpainting (After swapping)",
             id_prefix=f"{id_prefix}_face{unit_num}_postinpainting",
@@ -173,6 +273,7 @@ def faceswap_unit_ui(
             swap_in_generated,
         ]
         + pre_inpainting
+        + options
         + post_inpainting
     )
 

scripts/faceswaplab_utils/face_checkpoints_utils.py

@@ -7,21 +7,19 @@ import torch
 
 import modules.scripts as scripts
 from modules import scripts
+from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions
 from scripts.faceswaplab_utils.faceswaplab_logging import logger
 from scripts.faceswaplab_utils.typing import *
 from scripts.faceswaplab_utils import imgutils
-from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
-from scripts.faceswaplab_postprocessing.postprocessing_options import (
-    PostProcessingOptions,
-)
 from scripts.faceswaplab_utils.models_utils import get_models
-from modules.shared import opts
 import traceback
 
 import dill as pickle  # will be removed in future versions
 from scripts.faceswaplab_swapping import swapper
 from pprint import pformat
 import re
+from client_api import api_utils
+import tempfile
 
 
 def sanitize_name(name: str) -> str:
@@ -93,16 +91,9 @@ def build_face_checkpoint_and_save(
                     source_face=blended_face,
                     target_img=reference_preview_img,
                     model=get_models()[0],
-                    upscaled_swapper=opts.data.get(
-                        "faceswaplab_upscaled_swapper", False
-                    ),
-                )
-                preview_image = enhance_image(
-                    result.image,
-                    PostProcessingOptions(
-                        face_restorer_name="CodeFormer", restorer_visibility=1
-                    ),
+                    swapping_options=InswappperOptions(face_restorer_name="Codeformer"),
                 )
+                preview_image = result.image
 
             file_path = os.path.join(get_checkpoint_path(), f"{name}.safetensors")
             if not overwrite:
@@ -147,6 +138,16 @@ def save_face(face: Face, filename: str) -> None:
 
 
 def load_face(name: str) -> Face:
+    if name.startswith("data:application/face;base64,"):
+        with tempfile.NamedTemporaryFile(delete=True) as temp_file:
+            api_utils.base64_to_safetensors(name, temp_file.name)
+            face = {}
+            with safe_open(temp_file.name, framework="pt", device="cpu") as f:
+                for k in f.keys():
+                    logger.debug("load key %s", k)
+                    face[k] = f.get_tensor(k).numpy()
+            return Face(face)
+
     filename = matching_checkpoint(name)
     if filename is None:
         return None