from diffusers import StableDiffusionPipeline
import gradio as gr
import torch

# Load the Stable Diffusion model
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Define the image generation function
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Build the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=2, placeholder="Enter your creative prompt here..."),
    outputs="image",
    title="Text-to-Image Generator",
    description="Enter a creative prompt and see it turned into an image!"
)

# Launch the app
interface.launch()