import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Load the Stable Diffusion model from Hugging Face
model_id = "stabilityai/stable-diffusion-3.5-large"
device = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize the pipeline
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to(device)

# Define a function to generate images
def generate_image(prompt):
    with torch.autocast(device):
        image = pipe(prompt).images[0]
    return image

# Set up the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter your prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="Stable Diffusion 3.5 Demo"
)

interface.launch()