Merge pull request #96 from glucauze/v1.2.7

v1.2.7 

    remove dill
    add warnings in model install and checking
    add gender selection in build

.pre-commit-config.yaml

@@ -1,8 +1,4 @@
 repos:
-    - repo: https://github.com/psf/black
-      rev: 23.7.0
-      hooks:
-        - id: black
     - repo: https://github.com/pre-commit/pre-commit-hooks
       rev: v4.4.0
       hooks:
@@ -10,4 +6,8 @@ repos:
         - id: check-case-conflict
         - id: check-docstring-first
         - id: detect-private-key
-        - id: fix-byte-order-marker
+        - id: fix-byte-order-marker
+    - repo: https://github.com/psf/black
+      rev: 23.7.0
+      hooks:
+        - id: black

CHANGELOG.md

@@ -1,7 +1,61 @@
+# 1.2.5
+
+Allow seed selection in inpainting.
+
+# 1.2.4
+
+Fix default settings by marking only managed field as do_not_save.
+
+See the discussion here : https://github.com/glucauze/sd-webui-faceswaplab/issues/62
+
+# 1.2.3
+
+Speed up ui : change the way default settings are manage by not storing them in ui-config.json
+
+Migration : YOU NEED TO recreate ui-config.json (delete) or at least remove any faceswaplab reference to be able to use default settings again.
+
+See this for explainations : https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/6109
+
+# 1.2.2
+
++ Add NSFW filter option in settings (1 == disable)
++ Improve install speed
++ Install gpu requirements by default if --use-cpu is not used
++ Fix improved mask + color correction
++ Remove javascript, use https://github.com/w-e-w/sdwebui-close-confirmation-dialogue.git instead to prevent gradio from closing.
+
+# 1.2.1 :
+
+Add GPU support option : see https://github.com/glucauze/sd-webui-faceswaplab/pull/24
+
+# 1.2.0 :
+
+This version changes quite a few things.
+
++ The upscaled inswapper options are now moved to each face unit. This makes it possible to fine-tune the settings for each face.
+
++ Upscaled inswapper configuration in sd now concerns default values in each unit's interface.
+
++ Pre- and post-inpainting is now possible for each face. Here too, default options are set in the main sd settings.
+
++ Codeformer is no longer the default in post-processing. Don't be surprised if you get bad results by default. You can set it to default in the application's global settings
+
+Bug fixes :
+
++ The problem of saving the grid should be solved.
++ The downscaling problem for inpainting should be solved.
++ Change model download logic and add checksum. This should prevent some bugs.
+
+In terms of the API, it is now possible to create a remote checkpoint and use it in units. See the example in client_api or the tests in the tests directory.
+
+See https://github.com/glucauze/sd-webui-faceswaplab/pull/19
+
 # 1.1.2 :
 
 + Switch face checkpoint format from pkl to safetensors
 
+See https://github.com/glucauze/sd-webui-faceswaplab/pull/4
+
 ## 1.1.1 :
 
 + Add settings for default inpainting prompts

README.md

@@ -1,6 +1,12 @@
 # FaceSwapLab for a1111/Vlad
 
-Please read the documentation here : https://glucauze.github.io/sd-webui-faceswaplab/
+V1.2.3 : Breaking change for settings, please read changelog.
+
+Please read the documentation here : https://glucauze.github.io/sd-webui-faceswaplab/ 
+
+You can also read the [doc discussion section](https://github.com/glucauze/sd-webui-faceswaplab/discussions/categories/guide-doc)
+
+See [CHANGELOG.md](CHANGELOG.md) for changes in last versions.
 
 FaceSwapLab is an extension for Stable Diffusion that simplifies face-swapping. It has evolved from sd-webui-faceswap and some part of sd-webui-roop. However, a substantial amount of the code has been rewritten to improve performance and to better manage masks.
 
@@ -14,12 +20,71 @@ While FaceSwapLab is still under development, it has reached a good level of sta
 
 In short:
 
-+ **Ethical Guideline:** This extension should not be forked to create a public, easy way to circumvent NSFW filtering.
++ **Ethical Guideline:**  NSFW is now configurable due to performance issue. Please don't use this to do harm.
 + **License:** This software is distributed under the terms of the GNU Affero General Public License (AGPL), version 3 or later.
 + **Model License:** This software uses InsightFace's pre-trained models, which are available for non-commercial research purposes only.
 
 More on this here : https://glucauze.github.io/sd-webui-faceswaplab/ 
 
+### Known problems (wontfix):
+
++ Older versions of gradio don't work well with the extension. See this bug : https://github.com/glucauze/sd-webui-faceswaplab/issues/5
+
+## Quick Start
+
+Here are some gifs to explain (non cherry picked, just random pictures) : 
+
+## Simple Usage (roop like)
+
+This use codeformer on all faces (including non swapped)
+
+[simple.webm](https://github.com/glucauze/sd-webui-faceswaplab/assets/137925069/de00b685-d441-44f9-bae3-71cd7abef113)
+
+## Advanced options 
+
+This is use to improve results. This use upscaling and codeformer only on swapped faces
+
+[advanced.webm](https://github.com/glucauze/sd-webui-faceswaplab/assets/137925069/50630311-bd25-487f-871b-0a44eecd435d)
+
+## Inpainting 
+
+This add inpainting on faces : 
+
+[inpainting.webm](https://github.com/glucauze/sd-webui-faceswaplab/assets/137925069/3d3508e9-5be4-4566-8c41-8301b2d08355)
+
+## Build and use checkpoints : 
+
+[build.webm](https://github.com/glucauze/sd-webui-faceswaplab/assets/137925069/e84e9a3c-840d-4536-9fbb-09ed256406d7)
+
+
+
+### Simple
+
+1. Put a face in the reference.
+2. Select a face number.
+3. Select "Enable."
+4. Select "CodeFormer" in **Global Post-Processing** tab.
+
+Once you're happy with some results but want to improve, the next steps are to:
+
++ Use advanced settings in face units (which are not as complex as they might seem, it's basically fine tuning post-processing for each faces).
++ Use pre/post inpainting to tweak the image a bit for more natural results.
+
+### Better 
+
+1. Put a face in the reference.
+2. Select a face number.
+3. Select "Enable."
+
+4. In **Post-Processing** accordeon:
+    + Select "CodeFormer" 
+    + Select "LDSR" or a faster model "003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN" in upscaler. See [here for a list of upscalers](https://github.com/glucauze/sd-webui-faceswaplab/discussions/29). 
+    + Use sharpen, color_correction and improved mask
+
+5. Disable "CodeFormer" in **Global Post-Processing** tab (otherwise it will be applied twice)
+
+Don't hesitate to share config in the [discussion section](https://github.com/glucauze/sd-webui-faceswaplab/discussions).
+
 ### Features
 
 + **Face Unit Concept**: Similar to controlNet, the program introduces the concept of a face unit. You can configure up to 10 units (3 units are the default setting) in the program settings (sd).
@@ -28,6 +93,8 @@ More on this here : https://glucauze.github.io/sd-webui-faceswaplab/
 
 + **Batch Processing**
 
++ **GPU**
+
 + **Inpainting Fixes** : supports “only masked” and mask inpainting.
 
 + **Performance Improvements**: The overall performance of the software has been enhanced.
@@ -58,7 +125,7 @@ More on this here : https://glucauze.github.io/sd-webui-faceswaplab/
 
 + **Upscaled Inswapper**: The program now includes an upscaled inswapper option, which improves results by incorporating upsampling, sharpness adjustment, and color correction before face is merged to the original image.
 
-+ **API with typing support** :
++ **API with typing support**
 
 
 ## Installation

check.sh

@@ -1,4 +1,4 @@
 #!/bin/bash
 autoflake --in-place --remove-unused-variables -r --remove-all-unused-imports .
-mypy --install-types
+mypy --non-interactive --install-types
 pre-commit run --all-files

client_api/api_utils.py

@@ -9,6 +9,7 @@ from io import BytesIO
 from typing import List, Tuple, Optional
 import numpy as np
 import requests
+import safetensors
 
 
 class InpaintingWhen(Enum):
@@ -18,6 +19,55 @@ class InpaintingWhen(Enum):
     AFTER_ALL = "After All"
 
 
+class InpaintingOptions(BaseModel):
+    inpainting_denoising_strengh: float = Field(
+        description="Inpainting denoising strenght", default=0, lt=1, ge=0
+    )
+    inpainting_prompt: str = Field(
+        description="Inpainting denoising strenght",
+        examples=["Portrait of a [gender]"],
+        default="Portrait of a [gender]",
+    )
+    inpainting_negative_prompt: str = Field(
+        description="Inpainting denoising strenght",
+        examples=[
+            "Deformed, blurry, bad anatomy, disfigured, poorly drawn face, mutation"
+        ],
+        default="",
+    )
+    inpainting_steps: int = Field(
+        description="Inpainting steps",
+        examples=["Portrait of a [gender]"],
+        ge=1,
+        le=150,
+        default=20,
+    )
+    inpainting_sampler: str = Field(
+        description="Inpainting sampler", examples=["Euler"], default="Euler"
+    )
+    inpainting_model: str = Field(
+        description="Inpainting model", examples=["Current"], default="Current"
+    )
+    inpainting_seed: int = Field(description="Inpainting Seed", ge=-1, default=-1)
+
+
+class InswappperOptions(BaseModel):
+    face_restorer_name: str = Field(
+        description="face restorer name", default="CodeFormer"
+    )
+    restorer_visibility: float = Field(
+        description="face restorer visibility", default=1, le=1, ge=0
+    )
+    codeformer_weight: float = Field(
+        description="face restorer codeformer weight", default=1, le=1, ge=0
+    )
+    upscaler_name: str = Field(description="upscaler name", default=None)
+    improved_mask: bool = Field(description="Use Improved Mask", default=False)
+    color_corrections: bool = Field(description="Use Color Correction", default=False)
+    sharpen: bool = Field(description="Sharpen Image", default=False)
+    erosion_factor: float = Field(description="Erosion Factor", default=1, le=10, ge=0)
+
+
 class FaceSwapUnit(BaseModel):
     # The image given in reference
     source_img: str = Field(
@@ -82,6 +132,21 @@ class FaceSwapUnit(BaseModel):
         default=0,
     )
 
+    pre_inpainting: Optional[InpaintingOptions] = Field(
+        description="Inpainting options",
+        default=None,
+    )
+
+    swapping_options: Optional[InswappperOptions] = Field(
+        description="PostProcessing & Mask options",
+        default=None,
+    )
+
+    post_inpainting: Optional[InpaintingOptions] = Field(
+        description="Inpainting options",
+        default=None,
+    )
+
     def get_batch_images(self) -> List[Image.Image]:
         images = []
         if self.batch_images:
@@ -104,39 +169,15 @@ class PostProcessingOptions(BaseModel):
     upscaler_visibility: float = Field(
         description="upscaler visibility", default=1, le=1, ge=0
     )
-
-    inpainting_denoising_strengh: float = Field(
-        description="Inpainting denoising strenght", default=0, lt=1, ge=0
-    )
-    inpainting_prompt: str = Field(
-        description="Inpainting denoising strenght",
-        examples=["Portrait of a [gender]"],
-        default="Portrait of a [gender]",
-    )
-    inpainting_negative_prompt: str = Field(
-        description="Inpainting denoising strenght",
-        examples=[
-            "Deformed, blurry, bad anatomy, disfigured, poorly drawn face, mutation"
-        ],
-        default="",
-    )
-    inpainting_steps: int = Field(
-        description="Inpainting steps",
-        examples=["Portrait of a [gender]"],
-        ge=1,
-        le=150,
-        default=20,
-    )
-    inpainting_sampler: str = Field(
-        description="Inpainting sampler", examples=["Euler"], default="Euler"
-    )
     inpainting_when: InpaintingWhen = Field(
         description="When inpainting happens",
         examples=[e.value for e in InpaintingWhen.__members__.values()],
         default=InpaintingWhen.NEVER,
     )
-    inpainting_model: str = Field(
+
-        description="Inpainting model", examples=["Current"], default="Current"
+    inpainting_options: Optional[InpaintingOptions] = Field(
+        description="Inpainting options",
+        default=None,
     )
 
 
@@ -147,7 +188,7 @@ class FaceSwapRequest(BaseModel):
         default=None,
     )
     units: List[FaceSwapUnit]
-    postprocessing: Optional[PostProcessingOptions]
+    postprocessing: Optional[PostProcessingOptions] = None
 
 
 class FaceSwapResponse(BaseModel):
@@ -227,3 +268,26 @@ def compare_faces(
     )
 
     return float(result.text)
+
+
+def safetensors_to_base64(file_path: str) -> str:
+    with open(file_path, "rb") as file:
+        file_bytes = file.read()
+        return "data:application/face;base64," + base64.b64encode(file_bytes).decode(
+            "utf-8"
+        )
+
+
+def base64_to_safetensors(base64str: str, output_path: str) -> None:
+    try:
+        base64_data = base64str.split("base64,")[-1]
+        file_bytes = base64.b64decode(base64_data)
+        with open(output_path, "wb") as file:
+            file.write(file_bytes)
+        with safetensors.safe_open(output_path, framework="pt") as f:
+            print(output_path, "keys =", f.keys())
+    except Exception as e:
+        print("Error : failed to convert base64 string to safetensor", e)
+        import traceback
+
+        traceback.print_exc()

client_api/faceswaplab_api_example.py

@@ -1,18 +1,25 @@
+from typing import List
 import requests
 from api_utils import (
-    FaceSwapRequest,
     FaceSwapUnit,
-    PostProcessingOptions,
+    InswappperOptions,
-    FaceSwapResponse,
+    base64_to_safetensors,
     pil_to_base64,
+    PostProcessingOptions,
     InpaintingWhen,
-    FaceSwapCompareRequest,
+    InpaintingOptions,
+    FaceSwapRequest,
+    FaceSwapResponse,
     FaceSwapExtractRequest,
+    FaceSwapCompareRequest,
     FaceSwapExtractResponse,
+    safetensors_to_base64,
 )
 
 address = "http://127.0.0.1:7860"
 
+# This has been tested on Linux platforms. This might requires some minor adaptations for windows.
+
 
 #############################
 # FaceSwap
@@ -37,9 +44,11 @@ pp = PostProcessingOptions(
     restorer_visibility=1,
     upscaler_name="Lanczos",
     scale=4,
-    inpainting_steps=30,
-    inpainting_denoising_strengh=0.1,
     inpainting_when=InpaintingWhen.BEFORE_RESTORE_FACE,
+    inpainting_options=InpaintingOptions(
+        inpainting_steps=30,
+        inpainting_denoising_strengh=0.1,
+    ),
 )
 
 # Prepare the request
@@ -91,3 +100,52 @@ response = FaceSwapExtractResponse.parse_obj(result.json())
 
 for img in response.pil_images:
     img.show()
+
+
+#############################
+# Build checkpoint
+
+source_images: List[str] = [
+    pil_to_base64("../references/man.png"),
+    pil_to_base64("../references/woman.png"),
+]
+
+result = requests.post(
+    url=f"{address}/faceswaplab/build",
+    json=source_images,
+    headers={"Content-Type": "application/json; charset=utf-8"},
+)
+
+base64_to_safetensors(result.json(), output_path="test.safetensors")
+
+#############################
+# FaceSwap with local safetensors
+
+# First face unit :
+unit1 = FaceSwapUnit(
+    source_face=safetensors_to_base64(
+        "test.safetensors"
+    ),  # convert the checkpoint to base64
+    faces_index=(0,),  # Replace first face
+    swapping_options=InswappperOptions(
+        face_restorer_name="CodeFormer",
+        upscaler_name="LDSR",
+        improved_mask=True,
+        sharpen=True,
+        color_corrections=True,
+    ),
+)
+
+# Prepare the request
+request = FaceSwapRequest(image=pil_to_base64("test_image.png"), units=[unit1])
+
+# Face Swap
+result = requests.post(
+    url=f"{address}/faceswaplab/swap_face",
+    data=request.json(),
+    headers={"Content-Type": "application/json; charset=utf-8"},
+)
+response = FaceSwapResponse.parse_obj(result.json())
+
+for img in response.pil_images:
+    img.show()

client_api/requirements.txt

@@ -0,0 +1,5 @@
+numpy
+Pillow
+pydantic
+Requests
+safetensors>=0.3.1

docs/Gemfile

@@ -16,6 +16,7 @@ gem "github-pages", "~> 228", group: :jekyll_plugins
 
 group :jekyll_plugins do
   gem "webrick"
+  gem 'jekyll-toc'
 end
 
 # Windows and JRuby does not include zoneinfo files, so bundle the tzinfo-data gem

docs/Gemfile.lock

@@ -190,6 +190,9 @@ GEM
       jekyll-seo-tag (~> 2.0)
     jekyll-titles-from-headings (0.5.3)
       jekyll (>= 3.3, < 5.0)
+    jekyll-toc (0.18.0)
+      jekyll (>= 3.9)
+      nokogiri (~> 1.12)
     jekyll-watch (2.2.1)
       listen (~> 3.0)
     jemoji (0.12.0)
@@ -256,6 +259,7 @@ DEPENDENCIES
   github-pages (~> 228)
   http_parser.rb (~> 0.6.0)
   jekyll (~> 3.9.3)
+  jekyll-toc
   minima (~> 2.5.1)
   tzinfo (>= 1, < 3)
   tzinfo-data

docs/_config.yml

@@ -37,6 +37,9 @@ author:
 minima:
   skin: dark
 
+plugins:
+  - jekyll-toc
+
 # Exclude from processing.
 # The following items will not be processed, by default.
 # Any item listed under the `exclude:` key here will be automatically added to

docs/_layouts/page.html

@@ -0,0 +1,14 @@
+---
+layout: default
+---
+<article class="post">
+
+    <header class="post-header">
+        <h1 class="post-title">{{ page.title | escape }}</h1>
+    </header>
+
+    <div class="post-content">
+        {{ content | toc }}
+    </div>
+
+</article>

docs/documentation.markdown

@@ -2,9 +2,39 @@
 layout: page
 title: Documentation
 permalink: /doc/
+toc: true
 ---
 
-# Main Interface
+You can also read the [doc discussion section](https://github.com/glucauze/sd-webui-faceswaplab/discussions/categories/guide-doc)
+
+## TLDR: I Just Want Good Results:
+
+1. Put a face in the reference.
+2. Select a face number.
+3. Select "Enable."
+4. Select "CodeFormer" in global Post-Processing.
+
+Once you're happy with some results but want to improve, the next steps are to:
+
++ Use advanced settings in face units (which are not as complex as they might seem, it's basically fine tuning post-processing for each faces).
++ Use pre/post inpainting to tweak the image a bit for more natural results.
+
+### Getting better results
+
+1. Put a face in the reference.
+2. Select a face number.
+3. Select "Enable."
+
+4. In **Post-Processing** accordeon:
+    + Select "CodeFormer" 
+    + Select "LDSR" or a faster model "003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN" in upscaler. See [here for a list of upscalers](https://github.com/glucauze/sd-webui-faceswaplab/discussions/29). 
+    + Use sharpen, color_correction and improved mask
+
+5. Disable "CodeFormer" in **Global Post-Processing** tab (otherwise it will be applied twice)
+
+Don't hesitate to share config in the [discussion section](https://github.com/glucauze/sd-webui-faceswaplab/discussions).
+
+## Main Interface
 
 Here is the interface for FaceSwap Lab. It is available in the form of an accordion in both img2img and txt2img.
 
@@ -12,7 +42,7 @@ You can configure several units, each allowing you to replace a face. Here, 3 un
 
 ![](/assets/images/doc_mi.png)
 
-#### Face Unit
+### Face Unit
 
 The first thing to do is to activate the unit with **'enable'** if you want to use it.
 
@@ -25,6 +55,8 @@ Here are the main options for configuring a unit:
 
 **You must always have at least one reference face OR a checkpoint. If both are selected, the checkpoint will be used and the reference ignored.**
 
+### Similarity
+
 Always check for errors in the SD console. In particular, the absence of a reference face or a checkpoint can trigger errors.
 
 + **Comparison of faces** with the obtained swapped face: The swapped face can be compared to the original face using a distance function. The higher this value (from 1 to 0), the more similar the faces are. This calculation is performed if you activate **"Compute Similarity"** or **"Check Similarity"**. If you check the latter, you will have the opportunity to filter the output images with:
@@ -35,7 +67,43 @@ Always check for errors in the SD console. In particular, the absence of a refer
     + **Same gender:** the gender of the source face will be determined and only faces of the same gender will be considered.
     + **Sort by size:** faces will be sorted from largest to smallest.
 
-#### Post-processing
+### Pre-Inpainting
+
+This part is applied BEFORE face swapping and only on matching faces.
+
+The inpainting part works in the same way as adetailer. It sends each face to img2img for transformation. This is useful for transforming the face before swapping. For example, using a Lora model before swapping.
+
+You can use a specific model for the replacement, different from the model used for the generation.
+
+For inpainting to be active, denoising must be greater than 0 and the Inpainting When option must be set to:
+
+### Post-Processing & Advanced Masks Options : (upscaled inswapper)
+
+By default, these settings are disabled, but you can use the global settings to modify the default behavior. These options are called "Default Upscaled swapper..."
+
+The 'Upscaled Inswapper' is an option in SD FaceSwapLab which allows for upscaling of each face using an upscaller prior to its integration into the image. This is achieved by modifying a small segment of the InsightFace code.
+
+The purpose of this feature is to enhance the quality of the face in the final image. While this process might slightly increase the processing time, it can deliver improved results. In certain cases, this could even eliminate the need for additional tools such as Codeformer or GFPGAN in postprocessing. See the processing order section to understand when and how it is used.
+
+![](/assets/images/upscaled_settings.png)
+
+The upscaled inswapper is disabled by default. It can be enabled in the sd options. Understanding the various steps helps explain why results may be unsatisfactory and how to address this issue.
+
++ **upscaler** : LDSR if None. The LDSR option generally gives the best results but at the expense of a lot of computational time. You should test other models to form an opinion. The [003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN](https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth) model seems to give good results in a reasonable amount of time. It's not possible to disable upscaling, but it is possible to choose LANCZOS for speed if Codeformer is enabled in the upscaled inswapper. The result is generally satisfactory. You can check [here for an upscaler database](https://upscale.wiki/wiki/Model_Database) and [here for some comparison](https://phhofm.github.io/upscale/favorites.html). It is a test and try process.
++ **restorer** : The face restorer to be used if necessary. Codeformer generally gives good results.
++ **sharpening** can provide more natural results, but it may also add artifacts. The same goes for **color correction**. By default, these options are set to False.
++ **improved mask:** The segmentation mask for the upscaled swapper is designed to avoid the square mask and prevent degradation of the non-face parts of the image. It is based on the Codeformer implementation. If "Use improved segmented mask (use pastenet to mask only the face)" and "upscaled inswapper" are checked in the settings, the mask will only cover the face, and will not be squared. However, depending on the image, this might introduce different types of problems such as artifacts on the border of the face.
++ **erosion factor:** it is possible to adjust the mask erosion parameters using the erosion settings. The higher this setting is, the more the mask is reduced.
+
+### Post-Inpainting
+
+This part is applied AFTER face swapping and only on matching faces.
+
+This is useful for adding details to faces. The stronger the denoising, the more likely you are to lose the resemblance of the face. Some samplers (DPM variants for instance) seem to better preserve this resemblance than others.
+
+## Global Post-processing
+
+By default, these settings are disabled, but you can use the global settings to modify the default behavior. These options are called default "UI Default global post processing..."
 
 The post-processing window looks very much like what you might find in the extra options, except for the inpainting part. The process takes place after all units have swapped faces.
 
@@ -82,23 +150,9 @@ The checkpoint can then be used in the main interface (use refresh button)
 ![](/assets/images/checkpoints_use.png)
 
 
-## Upscaled-inswapper
 
-The 'Upscaled Inswapper' is an option in SD FaceSwapLab which allows for upscaling of each face using an upscaller prior to its integration into the image. This is achieved by modifying a small segment of the InsightFace code.
 
-The purpose of this feature is to enhance the quality of the face in the final image. While this process might slightly increase the processing time, it can deliver improved results. In certain cases, this could even eliminate the need for additional tools such as Codeformer or GFPGAN in postprocessing. See the processing order section to understand when and how it is used.
+## Processing order
-
-![](/assets/images/upscaled_settings.png)
-
-The upscaled inswapper is disabled by default. It can be enabled in the sd options. Understanding the various steps helps explain why results may be unsatisfactory and how to address this issue.
-
-+ **upscaler** : LDSR if None. The LDSR option generally gives the best results but at the expense of a lot of computational time. You should test other models to form an opinion. The 003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN model seems to give good results in a reasonable amount of time. It's not possible to disable upscaling, but it is possible to choose LANCZOS for speed if Codeformer is enabled in the upscaled inswapper. The result is generally satisfactory.
-+ **restorer** : The face restorer to be used if necessary. Codeformer generally gives good results.
-+ **sharpening** can provide more natural results, but it may also add artifacts. The same goes for **color correction**. By default, these options are set to False.
-+ **improved mask:** The segmentation mask for the upscaled swapper is designed to avoid the square mask and prevent degradation of the non-face parts of the image. It is based on the Codeformer implementation. If "Use improved segmented mask (use pastenet to mask only the face)" and "upscaled inswapper" are checked in the settings, the mask will only cover the face, and will not be squared. However, depending on the image, this might introduce different types of problems such as artifacts on the border of the face.
-+ **fthresh and erosion factor:** it is possible to adjust the mask erosion parameters using the fthresh and erosion settings. The higher these settings are (particularly erosion), the more the mask is reduced.
-
-## Processing order:
 
 The extension is activated after all other extensions have been processed.  During the execution, several steps take place.
 
@@ -123,42 +177,66 @@ The extension is activated after all other extensions have been processed.  Duri
 
 ![](/assets/images/step4.png)
 
+
+## API
+
+A specific API is available.  To understand how it works you can have a look at the example file in `client_utils`. You can also view the application's tests in the `tests` directory.
+
+The API is documented in the FaceSwapLab tags in the http://localhost:7860/docs docs.
+
+You don't have to use the api_utils.py file and pydantic types, but it can save time.
+
+
+## Experimental GPU support
+
+You need a sufficiently recent version of your SD environment. Using the GPU has a lot of little drawbacks to understand, but the performance gain is substantial.
+
+In Version 1.2.1, the ability to use the GPU has been added, a setting that can be configured in SD at startup. Currently, this feature is only supported on Windows and Linux, as the necessary dependencies for Mac have not been included.
+
+The `--faceswaplab_gpu` option in SD can be added to the args in webui-user.sh or webui-user.bat. **There is also an option in SD settings**.
+
+The model stays loaded in VRAM and won't be unloaded after each use. As of now, I don't know a straightforward way to handle this, so it will occupy space continuously. If your system's VRAM is limited, enabling this option might not be advisable.
+
+A change has also been made that could lead to some ripple effects. Previously, detection parameters such as det_size and det_thresh were automatically adjusted when a second model was loaded. This is no longer possible, so these parameters have been moved to the global settings to enable face detection.
+
+The `auto_det_size` option emulates the old behavior. It has no difference on CPU. BUT it will load the model twice if you use GPU. That means more VRAM comsumption and twice the initial load time. If you don't want that, you can use a det_size of 320, read below.
+
+If you enabled GPU and you are sure you avec a CUDA compatible card and the model keep using CPU provider, please checks that you have onnxruntime-gpu installed.
+
+### SD.NEXT and GPU
+
+Please read carefully.
+
+Using the GPU requires the use of the onnxruntime-gpu>=1.15.0 dependency. For the moment, this conflicts with older SD.Next dependencies (tensorflow, which uses numpy and potentially rembg). You will need to check numpy>=1.24.2 and tensorflow>=2.13.0.
+
+You should therefore be able to debug a little before activating the option. If you don't feel up to it, it's best not to use it.
+
+The first time the swap is used, the program will continue to use the CPU, but will offer to install the GPU. You will then need to restart. This is due to the optimizations made by SD.Next to the installation scripts.
+
+For SD.Next, the best is to install dependencies manually :
+
+on windows :
+
+```shell
+.\venv\Scripts\activate
+cd .\extensions\sd-webui-faceswaplab\
+ pip install .\requirements-gpu.txt
+```
+
 ## Settings
 
-Here are the parameters that can be configured in sd settings and their default values
+You can change the program's default behavior in your webui's global settings (FaceSwapLab section in settings). This is particularly useful if you want to have default options for inpainting or for post-processsing, for example.
 
-### General Settings :
+The interface must be restarted to take the changes into account. Sometimes you have to reboot the entire webui server.
 
- Name | Description | Default Value
+There may be display bugs on some radio buttons that may not display the value (Codeformer might look disabled for instance). Check the logs to ensure that the transformation has been applied.
----|---|---
- faceswaplab_model | Insightface model to use| models[0] if len(models) > 0 else "None"
- faceswaplab_keep_original | keep original image before swapping. It true, will show original image | False
- faceswaplab_units_count | How many faces units to use(requires restart) | 3
- faceswaplab_detection_threshold | Detection threshold to use to detect face, if low will detect non human face as face | 0.5
 
-### Default Settings :
+### det_size and det_thresh (detection accuracy and performances)
 
-These parameters are used to configure the default settings displayed in post-processing.
+V1.2.1 : A change has been made that could lead to some ripple effects. Previously, detection parameters such as det_size and det_thresh were automatically adjusted when a second model was loaded. This is no longer possible, so these parameters have been moved to the global settings to enable face detection.
 
- Name | Description | Default Value
+The `auto_det_size` option emulates the old behavior. It has no difference on CPU. BUT it will load the model twice if you use GPU. That means more VRAM comsumption and twice the initial load time. If you don't want that, you can use a det_size of 320, read below.
- faceswaplab_pp_default_face_restorer | UI Default post processing face restorer (requires restart) | None
- faceswaplab_pp_default_face_restorer_visibility | UI Default post processing face restorer visibility (requires restart) | 1
- faceswaplab_pp_default_face_restorer_weight | UI Default post processing face restorer weight (requires restart) | 1
- faceswaplab_pp_default_upscaler | UI Default post processing upscaler (requires restart) | None
- faceswaplab_pp_default_upscaler_visibility | UI Default post processing upscaler visibility(requires restart) | 1
 
-### Upscaled inswapper Settings :
+The `det_size` parameter defines the size of the detection area, controlling the spatial resolution at which faces are detected within an image. A larger detection size might capture more facial details, enhancing accuracy but potentially impacting processing speed. Conversely, the `det_thresh` parameter represents the detection threshold, serving as a sensitivity control for face detection. A higher threshold value leads to more conservative detection, capturing only the most prominent faces, while a lower threshold might detect more faces but could also result in more false positives.
 
-These parameters are used to control the upscaled inswapper, see above.
+It has been observed that a det_size value of 320 is more effective at detecting large faces. If there are issues with detecting large faces, switching to this value is recommended, though it might result in a loss of some quality.
-
- Name | Description | Default Value
- faceswaplab_upscaled_swapper | Upscaled swapper. Applied only to the swapped faces. Apply transformations before merging with the original image | False
- faceswaplab_upscaled_swapper_upscaler | Upscaled swapper upscaler (Recommended : LDSR but slow) | None
- faceswaplab_upscaled_swapper_sharpen | Upscaled swapper sharpen | False
- faceswaplab_upscaled_swapper_fixcolor | Upscaled swapper color correction | False
- faceswaplab_upscaled_improved_mask | Use improved segmented mask (use pastenet to mask only the face) | True
- faceswaplab_upscaled_swapper_face_restorer | Upscaled swapper face restorer | None
- faceswaplab_upscaled_swapper_face_restorer_visibility | Upscaled swapper face restorer visibility | 1
- faceswaplab_upscaled_swapper_face_restorer_weight | Upscaled swapper face restorer weight (codeformer) | 1
- faceswaplab_upscaled_swapper_fthresh | Upscaled swapper fthresh (diff sensitivity) 10 = default behaviour. Low impact | 10
- faceswaplab_upscaled_swapper_erosion | 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 | 1

docs/faq.markdown

@@ -2,6 +2,7 @@
 layout: page
 title: FAQ
 permalink: /faq/
+toc: true
 ---
 
 Our issue tracker often contains requests that may originate from a misunderstanding of the software's functionality. We aim to address these queries; however, due to time constraints, we may not be able to respond to each request individually. This FAQ section serves as a preliminary source of information for commonly raised concerns. We recommend reviewing these before submitting an issue.
@@ -71,6 +72,16 @@ The quality of results is inherently tied to the capabilities of the model and c
 
 Consider this extension as a low-cost alternative to more sophisticated tools like Lora, or as an addition to such tools. It's important to **maintain realistic expectations of the results** provided by this extension.
 
+#### Why is a face not detected?
+
+Face detection might be influenced by various factors and settings, particularly the det_size and det_thresh parameters. Here's how these could affect detection:
+
++ Detection Size (det_size): If the detection size is set too small, it may not capture large faces adequately. A value of 320 has been found to be more effective for detecting large faces, though it might result in a loss of some quality.
+
++ Detection Threshold (det_thresh): If the threshold is set too high, it can make the detection more conservative, capturing only the most prominent faces. A lower threshold might detect more faces but could also result in more false positives.
+
+If a face is not being detected, adjusting these parameters might solve the issue. Try increasing the det_size if large faces are the problem, or experiment with different det_thresh values to find the balance that works best for your specific case.
+
 
 #### Issue: Incorrect Gender Detection
 
@@ -78,11 +89,7 @@ The gender detection functionality is handled by the underlying analysis model.
 
 #### Why isn't GPU support included?
 
-While implementing GPU support may seem straightforward, simply requiring a modification to the onnxruntime implementation and a change in providers in the swapper, there are reasons we haven't included it as a standard option.
+GPU is supported via an option see [documentation](../doc/). This is expermental, use it carefully.
-
-The primary consideration is the substantial VRAM usage of the SD models. Integrating the model on the GPU doesn't result in significant performance gains with the current state of the software. Moreover, the GPU support becomes truly beneficial when processing large numbers of frames or video. However, our experience indicates that this tends to cause more issues than it resolves.
-
-Consequently, requests for GPU support as a standard feature will not be considered.
 
 #### What is the 'Upscaled Inswapper' Option in SD FaceSwapLab?
 
@@ -112,7 +119,7 @@ A face checkpoint is a saved embedding of a face, generated from multiple images
 
 The primary advantage of face checkpoints is their size. An embedding is only around 2KB, meaning it's lightweight and can be reused later without requiring additional calculations.
 
-Face checkpoints are saved as `.safetensors` files. Please be aware that exchanging `.safetensors` files carries potential security risks. These files, by default, are not secure and could potentially execute malicious code when opened. Therefore, extreme caution should be exercised when sharing or receiving this type of file.
+Face checkpoints are saved as `.safetensors` files.
 
 #### How is similarity determined?
 
@@ -133,3 +140,25 @@ The model generates faces with a resolution of 128x128, which is relatively low.
 SimSwap models are based on older InsightFace architectures, and SimSwap has not been released as a Python package. Its incorporation would complicate the process, and it does not guarantee any substantial gain.
 
 If you manage to implement SimSwap successfully, feel free to submit a pull request.
+
+
+#### Shasum of inswapper model
+
+Check that your model is correct and not corrupted :
+
+```shell
+$>sha1sum inswapper_128.onnx 
+17a64851eaefd55ea597ee41e5c18409754244c5  inswapper_128.onnx
+
+$>sha256sum inswapper_128.onnx 
+e4a3f08c753cb72d04e10aa0f7dbe3deebbf39567d4ead6dce08e98aa49e16af  inswapper_128.onnx
+
+$>sha512sum inswapper_128.onnx 
+4311f4ccd9da58ec544e912b32ac0cba95f5ab4b1a06ac367efd3e157396efbae1097f624f10e77dd811fbba0917fa7c96e73de44563aa6099e5f46830965069  inswapper_128.onnx
+```
+
+#### Gradio errors  (issubclass() arg 1 must be a class)
+
+Older versions of gradio don't work well with the extension. See this bug report : https://github.com/glucauze/sd-webui-faceswaplab/issues/5
+
+It has been tested on 3.32.0 

docs/index.markdown

@@ -20,7 +20,7 @@ While FaceSwapLab is still under development, it has reached a good level of sta
 
 In short:
 
-+ **Ethical Guideline:** This extension should not be forked to create a public, easy way to circumvent NSFW filtering.
++ **Ethical Guideline:** This extension is **not intended to facilitate the creation of not safe for work (NSFW) or non-consensual deepfake content**. Its purpose is to bring consistency to image creation, making it easier to repair existing images, or bring characters back to life.
 + **License:** This software is distributed under the terms of the GNU Affero General Public License (AGPL), version 3 or later.
 + **Model License:** This software uses InsightFace's pre-trained models, which are available for non-commercial research purposes only.
 
@@ -28,8 +28,6 @@ In short:
 
 This extension is **not intended to facilitate the creation of not safe for work (NSFW) or non-consensual deepfake content**. Its purpose is to bring consistency to image creation, making it easier to repair existing images, or bring characters back to life.
 
-While the code for this extension is licensed under the AGPL in compliance with models and other source materials, it's important to stress that **we strongly discourage any attempts to fork this project to create an uncensored version**. Any modifications to the code to enable the production of such content would be contrary to the ethical guidelines we advocate for.
-
 We will comply with European regulations regarding this type of software. As required by law, the code may include both visible and invisible watermarks. If your local laws prohibit the use of this extension, you should not use it.
 
 From an ethical perspective, the main goal of this extension is to generate consistent images by swapping faces. It's important to note that we've done our best to integrate censorship features. However, when users can access the source code, they might bypass these censorship measures. That's why we urge users to use this extension responsibly and avoid any malicious use. We emphasize the importance of respecting people's privacy and consent when swapping faces in images. We discourage any activities that could harm others, invade their privacy, or negatively affect their well-being.

docs/install.markdown

@@ -8,6 +8,8 @@ permalink: /install/
 
 The extension runs mainly on the CPU to avoid the use of VRAM. However, it is recommended to follow the specifications recommended by sd/a1111 with regard to prerequisites. At the time of writing, a version of python lower than 11 is preferable (even if it works with python 3.11, model loading and performance may fall short of expectations).
 
+Older versions of gradio don’t work well with the extension. See this bug report : https://github.com/glucauze/sd-webui-faceswaplab/issues/5. It has been tested on 3.32.0
+
 ### Windows-User : Visual Studio ! Don't neglect this !
 
 Before beginning the installation process, if you are using Windows, you need to install this requirement:
@@ -18,6 +20,12 @@ Before beginning the installation process, if you are using Windows, you need to
 
 3. OR if you don't want to install either the full Visual Studio suite or the VS C++ Build Tools: Follow the instructions provided in section VIII of the documentation.
 
+## SD.Next / Vladmantic
+
+SD.Next loading optimizations in relation to extension installation scripts can sometimes cause problems. This is particularly the case if you copy the script without installing it via the interface.
+
+If you get an error after startup, try restarting the server.
+
 ## Manual Install
 
 To install the extension, follow the steps below:

install.py

@@ -1,66 +1,72 @@
 import launch
 import os
-import pkg_resources
 import sys
-from tqdm import tqdm
+import pkg_resources
-import urllib.request
+from packaging.version import parse
-
-req_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "requirements.txt")
-
-models_dir = os.path.abspath("models/faceswaplab")
-faces_dir = os.path.abspath(os.path.join("models", "faceswaplab", "faces"))
-model_url = "https://huggingface.co/henryruhs/roop/resolve/main/inswapper_128.onnx"
-model_name = os.path.basename(model_url)
-model_path = os.path.join(models_dir, model_name)
 
 
-def download(url: str, path: str) -> None:
+def check_install() -> None:
-    request = urllib.request.urlopen(url)
+    use_gpu = True
-    total = int(request.headers.get("Content-Length", 0))
+
-    with tqdm(
+    if use_gpu and sys.platform != "darwin":
-        total=total, desc="Downloading", unit="B", unit_scale=True, unit_divisor=1024
+        print("Faceswaplab : Use GPU requirements")
-    ) as progress:
+        req_file = os.path.join(
-        urllib.request.urlretrieve(
+            os.path.dirname(os.path.realpath(__file__)), "requirements-gpu.txt"
-            url,
+        )
-            path,
+    else:
-            reporthook=lambda count, block_size, total_size: progress.update(
+        print("Faceswaplab : Use CPU requirements")
-                block_size
+        req_file = os.path.join(
-            ),
+            os.path.dirname(os.path.realpath(__file__)), "requirements.txt"
         )
 
-
+    def is_installed(package: str) -> bool:
-os.makedirs(models_dir, exist_ok=True)
+        package_name = package.split("==")[0].split(">=")[0].strip()
-os.makedirs(faces_dir, exist_ok=True)
-
-if not os.path.exists(model_path):
-    download(model_url, model_path)
-
-
-print("Checking faceswaplab requirements")
-with open(req_file) as file:
-    for package in file:
         try:
-            python = sys.executable
+            installed_version = parse(
-            package = package.strip()
+                pkg_resources.get_distribution(package_name).version
+            )
+        except pkg_resources.DistributionNotFound:
+            return False
 
-            if not launch.is_installed(package.split("==")[0]):
+        if "==" in package:
-                print(f"Install {package}")
+            required_version = parse(package.split("==")[1])
-                launch.run_pip(
+            return installed_version == required_version
-                    f"install {package}", f"sd-webui-faceswaplab requirement: {package}"
+        elif ">=" in package:
-                )
+            required_version = parse(package.split(">=")[1])
-            elif "==" in package:
+            return installed_version >= required_version
-                package_name, package_version = package.split("==")
+        else:
-                installed_version = pkg_resources.get_distribution(package_name).version
+            if package_name == "opencv-python":
-                if installed_version != package_version:
+                return launch.is_installed(package_name) or launch.is_installed("cv2")
-                    print(
+            return launch.is_installed(package_name)
-                        f"Install {package}, {installed_version} vs {package_version}"
+
-                    )
+    print("Checking faceswaplab requirements")
+    with open(req_file) as file:
+        for package in file:
+            try:
+                package = package.strip()
+
+                if not is_installed(package):
+                    print(f"Install {package}")
                     launch.run_pip(
                         f"install {package}",
-                        f"sd-webui-faceswaplab requirement: changing {package_name} version from {installed_version} to {package_version}",
+                        f"sd-webui-faceswaplab requirement: {package}",
                     )
 
-        except Exception as e:
+            except Exception as e:
-            print(e)
+                print(e)
-            print(f"Warning: Failed to install {package}, faceswaplab will not work.")
+                print(
-            raise e
+                    f"Warning: Failed to install {package}, faceswaplab may not work. Try to restart server or install dependencies manually."
+                )
+                raise e
+
+
+import timeit
+
+try:
+    check_time = timeit.timeit(check_install, number=1)
+    print(check_time)
+except Exception as e:
+    print("FaceswapLab install failed", e)
+    print(
+        "You can try to install dependencies manually by activating venv and installing requirements.txt or requirements-gpu.txt"
+    )

javascript/faceswap.js

@@ -1,4 +0,0 @@
-window.onbeforeunload = function() {
-    // Prevent the stable diffusion window from being closed by mistake
-    return "Are you sure ?";
-};

models.json

@@ -0,0 +1 @@
+[{"analyzerName":"intellisense-members-lstm-pylance","languageName":"python","identity":{"modelId":"E61945A9A512ED5E1A3EE3F1A2365B88F8FE","outputId":"E4E9EADA96734F01970E616FAB2FAC19","modifiedTimeUtc":"2020-08-11T14:06:50.811Z"},"filePath":"E61945A9A512ED5E1A3EE3F1A2365B88F8FE_E4E9EADA96734F01970E616FAB2FAC19","lastAccessTimeUtc":"2023-08-14T21:58:14.988Z"}]

requirements-gpu.txt

@@ -0,0 +1,11 @@
+cython
+ifnude
+insightface==0.7.3
+onnx>=1.14.0
+protobuf>=3.20.2
+opencv-python
+pandas
+pydantic
+safetensors
+onnxruntime>=1.15.0
+onnxruntime-gpu>=1.15.0

requirements.txt

@@ -1,9 +1,10 @@
+protobuf>=3.20.2
 cython
 ifnude
 insightface==0.7.3
-onnx==1.14.0
+onnx>=1.14.0
-onnxruntime==1.15.0
+onnxruntime>=1.15.0
-opencv-python==4.7.0.72
+opencv-python
 pandas
-pydantic==1.10.9
+pydantic
-dill==0.3.6
+safetensors

scripts/configure.py

@@ -0,0 +1,65 @@
+import os
+from tqdm import tqdm
+import urllib.request
+from scripts.faceswaplab_utils.faceswaplab_logging import logger
+from scripts.faceswaplab_globals import *
+from packaging import version
+import pkg_resources
+from scripts.faceswaplab_utils.models_utils import check_model
+
+ALREADY_DONE = False
+
+
+def check_configuration() -> None:
+    global ALREADY_DONE
+
+    if ALREADY_DONE:
+        return
+
+    # This has been moved here due to pb with sdnext in install.py not doing what a1111 is doing.
+    models_dir = MODELS_DIR
+    faces_dir = FACES_DIR
+
+    model_url = "https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx"
+    model_name = os.path.basename(model_url)
+    model_path = os.path.join(models_dir, model_name)
+
+    def download(url: str, path: str) -> None:
+        try:
+            request = urllib.request.urlopen(url)
+            total = int(request.headers.get("Content-Length", 0))
+            with tqdm(
+                total=total,
+                desc="Downloading inswapper model",
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+            ) as progress:
+                urllib.request.urlretrieve(
+                    url,
+                    path,
+                    reporthook=lambda count, block_size, total_size: progress.update(
+                        block_size
+                    ),
+                )
+        except:
+            logger.error(
+                "Failed to download inswapper_128.onnx model, please download it manually and put it in the (<sdwebui>/models/faceswaplab/inswapper_128.onnx) directory"
+            )
+
+    os.makedirs(models_dir, exist_ok=True)
+    os.makedirs(faces_dir, exist_ok=True)
+
+    if not os.path.exists(model_path):
+        download(model_url, model_path)
+        check_model()
+
+    gradio_version = pkg_resources.get_distribution("gradio").version
+
+    if version.parse(gradio_version) < version.parse("3.32.0"):
+        logger.warning(
+            "Errors may occur with gradio versions lower than 3.32.0. Your version : %s",
+            gradio_version,
+        )
+
+    ALREADY_DONE = True

scripts/faceswaplab.py

@@ -1,38 +1,45 @@
+from scripts.configure import check_configuration
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
+
+check_configuration()
+
 import importlib
-from scripts.faceswaplab_api import faceswaplab_api
+import traceback
-from scripts.faceswaplab_settings import faceswaplab_settings
-from scripts.faceswaplab_ui import faceswaplab_tab, faceswaplab_unit_ui
-from scripts.faceswaplab_utils.models_utils import (
-    get_current_model,
-)
 
 from scripts import faceswaplab_globals
-from scripts.faceswaplab_swapping import swapper
+from scripts.faceswaplab_api import faceswaplab_api
-from scripts.faceswaplab_utils import faceswaplab_logging, imgutils
-from scripts.faceswaplab_utils import models_utils
 from scripts.faceswaplab_postprocessing import upscaling
+from scripts.faceswaplab_settings import faceswaplab_settings
+from scripts.faceswaplab_swapping import swapper
+from scripts.faceswaplab_ui import faceswaplab_tab, faceswaplab_unit_ui
+from scripts.faceswaplab_utils import faceswaplab_logging, imgutils, models_utils
+from scripts.faceswaplab_utils.models_utils import get_current_swap_model
+from scripts.faceswaplab_utils.typing import *
+from scripts.faceswaplab_utils.ui_utils import dataclasses_from_flat_list
+from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug
 
 # Reload all the modules when using "apply and restart"
 # This is mainly done for development purposes
-importlib.reload(swapper)
+import logging
-importlib.reload(faceswaplab_logging)
+
-importlib.reload(faceswaplab_globals)
+if logger.getEffectiveLevel() <= logging.DEBUG:
-importlib.reload(imgutils)
+    importlib.reload(swapper)
-importlib.reload(upscaling)
+    importlib.reload(faceswaplab_logging)
-importlib.reload(faceswaplab_settings)
+    importlib.reload(faceswaplab_globals)
-importlib.reload(models_utils)
+    importlib.reload(imgutils)
-importlib.reload(faceswaplab_unit_ui)
+    importlib.reload(upscaling)
-importlib.reload(faceswaplab_api)
+    importlib.reload(faceswaplab_settings)
+    importlib.reload(models_utils)
+    importlib.reload(faceswaplab_unit_ui)
+    importlib.reload(faceswaplab_api)
 
 import os
-from dataclasses import fields
 from pprint import pformat
 from typing import Any, List, Optional, Tuple
 
 import gradio as gr
 import modules.scripts as scripts
-from modules import script_callbacks, scripts
+from modules import script_callbacks, scripts, shared
-from modules import scripts, shared
 from modules.images import save_image
 from modules.processing import (
     Processed,
@@ -40,17 +47,14 @@ from modules.processing import (
     StableDiffusionProcessingImg2Img,
 )
 from modules.shared import opts
-from PIL import Image
 
-from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug
 from scripts.faceswaplab_globals import VERSION_FLAG
+from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
 from scripts.faceswaplab_postprocessing.postprocessing_options import (
     PostProcessingOptions,
 )
-from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
 from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
 
-
 EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab")
 
 
@@ -62,25 +66,18 @@ try:
 
     script_callbacks.on_app_started(faceswaplab_api.faceswaplab_api)
 except:
-    pass
+    logger.error("Failed to register API")
+
+    traceback.print_exc()
 
 
 class FaceSwapScript(scripts.Script):
     def __init__(self) -> None:
-        logger.info(f"FaceSwapLab {VERSION_FLAG}")
         super().__init__()
 
     @property
     def units_count(self) -> int:
-        return opts.data.get("faceswaplab_units_count", 3)
+        return get_sd_option("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:
@@ -89,7 +86,7 @@ class FaceSwapScript(scripts.Script):
 
     @property
     def keep_original_images(self) -> bool:
-        return opts.data.get("faceswaplab_keep_original", False)
+        return get_sd_option("faceswaplab_keep_original", False)
 
     @property
     def swap_in_generated_units(self) -> List[FaceSwapUnitSettings]:
@@ -103,48 +100,46 @@ class FaceSwapScript(scripts.Script):
         return f"faceswaplab"
 
     def show(self, is_img2img: bool) -> bool:
-        return scripts.AlwaysVisible
+        return scripts.AlwaysVisible  # type: ignore
 
     def ui(self, is_img2img: bool) -> List[gr.components.Component]:
         with gr.Accordion(f"FaceSwapLab {VERSION_FLAG}", open=False):
-            components = []
+            components: List[gr.components.Component] = []
             for i in range(1, self.units_count + 1):
                 components += faceswaplab_unit_ui.faceswap_unit_ui(is_img2img, i)
-            upscaler = faceswaplab_tab.postprocessing_ui()
+            post_processing = faceswaplab_tab.postprocessing_ui()
         # If the order is modified, the before_process should be changed accordingly.
-        return components + upscaler
+
+        components = components + post_processing
+
+        return components
 
     def read_config(
-        self, p: StableDiffusionProcessing, *components: List[gr.components.Component]
+        self, p: StableDiffusionProcessing, *components: Tuple[Any, ...]
     ) -> None:
+        for i, c in enumerate(components):
+            logger.debug("%s>%s", i, pformat(c))
+
         # The order of processing for the components is important
         # The method first process faceswap units then postprocessing units
-
+        classes: List[Any] = dataclasses_from_flat_list(
-        # self.make_first_script(p)
+            [FaceSwapUnitSettings] * self.units_count + [PostProcessingOptions],
-
+            components,
+        )
         self.units: List[FaceSwapUnitSettings] = []
-
+        self.units += [u for u in classes if isinstance(u, FaceSwapUnitSettings)]
-        # Parse and convert units flat components into FaceSwapUnitSettings
+        self.postprocess_options = classes[-1]
-        for i in range(0, self.units_count):
-            self.units += [FaceSwapUnitSettings.get_unit_configuration(i, components)]
 
         for i, u in enumerate(self.units):
             logger.debug("%s, %s", pformat(i), pformat(u))
 
-        # Parse the postprocessing options
-        # We must first find where to start from (after face swapping units)
-        len_conf: int = len(fields(FaceSwapUnitSettings))
-        shift: int = self.units_count * len_conf
-        self.postprocess_options = PostProcessingOptions(
-            *components[shift : shift + len(fields(PostProcessingOptions))]  # type: ignore
-        )
         logger.debug("%s", pformat(self.postprocess_options))
 
         if self.enabled:
             p.do_not_save_samples = not self.keep_original_images
 
     def process(
-        self, p: StableDiffusionProcessing, *components: List[gr.components.Component]
+        self, p: StableDiffusionProcessing, *components: Tuple[Any, ...]
     ) -> None:
         try:
             self.read_config(p, *components)
@@ -152,14 +147,13 @@ class FaceSwapScript(scripts.Script):
             # If is instance of img2img, we check if face swapping in source is required.
             if isinstance(p, StableDiffusionProcessingImg2Img):
                 if self.enabled and len(self.swap_in_source_units) > 0:
-                    init_images: List[Tuple[Optional[Image.Image], Optional[str]]] = [
+                    init_images: List[Tuple[Optional[PILImage], Optional[str]]] = [
                         (img, None) for img in p.init_images
                     ]
                     new_inits = swapper.process_images_units(
-                        get_current_model(),
+                        get_current_swap_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)}")
@@ -167,6 +161,7 @@ class FaceSwapScript(scripts.Script):
                         p.init_images = [img[0] for img in new_inits]
         except Exception as e:
             logger.info("Failed to process : %s", e)
+            traceback.print_exc()
 
     def postprocess(
         self, p: StableDiffusionProcessing, processed: Processed, *args: List[Any]
@@ -174,7 +169,7 @@ class FaceSwapScript(scripts.Script):
         try:
             if self.enabled:
                 # Get the original images without the grid
-                orig_images: List[Image.Image] = processed.images[
+                orig_images: List[PILImage] = processed.images[
                     processed.index_of_first_image :
                 ]
                 orig_infotexts: List[str] = processed.infotexts[
@@ -190,10 +185,9 @@ class FaceSwapScript(scripts.Script):
                     for i, (img, info) in enumerate(zip(orig_images, orig_infotexts)):
                         batch_index = i % p.batch_size
                         swapped_images = swapper.process_images_units(
-                            get_current_model(),
+                            get_current_swap_model(),
                             self.swap_in_generated_units,
                             images=[(img, info)],
-                            upscaled_swapper=self.upscaled_swapper_in_generated,
                         )
                         if swapped_images is None:
                             continue
@@ -223,8 +217,8 @@ class FaceSwapScript(scripts.Script):
                                         swp_img,
                                         p.outpath_samples,
                                         "",
-                                        p.all_seeds[batch_index],
+                                        p.all_seeds[batch_index],  # type: ignore
-                                        p.all_prompts[batch_index],
+                                        p.all_prompts[batch_index],  # type: ignore
                                         opts.samples_format,
                                         info=new_info,
                                         p=p,
@@ -237,14 +231,27 @@ class FaceSwapScript(scripts.Script):
 
                 # Generate grid :
                 if opts.return_grid and len(images) > 1:
-                    # FIXME :Use sd method, not that if blended is not active, the result will be a bit messy.
                     grid = imgutils.create_square_image(images)
                     text = processed.infotexts[0]
                     infotexts.insert(0, text)
                     if opts.enable_pnginfo:
-                        grid.info["parameters"] = text
+                        grid.info["parameters"] = text  # type: ignore
                     images.insert(0, grid)
 
+                    if opts.grid_save:
+                        save_image(
+                            grid,
+                            p.outpath_grids,
+                            "swapped-grid",
+                            p.all_seeds[0],  # type: ignore
+                            p.all_prompts[0],  # type: ignore
+                            opts.grid_format,
+                            info=text,
+                            short_filename=not opts.grid_extended_filename,
+                            p=p,
+                            grid=True,
+                        )
+
                 if keep_original:
                     # If we want to keep original images, we add all existing (including grid this time)
                     images += processed.images
@@ -254,3 +261,4 @@ class FaceSwapScript(scripts.Script):
                 processed.infotexts = infotexts
         except Exception as e:
             logger.error("Failed to swap face in postprocess method : %s", e)
+            traceback.print_exc()

scripts/faceswaplab_api/faceswaplab_api.py

@@ -1,3 +1,4 @@
+import tempfile
 from PIL import Image
 import numpy as np
 from fastapi import FastAPI
@@ -17,7 +18,10 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
     PostProcessingOptions,
 )
 from client_api import api_utils
-from scripts.faceswaplab_postprocessing.postprocessing_options import InpaintingWhen
+from scripts.faceswaplab_swapping.face_checkpoints import (
+    build_face_checkpoint_and_save,
+)
+from scripts.faceswaplab_utils.typing import PILImage
 
 
 def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str:  # type: ignore
@@ -58,58 +62,12 @@ def encode_np_to_base64(image: np.ndarray) -> str:  # type: ignore
     return api.encode_pil_to_base64(pil)
 
 
-def get_postprocessing_options(
-    options: api_utils.PostProcessingOptions,
-) -> PostProcessingOptions:
-    pp_options = PostProcessingOptions(
-        face_restorer_name=options.face_restorer_name,
-        restorer_visibility=options.restorer_visibility,
-        codeformer_weight=options.codeformer_weight,
-        upscaler_name=options.upscaler_name,
-        scale=options.scale,
-        upscale_visibility=options.upscaler_visibility,
-        inpainting_denoising_strengh=options.inpainting_denoising_strengh,
-        inpainting_prompt=options.inpainting_prompt,
-        inpainting_negative_prompt=options.inpainting_negative_prompt,
-        inpainting_steps=options.inpainting_steps,
-        inpainting_sampler=options.inpainting_sampler,
-        # hacky way to prevent having a separate file for Inpainting when (2 classes)
-        # therfore a conversion is required from api IW to server side IW
-        inpainting_when=InpaintingWhen(options.inpainting_when.value),
-        inpainting_model=options.inpainting_model,
-    )
-
-    assert isinstance(
-        pp_options.inpainting_when, InpaintingWhen
-    ), "Value is not a valid InpaintingWhen enum"
-
-    return pp_options
-
-
 def get_faceswap_units_settings(
     api_units: List[api_utils.FaceSwapUnit],
 ) -> List[FaceSwapUnitSettings]:
     units = []
     for u in api_units:
-        units.append(
+        units.append(FaceSwapUnitSettings.from_api_dto(u))
-            FaceSwapUnitSettings(
-                source_img=base64_to_pil(u.source_img),
-                source_face=u.source_face,
-                _batch_files=u.get_batch_images(),
-                blend_faces=u.blend_faces,
-                enable=True,
-                same_gender=u.same_gender,
-                sort_by_size=u.sort_by_size,
-                check_similarity=u.check_similarity,
-                _compute_similarity=u.compute_similarity,
-                min_ref_sim=u.min_ref_sim,
-                min_sim=u.min_sim,
-                _faces_index=",".join([str(i) for i in (u.faces_index)]),
-                reference_face_index=u.reference_face_index,
-                swap_in_generated=True,
-                swap_in_source=False,
-            )
-        )
     return units
 
 
@@ -137,10 +95,12 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None:
 
         if src_image is not None:
             if request.postprocessing:
-                pp_options = get_postprocessing_options(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(
+            swapped_images: Optional[List[PILImage]] = swapper.batch_process(
                 [src_image], None, units=units, postprocess_options=pp_options
             )
 
@@ -172,7 +132,7 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None:
     ) -> api_utils.FaceSwapExtractResponse:
         pp_options = None
         if request.postprocessing:
-            pp_options = get_postprocessing_options(request.postprocessing)
+            pp_options = PostProcessingOptions.from_api_dto(request.postprocessing)
         images = [base64_to_pil(img) for img in request.images]
         faces = swapper.extract_faces(
             images, extract_path=None, postprocess_options=pp_options
@@ -180,3 +140,23 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None:
         result_images = [encode_to_base64(img) for img in faces]
         response = api_utils.FaceSwapExtractResponse(images=result_images)
         return response
+
+    @app.post(
+        "/faceswaplab/build",
+        tags=["faceswaplab"],
+        description="Build a face checkpoint using base64 images, return base64 satetensors",
+    )
+    async def build(base64_images: List[str]) -> Optional[str]:
+        if len(base64_images) > 0:
+            pil_images = [base64_to_pil(img) for img in base64_images]
+            with tempfile.NamedTemporaryFile(
+                delete=True, suffix=".safetensors"
+            ) as temp_file:
+                build_face_checkpoint_and_save(
+                    images=pil_images,
+                    name="api_ckpt",
+                    overwrite=True,
+                    path=temp_file.name,
+                )
+                return api_utils.safetensors_to_base64(temp_file.name)
+        return None

scripts/faceswaplab_globals.py

@@ -1,15 +1,24 @@
 import os
 from modules import scripts
 
+# Defining the absolute path for the 'faceswaplab' directory inside 'models' directory
 MODELS_DIR = os.path.abspath(os.path.join("models", "faceswaplab"))
+# Defining the absolute path for the 'analysers' directory inside 'MODELS_DIR'
 ANALYZER_DIR = os.path.abspath(os.path.join(MODELS_DIR, "analysers"))
+# Defining the absolute path for the 'parser' directory inside 'MODELS_DIR'
 FACE_PARSER_DIR = os.path.abspath(os.path.join(MODELS_DIR, "parser"))
+# Defining the absolute path for the 'faces' directory inside 'MODELS_DIR'
+FACES_DIR = os.path.abspath(os.path.join(MODELS_DIR, "faces"))
+
+# Constructing the path for 'references' directory inside the 'extensions' and 'sd-webui-faceswaplab' directories, based on the base directory of scripts
 REFERENCE_PATH = os.path.join(
     scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
 )
 
-VERSION_FLAG: str = "v1.1.2"
+# Defining the version flag for the application
+VERSION_FLAG: str = "v1.2.7"
+# Defining the path for 'sd-webui-faceswaplab' inside the 'extensions' directory
 EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab")
 
-# The NSFW score threshold. If any part of the image has a score greater than this threshold, the image will be considered NSFW.
+# Defining the expected SHA1 hash value for 'INSWAPPER'
-NSFW_SCORE_THRESHOLD: float = 0.7
+EXPECTED_INSWAPPER_SHA1 = "17a64851eaefd55ea597ee41e5c18409754244c5"

scripts/faceswaplab_inpainting/faceswaplab_inpainting.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+from typing import List, Optional
+import gradio as gr
+from client_api import api_utils
+
+
+@dataclass
+class InpaintingOptions:
+    inpainting_denoising_strengh: float = 0
+    inpainting_prompt: str = ""
+    inpainting_negative_prompt: str = ""
+    inpainting_steps: int = 20
+    inpainting_sampler: str = "Euler"
+    inpainting_model: str = "Current"
+    inpainting_seed: int = -1
+
+    @staticmethod
+    def from_gradio(components: List[gr.components.Component]) -> "InpaintingOptions":
+        return InpaintingOptions(*components)  # type: ignore
+
+    @staticmethod
+    def from_api_dto(dto: Optional[api_utils.InpaintingOptions]) -> "InpaintingOptions":
+        """
+        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 default values
+            return InpaintingOptions()
+
+        return InpaintingOptions(
+            inpainting_denoising_strengh=dto.inpainting_denoising_strengh,
+            inpainting_prompt=dto.inpainting_prompt,
+            inpainting_negative_prompt=dto.inpainting_negative_prompt,
+            inpainting_steps=dto.inpainting_steps,
+            inpainting_sampler=dto.inpainting_sampler,
+            inpainting_model=dto.inpainting_model,
+            inpainting_seed=dto.inpainting_seed,
+        )

scripts/faceswaplab_inpainting/i2i_pp.py

@@ -1,53 +1,61 @@
+from scripts.faceswaplab_inpainting.faceswaplab_inpainting import InpaintingOptions
 from scripts.faceswaplab_utils.faceswaplab_logging import logger
 from PIL import Image
 from modules import shared
 from scripts.faceswaplab_utils import imgutils
 from modules import shared, processing
 from modules.processing import StableDiffusionProcessingImg2Img
-from scripts.faceswaplab_postprocessing.postprocessing_options import (
-    PostProcessingOptions,
-)
 from modules import sd_models
-
+import traceback
 from scripts.faceswaplab_swapping import swapper
+from scripts.faceswaplab_utils.typing import *
+from typing import *
 
 
-def img2img_diffusion(img: Image.Image, pp: PostProcessingOptions) -> Image.Image:
+def img2img_diffusion(
-    if pp.inpainting_denoising_strengh == 0:
+    img: PILImage, options: InpaintingOptions, faces: Optional[List[Face]] = None
-        logger.info("Discard inpainting denoising strength is 0")
+) -> Image.Image:
+    if not options or options.inpainting_denoising_strengh == 0:
+        logger.info("Discard inpainting denoising strength is 0 or no inpainting")
         return img
 
     try:
         logger.info(
             f"""Inpainting face
-Sampler : {pp.inpainting_sampler}
+Sampler : {options.inpainting_sampler}
-inpainting_denoising_strength : {pp.inpainting_denoising_strengh}
+inpainting_denoising_strength : {options.inpainting_denoising_strengh}
-inpainting_steps : {pp.inpainting_steps}
+inpainting_steps : {options.inpainting_steps}
 """
         )
-        if not isinstance(pp.inpainting_sampler, str):
+        if not isinstance(options.inpainting_sampler, str):
-            pp.inpainting_sampler = "Euler"
+            options.inpainting_sampler = "Euler"
 
         logger.info("send faces to image to image")
         img = img.copy()
-        faces = swapper.get_faces(imgutils.pil_to_cv2(img))
+
+        if not faces:
+            faces = swapper.get_faces(imgutils.pil_to_cv2(img))
+
         if faces:
             for face in faces:
                 bbox = face.bbox.astype(int)
                 mask = imgutils.create_mask(img, bbox)
-                prompt = pp.inpainting_prompt.replace(
+                prompt = options.inpainting_prompt.replace(
                     "[gender]", "man" if face["gender"] == 1 else "woman"
                 )
-                negative_prompt = pp.inpainting_negative_prompt.replace(
+                negative_prompt = options.inpainting_negative_prompt.replace(
                     "[gender]", "man" if face["gender"] == 1 else "woman"
                 )
                 logger.info("Denoising prompt : %s", prompt)
-                logger.info("Denoising strenght : %s", pp.inpainting_denoising_strengh)
+                logger.info(
+                    "Denoising strenght : %s", options.inpainting_denoising_strengh
+                )
 
                 i2i_kwargs = {
-                    "sampler_name": pp.inpainting_sampler,
+                    "init_images": [img],
+                    "sampler_name": options.inpainting_sampler,
                     "do_not_save_samples": True,
-                    "steps": pp.inpainting_steps,
+                    "steps": options.inpainting_steps,
                     "width": img.width,
                     "inpainting_fill": 1,
                     "inpaint_full_res": True,
@@ -55,17 +63,19 @@ inpainting_steps : {pp.inpainting_steps}
                     "mask": mask,
                     "prompt": prompt,
                     "negative_prompt": negative_prompt,
-                    "denoising_strength": pp.inpainting_denoising_strengh,
+                    "denoising_strength": options.inpainting_denoising_strengh,
+                    "seed": options.inpainting_seed,
                 }
+
                 current_model_checkpoint = shared.opts.sd_model_checkpoint
-                if pp.inpainting_model and pp.inpainting_model != "Current":
+                if options.inpainting_model and options.inpainting_model != "Current":
                     # Change checkpoint
-                    shared.opts.sd_model_checkpoint = pp.inpainting_model
+                    shared.opts.sd_model_checkpoint = options.inpainting_model
                     sd_models.select_checkpoint
                     sd_models.load_model()
-                i2i_p = StableDiffusionProcessingImg2Img([img], **i2i_kwargs)
+                i2i_p = StableDiffusionProcessingImg2Img(**i2i_kwargs)
                 i2i_processed = processing.process_images(i2i_p)
-                if pp.inpainting_model and pp.inpainting_model != "Current":
+                if options.inpainting_model and options.inpainting_model != "Current":
                     # Restore checkpoint
                     shared.opts.sd_model_checkpoint = current_model_checkpoint
                     sd_models.select_checkpoint
@@ -76,8 +86,6 @@ inpainting_steps : {pp.inpainting_steps}
                     img = images[0]
         return img
     except Exception as e:
-        logger.error("Failed to apply img2img to face : %s", e)
+        logger.error("Failed to apply inpainting to face : %s", e)
-        import traceback
-
         traceback.print_exc()
         raise e

scripts/faceswaplab_postprocessing/postprocessing.py

@@ -4,8 +4,9 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
     PostProcessingOptions,
     InpaintingWhen,
 )
-from scripts.faceswaplab_postprocessing.i2i_pp import img2img_diffusion
+from scripts.faceswaplab_inpainting.i2i_pp import img2img_diffusion
 from scripts.faceswaplab_postprocessing.upscaling import upscale_img, restore_face
+import traceback
 
 
 def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image:
@@ -19,7 +20,9 @@ def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Imag
             or pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING
         ):
             logger.debug("Inpaint before upscale")
-            result_image = img2img_diffusion(result_image, pp_options)
+            result_image = img2img_diffusion(
+                img=result_image, options=pp_options.inpainting_options
+            )
         result_image = upscale_img(result_image, pp_options)
 
         if (
@@ -27,7 +30,9 @@ def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Imag
             or pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE
         ):
             logger.debug("Inpaint before restore")
-            result_image = img2img_diffusion(result_image, pp_options)
+            result_image = img2img_diffusion(
+                result_image, pp_options.inpainting_options
+            )
 
         result_image = restore_face(result_image, pp_options)
 
@@ -36,9 +41,11 @@ def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Imag
             or pp_options.inpainting_when == InpaintingWhen.AFTER_ALL
         ):
             logger.debug("Inpaint after all")
-            result_image = img2img_diffusion(result_image, pp_options)
+            result_image = img2img_diffusion(
+                result_image, pp_options.inpainting_options
+            )
 
     except Exception as e:
-        logger.error("Failed to upscale %s", e)
+        logger.error("Failed to post-process %s", e)
-
+        traceback.print_exc()
     return result_image

scripts/faceswaplab_postprocessing/postprocessing_options.py

@@ -1,8 +1,11 @@
+from typing import Optional
 from modules.face_restoration import FaceRestoration
 from modules.upscaler import UpscalerData
 from dataclasses import dataclass
 from modules import shared
 from enum import Enum
+from scripts.faceswaplab_inpainting.faceswaplab_inpainting import InpaintingOptions
+from client_api import api_utils
 
 
 class InpaintingWhen(Enum):
@@ -22,24 +25,46 @@ class PostProcessingOptions:
     scale: float = 1
     upscale_visibility: float = 0.5
 
-    inpainting_denoising_strengh: float = 0
-    inpainting_prompt: str = ""
-    inpainting_negative_prompt: str = ""
-    inpainting_steps: int = 20
-    inpainting_sampler: str = "Euler"
     inpainting_when: InpaintingWhen = InpaintingWhen.BEFORE_UPSCALING
-    inpainting_model: str = "Current"
+
+    # (Don't use optional for this or gradio parsing will fail) :
+    inpainting_options: InpaintingOptions = None  # type: ignore
 
     @property
-    def upscaler(self) -> UpscalerData:
+    def upscaler(self) -> Optional[UpscalerData]:
         for upscaler in shared.sd_upscalers:
             if upscaler.name == self.upscaler_name:
                 return upscaler
         return None
 
     @property
-    def face_restorer(self) -> FaceRestoration:
+    def face_restorer(self) -> Optional[FaceRestoration]:
         for face_restorer in shared.face_restorers:
             if face_restorer.name() == self.face_restorer_name:
                 return face_restorer
         return None
+
+    @staticmethod
+    def from_api_dto(
+        options: api_utils.PostProcessingOptions,
+    ) -> "PostProcessingOptions":
+        """
+        Converts a PostProcessingOptions object from an API DTO (Data Transfer Object).
+
+        :param options: An object of api_utils.PostProcessingOptions representing the
+                        post-processing options as received from the API.
+        :return: A PostProcessingOptions instance containing the translated values
+                from the API DTO.
+        """
+        return PostProcessingOptions(
+            face_restorer_name=options.face_restorer_name,
+            restorer_visibility=options.restorer_visibility,
+            codeformer_weight=options.codeformer_weight,
+            upscaler_name=options.upscaler_name,
+            scale=options.scale,
+            upscale_visibility=options.upscaler_visibility,
+            inpainting_when=InpaintingWhen(options.inpainting_when.value),
+            inpainting_options=InpaintingOptions.from_api_dto(
+                options.inpainting_options
+            ),
+        )

scripts/faceswaplab_postprocessing/upscaling.py

@@ -5,20 +5,26 @@ from scripts.faceswaplab_utils.faceswaplab_logging import logger
 from PIL import Image
 import numpy as np
 from modules import codeformer_model
+from scripts.faceswaplab_utils.typing import *
 
 
-def upscale_img(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image:
+def upscale_img(image: PILImage, pp_options: PostProcessingOptions) -> PILImage:
     if pp_options.upscaler is not None and pp_options.upscaler.name != "None":
-        original_image = image.copy()
+        original_image: PILImage = image.copy()
         logger.info(
             "Upscale with %s scale = %s",
             pp_options.upscaler.name,
             pp_options.scale,
         )
         result_image = pp_options.upscaler.scaler.upscale(
-            image, pp_options.scale, pp_options.upscaler.data_path
+            image, pp_options.scale, pp_options.upscaler.data_path  # type: ignore
         )
-        if pp_options.scale == 1:
+
+        # FIXME : Could be better (managing images whose dimensions are not multiples of 16)
+        if pp_options.scale == 1 and original_image.size == result_image.size:
+            logger.debug(
+                "Sizes orig=%s, result=%s", original_image.size, result_image.size
+            )
             result_image = Image.blend(
                 original_image, result_image, pp_options.upscale_visibility
             )

scripts/faceswaplab_settings/faceswaplab_settings.py

@@ -1,11 +1,11 @@
-from scripts.faceswaplab_utils.models_utils import get_models
+from scripts.faceswaplab_utils.models_utils import get_swap_models
 from modules import script_callbacks, shared
 import gradio as gr
 
 
 def on_ui_settings() -> None:
     section = ("faceswaplab", "FaceSwapLab")
-    models = get_models()
+    models = get_swap_models()
     shared.opts.add_option(
         "faceswaplab_model",
         shared.OptionInfo(
@@ -16,6 +16,16 @@ def on_ui_settings() -> None:
             section=section,
         ),
     )
+    shared.opts.add_option(
+        "faceswaplab_use_gpu",
+        shared.OptionInfo(
+            False,
+            "Use GPU, only for CUDA on Windows/Linux - experimental and risky, can messed up dependencies (requires restart)",
+            gr.Checkbox,
+            {"interactive": True},
+            section=section,
+        ),
+    )
     shared.opts.add_option(
         "faceswaplab_keep_original",
         shared.OptionInfo(
@@ -36,12 +46,44 @@ def on_ui_settings() -> None:
             section=section,
         ),
     )
+    shared.opts.add_option(
+        "faceswaplab_nsfw_threshold",
+        shared.OptionInfo(
+            0.7,
+            "NSFW score threshold. Any image part with a score above this value will be treated as NSFW (use extension responsibly !). 1=Disable filtering",
+            gr.Slider,
+            {"minimum": 0, "maximum": 1, "step": 0.01},
+            section=section,
+        ),
+    )
+
+    shared.opts.add_option(
+        "faceswaplab_det_size",
+        shared.OptionInfo(
+            640,
+            "det_size : Size of the detection area for face analysis. Higher values may improve quality but reduce speed. Low value may improve detection of very large face.",
+            gr.Slider,
+            {"minimum": 320, "maximum": 640, "step": 320},
+            section=section,
+        ),
+    )
+
+    shared.opts.add_option(
+        "faceswaplab_auto_det_size",
+        shared.OptionInfo(
+            True,
+            "Auto det_size : Will load model twice and test faces on each if needed (old behaviour). Takes more VRAM. Precedence over fixed det_size",
+            gr.Checkbox,
+            {"interactive": True},
+            section=section,
+        ),
+    )
 
     shared.opts.add_option(
         "faceswaplab_detection_threshold",
         shared.OptionInfo(
             0.5,
-            "Face Detection threshold",
+            "det_thresh : Face Detection threshold",
             gr.Slider,
             {"minimum": 0.1, "maximum": 0.99, "step": 0.001},
             section=section,
@@ -54,7 +96,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 +109,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 +119,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 +129,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 +142,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 +176,10 @@ def on_ui_settings() -> None:
     # UPSCALED SWAPPER
 
     shared.opts.add_option(
-        "faceswaplab_upscaled_swapper",
+        "faceswaplab_default_upscaled_swapper_upscaler",
-        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",
         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 +189,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,
+            False,
-            "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 +232,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",
+        "faceswaplab_default_upscaled_swapper_erosion",
-        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",
         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/face_checkpoints.py

@@ -0,0 +1,250 @@
+import glob
+import os
+from typing import *
+from insightface.app.common import Face
+from safetensors.torch import save_file, safe_open
+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_utils.models_utils import get_swap_models
+import traceback
+
+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:
+    """
+    Sanitize the input name by removing special characters and replacing spaces with underscores.
+
+    Parameters:
+        name (str): The input name to be sanitized.
+
+    Returns:
+        str: The sanitized name with special characters removed and spaces replaced by underscores.
+    """
+    name = re.sub("[^A-Za-z0-9_. ]+", "", name)
+    name = name.replace(" ", "_")
+    return name[:255]
+
+
+def build_face_checkpoint_and_save(
+    images: List[PILImage],
+    name: str,
+    gender: Gender = Gender.AUTO,
+    overwrite: bool = False,
+    path: Optional[str] = None,
+) -> Optional[PILImage]:
+    """
+    Builds a face checkpoint using the provided image files, performs face swapping,
+    and saves the result to a file. If a blended face is successfully obtained and the face swapping
+    process succeeds, the resulting image is returned. Otherwise, None is returned.
+
+    Args:
+        batch_files (list): List of image file paths used to create the face checkpoint.
+        name (str): The name assigned to the face checkpoint.
+
+    Returns:
+        PIL.PILImage or None: The resulting swapped face image if the process is successful; None otherwise.
+    """
+
+    try:
+        name = sanitize_name(name)
+        images = images or []
+        logger.info("Build %s with %s images", name, len(images))
+        faces: List[Face] = swapper.get_faces_from_img_files(images=images)
+        if faces is None or len(faces) == 0:
+            logger.error("No source faces found")
+            return None
+
+        blended_face: Optional[Face] = swapper.blend_faces(faces, gender=gender)
+        preview_path = os.path.join(
+            scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
+        )
+
+        reference_preview_img: PILImage
+        if blended_face:
+            if blended_face["gender"] == 0:
+                reference_preview_img = Image.open(
+                    os.path.join(preview_path, "woman.png")
+                )
+            else:
+                reference_preview_img = Image.open(
+                    os.path.join(preview_path, "man.png")
+                )
+
+            if name == "":
+                name = "default_name"
+            logger.debug("Face %s", pformat(blended_face))
+            target_face = swapper.get_or_default(
+                swapper.get_faces(imgutils.pil_to_cv2(reference_preview_img)), 0, None
+            )
+            if target_face is None:
+                logger.error(
+                    "Failed to open reference image, cannot create preview : That should not happen unless you deleted the references folder or change the detection threshold."
+                )
+            else:
+                result: swapper.ImageResult = swapper.swap_face(
+                    target_faces=[target_face],
+                    source_face=blended_face,
+                    target_img=reference_preview_img,
+                    model=get_swap_models()[0],
+                    swapping_options=InswappperOptions(
+                        face_restorer_name="CodeFormer",
+                        restorer_visibility=1,
+                        upscaler_name="Lanczos",
+                        codeformer_weight=1,
+                        improved_mask=True,
+                        color_corrections=False,
+                        sharpen=True,
+                    ),
+                )
+                preview_image = result.image
+
+                if path:
+                    file_path = path
+                else:
+                    file_path = os.path.join(
+                        get_checkpoint_path(), f"{name}.safetensors"
+                    )
+                    if not overwrite:
+                        file_number = 1
+                        while os.path.exists(file_path):
+                            file_path = os.path.join(
+                                get_checkpoint_path(),
+                                f"{name}_{file_number}.safetensors",
+                            )
+                            file_number += 1
+                save_face(filename=file_path, face=blended_face)
+                preview_image.save(file_path + ".png")
+                try:
+                    data = load_face(file_path)
+                    logger.debug(data)
+                except Exception as e:
+                    logger.error("Error loading checkpoint, after creation %s", e)
+                    traceback.print_exc()
+
+                return preview_image
+
+        else:
+            logger.error("No face found")
+            return None  # type: ignore
+    except Exception as e:
+        logger.error("Failed to build checkpoint %s", e)
+        traceback.print_exc()
+        return None
+
+
+def save_face(face: Face, filename: str) -> None:
+    try:
+        tensors = {
+            "embedding": torch.tensor(face["embedding"]),
+            "gender": torch.tensor(face["gender"]),
+            "age": torch.tensor(face["age"]),
+        }
+        save_file(tensors, filename)
+    except Exception as e:
+        traceback.print_exc
+        logger.error("Failed to save checkpoint %s", e)
+        raise e
+
+
+def load_face(name: str) -> Optional[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
+
+    if filename.endswith(".pkl"):
+        logger.warning(
+            "Pkl files for faces are deprecated to enhance safety, you need to convert them"
+        )
+        logger.warning("The file will be converted to .safetensors")
+        logger.warning(
+            "You can also use this script https://gist.github.com/glucauze/4a3c458541f2278ad801f6625e5b9d3d"
+        )
+        return None
+
+    elif filename.endswith(".safetensors"):
+        face = {}
+        with safe_open(filename, 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)
+
+    raise NotImplementedError("Unknown file type, face extraction not implemented")
+
+
+def get_checkpoint_path() -> str:
+    checkpoint_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces")
+    os.makedirs(checkpoint_path, exist_ok=True)
+    return checkpoint_path
+
+
+def matching_checkpoint(name: str) -> Optional[str]:
+    """
+    Retrieve the full path of a checkpoint file matching the given name.
+
+    If the name already includes a path separator, it is returned as-is. Otherwise, the function looks for a matching
+    file with the extensions ".safetensors" or ".pkl" in the checkpoint directory.
+
+    Args:
+        name (str): The name or path of the checkpoint file.
+
+    Returns:
+        Optional[str]: The full path of the matching checkpoint file, or None if no match is found.
+    """
+
+    # If the name already includes a path separator, return it as is
+    if os.path.sep in name:
+        return name
+
+    # If the name doesn't end with the specified extensions, look for a matching file
+    if not (name.endswith(".safetensors") or name.endswith(".pkl")):
+        # Try appending each extension and check if the file exists in the checkpoint path
+        for ext in [".safetensors", ".pkl"]:
+            full_path = os.path.join(get_checkpoint_path(), name + ext)
+            if os.path.exists(full_path):
+                return full_path
+        # If no matching file is found, return None
+        return None
+
+    # If the name already ends with the specified extensions, simply complete the path
+    return os.path.join(get_checkpoint_path(), name)
+
+
+def get_face_checkpoints() -> List[str]:
+    """
+    Retrieve a list of face checkpoint paths.
+
+    This function searches for face files with the extension ".safetensors" in the specified directory and returns a list
+    containing the paths of those files.
+
+    Returns:
+        list: A list of face paths, including the string "None" as the first element.
+    """
+    faces_path = os.path.join(get_checkpoint_path(), "*.safetensors")
+    faces = glob.glob(faces_path)
+
+    faces_path = os.path.join(get_checkpoint_path(), "*.pkl")
+    faces += glob.glob(faces_path)
+
+    return ["None"] + [os.path.basename(face) for face in sorted(faces)]

scripts/faceswaplab_swapping/facemask.py

@@ -20,7 +20,7 @@ def get_parsing_model(device: torch_device) -> torch.nn.Module:
     Returns:
         The parsing model.
     """
-    return init_parsing_model(device=device)
+    return init_parsing_model(device=device)  # type: ignore
 
 
 def convert_image_to_tensor(
@@ -83,7 +83,7 @@ def generate_face_mask(face_image: np.ndarray, device: torch.device) -> np.ndarr
         convert_bgr_to_rgb=True,
         use_float32=True,
     )
-    normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+    normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)  # type: ignore
     assert isinstance(face_input, torch.Tensor)
     face_input = torch.unsqueeze(face_input, 0).to(device)
 

scripts/faceswaplab_swapping/parsing/__init__.py

@@ -50,7 +50,7 @@ from scripts.faceswaplab_globals import FACE_PARSER_DIR
 ROOT_DIR = FACE_PARSER_DIR
 
 
-def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
+def load_file_from_url(url: str, model_dir=None, progress=True, file_name=None):
     """Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py"""
     if model_dir is None:
         hub_dir = get_dir()

scripts/faceswaplab_swapping/swapper.py

@@ -1,36 +1,87 @@
 import copy
 import os
 from dataclasses import dataclass
-from typing import Any, Dict, List, Set, Tuple, Optional
+from pprint import pformat
+import traceback
+from typing import Any, Dict, Generator, List, Set, Tuple, Optional, Union
 import tempfile
+from tqdm import tqdm
+import sys
+from io import StringIO
+from contextlib import contextmanager
 
 import cv2
 import insightface
 import numpy as np
-from insightface.app.common import Face
+from insightface.app.common import Face as ISFace
 
 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,
 )
 from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug
 from scripts import faceswaplab_globals
-from modules.shared import opts
 from functools import lru_cache
 from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
 from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
 from scripts.faceswaplab_postprocessing.postprocessing_options import (
     PostProcessingOptions,
 )
-from scripts.faceswaplab_utils.models_utils import get_current_model
+from scripts.faceswaplab_utils.models_utils import get_current_swap_model
-import gradio as gr
+from scripts.faceswaplab_utils.typing import CV2ImgU8, Gender, PILImage, Face
+from scripts.faceswaplab_inpainting.i2i_pp import img2img_diffusion
+from modules import shared
+import onnxruntime
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
 
 
-providers = ["CPUExecutionProvider"]
+def use_gpu() -> bool:
+    return (
+        getattr(shared.cmd_opts, "faceswaplab_gpu", False)
+        or get_sd_option("faceswaplab_use_gpu", False)
+    ) and sys.platform != "darwin"
+
+
+@lru_cache
+def force_install_gpu_providers() -> None:
+    # Ugly Ugly hack due to SDNEXT :
+    try:
+        from scripts.faceswaplab_utils.install_utils import check_install
+
+        logger.warning("Try to reinstall gpu dependencies")
+        check_install()
+        logger.warning("IF onnxruntime-gpu has been installed successfully, RESTART")
+        logger.warning(
+            "On SD.NEXT/vladmantic you will also need to check numpy>=1.24.2 and tensorflow>=2.13.0"
+        )
+    except:
+        logger.error(
+            "Reinstall has failed (which is normal on windows), please install requirements-gpu.txt manually to enable gpu."
+        )
+
+
+def get_providers() -> List[str]:
+    providers = ["CPUExecutionProvider"]
+    if use_gpu():
+        if "CUDAExecutionProvider" in onnxruntime.get_available_providers():
+            providers = ["CUDAExecutionProvider"]
+        else:
+            logger.error(
+                "CUDAExecutionProvider not found in onnxruntime.available_providers : %s, use CPU instead. Check onnxruntime-gpu is installed.",
+                onnxruntime.get_available_providers(),
+            )
+            force_install_gpu_providers()
+
+    return providers
+
+
+def is_cpu_provider() -> bool:
+    return get_providers() == ["CPUExecutionProvider"]
 
 
 def cosine_similarity_face(face1: Face, face2: Face) -> float:
@@ -50,8 +101,9 @@ def cosine_similarity_face(face1: Face, face2: Face) -> float:
         non-negative similarity score.
     """
     # Reshape the face embeddings to have a shape of (1, -1)
-    vec1 = face1.embedding.reshape(1, -1)
+    assert face1.normed_embedding is not None and face2.normed_embedding is not None
-    vec2 = face2.embedding.reshape(1, -1)
+    vec1 = face1.normed_embedding.reshape(1, -1)
+    vec2 = face2.normed_embedding.reshape(1, -1)
 
     # Calculate the cosine similarity between the reshaped embeddings
     similarity = cosine_similarity(vec1, vec2)
@@ -60,7 +112,7 @@ def cosine_similarity_face(face1: Face, face2: Face) -> float:
     return max(0, similarity[0, 0])
 
 
-def compare_faces(img1: Image.Image, img2: Image.Image) -> float:
+def compare_faces(img1: PILImage, img2: PILImage) -> float:
     """
     Compares the similarity between two faces extracted from images using cosine similarity.
 
@@ -87,22 +139,22 @@ def compare_faces(img1: Image.Image, img2: Image.Image) -> float:
 
 
 def batch_process(
-    src_images: List[Image.Image],
+    src_images: List[Union[PILImage, str]],  # image or filename
     save_path: Optional[str],
     units: List[FaceSwapUnitSettings],
-    postprocess_options: PostProcessingOptions,
+    postprocess_options: Optional[PostProcessingOptions],
-) -> Optional[List[Image.Image]]:
+) -> Optional[List[PILImage]]:
     """
     Process a batch of images, apply face swapping according to the given settings, and optionally save the resulting images to a specified path.
 
     Args:
-        src_images (List[Image.Image]): List of source PIL Images to process.
+        src_images (List[Union[PILImage, str]]): List of source PIL Images to process or list of images file names
         save_path (Optional[str]): Destination path where the processed images will be saved. If None, no images are saved.
         units (List[FaceSwapUnitSettings]): List of FaceSwapUnitSettings to apply to the images.
         postprocess_options (PostProcessingOptions): Post-processing settings to be applied to the images.
 
     Returns:
-        Optional[List[Image.Image]]: List of processed images, or None in case of an exception.
+        Optional[List[PILImage]]: List of processed images, or None in case of an exception.
 
     Raises:
         Any exceptions raised by the underlying process will be logged and the function will return None.
@@ -115,27 +167,34 @@ def batch_process(
         if src_images is not None and len(units) > 0:
             result_images = []
             for src_image in src_images:
+                path: str = ""
+                if isinstance(src_image, str):
+                    if save_path:
+                        path = os.path.join(
+                            save_path, "swapped_" + os.path.basename(src_image)
+                        )
+                    src_image = Image.open(src_image)
+                elif save_path:
+                    path = tempfile.NamedTemporaryFile(
+                        delete=False, suffix=".png", dir=save_path
+                    ).name
+                assert isinstance(src_image, Image.Image)
+
                 current_images = []
                 swapped_images = process_images_units(
-                    get_current_model(),
+                    get_current_swap_model(), images=[(src_image, None)], units=units
-                    images=[(src_image, None)],
-                    units=units,
-                    upscaled_swapper=opts.data.get(
-                        "faceswaplab_upscaled_swapper", False
-                    ),
                 )
-                if len(swapped_images) > 0:
+                if swapped_images and 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:
-                        path = tempfile.NamedTemporaryFile(
-                            delete=False, suffix=".png", dir=save_path
-                        ).name
                         img.save(path)
 
                 result_images += current_images
@@ -149,10 +208,10 @@ def batch_process(
 
 
 def extract_faces(
-    images: List[Image.Image],
+    images: List[PILImage],
     extract_path: Optional[str],
     postprocess_options: PostProcessingOptions,
-) -> Optional[List[str]]:
+) -> Optional[List[PILImage]]:
     """
     Extracts faces from a list of image files.
 
@@ -175,14 +234,14 @@ def extract_faces(
             os.makedirs(extract_path, exist_ok=True)
 
         if images:
-            result_images = []
+            result_images: list[PILImage] = []
             for img in images:
                 faces = get_faces(pil_to_cv2(img))
 
                 if faces:
                     face_images = []
                     for face in faces:
-                        bbox = face.bbox.astype(int)
+                        bbox = face.bbox.astype(int)  # type: ignore
                         x_min, y_min, x_max, y_max = bbox
                         face_image = img.crop((x_min, y_min, x_max, y_max))
 
@@ -206,7 +265,7 @@ def extract_faces(
 
             return result_images
     except Exception as e:
-        logger.info("Failed to extract : %s", e)
+        logger.error("Failed to extract : %s", e)
         import traceback
 
         traceback.print_exc()
@@ -225,8 +284,39 @@ class FaceModelException(Exception):
         super().__init__(self.message)
 
 
-@lru_cache(maxsize=1)
+@contextmanager
-def getAnalysisModel() -> insightface.app.FaceAnalysis:
+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=3)
+def getAnalysisModel(
+    det_size: Tuple[int, int] = (640, 640),
+    det_thresh: float = 0.5,
+    use_gpu: bool = False,
+) -> insightface.app.FaceAnalysis:
     """
     Retrieves the analysis model for face analysis.
 
@@ -237,11 +327,30 @@ 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.")
+        providers = get_providers()
-        # Initialize the analysis model with the specified name and providers
+        logger.info(
-        return insightface.app.FaceAnalysis(
+            f"Load analysis model det_size={det_size}, det_thresh={det_thresh}, providers = {providers}, will take some time. (> 30s)"
-            name="buffalo_l", providers=providers, root=faceswaplab_globals.ANALYZER_DIR
         )
+        # Initialize the analysis model with the specified name and providers
+
+        with tqdm(
+            total=1,
+            desc=f"Loading {det_size} analysis model (first time is slow)",
+            unit="model",
+        ) as pbar:
+            with capture_stdout() as captured:
+                model = insightface.app.FaceAnalysis(
+                    name="buffalo_l",
+                    providers=providers,
+                    root=faceswaplab_globals.ANALYZER_DIR,
+                )
+
+                # Prepare the analysis model for face detection with the specified detection size
+                model.prepare(ctx_id=0, det_thresh=det_thresh, det_size=det_size)
+            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.)"
@@ -250,7 +359,9 @@ def getAnalysisModel() -> insightface.app.FaceAnalysis:
 
 
 @lru_cache(maxsize=1)
-def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper:
+def getFaceSwapModel(
+    model_path: str, use_gpu: bool = False
+) -> upscaled_inswapper.UpscaledINSwapper:
     """
     Retrieves the face swap model and initializes it if necessary.
 
@@ -261,28 +372,34 @@ def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper:
         insightface.model_zoo.FaceModel: The face swap model.
     """
     try:
-        # Initializes the face swap model using the specified model path.
+        providers = get_providers()
-        return upscaled_inswapper.UpscaledINSwapper(
+        with tqdm(total=1, desc="Loading swap model", unit="model") as pbar:
-            insightface.model_zoo.get_model(model_path, providers=providers)
+            with capture_stdout() as captured:
-        )
+                model = upscaled_inswapper.UpscaledINSwapper(
+                    insightface.model_zoo.get_model(model_path, providers=providers)  # type: ignore
+                )
+            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.)"
         )
+        traceback.print_exc()
         raise FaceModelException("Loading of swapping model failed")
 
 
 def get_faces(
-    img_data: np.ndarray,  # type: ignore
+    img_data: CV2ImgU8,
-    det_size: Tuple[int, int] = (640, 640),
     det_thresh: Optional[float] = None,
-    sort_by_face_size: bool = False,
+    det_size: Tuple[int, int] = (640, 640),
 ) -> List[Face]:
     """
     Detects and retrieves faces from an image using an analysis model.
 
     Args:
-        img_data (np.ndarray): The image data as a NumPy array.
+        img_data (CV2ImgU8): The image data as a NumPy array.
         det_size (tuple): The desired detection size (width, height). Defaults to (640, 640).
         sort_by_face_size (bool) : Will sort the faces by their size from larger to smaller face
 
@@ -291,44 +408,99 @@ def get_faces(
     """
 
     if det_thresh is None:
-        det_thresh = opts.data.get("faceswaplab_detection_threshold", 0.5)
+        det_thresh = get_sd_option("faceswaplab_detection_threshold", 0.5)
 
-    # Create a deep copy of the analysis model (otherwise det_size is attached to the analysis model and can't be changed)
+    auto_det_size = get_sd_option("faceswaplab_auto_det_size", True)
-    face_analyser = copy.deepcopy(getAnalysisModel())
+    if not auto_det_size:
+        x = get_sd_option("faceswaplab_det_size", 640)
+        det_size = (x, x)
 
-    # Prepare the analysis model for face detection with the specified detection size
+    face_analyser = getAnalysisModel(
-    face_analyser.prepare(ctx_id=0, det_thresh=det_thresh, det_size=det_size)
+        det_size=det_size, det_thresh=det_thresh, use_gpu=not is_cpu_provider()
+    )
 
     # Get the detected faces from the image using the analysis model
-    face = face_analyser.get(img_data)
+    faces = face_analyser.get(img_data)
 
     # If no faces are detected and the detection size is larger than 320x320,
     # recursively call the function with a smaller detection size
-    if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320:
+    if len(faces) == 0:
-        det_size_half = (det_size[0] // 2, det_size[1] // 2)
+        if auto_det_size:
-        return get_faces(img_data, det_size=det_size_half, det_thresh=det_thresh)
+            if det_size[0] > 320 and det_size[1] > 320:
+                det_size_half = (det_size[0] // 2, det_size[1] // 2)
+                return get_faces(
+                    img_data, det_size=det_size_half, det_thresh=det_thresh
+                )
+
+        # If no faces are detected print a warning to user about change in detection
+        else:
+            if det_size[0] > 320:
+                logger.warning(
+                    "No faces detected, you might want to play with det_size by reducing it (in sd global settings). Lower (320) means more detection but less precise. Or activate auto-det-size."
+                )
 
     try:
-        if sort_by_face_size:
-            return sorted(
-                face,
-                reverse=True,
-                key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]),
-            )
-
         # Sort the detected faces based on their x-coordinate of the bounding box
-        return sorted(face, key=lambda x: x.bbox[0])
+        return sorted(faces, key=lambda x: x.bbox[0])  # type: ignore
     except Exception as e:
+        logger.error("Failed to get faces %s", e)
+        traceback.print_exc()
         return []
 
 
+@dataclass
+class FaceFilteringOptions:
+    faces_index: Set[int]
+    source_gender: Optional[int] = None  # if none will not use same gender
+    sort_by_face_size: bool = False
+
+
+def filter_faces(
+    all_faces: List[Face], filtering_options: FaceFilteringOptions
+) -> List[Face]:
+    """
+    Sorts and filters a list of faces based on specified criteria.
+
+    This function takes a list of Face objects and can sort them by face size and filter them by gender.
+    Sorting by face size is performed if sort_by_face_size is set to True, and filtering by gender is
+    performed if source_gender is provided.
+
+    :param faces: A list of Face objects representing the faces to be sorted and filtered.
+    :param faces_index: A set of faces index
+    :param source_gender: An optional integer representing the gender by which to filter the faces.
+                          If provided, only faces with the specified gender will be included in the result.
+    :param sort_by_face_size: A boolean indicating whether to sort the faces by size. If True, faces are
+                              sorted in descending order by size, calculated as the area of the bounding box.
+    :return: A list of Face objects sorted and filtered according to the specified criteria.
+    """
+    filtered_faces = copy.copy(all_faces)
+    if filtering_options.sort_by_face_size:
+        filtered_faces = sorted(
+            all_faces,
+            reverse=True,
+            key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]),  # type: ignore
+        )
+
+    if filtering_options.source_gender is not None:
+        filtered_faces = [
+            face
+            for face in filtered_faces
+            if face["gender"] == filtering_options.source_gender
+        ]
+    return [
+        face
+        for i, face in enumerate(filtered_faces)
+        if i in filtering_options.faces_index
+    ]
+
+
 @dataclass
 class ImageResult:
     """
     Represents the result of an image swap operation
     """
 
-    image: Image.Image
+    image: PILImage
     """
     The image object with the swapped face
     """
@@ -362,23 +534,22 @@ def get_or_default(l: List[Any], index: int, default: Any) -> Any:
     return l[index] if index < len(l) else default
 
 
-def get_faces_from_img_files(files: List[gr.File]) -> List[Optional[np.ndarray]]:  # type: ignore
+def get_faces_from_img_files(images: List[PILImage]) -> List[Face]:
     """
     Extracts faces from a list of image files.
 
     Args:
-        files (list): A list of file objects representing image files.
+        images (list): A list of PILImage objects representing image files.
 
     Returns:
         list: A list of detected faces.
 
     """
 
-    faces = []
+    faces: List[Face] = []
 
-    if len(files) > 0:
+    if len(images) > 0:
-        for file in files:
+        for img in images:
-            img = Image.open(file.name)  # Open the image file
             face = get_or_default(
                 get_faces(pil_to_cv2(img)), 0, None
             )  # Extract faces from the image
@@ -388,7 +559,7 @@ def get_faces_from_img_files(files: List[gr.File]) -> List[Optional[np.ndarray]]
     return faces
 
 
-def blend_faces(faces: List[Face]) -> Face:
+def blend_faces(faces: List[Face], gender: Gender = Gender.AUTO) -> Optional[Face]:
     """
     Blends the embeddings of multiple faces into a single face.
 
@@ -403,7 +574,7 @@ def blend_faces(faces: List[Face]) -> Face:
         ValueError: If the embeddings have different shapes.
 
     """
-    embeddings = [face.embedding for face in faces]
+    embeddings: list[Any] = [face.embedding for face in faces]
 
     if len(embeddings) > 0:
         embedding_shape = embeddings[0].shape
@@ -416,18 +587,21 @@ def blend_faces(faces: List[Face]) -> Face:
         # Compute the mean of all embeddings
         blended_embedding = np.mean(embeddings, axis=0)
 
+        if gender == Gender.AUTO:
+            int_gender: int = faces[0].gender  # type: ignore
+        else:
+            int_gender: int = gender.value
+
+        assert -1 < int_gender < 2, "wrong gender"
+
+        logger.info("Int Gender : %s", int_gender)
+
         # Create a new Face object using the properties of the first face in the list
         # Assign the blended embedding to the blended Face object
-        blended = Face(
+        blended = ISFace(
-            embedding=blended_embedding, gender=faces[0].gender, age=faces[0].age
+            embedding=blended_embedding, gender=int_gender, age=faces[0].age
         )
 
-        assert (
-            not np.array_equal(blended.embedding, faces[0].embedding)
-            if len(faces) > 1
-            else True
-        ), "If len(faces)>0, the blended embedding should not be the same than the first image"
-
         return blended
 
     # Return None if the input list is empty
@@ -435,99 +609,95 @@ def blend_faces(faces: List[Face]) -> Face:
 
 
 def swap_face(
-    reference_face: np.ndarray,  # type: ignore
+    source_face: Face,
-    source_face: np.ndarray,  # type: ignore
+    target_img: PILImage,
-    target_img: Image.Image,
+    target_faces: List[Face],
     model: str,
-    faces_index: Set[int] = {0},
+    swapping_options: Optional[InswappperOptions],
-    same_gender: bool = True,
-    upscaled_swapper: bool = False,
-    compute_similarity: bool = True,
-    sort_by_face_size: bool = False,
 ) -> ImageResult:
     """
     Swaps faces in the target image with the source face.
 
     Args:
-        reference_face (np.ndarray): The reference face used for similarity comparison.
+        source_face (CV2ImgU8): The source face to be swapped.
-        source_face (np.ndarray): The source face to be swapped.
+        target_img (PILImage): The target image to swap faces in.
-        target_img (Image.Image): The target image to swap faces in.
         model (str): Path to the face swap model.
-        faces_index (Set[int], optional): Set of indices specifying which faces to swap. Defaults to {0}.
-        same_gender (bool, optional): If True, only swap faces with the same gender as the source face. Defaults to True.
 
     Returns:
         ImageResult: An object containing the swapped image and similarity scores.
 
     """
     return_result = ImageResult(target_img, {}, {})
+    target_img_cv2: CV2ImgU8 = cv2.cvtColor(
+        np.array(target_img), cv2.COLOR_RGB2BGR
+    ).astype("uint8")
     try:
-        target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
         gender = source_face["gender"]
         logger.info("Source Gender %s", gender)
         if source_face is not None:
-            result = target_img
+            result: CV2ImgU8 = target_img_cv2
             model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model)
-            face_swapper = getFaceSwapModel(model_path)
+            face_swapper = getFaceSwapModel(model_path, use_gpu=not is_cpu_provider())
-            target_faces = get_faces(target_img, sort_by_face_size=sort_by_face_size)
             logger.info("Target faces count : %s", len(target_faces))
 
-            if same_gender:
-                target_faces = [x for x in target_faces if x["gender"] == gender]
-                logger.info("Target Gender Matches count %s", len(target_faces))
-
             for i, swapped_face in enumerate(target_faces):
                 logger.info(f"swap face {i}")
-                if i in faces_index:
+
-                    # type : ignore
+                result = face_swapper.get(
-                    result = face_swapper.get(
+                    img=result,
-                        result, swapped_face, source_face, upscale=upscaled_swapper
+                    target_face=swapped_face,
-                    )
+                    source_face=source_face,
+                    options=swapping_options,
+                )  # type: ignore
 
             result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
             return_result.image = result_image
 
-            if compute_similarity:
-                try:
-                    result_faces = get_faces(
-                        cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR),
-                        sort_by_face_size=sort_by_face_size,
-                    )
-                    if same_gender:
-                        result_faces = [
-                            x for x in result_faces if x["gender"] == gender
-                        ]
-
-                    for i, swapped_face in enumerate(result_faces):
-                        logger.info(f"compare face {i}")
-                        if i in faces_index and i < len(target_faces):
-                            return_result.similarity[i] = cosine_similarity_face(
-                                source_face, swapped_face
-                            )
-                            return_result.ref_similarity[i] = cosine_similarity_face(
-                                reference_face, swapped_face
-                            )
-
-                        logger.info(f"similarity {return_result.similarity}")
-                        logger.info(f"ref similarity {return_result.ref_similarity}")
-
-                except Exception as e:
-                    logger.error("Similarity processing failed %s", e)
-                    raise e
     except Exception as e:
         logger.error("Conversion failed %s", e)
         raise e
     return return_result
 
 
+def compute_similarity(
+    reference_face: Face,
+    source_face: Face,
+    swapped_image: PILImage,
+    filtering: FaceFilteringOptions,
+) -> Tuple[Dict[int, float], Dict[int, float]]:
+    similarity: Dict[int, float] = {}
+    ref_similarity: Dict[int, float] = {}
+    try:
+        swapped_image_cv2: CV2ImgU8 = cv2.cvtColor(
+            np.array(swapped_image), cv2.COLOR_RGB2BGR
+        )
+        new_faces = filter_faces(get_faces(swapped_image_cv2), filtering)
+        if len(new_faces) == 0:
+            logger.error("compute_similarity : No faces to compare with !")
+            return None
+
+        for i, swapped_face in enumerate(new_faces):
+            logger.info(f"compare face {i}")
+            similarity[i] = cosine_similarity_face(source_face, swapped_face)
+            ref_similarity[i] = cosine_similarity_face(reference_face, swapped_face)
+
+            logger.info(f"similarity {similarity}")
+            logger.info(f"ref similarity {ref_similarity}")
+
+        return (similarity, ref_similarity)
+    except Exception as e:
+        logger.error("Similarity processing failed %s", e)
+        raise e
+    return None
+
+
 def process_image_unit(
     model: str,
     unit: FaceSwapUnitSettings,
-    image: Image.Image,
+    image: PILImage,
-    info: str = None,
+    info: Optional[str] = None,
-    upscaled_swapper: bool = False,
     force_blend: bool = False,
-) -> List[Tuple[Image.Image, str]]:
+) -> List[Tuple[PILImage, Optional[str]]]:
     """Process one image and return a List of (image, info) (one if blended, many if not).
 
     Args:
@@ -541,6 +711,8 @@ def process_image_unit(
 
     results = []
     if unit.enable:
+        faces = get_faces(pil_to_cv2(image))
+
         if check_against_nsfw(image):
             return [(image, info)]
         if not unit.blend_faces and not force_blend:
@@ -549,15 +721,10 @@ def process_image_unit(
         else:
             logger.info("blend all faces together")
             src_faces = [unit.blended_faces]
-            assert (
-                not np.array_equal(
-                    unit.reference_face.embedding, src_faces[0].embedding
-                )
-                if len(unit.faces) > 1
-                else True
-            ), "Reference face cannot be the same as blended"
 
         for i, src_face in enumerate(src_faces):
+            current_image = image
+
             logger.info(f"Process face {i}")
             if unit.reference_face is not None:
                 reference_face = unit.reference_face
@@ -565,20 +732,50 @@ def process_image_unit(
                 logger.info("Use source face as reference face")
                 reference_face = src_face
 
-            save_img_debug(image, "Before swap")
+            face_filtering_options = FaceFilteringOptions(
-            result: ImageResult = swap_face(
-                reference_face,
-                src_face,
-                image,
                 faces_index=unit.faces_index,
-                model=model,
+                source_gender=src_face["gender"] if unit.same_gender else None,
-                same_gender=unit.same_gender,
-                upscaled_swapper=upscaled_swapper,
-                compute_similarity=unit.compute_similarity,
                 sort_by_face_size=unit.sort_by_size,
             )
+
+            target_faces: List[Face] = filter_faces(
+                all_faces=faces, filtering_options=face_filtering_options
+            )
+
+            # Apply pre-inpainting to image
+            if unit.pre_inpainting.inpainting_denoising_strengh > 0:
+                current_image = img2img_diffusion(
+                    img=current_image, faces=target_faces, options=unit.pre_inpainting
+                )
+
+            save_img_debug(image, "Before swap")
+            result: ImageResult = swap_face(
+                source_face=src_face,
+                target_img=current_image,
+                target_faces=target_faces,
+                model=model,
+                swapping_options=unit.swapping_options,
+            )
+            # Apply post-inpainting to image
+            if unit.post_inpainting.inpainting_denoising_strengh > 0:
+                result.image = img2img_diffusion(
+                    img=result.image, faces=target_faces, options=unit.post_inpainting
+                )
+
             save_img_debug(result.image, "After swap")
 
+            if unit.compute_similarity:
+                similarities = compute_similarity(
+                    reference_face=reference_face,
+                    source_face=src_face,
+                    swapped_image=result.image,
+                    filtering=face_filtering_options,
+                )
+                if similarities:
+                    (result.similarity, result.ref_similarity) = similarities
+                else:
+                    logger.error("Failed to compute similarity")
+
             if result.image is None:
                 logger.error("Result image is None")
             if (
@@ -610,17 +807,16 @@ def process_image_unit(
 def process_images_units(
     model: str,
     units: List[FaceSwapUnitSettings],
-    images: List[Tuple[Optional[Image.Image], Optional[str]]],
+    images: List[Tuple[Optional[PILImage], Optional[str]]],
-    upscaled_swapper: bool = False,
     force_blend: bool = False,
-) -> Optional[List[Tuple[Image.Image, str]]]:
+) -> Optional[List[Tuple[PILImage, str]]]:
     """
     Process a list of images using a specified model and unit settings for face swapping.
 
     Args:
         model (str): The name of the model to use for processing.
         units (List[FaceSwapUnitSettings]): A list of settings for face swap units to apply on each image.
-        images (List[Tuple[Optional[Image.Image], Optional[str]]]): A list of tuples, each containing
+        images (List[Tuple[Optional[PILImage], Optional[str]]]): A list of tuples, each containing
             an image and its associated info string. If an image or info string is not available,
             its value can be None.
         upscaled_swapper (bool, optional): If True, uses an upscaled version of the face swapper.
@@ -629,7 +825,7 @@ def process_images_units(
             image. Defaults to False.
 
     Returns:
-        Optional[List[Tuple[Image.Image, str]]]: A list of tuples, each containing a processed image
+        Optional[List[Tuple[PILImage, str]]]: A list of tuples, each containing a processed image
             and its associated info string. If no units are provided for processing, returns None.
 
     """
@@ -642,13 +838,9 @@ def process_images_units(
     processed_images = []
     for i, (image, info) in enumerate(images):
         logger.debug("Processing image %s", i)
-        swapped = process_image_unit(
+        swapped = process_image_unit(model, units[0], image, info, force_blend)
-            model, units[0], image, info, upscaled_swapper, force_blend
-        )
         logger.debug("Image %s -> %s images", i, len(swapped))
-        nexts = process_images_units(
+        nexts = process_images_units(model, units[1:], swapped, force_blend)
-            model, units[1:], swapped, upscaled_swapper, force_blend
-        )
         if nexts:
             processed_images.extend(nexts)
         else:

scripts/faceswaplab_swapping/upcaled_inswapper_options.py

@@ -0,0 +1,39 @@
+from dataclasses import *
+from typing import Optional
+from client_api import api_utils
+
+
+@dataclass
+class InswappperOptions:
+    face_restorer_name: Optional[str] = None
+    restorer_visibility: float = 1
+    codeformer_weight: float = 1
+    upscaler_name: Optional[str] = None
+    improved_mask: bool = False
+    color_corrections: bool = False
+    sharpen: bool = False
+    erosion_factor: float = 1.0
+
+    @staticmethod
+    def from_api_dto(dto: Optional[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

@@ -1,9 +1,9 @@
+from typing import Any, Optional, Tuple, Union
 import cv2
 import numpy as np
 from insightface.model_zoo.inswapper import INSwapper
 from insightface.utils import face_align
 from modules import processing, shared
-from modules.shared import opts
 from modules.upscaler import UpscalerData
 
 from scripts.faceswaplab_postprocessing import upscaling
@@ -11,19 +11,41 @@ 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.sd_utils import get_sd_option
+from scripts.faceswaplab_utils.typing import CV2ImgU8, Face
+from scripts.faceswaplab_utils.faceswaplab_logging import logger
 
 
-def get_upscaler() -> UpscalerData:
+def get_upscaler() -> Optional[UpscalerData]:
     for upscaler in shared.sd_upscalers:
-        if upscaler.name == opts.data.get(
+        if upscaler.name == get_sd_option(
             "faceswaplab_upscaled_swapper_upscaler", "LDSR"
         ):
             return upscaler
     return None
 
 
-def merge_images_with_mask(image1, image2, mask):
+def merge_images_with_mask(
+    image1: CV2ImgU8, image2: CV2ImgU8, mask: CV2ImgU8
+) -> CV2ImgU8:
+    """
+    Merges two images using a given mask. The regions where the mask is set will be replaced with the corresponding
+    areas of the second image.
+
+    Args:
+        image1 (CV2Img): The base image, which must have the same shape as image2.
+        image2 (CV2Img): The image to be merged, which must have the same shape as image1.
+        mask (CV2Img): A binary mask specifying the regions to be merged. The mask shape should match image1's first two dimensions.
+
+    Returns:
+        CV2Img: The merged image.
+
+    Raises:
+        ValueError: If the shapes of the images and mask do not match.
+    """
+
     if image1.shape != image2.shape or image1.shape[:2] != mask.shape:
         raise ValueError("Img should have the same shape")
     mask = mask.astype(np.uint8)
@@ -34,64 +56,114 @@ def merge_images_with_mask(image1, image2, mask):
     return merged_image
 
 
-def erode_mask(mask, kernel_size=3, iterations=1):
+def erode_mask(mask: CV2ImgU8, kernel_size: int = 3, iterations: int = 1) -> CV2ImgU8:
+    """
+    Erodes a binary mask using a given kernel size and number of iterations.
+
+    Args:
+        mask (CV2Img): The binary mask to erode.
+        kernel_size (int, optional): The size of the kernel. Default is 3.
+        iterations (int, optional): The number of erosion iterations. Default is 1.
+
+    Returns:
+        CV2Img: The eroded mask.
+    """
     kernel = np.ones((kernel_size, kernel_size), np.uint8)
     eroded_mask = cv2.erode(mask, kernel, iterations=iterations)
     return eroded_mask
 
 
-def apply_gaussian_blur(mask, kernel_size=(5, 5), sigma_x=0):
+def apply_gaussian_blur(
+    mask: CV2ImgU8, kernel_size: Tuple[int, int] = (5, 5), sigma_x: int = 0
+) -> CV2ImgU8:
+    """
+    Applies a Gaussian blur to a mask.
+
+    Args:
+        mask (CV2Img): The mask to blur.
+        kernel_size (tuple, optional): The size of the kernel, e.g. (5, 5). Default is (5, 5).
+        sigma_x (int, optional): The standard deviation in the X direction. Default is 0.
+
+    Returns:
+        CV2Img: The blurred mask.
+    """
     blurred_mask = cv2.GaussianBlur(mask, kernel_size, sigma_x)
     return blurred_mask
 
 
-def dilate_mask(mask, kernel_size=5, iterations=1):
+def dilate_mask(mask: CV2ImgU8, kernel_size: int = 5, iterations: int = 1) -> CV2ImgU8:
+    """
+    Dilates a binary mask using a given kernel size and number of iterations.
+
+    Args:
+        mask (CV2Img): The binary mask to dilate.
+        kernel_size (int, optional): The size of the kernel. Default is 5.
+        iterations (int, optional): The number of dilation iterations. Default is 1.
+
+    Returns:
+        CV2Img: The dilated mask.
+    """
     kernel = np.ones((kernel_size, kernel_size), np.uint8)
     dilated_mask = cv2.dilate(mask, kernel, iterations=iterations)
     return dilated_mask
 
 
-def get_face_mask(aimg, bgr_fake):
+def get_face_mask(aimg: CV2ImgU8, bgr_fake: CV2ImgU8) -> CV2ImgU8:
+    """
+    Generates a face mask by performing bitwise OR on two face masks and then dilating the result.
+
+    Args:
+        aimg (CV2Img): Input image for generating the first face mask.
+        bgr_fake (CV2Img): Input image for generating the second face mask.
+
+    Returns:
+        CV2Img: The combined and dilated face mask.
+    """
     mask1 = generate_face_mask(aimg, device=shared.device)
     mask2 = generate_face_mask(bgr_fake, device=shared.device)
     mask = dilate_mask(cv2.bitwise_or(mask1, mask2))
     return mask
 
 
-class UpscaledINSwapper:
+class UpscaledINSwapper(INSwapper):
     def __init__(self, inswapper: INSwapper):
         self.__dict__.update(inswapper.__dict__)
 
-    def forward(self, img, latent):
+    def upscale_and_restore(
-        img = (img - self.input_mean) / self.input_std
+        self,
-        pred = self.session.run(
+        img: CV2ImgU8,
-            self.output_names, {self.input_names[0]: img, self.input_names[1]: latent}
+        k: int = 2,
-        )[0]
+        inswapper_options: Optional[InswappperOptions] = None,
-        return pred
+    ) -> CV2ImgU8:
+        if inswapper_options is None:
+            return img
 
-    def super_resolution(self, img, k=2):
         pil_img = cv2_to_pil(img)
-        options = PostProcessingOptions(
+        pp_options = PostProcessingOptions(
-            upscaler_name=opts.data.get(
+            upscaler_name=inswapper_options.upscaler_name,
-                "faceswaplab_upscaled_swapper_upscaler", "LDSR"
-            ),
             upscale_visibility=1,
             scale=k,
-            face_restorer_name=opts.data.get(
+            face_restorer_name=inswapper_options.face_restorer_name,
-                "faceswaplab_upscaled_swapper_face_restorer", ""
+            codeformer_weight=inswapper_options.codeformer_weight,
-            ),
+            restorer_visibility=inswapper_options.restorer_visibility,
-            codeformer_weight=opts.data.get(
-                "faceswaplab_upscaled_swapper_face_restorer_weight", 1
-            ),
-            restorer_visibility=opts.data.get(
-                "faceswaplab_upscaled_swapper_face_restorer_visibility", 1
-            ),
         )
-        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(self, img, target_face, source_face, paste_back=True, upscale=True):
+    def get(
+        self,
+        img: CV2ImgU8,
+        target_face: Face,
+        source_face: Face,
+        paste_back: bool = True,
+        options: Optional[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(
             aimg,
@@ -100,9 +172,10 @@ class UpscaledINSwapper:
             (self.input_mean, self.input_mean, self.input_mean),
             swapRB=True,
         )
-        latent = source_face.normed_embedding.reshape((1, -1))
+        latent = source_face.normed_embedding.reshape((1, -1))  # type: ignore
         latent = np.dot(latent, self.emap)
         latent /= np.linalg.norm(latent)
+        assert self.session is not None
         pred = self.session.run(
             self.output_names, {self.input_names[0]: blob, self.input_names[1]: latent}
         )[0]
@@ -116,7 +189,7 @@ class UpscaledINSwapper:
             else:
                 target_img = img
 
-                def compute_diff(bgr_fake, aimg):
+                def compute_diff(bgr_fake: CV2ImgU8, aimg: CV2ImgU8) -> CV2ImgU8:
                     fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
                     fake_diff = np.abs(fake_diff).mean(axis=2)
                     fake_diff[:2, :] = 0
@@ -125,43 +198,65 @@ class UpscaledINSwapper:
                     fake_diff[:, -2:] = 0
                     return fake_diff
 
-                if upscale:
+                if options:
-                    print("*" * 80)
+                    logger.info("*" * 80)
-                    print(
+                    logger.info(f"Inswapper")
-                        f"Upscaled inswapper using {opts.data.get('faceswaplab_upscaled_swapper_upscaler', 'LDSR')}"
-                    )
-                    print("*" * 80)
 
-                    k = 4
+                    if options.upscaler_name and options.upscaler_name != "None":
-                    aimg, M = face_align.norm_crop2(
+                        # Upscale original image
-                        img, target_face.kps, self.input_size[0] * k
+                        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):
+                    fake_diff: CV2ImgU8 = None  # type: ignore
-                        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)
+                    if not options.improved_mask:
-                    fake_diff = compute_diff(bgr_fake, aimg)
+                        # If improved mask is not used, we should compute before sharpen and color correction (better diff)
+                        fake_diff = compute_diff(bgr_fake, aimg=aimg)
 
-                    if opts.data.get("faceswaplab_upscaled_swapper_sharpen", True):
+                    if options.sharpen:
-                        print("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):
+                    if options.color_corrections:
-                        print("color correction")
+                        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)
 
+                    if options.improved_mask:
+                        if k == 1:
+                            logger.warning(
+                                "Please note that improved mask does not work well without upscaling. Set upscaling to Lanczos at least if you want speed and want to use improved mask."
+                            )
+
+                        logger.info("improved_mask")
+                        mask = get_face_mask(aimg, bgr_fake)
+                        # save_img_debug(cv2_to_pil(bgr_fake), "Before Mask")
+                        bgr_fake = merge_images_with_mask(aimg, bgr_fake, mask)
+                        # save_img_debug(cv2_to_pil(bgr_fake), "After Mask")
+
+                        fake_diff = compute_diff(bgr_fake, aimg=aimg)
+
+                    assert (
+                        fake_diff is not None
+                    ), "fake diff is None, this should not happen"
+
+                    logger.info("*" * 80)
+
                 else:
                     fake_diff = compute_diff(bgr_fake, aimg)
 
@@ -182,6 +277,7 @@ class UpscaledINSwapper:
                     (target_img.shape[1], target_img.shape[0]),
                     borderValue=0.0,
                 )
+
                 fake_diff = cv2.warpAffine(
                     fake_diff,
                     IM,
@@ -189,8 +285,7 @@ class UpscaledINSwapper:
                     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
                 img_mask = img_white
@@ -198,9 +293,8 @@ class UpscaledINSwapper:
                 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(
+                erosion_factor = options.erosion_factor if options else 1
-                    "faceswaplab_upscaled_swapper_erosion", 1
+
-                )
                 k = max(int(mask_size // 10 * erosion_factor), int(10 * erosion_factor))
 
                 kernel = np.ones((k, k), np.uint8)

scripts/faceswaplab_ui/faceswaplab_inpainting_ui.py

@@ -0,0 +1,74 @@
+from typing import List
+import gradio as gr
+from modules import sd_models, sd_samplers
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
+
+
+def face_inpainting_ui(id_prefix: str = "faceswaplab") -> List[gr.components.Component]:
+    inpainting_denoising_strength = gr.Slider(
+        0,
+        1,
+        0,
+        step=0.01,
+        elem_id=f"{id_prefix}_pp_inpainting_denoising_strength",
+        label="Denoising strenght",
+    )
+
+    inpainting_denoising_prompt = gr.Textbox(
+        get_sd_option(
+            "faceswaplab_pp_default_inpainting_prompt", "Portrait of a [gender]"
+        ),
+        elem_id=f"{id_prefix}_pp_inpainting_denoising_prompt",
+        label="Inpainting prompt use [gender] instead of men or woman",
+    )
+    inpainting_denoising_negative_prompt = gr.Textbox(
+        get_sd_option("faceswaplab_pp_default_inpainting_negative_prompt", "blurry"),
+        elem_id=f"{id_prefix}_pp_inpainting_denoising_neg_prompt",
+        label="Inpainting negative prompt use [gender] instead of men or woman",
+    )
+    with gr.Row():
+        samplers_names = [s.name for s in sd_samplers.all_samplers]
+        inpainting_sampler = gr.Dropdown(
+            choices=samplers_names,
+            value=[samplers_names[0]],
+            label="Inpainting Sampler",
+            elem_id=f"{id_prefix}_pp_inpainting_sampler",
+        )
+        inpainting_denoising_steps = gr.Slider(
+            1,
+            150,
+            20,
+            step=1,
+            label="Inpainting steps",
+            elem_id=f"{id_prefix}_pp_inpainting_steps",
+        )
+
+    inpaiting_model = gr.Dropdown(
+        choices=["Current"] + sd_models.checkpoint_tiles(),
+        default="Current",
+        label="sd model (experimental)",
+        elem_id=f"{id_prefix}_pp_inpainting_sd_model",
+    )
+
+    inpaiting_seed = gr.Number(
+        label="Inpainting seed",
+        value=0,
+        minimum=0,
+        precision=0,
+        elem_id=f"{id_prefix}_pp_inpainting_seed",
+    )
+
+    gradio_components: List[gr.components.Component] = [
+        inpainting_denoising_strength,
+        inpainting_denoising_prompt,
+        inpainting_denoising_negative_prompt,
+        inpainting_denoising_steps,
+        inpainting_sampler,
+        inpaiting_model,
+        inpaiting_seed,
+    ]
+
+    for component in gradio_components:
+        setattr(component, "do_not_save_to_config", True)
+
+    return gradio_components

scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py

@@ -1,32 +1,33 @@
 from typing import List
 import gradio as gr
-import modules
+from modules import shared
-from modules import shared, sd_models
-from modules.shared import opts
 from scripts.faceswaplab_postprocessing.postprocessing_options import InpaintingWhen
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
+from scripts.faceswaplab_ui.faceswaplab_inpainting_ui import face_inpainting_ui
 
 
 def postprocessing_ui() -> List[gr.components.Component]:
-    with gr.Tab(f"Post-Processing"):
+    with gr.Tab(f"Global Post-Processing"):
         gr.Markdown(
-            """Upscaling is performed on the whole image. Upscaling happens before face restoration."""
+            """Upscaling is performed on the whole image and all faces (including not swapped). Upscaling happens before face restoration."""
         )
         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(
+                value=get_sd_option(
                     "faceswaplab_pp_default_face_restorer",
                     shared.face_restorers[0].name(),
                 ),
                 type="value",
                 elem_id="faceswaplab_pp_face_restorer",
             )
+
             with gr.Column():
                 face_restorer_visibility = gr.Slider(
                     0,
                     1,
-                    value=lambda: opts.data.get(
+                    value=get_sd_option(
                         "faceswaplab_pp_default_face_restorer_visibility", 1
                     ),
                     step=0.001,
@@ -36,7 +37,7 @@ def postprocessing_ui() -> List[gr.components.Component]:
                 codeformer_weight = gr.Slider(
                     0,
                     1,
-                    value=lambda: opts.data.get(
+                    value=get_sd_option(
                         "faceswaplab_pp_default_face_restorer_weight", 1
                     ),
                     step=0.001,
@@ -45,7 +46,7 @@ def postprocessing_ui() -> List[gr.components.Component]:
                 )
         upscaler_name = gr.Dropdown(
             choices=[upscaler.name for upscaler in shared.sd_upscalers],
-            value=lambda: opts.data.get("faceswaplab_pp_default_upscaler", "None"),
+            value=get_sd_option("faceswaplab_pp_default_upscaler", "None"),
             label="Upscaler",
             elem_id="faceswaplab_pp_upscaler",
         )
@@ -60,16 +61,15 @@ def postprocessing_ui() -> List[gr.components.Component]:
         upscaler_visibility = gr.Slider(
             0,
             1,
-            value=lambda: opts.data.get(
+            value=get_sd_option("faceswaplab_pp_default_upscaler_visibility", 1),
-                "faceswaplab_pp_default_upscaler_visibility", 1
-            ),
             step=0.1,
             label="Upscaler visibility (if scale = 1)",
             elem_id="faceswaplab_pp_upscaler_visibility",
         )
-        with gr.Accordion(f"Post Inpainting", open=True):
+
+        with gr.Accordion(label="Global-Inpainting (all faces)", open=False):
             gr.Markdown(
-                """Inpainting sends image to inpainting with a mask on face (once for each faces)."""
+                "Inpainting sends image to inpainting with a mask on face (once for each faces)."
             )
             inpainting_when = gr.Dropdown(
                 elem_id="faceswaplab_pp_inpainting_when",
@@ -77,64 +77,19 @@ def postprocessing_ui() -> List[gr.components.Component]:
                 value=[InpaintingWhen.BEFORE_RESTORE_FACE.value],
                 label="Enable/When",
             )
-            inpainting_denoising_strength = gr.Slider(
+            global_inpainting = face_inpainting_ui("faceswaplab_gpp")
-                0,
-                1,
-                0,
-                step=0.01,
-                elem_id="faceswaplab_pp_inpainting_denoising_strength",
-                label="Denoising strenght (will send face to img2img after processing)",
-            )
 
-            inpainting_denoising_prompt = gr.Textbox(
+    components = [
-                opts.data.get(
-                    "faceswaplab_pp_default_inpainting_prompt", "Portrait of a [gender]"
-                ),
-                elem_id="faceswaplab_pp_inpainting_denoising_prompt",
-                label="Inpainting prompt use [gender] instead of men or woman",
-            )
-            inpainting_denoising_negative_prompt = gr.Textbox(
-                opts.data.get(
-                    "faceswaplab_pp_default_inpainting_negative_prompt", "blurry"
-                ),
-                elem_id="faceswaplab_pp_inpainting_denoising_neg_prompt",
-                label="Inpainting negative prompt use [gender] instead of men or woman",
-            )
-            with gr.Row():
-                samplers_names = [s.name for s in modules.sd_samplers.all_samplers]
-                inpainting_sampler = gr.Dropdown(
-                    choices=samplers_names,
-                    value=[samplers_names[0]],
-                    label="Inpainting Sampler",
-                    elem_id="faceswaplab_pp_inpainting_sampler",
-                )
-                inpainting_denoising_steps = gr.Slider(
-                    1,
-                    150,
-                    20,
-                    step=1,
-                    label="Inpainting steps",
-                    elem_id="faceswaplab_pp_inpainting_steps",
-                )
-
-            inpaiting_model = gr.Dropdown(
-                choices=["Current"] + sd_models.checkpoint_tiles(),
-                default="Current",
-                label="sd model (experimental)",
-                elem_id="faceswaplab_pp_inpainting_sd_model",
-            )
-    return [
         face_restorer_name,
         face_restorer_visibility,
         codeformer_weight,
         upscaler_name,
         upscaler_scale,
         upscaler_visibility,
-        inpainting_denoising_strength,
-        inpainting_denoising_prompt,
-        inpainting_denoising_negative_prompt,
-        inpainting_denoising_steps,
-        inpainting_sampler,
         inpainting_when,
-        inpaiting_model,
+    ] + global_inpainting
-    ]
+
+    # Ask sd to not store in ui-config.json
+    for component in components:
+        setattr(component, "do_not_save_to_config", True)
+    return components

scripts/faceswaplab_ui/faceswaplab_tab.py

@@ -1,31 +1,28 @@
-import os
+import traceback
-from pprint import pformat, pprint
+from pprint import pformat
-from scripts.faceswaplab_utils import face_utils
+from typing import *
+from scripts.faceswaplab_swapping import face_checkpoints
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
+from scripts.faceswaplab_utils.typing import *
 import gradio as gr
-import modules.scripts as scripts
 import onnx
 import pandas as pd
-from scripts.faceswaplab_ui.faceswaplab_unit_ui import faceswap_unit_ui
-from scripts.faceswaplab_ui.faceswaplab_postprocessing_ui import postprocessing_ui
-from modules import scripts
 from PIL import Image
-from modules.shared import opts
 
-from scripts.faceswaplab_utils import imgutils
-from scripts.faceswaplab_utils.models_utils import get_models
-from scripts.faceswaplab_utils.faceswaplab_logging import logger
 import scripts.faceswaplab_swapping.swapper as swapper
 from scripts.faceswaplab_postprocessing.postprocessing_options import (
     PostProcessingOptions,
 )
-from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
+from scripts.faceswaplab_ui.faceswaplab_postprocessing_ui import postprocessing_ui
-from dataclasses import fields
-from typing import Any, Dict, List, Optional
 from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
-import re
+from scripts.faceswaplab_ui.faceswaplab_unit_ui import faceswap_unit_ui
+from scripts.faceswaplab_utils import imgutils
+from scripts.faceswaplab_utils.faceswaplab_logging import logger
+from scripts.faceswaplab_utils.models_utils import get_swap_models
+from scripts.faceswaplab_utils.ui_utils import dataclasses_from_flat_list
 
 
-def compare(img1: Image.Image, img2: Image.Image) -> str:
+def compare(img1: PILImage, img2: PILImage) -> str:
     """
     Compares the similarity between two faces extracted from images using cosine similarity.
 
@@ -43,14 +40,15 @@ def compare(img1: Image.Image, img2: Image.Image) -> str:
     except Exception as e:
         logger.error("Fail to compare", e)
 
+        traceback.print_exc()
     return "You need 2 images to compare"
 
 
 def extract_faces(
     files: List[gr.File],
     extract_path: Optional[str],
-    *components: List[gr.components.Component],
+    *components: Tuple[gr.components.Component, ...],
-) -> Optional[List[Image.Image]]:
+) -> Optional[List[PILImage]]:
     """
     Extracts faces from a list of image files.
 
@@ -69,22 +67,34 @@ def extract_faces(
                          If no faces are found, None is returned.
     """
 
-    postprocess_options = PostProcessingOptions(*components)  # type: ignore
+    if files and len(files) == 0:
-    images = [
+        logger.error("You need at least one image file to extract")
-        Image.open(file.name) for file in files
+        return []
-    ]  # potentially greedy but Image.open is supposed to be lazy
+    try:
-    return swapper.extract_faces(
+        postprocess_options = dataclasses_from_flat_list(
-        images, extract_path=extract_path, postprocess_options=postprocess_options
+            [PostProcessingOptions], components
-    )
+        ).pop()
+        images = [
+            Image.open(file.name) for file in files  # type: ignore
+        ]  # potentially greedy but Image.open is supposed to be lazy
+        result_images = swapper.extract_faces(
+            images, extract_path=extract_path, postprocess_options=postprocess_options
+        )
+        return result_images
+    except Exception as e:
+        logger.error("Failed to extract : %s", e)
+
+        traceback.print_exc()
+    return None
 
 
-def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> Optional[str]:
+def analyse_faces(image: PILImage, det_threshold: float = 0.5) -> Optional[str]:
     """
     Function to analyze the faces in an image and provide a detailed report.
 
     Parameters
     ----------
-    image : PIL.Image.Image
+    image : PIL.PILImage
         The input image where faces will be detected. The image must be a PIL Image object.
 
     det_threshold : float, optional
@@ -122,27 +132,13 @@ def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> Optional[st
     except Exception as e:
         logger.error("Analysis Failed : %s", e)
 
+        traceback.print_exc()
     return None
 
 
-def sanitize_name(name: str) -> str:
-    """
-    Sanitize the input name by removing special characters and replacing spaces with underscores.
-
-    Parameters:
-        name (str): The input name to be sanitized.
-
-    Returns:
-        str: The sanitized name with special characters removed and spaces replaced by underscores.
-    """
-    name = re.sub("[^A-Za-z0-9_. ]+", "", name)
-    name = name.replace(" ", "_")
-    return name[:255]
-
-
 def build_face_checkpoint_and_save(
-    batch_files: gr.File, name: str
+    batch_files: List[gr.File], name: str, str_gender: str, overwrite: bool
-) -> Optional[Image.Image]:
+) -> PILImage:
     """
     Builds a face checkpoint using the provided image files, performs face swapping,
     and saves the result to a file. If a blended face is successfully obtained and the face swapping
@@ -153,66 +149,28 @@ def build_face_checkpoint_and_save(
         name (str): The name assigned to the face checkpoint.
 
     Returns:
-        PIL.Image.Image or None: The resulting swapped face image if the process is successful; None otherwise.
+        PIL.PILImage or None: The resulting swapped face image if the process is successful; None otherwise.
     """
 
     try:
-        name = sanitize_name(name)
+        if not batch_files:
-        batch_files = batch_files or []
+            logger.error("No face found")
-        logger.info("Build %s %s", name, [x.name for x in batch_files])
+            return None  # type: ignore (Optional not really supported by old gradio)
-        faces = swapper.get_faces_from_img_files(batch_files)
+
-        blended_face = swapper.blend_faces(faces)
+        gender = getattr(Gender, str_gender)
-        preview_path = os.path.join(
+        logger.info("Choosen gender : %s", gender)
-            scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
+        images: list[PILImage] = [Image.open(file.name) for file in batch_files]  # type: ignore
+        preview_image: PILImage | None = (
+            face_checkpoints.build_face_checkpoint_and_save(
+                images=images, name=name, overwrite=overwrite, gender=gender
+            )
         )
-
-        faces_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces")
-
-        os.makedirs(faces_path, exist_ok=True)
-
-        target_img = None
-        if blended_face:
-            if blended_face["gender"] == 0:
-                target_img = Image.open(os.path.join(preview_path, "woman.png"))
-            else:
-                target_img = Image.open(os.path.join(preview_path, "man.png"))
-
-            if name == "":
-                name = "default_name"
-            pprint(blended_face)
-            result = swapper.swap_face(
-                blended_face, blended_face, target_img, get_models()[0]
-            )
-            result_image = enhance_image(
-                result.image,
-                PostProcessingOptions(
-                    face_restorer_name="CodeFormer", restorer_visibility=1
-                ),
-            )
-
-            file_path = os.path.join(faces_path, f"{name}.safetensors")
-            file_number = 1
-            while os.path.exists(file_path):
-                file_path = os.path.join(
-                    faces_path, f"{name}_{file_number}.safetensors"
-                )
-                file_number += 1
-            result_image.save(file_path + ".png")
-
-            face_utils.save_face(filename=file_path, face=blended_face)
-            try:
-                data = face_utils.load_face(filename=file_path)
-                print(data)
-            except Exception as e:
-                print(e)
-            return result_image
-
-        print("No face found")
     except Exception as e:
         logger.error("Failed to build checkpoint %s", e)
-        return None
 
-    return target_img
+        traceback.print_exc()
+        return None  # type: ignore
+    return preview_image  # type: ignore
 
 
 def explore_onnx_faceswap_model(model_path: str) -> pd.DataFrame:
@@ -242,59 +200,55 @@ def explore_onnx_faceswap_model(model_path: str) -> pd.DataFrame:
 
         df = pd.DataFrame(data)
     except Exception as e:
-        logger.info("Failed to explore model %s", e)
+        logger.error("Failed to explore model %s", e)
-        return None
+
+        traceback.print_exc()
+        return None  # type: ignore
     return df
 
 
 def batch_process(
-    files: List[gr.File], save_path: str, *components: List[gr.components.Component]
+    files: List[gr.File], save_path: str, *components: Tuple[Any, ...]
-) -> Optional[List[Image.Image]]:
+) -> List[PILImage]:
     try:
-        units_count = opts.data.get("faceswaplab_units_count", 3)
+        units_count = get_sd_option("faceswaplab_units_count", 3)
-        units: List[FaceSwapUnitSettings] = []
 
-        # Parse and convert units flat components into FaceSwapUnitSettings
+        classes: List[Any] = dataclasses_from_flat_list(
-        for i in range(0, units_count):
+            [FaceSwapUnitSettings] * units_count + [PostProcessingOptions],
-            units += [FaceSwapUnitSettings.get_unit_configuration(i, components)]  # type: ignore
+            components,
-
-        for i, u in enumerate(units):
-            logger.debug("%s, %s", pformat(i), pformat(u))
-
-        # Parse the postprocessing options
-        # We must first find where to start from (after face swapping units)
-        len_conf: int = len(fields(FaceSwapUnitSettings))
-        shift: int = units_count * len_conf
-        postprocess_options = PostProcessingOptions(
-            *components[shift : shift + len(fields(PostProcessingOptions))]  # type: ignore
         )
-        logger.debug("%s", pformat(postprocess_options))
+        units: List[FaceSwapUnitSettings] = [
-        images = [
+            u for u in classes if isinstance(u, FaceSwapUnitSettings)
-            Image.open(file.name) for file in files
+        ]
-        ]  # potentially greedy but Image.open is supposed to be lazy
+        postprocess_options = classes[-1]
-        return swapper.batch_process(
+
-            images,
+        images_paths = [file.name for file in files]  # type: ignore
-            save_path=save_path,
+
-            units=units,
+        return (
-            postprocess_options=postprocess_options,
+            swapper.batch_process(
+                images_paths,
+                save_path=save_path,
+                units=units,
+                postprocess_options=postprocess_options,
+            )
+            or []
         )
     except Exception as e:
         logger.error("Batch Process error : %s", e)
-        import traceback
 
         traceback.print_exc()
-    return None
+    return []
 
 
 def tools_ui() -> None:
-    models = get_models()
+    models = get_swap_models()
     with gr.Tab("Tools"):
         with gr.Tab("Build"):
             gr.Markdown(
                 """Build a face based on a batch list of images. Will blend the resulting face and store the checkpoint in the faceswaplab/faces directory."""
             )
             with gr.Row():
-                batch_files = gr.components.File(
+                build_batch_files = gr.components.File(
                     type="file",
                     file_count="multiple",
                     label="Batch Sources Images",
@@ -304,15 +258,30 @@ def tools_ui() -> None:
                 preview = gr.components.Image(
                     type="pil",
                     label="Preview",
+                    width=512,
+                    height=512,
                     interactive=False,
                     elem_id="faceswaplab_build_preview_face",
                 )
-            name = gr.Textbox(
+            build_name = gr.Textbox(
                 value="Face",
                 placeholder="Name of the character",
                 label="Name of the character",
                 elem_id="faceswaplab_build_character_name",
             )
+            build_gender = gr.Dropdown(
+                value=Gender.AUTO.name,
+                choices=[e.name for e in Gender],
+                placeholder="Gender of the character",
+                label="Gender of the character",
+                elem_id="faceswaplab_build_character_gender",
+            )
+            build_overwrite = gr.Checkbox(
+                False,
+                placeholder="overwrite",
+                label="Overwrite Checkpoint if exist (else will add number)",
+                elem_id="faceswaplab_build_overwrite",
+            )
             generate_checkpoint_btn = gr.Button(
                 "Save", elem_id="faceswaplab_build_save_btn"
             )
@@ -351,7 +320,7 @@ def tools_ui() -> None:
                     label="Extracted faces",
                     show_label=False,
                     elem_id="faceswaplab_extract_results",
-                ).style(columns=[2], rows=[2])
+                )
             extract_save_path = gr.Textbox(
                 label="Destination Directory",
                 value="",
@@ -407,7 +376,7 @@ def tools_ui() -> None:
                     label="Batch result",
                     show_label=False,
                     elem_id="faceswaplab_batch_results",
-                ).style(columns=[2], rows=[2])
+                )
             batch_save_path = gr.Textbox(
                 label="Destination Directory",
                 value="outputs/faceswap/",
@@ -417,7 +386,7 @@ def tools_ui() -> None:
                 "Process & Save", elem_id="faceswaplab_extract_btn"
             )
         unit_components = []
-        for i in range(1, opts.data.get("faceswaplab_units_count", 3) + 1):
+        for i in range(1, get_sd_option("faceswaplab_units_count", 3) + 1):
             unit_components += faceswap_unit_ui(False, i, id_prefix="faceswaplab_tab")
 
     upscale_options = postprocessing_ui()
@@ -427,7 +396,9 @@ def tools_ui() -> None:
     )
     compare_btn.click(compare, inputs=[img1, img2], outputs=[compare_result_text])
     generate_checkpoint_btn.click(
-        build_face_checkpoint_and_save, inputs=[batch_files, name], outputs=[preview]
+        build_face_checkpoint_and_save,
+        inputs=[build_batch_files, build_name, build_gender, build_overwrite],
+        outputs=[preview],
     )
     extract_btn.click(
         extract_faces,

scripts/faceswaplab_ui/faceswaplab_unit_settings.py

@@ -1,15 +1,17 @@
 from scripts.faceswaplab_swapping import swapper
-import numpy as np
 import base64
 import io
-from dataclasses import dataclass, fields
+from dataclasses import dataclass
-from typing import Any, List, Optional, Set, Union
+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_utils
+from scripts.faceswaplab_swapping import face_checkpoints
+from scripts.faceswaplab_inpainting.faceswaplab_inpainting import InpaintingOptions
+from client_api import api_utils
 
 
 @dataclass
@@ -17,11 +19,11 @@ class FaceSwapUnitSettings:
     # ORDER of parameters is IMPORTANT. It should match the result of faceswap_unit_ui
 
     # The image given in reference
-    source_img: Union[Image.Image, str]
+    source_img: Optional[Union[Image.Image, str]]
     # The checkpoint file
-    source_face: str
+    source_face: Optional[str]
     # The batch source images
-    _batch_files: Union[gr.components.File, List[Image.Image]]
+    _batch_files: Optional[Union[gr.components.File, List[Image.Image]]]
     # Will blend faces if True
     blend_faces: bool
     # Enable this unit
@@ -48,14 +50,42 @@ class FaceSwapUnitSettings:
     swap_in_source: bool
     # Swap in the generated image in img2img (always on for txt2img)
     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
 
     @staticmethod
-    def get_unit_configuration(
+    def from_api_dto(dto: api_utils.FaceSwapUnit) -> "FaceSwapUnitSettings":
-        unit: int, components: List[gr.components.Component]
+        """
-    ) -> Any:
+        Converts a InpaintingOptions object from an API DTO (Data Transfer Object).
-        fields_count = len(fields(FaceSwapUnitSettings))
+
+        :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.
+        """
         return FaceSwapUnitSettings(
-            *components[unit * fields_count : unit * fields_count + fields_count]
+            source_img=api_utils.base64_to_pil(dto.source_img),
+            source_face=dto.source_face,
+            _batch_files=dto.get_batch_images(),
+            blend_faces=dto.blend_faces,
+            enable=True,
+            same_gender=dto.same_gender,
+            sort_by_size=dto.sort_by_size,
+            check_similarity=dto.check_similarity,
+            _compute_similarity=dto.compute_similarity,
+            min_ref_sim=dto.min_ref_sim,
+            min_sim=dto.min_sim,
+            _faces_index=",".join([str(i) for i in (dto.faces_index)]),
+            reference_face_index=dto.reference_face_index,
+            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),
         )
 
     @property
@@ -92,14 +122,13 @@ class FaceSwapUnitSettings:
         """
         if not hasattr(self, "_reference_face"):
             if self.source_face and self.source_face != "None":
-                with open(self.source_face, "rb") as file:
+                try:
-                    try:
+                    logger.info(f"loading face {self.source_face}")
-                        logger.info(f"loading face {file.name}")
+                    face = face_checkpoints.load_face(self.source_face)
-                        face = face_utils.load_face(file.name)
+                    self._reference_face = face
-                        self._reference_face = face
+                except Exception as e:
-                    except Exception as e:
+                    logger.error("Failed to load checkpoint  : %s", e)
-                        logger.error("Failed to load checkpoint  : %s", e)
+                    raise e
-                        raise e
             elif self.source_img is not None:
                 if isinstance(self.source_img, str):  # source_img is a base64 string
                     if (
@@ -140,7 +169,7 @@ class FaceSwapUnitSettings:
                 if isinstance(file, Image.Image):
                     img = file
                 else:
-                    img = Image.open(file.name)
+                    img = Image.open(file.name)  # type: ignore
 
                 face = swapper.get_or_default(
                     swapper.get_faces(pil_to_cv2(img)), 0, None
@@ -156,24 +185,5 @@ class FaceSwapUnitSettings:
         """
         if not hasattr(self, "_blended_faces"):
             self._blended_faces = swapper.blend_faces(self.faces)
-            assert (
-                all(
-                    [
-                        not np.array_equal(
-                            self._blended_faces.embedding, face.embedding
-                        )
-                        for face in self.faces
-                    ]
-                )
-                if len(self.faces) > 1
-                else True
-            ), "Blended faces cannot be the same as one of the face if len(face)>0"
-            assert (
-                not np.array_equal(
-                    self._blended_faces.embedding, self.reference_face.embedding
-                )
-                if len(self.faces) > 1
-                else True
-            ), "Blended faces cannot be the same as reference face if len(face)>0"
 
         return self._blended_faces

scripts/faceswaplab_ui/faceswaplab_unit_ui.py

@@ -1,6 +1,100 @@
 from typing import List
-from scripts.faceswaplab_utils.face_utils import get_face_checkpoints
+from scripts.faceswaplab_ui.faceswaplab_inpainting_ui import face_inpainting_ui
+from scripts.faceswaplab_swapping.face_checkpoints import get_face_checkpoints
 import gradio as gr
+from modules import shared
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
+
+
+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. Best result : checks all, use LDSR, use Codeformer"""
+        )
+        with gr.Row():
+            face_restorer_name = gr.Radio(
+                label="Restore Face",
+                choices=["None"] + [x.name() for x in shared.face_restorers],
+                value=get_sd_option(
+                    "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=get_sd_option(
+                        "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=get_sd_option(
+                        "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=get_sd_option("faceswaplab_default_upscaled_swapper_upscaler", ""),
+            label="Upscaler",
+            elem_id=f"{id_prefix}_face{unit_num}_upscaler",
+        )
+
+        improved_mask = gr.Checkbox(
+            get_sd_option("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(
+            get_sd_option("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(
+            get_sd_option("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,
+            get_sd_option("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",
+        )
+
+    components = [
+        face_restorer_name,
+        face_restorer_visibility,
+        codeformer_weight,
+        upscaler_name,
+        improved_mask,
+        color_corrections,
+        sharpen_face,
+        erosion_factor,
+    ]
+
+    for component in components:
+        setattr(component, "do_not_save_to_config", True)
+
+    return components
 
 
 def faceswap_unit_ui(
@@ -40,7 +134,7 @@ def faceswap_unit_ui(
                     elem_id=f"{id_prefix}_face{unit_num}_refresh_checkpoints",
                 )
 
-                def refresh_fn(selected: str) -> None:
+                def refresh_fn(selected: str):
                     return gr.Dropdown.update(
                         value=selected, choices=get_face_checkpoints()
                     )
@@ -61,35 +155,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:"""
@@ -142,22 +207,85 @@ def faceswap_unit_ui(
                 visible=is_img2img,
                 elem_id=f"{id_prefix}_face{unit_num}_swap_in_generated",
             )
+
+        gr.Markdown(
+            """                                        
+## Advanced Options 
+
+**Simple :** If you have bad results and don't want to fine-tune here, just enable Codeformer in "Global Post-Processing".
+Otherwise, read the [doc](https://glucauze.github.io/sd-webui-faceswaplab/doc/) to understand following options.              
+
+"""
+        )
+
+        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",
+            )
+
+        with gr.Accordion(label="Pre-Inpainting (before swapping)", open=False):
+            gr.Markdown("Pre-inpainting sends face to inpainting before swapping")
+            pre_inpainting = face_inpainting_ui(
+                id_prefix=f"{id_prefix}_face{unit_num}_preinpainting",
+            )
+
+        options = faceswap_unit_advanced_options(is_img2img, unit_num, id_prefix)
+
+        with gr.Accordion(label="Post-Inpainting (After swapping)", open=False):
+            gr.Markdown("Pre-inpainting sends face to inpainting before swapping")
+            post_inpainting = face_inpainting_ui(
+                id_prefix=f"{id_prefix}_face{unit_num}_postinpainting",
+            )
+
+    gradio_components: List[gr.components.Component] = (
+        [
+            img,
+            face,
+            batch_files,
+            blend_faces,
+            enable,
+            same_gender,
+            sort_by_size,
+            check_similarity,
+            compute_similarity,
+            min_sim,
+            min_ref_sim,
+            target_faces_index,
+            reference_faces_index,
+            swap_in_source,
+            swap_in_generated,
+        ]
+        + pre_inpainting
+        + options
+        + post_inpainting
+    )
+
     # If changed, you need to change FaceSwapUnitSettings accordingly
     # ORDER of parameters is IMPORTANT. It should match the result of FaceSwapUnitSettings
-    return [
+    return gradio_components
-        img,
-        face,
-        batch_files,
-        blend_faces,
-        enable,
-        same_gender,
-        sort_by_size,
-        check_similarity,
-        compute_similarity,
-        min_sim,
-        min_ref_sim,
-        target_faces_index,
-        reference_faces_index,
-        swap_in_source,
-        swap_in_generated,
-    ]

scripts/faceswaplab_utils/face_utils.py

@@ -1,72 +0,0 @@
-import glob
-import os
-from typing import List
-from insightface.app.common import Face
-from safetensors.torch import save_file, safe_open
-import torch
-
-import modules.scripts as scripts
-from modules import scripts
-from scripts.faceswaplab_utils.faceswaplab_logging import logger
-import dill as pickle  # will be removed in future versions
-
-
-def save_face(face: Face, filename: str) -> None:
-    tensors = {
-        "embedding": torch.tensor(face["embedding"]),
-        "gender": torch.tensor(face["gender"]),
-        "age": torch.tensor(face["age"]),
-    }
-    save_file(tensors, filename)
-
-
-def load_face(filename: str) -> Face:
-    if filename.endswith(".pkl"):
-        logger.warning(
-            "Pkl files for faces are deprecated to enhance safety, they will be unsupported in future versions."
-        )
-        logger.warning("The file will be converted to .safetensors")
-        logger.warning(
-            "You can also use this script https://gist.github.com/glucauze/4a3c458541f2278ad801f6625e5b9d3d"
-        )
-        with open(filename, "rb") as file:
-            logger.info("Load pkl")
-            face = Face(pickle.load(file))
-            logger.warning(
-                "Convert to safetensors, you can remove the pkl version once you have ensured that the safetensor is working"
-            )
-            save_face(face, filename.replace(".pkl", ".safetensors"))
-        return face
-
-    elif filename.endswith(".safetensors"):
-        face = {}
-        with safe_open(filename, 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)
-
-    raise NotImplementedError("Unknown file type, face extraction not implemented")
-
-
-def get_face_checkpoints() -> List[str]:
-    """
-    Retrieve a list of face checkpoint paths.
-
-    This function searches for face files with the extension ".safetensors" in the specified directory and returns a list
-    containing the paths of those files.
-
-    Returns:
-        list: A list of face paths, including the string "None" as the first element.
-    """
-    faces_path = os.path.join(
-        scripts.basedir(), "models", "faceswaplab", "faces", "*.safetensors"
-    )
-    faces = glob.glob(faces_path)
-
-    faces_path = os.path.join(
-        scripts.basedir(), "models", "faceswaplab", "faces", "*.pkl"
-    )
-    faces += glob.glob(faces_path)
-
-    return ["None"] + sorted(faces)

scripts/faceswaplab_utils/imgutils.py

@@ -1,64 +1,76 @@
 import io
-from typing import List, Optional, Tuple, Union, Dict
+from typing import List, Optional, Union, Dict
 from PIL import Image
 import cv2
 import numpy as np
 from math import isqrt, ceil
 import torch
-from ifnude import detect
-from scripts.faceswaplab_globals import NSFW_SCORE_THRESHOLD
 from modules import processing
 import base64
 from collections import Counter
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
+from scripts.faceswaplab_utils.typing import BoxCoords, CV2ImgU8, PILImage
+from scripts.faceswaplab_utils.faceswaplab_logging import logger
 
 
-def check_against_nsfw(img: Image.Image) -> bool:
+def check_against_nsfw(img: PILImage) -> bool:
     """
     Check if an image exceeds the Not Safe for Work (NSFW) score.
 
     Parameters:
-    img (PIL.Image.Image): The image to be checked.
+    img (PILImage): The image to be checked.
 
     Returns:
     bool: True if any part of the image is considered NSFW, False otherwise.
     """
 
+    NSFW_SCORE_THRESHOLD = get_sd_option("faceswaplab_nsfw_threshold", 0.7)
+
+    # For testing purpose :
+    if NSFW_SCORE_THRESHOLD >= 1:
+        return False
+
+    from ifnude import detect
+
     shapes: List[bool] = []
     chunks: List[Dict[str, Union[int, float]]] = detect(img)
 
     for chunk in chunks:
+        logger.debug(
+            f"chunck score {chunk['score']}, threshold : {NSFW_SCORE_THRESHOLD}"
+        )
         shapes.append(chunk["score"] > NSFW_SCORE_THRESHOLD)
 
     return any(shapes)
 
 
-def pil_to_cv2(pil_img: Image.Image) -> np.ndarray:  # type: ignore
+def pil_to_cv2(pil_img: PILImage) -> CV2ImgU8:
     """
     Convert a PIL Image into an OpenCV image (cv2).
 
     Args:
-        pil_img (PIL.Image.Image): An image in PIL format.
+        pil_img (PILImage): An image in PIL format.
 
     Returns:
-        np.ndarray: The input image converted to OpenCV format (BGR).
+        CV2ImgU8: The input image converted to OpenCV format (BGR).
     """
-    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
+    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR).astype("uint8")
 
 
-def cv2_to_pil(cv2_img: np.ndarray) -> Image.Image:  # type: ignore
+def cv2_to_pil(cv2_img: CV2ImgU8) -> PILImage:  # type: ignore
     """
     Convert an OpenCV image (cv2) into a PIL Image.
 
     Args:
-        cv2_img (np.ndarray): An image in OpenCV format (BGR).
+        cv2_img (CV2ImgU8): An image in OpenCV format (BGR).
 
     Returns:
-        PIL.Image.Image: The input image converted to PIL format (RGB).
+        PILImage: The input image converted to PIL format (RGB).
     """
     return Image.fromarray(cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB))
 
 
-def torch_to_pil(images: torch.Tensor) -> List[Image.Image]:
+def torch_to_pil(tensor: torch.Tensor) -> List[PILImage]:
     """
     Converts a tensor image or a batch of tensor images to a PIL image or a list of PIL images.
 
@@ -72,7 +84,7 @@ def torch_to_pil(images: torch.Tensor) -> List[Image.Image]:
     list
         A list of PIL images.
     """
-    images = images.cpu().permute(0, 2, 3, 1).numpy()
+    images: CV2ImgU8 = tensor.cpu().permute(0, 2, 3, 1).numpy()
     if images.ndim == 3:
         images = images[None, ...]
     images = (images * 255).round().astype("uint8")
@@ -80,13 +92,13 @@ def torch_to_pil(images: torch.Tensor) -> List[Image.Image]:
     return pil_images
 
 
-def pil_to_torch(pil_images: Union[Image.Image, List[Image.Image]]) -> torch.Tensor:
+def pil_to_torch(pil_images: Union[PILImage, List[PILImage]]) -> torch.Tensor:
     """
     Converts a PIL image or a list of PIL images to a torch tensor or a batch of torch tensors.
 
     Parameters
     ----------
-    pil_images : Union[Image.Image, List[Image.Image]]
+    pil_images : Union[PILImage, List[PILImage]]
         A PIL image or a list of PIL images.
 
     Returns
@@ -104,7 +116,7 @@ def pil_to_torch(pil_images: Union[Image.Image, List[Image.Image]]) -> torch.Ten
     return torch_image
 
 
-def create_square_image(image_list: List[Image.Image]) -> Optional[Image.Image]:
+def create_square_image(image_list: List[PILImage]) -> Optional[PILImage]:
     """
     Creates a square image by combining multiple images in a grid pattern.
 
@@ -156,33 +168,21 @@ def create_square_image(image_list: List[Image.Image]) -> Optional[Image.Image]:
     return None
 
 
-# def create_mask(image : Image.Image, box_coords : Tuple[int, int, int, int]) -> Image.Image:
-#     width, height = image.size
-#     mask = Image.new("L", (width, height), 255)
-#     x1, y1, x2, y2 = box_coords
-#     for x in range(width):
-#         for y in range(height):
-#             if x1 <= x <= x2 and y1 <= y <= y2:
-#                 mask.putpixel((x, y), 255)
-#             else:
-#                 mask.putpixel((x, y), 0)
-#     return mask
-
-
 def create_mask(
-    image: Image.Image, box_coords: Tuple[int, int, int, int]
+    image: PILImage,
-) -> Image.Image:
+    box_coords: BoxCoords,
+) -> PILImage:
     """
     Create a binary mask for a given image and bounding box coordinates.
 
     Args:
-        image (PIL.Image.Image): The input image.
+        image (PILImage): The input image.
         box_coords (Tuple[int, int, int, int]): A tuple of 4 integers defining the bounding box.
         It follows the pattern (x1, y1, x2, y2), where (x1, y1) is the top-left coordinate of the
         box and (x2, y2) is the bottom-right coordinate of the box.
 
     Returns:
-        PIL.Image.Image: A binary mask of the same size as the input image, where pixels within
+        PILImage: A binary mask of the same size as the input image, where pixels within
         the bounding box are white (255) and pixels outside the bounding box are black (0).
     """
     width, height = image.size
@@ -195,8 +195,8 @@ def create_mask(
 
 
 def apply_mask(
-    img: Image.Image, p: processing.StableDiffusionProcessing, batch_index: int
+    img: PILImage, p: processing.StableDiffusionProcessing, batch_index: int
-) -> Image.Image:
+) -> PILImage:
     """
     Apply mask overlay and color correction to an image if enabled
 
@@ -213,8 +213,10 @@ def apply_mask(
             overlays = p.overlay_images
             if overlays is None or batch_index >= len(overlays):
                 return img
-            overlay: Image.Image = overlays[batch_index]
+            overlay: PILImage = overlays[batch_index]
-            overlay = overlay.resize((img.size), resample=Image.Resampling.LANCZOS)
+            logger.debug("Overlay size %s, Image size %s", overlay.size, img.size)
+            if overlay.size != img.size:
+                overlay = overlay.resize((img.size), resample=Image.Resampling.LANCZOS)
             img = img.copy()
             img.paste(overlay, (0, 0), overlay)
             return img
@@ -227,9 +229,7 @@ def apply_mask(
     return img
 
 
-def prepare_mask(
+def prepare_mask(mask: PILImage, p: processing.StableDiffusionProcessing) -> PILImage:
-    mask: Image.Image, p: processing.StableDiffusionProcessing
-) -> Image.Image:
     """
     Prepare an image mask for the inpainting process. (This comes from controlnet)
 
@@ -243,12 +243,12 @@ def prepare_mask(
        apply a Gaussian blur to the mask with a radius equal to 'mask_blur'.
 
     Args:
-        mask (Image.Image): The input mask as a PIL Image object.
+        mask (PILImage): The input mask as a PIL Image object.
         p (processing.StableDiffusionProcessing): An instance of the StableDiffusionProcessing class
                                                    containing the processing parameters.
 
     Returns:
-        mask (Image.Image): The prepared mask as a PIL Image object.
+        mask (PILImage): The prepared mask as a PIL Image object.
     """
     mask = mask.convert("L")
     # FIXME : Properly fix blur
@@ -257,7 +257,7 @@ def prepare_mask(
     return mask
 
 
-def base64_to_pil(base64str: Optional[str]) -> Optional[Image.Image]:
+def base64_to_pil(base64str: Optional[str]) -> Optional[PILImage]:
     """
     Converts a base64 string to a PIL Image object.
 
@@ -267,7 +267,7 @@ def base64_to_pil(base64str: Optional[str]) -> Optional[Image.Image]:
     will return None.
 
     Returns:
-    Optional[Image.Image]: A PIL Image object created from the base64 string. If the input is None,
+    Optional[PILImage]: A PIL Image object created from the base64 string. If the input is None,
     the function returns None.
 
     Raises:

scripts/faceswaplab_utils/install_utils.py

@@ -0,0 +1,18 @@
+from types import ModuleType
+
+
+def check_install() -> None:
+    # Very ugly hack :( due to sdnext optimization not calling install.py every time if git log has not changed
+    import importlib.util
+    import sys
+    import os
+
+    current_dir = os.path.dirname(os.path.realpath(__file__))
+    check_install_path = os.path.join(current_dir, "..", "..", "install.py")
+    spec = importlib.util.spec_from_file_location("check_install", check_install_path)
+    if spec != None:
+        check_install: ModuleType = importlib.util.module_from_spec(spec)
+        sys.modules["check_install"] = check_install
+        spec.loader.exec_module(check_install)  # type: ignore
+        check_install.check_install()  # type: ignore
+        #### End of ugly hack :( !

scripts/faceswaplab_utils/models_utils.py

@@ -3,12 +3,47 @@ import os
 from typing import List
 import modules.scripts as scripts
 from modules import scripts
-from scripts.faceswaplab_globals import EXTENSION_PATH
+from scripts.faceswaplab_globals import EXPECTED_INSWAPPER_SHA1, EXTENSION_PATH
 from modules.shared import opts
 from scripts.faceswaplab_utils.faceswaplab_logging import logger
+import traceback
+import hashlib
 
 
-def get_models() -> List[str]:
+def is_sha1_matching(file_path: str, expected_sha1: str) -> bool:
+    sha1_hash = hashlib.sha1(usedforsecurity=False)
+    try:
+        with open(file_path, "rb") as file:
+            for byte_block in iter(lambda: file.read(4096), b""):
+                sha1_hash.update(byte_block)
+            if sha1_hash.hexdigest() == expected_sha1:
+                return True
+            else:
+                return False
+    except Exception as e:
+        logger.error(
+            "Failed to check model hash, check the model is valid or has been downloaded adequately : %e",
+            e,
+        )
+        traceback.print_exc()
+        return False
+
+
+def check_model() -> bool:
+    model_path = get_current_swap_model()
+    if not is_sha1_matching(
+        file_path=model_path, expected_sha1=EXPECTED_INSWAPPER_SHA1
+    ):
+        logger.error(
+            "Suspicious sha1 for model %s, check the model is valid or has been downloaded adequately. Should be %s",
+            model_path,
+            EXPECTED_INSWAPPER_SHA1,
+        )
+        return False
+    return True
+
+
+def get_swap_models() -> List[str]:
     """
     Retrieve a list of swap model files.
 
@@ -31,15 +66,22 @@ def get_models() -> List[str]:
     return models
 
 
-def get_current_model() -> str:
+def get_current_swap_model() -> str:
-    model = opts.data.get("faceswaplab_model", None)
+    model = opts.data.get("faceswaplab_model", None)  # type: ignore
     if model is None:
-        models = get_models()
+        models = get_swap_models()
         model = models[0] if len(models) else None
     logger.info("Try to use model : %s", model)
-    if not os.path.isfile(model):
+    try:
-        logger.error("The model %s cannot be found or loaded", model)
+        if not model or not os.path.isfile(model):  # type: ignore
+            logger.error("The model %s cannot be found or loaded", model)
+            raise FileNotFoundError(
+                "No faceswap model found. Please add it to the faceswaplab directory. Ensure the model is in the proper directory (<sdwebui>/models/faceswaplab/inswapper_128.onnx)"
+            )
+    except:
         raise FileNotFoundError(
-            "No faceswap model found. Please add it to the faceswaplab directory."
+            "No faceswap model found. Please add it to the faceswaplab directory. Ensure the model is in the proper directory (<sdwebui>/models/faceswaplab/inswapper_128.onnx)"
         )
+
+    assert model is not None
     return model

scripts/faceswaplab_utils/sd_utils.py

@@ -0,0 +1,7 @@
+from typing import Any
+from modules.shared import opts
+
+
+def get_sd_option(name: str, default: Any) -> Any:
+    assert opts.data is not None
+    return opts.data.get(name, default)

scripts/faceswaplab_utils/typing.py

@@ -0,0 +1,17 @@
+from typing import Tuple
+from numpy import uint8
+from insightface.app.common import Face as IFace
+from PIL import Image
+import numpy as np
+from enum import Enum
+
+PILImage = Image.Image
+CV2ImgU8 = np.ndarray[int, np.dtype[uint8]]
+Face = IFace
+BoxCoords = Tuple[int, int, int, int]
+
+
+class Gender(Enum):
+    AUTO = -1
+    FEMALE = 0
+    MALE = 1

scripts/faceswaplab_utils/ui_utils.py

@@ -0,0 +1,43 @@
+from dataclasses import fields, is_dataclass
+from typing import *
+
+
+def dataclass_from_flat_list(cls: type, values: Tuple[Any, ...]) -> Any:
+    if not is_dataclass(cls):
+        raise TypeError(f"{cls} is not a dataclass")
+
+    idx = 0
+    init_values = {}
+    for field in fields(cls):
+        if is_dataclass(field.type):
+            inner_values = [values[idx + i] for i in range(len(fields(field.type)))]
+            init_values[field.name] = field.type(*inner_values)
+            idx += len(inner_values)
+        else:
+            if idx >= len(values):
+                raise IndexError(
+                    f"Expected more values for dataclass {cls}. Current index: {idx}, values length: {len(values)}"
+                )
+            value = values[idx]
+            init_values[field.name] = value
+            idx += 1
+    return cls(**init_values)
+
+
+def dataclasses_from_flat_list(
+    classes_mapping: List[type], values: Tuple[Any, ...]
+) -> List[Any]:
+    instances = []
+    idx = 0
+    for cls in classes_mapping:
+        num_fields = sum(
+            len(fields(field.type)) if is_dataclass(field.type) else 1
+            for field in fields(cls)
+        )
+        instance = dataclass_from_flat_list(cls, values[idx : idx + num_fields])
+        instances.append(instance)
+        idx += num_fields
+    assert [
+        isinstance(i, t) for i, t in zip(instances, classes_mapping)
+    ], "Instances should match types"
+    return instances

tests/test_api.py

@@ -2,21 +2,28 @@ from typing import List
 import pytest
 import requests
 import sys
+import tempfile
+import safetensors
 
 sys.path.append(".")
 
+import requests
 from client_api.api_utils import (
     FaceSwapUnit,
-    FaceSwapResponse,
+    InswappperOptions,
-    PostProcessingOptions,
-    FaceSwapRequest,
-    base64_to_pil,
     pil_to_base64,
+    PostProcessingOptions,
     InpaintingWhen,
-    FaceSwapCompareRequest,
+    InpaintingOptions,
+    FaceSwapRequest,
+    FaceSwapResponse,
     FaceSwapExtractRequest,
+    FaceSwapCompareRequest,
     FaceSwapExtractResponse,
     compare_faces,
+    base64_to_pil,
+    base64_to_safetensors,
+    safetensors_to_base64,
 )
 from PIL import Image
 
@@ -36,6 +43,13 @@ def face_swap_request() -> FaceSwapRequest:
         source_img=pil_to_base64("references/woman.png"),  # The face you want to use
         same_gender=True,
         faces_index=(0,),  # Replace first woman since same gender is on
+        swapping_options=InswappperOptions(
+            face_restorer_name="CodeFormer",
+            upscaler_name="LDSR",
+            improved_mask=True,
+            sharpen=True,
+            color_corrections=True,
+        ),
     )
 
     # Post-processing config
@@ -45,11 +59,12 @@ def face_swap_request() -> FaceSwapRequest:
         restorer_visibility=1,
         upscaler_name="Lanczos",
         scale=4,
-        inpainting_steps=30,
-        inpainting_denoising_strengh=0.1,
         inpainting_when=InpaintingWhen.BEFORE_RESTORE_FACE,
+        inpainting_options=InpaintingOptions(
+            inpainting_steps=30,
+            inpainting_denoising_strengh=0.1,
+        ),
     )
-
     # Prepare the request
     request = FaceSwapRequest(
         image=pil_to_base64("tests/test_image.png"),
@@ -149,3 +164,114 @@ def test_faceswap(face_swap_request: FaceSwapRequest) -> None:
     assert response.status_code == 200
     similarity = float(response.text)
     assert similarity > 0.50
+
+
+def test_faceswap_inpainting(face_swap_request: FaceSwapRequest) -> None:
+    face_swap_request.units[0].pre_inpainting = InpaintingOptions(
+        inpainting_denoising_strengh=0.4,
+        inpainting_prompt="Photo of a funny man",
+        inpainting_negative_prompt="blurry, bad art",
+        inpainting_steps=100,
+    )
+
+    face_swap_request.units[0].post_inpainting = InpaintingOptions(
+        inpainting_denoising_strengh=0.4,
+        inpainting_prompt="Photo of a funny man",
+        inpainting_negative_prompt="blurry, bad art",
+        inpainting_steps=20,
+        inpainting_sampler="Euler a",
+    )
+
+    response = requests.post(
+        f"{base_url}/faceswaplab/swap_face",
+        data=face_swap_request.json(),
+        headers={"Content-Type": "application/json; charset=utf-8"},
+    )
+
+    assert response.status_code == 200
+    data = response.json()
+    assert "images" in data
+    assert "infos" in data
+
+
+def test_faceswap_checkpoint_building() -> None:
+    source_images: List[str] = [
+        pil_to_base64("references/man.png"),
+        pil_to_base64("references/woman.png"),
+    ]
+
+    response = requests.post(
+        url=f"{base_url}/faceswaplab/build",
+        json=source_images,
+        headers={"Content-Type": "application/json; charset=utf-8"},
+    )
+
+    assert response.status_code == 200
+
+    with tempfile.NamedTemporaryFile(delete=True) as temp_file:
+        base64_to_safetensors(response.json(), output_path=temp_file.name)
+        with safetensors.safe_open(temp_file.name, framework="pt") as f:
+            assert "age" in f.keys()
+            assert "gender" in f.keys()
+            assert "embedding" in f.keys()
+
+
+def test_faceswap_checkpoint_building_and_using() -> None:
+    source_images: List[str] = [
+        pil_to_base64("references/man.png"),
+    ]
+
+    response = requests.post(
+        url=f"{base_url}/faceswaplab/build",
+        json=source_images,
+        headers={"Content-Type": "application/json; charset=utf-8"},
+    )
+
+    assert response.status_code == 200
+
+    with tempfile.NamedTemporaryFile(delete=True) as temp_file:
+        base64_to_safetensors(response.json(), output_path=temp_file.name)
+        with safetensors.safe_open(temp_file.name, framework="pt") as f:
+            assert "age" in f.keys()
+            assert "gender" in f.keys()
+            assert "embedding" in f.keys()
+
+        # First face unit :
+        unit1 = FaceSwapUnit(
+            source_face=safetensors_to_base64(
+                temp_file.name
+            ),  # convert the checkpoint to base64
+            faces_index=(0,),  # Replace first face
+            swapping_options=InswappperOptions(
+                face_restorer_name="CodeFormer",
+                upscaler_name="LDSR",
+                improved_mask=True,
+                sharpen=True,
+                color_corrections=True,
+            ),
+        )
+
+        # Prepare the request
+        request = FaceSwapRequest(
+            image=pil_to_base64("tests/test_image.png"), units=[unit1]
+        )
+
+        # Face Swap
+        response = requests.post(
+            url=f"{base_url}/faceswaplab/swap_face",
+            data=request.json(),
+            headers={"Content-Type": "application/json; charset=utf-8"},
+        )
+        assert response.status_code == 200
+        fsr = FaceSwapResponse.parse_obj(response.json())
+        data = response.json()
+        assert "images" in data
+        assert "infos" in data
+
+        # First face is the man
+        assert (
+            compare_faces(
+                fsr.pil_images[0], Image.open("references/man.png"), base_url=base_url
+            )
+            > 0.5
+        )