From ee7f7d09d2cb9220aba3dbeda6d76e35b5af7ddc Mon Sep 17 00:00:00 2001 From: Tran Xen <137925069+glucauze@users.noreply.github.com> Date: Tue, 1 Aug 2023 20:17:43 +0200 Subject: [PATCH] huge changes, inpainting in faces unit, change faces processing, change api, refactor, requires further testing --- client_api/api_utils.py | 73 +++--- client_api/faceswaplab_api_example.py | 15 +- requirements.txt | 3 +- scripts/faceswaplab.py | 86 ++++--- scripts/faceswaplab_api/faceswaplab_api.py | 53 +--- scripts/faceswaplab_globals.py | 2 +- .../faceswaplab_inpainting.py | 41 ++++ .../i2i_pp.py | 64 +++-- .../postprocessing.py | 19 +- .../postprocessing_options.py | 36 ++- .../faceswaplab_postprocessing/upscaling.py | 12 +- scripts/faceswaplab_swapping/facemask.py | 2 +- .../faceswaplab_swapping/parsing/__init__.py | 2 +- scripts/faceswaplab_swapping/swapper.py | 228 ++++++++++-------- .../upscaled_inswapper.py | 97 ++++++-- .../faceswaplab_inpainting_ui.py | 68 ++++++ .../faceswaplab_postprocessing_ui.py | 6 +- scripts/faceswaplab_ui/faceswaplab_tab.py | 135 ++++++----- .../faceswaplab_unit_settings.py | 67 ++--- scripts/faceswaplab_ui/faceswaplab_unit_ui.py | 52 ++-- scripts/faceswaplab_utils/imgutils.py | 58 ++--- scripts/faceswaplab_utils/typing.py | 10 + scripts/faceswaplab_utils/ui_utils.py | 39 +++ tests/test_api.py | 36 ++- 24 files changed, 786 insertions(+), 418 deletions(-) create mode 100644 scripts/faceswaplab_inpainting/faceswaplab_inpainting.py rename scripts/{faceswaplab_postprocessing => faceswaplab_inpainting}/i2i_pp.py (50%) create mode 100644 scripts/faceswaplab_ui/faceswaplab_inpainting_ui.py create mode 100644 scripts/faceswaplab_utils/typing.py create mode 100644 scripts/faceswaplab_utils/ui_utils.py diff --git a/client_api/api_utils.py b/client_api/api_utils.py index 4369ef8..8499306 100644 --- a/client_api/api_utils.py +++ b/client_api/api_utils.py @@ -18,6 +18,37 @@ class InpaintingWhen(Enum): AFTER_ALL = "After All" +class InpaintingOptions(BaseModel): + inpainting_denoising_strengh: float = Field( + description="Inpainting denoising strenght", default=0, lt=1, ge=0 + ) + inpainting_prompt: str = Field( + description="Inpainting denoising strenght", + examples=["Portrait of a [gender]"], + default="Portrait of a [gender]", + ) + inpainting_negative_prompt: str = Field( + description="Inpainting denoising strenght", + examples=[ + "Deformed, blurry, bad anatomy, disfigured, poorly drawn face, mutation" + ], + default="", + ) + inpainting_steps: int = Field( + description="Inpainting steps", + examples=["Portrait of a [gender]"], + ge=1, + le=150, + default=20, + ) + inpainting_sampler: str = Field( + description="Inpainting sampler", examples=["Euler"], default="Euler" + ) + inpainting_model: str = Field( + description="Inpainting model", examples=["Current"], default="Current" + ) + + class FaceSwapUnit(BaseModel): # The image given in reference source_img: str = Field( @@ -82,6 +113,16 @@ class FaceSwapUnit(BaseModel): default=0, ) + pre_inpainting: Optional[InpaintingOptions] = Field( + description="Inpainting options", + default=None, + ) + + post_inpainting: Optional[InpaintingOptions] = Field( + description="Inpainting options", + default=None, + ) + def get_batch_images(self) -> List[Image.Image]: images = [] if self.batch_images: @@ -104,39 +145,15 @@ class PostProcessingOptions(BaseModel): upscaler_visibility: float = Field( description="upscaler visibility", default=1, le=1, ge=0 ) - - inpainting_denoising_strengh: float = Field( - description="Inpainting denoising strenght", default=0, lt=1, ge=0 - ) - inpainting_prompt: str = Field( - description="Inpainting denoising strenght", - examples=["Portrait of a [gender]"], - default="Portrait of a [gender]", - ) - inpainting_negative_prompt: str = Field( - description="Inpainting denoising strenght", - examples=[ - "Deformed, blurry, bad anatomy, disfigured, poorly drawn face, mutation" - ], - default="", - ) - inpainting_steps: int = Field( - description="Inpainting steps", - examples=["Portrait of a [gender]"], - ge=1, - le=150, - default=20, - ) - inpainting_sampler: str = Field( - description="Inpainting sampler", examples=["Euler"], default="Euler" - ) inpainting_when: InpaintingWhen = Field( description="When inpainting happens", examples=[e.value for e in InpaintingWhen.__members__.values()], default=InpaintingWhen.NEVER, ) - inpainting_model: str = Field( - description="Inpainting model", examples=["Current"], default="Current" + + inpainting_options: Optional[InpaintingOptions] = Field( + description="Inpainting options", + default=None, ) diff --git a/client_api/faceswaplab_api_example.py b/client_api/faceswaplab_api_example.py index fdeffbe..26b2975 100644 --- a/client_api/faceswaplab_api_example.py +++ b/client_api/faceswaplab_api_example.py @@ -1,13 +1,14 @@ import requests from api_utils import ( - FaceSwapRequest, FaceSwapUnit, - PostProcessingOptions, - FaceSwapResponse, pil_to_base64, + PostProcessingOptions, InpaintingWhen, - FaceSwapCompareRequest, + InpaintingOptions, + FaceSwapRequest, + FaceSwapResponse, FaceSwapExtractRequest, + FaceSwapCompareRequest, FaceSwapExtractResponse, ) @@ -37,9 +38,11 @@ pp = PostProcessingOptions( restorer_visibility=1, upscaler_name="Lanczos", scale=4, - inpainting_steps=30, - inpainting_denoising_strengh=0.1, inpainting_when=InpaintingWhen.BEFORE_RESTORE_FACE, + inpainting_options=InpaintingOptions( + inpainting_steps=30, + inpainting_denoising_strengh=0.1, + ), ) # Prepare the request diff --git a/requirements.txt b/requirements.txt index e48dc85..2999cc3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ onnxruntime==1.15.0 opencv-python==4.7.0.72 pandas pydantic==1.10.9 -dill==0.3.6 \ No newline at end of file +dill==0.3.6 +safetensors \ No newline at end of file diff --git a/scripts/faceswaplab.py b/scripts/faceswaplab.py index e730d56..8c49032 100644 --- a/scripts/faceswaplab.py +++ b/scripts/faceswaplab.py @@ -1,16 +1,16 @@ import importlib -from scripts.faceswaplab_api import faceswaplab_api -from scripts.faceswaplab_settings import faceswaplab_settings -from scripts.faceswaplab_ui import faceswaplab_tab, faceswaplab_unit_ui -from scripts.faceswaplab_utils.models_utils import ( - get_current_model, -) +import traceback from scripts import faceswaplab_globals -from scripts.faceswaplab_swapping import swapper -from scripts.faceswaplab_utils import faceswaplab_logging, imgutils -from scripts.faceswaplab_utils import models_utils +from scripts.faceswaplab_api import faceswaplab_api from scripts.faceswaplab_postprocessing import upscaling +from scripts.faceswaplab_settings import faceswaplab_settings +from scripts.faceswaplab_swapping import swapper +from scripts.faceswaplab_ui import faceswaplab_tab, faceswaplab_unit_ui +from scripts.faceswaplab_utils import faceswaplab_logging, imgutils, models_utils +from scripts.faceswaplab_utils.models_utils import get_current_model +from scripts.faceswaplab_utils.typing import * +from scripts.faceswaplab_utils.ui_utils import dataclasses_from_flat_list # Reload all the modules when using "apply and restart" # This is mainly done for development purposes @@ -25,14 +25,12 @@ importlib.reload(faceswaplab_unit_ui) importlib.reload(faceswaplab_api) import os -from dataclasses import fields from pprint import pformat from typing import Any, List, Optional, Tuple import gradio as gr import modules.scripts as scripts -from modules import script_callbacks, scripts -from modules import scripts, shared +from modules import script_callbacks, scripts, shared from modules.images import save_image from modules.processing import ( Processed, @@ -40,16 +38,14 @@ from modules.processing import ( StableDiffusionProcessingImg2Img, ) from modules.shared import opts -from PIL import Image -from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug from scripts.faceswaplab_globals import VERSION_FLAG +from scripts.faceswaplab_postprocessing.postprocessing import enhance_image from scripts.faceswaplab_postprocessing.postprocessing_options import ( PostProcessingOptions, ) -from scripts.faceswaplab_postprocessing.postprocessing import enhance_image from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings - +from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab") @@ -62,7 +58,9 @@ try: script_callbacks.on_app_started(faceswaplab_api.faceswaplab_api) except: - pass + logger.error("Failed to register API") + + traceback.print_exc() class FaceSwapScript(scripts.Script): @@ -107,44 +105,39 @@ class FaceSwapScript(scripts.Script): def ui(self, is_img2img: bool) -> List[gr.components.Component]: with gr.Accordion(f"FaceSwapLab {VERSION_FLAG}", open=False): - components = [] + components: List[gr.components.Component] = [] for i in range(1, self.units_count + 1): components += faceswaplab_unit_ui.faceswap_unit_ui(is_img2img, i) - upscaler = faceswaplab_tab.postprocessing_ui() + post_processing = faceswaplab_tab.postprocessing_ui() # If the order is modified, the before_process should be changed accordingly. - return components + upscaler + return components + post_processing def read_config( - self, p: StableDiffusionProcessing, *components: List[gr.components.Component] + self, p: StableDiffusionProcessing, *components: Tuple[Any, ...] ) -> None: + for i, c in enumerate(components): + logger.debug("%s>%s", i, pformat(c)) + # The order of processing for the components is important # The method first process faceswap units then postprocessing units - - # self.make_first_script(p) - + classes: List[Any] = dataclasses_from_flat_list( + [FaceSwapUnitSettings] * self.units_count + [PostProcessingOptions], + components, + ) self.units: List[FaceSwapUnitSettings] = [] - - # Parse and convert units flat components into FaceSwapUnitSettings - for i in range(0, self.units_count): - self.units += [FaceSwapUnitSettings.get_unit_configuration(i, components)] + self.units += [u for u in classes if isinstance(u, FaceSwapUnitSettings)] + self.postprocess_options = classes[-1] for i, u in enumerate(self.units): logger.debug("%s, %s", pformat(i), pformat(u)) - # Parse the postprocessing options - # We must first find where to start from (after face swapping units) - len_conf: int = len(fields(FaceSwapUnitSettings)) - shift: int = self.units_count * len_conf - self.postprocess_options = PostProcessingOptions( - *components[shift : shift + len(fields(PostProcessingOptions))] # type: ignore - ) logger.debug("%s", pformat(self.postprocess_options)) if self.enabled: p.do_not_save_samples = not self.keep_original_images def process( - self, p: StableDiffusionProcessing, *components: List[gr.components.Component] + self, p: StableDiffusionProcessing, *components: Tuple[Any, ...] ) -> None: try: self.read_config(p, *components) @@ -152,7 +145,7 @@ class FaceSwapScript(scripts.Script): # 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]]] = [ + init_images: List[Tuple[Optional[PILImage], Optional[str]]] = [ (img, None) for img in p.init_images ] new_inits = swapper.process_images_units( @@ -167,6 +160,7 @@ class FaceSwapScript(scripts.Script): p.init_images = [img[0] for img in new_inits] except Exception as e: logger.info("Failed to process : %s", e) + traceback.print_exc() def postprocess( self, p: StableDiffusionProcessing, processed: Processed, *args: List[Any] @@ -174,7 +168,7 @@ class FaceSwapScript(scripts.Script): try: if self.enabled: # Get the original images without the grid - orig_images: List[Image.Image] = processed.images[ + orig_images: List[PILImage] = processed.images[ processed.index_of_first_image : ] orig_infotexts: List[str] = processed.infotexts[ @@ -237,7 +231,6 @@ class FaceSwapScript(scripts.Script): # 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) @@ -245,6 +238,20 @@ class FaceSwapScript(scripts.Script): grid.info["parameters"] = text images.insert(0, grid) + if opts.grid_save: + save_image( + grid, + p.outpath_grids, + "swapped-grid", + p.all_seeds[0], + p.all_prompts[0], + opts.grid_format, + info=text, + short_filename=not opts.grid_extended_filename, + p=p, + grid=True, + ) + if keep_original: # If we want to keep original images, we add all existing (including grid this time) images += processed.images @@ -254,3 +261,4 @@ class FaceSwapScript(scripts.Script): processed.infotexts = infotexts except Exception as e: logger.error("Failed to swap face in postprocess method : %s", e) + traceback.print_exc() diff --git a/scripts/faceswaplab_api/faceswaplab_api.py b/scripts/faceswaplab_api/faceswaplab_api.py index 34c4597..de280d2 100644 --- a/scripts/faceswaplab_api/faceswaplab_api.py +++ b/scripts/faceswaplab_api/faceswaplab_api.py @@ -17,7 +17,6 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import ( PostProcessingOptions, ) from client_api import api_utils -from scripts.faceswaplab_postprocessing.postprocessing_options import InpaintingWhen def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str: # type: ignore @@ -58,58 +57,12 @@ def encode_np_to_base64(image: np.ndarray) -> str: # type: ignore return api.encode_pil_to_base64(pil) -def get_postprocessing_options( - options: api_utils.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, - # hacky way to prevent having a separate file for Inpainting when (2 classes) - # therfore a conversion is required from api IW to server side IW - inpainting_when=InpaintingWhen(options.inpainting_when.value), - 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[api_utils.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, - ) - ) + units.append(FaceSwapUnitSettings.from_api_dto(u)) return units @@ -137,7 +90,7 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None: if src_image is not None: if request.postprocessing: - pp_options = get_postprocessing_options(request.postprocessing) + pp_options = PostProcessingOptions.from_api_dto(request.postprocessing) units = get_faceswap_units_settings(request.units) swapped_images = swapper.batch_process( @@ -172,7 +125,7 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None: ) -> api_utils.FaceSwapExtractResponse: pp_options = None if request.postprocessing: - pp_options = get_postprocessing_options(request.postprocessing) + pp_options = PostProcessingOptions.from_api_dto(request.postprocessing) images = [base64_to_pil(img) for img in request.images] faces = swapper.extract_faces( images, extract_path=None, postprocess_options=pp_options diff --git a/scripts/faceswaplab_globals.py b/scripts/faceswaplab_globals.py index a93249c..11f70f4 100644 --- a/scripts/faceswaplab_globals.py +++ b/scripts/faceswaplab_globals.py @@ -8,7 +8,7 @@ REFERENCE_PATH = os.path.join( scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references" ) -VERSION_FLAG: str = "v1.1.2" +VERSION_FLAG: str = "v1.2.0" EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab") # The NSFW score threshold. If any part of the image has a score greater than this threshold, the image will be considered NSFW. diff --git a/scripts/faceswaplab_inpainting/faceswaplab_inpainting.py b/scripts/faceswaplab_inpainting/faceswaplab_inpainting.py new file mode 100644 index 0000000..716d1b3 --- /dev/null +++ b/scripts/faceswaplab_inpainting/faceswaplab_inpainting.py @@ -0,0 +1,41 @@ +from dataclasses import dataclass +from typing import List +import gradio as gr +from client_api import api_utils + + +@dataclass +class InpaintingOptions: + inpainting_denoising_strengh: float = 0 + inpainting_prompt: str = "" + inpainting_negative_prompt: str = "" + inpainting_steps: int = 20 + inpainting_sampler: str = "Euler" + inpainting_model: str = "Current" + + @staticmethod + def from_gradio(components: List[gr.components.Component]) -> "InpaintingOptions": + return InpaintingOptions(*components) + + @staticmethod + def from_api_dto(dto: api_utils.InpaintingOptions) -> "InpaintingOptions": + """ + Converts a InpaintingOptions object from an API DTO (Data Transfer Object). + + :param options: An object of api_utils.InpaintingOptions representing the + post-processing options as received from the API. + :return: A InpaintingOptions instance containing the translated values + from the API DTO. + """ + if dto is None: + # Return default values + return InpaintingOptions() + + return InpaintingOptions( + inpainting_denoising_strengh=dto.inpainting_denoising_strengh, + inpainting_prompt=dto.inpainting_prompt, + inpainting_negative_prompt=dto.inpainting_negative_prompt, + inpainting_steps=dto.inpainting_steps, + inpainting_sampler=dto.inpainting_sampler, + inpainting_model=dto.inpainting_model, + ) diff --git a/scripts/faceswaplab_postprocessing/i2i_pp.py b/scripts/faceswaplab_inpainting/i2i_pp.py similarity index 50% rename from scripts/faceswaplab_postprocessing/i2i_pp.py rename to scripts/faceswaplab_inpainting/i2i_pp.py index 09ebce2..61a4850 100644 --- a/scripts/faceswaplab_postprocessing/i2i_pp.py +++ b/scripts/faceswaplab_inpainting/i2i_pp.py @@ -1,53 +1,60 @@ +from scripts.faceswaplab_inpainting.faceswaplab_inpainting import InpaintingOptions from scripts.faceswaplab_utils.faceswaplab_logging import logger from PIL import Image from modules import shared from scripts.faceswaplab_utils import imgutils from modules import shared, processing from modules.processing import StableDiffusionProcessingImg2Img -from scripts.faceswaplab_postprocessing.postprocessing_options import ( - PostProcessingOptions, -) from modules import sd_models - +import traceback from scripts.faceswaplab_swapping import swapper +from scripts.faceswaplab_utils.typing import * +from typing import * -def img2img_diffusion(img: Image.Image, pp: PostProcessingOptions) -> Image.Image: - if pp.inpainting_denoising_strengh == 0: - logger.info("Discard inpainting denoising strength is 0") +def img2img_diffusion( + img: PILImage, options: InpaintingOptions, faces: Optional[List[Face]] = None +) -> Image.Image: + if not options or options.inpainting_denoising_strengh == 0: + logger.info("Discard inpainting denoising strength is 0 or no inpainting") return img try: logger.info( f"""Inpainting face -Sampler : {pp.inpainting_sampler} -inpainting_denoising_strength : {pp.inpainting_denoising_strengh} -inpainting_steps : {pp.inpainting_steps} +Sampler : {options.inpainting_sampler} +inpainting_denoising_strength : {options.inpainting_denoising_strengh} +inpainting_steps : {options.inpainting_steps} """ ) - if not isinstance(pp.inpainting_sampler, str): - pp.inpainting_sampler = "Euler" + if not isinstance(options.inpainting_sampler, str): + options.inpainting_sampler = "Euler" logger.info("send faces to image to image") img = img.copy() - faces = swapper.get_faces(imgutils.pil_to_cv2(img)) + + if not faces: + faces = swapper.get_faces(imgutils.pil_to_cv2(img)) + if faces: for face in faces: bbox = face.bbox.astype(int) mask = imgutils.create_mask(img, bbox) - prompt = pp.inpainting_prompt.replace( + prompt = options.inpainting_prompt.replace( "[gender]", "man" if face["gender"] == 1 else "woman" ) - negative_prompt = pp.inpainting_negative_prompt.replace( + negative_prompt = options.inpainting_negative_prompt.replace( "[gender]", "man" if face["gender"] == 1 else "woman" ) logger.info("Denoising prompt : %s", prompt) - logger.info("Denoising strenght : %s", pp.inpainting_denoising_strengh) + logger.info( + "Denoising strenght : %s", options.inpainting_denoising_strengh + ) i2i_kwargs = { - "sampler_name": pp.inpainting_sampler, + "sampler_name": options.inpainting_sampler, "do_not_save_samples": True, - "steps": pp.inpainting_steps, + "steps": options.inpainting_steps, "width": img.width, "inpainting_fill": 1, "inpaint_full_res": True, @@ -55,17 +62,26 @@ inpainting_steps : {pp.inpainting_steps} "mask": mask, "prompt": prompt, "negative_prompt": negative_prompt, - "denoising_strength": pp.inpainting_denoising_strengh, + "denoising_strength": options.inpainting_denoising_strengh, + "override_settings": { + "return_mask_composite": False, + "save_images_before_face_restoration": False, + "save_images_before_highres_fix": False, + "save_images_before_color_correction": False, + "save_mask": False, + "save_mask_composite": False, + "samples_save": False, + }, } current_model_checkpoint = shared.opts.sd_model_checkpoint - if pp.inpainting_model and pp.inpainting_model != "Current": + if options.inpainting_model and options.inpainting_model != "Current": # Change checkpoint - shared.opts.sd_model_checkpoint = pp.inpainting_model + shared.opts.sd_model_checkpoint = options.inpainting_model sd_models.select_checkpoint sd_models.load_model() i2i_p = StableDiffusionProcessingImg2Img([img], **i2i_kwargs) i2i_processed = processing.process_images(i2i_p) - if pp.inpainting_model and pp.inpainting_model != "Current": + if options.inpainting_model and options.inpainting_model != "Current": # Restore checkpoint shared.opts.sd_model_checkpoint = current_model_checkpoint sd_models.select_checkpoint @@ -76,8 +92,6 @@ inpainting_steps : {pp.inpainting_steps} img = images[0] return img except Exception as e: - logger.error("Failed to apply img2img to face : %s", e) - import traceback - + logger.error("Failed to apply inpainting to face : %s", e) traceback.print_exc() raise e diff --git a/scripts/faceswaplab_postprocessing/postprocessing.py b/scripts/faceswaplab_postprocessing/postprocessing.py index 813a5dc..cba191e 100644 --- a/scripts/faceswaplab_postprocessing/postprocessing.py +++ b/scripts/faceswaplab_postprocessing/postprocessing.py @@ -4,8 +4,9 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import ( PostProcessingOptions, InpaintingWhen, ) -from scripts.faceswaplab_postprocessing.i2i_pp import img2img_diffusion +from scripts.faceswaplab_inpainting.i2i_pp import img2img_diffusion from scripts.faceswaplab_postprocessing.upscaling import upscale_img, restore_face +import traceback def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image: @@ -19,7 +20,9 @@ def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Imag or pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING ): logger.debug("Inpaint before upscale") - result_image = img2img_diffusion(result_image, pp_options) + result_image = img2img_diffusion( + img=result_image, options=pp_options.inpainting_options + ) result_image = upscale_img(result_image, pp_options) if ( @@ -27,7 +30,9 @@ def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Imag or pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE ): logger.debug("Inpaint before restore") - result_image = img2img_diffusion(result_image, pp_options) + result_image = img2img_diffusion( + result_image, pp_options.inpainting_options + ) result_image = restore_face(result_image, pp_options) @@ -36,9 +41,11 @@ def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Imag or pp_options.inpainting_when == InpaintingWhen.AFTER_ALL ): logger.debug("Inpaint after all") - result_image = img2img_diffusion(result_image, pp_options) + result_image = img2img_diffusion( + result_image, pp_options.inpainting_options + ) except Exception as e: - logger.error("Failed to upscale %s", e) - + logger.error("Failed to post-process %s", e) + traceback.print_exc() return result_image diff --git a/scripts/faceswaplab_postprocessing/postprocessing_options.py b/scripts/faceswaplab_postprocessing/postprocessing_options.py index e2f6f39..b7f7bcd 100644 --- a/scripts/faceswaplab_postprocessing/postprocessing_options.py +++ b/scripts/faceswaplab_postprocessing/postprocessing_options.py @@ -3,6 +3,8 @@ from modules.upscaler import UpscalerData from dataclasses import dataclass from modules import shared from enum import Enum +from scripts.faceswaplab_inpainting.faceswaplab_inpainting import InpaintingOptions +from client_api import api_utils class InpaintingWhen(Enum): @@ -22,13 +24,10 @@ class PostProcessingOptions: scale: float = 1 upscale_visibility: float = 0.5 - inpainting_denoising_strengh: float = 0 - inpainting_prompt: str = "" - inpainting_negative_prompt: str = "" - inpainting_steps: int = 20 - inpainting_sampler: str = "Euler" inpainting_when: InpaintingWhen = InpaintingWhen.BEFORE_UPSCALING - inpainting_model: str = "Current" + + # (Don't use optional for this or gradio parsing will fail) : + inpainting_options: InpaintingOptions = None @property def upscaler(self) -> UpscalerData: @@ -43,3 +42,28 @@ class PostProcessingOptions: if face_restorer.name() == self.face_restorer_name: return face_restorer return None + + @staticmethod + def from_api_dto( + options: api_utils.PostProcessingOptions, + ) -> "PostProcessingOptions": + """ + Converts a PostProcessingOptions object from an API DTO (Data Transfer Object). + + :param options: An object of api_utils.PostProcessingOptions representing the + post-processing options as received from the API. + :return: A PostProcessingOptions instance containing the translated values + from the API DTO. + """ + return 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_when=InpaintingWhen(options.inpainting_when.value), + inpainting_options=InpaintingOptions.from_api_dto( + options.inpainting_options + ), + ) diff --git a/scripts/faceswaplab_postprocessing/upscaling.py b/scripts/faceswaplab_postprocessing/upscaling.py index 04ba5fb..746f229 100644 --- a/scripts/faceswaplab_postprocessing/upscaling.py +++ b/scripts/faceswaplab_postprocessing/upscaling.py @@ -5,11 +5,12 @@ from scripts.faceswaplab_utils.faceswaplab_logging import logger from PIL import Image import numpy as np from modules import codeformer_model +from scripts.faceswaplab_utils.typing import * -def upscale_img(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image: +def upscale_img(image: PILImage, pp_options: PostProcessingOptions) -> PILImage: if pp_options.upscaler is not None and pp_options.upscaler.name != "None": - original_image = image.copy() + original_image: PILImage = image.copy() logger.info( "Upscale with %s scale = %s", pp_options.upscaler.name, @@ -18,7 +19,12 @@ def upscale_img(image: Image.Image, pp_options: PostProcessingOptions) -> Image. result_image = pp_options.upscaler.scaler.upscale( image, pp_options.scale, pp_options.upscaler.data_path ) - if pp_options.scale == 1: + + # FIXME : Could be better (managing images whose dimensions are not multiples of 16) + if pp_options.scale == 1 and original_image.size == result_image.size: + logger.debug( + "Sizes orig=%s, result=%s", original_image.size, result_image.size + ) result_image = Image.blend( original_image, result_image, pp_options.upscale_visibility ) diff --git a/scripts/faceswaplab_swapping/facemask.py b/scripts/faceswaplab_swapping/facemask.py index 8af364c..d49b713 100644 --- a/scripts/faceswaplab_swapping/facemask.py +++ b/scripts/faceswaplab_swapping/facemask.py @@ -20,7 +20,7 @@ def get_parsing_model(device: torch_device) -> torch.nn.Module: Returns: The parsing model. """ - return init_parsing_model(device=device) + return init_parsing_model(device=device) # type: ignore def convert_image_to_tensor( diff --git a/scripts/faceswaplab_swapping/parsing/__init__.py b/scripts/faceswaplab_swapping/parsing/__init__.py index 39fe887..e045dec 100644 --- a/scripts/faceswaplab_swapping/parsing/__init__.py +++ b/scripts/faceswaplab_swapping/parsing/__init__.py @@ -50,7 +50,7 @@ from scripts.faceswaplab_globals import FACE_PARSER_DIR ROOT_DIR = FACE_PARSER_DIR -def load_file_from_url(url, model_dir=None, progress=True, file_name=None): +def load_file_from_url(url: str, model_dir=None, progress=True, file_name=None): """Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py""" if model_dir is None: hub_dir = get_dir() diff --git a/scripts/faceswaplab_swapping/swapper.py b/scripts/faceswaplab_swapping/swapper.py index 27496b6..d2dd2d4 100644 --- a/scripts/faceswaplab_swapping/swapper.py +++ b/scripts/faceswaplab_swapping/swapper.py @@ -7,7 +7,7 @@ import tempfile import cv2 import insightface import numpy as np -from insightface.app.common import Face +from insightface.app.common import Face as ISFace from PIL import Image from sklearn.metrics.pairwise import cosine_similarity @@ -28,7 +28,8 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import ( ) from scripts.faceswaplab_utils.models_utils import get_current_model import gradio as gr - +from scripts.faceswaplab_utils.typing import CV2ImgU8, PILImage, Face +from scripts.faceswaplab_inpainting.i2i_pp import img2img_diffusion providers = ["CPUExecutionProvider"] @@ -60,7 +61,7 @@ def cosine_similarity_face(face1: Face, face2: Face) -> float: return max(0, similarity[0, 0]) -def compare_faces(img1: Image.Image, img2: Image.Image) -> float: +def compare_faces(img1: PILImage, img2: PILImage) -> float: """ Compares the similarity between two faces extracted from images using cosine similarity. @@ -87,22 +88,22 @@ def compare_faces(img1: Image.Image, img2: Image.Image) -> float: def batch_process( - src_images: List[Image.Image], + src_images: List[PILImage], save_path: Optional[str], units: List[FaceSwapUnitSettings], postprocess_options: PostProcessingOptions, -) -> Optional[List[Image.Image]]: +) -> 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. Args: - src_images (List[Image.Image]): List of source PIL Images to process. + src_images (List[PILImage]): List of source PIL Images to process. save_path (Optional[str]): Destination path where the processed images will be saved. If None, no images are saved. units (List[FaceSwapUnitSettings]): List of FaceSwapUnitSettings to apply to the images. postprocess_options (PostProcessingOptions): Post-processing settings to be applied to the images. Returns: - Optional[List[Image.Image]]: List of processed images, or None in case of an exception. + Optional[List[PILImage]]: List of processed images, or None in case of an exception. Raises: Any exceptions raised by the underlying process will be logged and the function will return None. @@ -149,7 +150,7 @@ def batch_process( def extract_faces( - images: List[Image.Image], + images: List[PILImage], extract_path: Optional[str], postprocess_options: PostProcessingOptions, ) -> Optional[List[str]]: @@ -206,7 +207,7 @@ def extract_faces( return result_images except Exception as e: - logger.info("Failed to extract : %s", e) + logger.error("Failed to extract : %s", e) import traceback traceback.print_exc() @@ -273,16 +274,15 @@ def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper: def get_faces( - img_data: np.ndarray, # type: ignore + img_data: CV2ImgU8, det_size: Tuple[int, int] = (640, 640), det_thresh: Optional[float] = None, - sort_by_face_size: bool = False, ) -> List[Face]: """ Detects and retrieves faces from an image using an analysis model. Args: - img_data (np.ndarray): The image data as a NumPy array. + img_data (CV2ImgU8): The image data as a NumPy array. det_size (tuple): The desired detection size (width, height). Defaults to (640, 640). sort_by_face_size (bool) : Will sort the faces by their size from larger to smaller face @@ -309,26 +309,55 @@ def get_faces( return get_faces(img_data, det_size=det_size_half, det_thresh=det_thresh) try: - if sort_by_face_size: - return sorted( - face, - reverse=True, - key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), - ) - # Sort the detected faces based on their x-coordinate of the bounding box return sorted(face, key=lambda x: x.bbox[0]) except Exception as e: return [] +def filter_faces( + all_faces: List[Face], + faces_index: Set[int], + source_gender: int = None, + sort_by_face_size: bool = False, +) -> List[Face]: + """ + Sorts and filters a list of faces based on specified criteria. + + This function takes a list of Face objects and can sort them by face size and filter them by gender. + Sorting by face size is performed if sort_by_face_size is set to True, and filtering by gender is + performed if source_gender is provided. + + :param faces: A list of Face objects representing the faces to be sorted and filtered. + :param faces_index: A set of faces index + :param source_gender: An optional integer representing the gender by which to filter the faces. + If provided, only faces with the specified gender will be included in the result. + :param sort_by_face_size: A boolean indicating whether to sort the faces by size. If True, faces are + sorted in descending order by size, calculated as the area of the bounding box. + :return: A list of Face objects sorted and filtered according to the specified criteria. + """ + filtered_faces = copy.copy(all_faces) + if sort_by_face_size: + filtered_faces = sorted( + all_faces, + reverse=True, + key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), + ) + + if source_gender is not None: + filtered_faces = [ + face for face in filtered_faces if face["gender"] == source_gender + ] + return [face for i, face in enumerate(filtered_faces) if i in faces_index] + + @dataclass class ImageResult: """ Represents the result of an image swap operation """ - image: Image.Image + image: PILImage """ The image object with the swapped face """ @@ -362,7 +391,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(files: List[gr.File]) -> List[Optional[np.ndarray]]: # type: ignore +def get_faces_from_img_files(files: List[gr.File]) -> List[Optional[CV2ImgU8]]: """ Extracts faces from a list of image files. @@ -388,7 +417,7 @@ def get_faces_from_img_files(files: List[gr.File]) -> List[Optional[np.ndarray]] return faces -def blend_faces(faces: List[Face]) -> Face: +def blend_faces(faces: List[Face]) -> Optional[Face]: """ Blends the embeddings of multiple faces into a single face. @@ -418,16 +447,10 @@ def blend_faces(faces: List[Face]) -> Face: # Create a new Face object using the properties of the first face in the list # Assign the blended embedding to the blended Face object - blended = Face( + blended = ISFace( embedding=blended_embedding, gender=faces[0].gender, age=faces[0].age ) - assert ( - not np.array_equal(blended.embedding, faces[0].embedding) - if len(faces) > 1 - else True - ), "If len(faces)>0, the blended embedding should not be the same than the first image" - return blended # Return None if the input list is empty @@ -435,85 +458,80 @@ def blend_faces(faces: List[Face]) -> Face: def swap_face( - reference_face: np.ndarray, # type: ignore - source_face: np.ndarray, # type: ignore - target_img: Image.Image, + reference_face: CV2ImgU8, + source_face: Face, + target_img: PILImage, + target_faces: List[Face], model: str, - faces_index: Set[int] = {0}, - same_gender: bool = True, upscaled_swapper: bool = False, compute_similarity: bool = True, - sort_by_face_size: bool = False, ) -> ImageResult: """ Swaps faces in the target image with the source face. Args: - reference_face (np.ndarray): The reference face used for similarity comparison. - source_face (np.ndarray): The source face to be swapped. - target_img (Image.Image): The target image to swap faces in. + 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. - faces_index (Set[int], optional): Set of indices specifying which faces to swap. Defaults to {0}. - same_gender (bool, optional): If True, only swap faces with the same gender as the source face. Defaults to True. Returns: ImageResult: An object containing the swapped image and similarity scores. """ return_result = ImageResult(target_img, {}, {}) + target_img_cv2: CV2ImgU8 = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) try: - target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) gender = source_face["gender"] logger.info("Source Gender %s", gender) if source_face is not None: - result = target_img + result = target_img_cv2 model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model) face_swapper = getFaceSwapModel(model_path) - target_faces = get_faces(target_img, sort_by_face_size=sort_by_face_size) logger.info("Target faces count : %s", len(target_faces)) - if same_gender: - target_faces = [x for x in target_faces if x["gender"] == gender] - logger.info("Target Gender Matches count %s", len(target_faces)) - for i, swapped_face in enumerate(target_faces): logger.info(f"swap face {i}") - if i in faces_index: - # type : ignore - result = face_swapper.get( - result, swapped_face, source_face, upscale=upscaled_swapper - ) + + result = face_swapper.get( + img=result, + target_face=swapped_face, + source_face=source_face, + upscale=upscaled_swapper, + ) # type: ignore result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) return_result.image = result_image - if compute_similarity: - try: - result_faces = get_faces( - cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR), - sort_by_face_size=sort_by_face_size, - ) - if same_gender: - result_faces = [ - x for x in result_faces if x["gender"] == gender - ] - - for i, swapped_face in enumerate(result_faces): - logger.info(f"compare face {i}") - if i in faces_index and i < len(target_faces): - return_result.similarity[i] = cosine_similarity_face( - source_face, swapped_face - ) - return_result.ref_similarity[i] = cosine_similarity_face( - reference_face, swapped_face - ) - - logger.info(f"similarity {return_result.similarity}") - logger.info(f"ref similarity {return_result.ref_similarity}") - - except Exception as e: - logger.error("Similarity processing failed %s", e) - raise e + # FIXME : recompute similarity + + # if compute_similarity: + # try: + # result_faces = get_faces( + # cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR), + # sort_by_face_size=sort_by_face_size, + # ) + # if same_gender: + # result_faces = [ + # x for x in result_faces if x["gender"] == gender + # ] + + # for i, swapped_face in enumerate(result_faces): + # logger.info(f"compare face {i}") + # if i in faces_index and i < len(target_faces): + # return_result.similarity[i] = cosine_similarity_face( + # source_face, swapped_face + # ) + # return_result.ref_similarity[i] = cosine_similarity_face( + # reference_face, swapped_face + # ) + + # logger.info(f"similarity {return_result.similarity}") + # logger.info(f"ref similarity {return_result.ref_similarity}") + + # except Exception as e: + # logger.error("Similarity processing failed %s", e) + # raise e except Exception as e: logger.error("Conversion failed %s", e) raise e @@ -523,11 +541,11 @@ def swap_face( def process_image_unit( model: str, unit: FaceSwapUnitSettings, - image: Image.Image, + image: PILImage, info: str = None, upscaled_swapper: bool = False, force_blend: bool = False, -) -> List[Tuple[Image.Image, str]]: +) -> List[Tuple[PILImage, str]]: """Process one image and return a List of (image, info) (one if blended, many if not). Args: @@ -541,6 +559,8 @@ def process_image_unit( results = [] if unit.enable: + faces = get_faces(pil_to_cv2(image)) + if check_against_nsfw(image): return [(image, info)] if not unit.blend_faces and not force_blend: @@ -549,15 +569,10 @@ def process_image_unit( else: logger.info("blend all faces together") src_faces = [unit.blended_faces] - assert ( - not np.array_equal( - unit.reference_face.embedding, src_faces[0].embedding - ) - if len(unit.faces) > 1 - else True - ), "Reference face cannot be the same as blended" for i, src_face in enumerate(src_faces): + current_image = image + logger.info(f"Process face {i}") if unit.reference_face is not None: reference_face = unit.reference_face @@ -565,18 +580,35 @@ def process_image_unit( logger.info("Use source face as reference face") reference_face = src_face + target_faces = filter_faces( + faces, + faces_index=unit.faces_index, + source_gender=src_face["gender"] if unit.same_gender else None, + sort_by_face_size=unit.sort_by_size, + ) + + # Apply pre-inpainting to image + if unit.pre_inpainting.inpainting_denoising_strengh > 0: + current_image = img2img_diffusion( + img=current_image, faces=target_faces, options=unit.pre_inpainting + ) + save_img_debug(image, "Before swap") result: ImageResult = swap_face( - reference_face, - src_face, - image, - faces_index=unit.faces_index, + reference_face=reference_face, + source_face=src_face, + target_img=current_image, + target_faces=target_faces, model=model, - same_gender=unit.same_gender, upscaled_swapper=upscaled_swapper, compute_similarity=unit.compute_similarity, - sort_by_face_size=unit.sort_by_size, ) + # Apply post-inpainting to image + if unit.post_inpainting.inpainting_denoising_strengh > 0: + result.image = img2img_diffusion( + img=result.image, faces=target_faces, options=unit.post_inpainting + ) + save_img_debug(result.image, "After swap") if result.image is None: @@ -610,17 +642,17 @@ def process_image_unit( def process_images_units( model: str, units: List[FaceSwapUnitSettings], - images: List[Tuple[Optional[Image.Image], Optional[str]]], + images: List[Tuple[Optional[PILImage], Optional[str]]], upscaled_swapper: bool = False, force_blend: bool = False, -) -> Optional[List[Tuple[Image.Image, str]]]: +) -> Optional[List[Tuple[PILImage, str]]]: """ Process a list of images using a specified model and unit settings for face swapping. Args: model (str): The name of the model to use for processing. units (List[FaceSwapUnitSettings]): A list of settings for face swap units to apply on each image. - images (List[Tuple[Optional[Image.Image], Optional[str]]]): A list of tuples, each containing + images (List[Tuple[Optional[PILImage], Optional[str]]]): A list of tuples, each containing an image and its associated info string. If an image or info string is not available, its value can be None. upscaled_swapper (bool, optional): If True, uses an upscaled version of the face swapper. @@ -629,7 +661,7 @@ def process_images_units( image. Defaults to False. Returns: - Optional[List[Tuple[Image.Image, str]]]: A list of tuples, each containing a processed image + Optional[List[Tuple[PILImage, str]]]: A list of tuples, each containing a processed image and its associated info string. If no units are provided for processing, returns None. """ diff --git a/scripts/faceswaplab_swapping/upscaled_inswapper.py b/scripts/faceswaplab_swapping/upscaled_inswapper.py index 2726e1b..a441dc0 100644 --- a/scripts/faceswaplab_swapping/upscaled_inswapper.py +++ b/scripts/faceswaplab_swapping/upscaled_inswapper.py @@ -1,3 +1,4 @@ +from typing import Any, Tuple, Union import cv2 import numpy as np from insightface.model_zoo.inswapper import INSwapper @@ -12,6 +13,7 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import ( ) from scripts.faceswaplab_swapping.facemask import generate_face_mask from scripts.faceswaplab_utils.imgutils import cv2_to_pil, pil_to_cv2 +from scripts.faceswaplab_utils.typing import CV2ImgU8, Face def get_upscaler() -> UpscalerData: @@ -23,7 +25,25 @@ def get_upscaler() -> UpscalerData: return None -def merge_images_with_mask(image1, image2, mask): +def merge_images_with_mask( + image1: CV2ImgU8, image2: CV2ImgU8, mask: CV2ImgU8 +) -> CV2ImgU8: + """ + Merges two images using a given mask. The regions where the mask is set will be replaced with the corresponding + areas of the second image. + + Args: + image1 (CV2Img): The base image, which must have the same shape as image2. + image2 (CV2Img): The image to be merged, which must have the same shape as image1. + mask (CV2Img): A binary mask specifying the regions to be merged. The mask shape should match image1's first two dimensions. + + Returns: + CV2Img: The merged image. + + Raises: + ValueError: If the shapes of the images and mask do not match. + """ + if image1.shape != image2.shape or image1.shape[:2] != mask.shape: raise ValueError("Img should have the same shape") mask = mask.astype(np.uint8) @@ -34,42 +54,80 @@ def merge_images_with_mask(image1, image2, mask): return merged_image -def erode_mask(mask, kernel_size=3, iterations=1): +def erode_mask(mask: CV2ImgU8, kernel_size: int = 3, iterations: int = 1) -> CV2ImgU8: + """ + Erodes a binary mask using a given kernel size and number of iterations. + + Args: + mask (CV2Img): The binary mask to erode. + kernel_size (int, optional): The size of the kernel. Default is 3. + iterations (int, optional): The number of erosion iterations. Default is 1. + + Returns: + CV2Img: The eroded mask. + """ kernel = np.ones((kernel_size, kernel_size), np.uint8) eroded_mask = cv2.erode(mask, kernel, iterations=iterations) return eroded_mask -def apply_gaussian_blur(mask, kernel_size=(5, 5), sigma_x=0): +def apply_gaussian_blur( + mask: CV2ImgU8, kernel_size: Tuple[int, int] = (5, 5), sigma_x: int = 0 +) -> CV2ImgU8: + """ + Applies a Gaussian blur to a mask. + + Args: + mask (CV2Img): The mask to blur. + kernel_size (tuple, optional): The size of the kernel, e.g. (5, 5). Default is (5, 5). + sigma_x (int, optional): The standard deviation in the X direction. Default is 0. + + Returns: + CV2Img: The blurred mask. + """ blurred_mask = cv2.GaussianBlur(mask, kernel_size, sigma_x) return blurred_mask -def dilate_mask(mask, kernel_size=5, iterations=1): +def dilate_mask(mask: CV2ImgU8, kernel_size: int = 5, iterations: int = 1) -> CV2ImgU8: + """ + Dilates a binary mask using a given kernel size and number of iterations. + + Args: + mask (CV2Img): The binary mask to dilate. + kernel_size (int, optional): The size of the kernel. Default is 5. + iterations (int, optional): The number of dilation iterations. Default is 1. + + Returns: + CV2Img: The dilated mask. + """ kernel = np.ones((kernel_size, kernel_size), np.uint8) dilated_mask = cv2.dilate(mask, kernel, iterations=iterations) return dilated_mask -def get_face_mask(aimg, bgr_fake): +def get_face_mask(aimg: CV2ImgU8, bgr_fake: CV2ImgU8) -> CV2ImgU8: + """ + Generates a face mask by performing bitwise OR on two face masks and then dilating the result. + + Args: + aimg (CV2Img): Input image for generating the first face mask. + bgr_fake (CV2Img): Input image for generating the second face mask. + + Returns: + CV2Img: The combined and dilated face mask. + """ mask1 = generate_face_mask(aimg, device=shared.device) mask2 = generate_face_mask(bgr_fake, device=shared.device) mask = dilate_mask(cv2.bitwise_or(mask1, mask2)) return mask -class UpscaledINSwapper: +class UpscaledINSwapper(INSwapper): def __init__(self, inswapper: INSwapper): self.__dict__.update(inswapper.__dict__) - def forward(self, img, latent): - img = (img - self.input_mean) / self.input_std - pred = self.session.run( - self.output_names, {self.input_names[0]: img, self.input_names[1]: latent} - )[0] - return pred - - def super_resolution(self, img, k=2): + def super_resolution(self, img: CV2ImgU8, k: int = 2) -> CV2ImgU8: pil_img = cv2_to_pil(img) options = PostProcessingOptions( upscaler_name=opts.data.get( @@ -91,7 +149,14 @@ class UpscaledINSwapper: upscaled = upscaling.restore_face(upscaled, options) return pil_to_cv2(upscaled) - def get(self, img, target_face, source_face, paste_back=True, upscale=True): + def get( + self, + img: CV2ImgU8, + target_face: Face, + source_face: Face, + paste_back: bool = True, + upscale: bool = True, + ) -> Union[CV2ImgU8, Tuple[CV2ImgU8, Any]]: aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0]) blob = cv2.dnn.blobFromImage( aimg, @@ -116,7 +181,7 @@ class UpscaledINSwapper: else: target_img = img - def compute_diff(bgr_fake, aimg): + def compute_diff(bgr_fake: CV2ImgU8, aimg: CV2ImgU8) -> CV2ImgU8: fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32) fake_diff = np.abs(fake_diff).mean(axis=2) fake_diff[:2, :] = 0 diff --git a/scripts/faceswaplab_ui/faceswaplab_inpainting_ui.py b/scripts/faceswaplab_ui/faceswaplab_inpainting_ui.py new file mode 100644 index 0000000..bb4dd93 --- /dev/null +++ b/scripts/faceswaplab_ui/faceswaplab_inpainting_ui.py @@ -0,0 +1,68 @@ +from typing import List +import gradio as gr +from modules.shared import opts +from modules import sd_models, sd_samplers + + +def face_inpainting_ui( + name: str, id_prefix: str = "faceswaplab", description: str = "" +) -> List[gr.components.Component]: + with gr.Accordion(name, open=False): + gr.Markdown(description) + inpainting_denoising_strength = gr.Slider( + 0, + 1, + 0, + step=0.01, + elem_id=f"{id_prefix}_pp_inpainting_denoising_strength", + label="Denoising strenght", + ) + + inpainting_denoising_prompt = gr.Textbox( + opts.data.get( + "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( + "faceswaplab_pp_default_inpainting_negative_prompt", "blurry" + ), + elem_id=f"{id_prefix}_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 sd_samplers.all_samplers] + inpainting_sampler = gr.Dropdown( + choices=samplers_names, + value=[samplers_names[0]], + label="Inpainting Sampler", + elem_id=f"{id_prefix}_pp_inpainting_sampler", + ) + inpainting_denoising_steps = gr.Slider( + 1, + 150, + 20, + step=1, + label="Inpainting steps", + elem_id=f"{id_prefix}_pp_inpainting_steps", + ) + + inpaiting_model = gr.Dropdown( + choices=["Current"] + sd_models.checkpoint_tiles(), + default="Current", + label="sd model (experimental)", + elem_id=f"{id_prefix}_pp_inpainting_sd_model", + ) + + gradio_components: List[gr.components.Component] = [ + inpainting_denoising_strength, + inpainting_denoising_prompt, + inpainting_denoising_negative_prompt, + inpainting_denoising_steps, + inpainting_sampler, + inpaiting_model, + ] + + return gradio_components diff --git a/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py b/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py index 62427b6..ec12a82 100644 --- a/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py +++ b/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py @@ -7,9 +7,9 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import Inpainting def postprocessing_ui() -> List[gr.components.Component]: - with gr.Tab(f"Post-Processing"): + with gr.Tab(f"Global Post-Processing"): gr.Markdown( - """Upscaling is performed on the whole image. Upscaling happens before face restoration.""" + """Upscaling is performed on the whole image and all faces (including not swapped). Upscaling happens before face restoration.""" ) with gr.Row(): face_restorer_name = gr.Radio( @@ -130,11 +130,11 @@ def postprocessing_ui() -> List[gr.components.Component]: upscaler_name, upscaler_scale, upscaler_visibility, + inpainting_when, inpainting_denoising_strength, inpainting_denoising_prompt, inpainting_denoising_negative_prompt, inpainting_denoising_steps, inpainting_sampler, - inpainting_when, inpaiting_model, ] diff --git a/scripts/faceswaplab_ui/faceswaplab_tab.py b/scripts/faceswaplab_ui/faceswaplab_tab.py index cea0935..622c48f 100644 --- a/scripts/faceswaplab_ui/faceswaplab_tab.py +++ b/scripts/faceswaplab_ui/faceswaplab_tab.py @@ -1,31 +1,32 @@ import os +import re +import traceback from pprint import pformat, pprint -from scripts.faceswaplab_utils import face_utils +from typing import * +from scripts.faceswaplab_utils.typing import * import gradio as gr import modules.scripts as scripts import onnx import pandas as pd -from scripts.faceswaplab_ui.faceswaplab_unit_ui import faceswap_unit_ui -from scripts.faceswaplab_ui.faceswaplab_postprocessing_ui import postprocessing_ui from modules import scripts -from PIL import Image from modules.shared import opts +from PIL import Image -from scripts.faceswaplab_utils import imgutils -from scripts.faceswaplab_utils.models_utils import get_models -from scripts.faceswaplab_utils.faceswaplab_logging import logger import scripts.faceswaplab_swapping.swapper as swapper +from scripts.faceswaplab_postprocessing.postprocessing import enhance_image from scripts.faceswaplab_postprocessing.postprocessing_options import ( PostProcessingOptions, ) -from scripts.faceswaplab_postprocessing.postprocessing import enhance_image -from dataclasses import fields -from typing import Any, Dict, List, Optional +from scripts.faceswaplab_ui.faceswaplab_postprocessing_ui import postprocessing_ui from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings -import re +from scripts.faceswaplab_ui.faceswaplab_unit_ui import faceswap_unit_ui +from scripts.faceswaplab_utils import face_utils, imgutils +from scripts.faceswaplab_utils.faceswaplab_logging import logger +from scripts.faceswaplab_utils.models_utils import get_models +from scripts.faceswaplab_utils.ui_utils import dataclasses_from_flat_list -def compare(img1: Image.Image, img2: Image.Image) -> str: +def compare(img1: PILImage, img2: PILImage) -> str: """ Compares the similarity between two faces extracted from images using cosine similarity. @@ -43,14 +44,15 @@ def compare(img1: Image.Image, img2: Image.Image) -> str: except Exception as e: logger.error("Fail to compare", e) + traceback.print_exc() return "You need 2 images to compare" def extract_faces( files: List[gr.File], extract_path: Optional[str], - *components: List[gr.components.Component], -) -> Optional[List[Image.Image]]: + *components: Tuple[gr.components.Component, ...], +) -> Optional[List[PILImage]]: """ Extracts faces from a list of image files. @@ -69,22 +71,32 @@ def extract_faces( If no faces are found, None is returned. """ - postprocess_options = PostProcessingOptions(*components) # type: ignore - images = [ - Image.open(file.name) for file in files - ] # potentially greedy but Image.open is supposed to be lazy - return swapper.extract_faces( - images, extract_path=extract_path, postprocess_options=postprocess_options - ) + if files and len(files) == 0: + logger.error("You need at least one image file to extract") + return [] + try: + postprocess_options = PostProcessingOptions(*components) # type: ignore + images = [ + Image.open(file.name) for file in files + ] # potentially greedy but Image.open is supposed to be lazy + result_images = swapper.extract_faces( + images, extract_path=extract_path, postprocess_options=postprocess_options + ) + return result_images + except Exception as e: + logger.error("Failed to extract : %s", e) + + traceback.print_exc() + return None -def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> Optional[str]: +def analyse_faces(image: PILImage, 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 + image : PIL.PILImage The input image where faces will be detected. The image must be a PIL Image object. det_threshold : float, optional @@ -122,6 +134,7 @@ def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> Optional[st except Exception as e: logger.error("Analysis Failed : %s", e) + traceback.print_exc() return None @@ -142,7 +155,7 @@ def sanitize_name(name: str) -> str: def build_face_checkpoint_and_save( batch_files: gr.File, name: str -) -> Optional[Image.Image]: +) -> 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 @@ -153,7 +166,7 @@ def build_face_checkpoint_and_save( name (str): The name assigned to the face checkpoint. Returns: - PIL.Image.Image or None: The resulting swapped face image if the process is successful; None otherwise. + PIL.PILImage or None: The resulting swapped face image if the process is successful; None otherwise. """ try: @@ -170,7 +183,7 @@ def build_face_checkpoint_and_save( os.makedirs(faces_path, exist_ok=True) - target_img = None + target_img: PILImage = None if blended_face: if blended_face["gender"] == 0: target_img = Image.open(os.path.join(preview_path, "woman.png")) @@ -180,15 +193,30 @@ def build_face_checkpoint_and_save( 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 - ), + target_face = swapper.get_or_default( + swapper.get_faces(imgutils.pil_to_cv2(target_img)), 0, None ) + if target_face is None: + logger.error( + "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( + reference_face=blended_face, + target_faces=[target_face], + source_face=blended_face, + target_img=target_img, + model=get_models()[0], + upscaled_swapper=opts.data.get( + "faceswaplab_upscaled_swapper", False + ), + ) + result_image = enhance_image( + result.image, + PostProcessingOptions( + face_restorer_name="CodeFormer", restorer_visibility=1 + ), + ) file_path = os.path.join(faces_path, f"{name}.safetensors") file_number = 1 @@ -202,14 +230,16 @@ def build_face_checkpoint_and_save( face_utils.save_face(filename=file_path, face=blended_face) try: data = face_utils.load_face(filename=file_path) - print(data) + logger.debug(data) except Exception as e: print(e) return result_image - print("No face found") + logger.error("No face found") except Exception as e: logger.error("Failed to build checkpoint %s", e) + + traceback.print_exc() return None return target_img @@ -242,36 +272,32 @@ def explore_onnx_faceswap_model(model_path: str) -> pd.DataFrame: df = pd.DataFrame(data) except Exception as e: - logger.info("Failed to explore model %s", e) + logger.error("Failed to explore model %s", e) + + traceback.print_exc() return None return df def batch_process( - files: List[gr.File], save_path: str, *components: List[gr.components.Component] -) -> Optional[List[Image.Image]]: + files: List[gr.File], save_path: str, *components: Tuple[Any, ...] +) -> Optional[List[PILImage]]: try: units_count = opts.data.get("faceswaplab_units_count", 3) - units: List[FaceSwapUnitSettings] = [] - # Parse and convert units flat components into FaceSwapUnitSettings - for i in range(0, units_count): - units += [FaceSwapUnitSettings.get_unit_configuration(i, components)] # type: ignore - - for i, u in enumerate(units): - logger.debug("%s, %s", pformat(i), pformat(u)) - - # Parse the postprocessing options - # We must first find where to start from (after face swapping units) - len_conf: int = len(fields(FaceSwapUnitSettings)) - shift: int = units_count * len_conf - postprocess_options = PostProcessingOptions( - *components[shift : shift + len(fields(PostProcessingOptions))] # type: ignore + classes: List[Any] = dataclasses_from_flat_list( + [FaceSwapUnitSettings] * units_count + [PostProcessingOptions], + components, ) - logger.debug("%s", pformat(postprocess_options)) + units: List[FaceSwapUnitSettings] = [ + u for u in classes if isinstance(u, FaceSwapUnitSettings) + ] + postprocess_options = classes[-1] + images = [ Image.open(file.name) for file in files ] # potentially greedy but Image.open is supposed to be lazy + return swapper.batch_process( images, save_path=save_path, @@ -280,7 +306,6 @@ def batch_process( ) except Exception as e: logger.error("Batch Process error : %s", e) - import traceback traceback.print_exc() return None diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py index b2fb3a3..d3ecf7c 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py +++ b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py @@ -1,15 +1,16 @@ from scripts.faceswaplab_swapping import swapper -import numpy as np import base64 import io -from dataclasses import dataclass, fields -from typing import Any, List, Optional, Set, Union +from dataclasses import dataclass +from typing import List, Optional, Set, Union import gradio as gr from insightface.app.common import Face from PIL import Image from scripts.faceswaplab_utils.imgutils import pil_to_cv2 from scripts.faceswaplab_utils.faceswaplab_logging import logger from scripts.faceswaplab_utils import face_utils +from scripts.faceswaplab_inpainting.faceswaplab_inpainting import InpaintingOptions +from client_api import api_utils @dataclass @@ -17,11 +18,11 @@ class FaceSwapUnitSettings: # ORDER of parameters is IMPORTANT. It should match the result of faceswap_unit_ui # The image given in reference - source_img: Union[Image.Image, str] + source_img: Optional[Union[Image.Image, str]] # The checkpoint file - source_face: str + source_face: Optional[str] # The batch source images - _batch_files: Union[gr.components.File, List[Image.Image]] + _batch_files: Optional[Union[gr.components.File, List[Image.Image]]] # Will blend faces if True blend_faces: bool # Enable this unit @@ -48,14 +49,39 @@ class FaceSwapUnitSettings: swap_in_source: bool # Swap in the generated image in img2img (always on for txt2img) swap_in_generated: bool + # Pre inpainting configuration (Don't use optional for this or gradio parsing will fail) : + pre_inpainting: InpaintingOptions + # Post inpainting configuration (Don't use optional for this or gradio parsing will fail) : + post_inpainting: InpaintingOptions @staticmethod - def get_unit_configuration( - unit: int, components: List[gr.components.Component] - ) -> Any: - fields_count = len(fields(FaceSwapUnitSettings)) + def from_api_dto(dto: api_utils.FaceSwapUnit) -> "FaceSwapUnitSettings": + """ + Converts a InpaintingOptions object from an API DTO (Data Transfer Object). + + :param options: An object of api_utils.InpaintingOptions representing the + post-processing options as received from the API. + :return: A InpaintingOptions instance containing the translated values + from the API DTO. + """ return FaceSwapUnitSettings( - *components[unit * fields_count : unit * fields_count + fields_count] + source_img=api_utils.base64_to_pil(dto.source_img), + source_face=dto.source_face, + _batch_files=dto.get_batch_images(), + blend_faces=dto.blend_faces, + enable=True, + same_gender=dto.same_gender, + sort_by_size=dto.sort_by_size, + check_similarity=dto.check_similarity, + _compute_similarity=dto.compute_similarity, + min_ref_sim=dto.min_ref_sim, + min_sim=dto.min_sim, + _faces_index=",".join([str(i) for i in (dto.faces_index)]), + reference_face_index=dto.reference_face_index, + swap_in_generated=True, + swap_in_source=False, + pre_inpainting=InpaintingOptions.from_api_dto(dto.pre_inpainting), + post_inpainting=InpaintingOptions.from_api_dto(dto.post_inpainting), ) @property @@ -156,24 +182,5 @@ class FaceSwapUnitSettings: """ if not hasattr(self, "_blended_faces"): self._blended_faces = swapper.blend_faces(self.faces) - assert ( - all( - [ - not np.array_equal( - self._blended_faces.embedding, face.embedding - ) - for face in self.faces - ] - ) - if len(self.faces) > 1 - else True - ), "Blended faces cannot be the same as one of the face if len(face)>0" - assert ( - not np.array_equal( - self._blended_faces.embedding, self.reference_face.embedding - ) - if len(self.faces) > 1 - else True - ), "Blended faces cannot be the same as reference face if len(face)>0" return self._blended_faces diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py b/scripts/faceswaplab_ui/faceswaplab_unit_ui.py index 9516ca1..8cac3db 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py +++ b/scripts/faceswaplab_ui/faceswaplab_unit_ui.py @@ -1,4 +1,5 @@ from typing import List +from scripts.faceswaplab_ui.faceswaplab_inpainting_ui import face_inpainting_ui from scripts.faceswaplab_utils.face_utils import get_face_checkpoints import gradio as gr @@ -142,22 +143,39 @@ def faceswap_unit_ui( visible=is_img2img, elem_id=f"{id_prefix}_face{unit_num}_swap_in_generated", ) + pre_inpainting = face_inpainting_ui( + name="Pre-Inpainting (Before swapping)", + id_prefix=f"{id_prefix}_face{unit_num}_preinpainting", + description="Pre-inpainting sends face to inpainting before swapping", + ) + post_inpainting = face_inpainting_ui( + name="Post-Inpainting (After swapping)", + id_prefix=f"{id_prefix}_face{unit_num}_postinpainting", + description="Post-inpainting sends face to inpainting after swapping", + ) + + gradio_components: List[gr.components.Component] = ( + [ + img, + face, + batch_files, + blend_faces, + enable, + same_gender, + sort_by_size, + check_similarity, + compute_similarity, + min_sim, + min_ref_sim, + target_faces_index, + reference_faces_index, + swap_in_source, + swap_in_generated, + ] + + pre_inpainting + + post_inpainting + ) + # If changed, you need to change FaceSwapUnitSettings accordingly # ORDER of parameters is IMPORTANT. It should match the result of FaceSwapUnitSettings - return [ - img, - face, - batch_files, - blend_faces, - enable, - same_gender, - sort_by_size, - check_similarity, - compute_similarity, - min_sim, - min_ref_sim, - target_faces_index, - reference_faces_index, - swap_in_source, - swap_in_generated, - ] + return gradio_components diff --git a/scripts/faceswaplab_utils/imgutils.py b/scripts/faceswaplab_utils/imgutils.py index f286bdc..7b0e534 100644 --- a/scripts/faceswaplab_utils/imgutils.py +++ b/scripts/faceswaplab_utils/imgutils.py @@ -1,5 +1,5 @@ import io -from typing import List, Optional, Tuple, Union, Dict +from typing import List, Optional, Union, Dict from PIL import Image import cv2 import numpy as np @@ -10,14 +10,15 @@ from scripts.faceswaplab_globals import NSFW_SCORE_THRESHOLD from modules import processing import base64 from collections import Counter +from scripts.faceswaplab_utils.typing import BoxCoords, CV2ImgU8, PILImage -def check_against_nsfw(img: Image.Image) -> bool: +def check_against_nsfw(img: PILImage) -> bool: """ Check if an image exceeds the Not Safe for Work (NSFW) score. Parameters: - img (PIL.Image.Image): The image to be checked. + img (PILImage): The image to be checked. Returns: bool: True if any part of the image is considered NSFW, False otherwise. @@ -32,33 +33,33 @@ def check_against_nsfw(img: Image.Image) -> bool: return any(shapes) -def pil_to_cv2(pil_img: Image.Image) -> np.ndarray: # type: ignore +def pil_to_cv2(pil_img: PILImage) -> CV2ImgU8: # type: ignore """ Convert a PIL Image into an OpenCV image (cv2). Args: - pil_img (PIL.Image.Image): An image in PIL format. + pil_img (PILImage): An image in PIL format. Returns: - np.ndarray: The input image converted to OpenCV format (BGR). + CV2ImgU8: The input image converted to OpenCV format (BGR). """ return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) -def cv2_to_pil(cv2_img: np.ndarray) -> Image.Image: # type: ignore +def cv2_to_pil(cv2_img: CV2ImgU8) -> PILImage: # type: ignore """ Convert an OpenCV image (cv2) into a PIL Image. Args: - cv2_img (np.ndarray): An image in OpenCV format (BGR). + cv2_img (CV2ImgU8): An image in OpenCV format (BGR). Returns: - PIL.Image.Image: The input image converted to PIL format (RGB). + PILImage: The input image converted to PIL format (RGB). """ return Image.fromarray(cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)) -def torch_to_pil(images: torch.Tensor) -> List[Image.Image]: +def torch_to_pil(tensor: torch.Tensor) -> List[PILImage]: """ Converts a tensor image or a batch of tensor images to a PIL image or a list of PIL images. @@ -72,7 +73,7 @@ def torch_to_pil(images: torch.Tensor) -> List[Image.Image]: list A list of PIL images. """ - images = images.cpu().permute(0, 2, 3, 1).numpy() + images: CV2ImgU8 = tensor.cpu().permute(0, 2, 3, 1).numpy() if images.ndim == 3: images = images[None, ...] images = (images * 255).round().astype("uint8") @@ -80,13 +81,13 @@ def torch_to_pil(images: torch.Tensor) -> List[Image.Image]: return pil_images -def pil_to_torch(pil_images: Union[Image.Image, List[Image.Image]]) -> torch.Tensor: +def pil_to_torch(pil_images: Union[PILImage, List[PILImage]]) -> torch.Tensor: """ Converts a PIL image or a list of PIL images to a torch tensor or a batch of torch tensors. Parameters ---------- - pil_images : Union[Image.Image, List[Image.Image]] + pil_images : Union[PILImage, List[PILImage]] A PIL image or a list of PIL images. Returns @@ -104,7 +105,7 @@ def pil_to_torch(pil_images: Union[Image.Image, List[Image.Image]]) -> torch.Ten return torch_image -def create_square_image(image_list: List[Image.Image]) -> Optional[Image.Image]: +def create_square_image(image_list: List[PILImage]) -> Optional[PILImage]: """ Creates a square image by combining multiple images in a grid pattern. @@ -156,7 +157,7 @@ def create_square_image(image_list: List[Image.Image]) -> Optional[Image.Image]: return None -# def create_mask(image : Image.Image, box_coords : Tuple[int, int, int, int]) -> Image.Image: +# def create_mask(image : PILImage, box_coords : Tuple[int, int, int, int]) -> PILImage: # width, height = image.size # mask = Image.new("L", (width, height), 255) # x1, y1, x2, y2 = box_coords @@ -170,19 +171,20 @@ def create_square_image(image_list: List[Image.Image]) -> Optional[Image.Image]: def create_mask( - image: Image.Image, box_coords: Tuple[int, int, int, int] -) -> Image.Image: + image: PILImage, + box_coords: BoxCoords, +) -> PILImage: """ Create a binary mask for a given image and bounding box coordinates. Args: - image (PIL.Image.Image): The input image. + image (PILImage): The input image. box_coords (Tuple[int, int, int, int]): A tuple of 4 integers defining the bounding box. It follows the pattern (x1, y1, x2, y2), where (x1, y1) is the top-left coordinate of the box and (x2, y2) is the bottom-right coordinate of the box. Returns: - PIL.Image.Image: A binary mask of the same size as the input image, where pixels within + PILImage: A binary mask of the same size as the input image, where pixels within the bounding box are white (255) and pixels outside the bounding box are black (0). """ width, height = image.size @@ -195,8 +197,8 @@ def create_mask( def apply_mask( - img: Image.Image, p: processing.StableDiffusionProcessing, batch_index: int -) -> Image.Image: + img: PILImage, p: processing.StableDiffusionProcessing, batch_index: int +) -> PILImage: """ Apply mask overlay and color correction to an image if enabled @@ -213,7 +215,7 @@ def apply_mask( overlays = p.overlay_images if overlays is None or batch_index >= len(overlays): return img - overlay: Image.Image = overlays[batch_index] + overlay: PILImage = overlays[batch_index] overlay = overlay.resize((img.size), resample=Image.Resampling.LANCZOS) img = img.copy() img.paste(overlay, (0, 0), overlay) @@ -227,9 +229,7 @@ def apply_mask( return img -def prepare_mask( - mask: Image.Image, p: processing.StableDiffusionProcessing -) -> Image.Image: +def prepare_mask(mask: PILImage, p: processing.StableDiffusionProcessing) -> PILImage: """ Prepare an image mask for the inpainting process. (This comes from controlnet) @@ -243,12 +243,12 @@ def prepare_mask( apply a Gaussian blur to the mask with a radius equal to 'mask_blur'. Args: - mask (Image.Image): The input mask as a PIL Image object. + mask (PILImage): The input mask as a PIL Image object. p (processing.StableDiffusionProcessing): An instance of the StableDiffusionProcessing class containing the processing parameters. Returns: - mask (Image.Image): The prepared mask as a PIL Image object. + mask (PILImage): The prepared mask as a PIL Image object. """ mask = mask.convert("L") # FIXME : Properly fix blur @@ -257,7 +257,7 @@ def prepare_mask( return mask -def base64_to_pil(base64str: Optional[str]) -> Optional[Image.Image]: +def base64_to_pil(base64str: Optional[str]) -> Optional[PILImage]: """ Converts a base64 string to a PIL Image object. @@ -267,7 +267,7 @@ def base64_to_pil(base64str: Optional[str]) -> Optional[Image.Image]: will return None. Returns: - Optional[Image.Image]: A PIL Image object created from the base64 string. If the input is None, + Optional[PILImage]: A PIL Image object created from the base64 string. If the input is None, the function returns None. Raises: diff --git a/scripts/faceswaplab_utils/typing.py b/scripts/faceswaplab_utils/typing.py new file mode 100644 index 0000000..d102d63 --- /dev/null +++ b/scripts/faceswaplab_utils/typing.py @@ -0,0 +1,10 @@ +from typing import Tuple +from numpy import uint8 +from numpy.typing import NDArray +from insightface.app.common import Face as IFace +from PIL import Image + +PILImage = Image.Image +CV2ImgU8 = NDArray[uint8] +Face = IFace +BoxCoords = Tuple[int, int, int, int] diff --git a/scripts/faceswaplab_utils/ui_utils.py b/scripts/faceswaplab_utils/ui_utils.py new file mode 100644 index 0000000..fc02120 --- /dev/null +++ b/scripts/faceswaplab_utils/ui_utils.py @@ -0,0 +1,39 @@ +from dataclasses import fields, is_dataclass +from typing import * + + +def dataclass_from_flat_list(cls: type, values: Tuple[Any, ...]) -> Any: + if not is_dataclass(cls): + raise TypeError(f"{cls} is not a dataclass") + + idx = 0 + init_values = {} + for field in fields(cls): + if is_dataclass(field.type): + inner_values = [values[idx + i] for i in range(len(fields(field.type)))] + init_values[field.name] = field.type(*inner_values) + idx += len(inner_values) + else: + value = values[idx] + init_values[field.name] = value + idx += 1 + return cls(**init_values) + + +def dataclasses_from_flat_list( + classes_mapping: List[type], values: Tuple[Any, ...] +) -> List[Any]: + instances = [] + idx = 0 + for cls in classes_mapping: + num_fields = sum( + len(fields(field.type)) if is_dataclass(field.type) else 1 + for field in fields(cls) + ) + instance = dataclass_from_flat_list(cls, values[idx : idx + num_fields]) + instances.append(instance) + idx += num_fields + assert [ + isinstance(i, t) for i, t in zip(instances, classes_mapping) + ], "Instances should match types" + return instances diff --git a/tests/test_api.py b/tests/test_api.py index abf7bfd..166ee24 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -17,6 +17,7 @@ from client_api.api_utils import ( FaceSwapExtractRequest, FaceSwapExtractResponse, compare_faces, + InpaintingOptions, ) from PIL import Image @@ -45,11 +46,12 @@ def face_swap_request() -> FaceSwapRequest: restorer_visibility=1, upscaler_name="Lanczos", scale=4, - inpainting_steps=30, - inpainting_denoising_strengh=0.1, inpainting_when=InpaintingWhen.BEFORE_RESTORE_FACE, + inpainting_options=InpaintingOptions( + inpainting_steps=30, + inpainting_denoising_strengh=0.1, + ), ) - # Prepare the request request = FaceSwapRequest( image=pil_to_base64("tests/test_image.png"), @@ -149,3 +151,31 @@ def test_faceswap(face_swap_request: FaceSwapRequest) -> None: assert response.status_code == 200 similarity = float(response.text) assert similarity > 0.50 + + +def test_faceswap_inpainting(face_swap_request: FaceSwapRequest) -> None: + face_swap_request.units[0].pre_inpainting = InpaintingOptions( + inpainting_denoising_strengh=0.4, + inpainting_prompt="Photo of a funny man", + inpainting_negative_prompt="blurry, bad art", + inpainting_steps=100, + ) + + face_swap_request.units[0].post_inpainting = InpaintingOptions( + inpainting_denoising_strengh=0.4, + inpainting_prompt="Photo of a funny man", + inpainting_negative_prompt="blurry, bad art", + inpainting_steps=20, + inpainting_sampler="Euler a", + ) + + response = requests.post( + f"{base_url}/faceswaplab/swap_face", + data=face_swap_request.json(), + headers={"Content-Type": "application/json; charset=utf-8"}, + ) + + assert response.status_code == 200 + data = response.json() + assert "images" in data + assert "infos" in data