import streamlit as st
from diffusers import DiffusionPipeline
import torch

# Load the pipeline
pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", 
    torch_dtype=torch.float16, 
    use_safetensors=True, 
    variant="fp16"
)

# Streamlit UI
st.title("Stable Diffusion Image Generator")
prompt = st.text_area('Enter a prompt to generate an image...')

if prompt:  # Ensure the prompt is not empty
    with st.spinner("Generating image..."):
        out = pipe(
            prompt,
            height=1024,
            width=1024,
            guidance_scale=3.5,
            num_inference_steps=50 
        ).images[0]
    
    # Display the image in Streamlit
    st.image(out, caption="Generated Image", use_column_width=True)