fix bugs, fix api, clear code

2023-07-29 00:24:51 +02:00
parent a511214aaa
commit d755be1325
11 changed files with 450 additions and 254 deletions
@@ -5,6 +5,7 @@ from enum import Enum
 import base64, io
 from io import BytesIO
 from typing import List, Tuple, Optional
+import numpy as np


 class InpaintingWhen(Enum):
@@ -41,7 +42,10 @@ class FaceSwapUnit(BaseModel):
    blend_faces: bool = Field(description="Will blend faces if True", default=True)

    # Use same gender filtering
-    same_gender: bool = Field(description="Use same gender filtering", default=True)
+    same_gender: bool = Field(description="Use same gender filtering", default=False)
+
+    # Use same gender filtering
+    sort_by_size: bool = Field(description="Sort Faces by size", default=False)

    # If True, discard images with low similarity
    check_similarity: bool = Field(
@@ -70,6 +74,18 @@ class FaceSwapUnit(BaseModel):
        default=(0,),
    )

+    reference_face_index: int = Field(
+        description="The face index to use to extract face from reference",
+        default=0,
+    )
+
+    def get_batch_images(self) -> List[Image.Image]:
+        images = []
+        if self.batch_images:
+            for img in self.batch_images:
+                images.append(base64_to_pil(img))
+        return images
+

 class PostProcessingOptions(BaseModel):
    face_restorer_name: str = Field(description="face restorer name", default=None)
@@ -82,7 +98,7 @@ class PostProcessingOptions(BaseModel):

    upscaler_name: str = Field(description="upscaler name", default=None)
    scale: float = Field(description="upscaling scale", default=1, le=10, ge=0)
-    upscale_visibility: float = Field(
+    upscaler_visibility: float = Field(
        description="upscaler visibility", default=1, le=1, ge=0
    )

@@ -116,6 +132,9 @@ class PostProcessingOptions(BaseModel):
        examples=[e.value for e in InpaintingWhen.__members__.values()],
        default=InpaintingWhen.NEVER,
    )
+    inpainting_model: str = Field(
+        description="Inpainting model", examples=["Current"], default="Current"
+    )


 class FaceSwapRequest(BaseModel):
@@ -125,7 +144,7 @@ class FaceSwapRequest(BaseModel):
        default=None,
    )
    units: List[FaceSwapUnit]
-    postprocessing: PostProcessingOptions
+    postprocessing: Optional[PostProcessingOptions]


 class FaceSwapResponse(BaseModel):
@@ -133,11 +152,11 @@ class FaceSwapResponse(BaseModel):
    infos: List[str]

    @property
-    def pil_images(self):
+    def pil_images(self) -> Image.Image:
        return [base64_to_pil(img) for img in self.images]


-def pil_to_base64(img):
+def pil_to_base64(img: Image.Image) -> np.array:  # type:ignore
    if isinstance(img, str):
        img = Image.open(img)

@@ -5,6 +5,7 @@ from client_utils import (
    PostProcessingOptions,
    FaceSwapResponse,
    pil_to_base64,
+    InpaintingWhen,
 )

 address = "http://127.0.0.1:7860"
@@ -24,7 +25,14 @@ unit2 = FaceSwapUnit(

 # Post-processing config :
 pp = PostProcessingOptions(
-    face_restorer_name="CodeFormer", codeformer_weight=0.5, restorer_visibility=1
+    face_restorer_name="CodeFormer",
+    codeformer_weight=0.5,
+    restorer_visibility=1,
+    upscaler_name="Lanczos",
+    scale=4,
+    inpainting_steps=30,
+    inpainting_denoising_strengh=0.1,
+    inpainting_when=InpaintingWhen.BEFORE_RESTORE_FACE,
 )

 # Prepare the request
@@ -146,105 +146,111 @@ class FaceSwapScript(scripts.Script):
    def process(
        self, p: StableDiffusionProcessing, *components: List[gr.components.Component]
    ) -> None:
-        self.read_config(p, *components)
+        try:
+            self.read_config(p, *components)

-        # 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]]] = [
-                    (img, None) for img in p.init_images
-                ]
-                new_inits = swapper.process_images_units(
-                    get_current_model(),
-                    self.swap_in_source_units,
-                    images=init_images,
-                    upscaled_swapper=self.upscaled_swapper_in_source,
-                    force_blend=True,
-                )
-                logger.info(f"processed init images: {len(init_images)}")
-                if new_inits is not None:
-                    p.init_images = [img[0] for img in new_inits]
+            # 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]]] = [
+                        (img, None) for img in p.init_images
+                    ]
+                    new_inits = swapper.process_images_units(
+                        get_current_model(),
+                        self.swap_in_source_units,
+                        images=init_images,
+                        upscaled_swapper=self.upscaled_swapper_in_source,
+                        force_blend=True,
+                    )
+                    logger.info(f"processed init images: {len(init_images)}")
+                    if new_inits is not None:
+                        p.init_images = [img[0] for img in new_inits]
+        except Exception as e:
+            logger.info("Failed to process : %s", e)

    def postprocess(
        self, p: StableDiffusionProcessing, processed: Processed, *args: List[Any]
    ) -> None:
-        if self.enabled:
-            # Get the original images without the grid
-            orig_images: List[Image.Image] = processed.images[
-                processed.index_of_first_image :
-            ]
-            orig_infotexts: List[str] = processed.infotexts[
-                processed.index_of_first_image :
-            ]
+        try:
+            if self.enabled:
+                # Get the original images without the grid
+                orig_images: List[Image.Image] = processed.images[
+                    processed.index_of_first_image :
+                ]
+                orig_infotexts: List[str] = processed.infotexts[
+                    processed.index_of_first_image :
+                ]

-            keep_original = self.keep_original_images
+                keep_original = self.keep_original_images

-            # These are were images and infos of swapped images will be stored
-            images = []
-            infotexts = []
-            if (len(self.swap_in_generated_units)) > 0:
-                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(),
-                        self.swap_in_generated_units,
-                        images=[(img, info)],
-                        upscaled_swapper=self.upscaled_swapper_in_generated,
-                    )
-                    if swapped_images is None:
-                        continue
+                # These are were images and infos of swapped images will be stored
+                images = []
+                infotexts = []
+                if (len(self.swap_in_generated_units)) > 0:
+                    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(),
+                            self.swap_in_generated_units,
+                            images=[(img, info)],
+                            upscaled_swapper=self.upscaled_swapper_in_generated,
+                        )
+                        if swapped_images is None:
+                            continue

-                    logger.info(f"{len(swapped_images)} images swapped")
-                    for swp_img, new_info in swapped_images:
-                        img = swp_img  # Will only swap the last image in the batch in next units (FIXME : hard to fix properly but not really critical)
+                        logger.info(f"{len(swapped_images)} images swapped")
+                        for swp_img, new_info in swapped_images:
+                            img = swp_img  # Will only swap the last image in the batch in next units (FIXME : hard to fix properly but not really critical)

-                        if swp_img is not None:
-                            save_img_debug(swp_img, "Before apply mask")
-                            swp_img = imgutils.apply_mask(swp_img, p, batch_index)
-                            save_img_debug(swp_img, "After apply mask")
+                            if swp_img is not None:
+                                save_img_debug(swp_img, "Before apply mask")
+                                swp_img = imgutils.apply_mask(swp_img, p, batch_index)
+                                save_img_debug(swp_img, "After apply mask")

-                            try:
-                                if self.postprocess_options is not None:
-                                    swp_img = enhance_image(
-                                        swp_img, self.postprocess_options
+                                try:
+                                    if self.postprocess_options is not None:
+                                        swp_img = enhance_image(
+                                            swp_img, self.postprocess_options
+                                        )
+                                except Exception as e:
+                                    logger.error("Failed to upscale : %s", e)
+
+                                logger.info("Add swp image to processed")
+                                images.append(swp_img)
+                                infotexts.append(new_info)
+                                if p.outpath_samples and opts.samples_save:
+                                    save_image(
+                                        swp_img,
+                                        p.outpath_samples,
+                                        "",
+                                        p.all_seeds[batch_index],
+                                        p.all_prompts[batch_index],
+                                        opts.samples_format,
+                                        info=new_info,
+                                        p=p,
+                                        suffix="-swapped",
                                    )
-                            except Exception as e:
-                                logger.error("Failed to upscale : %s", e)
+                            else:
+                                logger.error("swp image is None")
+                else:
+                    keep_original = True

-                            logger.info("Add swp image to processed")
-                            images.append(swp_img)
-                            infotexts.append(new_info)
-                            if p.outpath_samples and opts.samples_save:
-                                save_image(
-                                    swp_img,
-                                    p.outpath_samples,
-                                    "",
-                                    p.all_seeds[batch_index],
-                                    p.all_prompts[batch_index],
-                                    opts.samples_format,
-                                    info=new_info,
-                                    p=p,
-                                    suffix="-swapped",
-                                )
-                        else:
-                            logger.error("swp image is None")
-            else:
-                keep_original = True
+                # 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
+                    images.insert(0, grid)

-            # 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
-                images.insert(0, grid)
+                if keep_original:
+                    # If we want to keep original images, we add all existing (including grid this time)
+                    images += processed.images
+                    infotexts += processed.infotexts

-            if keep_original:
-                # If we want to keep original images, we add all existing (including grid this time)
-                images += processed.images
-                infotexts += processed.infotexts
-
-            processed.images = images
-            processed.infotexts = infotexts
+                processed.images = images
+                processed.infotexts = infotexts
+        except Exception as e:
+            logger.error("Failed to swap face %s in postprocess method", e)
@@ -3,7 +3,6 @@ import numpy as np
 from fastapi import FastAPI
 from modules.api import api
 from scripts.faceswaplab_api.faceswaplab_api_types import (
-    FaceSwapRequest,
    FaceSwapResponse,
 )
 from scripts.faceswaplab_globals import VERSION_FLAG
@@ -16,9 +15,15 @@ from scripts.faceswaplab_utils.imgutils import (
 )
 from scripts.faceswaplab_utils.models_utils import get_current_model
 from modules.shared import opts
+from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
+from scripts.faceswaplab_postprocessing.postprocessing_options import (
+    PostProcessingOptions,
+)
+from scripts.faceswaplab_api import faceswaplab_api_types
+from scripts.faceswaplab_postprocessing.postprocessing_options import InpaintingWhen


-def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str:
+def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str:  # type: ignore
    """
    Encode an image to a base64 string.

@@ -40,7 +45,7 @@ def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str:
        return ""


-def encode_np_to_base64(image: np.ndarray) -> str:
+def encode_np_to_base64(image: np.ndarray) -> str:  # type: ignore
    """
    Encode a NumPy array to a base64 string.

@@ -56,6 +61,59 @@ def encode_np_to_base64(image: np.ndarray) -> str:
    return api.encode_pil_to_base64(pil)


+def get_postprocessing_options(
+    options: faceswaplab_api_types.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,
+        inpainting_when=options.inpainting_when,
+        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[faceswaplab_api_types.FaceSwapUnit],
+) -> List[FaceSwapUnitSettings]:
+    units = []
+    for u in api_units:
+        units.append(
+            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
+
+
 def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None:
    @app.get(
        "/faceswaplab/version",
@@ -71,29 +129,17 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None:
        tags=["faceswaplab"],
        description="Swap a face in an image using units",
    )
-    async def swap_face(request: FaceSwapRequest) -> FaceSwapResponse:
+    async def swap_face(
+        request: faceswaplab_api_types.FaceSwapRequest,
+    ) -> faceswaplab_api_types.FaceSwapResponse:
        units: List[FaceSwapUnitSettings] = []
        src_image: Optional[Image.Image] = base64_to_pil(request.image)
        response = FaceSwapResponse(images=[], infos=[])
+        if request.postprocessing:
+            pp_options = get_postprocessing_options(request.postprocessing)
+
        if src_image is not None:
-            for u in request.units:
-                units.append(
-                    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,
-                        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)]),
-                        swap_in_generated=True,
-                        swap_in_source=False,
-                    )
-                )
+            units = get_faceswap_units_settings(request.units)

            swapped_images = swapper.process_images_units(
                get_current_model(),
@@ -102,6 +148,8 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None:
                upscaled_swapper=opts.data.get("faceswaplab_upscaled_swapper", False),
            )
            for img, info in swapped_images:
+                if pp_options:
+                    img = enhance_image(img, pp_options)
                response.images.append(encode_to_base64(img))
                response.infos.append(info)

@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 from PIL import Image
 from scripts.faceswaplab_utils.imgutils import (
    base64_to_pil,
@@ -34,7 +34,10 @@ class FaceSwapUnit(BaseModel):
    blend_faces: bool = Field(description="Will blend faces if True", default=True)

    # Use same gender filtering
-    same_gender: bool = Field(description="Use same gender filtering", default=True)
+    same_gender: bool = Field(description="Use same gender filtering", default=False)
+
+    # Use same gender filtering
+    sort_by_size: bool = Field(description="Sort Faces by size", default=False)

    # If True, discard images with low similarity
    check_similarity: bool = Field(
@@ -63,6 +66,11 @@ class FaceSwapUnit(BaseModel):
        default=(0,),
    )

+    reference_face_index: int = Field(
+        description="The face index to use to extract face from reference",
+        default=0,
+    )
+
    def get_batch_images(self) -> List[Image.Image]:
        images = []
        if self.batch_images:
@@ -82,7 +90,7 @@ class PostProcessingOptions(BaseModel):

    upscaler_name: str = Field(description="upscaler name", default=None)
    scale: float = Field(description="upscaling scale", default=1, le=10, ge=0)
-    upscale_visibility: float = Field(
+    upscaler_visibility: float = Field(
        description="upscaler visibility", default=1, le=1, ge=0
    )

@@ -116,6 +124,9 @@ class PostProcessingOptions(BaseModel):
        examples=[e.value for e in InpaintingWhen.__members__.values()],
        default=InpaintingWhen.NEVER,
    )
+    inpainting_model: str = Field(
+        description="Inpainting model", examples=["Current"], default="Current"
+    )


 class FaceSwapRequest(BaseModel):
@@ -125,7 +136,7 @@ class FaceSwapRequest(BaseModel):
        default=None,
    )
    units: List[FaceSwapUnit]
-    postprocessing: PostProcessingOptions
+    postprocessing: Optional[PostProcessingOptions]


 class FaceSwapResponse(BaseModel):
@@ -1,8 +1,12 @@
 import os
+from modules import scripts

 MODELS_DIR = os.path.abspath(os.path.join("models", "faceswaplab"))
 ANALYZER_DIR = os.path.abspath(os.path.join(MODELS_DIR, "analysers"))
 FACE_PARSER_DIR = os.path.abspath(os.path.join(MODELS_DIR, "parser"))
+REFERENCE_PATH = os.path.join(
+    scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
+)

 VERSION_FLAG: str = "v1.1.0"
 EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab")
@@ -14,6 +14,7 @@ from scripts.faceswaplab_swapping import swapper

 def img2img_diffusion(img: Image.Image, pp: PostProcessingOptions) -> Image.Image:
    if pp.inpainting_denoising_strengh == 0:
+        logger.info("Discard inpainting denoising strength is 0")
        return img

    try:
@@ -25,7 +26,7 @@ inpainting_steps : {pp.inpainting_steps}
 """
        )
        if not isinstance(pp.inpainting_sampler, str):
-            pass
+            pp.inpainting_sampler = "Euler"

        logger.info("send faces to image to image")
        img = img.copy()
@@ -11,17 +11,32 @@ from scripts.faceswaplab_postprocessing.upscaling import upscale_img, restore_fa
 def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image:
    result_image = image
    try:
-        if pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING.value:
-            result_image = img2img_diffusion(image, pp_options)
+        logger.debug("enhance_image, inpainting : %s", pp_options.inpainting_when)
+        result_image = image
+
+        if (
+            pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING.value
+            or pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING
+        ):
+            logger.debug("Inpaint before upscale")
+            result_image = img2img_diffusion(result_image, pp_options)
        result_image = upscale_img(result_image, pp_options)

-        if pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE.value:
-            result_image = img2img_diffusion(image, pp_options)
+        if (
+            pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE.value
+            or pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE
+        ):
+            logger.debug("Inpaint before restore")
+            result_image = img2img_diffusion(result_image, pp_options)

        result_image = restore_face(result_image, pp_options)

-        if pp_options.inpainting_when == InpaintingWhen.AFTER_ALL.value:
-            result_image = img2img_diffusion(image, pp_options)
+        if (
+            pp_options.inpainting_when == InpaintingWhen.AFTER_ALL.value
+            or pp_options.inpainting_when == InpaintingWhen.AFTER_ALL
+        ):
+            logger.debug("Inpaint after all")
+            result_image = img2img_diffusion(result_image, pp_options)

    except Exception as e:
        logger.error("Failed to upscale %s", e)
@@ -19,7 +19,7 @@ class PostProcessingOptions:
    codeformer_weight: float = 1

    upscaler_name: str = ""
-    scale: int = 1
+    scale: float = 1
    upscale_visibility: float = 0.5

    inpainting_denoising_strengh: float = 0
@@ -24,61 +24,132 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
 )
 from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
 from dataclasses import fields
-from typing import Any, List, Optional, Union
+from typing import Any, Dict, List, Optional
 from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
 from scripts.faceswaplab_utils.models_utils import get_current_model
 import re
+from scripts.faceswaplab_globals import REFERENCE_PATH


-def compare(img1: Image.Image, img2: Image.Image) -> Union[float, str]:
-    if img1 is not None and img2 is not None:
-        return swapper.compare_faces(img1, img2)
+def compare(img1: Image.Image, img2: Image.Image) -> str:
+    """
+    Compares the similarity between two faces extracted from images using cosine similarity.
+
+    Args:
+        img1: The first image containing a face.
+        img2: The second image containing a face.
+
+    Returns:
+        A str of a float value representing the similarity between the two faces (0 to 1).
+        Returns"You need 2 images to compare" if one or both of the images do not contain any faces.
+    """
+    try:
+        if img1 is not None and img2 is not None:
+            return str(swapper.compare_faces(img1, img2))
+    except Exception as e:
+        logger.error("Fail to compare", e)

    return "You need 2 images to compare"


 def extract_faces(
-    files,
-    extract_path,
+    files: Optional[List[str]],
+    extract_path: Optional[str],
    *components: List[gr.components.Component],
-):
-    postprocess_options = PostProcessingOptions(*components)  # type: ignore
+) -> Optional[List[str]]:
+    """
+    Extracts faces from a list of image files.
+
+    Given a list of image file paths, this function opens each image, extracts the faces,
+    and saves them in a specified directory. Post-processing is applied to each extracted face,
+    and the processed faces are saved as separate PNG files.
+
+    Parameters:
+    files (Optional[List[str]]): List of file paths to the images to extract faces from.
+    extract_path (Optional[str]): Path where the extracted faces will be saved.
+                                  If no path is provided, a temporary directory will be created.
+    components (List[gr.components.Component]): List of components for post-processing.
+
+    Returns:
+    Optional[List[str]]: List of file paths to the saved images of the extracted faces.
+                         If no faces are found, None is returned.
+    """
+
+    try:
+        postprocess_options = PostProcessingOptions(*components)  # type: ignore
+
+        if not extract_path:
+            extract_path = tempfile.mkdtemp()
+
+        if files:
+            images = []
+            for file in files:
+                img = Image.open(file).convert("RGB")
+                faces = swapper.get_faces(pil_to_cv2(img))
+
+                if faces:
+                    face_images = []
+                    for face in faces:
+                        bbox = face.bbox.astype(int)
+                        x_min, y_min, x_max, y_max = bbox
+                        face_image = img.crop((x_min, y_min, x_max, y_max))
+
+                        if (
+                            postprocess_options.face_restorer_name
+                            or postprocess_options.restorer_visibility
+                        ):
+                            postprocess_options.scale = (
+                                1 if face_image.width > 512 else 512 // face_image.width
+                            )
+                            face_image = enhance_image(face_image, postprocess_options)
+
+                        path = tempfile.NamedTemporaryFile(
+                            delete=False, suffix=".png", dir=extract_path
+                        ).name
+                        face_image.save(path)
+                        face_images.append(path)
+
+                    images += face_images
+
+            return images
+    except Exception as e:
+        logger.info("Failed to extract : %s", e)

-    if not extract_path:
-        tempfile.mkdtemp()
-    if files is not None:
-        images = []
-        for file in files:
-            img = Image.open(file.name).convert("RGB")
-            faces = swapper.get_faces(pil_to_cv2(img))
-            if faces:
-                face_images = []
-                for face in faces:
-                    bbox = face.bbox.astype(int)
-                    x_min, y_min, x_max, y_max = bbox
-                    face_image = img.crop((x_min, y_min, x_max, y_max))
-                    if (
-                        postprocess_options.face_restorer_name
-                        or postprocess_options.restorer_visibility
-                    ):
-                        postprocess_options.scale = (
-                            1 if face_image.width > 512 else 512 // face_image.width
-                        )
-                        face_image = enhance_image(
-                            face_image,
-                            postprocess_options,
-                        )
-                    path = tempfile.NamedTemporaryFile(
-                        delete=False, suffix=".png", dir=extract_path
-                    ).name
-                    face_image.save(path)
-                    face_images.append(path)
-                images += face_images
-        return images
    return None


-def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> str:
+def analyse_faces(image: Image.Image, 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
+        The input image where faces will be detected. The image must be a PIL Image object.
+
+    det_threshold : float, optional
+        The detection threshold for the face detection process, by default 0.5. It determines
+        the confidence level at which the function will consider a detected object as a face.
+        Value should be in the range [0, 1], with higher values indicating greater certainty.
+
+    Returns
+    -------
+    str or None
+        Returns a formatted string providing details about each face detected in the image.
+        For each face, the string will include an index and a set of facial details.
+        In the event of an exception (e.g., analysis failure), the function will log the error
+        and return None.
+
+    Raises
+    ------
+    This function handles exceptions internally and does not raise.
+
+    Examples
+    --------
+    >>> image = Image.open("test.jpg")
+    >>> print(analyse_faces(image, 0.7))
+    """
+
    try:
        faces = swapper.get_faces(imgutils.pil_to_cv2(image), det_thresh=det_threshold)
        result = ""
@@ -86,11 +157,12 @@ def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> str:
            result += f"\nFace {i} \n" + "=" * 40 + "\n"
            result += pformat(face) + "\n"
            result += "=" * 40
-        return result
+        return result if result else None

    except Exception as e:
        logger.error("Analysis Failed : %s", e)
-        return "Analysis Failed"
+
+    return None


 def sanitize_name(name: str) -> str:
@@ -116,96 +188,104 @@ def build_face_checkpoint_and_save(
    Returns:
        PIL.Image.Image or None: The resulting swapped face image if the process is successful; None otherwise.
    """
-    name = sanitize_name(name)
-    batch_files = batch_files or []
-    logger.info("Build %s %s", name, [x.name for x in batch_files])
-    faces = swapper.get_faces_from_img_files(batch_files)
-    blended_face = swapper.blend_faces(faces)
-    preview_path = os.path.join(
-        scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
-    )
-    faces_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces")

-    os.makedirs(faces_path, exist_ok=True)
+    try:
+        name = sanitize_name(name)
+        batch_files = batch_files or []
+        logger.info("Build %s %s", name, [x.name for x in batch_files])
+        faces = swapper.get_faces_from_img_files(batch_files)
+        blended_face = swapper.blend_faces(faces)
+        preview_path = REFERENCE_PATH

-    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"))
+        faces_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces")

-        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
-            ),
-        )
+        os.makedirs(faces_path, exist_ok=True)

-        file_path = os.path.join(faces_path, f"{name}.pkl")
-        file_number = 1
-        while os.path.exists(file_path):
-            file_path = os.path.join(faces_path, f"{name}_{file_number}.pkl")
-            file_number += 1
-        result_image.save(file_path + ".png")
-        with open(file_path, "wb") as file:
-            pickle.dump(
-                {
-                    "embedding": blended_face.embedding,
-                    "gender": blended_face.gender,
-                    "age": blended_face.age,
-                },
-                file,
+        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
+                ),
            )
-        try:
-            with open(file_path, "rb") as file:
-                data = Face(pickle.load(file))
-                print(data)
-        except Exception as e:
-            print(e)
-        return result_image

-    print("No face found")
+            file_path = os.path.join(faces_path, f"{name}.pkl")
+            file_number = 1
+            while os.path.exists(file_path):
+                file_path = os.path.join(faces_path, f"{name}_{file_number}.pkl")
+                file_number += 1
+            result_image.save(file_path + ".png")
+            with open(file_path, "wb") as file:
+                pickle.dump(
+                    {
+                        "embedding": blended_face.embedding,
+                        "gender": blended_face.gender,
+                        "age": blended_face.age,
+                    },
+                    file,
+                )
+            try:
+                with open(file_path, "rb") as file:
+                    data = Face(pickle.load(file))
+                    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


-def explore_onnx_faceswap_model(model_path):
-    data = {
-        "Node Name": [],
-        "Op Type": [],
-        "Inputs": [],
-        "Outputs": [],
-        "Attributes": [],
-    }
-    if model_path:
-        model = onnx.load(model_path)
-        for node in model.graph.node:
-            data["Node Name"].append(pformat(node.name))
-            data["Op Type"].append(pformat(node.op_type))
-            data["Inputs"].append(pformat(node.input))
-            data["Outputs"].append(pformat(node.output))
-            attributes = []
-            for attr in node.attribute:
-                attr_name = attr.name
-                attr_value = attr.t
-                attributes.append(
-                    "{} = {}".format(pformat(attr_name), pformat(attr_value))
-                )
-            data["Attributes"].append(attributes)
+def explore_onnx_faceswap_model(model_path: str) -> pd.DataFrame:
+    try:
+        data: Dict[str, Any] = {
+            "Node Name": [],
+            "Op Type": [],
+            "Inputs": [],
+            "Outputs": [],
+            "Attributes": [],
+        }
+        if model_path:
+            model = onnx.load(model_path)
+            for node in model.graph.node:
+                data["Node Name"].append(pformat(node.name))
+                data["Op Type"].append(pformat(node.op_type))
+                data["Inputs"].append(pformat(node.input))
+                data["Outputs"].append(pformat(node.output))
+                attributes = []
+                for attr in node.attribute:
+                    attr_name = attr.name
+                    attr_value = attr.t
+                    attributes.append(
+                        "{} = {}".format(pformat(attr_name), pformat(attr_value))
+                    )
+                data["Attributes"].append(attributes)

-    df = pd.DataFrame(data)
+        df = pd.DataFrame(data)
+    except Exception as e:
+        logger.info("Failed to explore model %s", e)
+        return None
    return df


 def batch_process(
-    files, save_path, *components: List[gr.components.Component]
+    files: List[gr.File], save_path: str, *components: List[gr.components.Component]
 ) -> Optional[List[Image.Image]]:
    try:
        if save_path is not None:
@@ -216,7 +296,7 @@ def batch_process(

        # Parse and convert units flat components into FaceSwapUnitSettings
        for i in range(0, units_count):
-            units += [FaceSwapUnitSettings.get_unit_configuration(i, components)]
+            units += [FaceSwapUnitSettings.get_unit_configuration(i, components)]  # type: ignore

        for i, u in enumerate(units):
            logger.debug("%s, %s", pformat(i), pformat(u))
@@ -3,7 +3,7 @@ import numpy as np
 import base64
 import io
 from dataclasses import dataclass, fields
-from typing import List, Union
+from typing import Any, List, Optional, Set, Union
 import dill as pickle
 import gradio as gr
 from insightface.app.common import Face
@@ -50,14 +50,16 @@ class FaceSwapUnitSettings:
    swap_in_generated: bool

    @staticmethod
-    def get_unit_configuration(unit: int, components):
+    def get_unit_configuration(
+        unit: int, components: List[gr.components.Component]
+    ) -> Any:
        fields_count = len(fields(FaceSwapUnitSettings))
        return FaceSwapUnitSettings(
            *components[unit * fields_count : unit * fields_count + fields_count]
        )

    @property
-    def faces_index(self):
+    def faces_index(self) -> Set[int]:
        """
        Convert _faces_index from str to int
        """
@@ -72,18 +74,18 @@ class FaceSwapUnitSettings:
        return faces_index

    @property
-    def compute_similarity(self):
+    def compute_similarity(self) -> bool:
        return self._compute_similarity or self.check_similarity

    @property
-    def batch_files(self):
+    def batch_files(self) -> List[gr.File]:
        """
        Return empty array instead of None for batch files
        """
        return self._batch_files or []

    @property
-    def reference_face(self):
+    def reference_face(self) -> Optional[Face]:
        """
        Extract reference face (only once and store it for the rest of processing).
        Reference face is the checkpoint or the source image or the first image in the batch in that order.
@@ -97,6 +99,7 @@ class FaceSwapUnitSettings:
                        self._reference_face = face
                    except Exception as e:
                        logger.error("Failed to load checkpoint  : %s", e)
+                        raise e
            elif self.source_img is not None:
                if isinstance(self.source_img, str):  # source_img is a base64 string
                    if (
@@ -119,11 +122,12 @@ class FaceSwapUnitSettings:

        if self._reference_face is None:
            logger.error("You need at least one reference face")
+            raise Exception("No reference face found")

        return self._reference_face

    @property
-    def faces(self):
+    def faces(self) -> List[Face]:
        """_summary_
        Extract all faces (including reference face) to provide an array of faces
        Only processed once.
@@ -146,7 +150,7 @@ class FaceSwapUnitSettings:
        return self._faces

    @property
-    def blended_faces(self):
+    def blended_faces(self) -> Face:
        """
        Blend the faces using the mean of all embeddings
        """