app.py

from base64 import b64encode, b64decode
from io import BytesIO
from pathlib import Path

import numpy as np
from basicsr.archs.rrdbnet_arch import RRDBNet
from PIL import Image
from realesrgan import RealESRGANer
from huggingface_hub import hf_hub_download
import gradio as gr
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
half = True if device == "cuda" else False
model = RRDBNet(num_in_ch=3, num_out_ch=3)
upsampler = RealESRGANer(
    scale=4,
    model_path=hf_hub_download('nateraw/real-esrgan', 'RealESRGAN_x4plus.pth'),
    model=model,
    tile=0,
    pre_pad=0,
    half=half,
)

def upsample(image, outscale=4):
    # image = np.array(image)
    image = image[:, :, ::-1]  # RGB -> BGR
    image, _ = upsampler.enhance(image, outscale=outscale)
    image = image[:, :, ::-1]  # BGR -> RGB
    image = Image.fromarray(image)
    return image

interface = gr.Interface(
    upsample,
    inputs=["image", gr.Dropdown([4, 2])],
    outputs=["image"]
)

if __name__ == '__main__':
    interface.launch()