Update app.py

app.py

@@ -151,6 +151,37 @@ async def log_faceswap_hit(token: str, status: str = "success"):
 # --------------------- Face Swap Pipeline ---------------------
 swap_lock = threading.Lock()
 
+def multi_face_swap(src_img, tgt_img):
+    """
+    Swap multiple faces from source image into target image.
+    Faces are matched left-to-right.
+    """
+    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 len(src_faces) < 1 or len(tgt_faces) < 1:
+        raise ValueError("No faces detected in one of the images")
+
+    # Sort faces left → right
+    src_faces = sorted(src_faces, key=lambda f: f.bbox[0])
+    tgt_faces = sorted(tgt_faces, key=lambda f: f.bbox[0])
+
+    face_count = min(len(src_faces), len(tgt_faces))
+    result_img = tgt_bgr.copy()
+
+    for i in range(face_count):
+        result_img = swapper.get(
+            result_img,
+            tgt_faces[i],
+            src_faces[i],
+            paste_back=True
+        )
+
+    return cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
+
 def face_swap_and_enhance(src_img, tgt_img, temp_dir=None):
     try:
         with swap_lock:
@@ -265,6 +296,29 @@ def download_from_spaces(key):
     obj = client.get_object(Bucket=DO_SPACES_BUCKET, Key=key)
     return obj['Body'].read()
 
+def build_multi_faceswap_gradio():
+    with gr.Blocks() as demo:
+        gr.Markdown("## 👩‍❤️‍👨 Multi Face Swap (Couple → Couple)")
+
+        with gr.Row():
+            src = gr.Image(type="numpy", label="Source Image (2 Faces)")
+            tgt = gr.Image(type="numpy", label="Target Image (2 Faces)")
+
+        out = gr.Image(type="numpy", label="Swapped Result")
+        error = gr.Textbox(label="Logs", interactive=False)
+
+        def process(src_img, tgt_img):
+            try:
+                result = multi_face_swap(src_img, tgt_img)
+                return result, ""
+            except Exception as e:
+                return None, str(e)
+
+        btn = gr.Button("Swap Faces")
+        btn.click(process, [src, tgt], [out, error])
+
+    return demo
+
 # --------------------- API Endpoints ---------------------
 @fastapi_app.get("/")
 def root():
@@ -664,10 +718,17 @@ async def preview_result(result_key: str):
     )
 # --------------------- Mount Gradio ---------------------
 
+multi_faceswap_app = build_multi_faceswap_gradio()
+fastapi_app = mount_gradio_app(
+    fastapi_app,
+    multi_faceswap_app,
+    path="/gradio-multi"
+)
 
 if __name__ == "__main__":
     uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)
 
+
 # # --------------------- List Images Endpoint ---------------------
 # import os
 # os.environ["OMP_NUM_THREADS"] = "1"