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

# Initialize pipeline
@st.cache_resource
def load_pipeline():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "segmind/SSD-1B",
        torch_dtype=torch.float16,
        use_safetensors=True,
        variant="fp16"
    )
    pipe.to("cpu")  # Change to "cpu" if no GPU available
    return pipe

pipe = load_pipeline()

# Streamlit app interface
st.title("Stable Diffusion XL Prompt-to-Image Generator")
st.write("Enter a prompt to generate an image using the Stable Diffusion XL model.")

# User inputs
user_prompt = st.text_input("Enter your prompt:")
negative_prompt = st.text_input("Enter a negative prompt (ugly, blurry, poor quality):")

# Generate button
if st.button("Generate Image"):
    with st.spinner("Generating image..."):
        # Generate image
        generated_image = pipe(prompt=user_prompt, negative_prompt=negative_prompt).images[0]
        st.image(generated_image, caption="Generated Image", use_column_width=True)