# Code taken and  adopted from https://huggingface.co/spaces/havas79/Real-ESRGAN_Demo/blob/main/app.py - credit where credit is due.

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

# global variables
last_file = None
img_mode = "RGB"

# Upscale function
def upscale(img, choice):
    
    # variables
    global last_file
    model_path = ""

    # remove last upscale when doing this new upscale to prevent memory being full   
    if last_file:
        os.remove(last_file)
        last_file = None
        
    # There is no input image to upscale
    if not img:
        return error("Input Image not detected")
    
    # Get image dimenstions
    imgwidth, imgheight = img.size

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

    # Define model parameters
    if choice == '2x Fast Upscale':
        model_path = os.path.join('weights', '2xNomosUni_compact_multijpg.pth')
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')
        
    elif choice == '2x Upscale':
        model_path = os.path.join('weights', '2xNomosUni_esrgan_multijpg.pth')
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)

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

    # 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:
        output, _ = upsampler.enhance(img, 2)
    except RuntimeError as error:
        print('Error when upscaling', error)
    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
        
# Get random image file name for newly created image
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

# Reset function to reset inputs and also delete last file
def reset():
    """
    Resets the Image components of the Gradio interface and deletes
    the last processed image
    """
    global last_file
    if last_file:
        os.remove(last_file)
        last_file = None
    return gr.update(value=None), gr.update(value=None)

# Check for transparency function
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

# Get image properties function
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

# Gradio Interface, Event Listeners, launch command
def main():
    # Gradio Interface
    with gr.Blocks(title="Self-trained 2x general upscaler models") as demo:
        
        gr.Markdown(
            """# <div align="center"> Upscale Image </div> 
            Here I demo two of my self-trained general 2x upscaler models which can be used to upscale an image. Use the Fast option, and if you like the result, you can also try out the normal upscale, which might give a better result, but will also take longer to upscale.
            """
        )

        with gr.Group():
            choice = gr.Radio(choices=["2x Fast Upscale", "2x Upscale"], label="Model Selection", info="See infos at the bottom of this page", value="2x Fast Upscale")

        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 these 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, choice], outputs=output_image, fn=upscale, cache_examples=False)
            gr.Markdown(
                """
                **Details**
                These two 2x models are a Compact(SRVGGNet) for the 'Fast Upscale' and an ESRGAN(RRDBNet) for the 'Upscale' upscaling model which I trained and released dec 23, called  
                2x Fast Upscale: [2xNomosUni_compact_multijpg](https://openmodeldb.info/models/2x-NomosUni-compact-multijpg)  
                2x Upscale: 2xNomosUni_esrgan_multijpg  
                
                These two models are general upscalers which can handle jpg compression and preserve depth of field for the most part.  
                
                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=upscale, inputs=[input_image, choice], outputs=output_image)
        reset_btn.click(fn=reset, inputs=[], outputs=[input_image, output_image])
            
    demo.launch()

if __name__ == "__main__":
    main()