From c3cbf8d4a2f3568195cf3125758c7f60cfad6b1b Mon Sep 17 00:00:00 2001 From: Qing Date: Fri, 31 Mar 2023 22:17:38 +0800 Subject: [PATCH] update --- lama_cleaner/plugins/gfpgan_plugin.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/lama_cleaner/plugins/gfpgan_plugin.py b/lama_cleaner/plugins/gfpgan_plugin.py index 1b826f2..12e66b5 100644 --- a/lama_cleaner/plugins/gfpgan_plugin.py +++ b/lama_cleaner/plugins/gfpgan_plugin.py @@ -17,9 +17,11 @@ class GFPGANPlugin(BasePlugin): model_path = download_model(url, model_md5) logger.info(f"GFPGAN model path: {model_path}") + face_det_device = "cpu" if "cuda" in str(device) else device import facexlib + if hasattr(facexlib.detection.retinaface, "device"): - facexlib.detection.retinaface.device = "cpu" + facexlib.detection.retinaface.device = device # Use GFPGAN for face enhancement self.face_enhancer = MyGFPGANer( @@ -30,9 +32,10 @@ class GFPGANPlugin(BasePlugin): device=device, bg_upsampler=upscaler.model if upscaler is not None else None, ) - self.face_enhancer.face_helper.face_det.mean_tensor.to("cpu") - self.face_enhancer.face_helper.face_det = self.face_enhancer.face_helper.face_det.to("cpu") - + self.face_enhancer.face_helper.face_det.mean_tensor.to(device) + self.face_enhancer.face_helper.face_det = ( + self.face_enhancer.face_helper.face_det.to(device) + ) def __call__(self, rgb_np_img, files, form): weight = 0.5