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

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

# ---------- Check version ----------
assert insightface.__version__ >= '0.7'

# ---------- Initialize models ----------
app = FaceAnalysis(name='buffalo_l')
app.prepare(ctx_id=0, det_size=(640, 640))

# ensure model downloads fully; safer with download=True and zip=True like your working version
swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True)

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

    if len(src_faces) == 0 or len(dest_faces) == 0:
        raise gr.Error("No faces detected in one of the images.")

    # Just swap first detected face from each image
    source_face = src_faces[0]
    dest_face = dest_faces[0]
    result = swapper.get(dest_img, dest_face, source_face, paste_back=True)

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

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