import subprocess
from PIL import Image
import gradio as gr
import os
import random
 
# First, running the inversion command to obtain the input noise that will reconstruct the image. 
# It will save the inversion as output/test_cat/inversion/image-name.pt 
# BLIP-generated caption prompt is saved as output/test_cat/prompt/image-name.txt - eg., a painting of a cat sitting on top of a ball
def inversion(image_in):  #, progress=gr.Progress(track_tqdm=True)):
  #progress(0, desc="Starting...")
  seed = random.randint(0, 1000000)
  img_name = f"input_image_{seed}.png"
  img_label = f"input_image_{seed}"
  # saving the input image
  image_in.save(img_name)
  #image_in.save("input_image.png") #("assets/test_images/cats/input_image.png")
  # Run the script file
  subprocess.run(["python", "src/inversion.py", "--input_image", f"input_image_{seed}.png", "--results_folder", "output/test_cat"])
  # Open the text file with blip caption
  with open(f"output/test_cat/prompt/{img_label}.txt", "r") as file:
      # Read the file 
      prompt = file.read()
  return f"output/test_cat/inversion/{img_label}.pt", prompt, seed


# Performing image editing with the editing directions 
# This will save the edited image as output/test_cat/edit/image-name.png
def image_edit(task_name, seed): #, progress=gr.Progress(track_tqdm=True)):
  #progress(0, desc="Starting...")
  seed = int(seed)
  img_label = f"input_image_{seed}"
  # Run the script file
  subprocess.run(["python", "src/edit_real.py", "--inversion", f"output/test_cat/inversion/{img_label}.pt", 
                  "--prompt", f"output/test_cat/prompt/{img_label}.txt", "--task_name", task_name,
                  "--results_folder", "output/test_cat/"])
  return f"output/test_cat/edit/{img_label}.png"


#Similarly, we can edit the synthetic images generated by Stable Diffusion with the following command.
def synthetic_image_edit(prompt, task_name): #, progress=gr.Progress(track_tqdm=True)):
  #progress(0, desc="Starting...")
  seed = random.randint(0, 1000000)
  #img_name = f"./input_image_{seed}.png"
  #img_label = f"./input_image_{seed}"
  # Run the script file
  subprocess.run(["python", "src/edit_synthetic.py", "--prompt_str", prompt, "--seed", str(seed),
                  "--task", task_name, "--results_folder", "output/synth_editing"])
  return f"output/synth_editing/reconstruction{seed}.png", f"output/synth_editing/edit{seed}.png"

    
def set_visible_true():
    return gr.update(visible=True)

def set_visible_False():
    return gr.update(visible=False)

#Gradio Blocks API
with gr.Blocks() as demo:
  gr.HTML("""<div style="text-align: center; max-width: 700px; margin: 0 auto;">
                    <div
                style="
                    display: inline-flex;
                    align-items: center;
                    gap: 0.8rem;
                    font-size: 1.75rem;
                "
                >
                <h1 style="font-weight: 900; margin-bottom: 7px; margin-top: 5px;">
                    Pix2Pix - Zero-shot Image-to-Image Translation
                </h1>
                </div>
                <p style="margin-bottom: 10px; font-size: 94%">
                This is an unofficial demo for <a href="https://pix2pixzero.github.io/" target="_blank">Pix2PixZero</a>. 
                Please visit their website and <a href="https://github.com/pix2pixzero/pix2pix-zero" target="_blank">github repo</a> for more details.<br>
                If you encounter an error at any step, it can be due to Cuda Memory Error as multiple users are trying to access the app at same time. In such cases, it is suggested to start over with same/different image.</p></div>""")
  gr.HTML("""<a href="https://huggingface.co/spaces/ysharma/pix2pix-zero-01?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate this Space and upgrade to a GPU for fast Inference & no queue<br>""")
  
  direction_html = gr.HTML(value="<h4>👇Upload a Dog or a Cat image here to get started👇</h4>", visible=True)
  refresh_direction_html = gr.HTML(value="<h4>👆Click the 'App' button to Refresh the space and try translation again with another Image</h4>", visible=False)
  with gr.Row():
    image_in = gr.Image(type="pil", label="Start by uploading an image of a Cat or a Dog that you want to translate")
    seed = gr.Number(visible=False, Precision=0)
    #with gr.Box():
    #    gr.Examples( examples=[os.path.join(os.path.dirname(__file__), "dog11.jpg"), os.path.join(os.path.dirname(__file__), "cat11.jpg")],
    #    inputs=image_in,)
    #    #fn=inversion,)
    with gr.Column():
      btn_inversion = gr.Button("Get input noise and image caption",visible=False )
      with gr.Row():
        blip_prompt = gr.Textbox(visible=False)
        inversion_file = gr.File(visible=False)
        #task_name = gr.Textbox()
  with gr.Row():
    image_out = gr.Image(visible=False, label="Translated Image output")
    with gr.Column():
      task_name_radio  = gr.Radio(choices = ["cat2dog", "dog2cat",], type="value", visible=False, label="Select a task that you want to accomplish") #, value="cat2dog"),
      btn_imageedit = gr.Button(value="Translate the image!",visible=False)
      html_tag = gr.HTML(value="""<h3 style="font-weight: 900; margin-bottom: 7px; margin-top: 5px;">
              🤩You can also Generate images with Stable Diffusion and 🚀Translate them on the fly🔥 (zero-shot) using Pix2PixZero. Try this below -</h3><br>
              Example - type a prompt like 'A small cat sitting on a blue ball', select the task as 'cat->dog' in this case, and press the button.""", visible=False)
      prompt_synth = gr.Textbox(visible=False, label="Type in a prompt to generate an Image using SD", placeholder="A small cat sitting on a blue ball")
      btn_synth_image = gr.Button(value="Generate & Translate the SD image",visible=False)
      with gr.Row():
        image_synth = gr.Image(visible=False, label="Synthetic image generated by Stable Diffusion on the fly")
        image_synth_translated = gr.Image(visible=False, label="Translated synthetic image")

    image_in.change(set_visible_False, [], direction_html)
      
    btn_inversion.click(inversion,[image_in],[inversion_file, blip_prompt, seed])
    #btn_inversion.click(set_visible_true, [], task_name_radio)  #inversion_file, blip_prompt,
    btn_inversion.click(set_visible_False, [], btn_inversion)

    inversion_file.change(set_visible_true, [], task_name_radio)  #inversion_file, blip_prompt,
    #task_initial_radio.change(set_visible_true, [], btn_imageedit)
    
    task_name_radio.change(set_visible_true, [], btn_imageedit)
    task_name_radio.change(set_visible_true, [], image_out)
    #task_name_radio.change(set_visible_true, [], html_tag)

    btn_imageedit.click(image_edit,[task_name_radio, seed],[image_out])
    btn_imageedit.click(set_visible_False, [], btn_imageedit)

    btn_imageedit.click(set_visible_true, [], html_tag)
    btn_imageedit.click(set_visible_true, [], prompt_synth)
    btn_imageedit.click(set_visible_true, [], btn_synth_image)
    btn_imageedit.click(set_visible_true, [], image_synth)
    btn_imageedit.click(set_visible_true, [], image_synth_translated)
    
    image_out.change(set_visible_true, [], refresh_direction_html)
    btn_synth_image.click(synthetic_image_edit,[prompt_synth, task_name_radio],[image_synth, image_synth_translated])

    image_in.clear(set_visible_true, [], btn_inversion)
    image_in.change(set_visible_true, [], btn_inversion)
    image_in.change(set_visible_true, [], blip_prompt)
    image_in.change(set_visible_true, [], inversion_file)

demo.queue(concurrency_count=8)
demo.launch(debug=True)