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

# Title of the app
st.title("Text-to-Image Generator")
st.markdown(
    "Generate images from text with Stable Diffusion v1.4 model!"
)

@st.cache_resource
def load_pipeline():
    # Load Stable Diffusion v1.4 model
    model_id = "CompVis/stable-diffusion-v-1-4-original"
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id, torch_dtype=torch.float16
    )
    pipe.to("cpu")  # Ensure it runs on CPU
    return pipe

# Load the pipeline
pipe = load_pipeline()

# Input field for the prompt
prompt = st.text_input("Enter your prompt:")

# Negative prompt for optional customization
neg_prompt = st.text_input(
    "Optional: Enter a negative prompt (to exclude elements, e.g., blurry, low quality):"
)

# Button to generate the image
if st.button("Generate Image"):
    if not prompt.strip():
        st.error("Please enter a valid prompt.")
    else:
        with st.spinner("Generating image..."):
            # Generate the image
            result = pipe(prompt=prompt, negative_prompt=neg_prompt, num_inference_steps=25)
            image = result.images[0]

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

        # Option to download the image
        img_buffer = Image.new("RGB", image.size, (255, 255, 255))
        img_buffer.paste(image, (0, 0))
        img_buffer.save("generated_image.png", "PNG")
        with open("generated_image.png", "rb") as file:
            st.download_button(
                label="Download Image",
                data=file,
                file_name="generated_image.png",
                mime="image/png",
            )