Update app.py

app.py

@@ -151,6 +151,37 @@ async def log_faceswap_hit(token: str, status: str = "success"):
 # --------------------- Face Swap Pipeline ---------------------
 swap_lock = threading.Lock()
 
+def enhance_image_with_codeformer(rgb_img, temp_dir=None):
+    if temp_dir is None:
+        temp_dir = os.path.join(tempfile.gettempdir(), f"enhance_{uuid.uuid4().hex[:8]}")
+    os.makedirs(temp_dir, exist_ok=True)
+
+    input_path = os.path.join(temp_dir, "input.jpg")
+    cv2.imwrite(input_path, cv2.cvtColor(rgb_img, cv2.COLOR_RGB2BGR))
+
+    python_cmd = sys.executable if sys.executable else "python3"
+    cmd = (
+        f"{python_cmd} {CODEFORMER_PATH} "
+        f"-w 0.7 "
+        f"--input_path {input_path} "
+        f"--output_path {temp_dir} "
+        f"--bg_upsampler realesrgan "
+        f"--face_upsample"
+    )
+
+    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(result.stderr)
+
+    final_dir = os.path.join(temp_dir, "final_results")
+    files = [f for f in os.listdir(final_dir) if f.endswith(".png")]
+    if not files:
+        raise RuntimeError("No enhanced output")
+
+    final_path = os.path.join(final_dir, files[0])
+    enhanced = cv2.imread(final_path)
+    return cv2.cvtColor(enhanced, cv2.COLOR_BGR2RGB)
+    
 def multi_face_swap(src_img, tgt_img):
     src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
     tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
@@ -348,8 +379,9 @@ def build_multi_faceswap_gradio():
 
         def process(src_img, tgt_img):
             try:
-                result = multi_face_swap(src_img, tgt_img)
-                return result, ""
+                swapped = multi_face_swap(src_img, tgt_img)
+                enhanced = enhance_image_with_codeformer(swapped)
+                return enhanced, ""
             except Exception as e:
                 return None, str(e)