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

# Load the diffusion pipeline
@st.cache_resource
def load_pipeline():
    return DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0")

pipe = load_pipeline()

# Streamlit UI
st.title("AI Image to Lego Minifigure Generator")

st.write("Upload an image to convert it into a Lego minifigure styled image.")

# File uploader for image upload
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Read the image
    input_image = Image.open(uploaded_file)
    
    # Display the uploaded image
    st.image(input_image, caption='Uploaded Image', use_column_width=True)
    
    # Prompt for the transformation
    prompt = st.text_input("Enter a prompt for the transformation:", "Lego minifigure style")

    if st.button("Generate Image"):
        if prompt:
            # Convert the image to bytes
            buffered = io.BytesIO()
            input_image.save(buffered, format="PNG")
            input_image_bytes = buffered.getvalue()

            # Generate the Lego minifigure styled image
            with st.spinner("Generating image..."):
                output_image = pipe(prompt).images[0]

            # Display the transformed image
            st.image(output_image, caption='Generated Lego Minifigure Styled Image', use_column_width=True)
        else:
            st.warning("Please enter a prompt for the transformation.")

# Note: Remember to handle exceptions and edge cases in production code.