import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# --- CONFIGURATION ---
SYSTEM_NAME = "DeepVision"
ENGINE_SOURCE = "dima806/deepfake_vs_real_image_detection"

print(f"Booting {SYSTEM_NAME}...")

# --- LOAD ENGINE ---
try:
    processor = AutoImageProcessor.from_pretrained(ENGINE_SOURCE)
    model = AutoModelForImageClassification.from_pretrained(ENGINE_SOURCE)
    print("Engine Online.")
except Exception as e:
    print(f"Error loading model: {e}")

# --- LOGIC ---
def analyze_image(image):
    if image is None:
        return None
    
    # Preprocess
    inputs = processor(images=image, return_tensors="pt")
    
    # Predict
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
    
    # Format
    id2label = model.config.id2label
    return {id2label[i]: float(probs[i]) for i in range(len(probs))}

# --- INTERFACE (BARE BONES) ---
# Removed: allow_flagging, css, theme (These caused the errors)
interface = gr.Interface(
    fn=analyze_image,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Label(num_top_classes=2, label="Result"),
    title=SYSTEM_NAME,
    description="Deepfake Detection System. Upload an image to analyze."
)

if __name__ == "__main__":
    interface.launch()