Spaces:
Running
Running
| import os | |
| import cv2 | |
| import tempfile | |
| import time | |
| import gradio as gr | |
| import insightface | |
| import onnxruntime | |
| from gfpgan import GFPGANer | |
| class Predictor: | |
| def __init__(self): | |
| self.setup() | |
| def setup(self): | |
| model_dir = "models" | |
| print("Loading models...") | |
| self.face_swapper = insightface.model_zoo.get_model( | |
| f"{model_dir}/inswapper_128.onnx", | |
| providers=onnxruntime.get_available_providers() | |
| ) | |
| self.face_enhancer = GFPGANer( | |
| model_path=f"{model_dir}/GFPGANv1.4.pth", | |
| upscale=1, | |
| ) | |
| self.face_analyser = insightface.app.FaceAnalysis(name="buffalo_l") | |
| self.face_analyser.prepare(ctx_id=0, det_size=(640, 640)) | |
| def get_face(self, img_data): | |
| """Return the largest detected face""" | |
| faces = self.face_analyser.get(img_data) | |
| if len(faces) == 0: | |
| return None | |
| return max(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1])) | |
| def predict(self, target_image, swap_image): | |
| """Face swap function""" | |
| try: | |
| frame = cv2.imread(target_image) | |
| source = cv2.imread(swap_image) | |
| if frame is None or source is None: | |
| return None, "❌ Could not read one of the images." | |
| target_face = self.get_face(frame) | |
| source_face = self.get_face(source) | |
| if target_face is None or source_face is None: | |
| return None, "❌ Could not detect a clear face in one or both images." | |
| # Perform face swap | |
| result = self.face_swapper.get(frame, target_face, source_face, paste_back=True) | |
| # Enhance with GFPGAN | |
| _, _, result = self.face_enhancer.enhance(result, paste_back=True) | |
| # Save result | |
| out_path = tempfile.mkdtemp() + f"/swapped_{int(time.time())}.jpg" | |
| cv2.imwrite(out_path, result) | |
| return out_path, "✅ Face swap completed successfully!" | |
| except Exception as e: | |
| print(f"Prediction error: {e}") | |
| return None, f"❌ Error: {str(e)}" | |
| # Load the predictor once at startup | |
| predictor = Predictor() | |
| # Clean UI without examples | |
| with gr.Blocks(title="Swap Face Model", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# 😻 Face Swap Model") | |
| gr.Markdown("Upload a **target image** and a **face to swap in**. Works best with clear, well-lit frontal faces.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| target_input = gr.Image( | |
| type="filepath", | |
| label="Target Image (the photo you want to change)", | |
| height=450 | |
| ) | |
| with gr.Column(): | |
| swap_input = gr.Image( | |
| type="filepath", | |
| label="Swap Face Image (the face you want to use)", | |
| height=450 | |
| ) | |
| with gr.Row(): | |
| swap_btn = gr.Button("🔄 Swap Faces", variant="primary", size="large") | |
| with gr.Row(): | |
| output_image = gr.Image( | |
| type="filepath", | |
| label="Result", | |
| height=500 | |
| ) | |
| status = gr.Textbox(label="Status", interactive=False) | |
| # Button action | |
| swap_btn.click( | |
| fn=predictor.predict, | |
| inputs=[target_input, swap_input], | |
| outputs=[output_image, status] | |
| ) | |
| gr.Markdown("**Tips:**\n• Use high-resolution, well-lit photos\n• Frontal or near-frontal faces work best\n• Results improve with similar lighting and angles.") | |
| if __name__ == "__main__": | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False, | |
| show_error=True | |
| ) |