Update app.py

app.py

@@ -20,7 +20,6 @@ try:
 except Exception:
     ONNX_AVAILABLE = False
 
-
 # ============================================================
 # GLOBALS
 # ============================================================
@@ -34,139 +33,102 @@ gfpgan_session = None
 TASKS = {}
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=6)
 
-
 # ============================================================
 # Helpers
 # ============================================================
-def download_file(url: str, dest_path: str, timeout: int = 120):
-    try:
-        r = requests.get(url, stream=True, timeout=timeout)
-        r.raise_for_status()
-        with open(dest_path, "wb") as f:
-            for chunk in r.iter_content(8192):
-                if chunk:
-                    f.write(chunk)
-        return True
-    except Exception as e:
-        print("Download failed:", e)
-        if os.path.exists(dest_path):
-            try:
-                os.remove(dest_path)
-            except Exception:
-                pass
-        return False
-
-
 def ensure_swapper_model():
     path = os.path.join(BASE_DIR, "inswapper_128.onnx")
     if not os.path.isfile(path):
         raise RuntimeError(
-            "inswapper_128.onnx not found. "
-            "Place it in the same directory as app.py."
+            "inswapper_128.onnx not found. Place it next to app.py"
         )
     print("Found swapper model:", path)
     return path
 
 
+def download_file(url, dest, timeout=180):
+    try:
+        r = requests.get(url, stream=True, timeout=timeout)
+        r.raise_for_status()
+        with open(dest, "wb") as f:
+            for c in r.iter_content(8192):
+                if c:
+                    f.write(c)
+        return True
+    except Exception as e:
+        print("Download failed:", e)
+        return False
+
 # ============================================================
 # GFPGAN-Lite
 # ============================================================
 GFPGAN_URL = "https://huggingface.co/ai-forever/GFPGAN-Lite/resolve/main/gfpgan_lite.onnx"
-GFPGAN_FILENAME = "gfpgan_lite.onnx"
-
 
 def try_load_gfpgan():
     global gfpgan_session
     if not ONNX_AVAILABLE:
-        print("onnxruntime not available — GFPGAN disabled")
         return
-
-    path = os.path.join(BASE_DIR, GFPGAN_FILENAME)
+    path = os.path.join(BASE_DIR, "gfpgan_lite.onnx")
     if not os.path.exists(path):
         print("Downloading GFPGAN-Lite...")
-        if not download_file(GFPGAN_URL, path, 180):
-            print("GFPGAN download failed")
+        if not download_file(GFPGAN_URL, path):
             return
-
     try:
         gfpgan_session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
-        print("Loaded GFPGAN-Lite:", path)
+        print("GFPGAN-Lite loaded")
     except Exception as e:
         print("GFPGAN load failed:", e)
         gfpgan_session = None
 
-
-def try_load_enhancer():
-    global enhancer_session
-    if not ONNX_AVAILABLE:
-        return
-    path = os.path.join(BASE_DIR, "small_enhancer.onnx")
-    if os.path.exists(path):
-        try:
-            enhancer_session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
-            print("Loaded small enhancer:", path)
-        except Exception as e:
-            enhancer_session = None
-            print("Small enhancer failed:", e)
-
-
 # ============================================================
 # Image helpers
 # ============================================================
-def swap_faces(target_img, target_face, source_face):
-    return swapper.get(target_img, target_face, source_face, paste_back=True)
-
+def swap_faces(target, t_face, s_face):
+    return swapper.get(target, t_face, s_face, paste_back=True)
 
-def clarity_enhance(img):
+def clarity(img):
     blur = cv2.GaussianBlur(img, (0, 0), 1.6)
     return cv2.addWeighted(img, 1.25, blur, -0.25, 0)
 
-
-def resize_keep_aspect_max512(img):
+def resize_max512(img):
     h, w = img.shape[:2]
     if max(h, w) <= 512:
         return img
-    scale = 512 / max(h, w)
-    return cv2.resize(img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_CUBIC)
-
+    s = 512 / max(h, w)
+    return cv2.resize(img, (int(w * s), int(h * s)), interpolation=cv2.INTER_CUBIC)
 
 # ============================================================
 # Worker
 # ============================================================
-def run_task(task_id, src_bytes, tgt_bytes):
-    TASKS[task_id]["status"] = "processing"
+def run_task(tid, src_bytes, tgt_bytes):
+    TASKS[tid]["status"] = "processing"
     try:
         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:
-            raise ValueError("Invalid image")
-
-        src_faces = face_app.get(src)
-        tgt_faces = face_app.get(tgt)
+        s_faces = face_app.get(src)
+        t_faces = face_app.get(tgt)
 
-        if not src_faces or not tgt_faces:
+        if not s_faces or not t_faces:
             raise ValueError("Face not detected")
 
-        swapped = swap_faces(tgt, tgt_faces[0], src_faces[0])
-        final = clarity_enhance(swapped)
-        final = resize_keep_aspect_max512(final)
+        out = swap_faces(tgt, t_faces[0], s_faces[0])
+        out = clarity(out)
+        out = resize_max512(out)
 
-        out_path = f"/tmp/{task_id}.jpg"
-        cv2.imwrite(out_path, final)
+        out_path = f"/tmp/{tid}.jpg"
+        cv2.imwrite(out_path, out)
 
-        TASKS[task_id] = {"status": "done", "result": out_path}
+        TASKS[tid] = {"status": "done", "result": out_path}
 
     except Exception as e:
-        TASKS[task_id] = {"status": "failed", "error": str(e)}
+        TASKS[tid] = {"status": "failed", "error": str(e)}
         print(traceback.format_exc())
 
-
 # ============================================================
-# FastAPI + FULL UI
+# FastAPI
 # ============================================================
-app = FastAPI(title="FaceSwap (GFPGAN-Lite)")
-
+app = FastAPI(title="FaceSwap")
 
 @app.get("/", response_class=HTMLResponse)
 def home():
@@ -181,81 +143,119 @@ def home():
 body{background:#0e1525;color:#fff;font-family:Arial;margin:0;padding:20px}
 .container{max-width:1100px;margin:auto}
 .flex{display:flex;gap:16px;flex-wrap:wrap}
-.zone{flex:1;min-width:280px;border:2px dashed #ffffff33;padding:18px;text-align:center;border-radius:10px;cursor:pointer}
-.preview img{max-width:100%;border-radius:10px}
+.box{flex:1;min-width:280px;background:#ffffff10;padding:12px;border-radius:10px;text-align:center}
+.box img{max-width:100%;border-radius:10px}
+.zone{border:2px dashed #ffffff33;padding:18px;border-radius:10px;cursor:pointer}
 button{padding:10px 18px;border-radius:10px;border:none;background:#6c5ce7;color:white;font-weight:600}
+a.download{display:inline-block;margin-top:10px;color:#00cec9;text-decoration:none}
 </style>
 </head>
 <body>
 <div class="container">
 <h2>FaceSwap</h2>
+
 <div class="flex">
-<label class="zone">
-<input type="file" id="src" hidden>
-Source Image
-</label>
-<label class="zone">
-<input type="file" id="tgt" hidden>
-Target Image
-</label>
+  <label class="zone">
+    <input type="file" id="src" hidden>
+    Select SOURCE
+  </label>
+  <label class="zone">
+    <input type="file" id="tgt" hidden>
+    Select TARGET
+  </label>
 </div>
-<button onclick="start()">Start</button>
-<img id="out">
+
+<div class="flex" style="margin-top:20px">
+  <div class="box">
+    <div>Source</div>
+    <img id="pSrc">
+  </div>
+  <div class="box">
+    <div>Target</div>
+    <img id="pTgt">
+  </div>
+  <div class="box">
+    <div>Output</div>
+    <img id="pOut">
+    <a id="dl" class="download" download="faceswap.jpg" style="display:none">Download</a>
+  </div>
+</div>
+
+<div style="margin-top:20px">
+  <button id="go">Start</button>
+  <span id="st" style="margin-left:10px"></span>
 </div>
+</div>
+
 <script>
-let src,tgt;
-document.getElementById("src").onchange=e=>src=e.target.files[0];
-document.getElementById("tgt").onchange=e=>tgt=e.target.files[0];
-async function start(){
- if(!src||!tgt){alert("Select images");return;}
- let fd=new FormData();
- fd.append("source",src);
- fd.append("target",tgt);
- let r=await fetch("/swap-image",{method:"POST",body:fd});
- let j=await r.json();
- let id=j.task_id;
- let t=setInterval(async()=>{
-   let s=await fetch("/task-status/"+id).then(r=>r.json());
-   if(s.status==="done"){
-     clearInterval(t);
-     document.getElementById("out").src="/task-result/"+id;
-   }
- },1000);
+const s=document.getElementById("src"),
+      t=document.getElementById("tgt"),
+      ps=document.getElementById("pSrc"),
+      pt=document.getElementById("pTgt"),
+      po=document.getElementById("pOut"),
+      dl=document.getElementById("dl"),
+      st=document.getElementById("st");
+
+s.onchange=()=>ps.src=URL.createObjectURL(s.files[0]);
+t.onchange=()=>pt.src=URL.createObjectURL(t.files[0]);
+
+document.getElementById("go").onclick=async()=>{
+  if(!s.files[0]||!t.files[0]) return alert("Select both images");
+  st.innerText="Processing...";
+  let fd=new FormData();
+  fd.append("source",s.files[0]);
+  fd.append("target",t.files[0]);
+  const r=await fetch("/swap-image",{method:"POST",body:fd});
+  const j=await r.json();
+  poll(j.task_id);
+};
+
+async function poll(id){
+  const r=await fetch("/task-status/"+id);
+  const j=await r.json();
+  if(j.status==="done"){
+    const img=await fetch("/task-result/"+id);
+    const b=await img.blob();
+    const u=URL.createObjectURL(b);
+    po.src=u;
+    dl.href=u;
+    dl.style.display="inline-block";
+    st.innerText="Done";
+  } else if(j.status==="failed"){
+    st.innerText="Failed: "+j.error;
+  } else {
+    setTimeout(()=>poll(id),1000);
+  }
 }
 </script>
 </body>
 </html>
 """
 
-
 @app.post("/swap-image")
 async def swap_image(source: UploadFile = File(...), target: UploadFile = File(...)):
-    task_id = str(uuid.uuid4())
-    TASKS[task_id] = {"status": "queued"}
-    executor.submit(run_task, task_id, await source.read(), await target.read())
-    return {"task_id": task_id}
-
-
-@app.get("/task-status/{task_id}")
-def task_status(task_id: str):
-    if task_id not in TASKS:
+    tid = str(uuid.uuid4())
+    TASKS[tid] = {"status": "queued"}
+    executor.submit(run_task, tid, await source.read(), await target.read())
+    return {"task_id": tid}
+
+@app.get("/task-status/{tid}")
+def status(tid: str):
+    if tid not in TASKS:
         raise HTTPException(404)
-    return TASKS[task_id]
-
+    return TASKS[tid]
 
-@app.get("/task-result/{task_id}")
-def task_result(task_id: str):
-    task = TASKS.get(task_id)
-    if not task or task["status"] != "done":
+@app.get("/task-result/{tid}")
+def result(tid: str):
+    t = TASKS.get(tid)
+    if not t or t["status"] != "done":
         raise HTTPException(404)
-    return StreamingResponse(open(task["result"], "rb"), media_type="image/jpeg")
-
+    return StreamingResponse(open(t["result"], "rb"), media_type="image/jpeg")
 
 # ============================================================
 # INIT
 # ============================================================
 print("Initializing models...")
-
 face_app = FaceAnalysis(name="buffalo_l")
 face_app.prepare(ctx_id=-1, det_size=(640, 640))
 
@@ -265,7 +265,6 @@ swapper = insightface.model_zoo.get_model(
     root=BASE_DIR
 )
 
-try_load_enhancer()
 try_load_gfpgan()
 
-print("FaceSwap API ready.")
+print("FaceSwap ready")