import torch
import streamlit as st
from diffusers import StableDiffusion3Pipeline

# Load the model
pipeline = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
pipeline = pipeline.to("cuda")  # Move model to GPU if available

# Streamlit UI
def main():
    st.title("Stable Diffusion 3 Medium Demo")
    prompt = st.text_input("Enter your prompt:", "A cat holding a sign that says hello world")

    if st.button("Generate Image"):
        with st.spinner("Generating..."):
            try:
                image = pipeline(prompt, negative_prompt="", num_inference_steps=28, guidance_scale=7.0).images[0]
                st.image(image, caption="Generated Image", use_column_width=True)
            except Exception as e:
                st.error(f"Error: {e}")

if __name__ == "__main__":
    main()