gradio+fastapi

app.py

@@ -8,6 +8,15 @@ import insightface
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
 import subprocess
+import tempfile
+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 gradio import mount_gradio_app
 
 # -------------------------------------------------
 # Paths
@@ -69,56 +78,71 @@ def ensure_codeformer():
 ensure_codeformer()
 
 # -------------------------------------------------
-# Pipeline Function
+# MongoDB setup
 # -------------------------------------------------
+MONGODB_URL = os.getenv(
+    "MONGODB_URL",
+    "mongodb+srv://harilogicgo_db_user:[email protected]/?retryWrites=true&w=majority&appName=Cluster0"
+)
+client = AsyncIOMotorClient(MONGODB_URL)
+database = client.FaceSwap
+target_images_collection = database.get_collection("Target_Images")
+source_images_collection = database.get_collection("Source_Images")
+results_collection = database.get_collection("Results")
+
+# -------------------------------------------------
+# Lock for face swap to prevent concurrent file operations
 # -------------------------------------------------
-# Pipeline Function (defaults applied)
+swap_lock = threading.Lock()
+
+# -------------------------------------------------
+# Pipeline Function (with lock for safety)
 # -------------------------------------------------
 def face_swap_and_enhance(src_img, tgt_img):
     try:
-        src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
-        tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
+        with swap_lock:
+            shutil.rmtree(UPLOAD_DIR, ignore_errors=True)
+            shutil.rmtree(RESULT_DIR, ignore_errors=True)
+            os.makedirs(UPLOAD_DIR, exist_ok=True)
+            os.makedirs(RESULT_DIR, exist_ok=True)
 
-        src_faces = app.get(src_bgr)
-        tgt_faces = app.get(tgt_bgr)
-        if not src_faces or not tgt_faces:
-            return None, None, "❌ Face not detected in one of the images."
+            src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
+            tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
 
-        shutil.rmtree(UPLOAD_DIR, ignore_errors=True)
-        shutil.rmtree(RESULT_DIR, ignore_errors=True)
-        os.makedirs(UPLOAD_DIR, exist_ok=True)
-        os.makedirs(RESULT_DIR, exist_ok=True)
+            src_faces = app.get(src_bgr)
+            tgt_faces = app.get(tgt_bgr)
+            if not src_faces or not tgt_faces:
+                return None, None, "❌ Face not detected in one of the images."
 
-        # ---------------- Save swapped face ----------------
-        unique_name = f"swapped_{uuid.uuid4().hex[:8]}.jpg"
-        swapped_path = os.path.join(UPLOAD_DIR, unique_name)
-        swapped_bgr = swapper.get(tgt_bgr, tgt_faces[0], src_faces[0])
-        cv2.imwrite(swapped_path, swapped_bgr)
+            # ---------------- Save swapped face ----------------
+            unique_name = f"swapped_{uuid.uuid4().hex[:8]}.jpg"
+            swapped_path = os.path.join(UPLOAD_DIR, unique_name)
+            swapped_bgr = swapper.get(tgt_bgr, tgt_faces[0], src_faces[0])
+            cv2.imwrite(swapped_path, swapped_bgr)
 
-        # ---------------- Run CodeFormer ----------------
-        cmd = f"python {CODEFORMER_PATH} -w 0.7 --input_path {swapped_path} --output_path {RESULT_DIR} --bg_upsampler realesrgan --face_upsample"
-        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
-        if result.returncode != 0:
-            return None, None, f"❌ CodeFormer failed:\n{result.stderr}"
+            # ---------------- Run CodeFormer ----------------
+            cmd = f"python {CODEFORMER_PATH} -w 0.7 --input_path {swapped_path} --output_path {RESULT_DIR} --bg_upsampler realesrgan --face_upsample"
+            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+            if result.returncode != 0:
+                return None, None, f"❌ CodeFormer failed:\n{result.stderr}"
 
-        # ---------------- Locate final result ----------------
-        final_results_dir = os.path.join(RESULT_DIR, "final_results")
-        if not os.path.exists(final_results_dir):
-            return None, None, "❌ CodeFormer did not produce final results."
+            # ---------------- Locate final result ----------------
+            final_results_dir = os.path.join(RESULT_DIR, "final_results")
+            if not os.path.exists(final_results_dir):
+                return None, None, "❌ CodeFormer did not produce final results."
 
-        final_files = [f for f in os.listdir(final_results_dir) if f.endswith(".png")]
-        if not final_files:
-            return None, None, "❌ No enhanced image found in final results."
+            final_files = [f for f in os.listdir(final_results_dir) if f.endswith(".png")]
+            if not final_files:
+                return None, None, "❌ No enhanced image found in final results."
 
-        final_path = os.path.join(final_results_dir, final_files[0])
-        final_img = cv2.cvtColor(cv2.imread(final_path), cv2.COLOR_BGR2RGB)
+            final_path = os.path.join(final_results_dir, final_files[0])
+            final_img = cv2.cvtColor(cv2.imread(final_path), cv2.COLOR_BGR2RGB)
 
-        return final_img, final_path, ""
+            return final_img, final_path, ""
 
     except Exception as e:
         return None, None, f"❌ Error: {str(e)}"
 
-
 # -------------------------------------------------
 # Gradio Interface (simplified)
 # -------------------------------------------------
@@ -141,4 +165,102 @@ with gr.Blocks() as demo:
 
     btn.click(process, [src_input, tgt_input], [output_img, download, error_box])
 
-demo.launch(share=True)
+# -------------------------------------------------
+# FastAPI App
+# -------------------------------------------------
+app = FastAPI()
+
+# Redirect root to Gradio for backward compatibility
+@app.get("/")
+def root():
+    return RedirectResponse("/gradio")
+
+# Health check
+@app.get("/health")
+def health():
+    return {"status": "healthy"}
+
+# Upload source image
+@app.post("/source")
+async def upload_source(image: UploadFile = File(...)):
+    contents = await image.read()
+    doc = {
+        "filename": image.filename,
+        "content_type": image.content_type,
+        "data": contents
+    }
+    result = await source_images_collection.insert_one(doc)
+    return {"source_id": str(result.inserted_id)}
+
+# Upload target image
+@app.post("/target")
+async def upload_target(image: UploadFile = File(...)):
+    contents = await image.read()
+    doc = {
+        "filename": image.filename,
+        "content_type": image.content_type,
+        "data": contents
+    }
+    result = await target_images_collection.insert_one(doc)
+    return {"target_id": str(result.inserted_id)}
+
+# Face swap request model
+class FaceSwapRequest(BaseModel):
+    source_id: str
+    target_id: str
+
+# Perform face swap
+@app.post("/faceswap")
+async def perform_faceswap(request: FaceSwapRequest):
+    source_doc = await source_images_collection.find_one({"_id": ObjectId(request.source_id)})
+    if not source_doc:
+        raise HTTPException(status_code=404, detail="Source image not found")
+
+    target_doc = await target_images_collection.find_one({"_id": ObjectId(request.target_id)})
+    if not target_doc:
+        raise HTTPException(status_code=404, detail="Target image not found")
+
+    # Decode source
+    source_array = np.frombuffer(source_doc["data"], np.uint8)
+    source_bgr = cv2.imdecode(source_array, cv2.IMREAD_COLOR)
+    source_rgb = cv2.cvtColor(source_bgr, cv2.COLOR_BGR2RGB)
+
+    # Decode target
+    target_array = np.frombuffer(target_doc["data"], np.uint8)
+    target_bgr = cv2.imdecode(target_array, cv2.IMREAD_COLOR)
+    target_rgb = cv2.cvtColor(target_bgr, cv2.COLOR_BGR2RGB)
+
+    # Perform swap
+    final_img, final_path, err = face_swap_and_enhance(source_rgb, target_rgb)
+    if err:
+        raise HTTPException(status_code=500, detail=err)
+
+    # Read final bytes
+    with open(final_path, "rb") as f:
+        final_bytes = f.read()
+
+    # Store result
+    result_doc = {
+        "source_id": request.source_id,
+        "target_id": request.target_id,
+        "filename": "enhanced.png",
+        "content_type": "image/png",
+        "data": final_bytes
+    }
+    result = await results_collection.insert_one(result_doc)
+    return {"result_id": str(result.inserted_id)}
+
+# Download result
+@app.get("/download/{result_id}")
+async def download_result(result_id: str):
+    doc = await results_collection.find_one({"_id": ObjectId(result_id)})
+    if not doc:
+        raise HTTPException(status_code=404, detail="Result not found")
+    return Response(
+        content=doc["data"],
+        media_type=doc["content_type"],
+        headers={"Content-Disposition": f"attachment; filename={doc['filename']}"}
+    )
+
+# Mount Gradio at /gradio
+app = mount_gradio_app(app, demo, path="/gradio")

requirements.txt

@@ -25,14 +25,15 @@ facexlib==0.2.5
 
 # Face swap / insightface
 insightface
-onnxruntime
 onnx
+onnxruntime-gpu  # Changed from onnxruntime to support GPU
 
 # App deps
 gradio
 huggingface_hub
 fastapi
 uvicorn
+motor  # Added for MongoDB async support
 
 # Optional but included in CodeFormer repo
-tb-nightly
+tb-nightly