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

# Stronger forensic model
model_name = "buildborderless/CommunityForensics-DeepfakeDet-ViT"

# Load once at startup
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
model.eval()

def predict(image):
    if image is None:
        return None, "Please upload an image."

    try:
        image = image.convert("RGB")

        # FORCE correct size
        inputs = processor(
            images=image,
            return_tensors="pt",
            size={"height": 384, "width": 384}
        )

        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.softmax(outputs.logits, dim=1)[0]

        real_prob = probs[0].item()
        fake_prob = probs[1].item()

        authenticity = real_prob * 100
        fake_score = fake_prob * 100

        if fake_prob > 0.5:
            verdict = f"""
            
🚨 LIKELY AI-GENERATED

Confidence: {fake_score:.2f}%
Authenticity Score: {authenticity:.2f}%
"""
        else:
            verdict = f"""
✅ LIKELY AUTHENTIC

Confidence: {authenticity:.2f}%
Fake Probability: {fake_score:.2f}%
"""

        return {
            "Authentic": real_prob,
            "AI-Generated": fake_prob
        }, verdict

    except Exception as e:
        return None, f"Error processing image: {str(e)}"

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Label(num_top_classes=2),
        gr.Textbox(label="Forensic Verdict")
    ],
    title="🔬 DeepTrust - Community Forensics AI",
    description="Upload an image to analyze whether it is real or AI-generated using a large-scale ViT forensic model."
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)