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