import gradio as gr
import requests
import os
from PIL import Image
from io import BytesIO
import base64

def error_str(error, title="Error"):
    return f"""#### {title}
            {error}"""  if error else ""

def inference(prompt, guidance, steps, image_size="Landscape", seed=0, img=None, strength=0.5, neg_prompt="", disable_auto_prompt_correction=False):
  try:
    response = requests.post(os.environ["BACKEND"], json={
    	"data": [
    		prompt,
    		guidance,
    		steps,
    		image_size,
    		seed,
    		img,
    		strength,
    		neg_prompt,
    		disable_auto_prompt_correction,
    	]
    }).json()

    data = response["data"]
    
    image=Image.open(BytesIO(base64.b64decode(data[0].split(',')[1])))

    return image,data[1],data[2]
  except Exception as e:
    print(error_str(e))
    return None, "Error", "Error"

css = """.main-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.main-div div h1{font-weight:900;margin-bottom:7px}.main-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}
"""
with gr.Blocks(css=css) as demo:
    gr.HTML(
        f"""
            <div class="main-div">
              <div>
                <h1>ChatEmi Beta デモ</h1>
              </div>
              <p>
              個人情報などは入れないでください。
              </p>
              <p>
              サンプルプロンプト１：黒い髪の美少女の顔アップ
              </p>
              <p>
              サンプルプロンプト２：白い髪の男性の上半身
              </p>
            </div>
        """
    )
    with gr.Row():
        
        with gr.Column(scale=55):
          with gr.Group():
              
              with gr.Row():
                prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2,placeholder="[your prompt]")
                generate = gr.Button(value="Generate")

              image_out = gr.Image(height=1024,width=1024)
          error_output = gr.Markdown()

        with gr.Column(scale=45):
          with gr.Tab("Options"):
            with gr.Group():
              neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")
              disable_auto_prompt_correction = gr.Checkbox(label="Disable auto prompt corretion.")
              with gr.Row():
                image_size=gr.Radio(["Portrait","Landscape","Square"])
                image_size.show_label=False
                image_size.value="Square"
                
              with gr.Row():
                guidance = gr.Slider(label="Guidance scale", value=7.5, maximum=25)
                steps = gr.Slider(label="Steps", value=20, minimum=2, maximum=30, step=1)

              seed = gr.Slider(0, 2147483647, label='Seed (0 = random)', value=0, step=1)
              prompt_display= gr.Textbox(label="Upsampled prompt", interactive=False)

          with gr.Tab("Image to image"):
              with gr.Group():
                image = gr.Image(label="Image", height=256, tool="editor", type="pil")
                strength = gr.Slider(label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5)
                  
    inputs = [prompt, guidance, steps, image_size, seed, image, strength, neg_prompt, disable_auto_prompt_correction]

    outputs = [image_out, error_output, prompt_display]
    prompt.submit(inference, inputs=inputs, outputs=outputs)
    generate.click(inference, inputs=inputs, outputs=outputs)
    
demo.queue(concurrency_count=1)
demo.launch()