Create app.py

app.py

@@ -0,0 +1,145 @@
+# mostly borrowed from TheStinger/Ilaria_Upscaler spaces.
+
+import gradio as gr
+import cv2
+import numpy
+import os
+import random
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.download_util import load_file_from_url
+
+
+from realesrgan import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+
+
+def model_params(model_name):
+    if model_name == 'RealESRGAN_x4plus':  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
+    elif model_name == 'RealESRNet_x4plus':  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
+    elif model_name == 'RealESRGAN_x4plus_anime_6B':  # x4 RRDBNet model with 6 blocks
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
+    elif model_name == 'RealESRGAN_x2plus':  # x2 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        netscale = 2
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
+    elif model_name == 'realesr-general-x4v3':  # x4 VGG-style model (S size)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+        netscale = 4
+        file_url = [
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
+        ]
+
+    model, netscale, file_url
+
+
+def upscale(image, model_name, tile, denoise, face_enhance, scale):
+    if not image: return
+
+    model, netscale, file_url = model_params(model_name)
+
+    model_path = os.path.join('weights', model_name + '.pth')
+    if not os.path.isfile(model_path):
+        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+        for url in file_url:
+            # model_path will be updated
+            model_path = load_file_from_url(
+                url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
+
+    dni_weight = None
+    if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
+        wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
+        model_path = [model_path, wdn_model_path]
+        dni_weight = [denoise_strength, 1 - denoise_strength]
+
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=tile,
+        tile_pad=10,
+        pre_pad=10,
+        half=False,
+        gpu_id=None
+    )
+
+    if face_enhance:
+        from gfpgan import GFPGANer
+        face_enhancer = GFPGANer(
+            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
+            upscale=outscale,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=upsampler)
+
+    ################
+
+    cv_img = numpy.array(image)
+    img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
+
+    try:
+        if face_enhance:
+            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+        else:
+            output, _ = upsampler.enhance(img, outscale=outscale)
+    except RuntimeError as error:
+        print('Error', error)
+        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
+
+    return output
+
+
+app = gr.Interface(
+    title='Real-ESRGAN Upscaler',
+    description='Yet another Real-ESRGAN upscaler that uses gradio `Interface`, because why not? It’s not like there are any other options for simplicity and backward compatibility. Oh wait, there are. Never mind.',
+    fn=upscale,
+    inputs=[
+        gr.Image(label='Source Image', type='pil', image_mode='RGBA'),
+        gr.Dropdown(
+            label='Model',
+            choices=["RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B","RealESRGAN_x2plus", "realesr-general-x4v3"],
+            show_label=True,
+            value='RealESRGAN_x4plus_anime_6B'
+        ),
+        gr.Slider(
+            label='Tile',
+            minimum=0,
+            maximum=1024,
+            step=32,
+            value=256
+        ),
+        gr.Slider(
+            label='Denoise Strength',
+            minimum=0,
+            maximum=1,
+            step=0.1,
+            value=0.5
+        ),
+        gr.Checkbox(
+            label='Face Enhancement (GFPGAN)',
+            value=False,
+            show_label=True
+        ),
+        gr.Slider(
+            minimum=1,
+            maximum=4,
+            step=1,
+            value=2,
+            show_label=True
+        )
+    ],
+    outputs=[
+        gr.Image(label='Upscaled Image', image_mode='RGBA')
+    ]
+)
+
+app.launch(show_api=True)