import os import sys import time import uuid import ssl import traceback import tempfile import subprocess from typing import Any import gradio as gr import cv2 import requests import torch from PIL import Image import insightface import onnxruntime from insightface.app import FaceAnalysis # Try importing GFPGAN safely so dependency issues are clearer try: import gfpgan except Exception as e: gfpgan = None gfpgan_import_error = e else: gfpgan_import_error = None from loggers import logger, request_id as _request_id ssl._create_default_https_context = ssl._create_unverified_context if sys.platform == 'darwin': cache_file_dir = '/tmp/file' else: cache_file_dir = '/src/file' os.makedirs(cache_file_dir, exist_ok=True) def img_url_to_local_path(img_url, file_path=None): filename = img_url.split('/')[-1] max_count = 3 count = 0 if file_path is None: temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(filename)[-1] or ".jpg") temp_file_name = temp_file.name temp_file.close() else: temp_file_name = file_path while True: count += 1 try: res = requests.get(img_url, timeout=60) res.raise_for_status() with open(temp_file_name, "wb") as f: f.write(res.content) return temp_file_name except Exception as e: logger.error(e) if count >= max_count: msg = f'request {max_count} times for url: {img_url} failed, please check' logger.error(msg) raise Exception(msg) def delete_files_day_ago(cache_days=10): command = f"find {cache_file_dir} -type f -ctime +{cache_days} -exec rm -f {{}} \\;" result = subprocess.run(command, shell=True, capture_output=True, text=True) if result.stdout: logger.info(result.stdout) if result.stderr: logger.warning(result.stderr) def image_format_by_path(image_path): image = Image.open(image_path) image_format = image.format if not image_format: image_format = 'jpg' elif image_format == "JPEG": image_format = 'jpg' else: image_format = image_format.lower() return image_format def local_file_for_url(url, cache_days=10): filename = url.split('/')[-1] file_path = os.path.join(cache_file_dir, filename) if not os.path.exists(file_path): img_url_to_local_path(url, file_path) logger.info(f'download file to {file_path}') delete_files_day_ago(cache_days) else: logger.info(f'cache file {file_path}') return file_path class Predictor: def __init__(self): self.det_thresh = 0.1 self.face_swapper = None self.face_enhancer = None self.face_analyser = None def setup(self): providers = onnxruntime.get_available_providers() self.face_swapper = insightface.model_zoo.get_model( 'cache/inswapper_128.onnx', providers=providers ) self.face_analyser = FaceAnalysis(name='buffalo_l') self.face_analyser.prepare(ctx_id=0, det_thresh=self.det_thresh) if gfpgan is None: raise ImportError( f"gfpgan import failed: {gfpgan_import_error}. " "This is usually caused by an incompatible torchvision/basicsr version." ) self.face_enhancer = gfpgan.GFPGANer( model_path='cache/GFPGANv1.4.pth', upscale=1 ) def get_face(self, img_data, image_type='target'): try: if self.face_analyser is None: raise RuntimeError("Face analyser is not initialized. Call setup() first.") if image_type == 'source': self.face_analyser.prepare(ctx_id=0, det_thresh=self.det_thresh) analysed = self.face_analyser.get(img_data) logger.info(f'face num: {len(analysed)}') if len(analysed) == 0: msg = 'no face' logger.error(msg) raise Exception(msg) largest = max( analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]) ) return largest except Exception as e: logger.error(str(e)) raise def enhance_face(self, target_face, target_frame, weight=0.5): if self.face_enhancer is None: raise RuntimeError("Face enhancer is not initialized. Call setup() first.") start_x, start_y, end_x, end_y = map(int, target_face['bbox']) padding_x = int((end_x - start_x) * 0.5) padding_y = int((end_y - start_y) * 0.5) start_x = max(0, start_x - padding_x) start_y = max(0, start_y - padding_y) end_x = min(target_frame.shape[1], end_x + padding_x) end_y = min(target_frame.shape[0], end_y + padding_y) temp_face = target_frame[start_y:end_y, start_x:end_x] if temp_face.size: _, _, temp_face = self.face_enhancer.enhance( temp_face, paste_back=True, weight=weight ) target_frame[start_y:end_y, start_x:end_x] = temp_face return target_frame def predict( self, source_image_path, target_image_path, enhance_face, ) -> Any: request_id = None det_thresh = 0.1 weight = 0.5 if torch.cuda.is_available(): device = 'cuda' elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): device = 'mps' else: device = 'cpu' logger.info(f'device: {device}, det_thresh:{det_thresh}') try: self.det_thresh = det_thresh start_time = time.time() if not request_id: request_id = str(uuid.uuid4()) _request_id.set(request_id) frame = cv2.imread(str(target_image_path)) source_frame = cv2.imread(str(source_image_path)) if frame is None: raise ValueError(f"Failed to read target image: {target_image_path}") if source_frame is None: raise ValueError(f"Failed to read source image: {source_image_path}") source_face = self.get_face(source_frame, image_type='source') target_face = self.get_face(frame) ext = image_format_by_path(target_image_path) size = os.path.getsize(target_image_path) logger.info(f'origin {size / 1024:.2f}k') result = self.face_swapper.get(frame, target_face, source_face, paste_back=True) if enhance_face: result = self.enhance_face(target_face, result, weight) out_dir = tempfile.mkdtemp() out_path = os.path.join(out_dir, f"{uuid.uuid4()}.{ext}") cv2.imwrite(str(out_path), result) out_size = os.path.getsize(out_path) logger.info(f'result {out_size / 1024:.2f}k') cost_time = time.time() - start_time logger.info(f'total time: {cost_time * 1000:.2f} ms') return Image.open(out_path) except Exception as e: logger.error(traceback.format_exc()) logger.error(str(e)) raise def swap_faces(source_image_path, target_image_path, enhance_face): predictor = Predictor() predictor.setup() return predictor.predict( source_image_path, target_image_path, enhance_face ) if __name__ == "__main__": demo = gr.Interface( fn=swap_faces, inputs=[ gr.Image(type="filepath"), gr.Image(type="filepath"), gr.Checkbox(label="Enhance Face", value=True), ], outputs=[ # Removed unsupported `show_download_button` for the installed Gradio version gr.Image(type="pil") ], title="Mar's Face Swap", #allow_flagging="never" ) demo.launch(share=True)