import gradio as gr
from PIL import Image
import numpy as np
from realesrgan import RealESRGANer
import os

# Проверка наличия модели
def check_model_path(model_path):
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Модель не найдена по пути: {model_path}. Скачайте модель и поместите её в папку 'weights'.")

# Инициализация модели Real-ESRGAN
def load_model():
    model_path = "weights/RealESRGAN_x4plus.pth"  # Путь к модели
    check_model_path(model_path)  # Проверка наличия модели
    
    model = RealESRGANer(
        scale=4,
        model_path=model_path,
        model=None,
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=False
    )
    return model

# Функция для увеличения разрешения изображения
def upscale_image(image):
    img = np.array(image)
    output, _ = model.enhance(img, outscale=4)
    return Image.fromarray(output)

# Загрузка модели
model = load_model()

# Создание интерфейса Gradio
iface = gr.Interface(
    fn=upscale_image,
    inputs=gr.Image(type="pil", label="Входное изображение"),
    outputs=gr.Image(type="pil", label="Увеличенное изображение"),
    title="AI Image Upscaler",
    description="Загрузите изображение, чтобы увеличить его разрешение с помощью Real-ESRGAN."
)

# Запуск интерфейса
iface.launch()