import torch
from diffusers import StableDiffusionXLPipeline
import gradio as gr

device = "cuda" if torch.cuda.is_available() else "cpu"

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

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

gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title= Made by David Thongam
).launch()