Update src/components/faceswap.py

src/components/faceswap.py

@@ -10,18 +10,16 @@ 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="cpu"):
-        model_dir = "/mannequin_to_model/.insightface/models"
+        model_dir = "/tmp/.insightface/models"
         os.makedirs(model_dir, exist_ok=True)
         self.device = device
         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.app.prepare(ctx_id=0 if device == "cuda" else -1, det_size=det_size)
         self.swapper = None
         self.enhancer_model = None
         self.model_runner = None
-
         logger.info('FaceSwapper initialized')
 
     def load_enhancer_model(self, enhancer, device):
@@ -33,7 +31,7 @@ class FaceSwapper:
 
     def load_swapper_model(self, model_url, model_path):
         # Set up the gdown cache directory
-        gdown_cache_dir = "/mannequin_to_model/.cache/gdown"
+        gdown_cache_dir = "/tmp/.cache/gdown"
         os.makedirs(gdown_cache_dir, exist_ok=True)
         os.environ['XDG_CACHE_HOME'] = gdown_cache_dir
 
@@ -43,7 +41,8 @@ class FaceSwapper:
         # Ensure cookies.txt exists
         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
+            with open(cookie_path, 'w') as f:
+                # Create an empty cookies.txt file
                 pass
         logger.info(f"Ensured cookies.txt exists at: {cookie_path}")
 
@@ -56,10 +55,8 @@ class FaceSwapper:
         logger.info('Swapper model loaded')
 
     def face_swap(self, img1, img2, enhance=False):
-
         face1 = self.app.get(img1)[0]
         face2 = self.app.get(img2)[0]
-
         img1_ = img1.copy()
         img1_ = self.swapper.get(img1_, face1, face2, paste_back=True)
         if enhance: