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

# Load the Hugging Face model and processor for deepfake detection.
processor = AutoImageProcessor.from_pretrained("Smogy/SMOGY-Ai-images-detector")
model = AutoModelForImageClassification.from_pretrained("Smogy/SMOGY-Ai-images-detector")

def detect_deepfake(image: Image.Image) -> str:
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    probs = torch.softmax(outputs.logits, dim=1)
    idx = probs.argmax(dim=1).item()
    label = model.config.id2label[idx]
    conf = probs[0, idx].item()
    return f"The image is {label} with confidence {conf:.2f}"

# Build Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Deepfake Detection App")
    gr.Markdown("### Upload an image to detect deepfake content.")
    img_in = gr.Image(type="pil", label="Upload Image")
    txt_out = gr.Textbox(label="Result")
    gr.Button("Detect").click(fn=detect_deepfake, inputs=img_in, outputs=txt_out)

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