import streamlit as st
import requests

# Hugging Face Inference API URL and Token
API_URL = "https://api-inference.huggingface.co/models/Organika/sdxl-detector"
API_TOKEN = st.secrets["HF_API_TOKEN"]

headers = {"Authorization": f"Bearer {API_TOKEN}"}

# Function to query the model
def query(image_bytes):
    # Prepare the payload with binary image data
    files = {"inputs": image_bytes}
    response = requests.post(API_URL, headers=headers, files=files)
    return response.json()

# Streamlit UI
st.title("AI Image Detector")

st.write("Upload an image, and we will check if it is AI-generated using the Hugging Face SDXL detector.")

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

if uploaded_file is not None:
    # Display the uploaded image
    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)

    st.write("Classifying...")

    # Read the uploaded file as bytes
    image_bytes = uploaded_file.read()  # Read the file as bytes

    # Send the image to the model
    result = query(image_bytes)

    # Display the result
    if "error" in result:
        st.error(f"Error: {result['error']}")
    else:
        label = result[0]["label"]
        if label == "AI-generated":
            st.success("This image is AI-generated.")
        else:
            st.success("This image is not AI-generated.")