wip, add nsfw option due to perf, still some mypy warnings

scripts/faceswaplab.py

@@ -1,4 +1,5 @@
 from scripts.configure import check_configuration
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
 
 check_configuration()
 
@@ -76,7 +77,7 @@ def __init__(self) -> None:
 
     @property
     def units_count(self) -> int:
-        return opts.data.get("faceswaplab_units_count", 3)
+        return get_sd_option("faceswaplab_units_count", 3)
 
     @property
     def enabled(self) -> bool:
@@ -85,7 +86,7 @@ def enabled(self) -> bool:
 
     @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]:

scripts/faceswaplab_api/faceswaplab_api.py

@@ -18,7 +18,7 @@
     PostProcessingOptions,
 )
 from client_api import api_utils
-from scripts.faceswaplab_utils.face_checkpoints_utils import (
+from scripts.faceswaplab_swapping.face_checkpoints import (
     build_face_checkpoint_and_save,
 )
 

scripts/faceswaplab_globals.py

@@ -1,6 +1,5 @@
 import os
 from modules import scripts
-from modules.shared import opts
 
 # Defining the absolute path for the 'faceswaplab' directory inside 'models' directory
 MODELS_DIR = os.path.abspath(os.path.join("models", "faceswaplab"))
@@ -21,7 +20,5 @@
 # Defining the path for 'sd-webui-faceswaplab' inside the 'extensions' directory
 EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab")
 
-# Defining the NSFW score threshold. Any image part with a score above this value will be treated as NSFW (Not Safe For Work)
-NSFW_SCORE_THRESHOLD: float = opts.data.get("faceswaplab_nsfw_threshold", 0.7)  # type: ignore
 # Defining the expected SHA1 hash value for 'INSWAPPER'
 EXPECTED_INSWAPPER_SHA1 = "17a64851eaefd55ea597ee41e5c18409754244c5"

scripts/faceswaplab_inpainting/faceswaplab_inpainting.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import List
+from typing import List, Optional
 import gradio as gr
 from client_api import api_utils
 
@@ -15,10 +15,10 @@ class InpaintingOptions:
 
     @staticmethod
     def from_gradio(components: List[gr.components.Component]) -> "InpaintingOptions":
-        return InpaintingOptions(*components)
+        return InpaintingOptions(*components)  # type: ignore
 
     @staticmethod
-    def from_api_dto(dto: api_utils.InpaintingOptions) -> "InpaintingOptions":
+    def from_api_dto(dto: Optional[api_utils.InpaintingOptions]) -> "InpaintingOptions":
         """
         Converts a InpaintingOptions object from an API DTO (Data Transfer Object).
 

scripts/faceswaplab_postprocessing/postprocessing_options.py

@@ -1,3 +1,4 @@
+from typing import Optional
 from modules.face_restoration import FaceRestoration
 from modules.upscaler import UpscalerData
 from dataclasses import dataclass
@@ -27,17 +28,17 @@ class PostProcessingOptions:
     inpainting_when: InpaintingWhen = InpaintingWhen.BEFORE_UPSCALING
 
     # (Don't use optional for this or gradio parsing will fail) :
-    inpainting_options: InpaintingOptions = None
+    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

scripts/faceswaplab_postprocessing/upscaling.py

@@ -17,7 +17,7 @@ def upscale_img(image: PILImage, pp_options: PostProcessingOptions) -> PILImage:
             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
         )
 
         # FIXME : Could be better (managing images whose dimensions are not multiples of 16)

scripts/faceswaplab_settings/faceswaplab_settings.py

@@ -46,6 +46,16 @@ 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 !)",
+            gr.Slider,
+            {"minimum": 0, "maximum": 1, "step": 0.01},
+            section=section,
+        ),
+    )
 
     shared.opts.add_option(
         "faceswaplab_det_size",

scripts/faceswaplab_utils/face_checkpoints_utils.py → scripts/faceswaplab_swapping/face_checkpoints.py

@@ -63,7 +63,7 @@ def build_face_checkpoint_and_save(
             scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
         )
 
-        reference_preview_img: PILImage = None
+        reference_preview_img: PILImage
         if blended_face:
             if blended_face["gender"] == 0:
                 reference_preview_img = Image.open(
@@ -86,7 +86,6 @@ def build_face_checkpoint_and_save(
                 )
             else:
                 result = swapper.swap_face(
-                    reference_face=blended_face,
                     target_faces=[target_face],
                     source_face=blended_face,
                     target_img=reference_preview_img,

scripts/faceswaplab_swapping/facemask.py

@@ -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/swapper.py

@@ -26,7 +26,6 @@
 )
 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
@@ -38,12 +37,13 @@
 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
 
 
 def use_gpu() -> bool:
     return (
         getattr(shared.cmd_opts, "faceswaplab_gpu", False)
-        or opts.data.get("faceswaplab_use_gpu", False)
+        or get_sd_option("faceswaplab_use_gpu", False)
     ) and sys.platform != "darwin"
 
 
@@ -166,6 +166,7 @@ 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(
@@ -182,7 +183,7 @@ def batch_process(
                 swapped_images = process_images_units(
                     get_current_swap_model(), images=[(src_image, None)], units=units
                 )
-                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))
@@ -209,7 +210,7 @@ def extract_faces(
     images: List[PILImage],
     extract_path: Optional[str],
     postprocess_options: PostProcessingOptions,
-) -> Optional[List[str]]:
+) -> Optional[List[PILImage]]:
     """
     Extracts faces from a list of image files.
 
@@ -232,14 +233,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))
 
@@ -370,7 +371,7 @@ def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper:
         with tqdm(total=1, desc="Loading swap model", unit="model") as pbar:
             with capture_stdout() as captured:
                 model = upscaled_inswapper.UpscaledINSwapper(
-                    insightface.model_zoo.get_model(model_path, providers=providers)
+                    insightface.model_zoo.get_model(model_path, providers=providers)  # type: ignore
                 )
             pbar.update(1)
         logger.info("%s", pformat(captured.getvalue()))
@@ -402,11 +403,11 @@ 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)
 
-    auto_det_size = opts.data.get("faceswaplab_auto_det_size", True)
+    auto_det_size = get_sd_option("faceswaplab_auto_det_size", True)
     if not auto_det_size:
-        x = opts.data.get("faceswaplab_det_size", 640)
+        x = get_sd_option("faceswaplab_det_size", 640)
         det_size = (x, x)
 
     face_analyser = getAnalysisModel(det_size, det_thresh)
@@ -433,7 +434,7 @@ def get_faces(
 
     try:
         # Sort the detected faces based on their x-coordinate of the bounding box
-        return sorted(faces, 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()
@@ -470,7 +471,7 @@ def filter_faces(
         filtered_faces = sorted(
             all_faces,
             reverse=True,
-            key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]),
+            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:
@@ -566,7 +567,7 @@ def blend_faces(faces: List[Face]) -> Optional[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
@@ -592,7 +593,6 @@ def blend_faces(faces: List[Face]) -> Optional[Face]:
 
 
 def swap_face(
-    reference_face: CV2ImgU8,
     source_face: Face,
     target_img: PILImage,
     target_faces: List[Face],
@@ -604,7 +604,6 @@ def swap_face(
     Swaps faces in the target image with the source face.
 
     Args:
-        reference_face (CV2ImgU8): The reference face used for similarity comparison.
         source_face (CV2ImgU8): The source face to be swapped.
         target_img (PILImage): The target image to swap faces in.
         model (str): Path to the face swap model.
@@ -614,7 +613,9 @@ def swap_face(
 
     """
     return_result = ImageResult(target_img, {}, {})
-    target_img_cv2: CV2ImgU8 = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
+    target_img_cv2: CV2ImgU8 = cv2.cvtColor(
+        np.array(target_img), cv2.COLOR_RGB2BGR
+    ).astype("uint8")
     try:
         gender = source_face["gender"]
         logger.info("Source Gender %s", gender)
@@ -732,7 +733,6 @@ def process_image_unit(
 
             save_img_debug(image, "Before swap")
             result: ImageResult = swap_face(
-                reference_face=reference_face,
                 source_face=src_face,
                 target_img=current_image,
                 target_faces=target_faces,

scripts/faceswaplab_swapping/upcaled_inswapper_options.py

@@ -1,20 +1,21 @@
 from dataclasses import *
+from typing import Optional
 from client_api import api_utils
 
 
 @dataclass
 class InswappperOptions:
-    face_restorer_name: str = None
+    face_restorer_name: Optional[str] = None
     restorer_visibility: float = 1
     codeformer_weight: float = 1
-    upscaler_name: str = None
+    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: api_utils.InswappperOptions) -> "InswappperOptions":
+    def from_api_dto(dto: Optional[api_utils.InswappperOptions]) -> "InswappperOptions":
         """
         Converts a InpaintingOptions object from an API DTO (Data Transfer Object).
 

scripts/faceswaplab_swapping/upscaled_inswapper.py

@@ -1,10 +1,9 @@
-from typing import Any, Tuple, Union
+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
@@ -14,13 +13,14 @@
 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
@@ -130,8 +130,14 @@ def __init__(self, inswapper: INSwapper):
         self.__dict__.update(inswapper.__dict__)
 
     def upscale_and_restore(
-        self, img: CV2ImgU8, k: int = 2, inswapper_options: InswappperOptions = None
+        self,
+        img: CV2ImgU8,
+        k: int = 2,
+        inswapper_options: Optional[InswappperOptions] = None,
     ) -> CV2ImgU8:
+        if inswapper_options is None:
+            return img
+
         pil_img = cv2_to_pil(img)
         pp_options = PostProcessingOptions(
             upscaler_name=inswapper_options.upscaler_name,
@@ -156,7 +162,7 @@ def get(
         target_face: Face,
         source_face: Face,
         paste_back: bool = True,
-        options: InswappperOptions = None,
+        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(
@@ -166,9 +172,10 @@ def get(
             (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]
@@ -274,7 +281,7 @@ def compute_diff(bgr_fake: CV2ImgU8, aimg: CV2ImgU8) -> CV2ImgU8:
                 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 = options.erosion_factor
+                erosion_factor = options.erosion_factor if options else 1
 
                 k = max(int(mask_size // 10 * erosion_factor), int(10 * erosion_factor))