Update app.py

app.py

@@ -10,6 +10,7 @@ import shutil
 import logging
 import requests
 from pathlib import Path
+import uvicorn
 
 app = FastAPI()
 
@@ -26,7 +27,15 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Global flag to prevent multiple downloads
+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"
@@ -35,61 +44,83 @@ def download_model():
         logger.info("Model not found. Downloading inswapper_128.onnx...")
         model_dir.mkdir(exist_ok=True)
         try:
-            response = requests.get(model_url, stream=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.")
+            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.")
-
-# Download model on startup
-download_model()
+    else:
+        logger.info("Model already exists at: %s", model_path)
+        MODEL_DOWNLOADED = True
+
+@app.on_event("startup")
+async def startup_event():
+    """Run startup tasks like downloading the model."""
+    logger.info("Starting up application...")
+    try:
+        download_model()
+        logger.info("Startup completed successfully.")
+    except Exception as e:
+        logger.error(f"Startup failed: {e}")
+        raise
 
 def get_many_faces(image):
-    """Simplified face detection using insightface (placeholder)."""
-    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 []
+    """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."""
-    from insightface.utils import face_align
-    from insightface.model_zoo import face_swapper
+    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=0, det_size=(640, 640))
+        # Initialize face analysis
+        face_analyzer = FaceAnalysis(name="buffalo_l")
+        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)
+        # 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.")
-    if len(source_faces) > 1 or len(target_faces) > 1:
-        raise ValueError("Multiple faces detected; only one face per image is supported.")
+        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_face = source_faces[0]
+        target_face = target_faces[0]
 
-    # Load the face swapper model
-    model_path = Path("models/inswapper_128.onnx")
-    swapper = face_swapper.FaceSwapper(model_path)
+        # 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(model_path)
 
-    # Perform face swap
-    result = swapper.get(target_img, target_face, source_face, paste_back=True)
+        # 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))
-    result_pil = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
-    result_pil = result_pil.resize(target_pil.size, Image.Resampling.LANCZOS)
+        # Resize to match target image size
+        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)
 
-    return cv2.cvtColor(np.array(result_pil), cv2.COLOR_RGB2BGR)
+        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(...), doFaceEnhancer: bool = True):
@@ -135,4 +166,8 @@ async def swap_face(source_file: UploadFile = File(...), target_file: UploadFile
 
     except Exception as e:
         logger.error("Error in swap_face: %s", str(e))
-        raise HTTPException(status_code=500, detail=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)