Update app.py

app.py

@@ -11,18 +11,19 @@ from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.responses import HTMLResponse, StreamingResponse
 
 # ============================================================
-# CONFIG (OPTIMIZED FOR SPEED + HF SPACES)
+# CONFIG
 # ============================================================
 MAX_FILE_MB = 10
-MAX_DIM = 640
-MAX_WORKERS = 3   # prevent CPU overload
+INPUT_MAX_DIM = 640
+OUTPUT_MAX_DIM = 1024
+MAX_WORKERS = 3
 CLEANUP_TIME = 300
 
 TASKS = {}
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
 
 # ============================================================
-# LOAD MODELS (ONLY ONCE)
+# MODELS
 # ============================================================
 face_app = insightface.app.FaceAnalysis(name="buffalo_l")
 face_app.prepare(ctx_id=-1, det_size=(640, 640))
@@ -30,42 +31,54 @@ face_app.prepare(ctx_id=-1, det_size=(640, 640))
 swapper = insightface.model_zoo.get_model("inswapper_128.onnx", root=".")
 
 # ============================================================
-# IMAGE OPTIMIZATION
+# IMAGE HELPERS
 # ============================================================
 
 def decode_image(file_bytes):
     img = cv2.imdecode(np.frombuffer(file_bytes, np.uint8), cv2.IMREAD_COLOR)
     if img is None:
-        raise ValueError("Invalid image format")
+        raise ValueError("Invalid image")
     return img
 
 
-def compress_and_resize(file_bytes):
-    img = decode_image(file_bytes)
+def compress_resize_input(img):
+    h, w = img.shape[:2]
 
-    size_mb = len(file_bytes) / (1024 * 1024)
+    if max(h, w) > INPUT_MAX_DIM:
+        scale = INPUT_MAX_DIM / max(h, w)
+        img = cv2.resize(img, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_LINEAR)
 
-    # compress if >10MB
-    if size_mb > MAX_FILE_MB:
-        img = cv2.resize(img, None, fx=0.6, fy=0.6, interpolation=cv2.INTER_AREA)
+    return img
 
-    # resize for speed
+
+def upscale_hd(img):
     h, w = img.shape[:2]
-    if max(h, w) > MAX_DIM:
-        scale = MAX_DIM / max(h, w)
-        img = cv2.resize(img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_LINEAR)
+
+    if max(h, w) < OUTPUT_MAX_DIM:
+        scale = OUTPUT_MAX_DIM / max(h, w)
+        img = cv2.resize(img, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_CUBIC)
 
     return img
 
 
-def enhance(img):
-    blur = cv2.GaussianBlur(img, (0, 0), 1.2)
-    return cv2.addWeighted(img, 1.2, blur, -0.2, 0)
+def sharpen(img):
+    kernel = np.array([[0,-1,0],[-1,5,-1],[0,-1,0]])
+    return cv2.filter2D(img, -1, kernel)
+
+
+def compress_if_needed(file_bytes):
+    size_mb = len(file_bytes) / (1024 * 1024)
+    img = decode_image(file_bytes)
+
+    if size_mb > MAX_FILE_MB:
+        img = cv2.resize(img, None, fx=0.6, fy=0.6, interpolation=cv2.INTER_AREA)
+
+    return img
 
 
 def cleanup():
     now = time.time()
-    delete_keys = []
+    remove = []
 
     for k, v in TASKS.items():
         if "time" in v and now - v["time"] > CLEANUP_TIME:
@@ -74,9 +87,9 @@ def cleanup():
                     os.remove(v["result"])
             except:
                 pass
-            delete_keys.append(k)
+            remove.append(k)
 
-    for k in delete_keys:
+    for k in remove:
         TASKS.pop(k, None)
 
 # ============================================================
@@ -87,20 +100,33 @@ def run_task(tid, src_bytes, tgt_bytes):
     TASKS[tid]["status"] = "processing"
 
     try:
-        src = compress_and_resize(src_bytes)
-        tgt = compress_and_resize(tgt_bytes)
+        # Decode + compress
+        src = compress_if_needed(src_bytes)
+        tgt = compress_if_needed(tgt_bytes)
+
+        # Resize input for speed
+        src = compress_resize_input(src)
+        tgt = compress_resize_input(tgt)
 
+        # Face detect
         s_faces = face_app.get(src)
         t_faces = face_app.get(tgt)
 
         if not s_faces or not t_faces:
             raise ValueError("Face not detected")
 
+        # Face swap
         result = swapper.get(tgt, t_faces[0], s_faces[0], paste_back=True)
-        result = enhance(result)
 
+        # 🔥 HD UPSCALE
+        result = upscale_hd(result)
+
+        # 🔥 SHARPEN
+        result = sharpen(result)
+
+        # Save
         out_path = f"/tmp/{tid}.webp"
-        cv2.imwrite(out_path, result, [cv2.IMWRITE_WEBP_QUALITY, 85])
+        cv2.imwrite(out_path, result, [cv2.IMWRITE_WEBP_QUALITY, 90])
 
         TASKS[tid] = {
             "status": "done",
@@ -119,7 +145,7 @@ def run_task(tid, src_bytes, tgt_bytes):
 app = FastAPI()
 
 # ============================================================
-# RESPONSIVE UI (UPLOAD + DOWNLOAD BUTTON)
+# UI (UPLOAD + DOWNLOAD)
 # ============================================================
 
 @app.get("/", response_class=HTMLResponse)
@@ -129,61 +155,17 @@ def home():
 <html>
 <head>
 <meta name="viewport" content="width=device-width, initial-scale=1">
-<title>Fast Face Swap</title>
+<title>HD Face Swap</title>
 
 <style>
-body{
-  margin:0;
-  font-family:sans-serif;
-  background:#0f172a;
-  color:white;
-  text-align:center;
-}
-
-.container{
-  max-width:900px;
-  margin:auto;
-  padding:20px;
-}
-
-.grid{
-  display:grid;
-  grid-template-columns:1fr 1fr;
-  gap:15px;
-}
-
-@media(max-width:700px){
-  .grid{grid-template-columns:1fr;}
-}
-
-.card{
-  background:#1e293b;
-  padding:15px;
-  border-radius:10px;
-}
-
-img{
-  width:100%;
-  max-height:250px;
-  object-fit:contain;
-  border-radius:10px;
-  background:black;
-}
-
-button{
-  padding:12px 18px;
-  margin:8px;
-  border:none;
-  border-radius:8px;
-  background:#6366f1;
-  color:white;
-  cursor:pointer;
-}
-
-.download{
-  display:none;
-  background:#10b981;
-}
+body{background:#0f172a;color:white;text-align:center;font-family:sans-serif}
+.container{max-width:900px;margin:auto;padding:20px}
+.grid{display:grid;grid-template-columns:1fr 1fr;gap:10px}
+@media(max-width:700px){.grid{grid-template-columns:1fr}}
+.card{background:#1e293b;padding:10px;border-radius:10px}
+img{width:100%;max-height:250px;object-fit:contain;border-radius:10px}
+button{padding:12px 18px;margin:10px;border:none;background:#6366f1;color:white;border-radius:8px}
+.download{display:none;background:#10b981}
 </style>
 
 </head>
@@ -192,90 +174,77 @@ button{
 
 <div class="container">
 
-<h2>⚡ Fast AI Face Swap</h2>
+<h2>⚡ HD Face Swap</h2>
 
 <div class="grid">
-
 <div class="card">
-<h3>Source</h3>
 <input type="file" id="src">
 <img id="p1">
 </div>
 
 <div class="card">
-<h3>Target</h3>
 <input type="file" id="tgt">
 <img id="p2">
 </div>
-
 </div>
 
-<br>
-
 <button onclick="start()">Upload & Swap</button>
 
 <p id="status"></p>
 
 <div class="card">
-<h3>Output</h3>
 <img id="out">
 <br>
-<a id="downloadBtn" class="download" download="faceswap.webp">Download</a>
+<a id="dl" class="download" download="faceswap_hd.webp">Download</a>
 </div>
 
 </div>
 
 <script>
-const src = document.getElementById("src");
-const tgt = document.getElementById("tgt");
-const statusEl = document.getElementById("status");
-const downloadBtn = document.getElementById("downloadBtn");
+const src=document.getElementById("src")
+const tgt=document.getElementById("tgt")
+const st=document.getElementById("status")
+const dl=document.getElementById("dl")
 
-src.onchange = ()=> p1.src = URL.createObjectURL(src.files[0]);
-tgt.onchange = ()=> p2.src = URL.createObjectURL(tgt.files[0]);
+src.onchange=()=>p1.src=URL.createObjectURL(src.files[0])
+tgt.onchange=()=>p2.src=URL.createObjectURL(tgt.files[0])
 
 async function start(){
-  if(!src.files[0] || !tgt.files[0]){
-    alert("Upload both images");
-    return;
-  }
+ if(!src.files[0]||!tgt.files[0]) return alert("upload both")
 
-  let fd = new FormData();
-  fd.append("source", src.files[0]);
-  fd.append("target", tgt.files[0]);
+ let fd=new FormData()
+ fd.append("source",src.files[0])
+ fd.append("target",tgt.files[0])
 
-  statusEl.innerText = "Uploading...";
+ st.innerText="Processing..."
 
-  let r = await fetch("/swap", {method:"POST", body:fd});
-  let j = await r.json();
+ let r=await fetch("/swap",{method:"POST",body:fd})
+ let j=await r.json()
 
-  poll(j.task_id);
+ poll(j.task_id)
 }
 
 async function poll(id){
-  let r = await fetch("/status/"+id);
-  let j = await r.json();
-
-  statusEl.innerText = j.status;
-
-  if(j.status === "done"){
-    let img = await fetch("/result/"+id);
-    let blob = await img.blob();
-
-    let url = URL.createObjectURL(blob);
-    out.src = url;
-
-    downloadBtn.href = url;
-    downloadBtn.style.display = "inline-block";
-
-    statusEl.innerText = "Done ✅";
-  }
-  else if(j.status === "failed"){
-    statusEl.innerText = "Error: " + j.error;
-  }
-  else{
-    setTimeout(()=>poll(id), 800);
-  }
+ let r=await fetch("/status/"+id)
+ let j=await r.json()
+
+ st.innerText=j.status
+
+ if(j.status==="done"){
+   let img=await fetch("/result/"+id)
+   let b=await img.blob()
+   let u=URL.createObjectURL(b)
+
+   out.src=u
+   dl.href=u
+   dl.style.display="inline-block"
+
+   st.innerText="Done ✅"
+ }else if(j.status==="failed"){
+   st.innerText="Error: "+j.error
+ }else{
+   setTimeout(()=>poll(id),800)
+ }
 }
 </script>
 
@@ -284,7 +253,7 @@ async function poll(id){
 """
 
 # ============================================================
-# API (UNCHANGED)
+# API
 # ============================================================
 
 @app.post("/swap")
@@ -293,12 +262,7 @@ async def swap(source: UploadFile = File(...), target: UploadFile = File(...)):
 
     TASKS[tid] = {"status": "queued", "time": time.time()}
 
-    executor.submit(
-        run_task,
-        tid,
-        await source.read(),
-        await target.read()
-    )
+    executor.submit(run_task, tid, await source.read(), await target.read())
 
     return {"task_id": tid}
 
@@ -306,17 +270,14 @@ async def swap(source: UploadFile = File(...), target: UploadFile = File(...)):
 @app.get("/status/{tid}")
 def status(tid: str):
     cleanup()
-
     if tid not in TASKS:
         raise HTTPException(404)
-
     return TASKS[tid]
 
 
 @app.get("/result/{tid}")
 def result(tid: str):
     task = TASKS.get(tid)
-
     if not task or task["status"] != "done":
         raise HTTPException(404)