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

# Title
st.title("Text-to-Image Generator")

# Sidebar Input
st.sidebar.header("Input Prompt")
text_prompt = st.sidebar.text_input("Enter a description:", "A fantasy landscape with mountains and rivers")

# Generate Button
if st.sidebar.button("Generate Image"):
    with st.spinner("Generating image..."):
        # Load the model (make sure you have access to GPUs for faster generation)
        pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
        pipe.to("cuda" if torch.cuda.is_available() else "cpu")

        # Generate image
        image = pipe(text_prompt).images[0]

        # Display the generated image
        st.image(image, caption="Generated Image", use_column_width=True)
        st.success("Image generated successfully!")

# Download Option
if 'image' in locals():
    with st.sidebar:
        st.download_button(
            "Download Image",
            data=image.tobytes(),
            file_name="generated_image.png",
            mime="image/png",
        )