Update src/components/faceswap.py

src/components/faceswap.py

@@ -1,68 +1,68 @@
-    import os
-    import cv2
-    import gdown
-    import insightface
-    from insightface.app import FaceAnalysis
-    from PIL import Image
-    from src.utils.logger import logger
-    import warnings
-    from src.components.face_enhancer import load_face_enhancer_model
+import os
+import cv2
+import gdown
+import insightface
+from insightface.app import FaceAnalysis
+from PIL import Image
+from src.utils.logger import logger
+import warnings
+from src.components.face_enhancer import load_face_enhancer_model
 
-    warnings.filterwarnings("ignore", category=FutureWarning)
+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"
-            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.swapper = None
-            self.enhancer_model = None
-            self.model_runner = None
+class FaceSwapper:
+    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.prepare(ctx_id=0 if device == device else -1, det_size=det_size)
+        self.swapper = None
+        self.enhancer_model = None
+        self.model_runner = None
 
-            logger.info('FaceSwapper initialized')
+        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_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):
-            # Set up the gdown cache directory
-            gdown_cache_dir = "/mannequin_to_model/.cache/gdown"
-            os.makedirs(gdown_cache_dir, exist_ok=True)
-            os.environ['XDG_CACHE_HOME'] = gdown_cache_dir
+    def load_swapper_model(self, model_url, model_path):
+        # Set up the gdown cache directory
+        gdown_cache_dir = "/mannequin_to_model/.cache/gdown"
+        os.makedirs(gdown_cache_dir, exist_ok=True)
+        os.environ['XDG_CACHE_HOME'] = gdown_cache_dir
 
-            # Log the cache directory and cookie path for debugging
-            logger.info(f"Setting XDG_CACHE_HOME to: {gdown_cache_dir}")
+        # Log the cache directory and cookie path for debugging
+        logger.info(f"Setting XDG_CACHE_HOME to: {gdown_cache_dir}")
 
-            # 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
-                    pass
-            logger.info(f"Ensured cookies.txt exists at: {cookie_path}")
+        # 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
+                pass
+        logger.info(f"Ensured cookies.txt exists at: {cookie_path}")
 
-            # Check if model path exists, and download if not
-            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)
+        # Check if model path exists, and download if not
+        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')
+        self.swapper = insightface.model_zoo.get_model(model_path, download=False, download_zip=False)
+        logger.info('Swapper model loaded')
 
-        def face_swap(self, img1, img2, enhance=False):
+    def face_swap(self, img1, img2, enhance=False):
 
-            face1 = self.app.get(img1)[0]
-            face2 = self.app.get(img2)[0]
+        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:
-                img1_ = self.model_runner(img1_, self.enhancer_model)
-            logger.info('Face swapped')
-            return img1_
+        img1_ = img1.copy()
+        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_