import gradio as gr
from PIL import Image
import torch
from torchvision import transforms

# Загрузка предобученной модели ESRGAN через torch.hub
def load_model():
    model = torch.hub.load('pytorch/vision:v0.10.0', 'esrgan', pretrained=True)
    model.eval()
    return model

# Функция для увеличения разрешения изображения
def upscale_image(image):
    # Преобразование изображения в тензор
    input_image = transforms.ToTensor()(image).unsqueeze(0)
    
    # Увеличение разрешения с помощью модели
    with torch.no_grad():
        output_image = model(input_image)
    
    # Преобразование тензора обратно в изображение
    output_image = output_image.squeeze(0).clamp(0, 1)
    output_image = transforms.ToPILImage()(output_image)
    
    return output_image

# Загрузка модели
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 (ESRGAN)",
    description="Загрузите изображение, чтобы увеличить его разрешение с помощью ESRGAN."
)

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