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')

def load_model(scale):
    model = RealESRGAN(device, scale=scale)
    model.load_weights(f'weights/RealESRGAN_x{scale}.pth', download=True)
    return model

def inference(image, size):
    if image is None:
        raise gr.Error("Image not uploaded")

    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()

    scale = int(size[0])
    model = load_model(scale)

    try:
        result = model.predict(image.convert('RGB'))
    except torch.cuda.OutOfMemoryError as e:
        print(e)
        model = load_model(scale)
        result = model.predict(image.convert('RGB'))

    print(f"Image size ({device}): {size} ... OK")
    return result

title = "RealESRGAN UpScale Model: 2x 4x 8x"
description = "This model running on CPU so it takes a bit of time, please be patient :)"

gr.Interface(
    inference,
    [gr.Image(type="pil"), gr.Radio(['2x', '4x', '8x'], type="value", value='2x', label='Resolution model')],
    gr.Image(type="pil", label="Output"),
    title=title,
    description=description,
    allow_flagging='never',
    cache_examples=False,
).queue(api_open=False).launch(show_error=True, show_api=False)