ishworrsubedii commited on
Commit
d9df3d5
·
verified ·
1 Parent(s): 8774faf

Update src/components/faceswap.py

Browse files
Files changed (1) hide show
  1. src/components/faceswap.py +8 -3
src/components/faceswap.py CHANGED
@@ -10,7 +10,6 @@ from src.components.face_enhancer import load_face_enhancer_model
10
 
11
  warnings.filterwarnings("ignore", category=FutureWarning)
12
 
13
-
14
  class FaceSwapper:
15
  def __init__(self, app_name='buffalo_l', det_size=(640, 640), device="cuda"):
16
  model_dir = "/mannequin_to_model/.insightface/models"
@@ -36,14 +35,19 @@ class FaceSwapper:
36
  os.makedirs(gdown_cache_dir, exist_ok=True)
37
  os.environ['XDG_CACHE_HOME'] = gdown_cache_dir
38
 
 
 
 
 
 
39
  if not os.path.exists(model_path):
40
  os.makedirs(os.path.dirname(model_path), exist_ok=True)
41
  gdown.download(model_url, model_path, quiet=False)
 
42
  self.swapper = insightface.model_zoo.get_model(model_path, download=False, download_zip=False)
43
  logger.info('Swapper model loaded')
44
 
45
  def face_swap(self, img1, img2, enhance=False):
46
-
47
  face1 = self.app.get(img1)[0]
48
  face2 = self.app.get(img2)[0]
49
 
@@ -51,5 +55,6 @@ class FaceSwapper:
51
  img1_ = self.swapper.get(img1_, face1, face2, paste_back=True)
52
  if enhance:
53
  img1_ = self.model_runner(img1_, self.enhancer_model)
 
54
  logger.info('Face swapped')
55
- return img1_
 
10
 
11
  warnings.filterwarnings("ignore", category=FutureWarning)
12
 
 
13
  class FaceSwapper:
14
  def __init__(self, app_name='buffalo_l', det_size=(640, 640), device="cuda"):
15
  model_dir = "/mannequin_to_model/.insightface/models"
 
35
  os.makedirs(gdown_cache_dir, exist_ok=True)
36
  os.environ['XDG_CACHE_HOME'] = gdown_cache_dir
37
 
38
+ # Ensure cookies.txt exists
39
+ cookie_path = os.path.join(gdown_cache_dir, "cookies.txt")
40
+ if not os.path.exists(cookie_path):
41
+ open(cookie_path, 'w').close()
42
+
43
  if not os.path.exists(model_path):
44
  os.makedirs(os.path.dirname(model_path), exist_ok=True)
45
  gdown.download(model_url, model_path, quiet=False)
46
+
47
  self.swapper = insightface.model_zoo.get_model(model_path, download=False, download_zip=False)
48
  logger.info('Swapper model loaded')
49
 
50
  def face_swap(self, img1, img2, enhance=False):
 
51
  face1 = self.app.get(img1)[0]
52
  face2 = self.app.get(img2)[0]
53
 
 
55
  img1_ = self.swapper.get(img1_, face1, face2, paste_back=True)
56
  if enhance:
57
  img1_ = self.model_runner(img1_, self.enhancer_model)
58
+
59
  logger.info('Face swapped')
60
+ return img1_