import streamlit as st
from diffusers import StableDiffusionPipeline
import torch

# Load the model
@st.cache_resource
def load_pipeline():
    model_name = "runwayml/stable-diffusion-v1-5"  # Replace with the desired Stable Diffusion model
    pipeline = StableDiffusionPipeline.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        use_auth_token=True  # Use Hugging Face token if required for the model
    )
    pipeline = pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
    return pipeline

pipeline = load_pipeline()

# App title
st.title("🎨 Open Source Text-to-Image Generator")
st.write("Generate images from text prompts using Stable Diffusion.")

# Input prompt
prompt = st.text_input("Enter your prompt:", placeholder="A futuristic cityscape at sunset")

# Image generation button
if st.button("Generate Image"):
    if prompt:
        with st.spinner("Generating image..."):
            try:
                # Generate image
                result = pipeline(prompt, num_inference_steps=50, guidance_scale=7.5)
                image = result.images[0]

                # Display the image
                st.image(image, caption="Generated Image", use_column_width=True)
            except Exception as e:
                st.error(f"An error occurred: {e}")
    else:
        st.warning("Please enter a prompt to generate an image.")

# Sidebar configuration
st.sidebar.title("Settings")
st.sidebar.write("Customize your generation:")
guidance_scale = st.sidebar.slider("Guidance Scale", 5.0, 15.0, 7.5)
num_inference_steps = st.sidebar.slider("Inference Steps", 10, 100, 50)

# Clear cache button
if st.sidebar.button("Clear Cache"):
    st.cache_resource.clear()
    st.success("Cache cleared!")