Update app.py

app.py

@@ -1,93 +1,67 @@
-# -*- coding: UTF-8 -*-
 import gradio as gr
-import cv2
 import numpy as np
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from contextlib import asynccontextmanager
+import cv2
+import onnxruntime
 from insightface.app import FaceAnalysis
-from insightface.model_zoo.face_swapper import FaceSwapper
 from pathlib import Path
-import requests
-import os
-import logging
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Initialize FastAPI
-app = FastAPI(title="Face Swap API", description="Gradio-based Face Swap App")
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-# Health Check Endpoint
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy"}
+# Initialize Face Analysis
+face_analyzer = FaceAnalysis(name="buffalo_l")
+face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
 
-# Download Model Function
+# Load Face Swapper Model
 MODEL_PATH = Path("models/inswapper_128.onnx")
-MODEL_URL = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
+if not MODEL_PATH.exists():
+    raise FileNotFoundError("Model file inswapper_128.onnx not found.")
 
-def download_model():
-    if not MODEL_PATH.exists():
-        logger.info("Downloading model...")
-        os.makedirs("models", exist_ok=True)
-        response = requests.get(MODEL_URL, stream=True, timeout=30)
-        response.raise_for_status()
-        with open(MODEL_PATH, 'wb') as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                f.write(chunk)
-        logger.info("Model downloaded successfully.")
+session = onnxruntime.InferenceSession(str(MODEL_PATH))
+
+def swap_faces(source_img, target_img):
+    """Perform face swapping using the ONNX model."""
+    try:
+        # Convert images to correct format
+        source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
+        target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
 
-download_model()
+        # Detect faces
+        source_faces = face_analyzer.get(source_img)
+        target_faces = face_analyzer.get(target_img)
 
-# Initialize Face Analyzer & Swapper
-face_analyzer = FaceAnalysis(name="buffalo_l")
-face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
-face_swapper = FaceSwapper(MODEL_PATH)
+        if not source_faces or not target_faces:
+            return "No faces detected in one or both images."
+        if len(source_faces) > 1 or len(target_faces) > 1:
+            return "Multiple faces detected; only one face per image is supported."
 
-# Face Swap Function
-def swap_faces(source_img, target_img):
-    source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
-    target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
+        source_face = source_faces[0]
+        target_face = target_faces[0]
 
-    source_faces = face_analyzer.get(source_img)
-    target_faces = face_analyzer.get(target_img)
+        # Prepare input data for ONNX model
+        input_data = {
+            "target_image": target_img,
+            "target_face": target_face.embedding,
+            "source_face": source_face.embedding
+        }
 
-    if not source_faces or not target_faces:
-        return "Error: No faces detected in one or both images."
-    if len(source_faces) > 1 or len(target_faces) > 1:
-        return "Error: Multiple faces detected; only one face per image is supported."
+        # Run the ONNX model
+        result = session.run(None, input_data)[0]
 
-    result_img = face_swapper.get(target_img, target_faces[0], source_faces[0], paste_back=True)
-    return cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
+        # Convert result to image format
+        result_img = np.clip(result * 255, 0, 255).astype(np.uint8)
+        result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
 
-# Gradio UI
-def gradio_face_swap(source_img, target_img):
-    try:
-        result = swap_faces(source_img, target_img)
-        if isinstance(result, str):
-            return result  # Error message
-        return result
+        return result_img
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Face swap failed: {e}"
 
-gr_interface = gr.Interface(
-    fn=gradio_face_swap,
-    inputs=[
-        gr.Image(label="Upload Source Face"),
-        gr.Image(label="Upload Target Image")
-    ],
-    outputs=gr.Image(label="Swapped Face"),
-    title="AI Face Swap",
-    description="Upload two images: one with the source face and another where you want the face swapped.",
-    live=True
-)
-
-# Launch Gradio
-app = gr.mount_gradio_app(app, gr_interface, path="/")
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# Face Swap Tool 🚀")
+    with gr.Row():
+        input_source = gr.Image(label="Source Face", type="pil")
+        input_target = gr.Image(label="Target Image", type="pil")
+    btn_swap = gr.Button("Swap Faces")
+    output_image = gr.Image(label="Swapped Face")
+    btn_swap.click(swap_faces, inputs=[input_source, input_target], outputs=output_image)
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+# Launch Gradio App
+demo.launch()