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):
    # Prepare the payload with binary image data and filename
    files = {
        "inputs": (image.name, image, "image/png" if image.name.endswith(".png") else "image/jpeg")
    }
    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...")

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

    # Debugging: Display the raw result from the Hugging Face API
    st.write(result)  # This will display the full API response for debugging

    # 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.")