Update src/components/faceswap.py

src/components/faceswap.py

@@ -10,6 +10,7 @@ from src.components.face_enhancer import load_face_enhancer_model
 
 warnings.filterwarnings("ignore", category=FutureWarning)
 
+
 class FaceSwapper:
     def __init__(self, app_name='buffalo_l', det_size=(640, 640), device="cuda"):
         model_dir = "/mannequin_to_model/.insightface/models"
@@ -23,12 +24,22 @@ class FaceSwapper:
 
         logger.info('FaceSwapper initialized')
 
-    def load_enhancer_model(self, enhancer, device):
-        model, model_runner = load_face_enhancer_model(enhancer, device)
-        logger.info(f'{enhancer} model loaded')
-        self.enhancer_model = model
-        self.model_runner = model_runner
-        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
+
+        cookie_path = os.path.join(gdown_cache_dir, "cookies.txt")
+        if not os.path.exists(cookie_path):
+            with open(cookie_path, 'w') as f:  # Create an empty cookies.txt file
+                pass
+
+        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)
+
+        self.swapper = insightface.model_zoo.get_model(model_path, download=False, download_zip=False)
+        logger.info('Swapper model loaded')
 
     def load_swapper_model(self, model_url, model_path):
         gdown_cache_dir = "/mannequin_to_model/.cache/gdown"
@@ -43,7 +54,7 @@ class FaceSwapper:
         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)
-        
+
         self.swapper = insightface.model_zoo.get_model(model_path, download=False, download_zip=False)
         logger.info('Swapper model loaded')
 
@@ -55,6 +66,6 @@ class FaceSwapper:
         img1_ = self.swapper.get(img1_, face1, face2, paste_back=True)
         if enhance:
             img1_ = self.model_runner(img1_, self.enhancer_model)
-        
+
         logger.info('Face swapped')
         return img1_