From 9bb1c65db671f56f6dba0bcd8032951ca744b980 Mon Sep 17 00:00:00 2001 From: Tran Xen <137925069+glucauze@users.noreply.github.com> Date: Fri, 4 Aug 2023 14:23:50 +0200 Subject: [PATCH] experimental gpu option --- install.py | 11 +++++- preload.py | 5 +++ requirements-gpu.txt | 11 ++++++ .../faceswaplab_settings.py | 13 ++++++- scripts/faceswaplab_swapping/swapper.py | 38 +++++++++++-------- .../upscaled_inswapper.py | 1 - 6 files changed, 61 insertions(+), 18 deletions(-) create mode 100644 requirements-gpu.txt diff --git a/install.py b/install.py index ba3aff8..172da70 100644 --- a/install.py +++ b/install.py @@ -2,9 +2,18 @@ import os import pkg_resources import sys +from modules import shared +use_gpu = getattr(shared.cmd_opts, "faceswaplab_gpu", False) -req_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "requirements.txt") +if use_gpu and sys.platform != "darwin": + req_file = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "requirements-gpu.txt" + ) +else: + req_file = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "requirements.txt" + ) print("Checking faceswaplab requirements") with open(req_file) as file: diff --git a/preload.py b/preload.py index 5c3eaf5..983e711 100644 --- a/preload.py +++ b/preload.py @@ -8,3 +8,8 @@ def preload(parser: ArgumentParser) -> None: choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], help="Set the log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)", ) + parser.add_argument( + "--faceswaplab_gpu", + action="store_true", + help="Enable GPU if set, disable if not set", + ) diff --git a/requirements-gpu.txt b/requirements-gpu.txt new file mode 100644 index 0000000..83b9209 --- /dev/null +++ b/requirements-gpu.txt @@ -0,0 +1,11 @@ +cython +dill +ifnude +insightface==0.7.3 +onnx==1.14.0 +opencv-python +pandas +pydantic +safetensors +onnxruntime==1.15.1 +onnxruntime-gpu==1.15.1 \ No newline at end of file diff --git a/scripts/faceswaplab_settings/faceswaplab_settings.py b/scripts/faceswaplab_settings/faceswaplab_settings.py index 71c315c..d826061 100644 --- a/scripts/faceswaplab_settings/faceswaplab_settings.py +++ b/scripts/faceswaplab_settings/faceswaplab_settings.py @@ -37,11 +37,22 @@ def on_ui_settings() -> None: ), ) + shared.opts.add_option( + "faceswaplab_det_size", + shared.OptionInfo( + 640, + "det_size : Size of the detection area for face analysis. Higher values may improve quality but reduce speed. Low value may improve detection of very large face.", + gr.Slider, + {"minimum": 320, "maximum": 640, "step": 320}, + section=section, + ), + ) + shared.opts.add_option( "faceswaplab_detection_threshold", shared.OptionInfo( 0.5, - "Face Detection threshold", + "det_thresh : Face Detection threshold", gr.Slider, {"minimum": 0.1, "maximum": 0.99, "step": 0.001}, section=section, diff --git a/scripts/faceswaplab_swapping/swapper.py b/scripts/faceswaplab_swapping/swapper.py index e61ee89..150cfc3 100644 --- a/scripts/faceswaplab_swapping/swapper.py +++ b/scripts/faceswaplab_swapping/swapper.py @@ -37,8 +37,19 @@ 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 -providers = ["CPUExecutionProvider"] + +use_gpu = getattr(shared.cmd_opts, "faceswaplab_gpu", False) + +if use_gpu and sys.platform != "darwin": + providers = [ + "TensorrtExecutionProvider", + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] +else: + providers = ["CPUExecutionProvider"] def cosine_similarity_face(face1: Face, face2: Face) -> float: @@ -258,7 +269,9 @@ def capture_stdout() -> Generator[StringIO, None, None]: @lru_cache(maxsize=1) -def getAnalysisModel() -> insightface.app.FaceAnalysis: +def getAnalysisModel( + det_size: Tuple[int, int] = (640, 640), det_thresh: float = 0.5 +) -> insightface.app.FaceAnalysis: """ Retrieves the analysis model for face analysis. @@ -269,7 +282,9 @@ 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. (> 30s)") + logger.info( + f"Load analysis model det_size={det_size}, det_thresh={det_thresh}, will take some time. (> 30s)" + ) # Initialize the analysis model with the specified name and providers with tqdm( @@ -281,6 +296,9 @@ def getAnalysisModel() -> insightface.app.FaceAnalysis: 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) pbar.update(1) logger.info("%s", pformat(captured.getvalue())) @@ -350,7 +368,6 @@ def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper: def get_faces( img_data: CV2ImgU8, - det_size: Tuple[int, int] = (640, 640), det_thresh: Optional[float] = None, ) -> List[Face]: """ @@ -368,21 +385,12 @@ def get_faces( if det_thresh is None: det_thresh = opts.data.get("faceswaplab_detection_threshold", 0.5) - # Create a deep copy of the analysis model (otherwise det_size is attached to the analysis model and can't be changed) - face_analyser = copy.deepcopy(getAnalysisModel()) - - # Prepare the analysis model for face detection with the specified detection size - face_analyser.prepare(ctx_id=0, det_thresh=det_thresh, det_size=det_size) + det_size = opts.data.get("faceswaplab_det_size", 640) + face_analyser = getAnalysisModel((det_size, det_size), det_thresh) # Get the detected faces from the image using the analysis model face = 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(face) == 0 and 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) - try: # Sort the detected faces based on their x-coordinate of the bounding box return sorted(face, key=lambda x: x.bbox[0]) diff --git a/scripts/faceswaplab_swapping/upscaled_inswapper.py b/scripts/faceswaplab_swapping/upscaled_inswapper.py index ec3fe7d..47f5e1e 100644 --- a/scripts/faceswaplab_swapping/upscaled_inswapper.py +++ b/scripts/faceswaplab_swapping/upscaled_inswapper.py @@ -262,7 +262,6 @@ def compute_diff(bgr_fake: CV2ImgU8, aimg: CV2ImgU8) -> CV2ImgU8: ) img_white[img_white > 20] = 255 fthresh = 10 - print("fthresh", fthresh) fake_diff[fake_diff < fthresh] = 0 fake_diff[fake_diff >= fthresh] = 255 img_mask = img_white