diff --git a/scripts/faceswaplab.py b/scripts/faceswaplab.py index a116fb5..4863d20 100644 --- a/scripts/faceswaplab.py +++ b/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]: diff --git a/scripts/faceswaplab_api/faceswaplab_api.py b/scripts/faceswaplab_api/faceswaplab_api.py index ef6803f..cb000f8 100644 --- a/scripts/faceswaplab_api/faceswaplab_api.py +++ b/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, ) diff --git a/scripts/faceswaplab_globals.py b/scripts/faceswaplab_globals.py index 7d50774..ca1f76b 100644 --- a/scripts/faceswaplab_globals.py +++ b/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" diff --git a/scripts/faceswaplab_inpainting/faceswaplab_inpainting.py b/scripts/faceswaplab_inpainting/faceswaplab_inpainting.py index 716d1b3..d4409a1 100644 --- a/scripts/faceswaplab_inpainting/faceswaplab_inpainting.py +++ b/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). diff --git a/scripts/faceswaplab_postprocessing/postprocessing_options.py b/scripts/faceswaplab_postprocessing/postprocessing_options.py index b7f7bcd..70f65fd 100644 --- a/scripts/faceswaplab_postprocessing/postprocessing_options.py +++ b/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 diff --git a/scripts/faceswaplab_postprocessing/upscaling.py b/scripts/faceswaplab_postprocessing/upscaling.py index 746f229..bc0c257 100644 --- a/scripts/faceswaplab_postprocessing/upscaling.py +++ b/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) diff --git a/scripts/faceswaplab_settings/faceswaplab_settings.py b/scripts/faceswaplab_settings/faceswaplab_settings.py index 5183869..c6c24ff 100644 --- a/scripts/faceswaplab_settings/faceswaplab_settings.py +++ b/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", diff --git a/scripts/faceswaplab_utils/face_checkpoints_utils.py b/scripts/faceswaplab_swapping/face_checkpoints.py similarity index 98% rename from scripts/faceswaplab_utils/face_checkpoints_utils.py rename to scripts/faceswaplab_swapping/face_checkpoints.py index 280359e..89ea6ff 100644 --- a/scripts/faceswaplab_utils/face_checkpoints_utils.py +++ b/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, diff --git a/scripts/faceswaplab_swapping/facemask.py b/scripts/faceswaplab_swapping/facemask.py index d49b713..e16f63e 100644 --- a/scripts/faceswaplab_swapping/facemask.py +++ b/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) diff --git a/scripts/faceswaplab_swapping/swapper.py b/scripts/faceswaplab_swapping/swapper.py index 146cd6b..ceb8e56 100644 --- a/scripts/faceswaplab_swapping/swapper.py +++ b/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, diff --git a/scripts/faceswaplab_swapping/upcaled_inswapper_options.py b/scripts/faceswaplab_swapping/upcaled_inswapper_options.py index efba4e4..759235e 100644 --- a/scripts/faceswaplab_swapping/upcaled_inswapper_options.py +++ b/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). diff --git a/scripts/faceswaplab_swapping/upscaled_inswapper.py b/scripts/faceswaplab_swapping/upscaled_inswapper.py index 691f929..30f4907 100644 --- a/scripts/faceswaplab_swapping/upscaled_inswapper.py +++ b/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)) diff --git a/scripts/faceswaplab_ui/faceswaplab_inpainting_ui.py b/scripts/faceswaplab_ui/faceswaplab_inpainting_ui.py index bb4dd93..052fa17 100644 --- a/scripts/faceswaplab_ui/faceswaplab_inpainting_ui.py +++ b/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", diff --git a/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py b/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py index ec12a82..8209854 100644 --- a/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py +++ b/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py @@ -2,8 +2,8 @@ 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]], diff --git a/scripts/faceswaplab_ui/faceswaplab_tab.py b/scripts/faceswaplab_ui/faceswaplab_tab.py index 6415a36..1e1e593 100644 --- a/scripts/faceswaplab_ui/faceswaplab_tab.py +++ b/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_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 - images = [Image.open(file.name) for file in batch_files] - preview_image = face_checkpoints_utils.build_face_checkpoint_and_save( + return None # type: ignore (Optional not really supported by old gradio) + images = [Image.open(file.name) for file in batch_files] # type: ignore + 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( - images_paths, - save_path=save_path, - units=units, - postprocess_options=postprocess_options, + return ( + swapper.batch_process( + images_paths, + save_path=save_path, + units=units, + postprocess_options=postprocess_options, + ) + or [] ) except Exception as e: logger.error("Batch Process error : %s", e) @@ -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() diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py index 84db6d3..b108d9a 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py +++ b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py @@ -9,7 +9,7 @@ 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 @@ def reference_face(self) -> Optional[Face]: 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 @@ def faces(self) -> List[Face]: 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 diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py b/scripts/faceswaplab_ui/faceswaplab_unit_ui.py index a76997e..127cec0 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py +++ b/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", diff --git a/scripts/faceswaplab_utils/imgutils.py b/scripts/faceswaplab_utils/imgutils.py index 3b55b6c..2306253 100644 --- a/scripts/faceswaplab_utils/imgutils.py +++ b/scripts/faceswaplab_utils/imgutils.py @@ -6,10 +6,10 @@ 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) diff --git a/scripts/faceswaplab_utils/models_utils.py b/scripts/faceswaplab_utils/models_utils.py index e235dec..2d8b8e9 100644 --- a/scripts/faceswaplab_utils/models_utils.py +++ b/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 diff --git a/scripts/faceswaplab_utils/sd_utils.py b/scripts/faceswaplab_utils/sd_utils.py new file mode 100644 index 0000000..fd797ff --- /dev/null +++ b/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)