import os
import json
import requests
import gradio as gr
from util import base64_to_img, img_to_base64, resize_image

url = os.getenv("REQUEST_URL")
headers = {'Content-Type': 'application/json',
           'Validation-Key': os.getenv("VALIDATION_KEY")}
names = ["input_image", "prompt", "neg_prompt", "maxlen", "step", "cfg", "seed", "up", "down", "left", "right"]
def run(*params):
    params = {k:v for k, v in zip(names, params)}
    image = params.pop("input_image")
    image = resize_image(image)
    params["image_base64"] = img_to_base64(image)
    try:
        response = requests.post(url, headers=headers, data=json.dumps(params), timeout=30)
        if response.status_code != 200:
            raise ValueError()
        data = response.json()
    except:
        raise gr.Error("Fail to generate")
    if data["code"] != 0:
        raise gr.Error(data["message"])
    result = base64_to_img(data["content"])
    return result

with gr.Blocks() as demo:
    gr.Markdown("# SDXL inpainting for Clothes2Human")
    with gr.Row().style(equal_height=True):
        with gr.Column():
            input_image = gr.Image(type="pil", height=300)
        with gr.Column():
            output_image = gr.Image(type="pil", height=300)

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt")
            neg_prompt = gr.Textbox(label="Negative Prompt")

            maxlen = gr.Slider(label="Max Edge Length", step=32, minimum=768, maximum=1536, value=1024)
            step = gr.Slider(label="Step", minimum=20, maximum=70, value=50, step=1)

        with gr.Column():
            up = gr.Slider(label="Scale Up Image", minimum=-0.3, maximum=0.5, value=0, step=0.1)
            down = gr.Slider(label="Scale Down Image", minimum=-0.3, maximum=0.5, value=0, step=0.1)
            left = gr.Slider(label="Scale Left Image", minimum=-0.3, maximum=0.5, value=0, step=0.1)
            right = gr.Slider(label="Scale Right Image", minimum=-0.3, maximum=0.5, value=0, step=0.1)
        with gr.Column():
            cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=9.0, value=5.0, step=0.5)
            seed = gr.Slider(label="Seed", minimum=-1, maximum=1000000, value=-1, step=1)
            inpaint_button = gr.Button()

            run_in = [input_image, prompt, neg_prompt, maxlen, step, cfg, seed, up, down, left, right]
            inpaint_button.click(run, inputs=run_in, outputs=[output_image])

    gr.Examples([["imgs/1.jpg","A man wearing a white T-shirt stands on the beach","", 1024, 50, 5.0, 333866, 0.3, 0.3, 0.1, 0.1],
                 ["imgs/2.jpg"," woman wearing a blue dress stands in a park, asian race","", 1280, 50, 5.0, 443652, 0.3, 0.3, 0.2, 0.2],
                 ["imgs/3.jpg","A woman wearing a white dress stands","", 1280, 50, 5.0, 306728, -0.1, -0.2, 0, 0]],
                inputs=run_in, outputs=[output_image], fn=run, cache_examples=True)

demo.queue(concurrency_count=2).launch()