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

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

model = None
scale = 4

def upscale_image(image, scale):
    global device
    global model
    
    if model is None or model.scale != scale:
        model = RealESRGAN(device, scale=scale)
        model.load_weights('weights/RealESRGAN_x{}.pth'.format(scale), download=True)
    
    sr_image = model.predict(image)
    return sr_image

inputs = [
    gr.Image(label="Input Image", type="pil"),
    gr.Slider(minimum=2, maximum=8, value=4, label="Upscale Scale", step=2)
]

output = gr.Image(label="Upscaled Image")

examples = [
    ['groot.jpg', '4'],
    ['woman.jpg', '8']
]

title = "Image upscaler 2k, 4k, 8k"

gr.Interface(fn=upscale_image, title=title inputs=inputs, outputs=output, title="Image Upscaler",
             examples=examples, theme="soft").launch()