import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load model
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    revision="fp16"
)
pipe.to("cuda")

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()