app.py revised with endpoints

app.py

@@ -1,14 +1,20 @@
 import os
 os.environ["OMP_NUM_THREADS"] = "1"
-import gradio as gr
+
 import cv2
 import shutil
 import uuid
+import subprocess
+import numpy as np
 import insightface
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
-import subprocess
-
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import StreamingResponse, JSONResponse
+import io
+import threading
+import gradio as gr
+import uvicorn
 
 # -------------------------------------------------
 # Paths
@@ -50,8 +56,8 @@ for f in buffalo_files:
 # -------------------------------------------------
 # Initialize face analysis and swapper
 # -------------------------------------------------
-app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=['CPUExecutionProvider'])
-app.prepare(ctx_id=0, det_size=(640, 640))
+app_face = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=['CPUExecutionProvider'])
+app_face.prepare(ctx_id=0, det_size=(640, 640))
 swapper = insightface.model_zoo.get_model(inswapper_path, providers=['CPUExecutionProvider'])
 
 # -------------------------------------------------
@@ -72,74 +78,106 @@ ensure_codeformer()
 # -------------------------------------------------
 # Pipeline Function
 # -------------------------------------------------
-# -------------------------------------------------
-# Pipeline Function (defaults applied)
-# -------------------------------------------------
 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)
 
-        src_faces = app.get(src_bgr)
-        tgt_faces = app.get(tgt_bgr)
+        src_faces = app_face.get(src_bgr)
+        tgt_faces = app_face.get(tgt_bgr)
         if not src_faces or not tgt_faces:
-            return None, None, "❌ Face not detected in one of the images."
+            return None, "❌ Face not detected in one of the images."
 
         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)
 
-        # ---------------- Save swapped face ----------------
+        # 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 ----------------
+        # 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}"
+            return 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."
+            return 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."
+            return 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)
 
-        return final_img, final_path, ""
+        return final_img, None
 
     except Exception as e:
-        return None, None, f"❌ Error: {str(e)}"
+        return None, f"❌ Error: {str(e)}"
+
+# -------------------------------------------------
+# FastAPI REST API
+# -------------------------------------------------
+api = FastAPI()
 
+@api.post("/api/faceswap")
+async def faceswap(src_img: UploadFile = File(...), tgt_img: UploadFile = File(...)):
+    try:
+        src_bytes = await src_img.read()
+        tgt_bytes = await tgt_img.read()
+
+        src = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
+        tgt = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
+
+        if src is None or tgt is None:
+            return JSONResponse({"error": "Invalid images"}, status_code=400)
+
+        img_rgb_src = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
+        img_rgb_tgt = cv2.cvtColor(tgt, cv2.COLOR_BGR2RGB)
+
+        result_img, err = face_swap_and_enhance(img_rgb_src, img_rgb_tgt)
+        if result_img is None:
+            return JSONResponse({"error": err}, status_code=500)
+
+        _, buffer = cv2.imencode(".png", cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR))
+        return StreamingResponse(io.BytesIO(buffer.tobytes()), media_type="image/png")
+
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
 
 # -------------------------------------------------
-# Gradio Interface (simplified)
+# Gradio UI (kept for demo)
 # -------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("Face Swap")
+    gr.Markdown("Face Swap Demo")
 
     with gr.Row():
         src_input = gr.Image(type="numpy", label="Upload Your Face")
         tgt_input = gr.Image(type="numpy", label="Upload Target Image")
 
     btn = gr.Button("Swap Face")
-
     output_img = gr.Image(type="numpy", label="Enhanced Output")
-    download = gr.File(label="⬇️ Download Enhanced Image")
     error_box = gr.Textbox(label="Logs / Errors", interactive=False)
 
     def process(src, tgt):
-        img, path, err = face_swap_and_enhance(src, tgt)
-        return img, path, err
+        img, err = face_swap_and_enhance(src, tgt)
+        if img is None:
+            return None, err
+        return img, ""
+
+    btn.click(process, [src_input, tgt_input], [output_img, error_box])
 
-    btn.click(process, [src_input, tgt_input], [output_img, download, error_box])
+# -------------------------------------------------
+# Run both FastAPI + Gradio
+# -------------------------------------------------
+def run_api():
+    uvicorn.run(api, host="0.0.0.0", port=7861)
 
-demo.launch()
+threading.Thread(target=run_api).start()
+demo.launch(server_port=7860, server_name="0.0.0.0")

requirements.txt

@@ -32,5 +32,9 @@ onnx
 gradio
 huggingface_hub
 
+# FastAPI REST API + server
+fastapi
+uvicorn
+
 # Optional but included in CodeFormer repo
 tb-nightly