import gradio as gr
import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
from PIL import Image
# Tambahkan pustaka resmi Hugging Face untuk mengunduh/membaca berkas repositori
from huggingface_hub import hf_hub_download

# 1. Inisialisasi model Face Analysis (Menggunakan CPU)
app = FaceAnalysis(name='buffalo_l')
app.prepare(ctx_id=-1, det_size=(640, 640)) 

# 2. MENGAMBIL BERKAS DARI REPOSITORI SPACE ANDA SENDIRI
# Kode ini akan otomatis mengambil berkas lokal jika sudah ada di Space
model_path = hf_hub_download(
    repo_id="Dentro/face-swap",
    filename="inswapper_128.onnx",
    repo_type="space"
)

# Memuat model face swapper menggunakan jalur berkas yang aman dari HF Hub
swapper = insightface.model_zoo.get_model(model_path, download=False)

def swap_faces(source_img, target_img):
    if source_img is None or target_img is None:
        return None
    
    src_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
    tgt_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
    
    src_faces = app.get(src_img)
    tgt_faces = app.get(tgt_img)
    
    if len(src_faces) == 0 or len(tgt_faces) == 0:
        return target_img 
    
    source_face = src_faces[0]
    
    result_img = tgt_img.copy()
    for face in tgt_faces:
        result_img = swapper.get(result_img, face, source_face, paste_back=True)
        
    result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
    return Image.fromarray(result_img)

# Tampilan GUI Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🔄 Face Swap App (HF Hub File Storage)")
    
    with gr.Row():
        with gr.Column():
            src_input = gr.Image(type="pil", label="Wajah Sumber")
            tgt_input = gr.Image(type="pil", label="Gambar Target")
            submit_btn = gr.Button("Mulai Swap Wajah", variant="primary")
        
        with gr.Column():
            output_img = gr.Image(label="Hasil Swap")
            
    submit_btn.click(fn=swap_faces, inputs=[src_input, tgt_input], outputs=output_img)

demo.launch()