import gradio as gr
import torch
from diffusers import AutoPipelineForText2Image

pipe = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/sdxl-turbo",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    variant="fp16" if torch.cuda.is_available() else None,
)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def generate_image(prompt):
    result = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
    return result.images[0]  # returns a PIL image

iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Image(type="pil", format="jpeg"),  # base64 encoded image
    title="Speech2Image - Turbo",
    description="Fast image gen using SDXL Turbo"
)

if __name__ == "__main__":
    iface.launch()