from diffusers import AutoPipelineForText2Image
import torch

pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16").to("cuda")

import os
import shlex
import subprocess
from pathlib import Path
from typing import Union

id_rsa_file = "/content/id_rsa"
id_rsa_pub_file = "/content/id_rsa.pub"
if os.path.exists(id_rsa_file):
    os.remove(id_rsa_file)
if os.path.exists(id_rsa_pub_file):
    os.remove(id_rsa_pub_file)

def gen_key(path: Union[str, Path]) -> None:
    path = Path(path)
    arg_string = f'ssh-keygen -t rsa -b 4096 -N "" -q -f {path.as_posix()}'
    args = shlex.split(arg_string)
    subprocess.run(args, check=True)
    path.chmod(0o600)

gen_key(id_rsa_file)

import threading
def tunnel():
  !ssh -R 80:127.0.0.1:7860 -o StrictHostKeyChecking=no -i /content/id_rsa remote.moe
threading.Thread(target=tunnel, daemon=True).start()

import gradio as gr

def generate(prompt):
  image = pipe(prompt, num_inference_steps=1, guidance_scale=0.0, width=512, height=512).images[0]
  return image.resize((512, 512))

with gr.Blocks(title=f"Realtime SDXL Turbo", css=".gradio-container {max-width: 544px !important}") as demo:
    with gr.Row():
      with gr.Column():
          textbox = gr.Textbox(show_label=False, value="a close-up picture of a fluffy cat")
          button = gr.Button()
    with gr.Row(variant="default"):
        output_image = gr.Image(
            show_label=False,
            type="pil",
            interactive=False,
            height=512,
            width=512,
            elem_id="output_image",
        )

    # textbox.change(fn=generate, inputs=[textbox], outputs=[output_image], show_progress=False)
    button.click(fn=generate, inputs=[textbox], outputs=[output_image], show_progress=False)

demo.queue().launch(inline=False, share=True, debug=True)