import streamlit as st
from diffusers import DiffusionPipeline
from PIL import Image

# Load the model
@st.cache_resource
def load_pipeline():
    return DiffusionPipeline.from_pretrained("IamCreateAI/Ruyi-Mini-7B")

pipe = load_pipeline()

# Streamlit App
st.title("AI Image Generator")
st.subheader("Generate stunning AI-powered images")

# Input for text prompt
prompt = st.text_input(
    "Enter your prompt:",
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
)

if st.button("Generate Image"):
    st.write("Generating image, please wait...")
    try:
        # Generate the image
        image = pipe(prompt).images[0]
        # Display the image
        st.image(image, caption="Generated Image", use_column_width=True)
    except Exception as e:
        st.error(f"An error occurred: {e}")

st.write("Enter a creative prompt above and click the button to generate an AI-powered image!")