import spaces
import gradio as gr
import cv2
import numpy as np


def _to_bgr(img):
    if img is None:
        raise ValueError("no image received")
    if isinstance(img, str):
        im = cv2.imread(img)
        if im is None:
            raise ValueError("unreadable path: %s" % img)
        return im
    if hasattr(img, "convert"):
        return cv2.cvtColor(np.array(img.convert("RGB")), cv2.COLOR_RGB2BGR)
    a = np.array(img)
    if a.ndim == 3 and a.shape[2] >= 3:
        return cv2.cvtColor(a[:, :, :3], cv2.COLOR_RGB2BGR)
    return a


@spaces.GPU
def swap_face(source_img, target_img):
    import traceback
    import insightface
    from insightface.app import FaceAnalysis
    try:
        src = _to_bgr(source_img)
        tgt = _to_bgr(target_img)
        app = FaceAnalysis(name="buffalo_l")
        try:
            app.prepare(ctx_id=0, det_size=(640, 640))
        except Exception:
            app.prepare(ctx_id=-1, det_size=(640, 640))
        import insightface.model_zoo
        import os as _os, urllib.request as _url
        _mdir="/tmp/ifmodels"; _os.makedirs(_mdir, exist_ok=True)
        _sp=_os.path.join(_mdir,"inswapper_128.onnx")
        if not _os.path.exists(_sp):
            _url.urlretrieve("https://github.com/deepinsight/insightface/releases/download/v0.7/inswapper_128.onnx", _sp)
        swapper = insightface.model_zoo.get_model(_sp)
        s_faces = app.get(src)
        t_faces = app.get(tgt)
        if len(s_faces) == 0 or len(t_faces) == 0:
            app2 = FaceAnalysis(name="buffalo_l")
            app2.prepare(ctx_id=-1, det_size=(1024, 1024))
            if len(s_faces) == 0:
                s_faces = app2.get(src)
            if len(t_faces) == 0:
                t_faces = app2.get(tgt)
        if len(s_faces) == 0 or len(t_faces) == 0:
            raise gr.Error("no faces detected s:%d t:%d" % (len(s_faces), len(t_faces)))
        res = swapper.get(tgt, t_faces[0], s_faces[0], paste_back=True)
        out = "/tmp/swapped.jpg"
        cv2.imwrite(out, res)
        return out
    except gr.Error:
        raise
    except Exception:
        raise gr.Error(traceback.format_exc()[-1500:])


with gr.Blocks() as demo:
    gr.Markdown("# FaceFusion Zero")
    with gr.Row():
        src = gr.Image(label="Source")
        tgt = gr.Image(label="Target")
    out = gr.Image(label="Result")
    btn = gr.Button("Swap")
    btn.click(swap_face, inputs=[src, tgt], outputs=out)

demo.launch()