From 9c935442ff8b7490c28db164166adcc68e79a342 Mon Sep 17 00:00:00 2001 From: Tran Xen <137925069+glucauze@users.noreply.github.com> Date: Sat, 5 Aug 2023 01:47:35 +0200 Subject: [PATCH] update install (risky but cleaner). Add auto_det_size param and try to emulate old behavior --- client_api/requirements.txt | 10 +-- docs/documentation.markdown | 14 ++- install.py | 35 +++++--- requirements-gpu.txt | 6 +- requirements.txt | 4 +- scripts/faceswaplab_globals.py | 2 +- .../faceswaplab_settings.py | 11 +++ scripts/faceswaplab_swapping/swapper.py | 86 ++++++++++++++----- .../faceswaplab_postprocessing_ui.py | 2 +- scripts/faceswaplab_ui/faceswaplab_unit_ui.py | 14 ++- 10 files changed, 136 insertions(+), 48 deletions(-) diff --git a/client_api/requirements.txt b/client_api/requirements.txt index e5e672f..4e2edfc 100644 --- a/client_api/requirements.txt +++ b/client_api/requirements.txt @@ -1,5 +1,5 @@ -numpy==1.25.1 -Pillow==10.0.0 -pydantic==1.10.9 -Requests==2.31.0 -safetensors==0.3.1 +numpy +Pillow +pydantic +Requests +safetensors>=0.3.1 diff --git a/docs/documentation.markdown b/docs/documentation.markdown index eb34431..fbc69ce 100644 --- a/docs/documentation.markdown +++ b/docs/documentation.markdown @@ -5,6 +5,18 @@ permalink: /doc/ toc: true --- +## TLDR: I Just Want Good Results: + +1. Put a face in the reference. +2. Select a face number. +3. Select "Enable." +4. Select "CodeFormer" in global Post-Processing. + +Once you're happy with some results but want to improve, the next steps are to: + ++ Use advanced settings in face units (which are not as complex as they might seem, it's basically fine tuning post-processing for each faces). ++ Use pre/post inpainting to tweak the image a bit for more natural results. + ## Main Interface Here is the interface for FaceSwap Lab. It is available in the form of an accordion in both img2img and txt2img. @@ -60,7 +72,7 @@ The purpose of this feature is to enhance the quality of the face in the final i The upscaled inswapper is disabled by default. It can be enabled in the sd options. Understanding the various steps helps explain why results may be unsatisfactory and how to address this issue. -+ **upscaler** : LDSR if None. The LDSR option generally gives the best results but at the expense of a lot of computational time. You should test other models to form an opinion. The 003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN model seems to give good results in a reasonable amount of time. It's not possible to disable upscaling, but it is possible to choose LANCZOS for speed if Codeformer is enabled in the upscaled inswapper. The result is generally satisfactory. ++ **upscaler** : LDSR if None. The LDSR option generally gives the best results but at the expense of a lot of computational time. You should test other models to form an opinion. The [003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN](https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth) model seems to give good results in a reasonable amount of time. It's not possible to disable upscaling, but it is possible to choose LANCZOS for speed if Codeformer is enabled in the upscaled inswapper. The result is generally satisfactory. You can check [here for an upscaler database](https://upscale.wiki/wiki/Model_Database) and [here for some comparison](https://phhofm.github.io/upscale/favorites.html). It is a test and try process. + **restorer** : The face restorer to be used if necessary. Codeformer generally gives good results. + **sharpening** can provide more natural results, but it may also add artifacts. The same goes for **color correction**. By default, these options are set to False. + **improved mask:** The segmentation mask for the upscaled swapper is designed to avoid the square mask and prevent degradation of the non-face parts of the image. It is based on the Codeformer implementation. If "Use improved segmented mask (use pastenet to mask only the face)" and "upscaled inswapper" are checked in the settings, the mask will only cover the face, and will not be squared. However, depending on the image, this might introduce different types of problems such as artifacts on the border of the face. diff --git a/install.py b/install.py index 172da70..169a434 100644 --- a/install.py +++ b/install.py @@ -1,8 +1,9 @@ import launch import os -import pkg_resources import sys +import pkg_resources from modules import shared +from packaging.version import parse use_gpu = getattr(shared.cmd_opts, "faceswaplab_gpu", False) @@ -15,29 +16,35 @@ else: os.path.dirname(os.path.realpath(__file__)), "requirements.txt" ) + +def is_installed(package: str) -> bool: + package_name = package.split("==")[0].split(">=")[0].strip() + try: + installed_version = parse(pkg_resources.get_distribution(package_name).version) + except pkg_resources.DistributionNotFound: + return False + + if "==" in package: + required_version = parse(package.split("==")[1]) + return installed_version == required_version + elif ">=" in package: + required_version = parse(package.split(">=")[1]) + return installed_version >= required_version + else: + return True + + print("Checking faceswaplab requirements") with open(req_file) as file: for package in file: try: - python = sys.executable package = package.strip() - if not launch.is_installed(package.split("==")[0]): + if not is_installed(package): print(f"Install {package}") launch.run_pip( f"install {package}", f"sd-webui-faceswaplab requirement: {package}" ) - elif "==" in package: - package_name, package_version = package.split("==") - installed_version = pkg_resources.get_distribution(package_name).version - if installed_version != package_version: - print( - f"Install {package}, {installed_version} vs {package_version}" - ) - launch.run_pip( - f"install {package}", - f"sd-webui-faceswaplab requirement: changing {package_name} version from {installed_version} to {package_version}", - ) except Exception as e: print(e) diff --git a/requirements-gpu.txt b/requirements-gpu.txt index 83b9209..645b201 100644 --- a/requirements-gpu.txt +++ b/requirements-gpu.txt @@ -2,10 +2,10 @@ cython dill ifnude insightface==0.7.3 -onnx==1.14.0 +onnx>=1.14.0 opencv-python pandas pydantic safetensors -onnxruntime==1.15.1 -onnxruntime-gpu==1.15.1 \ No newline at end of file +onnxruntime>=1.15.0 +onnxruntime-gpu>=1.15.0 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index b4fb1e8..f6581d5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,8 +2,8 @@ cython dill ifnude insightface==0.7.3 -onnx==1.14.0 -onnxruntime==1.15.1 +onnx>=1.14.0 +onnxruntime>=1.15.0 opencv-python pandas pydantic diff --git a/scripts/faceswaplab_globals.py b/scripts/faceswaplab_globals.py index 4b43407..7fa0b08 100644 --- a/scripts/faceswaplab_globals.py +++ b/scripts/faceswaplab_globals.py @@ -10,7 +10,7 @@ REFERENCE_PATH = os.path.join( scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references" ) -VERSION_FLAG: str = "v1.2.0" +VERSION_FLAG: str = "v1.2.1" 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_settings/faceswaplab_settings.py b/scripts/faceswaplab_settings/faceswaplab_settings.py index d826061..7d715ed 100644 --- a/scripts/faceswaplab_settings/faceswaplab_settings.py +++ b/scripts/faceswaplab_settings/faceswaplab_settings.py @@ -48,6 +48,17 @@ def on_ui_settings() -> None: ), ) + shared.opts.add_option( + "faceswaplab_auto_det_size", + shared.OptionInfo( + True, + "Auto det_size : Will load model twice and test faces on each if needed (old behaviour). Takes more VRAM. Precedence over fixed det_size", + gr.Checkbox, + {"interactive": True}, + section=section, + ), + ) + shared.opts.add_option( "faceswaplab_detection_threshold", shared.OptionInfo( diff --git a/scripts/faceswaplab_swapping/swapper.py b/scripts/faceswaplab_swapping/swapper.py index 75e951f..0b5a980 100644 --- a/scripts/faceswaplab_swapping/swapper.py +++ b/scripts/faceswaplab_swapping/swapper.py @@ -38,18 +38,22 @@ from scripts.faceswaplab_utils.models_utils import get_current_model from scripts.faceswaplab_utils.typing import CV2ImgU8, PILImage, Face from scripts.faceswaplab_inpainting.i2i_pp import img2img_diffusion from modules import shared +import onnxruntime +USE_GPU = ( + getattr(shared.cmd_opts, "faceswaplab_gpu", False) and sys.platform != "darwin" +) -USE_GPU = getattr(shared.cmd_opts, "faceswaplab_gpu", False) - +providers = ["CPUExecutionProvider"] if USE_GPU and sys.platform != "darwin": - providers = [ - "TensorrtExecutionProvider", - "CUDAExecutionProvider", - "CPUExecutionProvider", - ] -else: - providers = ["CPUExecutionProvider"] + if "CUDAExecutionProvider" in onnxruntime.get_available_providers(): + providers = ["CUDAExecutionProvider"] + else: + logger.error( + "CUDAExecutionProvider not found in onnxruntime.available_providers : %s, use CPU instead. Check onnxruntime-gpu is installed.", + onnxruntime.get_available_providers(), + ) + USE_GPU = False def cosine_similarity_face(face1: Face, face2: Face) -> float: @@ -268,7 +272,21 @@ def capture_stdout() -> Generator[StringIO, None, None]: sys.stdout = original_stdout # Type: ignore +# On GPU we can keep a non prepared model in ram and deepcopy it every time det_size change (old behaviour) @lru_cache(maxsize=1) +def get_cpu_analysis() -> insightface.app.FaceAnalysis: + return insightface.app.FaceAnalysis( + name="buffalo_l", + providers=providers, + root=faceswaplab_globals.ANALYZER_DIR, + ) + + +# FIXME : This function is way more complicated than it could be. +# It is done that way to preserve the original behavior with CPU. +# Most users don't reed the doc, so we need to keep the features as close as possible +# to original behavior. +@lru_cache(maxsize=3) def getAnalysisModel( det_size: Tuple[int, int] = (640, 640), det_thresh: float = 0.5 ) -> insightface.app.FaceAnalysis: @@ -291,14 +309,20 @@ def getAnalysisModel( total=1, desc="Loading analysis model (first time is slow)", unit="model" ) as pbar: with capture_stdout() as captured: - model = insightface.app.FaceAnalysis( - name="buffalo_l", - providers=providers, - root=faceswaplab_globals.ANALYZER_DIR, - ) + if USE_GPU: + model = insightface.app.FaceAnalysis( + name="buffalo_l", + providers=providers, + root=faceswaplab_globals.ANALYZER_DIR, + ) + + # Prepare the analysis model for face detection with the specified detection size + model.prepare(ctx_id=0, det_thresh=det_thresh, det_size=det_size) + else: + # This is a hacky way to speed up loading for gpu only + model = copy.deepcopy(get_cpu_analysis()) + model.prepare(ctx_id=0, det_thresh=det_thresh, det_size=det_size) - # Prepare the analysis model for face detection with the specified detection size - model.prepare(ctx_id=0, det_thresh=det_thresh, det_size=det_size) pbar.update(1) logger.info("%s", pformat(captured.getvalue())) @@ -369,6 +393,7 @@ def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper: def get_faces( img_data: CV2ImgU8, det_thresh: Optional[float] = None, + det_size: Tuple[int, int] = (640, 640), ) -> List[Face]: """ Detects and retrieves faces from an image using an analysis model. @@ -385,15 +410,36 @@ def get_faces( if det_thresh is None: det_thresh = opts.data.get("faceswaplab_detection_threshold", 0.5) - det_size = opts.data.get("faceswaplab_det_size", 640) - face_analyser = getAnalysisModel((det_size, det_size), det_thresh) + auto_det_size = opts.data.get("faceswaplab_auto_det_size", True) + if not auto_det_size: + x = opts.data.get("faceswaplab_det_size", 640) + det_size = (x, x) + + face_analyser = getAnalysisModel(det_size, det_thresh) # Get the detected faces from the image using the analysis model - face = face_analyser.get(img_data) + faces = face_analyser.get(img_data) + + # If no faces are detected and the detection size is larger than 320x320, + # recursively call the function with a smaller detection size + if len(faces) == 0: + if auto_det_size: + if det_size[0] > 320 and det_size[1] > 320: + det_size_half = (det_size[0] // 2, det_size[1] // 2) + return get_faces( + img_data, det_size=det_size_half, det_thresh=det_thresh + ) + + # If no faces are detected print a warning to user about change in detection + else: + if det_size[0] > 320: + logger.warning( + "No faces detected, you might want to play with det_size by reducing it (in sd global settings). Lower (320) means more detection but less precise. Or activate auto-det-size." + ) try: # Sort the detected faces based on their x-coordinate of the bounding box - return sorted(face, key=lambda x: x.bbox[0]) + return sorted(faces, key=lambda x: x.bbox[0]) except Exception as e: logger.error("Failed to get faces %s", e) traceback.print_exc() diff --git a/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py b/scripts/faceswaplab_ui/faceswaplab_postprocessing_ui.py index 6207b5b..ec12a82 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", - "None", + shared.face_restorers[0].name(), ), type="value", elem_id="faceswaplab_pp_face_restorer", diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py b/scripts/faceswaplab_ui/faceswaplab_unit_ui.py index e75b2cf..a76997e 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_ui.py +++ b/scripts/faceswaplab_ui/faceswaplab_unit_ui.py @@ -10,7 +10,9 @@ 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""") + gr.Markdown( + """Post-processing and mask settings for unit faces. Best result : checks all, use LDSR, use Codeformer""" + ) with gr.Row(): face_restorer_name = gr.Radio( label="Restore Face", @@ -209,6 +211,16 @@ def faceswap_unit_ui( elem_id=f"{id_prefix}_face{unit_num}_swap_in_generated", ) + gr.Markdown( + """ +## Advanced Options + +**Simple :** If you have bad results and don't want to fine-tune here, just enable Codeformer in "Global Post-Processing". +Otherwise, read the [doc](https://glucauze.github.io/sd-webui-faceswaplab/doc/) to understand following options. + +""" + ) + with gr.Accordion("Similarity", open=False): gr.Markdown("""Discard images with low similarity or no faces :""") with gr.Row():