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

# Load the pre-trained model
pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")

# 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("")
if st.button("Generate Image"):
    with st.spinner("Generating image..."):
        # Generate the image
        image = pipeline(prompt).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)