Update app.py

app.py

@@ -7,18 +7,17 @@ import cv2
 import numpy as np
 from PIL import Image
 import os
-import shutil
 import logging
 import requests
 from pathlib import Path
 import uvicorn
 
-# Initialize FastAPI with explicit docs settings
+# Initialize FastAPI
 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
+    docs_url="/docs",  
+    redoc_url="/redoc"
 )
 
 # Configure logging
@@ -28,37 +27,33 @@ logger = logging.getLogger(__name__)
 # Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Update with your Framer domain in production
+    allow_origins=["*"],  
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Add a root endpoint to confirm the app is running
+# Root endpoint
 @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."}
 
-# Add a health check endpoint
+# Health check
 @app.get("/health")
 async def health_check():
     return {"status": "healthy"}
 
-# Global flag to prevent multiple downloads
+# Global flag for model download
 MODEL_DOWNLOADED = False
 
 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...")
+        logger.info("Downloading inswapper_128.onnx...")
         model_dir.mkdir(exist_ok=True)
         try:
             response = requests.get(model_url, stream=True, timeout=30)
@@ -70,58 +65,43 @@ def download_model():
             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.")
+            raise RuntimeError("Could not download inswapper_128.onnx.")
     else:
-        logger.info("Model already exists at: %s", model_path)
+        logger.info("Model already exists.")
         MODEL_DOWNLOADED = True
 
-# Use lifespan event handler
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup code
-    logger.info("Starting up application...")
+    logger.info("Starting application...")
     try:
         download_model()
-        logger.info("Startup completed successfully.")
+        logger.info("Application started successfully.")
     except Exception as e:
         logger.error(f"Startup failed: {e}")
         raise
     yield
-    # Shutdown code (if any)
     logger.info("Shutting down application...")
 
 app.lifespan = lifespan
 
-def get_many_faces(image):
-    """Simplified face detection using insightface."""
-    try:
-        from insightface.app import FaceAnalysis
-        app = FaceAnalysis(name="buffalo_l")
-        app.prepare(ctx_id=0, det_size=(640, 640))
-        faces = app.get(image)
-        return faces if faces else []
-    except Exception as e:
-        logger.error(f"Face detection failed: {e}")
-        raise
-
 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
 
         # Initialize face analysis
         face_analyzer = FaceAnalysis(name="buffalo_l")
-        face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
+        face_analyzer.prepare(ctx_id=-1, 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 in one or both images.")
+            raise ValueError("No faces detected.")
         if len(source_faces) > 1 or len(target_faces) > 1:
-            raise ValueError("Multiple faces detected; only one face per image is supported.")
+            raise ValueError("Multiple faces detected; only one per image supported.")
 
         source_face = source_faces[0]
         target_face = target_faces[0]
@@ -130,14 +110,14 @@ def swap_faces(source_img, target_img):
         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)
+        swapper = face_swapper.FaceSwapper(str(model_path))
 
         # Perform face swap
         result = swapper.get(target_img, target_face, source_face, paste_back=True)
 
-        # Resize to match target image size
-        target_pil = Image.fromarray(cv2.cvtColor(target_img, cv2.COLOR_BGR2RGB))
+        # 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)
@@ -146,51 +126,32 @@ def swap_faces(source_img, target_img):
         raise
 
 @app.post("/swap-face/")
-async def swap_face(source_file: UploadFile = File(...), target_file: UploadFile = File(...), doFaceEnhancer: bool = True):
+async def swap_face(source_file: UploadFile = File(...), target_file: UploadFile = File(...)):
     try:
-        # Save uploaded files temporarily
-        source_path = "temp_source.jpg"
-        target_path = "temp_target.jpg"
-        output_path = "output.jpg"
-
-        # Read and save source image
+        # Read source image
         source_content = await source_file.read()
-        with open(source_path, "wb") as f:
-            f.write(source_content)
-        source_img = cv2.imread(source_path)
+        source_np = np.frombuffer(source_content, np.uint8)
+        source_img = cv2.imdecode(source_np, cv2.IMREAD_COLOR)
         if source_img is None:
-            raise ValueError("Failed to load source image.")
+            raise ValueError("Invalid source image.")
 
-        # Read and save target image
+        # Read target image
         target_content = await target_file.read()
-        with open(target_path, "wb") as f:
-            f.write(target_content)
-        target_img = cv2.imread(target_path)
+        target_np = np.frombuffer(target_content, np.uint8)
+        target_img = cv2.imdecode(target_np, cv2.IMREAD_COLOR)
         if target_img is None:
-            raise ValueError("Failed to load target image.")
+            raise ValueError("Invalid target image.")
 
         # Perform face swap
         result_img = swap_faces(source_img, target_img)
 
-        # Save the result
-        cv2.imwrite(output_path, result_img)
-
-        # Read the output image
-        with open(output_path, "rb") as f:
-            image_data = f.read()
-
-        # Clean up temporary files
-        for path in [source_path, target_path, output_path]:
-            if os.path.exists(path):
-                os.remove(path)
-
-        # Return the swapped image
-        return Response(content=image_data, media_type="image/jpeg")
+        # Convert result to bytes
+        _, img_encoded = cv2.imencode(".jpg", result_img)
+        return Response(content=img_encoded.tobytes(), media_type="image/jpeg")
 
     except Exception as e:
         logger.error("Error in swap_face: %s", str(e))
         raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    # Hugging Face Spaces expects the app to run on port 7860
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)