Update app.py

app.py

@@ -1,15 +1,29 @@
+import sys
+import types
+
+# Patch for newer torchvision compatibility with GFPGAN
+try:
+    import torchvision.transforms.functional as F
+    functional_tensor = types.ModuleType("torchvision.transforms.functional_tensor")
+    functional_tensor.rgb_to_grayscale = F.rgb_to_grayscale
+    sys.modules["torchvision.transforms.functional_tensor"] = functional_tensor
+except Exception:
+    pass
+
 import gradio as gr
 import cv2
 import insightface
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
+from gfpgan import GFPGANer
 import spaces
 
 app = None
 swapper = None
+restorer = None
 
 def load_models():
-    global app, swapper
+    global app, swapper, restorer
     if app is None:
         app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
         app.prepare(ctx_id=-1, det_size=(640, 640))
@@ -19,9 +33,21 @@ def load_models():
             filename="inswapper_128.onnx"
         )
         swapper = insightface.model_zoo.get_model(swapper_path, download=False)
+    if restorer is None:
+        gfpgan_path = hf_hub_download(
+            repo_id="gmk123/GFPGAN",
+            filename="GFPGANv1.4.pth"
+        )
+        restorer = GFPGANer(
+            model_path=gfpgan_path,
+            upscale=1,
+            arch="clean",
+            channel_multiplier=2,
+            bg_upsampler=None
+        )
 
 @spaces.GPU
-def swap_faces(source_img, target_img):
+def swap_faces(source_img, target_img, enhance_face):
     if source_img is None or target_img is None:
         return None
 
@@ -39,19 +65,31 @@ def swap_faces(source_img, target_img):
     source_face = sorted(source_faces, key=lambda x: x.bbox[0])[0]
     target_face = sorted(target_faces, key=lambda x: x.bbox[0])[0]
 
+    # 1. Perform the face swap
     result_bgr = swapper.get(target_bgr, target_face, source_face, paste_back=True)
+
+    # 2. Sharpen and restore the face with GFPGAN if checked
+    if enhance_face:
+        _, _, result_bgr = restorer.enhance(
+            result_bgr,
+            has_aligned=False,
+            only_center_face=False,
+            paste_back=True
+        )
+
     return cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎭 Private Face Swapper")
+    gr.Markdown("# 🎭 Private Face Swapper + GFPGAN Enhancer")
     with gr.Row():
         with gr.Column():
             src = gr.Image(label="Source Image (Face to Copy)", type="numpy")
             tgt = gr.Image(label="Target Image (Body / Scene)", type="numpy")
+            enhance = gr.Checkbox(label="Enhance & Sharpen Face (GFPGAN v1.4)", value=True)
             btn = gr.Button("Swap Face", variant="primary")
         with gr.Column():
             out = gr.Image(label="Swapped Result", type="numpy")
 
-    btn.click(fn=swap_faces, inputs=[src, tgt], outputs=out)
+    btn.click(fn=swap_faces, inputs=[src, tgt, enhance], outputs=out)
 
 demo.queue().launch()