import gradio as gr
import spaces
from gradio_imageslider import ImageSlider
from image_gen_aux import UpscaleWithModel
from image_gen_aux.utils import load_image

# This uses https://github.com/asomoza/image_gen_aux/blob/main/src/image_gen_aux/upscalers/README.md
# Also this space has been duplicated from their official huggingface space, https://huggingface.co/spaces/OzzyGT/basic_upscaler
# They did great work, and I was happy to see them to also use my models :) I thought Id duplicate it and extend it.
# It basically made me get a pro account so I can make a Zero GPU space. And I will also upload more of my models as a model card now to use here.

# Start out with my own models. If others like kim, sirosky, and other model trainers would like their models added here, then thats great.
# I simply want them to message me first so I know that everythings okay with having their model as a selection here since they are the author of that model. If they want their model on here or not basically.
# I load models from huggingface model cards though, so the model should be hosted on huggingface.
MODELS = {
    "4xNomos2_hq_drct-l": "Phips/4xNomos2_hq_drct-l",
    "4xNomosWebPhoto_RealPLKSR": "Phips/4xNomosWebPhoto_RealPLKSR",
    "4xRealWebPhoto_v4_dat2": "Phips/4xRealWebPhoto_v4_dat2",
    "4xRealWebPhoto_v3_atd": "Phips/4xRealWebPhoto_v3_atd",
    "4xNomos8k_atd_jpg": "Phips/4xNomos8k_atd_jpg",
    "4xNomosUni_rgt_multijpg": "Phips/4xNomosUni_rgt_multijpg",
    "4xLSDIRDAT": "Phips/4xLSDIRDAT",
    "4xSSDIRDAT": "Phips/4xSSDIRDAT",
    "4xNomos8kHAT-L_otf": "Phips/4xNomos8kHAT-L_otf",
    "4xNomosUniDAT_otf": "Phips/4xNomosUniDAT_otf",
    "4xNomosUniDAT_bokeh_jpg": "Phips/4xNomosUniDAT_bokeh_jpg",
    "4xNomos8kSCHAT-L": "Phips/4xNomos8kSCHAT-L",
    "4xFFHQDAT": "Phips/4xFFHQDAT",
    "4xFaceUpDAT": "Phips/4xFaceUpDAT",
    "4xTextures_GTAV_rgt-s_dither": "Phips/4xTextures_GTAV_rgt-s_dither",
    "4xTextureDAT2_otf": "Phips/4xTextureDAT2_otf",
    "4xLexicaDAT2_otf": "Phips/4xLexicaDAT2_otf",
    "2xHFA2k_LUDVAE_compact": "Phips/2xHFA2k_LUDVAE_compact",
    "2xHFA2kAVCCompact": "Phips/2xHFA2kAVCCompact",
    "2xHFA2kCompact": "Phips/2xHFA2kCompact",
    "2xEvangelion_dat2": "Phips/2xEvangelion_dat2",
    "1xDeJPG_realplksr_otf": "Phips/1xDeJPG_realplksr_otf",
    "1xDeH264_realplksr": "Phips/1xDeH264_realplksr",
    "1xDeNoise_realplksr_otf": "Phips/1xDeNoise_realplksr_otf",
    "1xExposureCorrection_compact": "Phips/1xExposureCorrection_compact",
    "1xUnderExposureCorrection_compact": "Phips/1xUnderExposureCorrection_compact",
    "1xOverExposureCorrection_compact": "Phips/1xOverExposureCorrection_compact",
}


@spaces.GPU
def upscale_image(image, model_selection):
    original = load_image(image)

    upscaler = UpscaleWithModel.from_pretrained(MODELS[model_selection]).to("cuda")
    image = upscaler(original, tiling=True, tile_width=1024, tile_height=1024)

    return original, image


def clear_result():
    return gr.update(value=None)


title = """<h1 align="center">Image Upscaler</h1>
<div align="center">Use this Space to upscale your images, makes use of the
<a href="https://github.com/asomoza/image_gen_aux">Image Generation Auxiliary Tools</a> library. <br> For now makes use of my self trained models, but can be extended to more models from other authors if they message me.</div>
"""

with gr.Blocks() as demo:
    gr.HTML(title)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")

            model_selection = gr.Dropdown(
                choices=list(MODELS.keys()),
                value="4xNomos2_hq_drct-l",
                label="Model",
            )

            run_button = gr.Button("Upscale")
        with gr.Column():
            result = ImageSlider(
                interactive=False,
                label="Generated Image",
            )

    run_button.click(
        fn=clear_result,
        inputs=None,
        outputs=result,
    ).then(
        fn=upscale_image,
        inputs=[input_image, model_selection],
        outputs=result,
    )


demo.launch(share=False)