Update app.py

app.py

@@ -16,66 +16,60 @@ import uvicorn
 app = FastAPI(
     title="Face Swap API",
     description="API for swapping faces in images.",
-    docs_url="/docs",  
-    redoc_url="/redoc"
+    docs_url="/docs",
+    redoc_url="/redoc",
 )
 
-# Configure logging
+# Logging setup
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Add CORS middleware
+# CORS setup
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  
+    allow_origins=["*"],  # Update with your domain in production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Root endpoint
+# Health check route
 @app.get("/")
 async def root():
-    return {"message": "Welcome to the Face Swap API! Use /swap-face/ to swap faces or /docs to test the API."}
+    return {"message": "Face Swap API is running. Use /docs to test the API."}
 
-# Health check
 @app.get("/health")
 async def health_check():
     return {"status": "healthy"}
 
-# Global flag for model download
-MODEL_DOWNLOADED = False
+# Prevent multiple downloads
+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
-    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("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.")
-    else:
-        logger.info("Model already exists.")
-        MODEL_DOWNLOADED = True
+    if MODEL_PATH.exists():
+        logger.info("Model already exists, skipping download.")
+        return
+    logger.info("Downloading model...")
+    MODEL_PATH.parent.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.")
+    except Exception as e:
+        logger.error(f"Failed to download model: {e}")
+        raise RuntimeError("Could not download inswapper_128.onnx.")
 
+# FastAPI startup event
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     logger.info("Starting application...")
     try:
         download_model()
-        logger.info("Application started successfully.")
+        logger.info("Startup completed successfully.")
     except Exception as e:
         logger.error(f"Startup failed: {e}")
         raise
@@ -84,43 +78,37 @@ async def lifespan(app: FastAPI):
 
 app.lifespan = lifespan
 
-def swap_faces(source_img, target_img):
+# Face detection and swap functions
+def get_faces(image):
     try:
         from insightface.app import FaceAnalysis
+        app = FaceAnalysis(name="buffalo_l")
+        app.prepare(ctx_id=0, det_size=(640, 640))
+        return app.get(image) or []
+    except Exception as e:
+        logger.error(f"Face detection failed: {e}")
+        raise
+
+def swap_faces(source_img, target_img):
+    try:
         from insightface.utils import face_align
         from insightface.model_zoo import face_swapper
 
-        # Initialize face analysis
         face_analyzer = FaceAnalysis(name="buffalo_l")
-        face_analyzer.prepare(ctx_id=-1, det_size=(640, 640))
+        face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
 
-        # Detect faces
         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.")
         if len(source_faces) > 1 or len(target_faces) > 1:
-            raise ValueError("Multiple faces detected; only one per image supported.")
-
-        source_face = source_faces[0]
-        target_face = target_faces[0]
+            raise ValueError("Multiple faces detected. Only one face per image is supported.")
 
-        # Load the face swapper model
-        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(str(model_path))
+        swapper = face_swapper.FaceSwapper(MODEL_PATH)
+        result = swapper.get(target_img, target_faces[0], source_faces[0], paste_back=True)
 
-        # Perform face swap
-        result = swapper.get(target_img, target_face, source_face, paste_back=True)
-
-        # Resize result to match target image size
-        result_pil = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
-        target_pil = Image.fromarray(cv2.cvtColor(target_img, cv2.COLOR_BGR2RGB))
-        result_pil = result_pil.resize(target_pil.size, Image.Resampling.LANCZOS)
-
-        return cv2.cvtColor(np.array(result_pil), cv2.COLOR_RGB2BGR)
+        return cv2.cvtColor(np.array(Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))), cv2.COLOR_RGB2BGR)
     except Exception as e:
         logger.error(f"Face swap failed: {e}")
         raise
@@ -128,26 +116,32 @@ def swap_faces(source_img, target_img):
 @app.post("/swap-face/")
 async def swap_face(source_file: UploadFile = File(...), target_file: UploadFile = File(...)):
     try:
-        # Read source image
-        source_content = await source_file.read()
-        source_np = np.frombuffer(source_content, np.uint8)
-        source_img = cv2.imdecode(source_np, cv2.IMREAD_COLOR)
-        if source_img is None:
-            raise ValueError("Invalid source image.")
-
-        # Read target image
-        target_content = await target_file.read()
-        target_np = np.frombuffer(target_content, np.uint8)
-        target_img = cv2.imdecode(target_np, cv2.IMREAD_COLOR)
-        if target_img is None:
-            raise ValueError("Invalid target image.")
-
-        # Perform face swap
+        source_path = "temp_source.jpg"
+        target_path = "temp_target.jpg"
+        output_path = "output.jpg"
+
+        with open(source_path, "wb") as f:
+            f.write(await source_file.read())
+        with open(target_path, "wb") as f:
+            f.write(await target_file.read())
+
+        source_img = cv2.imread(source_path)
+        target_img = cv2.imread(target_path)
+
+        if source_img is None or target_img is None:
+            raise ValueError("Invalid images provided.")
+
         result_img = swap_faces(source_img, target_img)
 
-        # Convert result to bytes
-        _, img_encoded = cv2.imencode(".jpg", result_img)
-        return Response(content=img_encoded.tobytes(), media_type="image/jpeg")
+        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")
 
     except Exception as e:
         logger.error("Error in swap_face: %s", str(e))