import torch
from PIL import Image
from RealESRGAN import RealESRGAN
import gradio as gr

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_scales = {'2x': 2, '4x': 4, '8x': 8}

# Load RealESRGAN models for different scales
models = {scale: RealESRGAN(device, scale=scale) for scale in model_scales.values()}

def inference(images, scale):
    results = []
    
    if images is None or len(images) == 0:
        raise gr.Error("No image uploaded. Please upload at least one image.")
    
    for image in images:
        width, height = image.size
        if width >= 5000 or height >= 5000:
            raise gr.Error("The image is too large.")
    
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        # Select the appropriate model based on the chosen scale
        model = models[model_scales[scale]]
        result = model.predict(image.convert('RGB'))
        print(f"Image size ({device}): {scale} ... OK")
        results.append(result)
    
    return results

title = "Advanced Real ESRGAN UpScale: 2x 4x 8x"
description = (
    "This advanced demo for Real-ESRGAN allows you to upscale multiple images "
    "with different models and resolutions. Choose the scale and upload images for high-resolution enhancement."
)
article = (
    "<div style='text-align: center;'>Twitter "
    "<a href='https://twitter.com/DoEvent' target='_blank'>Max Skobeev</a> | "
    "<a href='https://huggingface.co/sberbank-ai/Real-ESRGAN' target='_blank'>Model card</a></div>"
)

gr.Interface(
    inference,
    [
        gr.Image(type="pil", label="Upload Image", multiple=True),
        gr.Radio(
            list(model_scales.keys()),
            type="value",
            value='2x',
            label='Resolution model',
        ),
    ],
    gr.Image(type="pil", label="Output"),
    title=title,
    description=description,
    article=article,
    examples=[['groot.jpeg', '2x']],
    allow_flagging='never',
    cache_examples=False,
).launch()