import os
import warnings
from pathlib import Path

import gradio as gr
import numpy as np
from PIL import Image
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import get_image_colorizer
from huggingface_hub import snapshot_download

# تعطيل التحذيرات الغير ضرورية
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")

# تحميل نموذج DeOldify من Hugging Face
REPO_ID = "leonelhs/deoldify"
snapshot_folder = snapshot_download(repo_id=REPO_ID)

# تعيين الجهاز (CPU فقط، أو GPU لو متاح)
try:
    device.set(device=DeviceId.GPU0)  # تجربة استخدام GPU
except:
    device.set(device=DeviceId.CPU)  # الرجوع إلى CPU لو الـ GPU غير متاح

# تحميل النموذج
colorizer = get_image_colorizer(root_folder=Path(snapshot_folder), artistic=True)

def predict(image_path):
    """
    تحويل صورة الأبيض والأسود إلى صورة ملونة باستخدام DeOldify
    """
    try:
        # تشغيل النموذج على الصورة
        result = colorizer.get_transformed_image(image_path, render_factor=35, watermarked=False)

        # تحويل النتيجة إلى numpy array
        result_np = np.array(result)

        return result_np  # إرجاع الصورة الملونة
    except Exception as e:
        return f"Error: {str(e)}"

# إعداد Gradio API
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="filepath", label="Upload Black & White Image"),
    outputs=gr.Image(type="numpy", label="Colorized Image"),
    title="DeOldify - Image Colorization",
    description="Upload a grayscale image and get a colorized version using DeOldify.",
)

# تشغيل التطبيق
demo.queue().launch()