# Code taken (and slightly adopted) from https://huggingface.co/spaces/havas79/Real-ESRGAN_Demo/blob/main/app.py - credit where credit is due. I am not showcasing code here, but demoing my own trained models ;)

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

last_file = None
img_mode = "RGBA"

def realesrgan(img, model_name, face_enhance):
    global last_file

    # remove last upscale when doing this new upscale to prevent memory being full   
    if last_file:
        print(f"Deleting {last_file} ...")
        os.remove(last_file)
        last_file = None
            
    if not img:
        return

    imgwidth, imgheight = img.size

    if imgwidth > 1000 or imgheight > 1000:
        return error("Input Image too big")

    # Define model parameters
    if model_name == '2x Fast Upscale':
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')
        netscale = 2
        model_name = '2xNomosUni_compact_multijpg'
    elif model_name == '2x Upscale':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
        netscale = 2
        model_name = '2xNomosUni_esrgan_multijpg'

    # Determine model paths
    model_path = os.path.join('weights', model_name + '.pth')

    # Restorer Class
    upsampler = RealESRGANer(
        scale=netscale,
        model_path=model_path,
        dni_weight=None,
        model=model,
        tile=128,
        tile_pad=10,
        pre_pad=10,
        half=False,
        gpu_id=None,
    )

    # Use GFPGAN for face enhancement
    if face_enhance:
        from gfpgan import GFPGANer
        face_enhancer = GFPGANer(
            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth',
            upscale=netscale,
            arch='clean',
            channel_multiplier=2,
            bg_upsampler=upsampler)

    # Convert the input PIL image to cv2 image, so that it can be processed by realesrgan
    cv_img = numpy.array(img)
    img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)

    # Apply restoration
    try:
        if face_enhance:
            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
        else:
            output, _ = upsampler.enhance(img, netscale)
    except RuntimeError as error:
        print('Error', error)
        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
    else:
        # Save restored image and return it to the output Image component
        extension = 'jpg'
        out_filename = f"output_{rnd_string(16)}.{extension}"
        cv2.imwrite(out_filename, output)
        last_file = out_filename
        return out_filename
        

def rnd_string(x):
    """Returns a string of 'x' random characters
    """
    characters = "abcdefghijklmnopqrstuvwxyz_0123456789"
    result = "".join((random.choice(characters)) for i in range(x))
    return result


def reset():
    """Resets the Image components of the Gradio interface and deletes
    the last processed image
    """
    global last_file
    if last_file:
        print(f"Deleting {last_file} ...")
        os.remove(last_file)
        last_file = None
    return gr.update(value=None), gr.update(value=None)


def has_transparency(img):
    """This function works by first checking to see if a "transparency" property is defined
    in the image's info -- if so, we return "True". Then, if the image is using indexed colors
    (such as in GIFs), it gets the index of the transparent color in the palette
    (img.info.get("transparency", -1)) and checks if it's used anywhere in the canvas
    (img.getcolors()). If the image is in RGBA mode, then presumably it has transparency in
    it, but it double-checks by getting the minimum and maximum values of every color channel
    (img.getextrema()), and checks if the alpha channel's smallest value falls below 255.
    https://stackoverflow.com/questions/43864101/python-pil-check-if-image-is-transparent
    """
    if img.info.get("transparency", None) is not None:
        return True
    if img.mode == "P":
        transparent = img.info.get("transparency", -1)
        for _, index in img.getcolors():
            if index == transparent:
                return True
    elif img.mode == "RGBA":
        extrema = img.getextrema()
        if extrema[3][0] < 255:
            return True
    return False


def image_properties(img):
    """Returns the dimensions (width and height) and color mode of the input image and
    also sets the global img_mode variable to be used by the realesrgan function
    """
    global img_mode
    if img:
        if has_transparency(img):
            img_mode = "RGBA"
        else:
            img_mode = "RGB"
        properties = f"Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
        return properties


def main():
    # Gradio Interface
    with gr.Blocks(title="Self-trained ESRGAN models demo", theme="dark") as demo:
        
        gr.Markdown(
            """# <div align="center"> Upscale Image </div> 
        Here I demo two of my self-trained models which can be used to upscale an image.
        """
        )

        with gr.Group():
            with gr.Group():
                model_name = gr.Dropdown(label="Model to be used",
                                         choices=["2x Fast Upscale", "2x Upscale"], value="2x Fast Upscale", 
                                         info="See model infos at the bottom of this page")
                face_enhance = gr.Checkbox(label="Face Enhancement using GFPGAN (Doesn't work for anime images)",value=False, show_label=True)

        with gr.Group():
                input_image = gr.Image(label="Source Image", type="pil", image_mode="RGB")
                input_image_properties = gr.Textbox(label="Image Properties - Demo will throw error if input image has either width or height > 1000. Output download is jpg for smaller size. Use models locally to circument these limits.", max_lines=1)
        with gr.Group():
            output_image = gr.Image(label="Upscaled Image", type="pil", image_mode="RGB", interactive=False)
            output_image_properties = gr.Textbox(label="Image Properties", max_lines=1)
        with gr.Row():
            upscale_btn = gr.Button("Upscale")
            reset_btn = gr.Button("Reset")
        with gr.Group():
            gr.Markdown(""" **Examples are not pre-cached. You need to press the Upscale Button after selecting one**""")
            gr.Examples(examples="examples",inputs=[input_image, model_name, face_enhance],outputs=output_image,fn=realesrgan, cache_examples=False)
            gr.Markdown(
    """
    **Details**
    These two 2x models are a universal Compact and ESRGAN upscaling model which I trained and released dec 23, called
    2x Fast Upscale: 2xNomosUni_compact_multijpg
    2x Upscale: 2xNomosUni_esrgan_multijpg

    I trained these using musl's [neosr](https://github.com/muslll/neosr) and Kim's [Dataset Destroyer](https://github.com/Kim2091/helpful-scripts/tree/main/Dataset%20Destroyer) on musl's universal nomos_uni dataset.

    You can find more information about upscaling model training on the [training info repo](https://github.com/Upscale-Community/training-info). If you have questions or simply be up to date on new upscaling models released can of course also join our upscaling discord community [Enhance Everything](https://discord.gg/enhance-everything-547949405949657098) and watch the model-releases channel.


    You can also run these two and way more models locally on your PC with [chaiNNer](https://github.com/chaiNNer-org/chaiNNer).
    This is a [google drive folder](https://drive.google.com/drive/folders/1coYgan0VDo578MVO1LUpjpsxdY3LMyJW?usp=drive_link) with most of my self trained models.
    Otherwise you can also find more models on our [Open Model Database](https://openmodeldb.info/)

    I called myself 'Helaman' when joining the discord server because that was my gamer name I picked when joining games like league of legends, so I just used it for discord in general and releasing models with it.
    My real name is [Philip Hofmann](https://github.com/Phhofm), got into upscaling after trying out Midjourney when it came out, then found chaiNNer, then found the old database where there were only textual descriptions of upsclaing models,
    but no visual outputs, this is why I gatheres a lot of upscaling models, used them, then created a [youtube vid](https://youtu.be/0TYRDmQ5LZk) in oct 22, made [reddit posts](https://www.reddit.com/r/StableDiffusion/comments/yev37i/comparison_of_upscaling_models_for_ai_generated/), made a [vitepress website](https://phhofm.github.io/upscale/) to compare the visual outputs of over 300 different upsaling models
    then thought about training models myself instead of just using and comparing existing ones, so I got into upscale training, made [youtube videos](https://youtu.be/l_tZE5l90VY) about it, more [reddit posts](https://www.reddit.com/r/StableDiffusion/comments/17d88rt/selftrained_sd_upscaling_models/), released my self trained models in discord, helped with testing neosr, trained different networks/architectures like SRFormer, [DAT](https://github.com/zhengchen1999/DAT) where my model got features on the devs readme :D and so forth, just a lot of experiments/models, probably released around 50 models so far, it has been fun and fascinating.

    To keep up on latest sisr (single image super resolution) networks (architectures) one can also follow the [papers with code image super resolution task](https://paperswithcode.com/task/image-super-resolution/latest) where interesting papers with code bases get published frequently. Also the [Awesome Image Super Resolution Repo](https://github.com/ChaofWang/Awesome-Super-Resolution). Its a fascinating and huge field with a lot of stuff to learn about.
    
    Face enhance: [GFPGANv1.4](https://github.com/TencentARC/GFPGAN)
    """)

        # Event listeners:
        input_image.change(fn=image_properties, inputs=input_image, outputs=input_image_properties)
        output_image.change(fn=image_properties, inputs=output_image, outputs=output_image_properties)
        upscale_btn.click(fn=realesrgan, inputs=[input_image, model_name, face_enhance], outputs=output_image)
        reset_btn.click(fn=reset, inputs=[], outputs=[output_image, input_image])

    demo.launch()


if __name__ == "__main__":
    main()