diff --git a/README.md b/README.md index 3728c68..e97d8cf 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ While FaceSwapLab is still under development, it has reached a good level of sta In short: -+ **Ethical Guideline:** This extension should not be forked to create a public, easy way to circumvent NSFW filtering. ++ **Ethical Guideline:** This extension should not be forked to create a public, easy way to bypass NSFW filtering. If you modify it for this purpose, keep it private, or you'll be banned. + **License:** This software is distributed under the terms of the GNU Affero General Public License (AGPL), version 3 or later. + **Model License:** This software uses InsightFace's pre-trained models, which are available for non-commercial research purposes only. diff --git a/client_api/api_utils.py b/client_api/api_utils.py index 8499306..9190b0d 100644 --- a/client_api/api_utils.py +++ b/client_api/api_utils.py @@ -9,6 +9,7 @@ from io import BytesIO from typing import List, Tuple, Optional import numpy as np import requests +import safetensors class InpaintingWhen(Enum): @@ -49,6 +50,23 @@ class InpaintingOptions(BaseModel): ) +class InswappperOptions(BaseModel): + face_restorer_name: str = Field( + description="face restorer name", default="CodeFormer" + ) + restorer_visibility: float = Field( + description="face restorer visibility", default=1, le=1, ge=0 + ) + codeformer_weight: float = Field( + description="face restorer codeformer weight", default=1, le=1, ge=0 + ) + upscaler_name: str = Field(description="upscaler name", default=None) + improved_mask: bool = Field(description="Use Improved Mask", default=False) + color_corrections: bool = Field(description="Use Color Correction", default=False) + sharpen: bool = Field(description="Sharpen Image", default=False) + erosion_factor: float = Field(description="Erosion Factor", default=1, le=10, ge=0) + + class FaceSwapUnit(BaseModel): # The image given in reference source_img: str = Field( @@ -118,6 +136,11 @@ class FaceSwapUnit(BaseModel): default=None, ) + swapping_options: Optional[InswappperOptions] = Field( + description="PostProcessing & Mask options", + default=None, + ) + post_inpainting: Optional[InpaintingOptions] = Field( description="Inpainting options", default=None, @@ -244,3 +267,26 @@ def compare_faces( ) return float(result.text) + + +def safetensors_to_base64(file_path: str) -> str: + with open(file_path, "rb") as file: + file_bytes = file.read() + return "data:application/face;base64," + base64.b64encode(file_bytes).decode( + "utf-8" + ) + + +def base64_to_safetensors(base64str: str, output_path: str) -> None: + try: + base64_data = base64str.split("base64,")[-1] + file_bytes = base64.b64decode(base64_data) + with open(output_path, "wb") as file: + file.write(file_bytes) + with safetensors.safe_open(output_path, framework="pt") as f: + print(output_path, "keys =", f.keys()) + except Exception as e: + print("Error : failed to convert base64 string to safetensor", e) + import traceback + + traceback.print_exc() diff --git a/client_api/faceswaplab_api_example.py b/client_api/faceswaplab_api_example.py index 26b2975..b992f15 100644 --- a/client_api/faceswaplab_api_example.py +++ b/client_api/faceswaplab_api_example.py @@ -1,6 +1,7 @@ import requests from api_utils import ( FaceSwapUnit, + InswappperOptions, pil_to_base64, PostProcessingOptions, InpaintingWhen, @@ -10,6 +11,7 @@ from api_utils import ( FaceSwapExtractRequest, FaceSwapCompareRequest, FaceSwapExtractResponse, + safetensors_to_base64, ) address = "http://127.0.0.1:7860" @@ -94,3 +96,34 @@ response = FaceSwapExtractResponse.parse_obj(result.json()) for img in response.pil_images: img.show() + + +############################# +# FaceSwap with local safetensors + +# First face unit : +unit1 = FaceSwapUnit( + source_face=safetensors_to_base64("test.safetensors"), + faces_index=(0,), # Replace first face + swapping_options=InswappperOptions( + face_restorer_name="CodeFormer", + upscaler_name="LDSR", + improved_mask=True, + sharpen=True, + color_corrections=True, + ), +) + +# Prepare the request +request = FaceSwapRequest(image=pil_to_base64("test_image.png"), units=[unit1]) + +# Face Swap +result = requests.post( + url=f"{address}/faceswaplab/swap_face", + data=request.json(), + headers={"Content-Type": "application/json; charset=utf-8"}, +) +response = FaceSwapResponse.parse_obj(result.json()) + +for img in response.pil_images: + img.show() diff --git a/client_api/requirements.txt b/client_api/requirements.txt new file mode 100644 index 0000000..e5e672f --- /dev/null +++ b/client_api/requirements.txt @@ -0,0 +1,5 @@ +numpy==1.25.1 +Pillow==10.0.0 +pydantic==1.10.9 +Requests==2.31.0 +safetensors==0.3.1 diff --git a/client_api/test.safetensors b/client_api/test.safetensors new file mode 100644 index 0000000..c74264b Binary files /dev/null and b/client_api/test.safetensors differ diff --git a/docs/faq.markdown b/docs/faq.markdown index 115b9ad..3bfa1c3 100644 --- a/docs/faq.markdown +++ b/docs/faq.markdown @@ -133,3 +133,25 @@ The model generates faces with a resolution of 128x128, which is relatively low. SimSwap models are based on older InsightFace architectures, and SimSwap has not been released as a Python package. Its incorporation would complicate the process, and it does not guarantee any substantial gain. If you manage to implement SimSwap successfully, feel free to submit a pull request. + + +#### Shasum of inswapper model + +Check that your model is correct and not corrupted : + +```shell +$>sha1sum inswapper_128.onnx +17a64851eaefd55ea597ee41e5c18409754244c5 inswapper_128.onnx + +$>sha256sum inswapper_128.onnx +e4a3f08c753cb72d04e10aa0f7dbe3deebbf39567d4ead6dce08e98aa49e16af inswapper_128.onnx + +$>sha512sum inswapper_128.onnx +4311f4ccd9da58ec544e912b32ac0cba95f5ab4b1a06ac367efd3e157396efbae1097f624f10e77dd811fbba0917fa7c96e73de44563aa6099e5f46830965069 inswapper_128.onnx +``` + +#### Gradio errors (issubclass() arg 1 must be a class) + +Older versions of gradio don't work well with the extension. See this bug report : https://github.com/glucauze/sd-webui-faceswaplab/issues/5 + +It has been tested on 3.32.0 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 1c8637a..b4fb1e8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,10 @@ cython +dill ifnude insightface==0.7.3 onnx==1.14.0 onnxruntime==1.15.1 -opencv-python==4.7.0.72 +opencv-python pandas -pydantic==1.10.9 -dill==0.3.6 +pydantic safetensors \ No newline at end of file diff --git a/scripts/faceswaplab.py b/scripts/faceswaplab.py index 8c49032..82d66eb 100644 --- a/scripts/faceswaplab.py +++ b/scripts/faceswaplab.py @@ -72,14 +72,6 @@ class FaceSwapScript(scripts.Script): def units_count(self) -> int: return opts.data.get("faceswaplab_units_count", 3) - @property - def upscaled_swapper_in_generated(self) -> bool: - return opts.data.get("faceswaplab_upscaled_swapper", False) - - @property - def upscaled_swapper_in_source(self) -> bool: - return opts.data.get("faceswaplab_upscaled_swapper_in_source", False) - @property def enabled(self) -> bool: """Return True if any unit is enabled and the state is not interupted""" @@ -152,7 +144,6 @@ class FaceSwapScript(scripts.Script): get_current_model(), self.swap_in_source_units, images=init_images, - upscaled_swapper=self.upscaled_swapper_in_source, force_blend=True, ) logger.info(f"processed init images: {len(init_images)}") @@ -187,7 +178,6 @@ class FaceSwapScript(scripts.Script): get_current_model(), self.swap_in_generated_units, images=[(img, info)], - upscaled_swapper=self.upscaled_swapper_in_generated, ) if swapped_images is None: continue diff --git a/scripts/faceswaplab_api/faceswaplab_api.py b/scripts/faceswaplab_api/faceswaplab_api.py index de280d2..e4d9c3e 100644 --- a/scripts/faceswaplab_api/faceswaplab_api.py +++ b/scripts/faceswaplab_api/faceswaplab_api.py @@ -91,6 +91,8 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None: if src_image is not None: if request.postprocessing: pp_options = PostProcessingOptions.from_api_dto(request.postprocessing) + else: + pp_options = None units = get_faceswap_units_settings(request.units) swapped_images = swapper.batch_process( diff --git a/scripts/faceswaplab_settings/faceswaplab_settings.py b/scripts/faceswaplab_settings/faceswaplab_settings.py index 313816a..71c315c 100644 --- a/scripts/faceswaplab_settings/faceswaplab_settings.py +++ b/scripts/faceswaplab_settings/faceswaplab_settings.py @@ -54,7 +54,7 @@ def on_ui_settings() -> None: "faceswaplab_pp_default_face_restorer", shared.OptionInfo( None, - "UI Default post processing face restorer (requires restart)", + "UI Default global post processing face restorer (requires restart)", gr.Dropdown, { "interactive": True, @@ -67,7 +67,7 @@ def on_ui_settings() -> None: "faceswaplab_pp_default_face_restorer_visibility", shared.OptionInfo( 1, - "UI Default post processing face restorer visibility (requires restart)", + "UI Default global post processing face restorer visibility (requires restart)", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section, @@ -77,7 +77,7 @@ def on_ui_settings() -> None: "faceswaplab_pp_default_face_restorer_weight", shared.OptionInfo( 1, - "UI Default post processing face restorer weight (requires restart)", + "UI Default global post processing face restorer weight (requires restart)", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section, @@ -87,7 +87,7 @@ def on_ui_settings() -> None: "faceswaplab_pp_default_upscaler", shared.OptionInfo( None, - "UI Default post processing upscaler (requires restart)", + "UI Default global post processing upscaler (requires restart)", gr.Dropdown, { "interactive": True, @@ -100,13 +100,15 @@ def on_ui_settings() -> None: "faceswaplab_pp_default_upscaler_visibility", shared.OptionInfo( 1, - "UI Default post processing upscaler visibility(requires restart)", + "UI Default global post processing upscaler visibility(requires restart)", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section, ), ) + # Inpainting + shared.opts.add_option( "faceswaplab_pp_default_inpainting_prompt", shared.OptionInfo( @@ -132,20 +134,10 @@ def on_ui_settings() -> None: # UPSCALED SWAPPER shared.opts.add_option( - "faceswaplab_upscaled_swapper", - shared.OptionInfo( - False, - "Upscaled swapper. Applied only to the swapped faces. Apply transformations before merging with the original image.", - gr.Checkbox, - {"interactive": True}, - section=section, - ), - ) - shared.opts.add_option( - "faceswaplab_upscaled_swapper_upscaler", + "faceswaplab_default_upscaled_swapper_upscaler", shared.OptionInfo( None, - "Upscaled swapper upscaler (Recommanded : LDSR but slow)", + "Default Upscaled swapper upscaler (Recommanded : LDSR but slow) (requires restart)", gr.Dropdown, { "interactive": True, @@ -155,40 +147,40 @@ def on_ui_settings() -> None: ), ) shared.opts.add_option( - "faceswaplab_upscaled_swapper_sharpen", + "faceswaplab_default_upscaled_swapper_sharpen", shared.OptionInfo( False, - "Upscaled swapper sharpen", + "Default Upscaled swapper sharpen", gr.Checkbox, {"interactive": True}, section=section, ), ) shared.opts.add_option( - "faceswaplab_upscaled_swapper_fixcolor", + "faceswaplab_default_upscaled_swapper_fixcolor", shared.OptionInfo( False, - "Upscaled swapper color correction", + "Default Upscaled swapper color corrections (requires restart)", gr.Checkbox, {"interactive": True}, section=section, ), ) shared.opts.add_option( - "faceswaplab_upscaled_improved_mask", + "faceswaplab_default_upscaled_swapper_improved_mask", shared.OptionInfo( True, - "Use improved segmented mask (use pastenet to mask only the face)", + "Default Use improved segmented mask (use pastenet to mask only the face) (requires restart)", gr.Checkbox, {"interactive": True}, section=section, ), ) shared.opts.add_option( - "faceswaplab_upscaled_swapper_face_restorer", + "faceswaplab_default_upscaled_swapper_face_restorer", shared.OptionInfo( None, - "Upscaled swapper face restorer", + "Default Upscaled swapper face restorer (requires restart)", gr.Dropdown, { "interactive": True, @@ -198,40 +190,30 @@ def on_ui_settings() -> None: ), ) shared.opts.add_option( - "faceswaplab_upscaled_swapper_face_restorer_visibility", + "faceswaplab_default_upscaled_swapper_face_restorer_visibility", shared.OptionInfo( 1, - "Upscaled swapper face restorer visibility", + "Default Upscaled swapper face restorer visibility (requires restart)", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section, ), ) shared.opts.add_option( - "faceswaplab_upscaled_swapper_face_restorer_weight", + "faceswaplab_default_upscaled_swapper_face_restorer_weight", shared.OptionInfo( 1, - "Upscaled swapper face restorer weight (codeformer)", + "Default Upscaled swapper face restorer weight (codeformer) (requires restart)", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section, ), ) shared.opts.add_option( - "faceswaplab_upscaled_swapper_fthresh", - shared.OptionInfo( - 10, - "Upscaled swapper fthresh (diff sensitivity) 10 = default behaviour. Low impact.", - gr.Slider, - {"minimum": 5, "maximum": 250, "step": 1}, - section=section, - ), - ) - shared.opts.add_option( - "faceswaplab_upscaled_swapper_erosion", + "faceswaplab_default_upscaled_swapper_erosion", shared.OptionInfo( 1, - "Upscaled swapper mask erosion factor, 1 = default behaviour. The larger it is, the more blur is applied around the face. Too large and the facial change is no longer visible.", + "Default Upscaled swapper mask erosion factor, 1 = default behaviour. The larger it is, the more blur is applied around the face. Too large and the facial change is no longer visible. (requires restart)", gr.Slider, {"minimum": 0, "maximum": 10, "step": 0.001}, section=section, diff --git a/scripts/faceswaplab_swapping/swapper.py b/scripts/faceswaplab_swapping/swapper.py index 89eaf8e..e6fcae9 100644 --- a/scripts/faceswaplab_swapping/swapper.py +++ b/scripts/faceswaplab_swapping/swapper.py @@ -1,8 +1,14 @@ import copy import os from dataclasses import dataclass -from typing import Any, Dict, List, Set, Tuple, Optional +from pprint import pformat +from typing import Any, Dict, Generator, List, Set, Tuple, Optional import tempfile +from tqdm import tqdm +import sys +from io import StringIO +from contextlib import contextmanager +import hashlib import cv2 import insightface @@ -13,6 +19,7 @@ from PIL import Image from sklearn.metrics.pairwise import cosine_similarity from scripts.faceswaplab_swapping import upscaled_inswapper +from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions from scripts.faceswaplab_utils.imgutils import ( pil_to_cv2, check_against_nsfw, @@ -117,19 +124,16 @@ def batch_process( for src_image in src_images: current_images = [] swapped_images = process_images_units( - get_current_model(), - images=[(src_image, None)], - units=units, - upscaled_swapper=opts.data.get( - "faceswaplab_upscaled_swapper", False - ), + get_current_model(), images=[(src_image, None)], units=units ) if len(swapped_images) > 0: current_images += [img for img, _ in swapped_images] logger.info("%s images generated", len(current_images)) - for i, img in enumerate(current_images): - current_images[i] = enhance_image(img, postprocess_options) + + if postprocess_options: + for i, img in enumerate(current_images): + current_images[i] = enhance_image(img, postprocess_options) if save_path: for img in current_images: @@ -225,6 +229,33 @@ class FaceModelException(Exception): super().__init__(self.message) +@contextmanager +def capture_stdout() -> Generator[StringIO, None, None]: + """ + Capture and yield the printed messages to stdout. + + This context manager temporarily replaces sys.stdout with a StringIO object, + capturing all printed output. After the context block is exited, sys.stdout + is restored to its original value. + + Example usage: + with capture_stdout() as captured: + print("Hello, World!") + output = captured.getvalue() + # output now contains "Hello, World!\n" + + Returns: + A StringIO object containing the captured output. + """ + original_stdout = sys.stdout # Type: ignore + captured_stdout = StringIO() + sys.stdout = captured_stdout # Type: ignore + try: + yield captured_stdout + finally: + sys.stdout = original_stdout # Type: ignore + + @lru_cache(maxsize=1) def getAnalysisModel() -> insightface.app.FaceAnalysis: """ @@ -237,11 +268,20 @@ def getAnalysisModel() -> insightface.app.FaceAnalysis: if not os.path.exists(faceswaplab_globals.ANALYZER_DIR): os.makedirs(faceswaplab_globals.ANALYZER_DIR) - logger.info("Load analysis model, will take some time.") + logger.info("Load analysis model, will take some time. (> 30s)") # Initialize the analysis model with the specified name and providers - return insightface.app.FaceAnalysis( - name="buffalo_l", providers=providers, root=faceswaplab_globals.ANALYZER_DIR - ) + + with tqdm(total=1, desc="Loading analysis model", unit="model") as pbar: + with capture_stdout() as captured: + model = insightface.app.FaceAnalysis( + name="buffalo_l", + providers=providers, + root=faceswaplab_globals.ANALYZER_DIR, + ) + pbar.update(1) + logger.info("%s", pformat(captured.getvalue())) + + return model except Exception as e: logger.error( "Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)" @@ -249,11 +289,8 @@ def getAnalysisModel() -> insightface.app.FaceAnalysis: raise FaceModelException("Loading of analysis model failed") -import hashlib - - def is_sha1_matching(file_path: str, expected_sha1: str) -> bool: - sha1_hash = hashlib.sha1() + sha1_hash = hashlib.sha1(usedforsecurity=False) with open(file_path, "rb") as file: for byte_block in iter(lambda: file.read(4096), b""): @@ -284,10 +321,15 @@ def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper: expected_sha1, ) - # Initializes the face swap model using the specified model path. - return upscaled_inswapper.UpscaledINSwapper( - insightface.model_zoo.get_model(model_path, providers=providers) - ) + 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) + ) + pbar.update(1) + logger.info("%s", pformat(captured.getvalue())) + return model + except Exception as e: logger.error( "Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)" @@ -498,7 +540,7 @@ def swap_face( target_img: PILImage, target_faces: List[Face], model: str, - upscaled_swapper: bool = False, + swapping_options: Optional[InswappperOptions], compute_similarity: bool = True, ) -> ImageResult: """ @@ -532,7 +574,7 @@ def swap_face( img=result, target_face=swapped_face, source_face=source_face, - upscale=upscaled_swapper, + options=swapping_options, ) # type: ignore result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) @@ -581,7 +623,6 @@ def process_image_unit( unit: FaceSwapUnitSettings, image: PILImage, info: str = None, - upscaled_swapper: bool = False, force_blend: bool = False, ) -> List[Tuple[PILImage, str]]: """Process one image and return a List of (image, info) (one if blended, many if not). @@ -639,7 +680,7 @@ def process_image_unit( target_img=current_image, target_faces=target_faces, model=model, - upscaled_swapper=upscaled_swapper, + swapping_options=unit.swapping_options, compute_similarity=unit.compute_similarity, ) # Apply post-inpainting to image @@ -694,7 +735,6 @@ def process_images_units( model: str, units: List[FaceSwapUnitSettings], images: List[Tuple[Optional[PILImage], Optional[str]]], - upscaled_swapper: bool = False, force_blend: bool = False, ) -> Optional[List[Tuple[PILImage, str]]]: """ @@ -725,13 +765,9 @@ def process_images_units( processed_images = [] for i, (image, info) in enumerate(images): logger.debug("Processing image %s", i) - swapped = process_image_unit( - model, units[0], image, info, upscaled_swapper, force_blend - ) + swapped = process_image_unit(model, units[0], image, info, force_blend) logger.debug("Image %s -> %s images", i, len(swapped)) - nexts = process_images_units( - model, units[1:], swapped, upscaled_swapper, force_blend - ) + nexts = process_images_units(model, units[1:], swapped, force_blend) if nexts: processed_images.extend(nexts) else: diff --git a/scripts/faceswaplab_swapping/upcaled_inswapper_options.py b/scripts/faceswaplab_swapping/upcaled_inswapper_options.py new file mode 100644 index 0000000..efba4e4 --- /dev/null +++ b/scripts/faceswaplab_swapping/upcaled_inswapper_options.py @@ -0,0 +1,38 @@ +from dataclasses import * +from client_api import api_utils + + +@dataclass +class InswappperOptions: + face_restorer_name: str = None + restorer_visibility: float = 1 + codeformer_weight: float = 1 + upscaler_name: 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": + """ + 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 InswappperOptions() + + return InswappperOptions( + face_restorer_name=dto.face_restorer_name, + restorer_visibility=dto.restorer_visibility, + codeformer_weight=dto.codeformer_weight, + upscaler_name=dto.upscaler_name, + improved_mask=dto.improved_mask, + color_corrections=dto.color_corrections, + sharpen=dto.sharpen, + erosion_factor=dto.erosion_factor, + ) diff --git a/scripts/faceswaplab_swapping/upscaled_inswapper.py b/scripts/faceswaplab_swapping/upscaled_inswapper.py index a441dc0..dd66d82 100644 --- a/scripts/faceswaplab_swapping/upscaled_inswapper.py +++ b/scripts/faceswaplab_swapping/upscaled_inswapper.py @@ -12,8 +12,10 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import ( PostProcessingOptions, ) 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.typing import CV2ImgU8, Face +from scripts.faceswaplab_utils.faceswaplab_logging import logger def get_upscaler() -> UpscalerData: @@ -127,26 +129,25 @@ class UpscaledINSwapper(INSwapper): def __init__(self, inswapper: INSwapper): self.__dict__.update(inswapper.__dict__) - def super_resolution(self, img: CV2ImgU8, k: int = 2) -> CV2ImgU8: + def upscale_and_restore( + self, img: CV2ImgU8, k: int = 2, inswapper_options: InswappperOptions = None + ) -> CV2ImgU8: pil_img = cv2_to_pil(img) - options = PostProcessingOptions( - upscaler_name=opts.data.get( - "faceswaplab_upscaled_swapper_upscaler", "LDSR" - ), + pp_options = PostProcessingOptions( + upscaler_name=inswapper_options.upscaler_name, upscale_visibility=1, scale=k, - face_restorer_name=opts.data.get( - "faceswaplab_upscaled_swapper_face_restorer", "" - ), - codeformer_weight=opts.data.get( - "faceswaplab_upscaled_swapper_face_restorer_weight", 1 - ), - restorer_visibility=opts.data.get( - "faceswaplab_upscaled_swapper_face_restorer_visibility", 1 - ), + face_restorer_name=inswapper_options.face_restorer_name, + codeformer_weight=inswapper_options.codeformer_weight, + restorer_visibility=inswapper_options.restorer_visibility, ) - upscaled = upscaling.upscale_img(pil_img, options) - upscaled = upscaling.restore_face(upscaled, options) + + upscaled = pil_img + if pp_options.upscaler_name: + upscaled = upscaling.upscale_img(pil_img, pp_options) + if pp_options.face_restorer_name: + upscaled = upscaling.restore_face(upscaled, pp_options) + return pil_to_cv2(upscaled) def get( @@ -155,7 +156,7 @@ class UpscaledINSwapper(INSwapper): target_face: Face, source_face: Face, paste_back: bool = True, - upscale: bool = True, + options: 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( @@ -190,43 +191,48 @@ class UpscaledINSwapper(INSwapper): fake_diff[:, -2:] = 0 return fake_diff - if upscale: - print("*" * 80) - print( - f"Upscaled inswapper using {opts.data.get('faceswaplab_upscaled_swapper_upscaler', 'LDSR')}" - ) - print("*" * 80) + if options: + logger.info("*" * 80) + logger.info(f"Upscaled inswapper") - k = 4 - aimg, M = face_align.norm_crop2( - img, target_face.kps, self.input_size[0] * k - ) + if options.upscaler_name: + # Upscale original image + k = 4 + aimg, M = face_align.norm_crop2( + img, target_face.kps, self.input_size[0] * k + ) + else: + k = 1 # upscale and restore face : - bgr_fake = self.super_resolution(bgr_fake, k) + bgr_fake = self.upscale_and_restore( + bgr_fake, inswapper_options=options, k=k + ) - if opts.data.get("faceswaplab_upscaled_improved_mask", True): + if options.improved_mask: mask = get_face_mask(aimg, bgr_fake) bgr_fake = merge_images_with_mask(aimg, bgr_fake, mask) # compute fake_diff before sharpen and color correction (better result) fake_diff = compute_diff(bgr_fake, aimg) - if opts.data.get("faceswaplab_upscaled_swapper_sharpen", True): - print("sharpen") + if options.sharpen: + logger.info("sharpen") # Add sharpness blurred = cv2.GaussianBlur(bgr_fake, (0, 0), 3) bgr_fake = cv2.addWeighted(bgr_fake, 1.5, blurred, -0.5, 0) # Apply color corrections - if opts.data.get("faceswaplab_upscaled_swapper_fixcolor", True): - print("color correction") + if options.color_corrections: + logger.info("color correction") correction = processing.setup_color_correction(cv2_to_pil(aimg)) bgr_fake_pil = processing.apply_color_correction( correction, cv2_to_pil(bgr_fake) ) bgr_fake = pil_to_cv2(bgr_fake_pil) + logger.info("*" * 80) + else: fake_diff = compute_diff(bgr_fake, aimg) @@ -254,7 +260,7 @@ class UpscaledINSwapper(INSwapper): borderValue=0.0, ) img_white[img_white > 20] = 255 - fthresh = opts.data.get("faceswaplab_upscaled_swapper_fthresh", 10) + fthresh = 10 print("fthresh", fthresh) fake_diff[fake_diff < fthresh] = 0 fake_diff[fake_diff >= fthresh] = 255 @@ -263,9 +269,8 @@ class UpscaledINSwapper(INSwapper): mask_h = np.max(mask_h_inds) - np.min(mask_h_inds) mask_w = np.max(mask_w_inds) - np.min(mask_w_inds) mask_size = int(np.sqrt(mask_h * mask_w)) - erosion_factor = opts.data.get( - "faceswaplab_upscaled_swapper_erosion", 1 - ) + erosion_factor = options.erosion_factor + k = max(int(mask_size // 10 * erosion_factor), int(10 * erosion_factor)) kernel = np.ones((k, k), np.uint8) diff --git a/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py b/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py index ec12a82..6207b5b 100644 --- a/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py +++ b/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py @@ -17,7 +17,7 @@ def postprocessing_ui() -> List[gr.components.Component]: choices=["None"] + [x.name() for x in shared.face_restorers], value=lambda: opts.data.get( "faceswaplab_pp_default_face_restorer", - shared.face_restorers[0].name(), + "None", ), type="value", elem_id="faceswaplab_pp_face_restorer", diff --git a/scripts/faceswaplab_ui/faceswaplab_tab.py b/scripts/faceswaplab_ui/faceswaplab_tab.py index 0d66cec..5db3dd7 100644 --- a/scripts/faceswaplab_ui/faceswaplab_tab.py +++ b/scripts/faceswaplab_ui/faceswaplab_tab.py @@ -249,6 +249,8 @@ def tools_ui() -> None: preview = gr.components.Image( type="pil", label="Preview", + width=512, + height=512, interactive=False, elem_id="faceswaplab_build_preview_face", ) diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py index 9bfe4b0..84db6d3 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py +++ b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py @@ -6,6 +6,7 @@ from typing import List, Optional, Set, Union import gradio as gr from insightface.app.common import Face from PIL import Image +from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions from scripts.faceswaplab_utils.imgutils import pil_to_cv2 from scripts.faceswaplab_utils.faceswaplab_logging import logger from scripts.faceswaplab_utils import face_checkpoints_utils @@ -51,6 +52,8 @@ class FaceSwapUnitSettings: swap_in_generated: bool # Pre inpainting configuration (Don't use optional for this or gradio parsing will fail) : pre_inpainting: InpaintingOptions + # Configure swapping options + swapping_options: InswappperOptions # Post inpainting configuration (Don't use optional for this or gradio parsing will fail) : post_inpainting: InpaintingOptions @@ -81,6 +84,7 @@ class FaceSwapUnitSettings: swap_in_generated=True, swap_in_source=False, pre_inpainting=InpaintingOptions.from_api_dto(dto.pre_inpainting), + swapping_options=InswappperOptions.from_api_dto(dto.swapping_options), post_inpainting=InpaintingOptions.from_api_dto(dto.post_inpainting), ) diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py b/scripts/faceswaplab_ui/faceswaplab_unit_ui.py index a035c14..e75b2cf 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py +++ b/scripts/faceswaplab_ui/faceswaplab_unit_ui.py @@ -2,6 +2,100 @@ 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 import gradio as gr +from modules.shared import opts +from modules import shared + + +def faceswap_unit_advanced_options( + is_img2img: bool, unit_num: int = 1, id_prefix: str = "faceswaplab_" +) -> List[gr.components.Component]: + with gr.Accordion(f"Post-Processing & Advanced Mask Options", open=False): + gr.Markdown("""Post-processing and mask settings for unit faces""") + with gr.Row(): + face_restorer_name = gr.Radio( + label="Restore Face", + choices=["None"] + [x.name() for x in shared.face_restorers], + value=lambda: opts.data.get( + "faceswaplab_default_upscaled_swapper_face_restorer", + "None", + ), + type="value", + elem_id=f"{id_prefix}_face{unit_num}_face_restorer", + ) + with gr.Column(): + face_restorer_visibility = gr.Slider( + 0, + 1, + value=lambda: opts.data.get( + "faceswaplab_default_upscaled_swapper_face_restorer_visibility", + 1.0, + ), + step=0.001, + label="Restore visibility", + elem_id=f"{id_prefix}_face{unit_num}_face_restorer_visibility", + ) + codeformer_weight = gr.Slider( + 0, + 1, + value=lambda: opts.data.get( + "faceswaplab_default_upscaled_swapper_face_restorer_weight", 1.0 + ), + step=0.001, + label="codeformer weight", + elem_id=f"{id_prefix}_face{unit_num}_face_restorer_weight", + ) + upscaler_name = gr.Dropdown( + choices=[upscaler.name for upscaler in shared.sd_upscalers], + value=lambda: opts.data.get( + "faceswaplab_default_upscaled_swapper_upscaler", "" + ), + label="Upscaler", + elem_id=f"{id_prefix}_face{unit_num}_upscaler", + ) + + improved_mask = gr.Checkbox( + lambda: opts.data.get( + "faceswaplab_default_upscaled_swapper_improved_mask", False + ), + interactive=True, + label="Use improved segmented mask (use pastenet to mask only the face)", + elem_id=f"{id_prefix}_face{unit_num}_improved_mask", + ) + color_corrections = gr.Checkbox( + lambda: opts.data.get( + "faceswaplab_default_upscaled_swapper_fixcolor", False + ), + interactive=True, + label="Use color corrections", + elem_id=f"{id_prefix}_face{unit_num}_color_corrections", + ) + sharpen_face = gr.Checkbox( + lambda: opts.data.get( + "faceswaplab_default_upscaled_swapper_sharpen", False + ), + interactive=True, + label="sharpen face", + elem_id=f"{id_prefix}_face{unit_num}_sharpen_face", + ) + erosion_factor = gr.Slider( + 0.0, + 10.0, + lambda: opts.data.get("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", + ) + + return [ + face_restorer_name, + face_restorer_visibility, + codeformer_weight, + upscaler_name, + improved_mask, + color_corrections, + sharpen_face, + erosion_factor, + ] def faceswap_unit_ui( @@ -62,35 +156,6 @@ def faceswap_unit_ui( elem_id=f"{id_prefix}_face{unit_num}_blend_faces", interactive=True, ) - gr.Markdown("""Discard images with low similarity or no faces :""") - with gr.Row(): - check_similarity = gr.Checkbox( - False, - placeholder="discard", - label="Check similarity", - elem_id=f"{id_prefix}_face{unit_num}_check_similarity", - ) - compute_similarity = gr.Checkbox( - False, - label="Compute similarity", - elem_id=f"{id_prefix}_face{unit_num}_compute_similarity", - ) - min_sim = gr.Slider( - 0, - 1, - 0, - step=0.01, - label="Min similarity", - elem_id=f"{id_prefix}_face{unit_num}_min_similarity", - ) - min_ref_sim = gr.Slider( - 0, - 1, - 0, - step=0.01, - label="Min reference similarity", - elem_id=f"{id_prefix}_face{unit_num}_min_ref_similarity", - ) gr.Markdown( """Select the face to be swapped, you can sort by size or use the same gender as the desired face:""" @@ -143,11 +208,46 @@ def faceswap_unit_ui( visible=is_img2img, elem_id=f"{id_prefix}_face{unit_num}_swap_in_generated", ) + + with gr.Accordion("Similarity", open=False): + gr.Markdown("""Discard images with low similarity or no faces :""") + with gr.Row(): + check_similarity = gr.Checkbox( + False, + placeholder="discard", + label="Check similarity", + elem_id=f"{id_prefix}_face{unit_num}_check_similarity", + ) + compute_similarity = gr.Checkbox( + False, + label="Compute similarity", + elem_id=f"{id_prefix}_face{unit_num}_compute_similarity", + ) + min_sim = gr.Slider( + 0, + 1, + 0, + step=0.01, + label="Min similarity", + elem_id=f"{id_prefix}_face{unit_num}_min_similarity", + ) + min_ref_sim = gr.Slider( + 0, + 1, + 0, + step=0.01, + label="Min reference similarity", + elem_id=f"{id_prefix}_face{unit_num}_min_ref_similarity", + ) + 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", ) + + options = faceswap_unit_advanced_options(is_img2img, unit_num, id_prefix) + post_inpainting = face_inpainting_ui( name="Post-Inpainting (After swapping)", id_prefix=f"{id_prefix}_face{unit_num}_postinpainting", @@ -173,6 +273,7 @@ def faceswap_unit_ui( swap_in_generated, ] + pre_inpainting + + options + post_inpainting ) diff --git a/scripts/faceswaplab_utils/face_checkpoints_utils.py b/scripts/faceswaplab_utils/face_checkpoints_utils.py index de3ebf1..481718f 100644 --- a/scripts/faceswaplab_utils/face_checkpoints_utils.py +++ b/scripts/faceswaplab_utils/face_checkpoints_utils.py @@ -7,21 +7,19 @@ import torch import modules.scripts as scripts from modules import scripts +from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions from scripts.faceswaplab_utils.faceswaplab_logging import logger from scripts.faceswaplab_utils.typing import * from scripts.faceswaplab_utils import imgutils -from scripts.faceswaplab_postprocessing.postprocessing import enhance_image -from scripts.faceswaplab_postprocessing.postprocessing_options import ( - PostProcessingOptions, -) from scripts.faceswaplab_utils.models_utils import get_models -from modules.shared import opts import traceback import dill as pickle # will be removed in future versions from scripts.faceswaplab_swapping import swapper from pprint import pformat import re +from client_api import api_utils +import tempfile def sanitize_name(name: str) -> str: @@ -93,16 +91,9 @@ def build_face_checkpoint_and_save( source_face=blended_face, target_img=reference_preview_img, model=get_models()[0], - upscaled_swapper=opts.data.get( - "faceswaplab_upscaled_swapper", False - ), - ) - preview_image = enhance_image( - result.image, - PostProcessingOptions( - face_restorer_name="CodeFormer", restorer_visibility=1 - ), + swapping_options=InswappperOptions(face_restorer_name="Codeformer"), ) + preview_image = result.image file_path = os.path.join(get_checkpoint_path(), f"{name}.safetensors") if not overwrite: @@ -147,6 +138,16 @@ def save_face(face: Face, filename: str) -> None: def load_face(name: str) -> Face: + if name.startswith("data:application/face;base64,"): + with tempfile.NamedTemporaryFile(delete=True) as temp_file: + api_utils.base64_to_safetensors(name, temp_file.name) + face = {} + with safe_open(temp_file.name, framework="pt", device="cpu") as f: + for k in f.keys(): + logger.debug("load key %s", k) + face[k] = f.get_tensor(k).numpy() + return Face(face) + filename = matching_checkpoint(name) if filename is None: return None