app.py

# app.py

import gradio as gr
import cv2
import numpy as np
from PIL import Image
import os
import tempfile
from insightface.app import FaceAnalysis
from insightface.model_zoo.inswapper import INSwapper
from skimage.exposure import match_histograms
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from gfpgan import GFPGANer

# Setup face analyzer
face_analyzer = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
face_analyzer.prepare(ctx_id=0, det_size=(640, 640))

# Load face swapper
swapper = INSwapper("models/inswapper_128.onnx")

# Load GFPGAN and Real-ESRGAN
bg_upsampler = RealESRGANer(
    scale=2,
    model_path='models/RealESRGAN_x2plus.pth',
    model=RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
                  num_block=23, num_grow_ch=32, scale=2),
    tile=400, tile_pad=10, pre_pad=0, half=False
)

face_enhancer = GFPGANer(
    model_path='models/GFPGANv1.3.pth',
    upscale=2,
    arch='clean',
    channel_multiplier=2,
    bg_upsampler=bg_upsampler
)

def match_histogram(source_face, target_face):
    return match_histograms(source_face, target_face, channel_axis=-1)

def swap_faces(source_img, target_img, face_index_str, blend_alpha, enhance_face):
    if source_img is None or target_img is None:
        raise gr.Error("Please upload both source and target images.")

    source_np = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
    target_np = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)

    source_faces = face_analyzer.get(source_np)
    target_faces = face_analyzer.get(target_np)

    if len(source_faces) == 0:
        raise gr.Error("No face detected in source image.")
    if len(target_faces) == 0:
        raise gr.Error("No faces detected in target image.")

    face_index = int(face_index_str)
    if face_index >= len(target_faces):
        raise gr.Error(f"Face index {face_index} out of range (found {len(target_faces)} faces).")

    try:
        swapped = swapper.get(target_np, target_faces[face_index], source_faces[0], paste_back=True)

        box = target_faces[face_index].bbox.astype(int)
        x1, y1, x2, y2 = box
        original_face = target_np[y1:y2, x1:x2]
        swapped_face = swapped[y1:y2, x1:x2]
        matched = match_histogram(swapped_face, original_face)
        swapped[y1:y2, x1:x2] = matched.astype(np.uint8)

        mask = np.zeros_like(swapped[:, :, 0])
        cv2.rectangle(mask, (x1, y1), (x2, y2), 255, -1)
        mask = cv2.GaussianBlur(mask, (15, 15), 0)
        blended = np.where(mask[..., None] > 0,
                           cv2.addWeighted(swapped, blend_alpha, target_np, 1 - blend_alpha, 0),
                           target_np)

        if enhance_face:
            _, _, enhanced = face_enhancer.enhance(blended, has_aligned=False, only_center_face=False, paste_back=True)
            blended = enhanced

        final_rgb = cv2.cvtColor(blended, cv2.COLOR_BGR2RGB)
        output_img = Image.fromarray(final_rgb)

        tmp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
        output_img.save(tmp_file.name)
        return output_img, tmp_file.name, "✅ Swap & Enhancement Completed!"
    except Exception as e:
        raise gr.Error(f"Swap failed: {str(e)}")

def count_faces(image):
    if image is None:
        return "Please upload a target image."
    img_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    faces = face_analyzer.get(img_np)
    return f"{len(faces)} face(s) detected in the target image."

# Gradio UI
with gr.Blocks(title="💫 Face Swap & Enhance App") as demo:
    gr.Markdown("""
    # 🔄 Face Swapper + 👩‍🔧 GFPGAN Enhancer + 🧙 Real-ESRGAN
    Upload a face and swap it into another photo. Optional face enhancement with GFPGAN!
    """)

    with gr.Row():
        source_input = gr.Image(label="🧑 Source Face", type="pil")
        target_input = gr.Image(label="🧑 Target Image", type="pil")

    detect_button = gr.Button("🔍 Detect Faces in Target")
    face_count = gr.Textbox(label="Detected Faces", interactive=False)
    detect_button.click(count_faces, inputs=target_input, outputs=face_count)

    with gr.Row():
        face_index = gr.Textbox(label="🎯 Face Index", value="0")
        blend_slider = gr.Slider(label="💫 Blend Smoothness", minimum=0, maximum=1, value=0.8, step=0.05)
        enhance_face = gr.Checkbox(label="✨ Enhance Face with GFPGAN", value=True)

    swap_button = gr.Button("🔁 Run Face Swap")
    output_img = gr.Image(label="🖼️ Output")
    download_path = gr.File(label="📥 Download")
    status = gr.Textbox(label="Status", interactive=False)

    swap_button.click(swap_faces, inputs=[source_input, target_input, face_index, blend_slider, enhance_face],
                      outputs=[output_img, download_path, status])

if __name__ == '__main__':
    demo.launch()