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

load_dotenv()
hf_api_key = os.getenv("HUGGINGFACE_API_KEY")

# Function to generate AI-based images using Hugging Face API
def generate_images_using_huggingface_api(text):
    API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2"
    headers = {"Authorization": f"Bearer {hf_api_key}"}
    payload = {"inputs": text}

    response = requests.post(API_URL, headers=headers, json=payload)
    if response.status_code == 200:
        image = Image.open(BytesIO(response.content))
        return image
    else:
        st.error("Error generating image. Check your API key and model.")
        return None

# Streamlit Code
choice = st.sidebar.selectbox("Select your choice", ["Home", "Hugging Face API"])

if choice == "Home":
    st.title("AI Image Generation App")
    with st.expander("About the App"):
        st.write("This is a simple image generation app that uses AI to generate images from a text prompt.")

elif choice == "Hugging Face API":
    st.subheader("Image generation using Hugging Face API")
    input_prompt = st.text_input("Enter your text prompt")
    
    if input_prompt and st.button("Generate Image"):
        st.info("Generating image.....")
        image_output = generate_images_using_huggingface_api(input_prompt)
        if image_output:
            st.success("Image Generated Successfully")
            st.image(image_output, caption="Generated by Hugging Face API")