concurrency increased

app.py

@@ -1,6 +1,7 @@
 import os
-import shutil
 import uuid
+import shutil
+import glob
 import cv2
 import numpy as np
 import threading
@@ -127,59 +128,73 @@ async def shutdown_db():
         logger.info("MongoDB connection closed")
 
 # -------------------------------------------------
-# Lock for face swap
+# Multi-user safe face swap
 # -------------------------------------------------
-swap_lock = threading.Lock()
+MAX_CONCURRENT_SWAPS = 2
+swap_semaphore = threading.Semaphore(MAX_CONCURRENT_SWAPS)
+
+def cleanup_old_dirs(base_dir, age_seconds=3600):
+    for folder in glob.glob(os.path.join(base_dir, "*")):
+        if os.path.isdir(folder):
+            try:
+                shutil.rmtree(folder, ignore_errors=True)
+            except Exception:
+                pass
+
+def face_swap_and_enhance_multiuser(src_img, tgt_img):
+    session_id = uuid.uuid4().hex
+    upload_dir = os.path.join(UPLOAD_DIR, session_id)
+    result_dir = os.path.join(RESULT_DIR, session_id)
+    os.makedirs(upload_dir, exist_ok=True)
+    os.makedirs(result_dir, exist_ok=True)
+
+    # Try acquiring GPU slot
+    acquired = swap_semaphore.acquire(blocking=False)
+    if not acquired:
+        return None, None, "⚠️ GPU is busy. Please try again in a few seconds."
 
-# -------------------------------------------------
-# Pipeline Function
-# -------------------------------------------------
-def face_swap_and_enhance(src_img, tgt_img):
     try:
-        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_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
+        tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
 
-            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 in one of the images"
 
-            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 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])
+        if swapped_bgr is None:
+            return None, None, "❌ Face swap failed"
+        cv2.imwrite(swapped_path, swapped_bgr)
 
-            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])
-            if swapped_bgr is None:
-                return None, None, "❌ Face swap failed"
+        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}"
 
-            cv2.imwrite(swapped_path, swapped_bgr)
+        final_files = glob.glob(os.path.join(result_dir, "final_results", "*.png"))
+        if not final_files:
+            return None, None, "❌ No enhanced image found"
 
-            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}"
+        final_path = final_files[0]
+        final_img = cv2.cvtColor(cv2.imread(final_path), cv2.COLOR_BGR2RGB)
 
-            final_results_dir = os.path.join(RESULT_DIR, "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"
-
-            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)}"
 
+    finally:
+        swap_semaphore.release()
+        cleanup_old_dirs(UPLOAD_DIR)
+        cleanup_old_dirs(RESULT_DIR)
+
 # -------------------------------------------------
 # Gradio UI
 # -------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("Face Swap")
+    gr.Markdown("## Face Swap & Enhancement")
 
     with gr.Row():
         src_input = gr.Image(type="numpy", label="Upload Your Face")
@@ -191,7 +206,7 @@ with gr.Blocks() as demo:
     error_box = gr.Textbox(label="Logs / Errors", interactive=False)
 
     def process(src, tgt):
-        img, path, err = face_swap_and_enhance(src, tgt)
+        img, path, err = face_swap_and_enhance_multiuser(src, tgt)
         return img, path, err
 
     btn.click(process, [src_input, tgt_input], [output_img, download, error_box])
@@ -254,7 +269,7 @@ async def perform_faceswap(request: FaceSwapRequest):
     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)
+    final_img, final_path, err = face_swap_and_enhance_multiuser(src_rgb, tgt_rgb)
     if err:
         raise HTTPException(status_code=500, detail=err)