import torch
import gradio as gr
from diffusers import StableDiffusionPipeline

# Load the pretrained text-to-image pipeline
pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
).to("cuda" if torch.cuda.is_available() else "cpu")

# Automatically generate a UI from the pipeline
demo = gr.Interface.from_pipeline(pipe, title="Text-to-Image with Stable Diffusion")

demo.launch()