import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Load the Stable Diffusion model
model_id = "stabilityai/stable-diffusion-3-medium"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

def generate_image(prompt):
    # Generate an image from the prompt
    with torch.autocast("cuda"):
        image = pipe(prompt).images[0]
    return image

# Create a Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="Stable Diffusion 3 Image Generator",
    description="Generate images from text prompts using Stable Diffusion 3."
)

# Launch the interface
iface.launch()