import gradio as gr
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from io import BytesIO


def roundtrip(im, **options):
    out = BytesIO()
    im.save(out, "PNG", **options)
    out.seek(0)
    return Image.open(out)


def process_image(msg, img):
  targetImage = Image.open("1x1.png") if img is None else img
  metadata = PngInfo()
  metadata.add_text("parameters", msg+" Steps: 0, Hidden message added using: [promptinspector-abuser](https://huggingface.co/spaces/NoCrypt/promptinspector-abuser)")
  im = roundtrip(targetImage, pnginfo=metadata)
  status = "Hidden message added. Message: "+msg
  return im, status


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(label="Hidden Message")
            input_img = gr.Image(elem_id="input", label="Input", interactive=True, type="pil", visible=False)
            btn = gr.Button(value="Submit", variant="primary")
        with gr.Column():
            status = gr.Textbox(label="Status", value="Ready", interactive=False)
            output = gr.Image(elem_id="output_image", label="Output", interactive=False, type="pil", visible=False)
            download = gr.Button(value="Download", interactive=False)

    btn.click(process_image, inputs=[text, input_img], outputs=[output, status])
    download.click(lambda x: x, inputs=[output], outputs=[output], _js="""
    (e)=>{
      var a = document.createElement("a");
      a.href = e;
      a.download = "output_hidden.png";
      a.click(); 
    }
    """)


demo.launch()