import gradio as gr
import numpy as np
from PIL import Image
import os
os.system('pip install basicsr')
os.system('pip install realesrgan')
from gfpgan import GFPGANer

# installing version 1 of GFPGAN
os.system('wget https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/GFPGANv1.pth')
# installing version 1.2 of GFPGAN
os.system('wget https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth')
# installing version 1.3 of GFPGAN (latest)
os.system('wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth')


def interface(image: Image, model: str = "GFPGANv1.3.pth", useRealesrgan=False):
    if model == "":
        model = "GFPGANv1.3.pth"
    if model != "GFPGANv1.pth" and model != "GFPGANCleanv1-NoCE-C2.pth" and model != "GFPGANv1.3.pth":
        model = "GFPGANv1.3.pth"
    if useRealesrgan == True:
        from basicsr.archs.rrdbnet_arch import RRDBNet
        from realesrgan import RealESRGANer
        BGupscaler = RealESRGANer(
            scale=2,
            model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
            model=RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2),
            tile=0,
            tile_pad=10,
            pre_pad=0
        )
    else:
        BGupscaler = None
    restorer = GFPGANer(
        model_path=model,
        arch="original" if model == "GFPGANv1.pth" else "clean",
        bg_upsampler=BGupscaler,
        channel_multiplier=1 if model == "GFPGANv1.pth" else 2,
        upscale=2)
    img = np.array(image).copy()
    cropped_faces, restored_faces, restored_img = restorer.enhance(img)
    return restored_img


gr.Interface(
    interface,
    [
        gr.components.Image(
            type="pil",
            label="Image",
        ),
        gr.components.Radio([
            "GFPGANv1.pth",
            "GFPGANCleanv1-NoCE-C2.pth",
            "GFPGANv1.3.pth",
        ],
                            label="model",
                            default="GFPGANv1.3.pth",
                            type="value"),
        gr.Checkbox(label="realesrgan?"),
    ],
    [gr.components.Image(label="Enhanced Image")],
).launch()