import streamlit as st
from huggingface_hub import InferenceClient
import os
from io import BytesIO
from PIL import Image

# Load API key from secrets
api_key = os.getenv("HF_API_KEY")

# Initialize the InferenceClient
client = InferenceClient(api_key=api_key)

st.title("🤖 AI Image Generator")

# Text input for the image description
prompt = st.text_input("Enter a prompt for image generation:", "a robot on a stallion.")

if st.button("Generate Image"):
    if api_key:
        with st.spinner("Generating image... Please wait!"):
            # Generate image from Hugging Face model
            image = client.text_to_image(prompt, model="black-forest-labs/FLUX.1-dev")

            # Convert PIL Image to BytesIO
            img_buffer = BytesIO()
            image.save(img_buffer, format="PNG")
            img_bytes = img_buffer.getvalue()  # Convert to bytes

            # Display image
            st.image(image, caption="Generated Image", use_column_width=True)

            # Provide download button
            st.download_button(
                label="Download Image",
                data=img_bytes,
                file_name="generated_image.png",
                mime="image/png"
            )
    else:
        st.error("API key not found. Please set the `HF_API_KEY` in secrets.")