Update app.py

app.py

@@ -11,9 +11,9 @@ from insightface.app import FaceAnalysis
 from tempfile import NamedTemporaryFile
 
 from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.responses import HTMLResponse, Response, StreamingResponse
+from fastapi.responses import HTMLResponse, StreamingResponse
 
-# Optional ONNX enhancer import (only used if small_enhancer.onnx exists)
+# Optional ONNX enhancer import
 try:
     import onnxruntime as ort
     ONNX_AVAILABLE = True
@@ -25,17 +25,16 @@ except Exception:
 # -----------------------------------------------------------
 face_app = None
 swapper = None
-enhancer_session = None  # optional onnx session
+enhancer_session = None
 
 # task store (in-memory)
 TASKS = {}  # task_id -> {"status": "queued|processing|done|failed", "result": path, "error": text}
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=6)
 
 # -----------------------------------------------------------
-# MODEL DOWNLOAD (swapper)
+# MODELS - download swapper if needed
 # -----------------------------------------------------------
 def download_swapper_model():
-    """Download inswapper_128.onnx if not present (used by insightface model_zoo loader)."""
     url = "https://cdn.adikhanofficial.com/python/insightface/models/inswapper_128.onnx"
     filename = os.path.basename(url)
     save_path = os.path.join(os.path.dirname(__file__), filename)
@@ -49,19 +48,17 @@ def download_swapper_model():
     else:
         print("Swapper model already exists:", filename)
 
-# -----------------------------------------------------------
-# OPTIONAL: try to initialize a small ONNX enhancer if present
-# -----------------------------------------------------------
+
 def try_load_enhancer():
     global enhancer_session
     if not ONNX_AVAILABLE:
-        print("onnxruntime not installed — enhancer disabled (fallback to OpenCV).")
+        print("onnxruntime not installed — ONNX enhancer disabled (using OpenCV fallback).")
         return
-    path = os.path.join(os.path.dirname(__file__), "small_enhancer.onnx")
-    if os.path.exists(path):
+    p = os.path.join(os.path.dirname(__file__), "small_enhancer.onnx")
+    if os.path.exists(p):
         try:
-            enhancer_session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
-            print("Loaded ONNX enhancer from", path)
+            enhancer_session = ort.InferenceSession(p, providers=["CPUExecutionProvider"])
+            print("Loaded ONNX enhancer:", p)
         except Exception as e:
             print("Failed to load ONNX enhancer:", e)
             enhancer_session = None
@@ -70,47 +67,34 @@ def try_load_enhancer():
 
 
 # -----------------------------------------------------------
-# FACE SWAP HELPER (insightface)
+# SWAP HELPER (insightface swapper)
 # -----------------------------------------------------------
 def swap_faces(target_img, target_face, source_face):
-    # swapper.get is synchronous
     return swapper.get(target_img, target_face, source_face, paste_back=True)
 
 
 # -----------------------------------------------------------
-# ENHANCEMENT: ONNX (if present) OR OpenCV fallback
+# ENHANCERS
 # -----------------------------------------------------------
 def enhance_face_with_onnx(face_bgr):
-    """
-    Attempt to use enhancer_session. This code assumes the ONNX model accepts
-    a normalized float32 CHW input named whatever the first input is.
-    Because small ONNX models vary, this function tries a generic call.
-    If it fails, it will raise and caller will fallback to OpenCV enhancement.
-    """
     if enhancer_session is None:
         raise RuntimeError("Enhancer session not available")
-
-    # Generic preprocessing: resize to model input size if single static input size detected
     inp = enhancer_session.get_inputs()[0]
     name = inp.name
-    shape = inp.shape  # e.g. [1,3,256,256] or [None,3,512,512]
-    # decide target size
+    shape = inp.shape
+    # determine target size (safe fallback)
     try:
-        _, c, h, w = (shape if len(shape) == 4 else (1,3,512,512))
-        h = int(h) if (isinstance(h, int) or (isinstance(h, np.integer))) else 512
-        w = int(w) if (isinstance(w, int) or (isinstance(w, np.integer))) else 512
+        _, c, h, w = (shape if len(shape) == 4 else (1, 3, 512, 512))
+        h = int(h) if isinstance(h, (int, np.integer)) else 512
+        w = int(w) if isinstance(w, (int, np.integer)) else 512
     except Exception:
         h, w = 512, 512
 
     img = cv2.cvtColor(face_bgr, cv2.COLOR_BGR2RGB)
     img = cv2.resize(img, (w, h), interpolation=cv2.INTER_CUBIC)
     img = img.astype(np.float32) / 255.0
-    # CHW
     img = np.transpose(img, (2, 0, 1))[None].astype(np.float32)
-
-    # Run
     out = enhancer_session.run(None, {name: img})
-    # pick first output
     out_img = out[0][0]
     out_img = np.clip(out_img * 255.0, 0, 255).astype(np.uint8)
     out_img = np.transpose(out_img, (1, 2, 0))
@@ -119,57 +103,108 @@ def enhance_face_with_onnx(face_bgr):
 
 
 def enhance_face_opencv(face_bgr):
-    """
-    Lightweight, fast CPU-only enhancement:
-    - upscale x2 using bicubic
-    - denoise with bilateral filter
-    - sharpen via unsharp mask
-    This works well as a fallback for small faces and is deterministic.
-    """
     if face_bgr is None or face_bgr.size == 0:
         return face_bgr
-
-    # upscale x2
     h, w = face_bgr.shape[:2]
     target_h, target_w = max(64, h * 2), max(64, w * 2)
     up = cv2.resize(face_bgr, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
-
-    # bilateral filter to smooth artifacts
     denoised = cv2.bilateralFilter(up, d=5, sigmaColor=75, sigmaSpace=75)
-
-    # unsharp mask (sharpen)
-    blurred = cv2.GaussianBlur(denoised, (0,0), sigmaX=2, sigmaY=2)
+    blurred = cv2.GaussianBlur(denoised, (0, 0), sigmaX=2, sigmaY=2)
     sharpened = cv2.addWeighted(denoised, 1.4, blurred, -0.4, 0)
-
-    # optional mild contrast boost
     lab = cv2.cvtColor(sharpened, cv2.COLOR_BGR2LAB)
     l, a, b = cv2.split(lab)
     l = cv2.equalizeHist(l)
     lab = cv2.merge((l, a, b))
     result = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
-
-    # resize back to original face box size will be handled by caller
     return result
 
 
 def enhance_face(face_bgr):
-    # Try ONNX first if available
     if enhancer_session is not None:
         try:
             return enhance_face_with_onnx(face_bgr)
         except Exception as e:
-            print("ONNX enhancer failed, falling back to OpenCV enhancer:", e)
-            # fall through to OpenCV
+            print("ONNX enhancement failed:", e)
     return enhance_face_opencv(face_bgr)
 
 
+# -----------------------------------------------------------
+# COLOR HARMONIZATION (Reinhard-like in LAB)
+# -----------------------------------------------------------
+def color_transfer_reinhard(src_bgr, ref_bgr):
+    """
+    Map src colors to match ref using mean/std in LAB color space (fast).
+    Returns adjusted src (same size).
+    """
+    try:
+        src_lab = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2LAB).astype(np.float32)
+        ref_lab = cv2.cvtColor(ref_bgr, cv2.COLOR_BGR2LAB).astype(np.float32)
+
+        src_means, src_stds = cv2.meanStdDev(src_lab)
+        ref_means, ref_stds = cv2.meanStdDev(ref_lab)
+
+        src_means = src_means.flatten()
+        src_stds = src_stds.flatten()
+        ref_means = ref_means.flatten()
+        ref_stds = ref_stds.flatten()
+
+        # avoid division by zero
+        src_stds[src_stds < 1e-6] = 1.0
+
+        # transfer
+        lab = (src_lab - src_means) * (ref_stds / src_stds) + ref_means
+        lab = np.clip(lab, 0, 255).astype(np.uint8)
+        out = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
+        return out
+    except Exception:
+        return src_bgr
+
+
+# -----------------------------------------------------------
+# SOFT FEATHER BLEND
+# -----------------------------------------------------------
+def feather_blend(target_img, new_face, bbox, feather_amount=0.3):
+    """
+    Soft blend new_face (same size as bbox) into target_img at bbox using a
+    Gaussian feathered alpha mask.
+    - bbox: (x1,y1,x2,y2)
+    - feather_amount: fraction of bbox size used for feather radius
+    """
+    x1, y1, x2, y2 = bbox
+    h = y2 - y1
+    w = x2 - x1
+    if h <= 0 or w <= 0:
+        return target_img
+
+    # build mask
+    mask = np.ones((h, w), dtype=np.uint8) * 255
+    # feather sigma proportional to size
+    k = int(max(1, min(h, w) * feather_amount))
+    k = k if k % 2 == 1 else k + 1
+    mask_blur = cv2.GaussianBlur(mask, (k, k), 0)
+
+    # normalize alpha to [0,1]
+    alpha = (mask_blur.astype(np.float32) / 255.0)[:, :, None]
+
+    # ensure sizes match
+    face_resized = cv2.resize(new_face, (w, h), interpolation=cv2.INTER_CUBIC)
+
+    # region in target
+    target_region = target_img[y1:y2, x1:x2].astype(np.float32)
+    src_region = face_resized.astype(np.float32)
+
+    comp = (alpha * src_region + (1 - alpha) * target_region).astype(np.uint8)
+    out_img = target_img.copy()
+    out_img[y1:y2, x1:x2] = comp
+    return out_img
+
+
 # -----------------------------------------------------------
 # BACKGROUND TASK HANDLER
 # -----------------------------------------------------------
 def run_task(task_id, src_bytes, tgt_bytes):
     TASKS[task_id]["status"] = "processing"
     try:
-        # decode
         src = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
         tgt = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
 
@@ -178,7 +213,6 @@ def run_task(task_id, src_bytes, tgt_bytes):
         if tgt is None:
             raise ValueError("Invalid target image")
 
-        # detect faces
         src_faces = face_app.get(src)
         tgt_faces = face_app.get(tgt)
 
@@ -187,38 +221,44 @@ def run_task(task_id, src_bytes, tgt_bytes):
         if not tgt_faces:
             raise ValueError("No face detected in target image")
 
-        # choose first faces (you can enhance selection logic)
         s_face = src_faces[0]
         t_face = tgt_faces[0]
 
-        # perform swap (returns full image with pasted face)
         swapped = swap_faces(tgt, t_face, s_face)
 
-        # attempt to crop the face region (using bbox from detection)
+        # Attempt to extract bbox and do enhancement + color transfer + blend
         try:
-            # bbox could be float, ensure ints and clamp
             x1, y1, x2, y2 = map(int, map(round, t_face.bbox))
             h, w = swapped.shape[:2]
-            x1, x2 = max(0, min(x1, w-1)), max(0, min(x2, w))
-            y1, y2 = max(0, min(y1, h-1)), max(0, min(y2, h))
+            x1, x2 = max(0, min(x1, w - 1)), max(0, min(x2, w))
+            y1, y2 = max(0, min(y1, h - 1)), max(0, min(y2, h))
+
             if x2 - x1 > 10 and y2 - y1 > 10:
                 face_crop = swapped[y1:y2, x1:x2].copy()
 
-                # enhance the face crop
+                # enhancement (ONNX or OpenCV)
                 enhanced = enhance_face(face_crop)
 
-                # resize enhanced back to bbox size (in case upscaled)
-                enhanced_resized = cv2.resize(enhanced, (x2 - x1, y2 - y1), interpolation=cv2.INTER_CUBIC)
-                swapped[y1:y2, x1:x2] = enhanced_resized
+                # color transfer: match enhanced face colors to the target's surrounding region
+                # sample a slightly expanded patch from the original target to get ambient color
+                pad = int(0.15 * max(y2 - y1, x2 - x1))
+                rx1, ry1 = max(0, x1 - pad), max(0, y1 - pad)
+                rx2, ry2 = min(w, x2 + pad), min(h, y2 + pad)
+                ref_region = tgt[ry1:ry2, rx1:rx2]
+                if ref_region is None or ref_region.size == 0:
+                    ref_region = tgt[y1:y2, x1:x2]
+                # perform Reinhard-style mean/std transfer in LAB on enhanced crop
+                transferred = color_transfer_reinhard(enhanced, ref_region)
+
+                # blend transferred face back into swapped image softly
+                swapped = feather_blend(swapped, transferred, (x1, y1, x2, y2), feather_amount=0.35)
             else:
-                # bbox too small; skip enhancement
+                # bbox too small, keep swapped
                 pass
         except Exception as e:
-            # If any error occurs during enhancement, continue with swapped image
-            print("Face enhancement step failed:", e)
+            print("Enhancement/color-blend failed:", e)
             print(traceback.format_exc())
 
-        # write output
         out_path = f"/tmp/{task_id}.jpg"
         ok = cv2.imwrite(out_path, swapped)
         if not ok:
@@ -237,11 +277,10 @@ def run_task(task_id, src_bytes, tgt_bytes):
 # -----------------------------------------------------------
 # FASTAPI app & UI
 # -----------------------------------------------------------
-app = FastAPI(title="FaceSwap Async API (image-only)")
+app = FastAPI(title="FaceSwap Async API (image-only, color harmonized)")
 
 @app.get("/", response_class=HTMLResponse)
 def home():
-    # stylish async UI with fixed drag & drop (same as previous UI, improved)
     return """
 <!doctype html>
 <html lang="en">
@@ -250,191 +289,72 @@ def home():
 <title>FaceSwap Async API</title>
 <meta name="viewport" content="width=device-width,initial-scale=1">
 <style>
-    :root { --bg:#0e1525; --card:rgba(255,255,255,0.06); --accent:#6c5ce7; --accent2:#00cec9; --text:#e8ecf2; --muted:#9ba3b4; --danger:#ff6b6b; --success:#2ecc71; --radius:14px; }
-    body { background: linear-gradient(160deg,#111827,#0e1525); font-family: Inter, sans-serif; padding:25px; color:var(--text); margin:0; }
-    .card { background:var(--card); padding:22px; border-radius:var(--radius); box-shadow:0 8px 25px rgba(0,0,0,0.4); max-width:900px; margin:auto; }
-    h1{ font-size:24px; margin:0 0 6px 0; } .subtitle{ color:var(--muted); font-size:14px; margin-bottom:14px; }
-    .row{ display:flex; gap:12px; align-items:center; flex-wrap:wrap; }
-    .upload-zone { flex:1; min-width:220px; border:2px dashed rgba(255,255,255,0.12); padding:18px; border-radius:12px; text-align:center; cursor:pointer; transition:0.12s; }
-    .upload-zone.dragover{ border-color: var(--accent); background: rgba(255,255,255,0.02); transform: translateY(-4px); }
-    input[type=file]{ display:none; }
-    button{ background: linear-gradient(90deg,var(--accent),var(--accent2)); border:none; padding:12px 18px; border-radius:12px; color:white; font-weight:600; cursor:pointer; }
-    .status{ margin-top:12px; color:var(--muted); }
-    .task-id{ margin-top:8px; color:var(--accent2); font-weight:700; }
-    .preview{ margin-top:16px; text-align:center; }
-    .preview img{ max-width:100%; border-radius:10px; border:1px solid rgba(255,255,255,0.06); background:#021017; }
-    .loader { width:20px;height:20px;border:3px solid rgba(255,255,255,0.12); border-top-color:var(--accent2); border-radius:50%; animation:spin .9s linear infinite; display:inline-block; }
-    @keyframes spin{ to{ transform:rotate(360deg) } }
+:root{--bg:#0e1525;--card:rgba(255,255,255,0.06);--accent:#6c5ce7;--accent2:#00cec9;--text:#e8ecf2;--muted:#9ba3b4;--danger:#ff6b6b;--success:#2ecc71;--radius:14px}
+body{background:linear-gradient(160deg,#111827,#0e1525);font-family:Inter, sans-serif;padding:25px;color:var(--text);margin:0}
+.card{background:var(--card);padding:22px;border-radius:var(--radius);box-shadow:0 8px 25px rgba(0,0,0,0.4);max-width:900px;margin:auto}
+h1{font-size:24px;margin:0 0 6px 0}.subtitle{color:var(--muted);font-size:14px;margin-bottom:14px}.row{display:flex;gap:12px;align-items:center;flex-wrap:wrap}.upload-zone{flex:1;min-width:220px;border:2px dashed rgba(255,255,255,0.12);padding:18px;border-radius:12px;text-align:center;cursor:pointer;transition:0.12s}.upload-zone.dragover{border-color:var(--accent);background:rgba(255,255,255,0.02);transform:translateY(-4px)}input[type=file]{display:none}button{background:linear-gradient(90deg,var(--accent),var(--accent2));border:none;padding:12px 18px;border-radius:12px;color:white;font-weight:600;cursor:pointer}.status{margin-top:12px;color:var(--muted)}.task-id{margin-top:8px;color:var(--accent2);font-weight:700}.preview{margin-top:16px;text-align:center}.preview img{max-width:100%;border-radius:10px;border:1px solid rgba(255,255,255,0.06);background:#021017}.loader{width:20px;height:20px;border:3px solid rgba(255,255,255,0.12);border-top-color:var(--accent2);border-radius:50%;animation:spin .9s linear infinite;display:inline-block}@keyframes spin{to{transform:rotate(360deg)}}
 </style>
 </head>
 <body>
 <div class="card">
   <h1>Async FaceSwap</h1>
   <div class="subtitle">Upload a source face image and a target image. You get a task id immediately; result will appear when ready.</div>
-
   <div class="row">
-    <label class="upload-zone" id="zoneSource">
-      <div><strong id="labelSource">Click or drop source image</strong></div>
-      <input id="inputSource" type="file" accept="image/*">
-    </label>
-
-    <label class="upload-zone" id="zoneTarget">
-      <div><strong id="labelTarget">Click or drop target image</strong></div>
-      <input id="inputTarget" type="file" accept="image/*">
-    </label>
+    <label class="upload-zone" id="zoneSource"><div><strong id="labelSource">Click or drop source image</strong></div><input id="inputSource" type="file" accept="image/*"></label>
+    <label class="upload-zone" id="zoneTarget"><div><strong id="labelTarget">Click or drop target image</strong></div><input id="inputTarget" type="file" accept="image/*"></label>
   </div>
-
-  <div class="row" style="margin-top:12px;">
+  <div class="row" style="margin-top:12px">
     <button id="startBtn">Start FaceSwap</button>
-    <button id="resetBtn" style="background:transparent;border:1px solid rgba(255,255,255,0.06);color:var(--muted);">Reset</button>
+    <button id="resetBtn" style="background:transparent;border:1px solid rgba(255,255,255,0.06);color:var(--muted)">Reset</button>
     <div style="flex:1"></div>
     <div id="status" class="status"></div>
   </div>
-
   <div class="task-id" id="taskId"></div>
-
-  <div class="preview">
-    <img id="resultImg" style="display:none;">
-  </div>
-
+  <div class="preview"><img id="resultImg" style="display:none"></div>
 </div>
-
 <script>
 (function(){
-  const zoneSource = document.getElementById('zoneSource');
-  const zoneTarget = document.getElementById('zoneTarget');
-  const inputSource = document.getElementById('inputSource');
-  const inputTarget = document.getElementById('inputTarget');
-  const labelSource = document.getElementById('labelSource');
-  const labelTarget = document.getElementById('labelTarget');
-  const startBtn = document.getElementById('startBtn');
-  const resetBtn = document.getElementById('resetBtn');
-  const status = document.getElementById('status');
-  const taskIdEl = document.getElementById('taskId');
-  const resultImg = document.getElementById('resultImg');
-
-  let sourceFile = null;
-  let targetFile = null;
-  let pollHandle = null;
-
-  function prevent(e){ e.preventDefault(); e.stopPropagation(); }
-
-  function makeDrop(zone, input, label){
-    ['dragenter','dragover','dragleave','drop'].forEach(ev => {
-      zone.addEventListener(ev, prevent);
-    });
-    zone.addEventListener('dragover', ()=> zone.classList.add('dragover'));
-    zone.addEventListener('dragleave', ()=> zone.classList.remove('dragover'));
-    zone.addEventListener('drop', (ev)=>{
-      zone.classList.remove('dragover');
-      const f = ev.dataTransfer.files && ev.dataTransfer.files[0];
-      if (!f) return;
-      input.files = ev.dataTransfer.files;
-      label.innerText = f.name;
-      if (input === inputSource) sourceFile = f;
-      if (input === inputTarget) targetFile = f;
-    });
-    zone.addEventListener('click', ()=> input.click());
-    input.addEventListener('change', ()=> {
-      const f = input.files && input.files[0];
-      if (!f) return;
-      label.innerText = f.name;
-      if (input === inputSource) sourceFile = f;
-      if (input === inputTarget) targetFile = f;
-    });
+  const zoneSource=document.getElementById('zoneSource'), zoneTarget=document.getElementById('zoneTarget');
+  const inputSource=document.getElementById('inputSource'), inputTarget=document.getElementById('inputTarget');
+  const labelSource=document.getElementById('labelSource'), labelTarget=document.getElementById('labelTarget');
+  const startBtn=document.getElementById('startBtn'), resetBtn=document.getElementById('resetBtn');
+  const status=document.getElementById('status'), taskIdEl=document.getElementById('taskId'), resultImg=document.getElementById('resultImg');
+  let sourceFile=null, targetFile=null, pollHandle=null;
+  function prevent(e){e.preventDefault();e.stopPropagation();}
+  function makeDrop(zone,input,label){
+    ['dragenter','dragover','dragleave','drop'].forEach(ev=>zone.addEventListener(ev,prevent));
+    zone.addEventListener('dragover',()=>zone.classList.add('dragover'));
+    zone.addEventListener('dragleave',()=>zone.classList.remove('dragover'));
+    zone.addEventListener('drop',ev=>{zone.classList.remove('dragover');const f=ev.dataTransfer.files&&ev.dataTransfer.files[0];if(!f) return;input.files=ev.dataTransfer.files;label.innerText=f.name;if(input===inputSource) sourceFile=f; if(input===inputTarget) targetFile=f;});
+    zone.addEventListener('click',()=>input.click());
+    input.addEventListener('change',()=>{const f=input.files&&input.files[0]; if(!f) return; label.innerText=f.name; if(input===inputSource) sourceFile=f; if(input===inputTarget) targetFile=f; });
   }
-
-  makeDrop(zoneSource, inputSource, labelSource);
-  makeDrop(zoneTarget, inputTarget, labelTarget);
-
-  startBtn.addEventListener('click', async ()=>{
-    if (!sourceFile || !targetFile) { alert('Select source and target images'); return; }
-    status.innerHTML = 'Uploading <span class="loader"></span>';
-    startBtn.disabled = true;
-    try {
-      const fd = new FormData();
-      fd.append('source', sourceFile);
-      fd.append('target', targetFile);
-      const res = await fetch('/swap-image', { method:'POST', body: fd });
-      if (!res.ok) {
-        const txt = await res.text().catch(()=>res.statusText);
-        status.innerHTML = '<span style="color:#ff6b6b">Upload failed: ' + txt + '</span>';
-        startBtn.disabled = false;
-        return;
-      }
-      const data = await res.json();
-      taskIdEl.innerText = 'Task ID: ' + data.task_id;
-      pollStatus(data.task_id);
-    } catch (err) {
-      status.innerText = 'Network error';
-      console.error(err);
-    } finally {
-      // keep button disabled until task completes or reset
-    }
+  makeDrop(zoneSource,inputSource,labelSource); makeDrop(zoneTarget,inputTarget,labelTarget);
+  startBtn.addEventListener('click',async()=>{
+    if(!sourceFile||!targetFile){alert('Select source and target images');return;}
+    status.innerHTML='Uploading <span class="loader"></span>'; startBtn.disabled=true;
+    try{
+      const fd=new FormData(); fd.append('source', sourceFile); fd.append('target', targetFile);
+      const res=await fetch('/swap-image',{method:'POST',body:fd});
+      if(!res.ok){ const t=await res.text().catch(()=>res.statusText); status.innerHTML='<span style="color:#ff6b6b">Upload failed: '+t+'</span>'; startBtn.disabled=false; return; }
+      const data=await res.json(); taskIdEl.innerText='Task ID: '+data.task_id; pollStatus(data.task_id);
+    }catch(err){ console.error(err); status.innerText='Network error'; startBtn.disabled=false; }
   });
-
-  resetBtn.addEventListener('click', ()=>{
-    sourceFile = targetFile = null;
-    inputSource.value = '';
-    inputTarget.value = '';
-    labelSource.innerText = 'Click or drop source image';
-    labelTarget.innerText = 'Click or drop target image';
-    status.innerText = '';
-    taskIdEl.innerText = '';
-    resultImg.style.display = 'none';
-    if (pollHandle) { clearInterval(pollHandle); pollHandle = null; }
-    startBtn.disabled = false;
-  });
-
-  function pollStatus(taskId){
-    // poll every 1s
-    pollHandle = setInterval(async ()=>{
-      try {
-        const res = await fetch('/task-status/' + taskId);
-        const data = await res.json();
-        if (data.status === 'processing') {
-          status.innerHTML = 'Processing <span class="loader"></span>';
-        } else if (data.status === 'failed') {
-          status.innerHTML = '<span style="color:#ff6b6b">Failed: ' + (data.error || '') + '</span>';
-          clearInterval(pollHandle);
-          pollHandle = null;
-          startBtn.disabled = false;
-        } else if (data.status === 'done') {
-          status.innerHTML = '<span style="color:#2ecc71">Completed</span>';
-          clearInterval(pollHandle);
-          pollHandle = null;
-          fetchResult(taskId);
-        } else {
-          status.innerText = data.status;
-        }
-      } catch (err) {
-        console.error(err);
-        status.innerText = 'Status error';
-      }
-    }, 1000);
-  }
-
+  resetBtn.addEventListener('click',()=>{ sourceFile=targetFile=null; inputSource.value=''; inputTarget.value=''; labelSource.innerText='Click or drop source image'; labelTarget.innerText='Click or drop target image'; status.innerText=''; taskIdEl.innerText=''; resultImg.style.display='none'; if(pollHandle){clearInterval(pollHandle); pollHandle=null;} startBtn.disabled=false; });
+  function pollStatus(taskId){ pollHandle=setInterval(async()=>{
+    try{
+      const res=await fetch('/task-status/'+taskId); const data=await res.json();
+      if(data.status==='processing'){ status.innerHTML='Processing <span class="loader"></span>'; }
+      else if(data.status==='failed'){ status.innerHTML='<span style="color:#ff6b6b">Failed: '+(data.error||'')+'</span>'; clearInterval(pollHandle); pollHandle=null; startBtn.disabled=false; }
+      else if(data.status==='done'){ status.innerHTML='<span style="color:#2ecc71">Completed</span>'; clearInterval(pollHandle); pollHandle=null; fetchResult(taskId); } else { status.innerText=data.status; }
+    }catch(err){ console.error(err); status.innerText='Status error'; }
+  },1000); }
   async function fetchResult(taskId){
-    try {
-      const res = await fetch('/task-result/' + taskId);
-      if (!res.ok) {
-        const t = await res.text().catch(()=>res.statusText);
-        status.innerHTML = '<span style="color:#ff6b6b">Result fetch failed: ' + t + '</span>';
-        startBtn.disabled = false;
-        return;
-      }
-      const blob = await res.blob();
-      const url = URL.createObjectURL(blob);
-      resultImg.src = url;
-      resultImg.style.display = 'block';
-      startBtn.disabled = false;
-    } catch (err) {
-      console.error(err);
-      status.innerText = 'Fetch error';
-      startBtn.disabled = false;
-    }
+    try{
+      const res=await fetch('/task-result/'+taskId); if(!res.ok){ const t=await res.text().catch(()=>res.statusText); status.innerHTML='<span style="color:#ff6b6b">Result fetch failed: '+t+'</span>'; startBtn.disabled=false; return; }
+      const blob=await res.blob(); const url=URL.createObjectURL(blob); resultImg.src=url; resultImg.style.display='block'; startBtn.disabled=false;
+    }catch(err){ console.error(err); status.innerText='Fetch error'; startBtn.disabled=false; }
   }
-
 })();
 </script>
 </body>
@@ -452,7 +372,6 @@ async def swap_image(source: UploadFile = File(...), target: UploadFile = File(.
     task_id = str(uuid.uuid4())
     TASKS[task_id] = {"status": "queued"}
 
-    # submit background job
     executor.submit(run_task, task_id, src_bytes, tgt_bytes)
 
     return {"task_id": task_id, "status": "queued"}
@@ -478,7 +397,6 @@ def task_result(task_id: str):
     task = TASKS[task_id]
     if task["status"] != "done":
         return {"status": task["status"], "error": task.get("error")}
-    # stream result
     return StreamingResponse(open(task["result"], "rb"), media_type="image/jpeg")
 
 
@@ -486,15 +404,15 @@ def task_result(task_id: str):
 # INIT MODELS
 # -----------------------------------------------------------
 print("Initializing models...")
-# prepare detection model (CPU)
+# detection model
 face_app = FaceAnalysis(name="buffalo_l")
 face_app.prepare(ctx_id=-1, det_size=(640, 640))
 
-# download swapper model if needed and load swapper
+# swapper model
 download_swapper_model()
 swapper = insightface.model_zoo.get_model("inswapper_128.onnx", root=os.path.dirname(__file__))
 
-# try load optional small ONNX enhancer if you provided it
+# optional enhancer
 if ONNX_AVAILABLE:
     try_load_enhancer()