import os
import streamlit as st
from dotenv import load_dotenv
from huggingface_hub import InferenceApi
from PIL import Image
from io import BytesIO

# Load environment variables from the .env file
load_dotenv()

# Hugging Face API token
HUGGINGFACE_API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")

# Initialize the Hugging Face Inference API
inference = InferenceApi(repo_id="stabilityai/stable-diffusion-3.5-large", token=HUGGINGFACE_API_TOKEN)

# Streamlit App UI
st.set_page_config(page_title="Stable Diffusion Demo", page_icon="🖼️")
st.title("Stable Diffusion 3.5 - Text-to-Image")

# Text input for the prompt
prompt = st.text_input("Enter a prompt for the image:")

# Button to generate the image
if st.button("Generate Image"):
    if prompt:
        try:
            # Make request to the Hugging Face model
            output = inference(inputs=prompt)

            # Convert the output to an image
            image = Image.open(BytesIO(output))
            st.image(image, caption="Generated Image", use_column_width=True)
        except Exception as e:
            st.error(f"Error: {str(e)}")
    else:
        st.warning("Please enter a prompt.")