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

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)