Spaces:
Sleeping
Sleeping
File size: 2,041 Bytes
6888585 36cd99b 6888585 36cd99b 6888585 ede97cc 608091e 6888585 608091e 6888585 070b382 6888585 36cd99b a6b6ab1 6888585 608091e 6888585 d05216e a6b6ab1 8eed4f7 d05216e a6b6ab1 6888585 8eed4f7 6888585 608091e 6888585 36cd99b 6888585 759b1bf ff11707 6888585 d9df3d5 6888585 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import os
import gdown
import insightface
from insightface.app import FaceAnalysis
from src.utils.logger import logger
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)
class FaceSwapper:
def __init__(self, app_name='buffalo_l', det_size=(640, 640), device="cpu"):
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 == "cuda" else -1, det_size=det_size)
self.swapper = None
logger.info('FaceSwapper initialized')
def load_swapper_model(self, model_url, model_path):
# Set up the gdown cache directory
gdown_cache_dir = "/tmp/.cache/gdown"
os.makedirs(gdown_cache_dir, exist_ok=True)
os.environ['XDG_CACHE_HOME'] = "/tmp/.cache"
# Log the cache directory for debugging
logger.info(f"Setting XDG_CACHE_HOME to: {os.environ['XDG_CACHE_HOME']}")
# 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, use_cookies=False)
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):
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)
logger.info('Face swapped')
return img1_ |