import gradio as gr
import numpy as np
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url
from realesrgan import RealESRGANer
import gc
import os
import torch
from PIL import Image

def upscale_image(img: np.ndarray) -> np.ndarray:
    upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
    upsampler = RealESRGANer(
        scale=4,
        model_path=os.path.join('weights', 'RealESRGAN_x4plus.pth'),
        dni_weight=None,
        model=upscale_model,
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=False,
    )
    output, _ = upsampler.enhance(img, outscale=4)

    #cleanup
    torch.cuda.empty_cache()
    del(upsampler)
    del(upscale_model)
    gc.collect()
    return Image.fromarray(output)

with gr.Blocks() as demo:
    #gr.Markdown("Flip text or image files using this demo.")
    with gr.Tab("Upscale Image"):
        with gr.Row():
            image_input = gr.Image()
            image_output = gr.Image()
        image_button = gr.Button("Upscale")
        gr.Markdown("this is running on a CPU so it's gonna be VERY slow")

    gr.Markdown("Real esrgan implementation from [this repo](https://github.com/xinntao/Real-ESRGAN)")
    gr.Markdown("The idea behind this space was that you can clone it and use this model locally on your GPU")
    image_button.click(upscale_image, inputs=image_input, outputs=image_output)

if __name__ == "__main__":
    demo.launch(max_threads = 16)