Update app.py

app.py

@@ -1,163 +1,93 @@
 # -*- coding: UTF-8 -*-
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.responses import Response
-from fastapi.middleware.cors import CORSMiddleware
-from contextlib import asynccontextmanager
+import gradio as gr
 import cv2
 import numpy as np
-from PIL import Image
-import os
-import requests
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from contextlib import asynccontextmanager
+from insightface.app import FaceAnalysis
+from insightface.model_zoo.face_swapper import FaceSwapper
 from pathlib import Path
-import uvicorn
+import requests
+import os
 import logging
 
-# Initialize FastAPI with explicit docs settings
-app = FastAPI(
-    title="Face Swap API",
-    description="API for swapping faces in images.",
-    docs_url="/docs",  # Explicitly set docs URL
-    redoc_url="/redoc",  # Explicitly set redoc URL
-)
-
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Update with your Framer domain in production
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# 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"}
 
-# Global flag to prevent multiple downloads
-MODEL_DOWNLOADED = False
+# Download Model Function
+MODEL_PATH = Path("models/inswapper_128.onnx")
+MODEL_URL = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
 
 def download_model():
-    global MODEL_DOWNLOADED
-    if MODEL_DOWNLOADED:
-        logger.info("Model already downloaded, skipping.")
-        return
-
-    model_dir = Path("models")
-    model_path = model_dir / "inswapper_128.onnx"
-    model_url = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
-
-    if not model_path.exists():
-        logger.info("Model not found. Downloading inswapper_128.onnx...")
-        model_dir.mkdir(exist_ok=True)
-        try:
-            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.")
-            MODEL_DOWNLOADED = True
-        except Exception as e:
-            logger.error(f"Failed to download model: {e}")
-            raise RuntimeError("Could not download inswapper_128.onnx. Please check logs.")
-    else:
-        logger.info("Model already exists at: %s", model_path)
-        MODEL_DOWNLOADED = True
-
-# Use lifespan event handler
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    logger.info("Starting up application...")
-    try:
-        download_model()
-        logger.info("Startup completed successfully.")
-    except Exception as e:
-        logger.error(f"Startup failed: {e}")
-        raise
-    yield
-    logger.info("Shutting down application...")
-
-app.lifespan = lifespan
-
+    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.")
+
+download_model()
+
+# 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)
+
+# Face Swap Function
 def swap_faces(source_img, target_img):
-    """Perform face swapping using insightface and inswapper model."""
-    try:
-        from insightface.app import FaceAnalysis
-        from insightface.utils import face_align
-        from insightface.model_zoo import face_swapper
-
-        face_analyzer = FaceAnalysis(name="buffalo_l")
-        face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
-
-        source_faces = face_analyzer.get(source_img)
-        target_faces = face_analyzer.get(target_img)
-
-        if not source_faces or not target_faces:
-            raise ValueError("No faces detected in one or both images.")
-        if len(source_faces) > 1 or len(target_faces) > 1:
-            raise ValueError("Multiple faces detected; only one face per image is supported.")
-
-        source_face = source_faces[0]
-        target_face = target_faces[0]
+    source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
+    target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
 
-        model_path = Path("models/inswapper_128.onnx")
-        if not model_path.exists():
-            raise FileNotFoundError("Model file inswapper_128.onnx not found.")
-        swapper = face_swapper.FaceSwapper(model_path)
+    source_faces = face_analyzer.get(source_img)
+    target_faces = face_analyzer.get(target_img)
 
-        result = swapper.get(target_img, target_face, source_face, paste_back=True)
+    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."
 
-        target_pil = Image.fromarray(cv2.cvtColor(target_img, cv2.COLOR_BGR2RGB))
-        result_pil = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
-        result_pil = result_pil.resize(target_pil.size, Image.Resampling.LANCZOS)
+    result_img = face_swapper.get(target_img, target_faces[0], source_faces[0], paste_back=True)
+    return cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
 
-        return cv2.cvtColor(np.array(result_pil), cv2.COLOR_RGB2BGR)
-    except Exception as e:
-        logger.error(f"Face swap failed: {e}")
-        raise
-
-@app.post("/swap-face/")
-async def swap_face(source_file: UploadFile = File(...), target_file: UploadFile = File(...)):
+# Gradio UI
+def gradio_face_swap(source_img, target_img):
     try:
-        source_path = "temp_source.jpg"
-        target_path = "temp_target.jpg"
-        output_path = "output.jpg"
-
-        source_content = await source_file.read()
-        with open(source_path, "wb") as f:
-            f.write(source_content)
-        source_img = cv2.imread(source_path)
-        if source_img is None:
-            raise ValueError("Failed to load source image.")
-
-        target_content = await target_file.read()
-        with open(target_path, "wb") as f:
-            f.write(target_content)
-        target_img = cv2.imread(target_path)
-        if target_img is None:
-            raise ValueError("Failed to load target image.")
-
-        result_img = swap_faces(source_img, target_img)
-
-        cv2.imwrite(output_path, result_img)
-
-        with open(output_path, "rb") as f:
-            image_data = f.read()
-
-        for path in [source_path, target_path, output_path]:
-            if os.path.exists(path):
-                os.remove(path)
-
-        return Response(content=image_data, media_type="image/jpeg")
-
+        result = swap_faces(source_img, target_img)
+        if isinstance(result, str):
+            return result  # Error message
+        return result
     except Exception as e:
-        logger.error("Error in swap_face: %s", str(e))
-        raise HTTPException(status_code=500, detail=str(e))
+        return f"Error: {str(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="/")
 
 if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)