import torch
from PIL import Image
import gradio as gr
import download_models  # تحميل الموديلات تلقائيًا

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# قائمة الموديلات المتاحة
MODELS = {
    "RealESRGAN_x4plus": "weights/RealESRGAN_x4plus.pth",
    "RealESRGAN_x4plus_anime_6B": "weights/RealESRGAN_x4plus_anime_6B.pth",
    "realesr-animevideov3-x2": "weights/realesr-animevideov3-x2.pth",
}

def upscale_image(image, model_name, quality, output_size):
    # تحميل الموديل المختار
    model = RealESRGAN(device, scale=4)
    model.load_weights(MODELS[model_name], download=False)

    # تحديد إعدادات الجودة
    settings = {
        "Fast": {"tile": 512, "tile_pad": 10, "pre_pad": 0},
        "Balanced": {"tile": 256, "tile_pad": 5, "pre_pad": 5},
        "High Quality": {"tile": 128, "tile_pad": 0, "pre_pad": 10},
    }
    
    model.tile = settings[quality]["tile"]
    model.tile_pad = settings[quality]["tile_pad"]
    model.pre_pad = settings[quality]["pre_pad"]

    # تشغيل الموديل
    sr_image = model.predict(image)

    # تغيير حجم الصورة النهائية حسب اختيار المستخدم
    sr_image = sr_image.resize(output_size, Image.ANTIALIAS)

    return sr_image

# واجهة Gradio
interface = gr.Interface(
    fn=upscale_image,
    inputs=[
        gr.Image(type="pil"),
        gr.Radio(choices=list(MODELS.keys()), label="Choose Model", value="RealESRGAN_x4plus"),
        gr.Radio(choices=["Fast", "Balanced", "High Quality"], label="Processing Speed", value="Balanced"),
        gr.Slider(256, 2048, step=128, label="Output Size (px)", value=1024),
    ],
    outputs=gr.Image(type="pil"),
    title="AI Image & Video Upscaler",
    description="Upload an image to upscale it using Real-ESRGAN. Supports anime, videos, and normal images.",
)

interface.launch()