Update app.py

app.py

@@ -1,55 +1,70 @@
 import gradio as gr
+import os
 import cv2
+import torch
 import numpy as np
-from insightface.app import FaceAnalysis
-from insightface.model_zoo.inswapper import INSwapper
+from PIL import Image
+
 from gfpgan import GFPGANer
-import os
+from insightface.app import FaceAnalysis
+from insightface.model_zoo import model_zoo
+
+# Paths
+MODEL_DIR = "model"
+GFPGAN_MODEL_PATH = os.path.join(MODEL_DIR, "GFPGANv1.4.pth")
+INSWAPPER_PATH = os.path.join(MODEL_DIR, "inswapper_128.onnx")
 
-# Setup paths
-swapper_model_path = "models/inswapper_128.onnx"
-gfpgan_model_path = "models/GFPGANv1.4.pth"
-gfpganer = GFPGANer(model_path=gfpgan_model_path, upscale=1, arch='clean', channel_multiplier=2)
+# Initialize GFPGAN
+gfpganer = GFPGANer(
+    model_path=GFPGAN_MODEL_PATH,
+    upscale=2,
+    arch='clean',
+    channel_multiplier=2,
+    bg_upsampler=None
+)
 
-# Initialize FaceAnalysis and INSwapper
-face_app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
-face_app.prepare(ctx_id=0, det_size=(640, 640))
-swapper = INSwapper(model_file=swapper_model_path)
+# Initialize InsightFace (for face detection + swapping)
+faceapp = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
+faceapp.prepare(ctx_id=0)
 
-def enhance_face(img_np):
-    _, _, output = gfpganer.enhance(img_np, has_aligned=False, only_center_face=False, paste_back=True)
-    return output
+swapper = model_zoo.get_model(INSWAPPER_PATH, providers=['CPUExecutionProvider'])
 
-def swap_faces(source_img, target_img):
-    source = np.array(source_img)
-    target = np.array(target_img)
+# Function: Face Restoration + Swapping
+def restore_and_swap(input_img: Image.Image, target_img: Image.Image):
+    input_np = np.array(input_img.convert("RGB"))
+    target_np = np.array(target_img.convert("RGB"))
 
-    source_faces = face_app.get(source)
-    target_faces = face_app.get(target)
+    # Step 1: Face restore
+    cropped_faces, restored_img = gfpganer.enhance(input_np, has_aligned=False, only_center_face=False, paste_back=True)
 
-    if len(source_faces) == 0 or len(target_faces) == 0:
-        return target
+    # Step 2: Detect face in target
+    target_faces = faceapp.get(target_np)
+    if len(target_faces) == 0:
+        return Image.fromarray(restored_img), "No face found in target image."
 
-    source_face = source_faces[0]
-    for face in target_faces:
-        target = swapper.get(target, face, source_face)
+    # Step 3: Detect face in restored image
+    source_faces = faceapp.get(restored_img)
+    if len(source_faces) == 0:
+        return Image.fromarray(restored_img), "No face found in restored image."
 
-    target = enhance_face(target)
-    return target
+    # Step 4: Face swap using first face
+    swapped = swapper.get(restored_img, source_faces[0], target_faces[0], paste_back=True)
 
-def interface(source, target):
-    result = swap_faces(source, target)
-    return result
+    return Image.fromarray(swapped), "Success! Face restored and swapped."
 
+# Gradio Interface
 iface = gr.Interface(
-    fn=interface,
+    fn=restore_and_swap,
     inputs=[
-        gr.Image(label="Source Face", type="numpy"),
-        gr.Image(label="Target Image", type="numpy")
+        gr.Image(type="pil", label="Input Image (to restore + swap)"),
+        gr.Image(type="pil", label="Target Image (face to apply)")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Output Image"),
+        gr.Textbox(label="Status")
     ],
-    outputs=gr.Image(label="Swapped Output"),
-    title="Face Swap with Enhancement",
-    description="Upload a source face and a target image. The app swaps the face and enhances it using GFPGAN."
+    title="GFPGAN + InsightFace Face Swapper",
+    description="This app restores old faces using GFPGAN and swaps them with another face using InsightFace InSwapper."
 )
 
-iface.launch()
+iface.launch()