import gradio as gr
from diffusers import DiffusionPipeline
import torch

# Load model

pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")

def generate(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio Interface
gr.Interface(
    fn=generate,
    inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g. A dragon flying over a city"),
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="🖼️ Text to Image Generator",
    description="Enter a prompt and watch it turn into an AI-generated image using Stable Diffusion v1.5"
).launch()