Spaces:
Sleeping
Sleeping
Update src/components/faceswap.py
Browse files
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_
|