import fixed

app.py

@@ -1,26 +1,25 @@
 import os
-os.environ["OMP_NUM_THREADS"] = "1"
-import gradio as gr
-import cv2
 import shutil
 import uuid
+import cv2
+import numpy as np
+import threading
+import subprocess
+import logging
+
 import insightface
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
-import subprocess
-import numpy as np
-import threading
+
 from fastapi import FastAPI, UploadFile, File, HTTPException, Response
 from fastapi.responses import RedirectResponse
 from pydantic import BaseModel
-from motor.motor_asyncio import AsyncIOMotorClient
-from bson.objectid import ObjectId
-from motor.motor_asyncio import AsyncIOMotorGridFSBucket
+from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorGridFSBucket
+from bson import ObjectId
 
-from gradio import mount_gradio_app
 import uvicorn
-import logging
-import io
+import gradio as gr
+from gradio import mount_gradio_app
 
 # -------------------------------------------------
 # Logging
@@ -75,11 +74,9 @@ inswapper_path = download_models()
 # Face Analysis + Swapper
 # -------------------------------------------------
 providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
-logger.info(f"Initializing FaceAnalysis with providers: {providers}")
 face_analysis_app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=providers)
 face_analysis_app.prepare(ctx_id=0, det_size=(640, 640))
 swapper = insightface.model_zoo.get_model(inswapper_path, providers=providers)
-logger.info("FaceAnalysis and swapper initialized")
 
 # -------------------------------------------------
 # CodeFormer setup
@@ -88,28 +85,46 @@ CODEFORMER_PATH = "CodeFormer/inference_codeformer.py"
 
 def ensure_codeformer():
     if not os.path.exists("CodeFormer"):
-        logger.info("Cloning CodeFormer repository...")
         subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True, check=True)
         subprocess.run("pip install -r CodeFormer/requirements.txt", shell=True, check=True)
         subprocess.run("python CodeFormer/basicsr/setup.py develop", shell=True, check=True)
         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True, check=True)
         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True, check=True)
-        logger.info("CodeFormer setup complete")
 
 ensure_codeformer()
 
 # -------------------------------------------------
-# MongoDB + GridFS
+# MongoDB + GridFS setup
 # -------------------------------------------------
 MONGODB_URL = os.getenv(
     "MONGODB_URL",
     "mongodb+srv://harilogicgo_db_user:logicgoinfotech@cluster0.dcs1tnb.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0"
 )
-client = AsyncIOMotorClient(MONGODB_URL)
-database = client.FaceSwap
-fs_bucket = AsyncIOMotorGridFSBucket(database)
 
-logger.info("MongoDB + GridFS initialized")
+client: AsyncIOMotorClient = None
+database = None
+fs_bucket: AsyncIOMotorGridFSBucket = None
+
+# -------------------------------------------------
+# FastAPI App
+# -------------------------------------------------
+fastapi_app = FastAPI()
+
+@fastapi_app.on_event("startup")
+async def startup_db():
+    global client, database, fs_bucket
+    logger.info("Initializing MongoDB + GridFS...")
+    client = AsyncIOMotorClient(MONGODB_URL)
+    database = client.FaceSwap
+    fs_bucket = AsyncIOMotorGridFSBucket(database)
+    logger.info("MongoDB + GridFS initialized")
+
+@fastapi_app.on_event("shutdown")
+async def shutdown_db():
+    global client
+    if client:
+        client.close()
+        logger.info("MongoDB connection closed")
 
 # -------------------------------------------------
 # Lock for face swap
@@ -117,10 +132,9 @@ logger.info("MongoDB + GridFS initialized")
 swap_lock = threading.Lock()
 
 # -------------------------------------------------
-# Face Swap Pipeline
+# Pipeline Function
 # -------------------------------------------------
 def face_swap_and_enhance(src_img, tgt_img):
-    logger.info("Starting face swap and enhancement")
     try:
         with swap_lock:
             shutil.rmtree(UPLOAD_DIR, ignore_errors=True)
@@ -128,16 +142,13 @@ def face_swap_and_enhance(src_img, tgt_img):
             os.makedirs(UPLOAD_DIR, exist_ok=True)
             os.makedirs(RESULT_DIR, exist_ok=True)
 
-            if not isinstance(src_img, np.ndarray) or not isinstance(tgt_img, np.ndarray):
-                return None, None, "❌ Invalid input images"
-
             src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
             tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
 
             src_faces = face_analysis_app.get(src_bgr)
             tgt_faces = face_analysis_app.get(tgt_bgr)
             if not src_faces or not tgt_faces:
-                return None, None, "❌ Face not detected"
+                return None, None, "❌ Face not detected in one of the images"
 
             swapped_path = os.path.join(UPLOAD_DIR, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
             swapped_bgr = swapper.get(tgt_bgr, tgt_faces[0], src_faces[0])
@@ -165,7 +176,7 @@ def face_swap_and_enhance(src_img, tgt_img):
         return None, None, f"❌ Error: {str(e)}"
 
 # -------------------------------------------------
-# Gradio Interface
+# Gradio UI
 # -------------------------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("Face Swap")
@@ -186,10 +197,8 @@ with gr.Blocks() as demo:
     btn.click(process, [src_input, tgt_input], [output_img, download, error_box])
 
 # -------------------------------------------------
-# FastAPI App
+# API Endpoints
 # -------------------------------------------------
-fastapi_app = FastAPI()
-
 @fastapi_app.get("/")
 def root():
     return RedirectResponse("/gradio")
@@ -198,20 +207,18 @@ def root():
 async def health():
     return {"status": "healthy"}
 
-# -------- Upload Endpoints with GridFS --------
 @fastapi_app.post("/source")
 async def upload_source(image: UploadFile = File(...)):
     contents = await image.read()
-    file_id = await fs_bucket.upload_from_stream(image.filename, contents)
+    file_id = await fs_bucket.upload_from_stream(image.filename, contents, metadata={"type": "source"})
     return {"source_id": str(file_id)}
 
 @fastapi_app.post("/target")
 async def upload_target(image: UploadFile = File(...)):
     contents = await image.read()
-    file_id = await fs_bucket.upload_from_stream(image.filename, contents)
+    file_id = await fs_bucket.upload_from_stream(image.filename, contents, metadata={"type": "target"})
     return {"target_id": str(file_id)}
 
-# -------- Faceswap Endpoint --------
 class FaceSwapRequest(BaseModel):
     source_id: str
     target_id: str
@@ -219,50 +226,51 @@ class FaceSwapRequest(BaseModel):
 @fastapi_app.post("/faceswap")
 async def perform_faceswap(request: FaceSwapRequest):
     try:
-        # Read source
-        source_stream = await fs_bucket.open_download_stream(ObjectId(request.source_id))
-        source_bytes = await source_stream.read()
-        source_array = np.frombuffer(source_bytes, np.uint8)
-        source_bgr = cv2.imdecode(source_array, cv2.IMREAD_COLOR)
-        source_rgb = cv2.cvtColor(source_bgr, cv2.COLOR_BGR2RGB)
-
-        # Read target
-        target_stream = await fs_bucket.open_download_stream(ObjectId(request.target_id))
-        target_bytes = await target_stream.read()
-        target_array = np.frombuffer(target_bytes, np.uint8)
-        target_bgr = cv2.imdecode(target_array, cv2.IMREAD_COLOR)
-        target_rgb = cv2.cvtColor(target_bgr, cv2.COLOR_BGR2RGB)
-
-        # Run pipeline
-        final_img, final_path, err = face_swap_and_enhance(source_rgb, target_rgb)
-        if err:
-            raise HTTPException(status_code=500, detail=err)
-
-        # Store result in GridFS
-        with open(final_path, "rb") as f:
-            final_bytes = f.read()
-        result_id = await fs_bucket.upload_from_stream("enhanced.png", final_bytes)
-
-        return {"result_id": str(result_id)}
+        src_stream = await fs_bucket.open_download_stream(ObjectId(request.source_id))
+        src_bytes = await src_stream.read()
+        tgt_stream = await fs_bucket.open_download_stream(ObjectId(request.target_id))
+        tgt_bytes = await tgt_stream.read()
+    except Exception:
+        raise HTTPException(status_code=404, detail="Source or Target not found")
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    src_array = np.frombuffer(src_bytes, np.uint8)
+    src_bgr = cv2.imdecode(src_array, cv2.IMREAD_COLOR)
+    tgt_array = np.frombuffer(tgt_bytes, np.uint8)
+    tgt_bgr = cv2.imdecode(tgt_array, cv2.IMREAD_COLOR)
+
+    if src_bgr is None or tgt_bgr is None:
+        raise HTTPException(status_code=400, detail="Invalid image data")
+
+    src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
+    tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
+
+    final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
+    if err:
+        raise HTTPException(status_code=500, detail=err)
+
+    with open(final_path, "rb") as f:
+        final_bytes = f.read()
+
+    result_id = await fs_bucket.upload_from_stream("enhanced.png", final_bytes, metadata={"type": "result"})
+    return {"result_id": str(result_id)}
 
-# -------- Download Endpoint --------
 @fastapi_app.get("/download/{result_id}")
 async def download_result(result_id: str):
     try:
         stream = await fs_bucket.open_download_stream(ObjectId(result_id))
-        file_data = await stream.read()
-        return Response(
-            content=file_data,
-            media_type="image/png",
-            headers={"Content-Disposition": f"attachment; filename=enhanced.png"}
-        )
+        data = await stream.read()
     except Exception:
         raise HTTPException(status_code=404, detail="Result not found")
 
+    return Response(
+        content=data,
+        media_type="image/png",
+        headers={"Content-Disposition": "attachment; filename=result.png"}
+    )
+
+# -------------------------------------------------
 # Mount Gradio
+# -------------------------------------------------
 fastapi_app = mount_gradio_app(fastapi_app, demo, path="/gradio")
 
 if __name__ == "__main__":