import os
import gradio as gr
import util
from run_cmd import run_cmd
from random import randint
from PIL import Image

is_colab = util.is_google_colab()

run_cmd("pip install pngquant")

def inference(img, size, type):
    _id = randint(1, 10000)
    INPUT_DIR = "/tmp/input_image" + str(_id) + "/"
    OUTPUT_DIR = "/tmp/output_image" + str(_id) + "/"
    img_in_path = os.path.join(INPUT_DIR, "1.jpg")
    img_out_path = os.path.join(OUTPUT_DIR, "1_out.png")
    run_cmd(f"rm -rf {INPUT_DIR}")
    run_cmd(f"rm -rf {OUTPUT_DIR}")
    run_cmd(f"mkdir {INPUT_DIR}")
    run_cmd(f"mkdir {OUTPUT_DIR}")
    img.save(img_in_path, "PNG")

    if type == "Manga":
        run_cmd(f"python inference_manga_v2.py {img_in_path} {img_out_path}")
    else:
        run_cmd(f"python inference.py {img_in_path} {img_out_path} {type}")

    img_out = Image.open(img_out_path)

    if size == "x2":
        img_out = img_out.resize((img_out.width // 2, img_out.height // 2), resample=Image.BICUBIC)

    #img_out.save(img_out_path, optimize=True) # Add more optimizations
    #img_out = Image.open(img_out_path)

    # Remove input and output image
    run_cmd(f"rm -f {img_in_path}")
    run_cmd(f"rm -f {img_out_path}")

    return [img_out]

input_image = gr.Image(type="pil", label="Input")
upscale_type = gr.Radio(["Manga", "Anime", "General"], label="Select the type of picture you want to upscale:", value="Manga")
upscale_size = gr.Radio(["x4", "x2"], label="Upscale by:", value="x4")
output_image = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")

title = "ESRGAN Upscaling With Custom Models"
description = "This space uses old ESRGAN architecture to upscale images, using models made by the community."
article = "<p><a href='https://upscale.wiki/wiki/Model_Database'>Model Database</a></p>"

demo = gr.Interface(
    inference,
    inputs=[input_image, upscale_size, upscale_type],
    outputs=[output_image],
    title=title,
    description=description,
    article=article
)

demo.queue()
demo.launch(debug=is_colab, share=is_colab, inline=is_colab)