import gradio as gr
import numpy as np
from PIL import Image
import insightface
from insightface.app import FaceAnalysis

# ---------- Load models once ----------
app = FaceAnalysis(name="buffalo_l")
app.prepare(ctx_id=0, det_size=(640, 640))
swapper = insightface.model_zoo.get_model("inswapper_128.onnx",
                                          download=False, download_zip=False)

# ---------- HTML description ----------
description_html = """
    <h1>Face Swap</h1>
    <p>by <a href="https://www.tonyassi.com/" target="_blank" style="color:#97d8be;">Tony Assi</a></p>
    <p>Try out <a href="https://huggingface.co/spaces/tonyassi/video-face-swap" target="_blank">Video Face Swap</a> ❤️</p>
"""

# ---------- Swap logic ----------
def swap_faces(src_img, dest_img):
    src_faces = app.get(src_img)
    dest_faces = app.get(dest_img)

    res = dest_img.copy()
    if src_faces and dest_faces:
        res = swapper.get(res, dest_faces[0], src_faces[0], paste_back=True)

    return Image.fromarray(np.uint8(res)).convert("RGB")

# ---------- Interface ----------
gr.Interface(
    fn=swap_faces,
    inputs=[gr.Image(), gr.Image()],
    outputs=gr.Image(),
    description=description_html,
    examples=[
        ['./Images/kim.jpg', './Images/marilyn.jpg'],
        ['./Images/friends.jpg', './Images/friends.jpg']
    ],
).launch()