import streamlit as st
from diffusers import DiffusionPipeline
import base64
import io

# Load the pre-trained model
pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-512x512", use_safetensors=True)

# Create a Streamlit interface
st.title("Image Generation Model")
st.write("Enter a prompt for the image generation model:")

# Generate the image
prompt = st.text_input("Prompt", label_visibility="collapsed")
if st.button("Generate Image"):
    with st.spinner("Generating image..."):
        # Generate the image
        image = pipeline(prompt, num_inference_steps=20).images[0]

        # Convert the image to bytes
        buf = io.BytesIO()
        image.save(buf, format='JPEG')
        img_bytes = buf.getvalue()

        # Display the generated image
        st.write("Generated image:")
        st.image(img_bytes, caption="Generated image", use_column_width=True)