Update app.py

app.py

@@ -1,31 +1,30 @@
-from transformers import AutoImageProcessor, AutoModelForImageClassification
+import gradio as gr
+from transformers import ViTForImageClassification, ViTImageProcessor
 from PIL import Image
 import torch
 
-# ✅ Use a valid Hugging Face model repo
-MODEL_NAME = "selimsef/dfdc_deepfake_challenge"
-
-# Load processor & model
-processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
-model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
-
-# Load your image
-image_path = "your_image.jpg"  # replace with your raw image path
-image = Image.open(image_path).convert("RGB")
-
-# Preprocess
-inputs = processor(images=image, return_tensors="pt")
-
-# Inference
-with torch.no_grad():
-    outputs = model(**inputs)
-    logits = outputs.logits
-    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
-
-# Get labels
-labels = model.config.id2label
-real_prob = probs[labels["0"]] if "0" in labels else probs[0]
-fake_prob = probs[labels["1"]] if "1" in labels else probs[1]
-
-print(f"Prediction: Real={real_prob:.4f}, Fake={fake_prob:.4f}")
+MODEL_NAME = "prithivMLmods/Deep-Fake-Detector-v2-Model"
+
+processor = ViTImageProcessor.from_pretrained(MODEL_NAME)
+model = ViTForImageClassification.from_pretrained(MODEL_NAME)
+
+def classify_image(image):
+    image = Image.fromarray(image).convert("RGB")
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
+    labels = model.config.id2label
+    return { labels[i]: float(probs[i]) for i in range(len(probs)) }
+
+demo = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Label(num_top_classes=2, label="Prediction (Real vs Deepfake)"),
+    title="Deepfake Detector (ViT)",
+    description="Upload an image — model classifies it as Real or Deepfake."
+)
+
+if __name__ == "__main__":
+    demo.launch()