import streamlit as st
import requests
from io import BytesIO

# 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"]  # You'll store this in the Hugging Face secret

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

def query(image_bytes):
    response = requests.post(API_URL, headers=headers, files={"inputs": image_bytes})
    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.")

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Display the uploaded image
    image = uploaded_file.read()
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    st.write("Classifying...")
    
    # Send the image to the model
    result = query(image)
    
    # 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.")