fix preview in build

scripts/faceswaplab_api/faceswaplab_api.py

@@ -21,6 +21,7 @@ from client_api import api_utils
 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
@@ -99,7 +100,7 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None:
                 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
             )
 

scripts/faceswaplab_swapping/face_checkpoints.py

@@ -38,8 +38,11 @@ def sanitize_name(name: str) -> str:
 
 
 def build_face_checkpoint_and_save(
-    images: List[PILImage], name: str, overwrite: bool = False, path: str = None
+    images: List[PILImage],
-) -> PILImage:
+    name: str,
+    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
@@ -57,8 +60,12 @@ def build_face_checkpoint_and_save(
         name = sanitize_name(name)
         images = images or []
         logger.info("Build %s with %s images", name, len(images))
-        faces = swapper.get_faces_from_img_files(images)
+        faces: List[Face] = swapper.get_faces_from_img_files(images=images)
-        blended_face = swapper.blend_faces(faces)
+        if faces is None or len(faces) == 0:
+            logger.error("No source faces found")
+            return None
+
+        blended_face: Optional[Face] = swapper.blend_faces(faces)
         preview_path = os.path.join(
             scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
         )
@@ -85,40 +92,51 @@ def build_face_checkpoint_and_save(
                     "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.swap_face(
+                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"),
+                    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:
+                if path:
-                file_path = path
+                    file_path = path
-            else:
+                else:
-                file_path = os.path.join(get_checkpoint_path(), f"{name}.safetensors")
+                    file_path = os.path.join(
-                if not overwrite:
+                        get_checkpoint_path(), f"{name}.safetensors"
-                    file_number = 1
+                    )
-                    while os.path.exists(file_path):
+                    if not overwrite:
-                        file_path = os.path.join(
+                        file_number = 1
-                            get_checkpoint_path(), f"{name}_{file_number}.safetensors"
+                        while os.path.exists(file_path):
-                        )
+                            file_path = os.path.join(
-                        file_number += 1
+                                get_checkpoint_path(),
-            save_face(filename=file_path, face=blended_face)
+                                f"{name}_{file_number}.safetensors",
-            preview_image.save(file_path + ".png")
+                            )
-            try:
+                            file_number += 1
-                data = load_face(file_path)
+                save_face(filename=file_path, face=blended_face)
-                logger.debug(data)
+                preview_image.save(file_path + ".png")
-            except Exception as e:
+                try:
-                logger.error("Error loading checkpoint, after creation %s", e)
+                    data = load_face(file_path)
-                traceback.print_exc()
+                    logger.debug(data)
-
+                except Exception as e:
-            return preview_image
+                    logger.error("Error loading checkpoint, after creation %s", e)
+                    traceback.print_exc()
+
+                return preview_image
 
         else:
             logger.error("No face found")
-            return None
+            return None  # type: ignore
     except Exception as e:
         logger.error("Failed to build checkpoint %s", e)
         traceback.print_exc()
@@ -139,7 +157,7 @@ def save_face(face: Face, filename: str) -> None:
         raise e
 
 
-def load_face(name: str) -> Face:
+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)

scripts/faceswaplab_swapping/swapper.py

@@ -141,7 +141,7 @@ def batch_process(
     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[PILImage]]:
     """
     Process a batch of images, apply face swapping according to the given settings, and optionally save the resulting images to a specified path.
@@ -527,7 +527,7 @@ 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(images: List[PILImage]) -> List[Optional[CV2ImgU8]]:
+def get_faces_from_img_files(images: List[PILImage]) -> List[Face]:
     """
     Extracts faces from a list of image files.
 
@@ -539,7 +539,7 @@ def get_faces_from_img_files(images: List[PILImage]) -> List[Optional[CV2ImgU8]]
 
     """
 
-    faces = []
+    faces: List[Face] = []
 
     if len(images) > 0:
         for img in images:
@@ -598,7 +598,6 @@ def swap_face(
     target_faces: List[Face],
     model: str,
     swapping_options: Optional[InswappperOptions],
-    compute_similarity: bool = True,
 ) -> ImageResult:
     """
     Swaps faces in the target image with the source face.
@@ -680,9 +679,9 @@ def process_image_unit(
     model: str,
     unit: FaceSwapUnitSettings,
     image: PILImage,
-    info: str = None,
+    info: Optional[str] = None,
     force_blend: bool = False,
-) -> List[Tuple[PILImage, str]]:
+) -> List[Tuple[PILImage, Optional[str]]]:
     """Process one image and return a List of (image, info) (one if blended, many if not).
 
     Args:
@@ -723,7 +722,9 @@ def process_image_unit(
                 sort_by_face_size=unit.sort_by_size,
             )
 
-            target_faces = filter_faces(faces, filtering_options=face_filtering_options)
+            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:
@@ -738,7 +739,6 @@ def process_image_unit(
                 target_faces=target_faces,
                 model=model,
                 swapping_options=unit.swapping_options,
-                compute_similarity=unit.compute_similarity,
             )
             # Apply post-inpainting to image
             if unit.post_inpainting.inpainting_denoising_strengh > 0:

scripts/faceswaplab_ui/faceswaplab_tab.py

@@ -156,16 +156,18 @@ def build_face_checkpoint_and_save(
         if not batch_files:
             logger.error("No face found")
             return None  # type: ignore (Optional not really supported by old gradio)
-        images = [Image.open(file.name) for file in batch_files]  # type: ignore
+        images: list[PILImage] = [Image.open(file.name) for file in batch_files]  # type: ignore
-        preview_image = face_checkpoints.build_face_checkpoint_and_save(
+        preview_image: PILImage | None = (
-            images, name, overwrite=overwrite
+            face_checkpoints.build_face_checkpoint_and_save(
+                images=images, name=name, overwrite=overwrite
+            )
         )
     except Exception as e:
         logger.error("Failed to build checkpoint %s", e)
 
         traceback.print_exc()
         return None  # type: ignore
-    return preview_image
+    return preview_image  # type: ignore
 
 
 def explore_onnx_faceswap_model(model_path: str) -> pd.DataFrame: