import gradio as gr
import os
import uuid
import cv2
from mask import FaceSwapper  # الكلاس اللي فوق

# تحميل الموديل مرة واحدة
swapper = FaceSwapper(
    model_path="models/inswapper_128.onnx",
    gfpgan_model_path="gfpgan/weights/GFPGANv1.4.pth"
)

# دالة للتعامل مع Gradio
def swap_faces(source_img, target_img):
    try:
        source_path = f"temp_source_{uuid.uuid4().hex}.jpg"
        target_path = f"temp_target_{uuid.uuid4().hex}.jpg"
        output_path = f"img/result_{uuid.uuid4().hex}.jpg"

        cv2.imwrite(source_path, cv2.cvtColor(source_img, cv2.COLOR_RGB2BGR))
        cv2.imwrite(target_path, cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR))

        # دلوقتي الموديل بيرجع numpy مش path
        result_img = swapper.merge_face_into_image(source_path, target_path, output_path)

        # مسح الملفات المؤقتة
        if os.path.exists(source_path):
            os.remove(source_path)
        if os.path.exists(target_path):
            os.remove(target_path)

        return result_img  # numpy جاهز للعرض
    except Exception as e:
        return f"❌ حصل خطأ: {str(e)}"

# واجهة Gradio
demo = gr.Interface(
    fn=swap_faces,
    inputs=[
        gr.Image(type="numpy", label="Source Image (الطفل)"),
        gr.Image(type="numpy", label="Target Image (المشهد)")
    ],
    outputs=gr.Image(type="numpy", label="النتيجة"),
    title="FaceSwap with GFPGAN",
    description="ارفع صورتين: (1) صورة الطفل (2) المشهد اللي عايز تدخله فيه. وهنرجعلك صورة face swap محسّنة بـ GFPGAN."
)

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