fix gfpgan cpu/mps device

2023-03-28 16:36:41 +08:00
parent 8af5df491e
commit c4968dd0a9
3 changed files with 14 additions and 6 deletions
--- a/lama_cleaner/plugins/gfpgan_plugin.py
+++ b/lama_cleaner/plugins/gfpgan_plugin.py
@@ -17,6 +17,10 @@ class GFPGANPlugin(BasePlugin):
        model_path = download_model(url, model_md5)
        logger.info(f"GFPGAN model path: {model_path}")

+        import facexlib
+        if hasattr(facexlib.detection.retinaface, "device"):
+            facexlib.detection.retinaface.device = device
+
        # Use GFPGAN for face enhancement
        self.face_enhancer = MyGFPGANer(
            model_path=model_path,