|
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) |
|
|
|
|
|
torch.cuda.empty_cache() |
|
del(upsampler) |
|
del(upscale_model) |
|
gc.collect() |
|
return Image.fromarray(output) |
|
|
|
with gr.Blocks() as 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) |
