revised

app.py

@@ -8,7 +8,6 @@ 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
@@ -17,11 +16,12 @@ from pydantic import BaseModel
 from motor.motor_asyncio import AsyncIOMotorClient
 from bson.objectid import ObjectId
 from gradio import mount_gradio_app
+import uvicorn
 
 # -------------------------------------------------
 # Paths
 # -------------------------------------------------
-REPO_ID = "HariLogicgo/face_swap_models"  # <- your HF repo for models
+REPO_ID = "HariLogicgo/face_swap_models"
 BASE_DIR = "./workspace"
 UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
 RESULT_DIR = os.path.join(BASE_DIR, "results")
@@ -29,38 +29,43 @@ MODELS_DIR = "./models"
 
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(RESULT_DIR, exist_ok=True)
+os.makedirs(MODELS_DIR, exist_ok=True)
 
 # -------------------------------------------------
 # Download models once
 # -------------------------------------------------
-inswapper_path = hf_hub_download(
-    repo_id=REPO_ID,
-    filename="models/inswapper_128.onnx",
-    repo_type="model",
-    local_dir=MODELS_DIR
-)
-
-buffalo_files = [
-    "1k3d68.onnx",
-    "2d106det.onnx",
-    "genderage.onnx",
-    "det_10g.onnx",
-    "w600k_r50.onnx"
-]
-for f in buffalo_files:
-    hf_hub_download(
+def download_models():
+    inswapper_path = hf_hub_download(
         repo_id=REPO_ID,
-        filename=f"models/buffalo_l/{f}",
+        filename="models/inswapper_128.onnx",
         repo_type="model",
         local_dir=MODELS_DIR
     )
+    buffalo_files = [
+        "1k3d68.onnx",
+        "2d106det.onnx",
+        "genderage.onnx",
+        "det_10g.onnx",
+        "w600k_r50.onnx"
+    ]
+    for f in buffalo_files:
+        hf_hub_download(
+            repo_id=REPO_ID,
+            filename=f"models/buffalo_l/{f}",
+            repo_type="model",
+            local_dir=MODELS_DIR
+        )
+    return inswapper_path
+
+inswapper_path = download_models()
 
 # -------------------------------------------------
 # Initialize face analysis and swapper
 # -------------------------------------------------
-app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=['CPUExecutionProvider'])
+providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']  # Prefer CUDA
+app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=providers)
 app.prepare(ctx_id=0, det_size=(640, 640))
-swapper = insightface.model_zoo.get_model(inswapper_path, providers=['CPUExecutionProvider'])
+swapper = insightface.model_zoo.get_model(inswapper_path, providers=providers)
 
 # -------------------------------------------------
 # CodeFormer setup
@@ -69,11 +74,11 @@ CODEFORMER_PATH = "CodeFormer/inference_codeformer.py"
 
 def ensure_codeformer():
     if not os.path.exists("CodeFormer"):
-        subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True)
-        subprocess.run("pip install -r CodeFormer/requirements.txt", shell=True)
-        subprocess.run("python CodeFormer/basicsr/setup.py develop", shell=True)
-        subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True)
-        subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True)
+        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)
 
 ensure_codeformer()
 
@@ -91,12 +96,12 @@ source_images_collection = database.get_collection("Source_Images")
 results_collection = database.get_collection("Results")
 
 # -------------------------------------------------
-# Lock for face swap to prevent concurrent file operations
+# Lock for face swap
 # -------------------------------------------------
 swap_lock = threading.Lock()
 
 # -------------------------------------------------
-# Pipeline Function (with lock for safety)
+# Pipeline Function
 # -------------------------------------------------
 def face_swap_and_enhance(src_img, tgt_img):
     try:
@@ -114,19 +119,16 @@ def face_swap_and_enhance(src_img, tgt_img):
             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)
 
-            # ---------------- 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."
@@ -144,7 +146,7 @@ def face_swap_and_enhance(src_img, tgt_img):
         return None, None, f"❌ Error: {str(e)}"
 
 # -------------------------------------------------
-# Gradio Interface (simplified)
+# Gradio Interface
 # -------------------------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("Face Swap")
@@ -170,17 +172,14 @@ with gr.Blocks() as demo:
 # -------------------------------------------------
 app = FastAPI()
 
-# Redirect root to Gradio for backward compatibility
 @app.get("/")
 def root():
     return RedirectResponse("/gradio")
 
-# Health check
 @app.get("/health")
-def health():
+async def health():
     return {"status": "healthy"}
 
-# Upload source image
 @app.post("/source")
 async def upload_source(image: UploadFile = File(...)):
     contents = await image.read()
@@ -192,7 +191,6 @@ async def upload_source(image: UploadFile = File(...)):
     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()
@@ -204,12 +202,10 @@ async def upload_target(image: UploadFile = File(...)):
     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)})
@@ -220,26 +216,21 @@ async def perform_faceswap(request: FaceSwapRequest):
     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,
@@ -250,7 +241,6 @@ async def perform_faceswap(request: FaceSwapRequest):
     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)})
@@ -263,4 +253,8 @@ async def download_result(result_id: str):
     )
 
 # Mount Gradio at /gradio
-app = mount_gradio_app(app, demo, path="/gradio")
+app = mount_gradio_app(app, demo, path="/gradio")
+
+# Run the app with Uvicorn
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

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