Update app.py

app.py

@@ -2,32 +2,52 @@ from transformers import ViTForImageClassification, ViTImageProcessor
 import torch
 from PIL import Image
 import gradio as gr
+import warnings
 
-# Load pre-trained model and processor
-model_name = "facebook/deit-base-distilled-patch16-224"
-processor = ViTImageProcessor.from_pretrained(model_name)
-model = ViTForImageClassification.from_pretrained(model_name)
+# Suppress warnings (optional)
+warnings.filterwarnings('ignore')
 
-def detect_deepfake(image):
-    # Preprocess the image
-    inputs = processor(images=image, return_tensors="pt")
+try:
+    # Load model (smaller version for better performance)
+    model_name = "google/vit-base-patch16-224"
+    processor = ViTImageProcessor.from_pretrained(model_name)
+    model = ViTForImageClassification.from_pretrained(model_name)
     
-    # Make prediction
-    outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
+    print("Model loaded successfully!")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise
+
+def detect_deepfake(image):
+    try:
+        # Convert image to RGB
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+            
+        # Process image
+        inputs = processor(images=image, return_tensors="pt")
+        
+        # Predict
+        with torch.no_grad():
+            outputs = model(**inputs)
+        
+        # Get result
+        predicted_class = outputs.logits.argmax(-1).item()
+        return "Real" if predicted_class == 0 else "Fake (Possible Deepfake)"
     
-    # For demonstration, we'll assume class 0 is real and 1 is fake
-    # (In a real project, you'd need to verify this with your model)
-    return "Real" if predicted_class_idx == 0 else "Fake (Possible Deepfake)"
+    except Exception as e:
+        return f"Error processing image: {str(e)}"
 
-# Create a simple interface
+# Create interface
 iface = gr.Interface(
     fn=detect_deepfake,
     inputs=gr.Image(type="pil"),
     outputs="text",
     title="Deepfake Detection",
-    description="Upload an image to check if it might be a deepfake."
+    examples=[
+        ["real_example.jpg"],  # Add your example files
+        ["fake_example.jpg"]
+    ]
 )
 
-iface.launch()
+iface.launch(server_port=7860, share=False)  # Disable share for local use