Spaces:
Sleeping
Sleeping
| 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_ |