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 @@ class FaceSwapScript(scripts.Script):
 
     @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 @@ 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]:

scripts/faceswaplab_api/faceswaplab_api.py

@@ -18,7 +18,7 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
     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 @@ VERSION_FLAG: str = "v1.2.2"
 # 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_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.imgutils import (
 )
 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_utils.typing import CV2ImgU8, 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
 
 
 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_postprocessing.postprocessing_options import (
 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 @@ class UpscaledINSwapper(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 @@ class UpscaledINSwapper(INSwapper):
         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 @@ class UpscaledINSwapper(INSwapper):
             (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 @@ class UpscaledINSwapper(INSwapper):
                 mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
                 mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
                 mask_size = int(np.sqrt(mask_h * mask_w))
-                erosion_factor = options.erosion_factor
+                erosion_factor = options.erosion_factor if options else 1
 
                 k = max(int(mask_size // 10 * erosion_factor), int(10 * erosion_factor))
 

scripts/faceswaplab_ui/faceswaplab_inpainting_ui.py

@@ -1,7 +1,7 @@
 from typing import List
 import gradio as gr
-from modules.shared import opts
 from modules import sd_models, sd_samplers
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
 
 
 def face_inpainting_ui(
@@ -19,14 +19,14 @@ def face_inpainting_ui(
         )
 
         inpainting_denoising_prompt = gr.Textbox(
-            opts.data.get(
+            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(
-            opts.data.get(
+            get_sd_option(
                 "faceswaplab_pp_default_inpainting_negative_prompt", "blurry"
             ),
             elem_id=f"{id_prefix}_pp_inpainting_denoising_neg_prompt",

scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py

@@ -2,8 +2,8 @@ from typing import List
 import gradio as gr
 import modules
 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
 
 
 def postprocessing_ui() -> List[gr.components.Component]:
@@ -15,7 +15,7 @@ def postprocessing_ui() -> List[gr.components.Component]:
             face_restorer_name = gr.Radio(
                 label="Restore Face",
                 choices=["None"] + [x.name() for x in shared.face_restorers],
-                value=lambda: opts.data.get(
+                value=lambda: get_sd_option(
                     "faceswaplab_pp_default_face_restorer",
                     shared.face_restorers[0].name(),
                 ),
@@ -26,7 +26,7 @@ def postprocessing_ui() -> List[gr.components.Component]:
                 face_restorer_visibility = gr.Slider(
                     0,
                     1,
-                    value=lambda: opts.data.get(
+                    value=lambda: get_sd_option(
                         "faceswaplab_pp_default_face_restorer_visibility", 1
                     ),
                     step=0.001,
@@ -36,7 +36,7 @@ def postprocessing_ui() -> List[gr.components.Component]:
                 codeformer_weight = gr.Slider(
                     0,
                     1,
-                    value=lambda: opts.data.get(
+                    value=lambda: get_sd_option(
                         "faceswaplab_pp_default_face_restorer_weight", 1
                     ),
                     step=0.001,
@@ -45,7 +45,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=lambda: get_sd_option("faceswaplab_pp_default_upscaler", "None"),
             label="Upscaler",
             elem_id="faceswaplab_pp_upscaler",
         )
@@ -60,7 +60,7 @@ def postprocessing_ui() -> List[gr.components.Component]:
         upscaler_visibility = gr.Slider(
             0,
             1,
-            value=lambda: opts.data.get(
+            value=lambda: get_sd_option(
                 "faceswaplab_pp_default_upscaler_visibility", 1
             ),
             step=0.1,
@@ -87,21 +87,21 @@ def postprocessing_ui() -> List[gr.components.Component]:
             )
 
             inpainting_denoising_prompt = gr.Textbox(
-                opts.data.get(
+                get_sd_option(
                     "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(
+                get_sd_option(
                     "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]
+                samplers_names = [s.name for s in modules.sd_samplers.all_samplers]  # type: ignore
                 inpainting_sampler = gr.Dropdown(
                     choices=samplers_names,
                     value=[samplers_names[0]],

scripts/faceswaplab_ui/faceswaplab_tab.py

@@ -1,11 +1,12 @@
 import traceback
 from pprint import pformat
 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 onnx
 import pandas as pd
-from modules.shared import opts
 from PIL import Image
 
 import scripts.faceswaplab_swapping.swapper as swapper
@@ -15,7 +16,7 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
 from scripts.faceswaplab_ui.faceswaplab_postprocessing_ui import postprocessing_ui
 from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
 from scripts.faceswaplab_ui.faceswaplab_unit_ui import faceswap_unit_ui
-from scripts.faceswaplab_utils import face_checkpoints_utils, imgutils
+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
@@ -74,7 +75,7 @@ def extract_faces(
             [PostProcessingOptions], components
         ).pop()
         images = [
-            Image.open(file.name) for file in files
+            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
@@ -136,7 +137,7 @@ def analyse_faces(image: PILImage, det_threshold: float = 0.5) -> Optional[str]:
 
 
 def build_face_checkpoint_and_save(
-    batch_files: gr.File, name: str, overwrite: bool
+    batch_files: List[gr.File], name: str, overwrite: bool
 ) -> PILImage:
     """
     Builds a face checkpoint using the provided image files, performs face swapping,
@@ -154,16 +155,16 @@ def build_face_checkpoint_and_save(
     try:
         if not batch_files:
             logger.error("No face found")
-            return None
+            return None  # type: ignore (Optional not really supported by old gradio)
-        images = [Image.open(file.name) for file in batch_files]
+        images = [Image.open(file.name) for file in batch_files]  # type: ignore
-        preview_image = face_checkpoints_utils.build_face_checkpoint_and_save(
+        preview_image = face_checkpoints.build_face_checkpoint_and_save(
             images, name, overwrite=overwrite
         )
     except Exception as e:
         logger.error("Failed to build checkpoint %s", e)
 
         traceback.print_exc()
-        return None
+        return None  # type: ignore
     return preview_image
 
 
@@ -197,7 +198,7 @@ def explore_onnx_faceswap_model(model_path: str) -> pd.DataFrame:
         logger.error("Failed to explore model %s", e)
 
         traceback.print_exc()
-        return None
+        return None  # type: ignore
     return df
 
 
@@ -205,7 +206,7 @@ def batch_process(
     files: List[gr.File], save_path: str, *components: Tuple[Any, ...]
 ) -> List[PILImage]:
     try:
-        units_count = opts.data.get("faceswaplab_units_count", 3)
+        units_count = get_sd_option("faceswaplab_units_count", 3)
 
         classes: List[Any] = dataclasses_from_flat_list(
             [FaceSwapUnitSettings] * units_count + [PostProcessingOptions],
@@ -216,13 +217,16 @@ def batch_process(
         ]
         postprocess_options = classes[-1]
 
-        images_paths = [file.name for file in files]
+        images_paths = [file.name for file in files]  # type: ignore
 
-        return swapper.batch_process(
+        return (
-            images_paths,
+            swapper.batch_process(
-            save_path=save_path,
+                images_paths,
-            units=units,
+                save_path=save_path,
-            postprocess_options=postprocess_options,
+                units=units,
+                postprocess_options=postprocess_options,
+            )
+            or []
         )
     except Exception as e:
         logger.error("Batch Process error : %s", e)
@@ -304,7 +308,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="",
@@ -360,7 +364,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/",
@@ -370,7 +374,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()

scripts/faceswaplab_ui/faceswaplab_unit_settings.py

@@ -9,7 +9,7 @@ from PIL import Image
 from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions
 from scripts.faceswaplab_utils.imgutils import pil_to_cv2
 from scripts.faceswaplab_utils.faceswaplab_logging import logger
-from scripts.faceswaplab_utils import face_checkpoints_utils
+from scripts.faceswaplab_swapping import face_checkpoints
 from scripts.faceswaplab_inpainting.faceswaplab_inpainting import InpaintingOptions
 from client_api import api_utils
 
@@ -124,7 +124,7 @@ class FaceSwapUnitSettings:
             if self.source_face and self.source_face != "None":
                 try:
                     logger.info(f"loading face {self.source_face}")
-                    face = face_checkpoints_utils.load_face(self.source_face)
+                    face = face_checkpoints.load_face(self.source_face)
                     self._reference_face = face
                 except Exception as e:
                     logger.error("Failed to load checkpoint  : %s", e)
@@ -169,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

scripts/faceswaplab_ui/faceswaplab_unit_ui.py

@@ -1,9 +1,9 @@
 from typing import List
 from scripts.faceswaplab_ui.faceswaplab_inpainting_ui import face_inpainting_ui
-from scripts.faceswaplab_utils.face_checkpoints_utils import get_face_checkpoints
+from scripts.faceswaplab_swapping.face_checkpoints import get_face_checkpoints
 import gradio as gr
-from modules.shared import opts
 from modules import shared
+from scripts.faceswaplab_utils.sd_utils import get_sd_option
 
 
 def faceswap_unit_advanced_options(
@@ -17,7 +17,7 @@ def faceswap_unit_advanced_options(
             face_restorer_name = gr.Radio(
                 label="Restore Face",
                 choices=["None"] + [x.name() for x in shared.face_restorers],
-                value=lambda: opts.data.get(
+                value=lambda: get_sd_option(
                     "faceswaplab_default_upscaled_swapper_face_restorer",
                     "None",
                 ),
@@ -28,7 +28,7 @@ def faceswap_unit_advanced_options(
                 face_restorer_visibility = gr.Slider(
                     0,
                     1,
-                    value=lambda: opts.data.get(
+                    value=lambda: get_sd_option(
                         "faceswaplab_default_upscaled_swapper_face_restorer_visibility",
                         1.0,
                     ),
@@ -39,7 +39,7 @@ def faceswap_unit_advanced_options(
                 codeformer_weight = gr.Slider(
                     0,
                     1,
-                    value=lambda: opts.data.get(
+                    value=lambda: get_sd_option(
                         "faceswaplab_default_upscaled_swapper_face_restorer_weight", 1.0
                     ),
                     step=0.001,
@@ -48,7 +48,7 @@ def faceswap_unit_advanced_options(
                 )
         upscaler_name = gr.Dropdown(
             choices=[upscaler.name for upscaler in shared.sd_upscalers],
-            value=lambda: opts.data.get(
+            value=lambda: get_sd_option(
                 "faceswaplab_default_upscaled_swapper_upscaler", ""
             ),
             label="Upscaler",
@@ -56,7 +56,7 @@ def faceswap_unit_advanced_options(
         )
 
         improved_mask = gr.Checkbox(
-            lambda: opts.data.get(
+            lambda: get_sd_option(
                 "faceswaplab_default_upscaled_swapper_improved_mask", False
             ),
             interactive=True,
@@ -64,7 +64,7 @@ def faceswap_unit_advanced_options(
             elem_id=f"{id_prefix}_face{unit_num}_improved_mask",
         )
         color_corrections = gr.Checkbox(
-            lambda: opts.data.get(
+            lambda: get_sd_option(
                 "faceswaplab_default_upscaled_swapper_fixcolor", False
             ),
             interactive=True,
@@ -72,7 +72,7 @@ def faceswap_unit_advanced_options(
             elem_id=f"{id_prefix}_face{unit_num}_color_corrections",
         )
         sharpen_face = gr.Checkbox(
-            lambda: opts.data.get(
+            lambda: get_sd_option(
                 "faceswaplab_default_upscaled_swapper_sharpen", False
             ),
             interactive=True,
@@ -82,7 +82,7 @@ def faceswap_unit_advanced_options(
         erosion_factor = gr.Slider(
             0.0,
             10.0,
-            lambda: opts.data.get("faceswaplab_default_upscaled_swapper_erosion", 1.0),
+            lambda: 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",

scripts/faceswaplab_utils/imgutils.py

@@ -6,10 +6,10 @@ 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
 
@@ -25,6 +25,12 @@ def check_against_nsfw(img: PILImage) -> bool:
     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
+
     shapes: List[bool] = []
     chunks: List[Dict[str, Union[int, float]]] = detect(img)
 

scripts/faceswaplab_utils/models_utils.py

@@ -77,4 +77,5 @@ def get_current_swap_model() -> str:
         raise FileNotFoundError(
             "No faceswap model found. Please add it to the faceswaplab directory."
         )
+    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)