faceswap function rewritten

app.py

@@ -87,12 +87,14 @@ def face_swap_and_enhance(src_img, tgt_img, fidelity=0.7, background_enhance=Tru
         os.makedirs(UPLOAD_DIR, exist_ok=True)
         os.makedirs(RESULT_DIR, exist_ok=True)
 
+        # ---------------- 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)
 
-        cmd = f"python {CODEFORMER_PATH} -w {fidelity:.2f} --input_path {UPLOAD_DIR} --output_path {RESULT_DIR}"
+        # ---------------- Run CodeFormer ----------------
+        cmd = f"python {CODEFORMER_PATH} -w {fidelity:.2f} --input_path {swapped_path} --output_path {RESULT_DIR}"
         if background_enhance:
             cmd += " --bg_upsampler realesrgan"
         if face_upsample:
@@ -102,24 +104,24 @@ def face_swap_and_enhance(src_img, tgt_img, fidelity=0.7, background_enhance=Tru
         if result.returncode != 0:
             return None, None, f"❌ CodeFormer failed:\n{result.stderr}"
 
-        final_path = None
-        for root, _, files in os.walk(RESULT_DIR):
-            for f in files:
-                if f.endswith((".png", ".jpg")):
-                    final_path = os.path.join(root, f)
-                    break
-            if final_path:
-                break
+        # ---------------- 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."
 
-        if not final_path or not os.path.exists(final_path):
-            return None, None, "❌ CodeFormer output missing."
+        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."
 
+        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, ""
 
     except Exception as e:
         return None, None, f"❌ Error: {str(e)}"
 
+
 # -------------------------------------------------
 # Gradio Interface
 # -------------------------------------------------