Update src/components/faceswap.py

src/components/faceswap.py

@@ -12,11 +12,11 @@ warnings.filterwarnings("ignore", category=FutureWarning)
 
 
 class FaceSwapper:
-    def __init__(self, app_name='buffalo_l', det_size=(640, 640), device="cpu"):
+    def __init__(self, app_name='buffalo_l', det_size=(640, 640), device="cuda"):
         model_dir = "/mannequin_to_model/.insightface/models"
         os.makedirs(model_dir, exist_ok=True)
         self.device = device
-        self.app = FaceAnalysis(name=app_name,root=model_dir)
+        self.app = FaceAnalysis(name=app_name, root=model_dir)
         self.app.prepare(ctx_id=0 if device == device else -1, det_size=det_size)
         self.swapper = None
         self.enhancer_model = None
@@ -32,6 +32,10 @@ class FaceSwapper:
         logger.info('Enhancer model loaded')
 
     def load_swapper_model(self, model_url, model_path):
+        gdown_cache_dir = "/mannequin_to_model/.cache/gdown"
+        os.makedirs(gdown_cache_dir, exist_ok=True)
+        os.environ['XDG_CACHE_HOME'] = gdown_cache_dir
+
         if not os.path.exists(model_path):
             os.makedirs(os.path.dirname(model_path), exist_ok=True)
             gdown.download(model_url, model_path, quiet=False)