import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image

# Load the stable diffusion model
@st.cache_resource
def load_model():
    model_id = "runwayml/stable-diffusion-v1-5"
    
    # Load the model
    pipe = StableDiffusionPipeline.from_pretrained(model_id)
    
    # Check if CUDA is available and use it, otherwise fall back to CPU
    if torch.cuda.is_available():
        pipe = pipe.to("cuda")
    else:
        pipe = pipe.to("cpu")
        
    return pipe

pipe = load_model()

# Streamlit app title
st.title("Stable Diffusion Image Generator")

# User input for text prompt
prompt = st.text_input("Enter your image prompt:", "")

# Button to generate image
if st.button("Generate Image"):
    if prompt:
        # Generate the image
        with st.spinner("Generating..."):
            image = pipe(prompt).images[0]
            st.image(image, caption=prompt, use_column_width=True)
            
            # Save the image
            image.save(f"{prompt[:50].replace(' ', '_')}.png")
            st.success("Image generated successfully!")
    else:
        st.warning("Please enter a prompt to generate an image.")