Update app.py

app.py

@@ -3,13 +3,15 @@ from flask_cors import CORS
 import requests
 import os
 import re
+import io
 import hashlib
 import datetime
 import base64
 import time
 import torch
 import torch.nn as nn
-from transformers import AutoTokenizer, AutoConfig, AutoModel, PreTrainedModel
+from PIL import Image
+from transformers import AutoTokenizer, AutoConfig, AutoModel, PreTrainedModel, pipeline
 
 app = Flask(__name__)
 CORS(app)
@@ -71,6 +73,21 @@ except Exception as e:
     print(f"❌ Text model failed to load: {type(e).__name__}: {e}")
 
 
+# ── DEEPFAKE IMAGE MODEL (loaded locally inside the Space) ──
+# dima806/deepfake_vs_real_image_detection — a Vision Transformer for Real-vs-Fake
+# classification. Loaded locally (standard image-classification pipeline) instead
+# of the HF Inference API, which is what caused the old 500 errors.
+# NOTE: this model is trained on an older dataset and works best on human FACE
+# photos; it is less reliable on brand-new AI-generated images.
+IMAGE_MODEL_NAME = "dima806/deepfake_vs_real_image_detection"
+try:
+    image_classifier = pipeline("image-classification", model=IMAGE_MODEL_NAME)
+    print(f"✅ Image model loaded: {IMAGE_MODEL_NAME}")
+except Exception as e:
+    image_classifier = None
+    print(f"❌ Image model failed to load: {type(e).__name__}: {e}")
+
+
 def _ai_prob_for_chunk(chunk):
     enc = _tokenizer(chunk, truncation=True, max_length=512, return_tensors="pt")
     with torch.no_grad():
@@ -218,31 +235,31 @@ def build_signals(trust_score, ai_probability, used_hf):
 
 
 def detect_image_deepfake(image_data):
+    """Runs the locally-loaded ViT deepfake detector. No more HF Inference API,
+    so no more 500s from a dead endpoint."""
+    if image_classifier is None:
+        return {"success":False,"error":"Image model not loaded on server. Try again shortly.","deepfake_probability":0,"real_probability":0,"verdict":"UNKNOWN","confidence":0,"model":"error"}
     try:
-        headers = {"Authorization": f"Bearer {HF_API_KEY}", "Content-Type": "image/jpeg"}
-        response = requests.post(HF_IMAGE_MODEL_URL, headers=headers, data=image_data, timeout=60)
-        print(f"Image HF Status: {response.status_code}, Response: {response.text[:300]}")
-        result = response.json()
-        if isinstance(result, list) and len(result) > 0:
-            scores = result[0] if isinstance(result[0], list) else result
-            fake_score = 0
-            real_score = 0
-            for item in scores:
-                label = item.get('label', '').lower()
-                score = item.get('score', 0)
-                if label in ['fake', 'deepfake', 'label_1', 'artificial']:
-                    fake_score = score
-                elif label in ['real', 'genuine', 'label_0', 'authentic']:
-                    real_score = score
-            if fake_score == 0 and real_score == 0 and len(scores) >= 2:
-                real_score = scores[0]['score']
-                fake_score = scores[1]['score']
-            deepfake_prob = round(fake_score * 100)
-            real_prob = round(real_score * 100)
-            return {"success":True,"deepfake_probability":deepfake_prob,"real_probability":real_prob,"verdict":"DEEPFAKE" if deepfake_prob > 50 else "AUTHENTIC","confidence":max(deepfake_prob, real_prob),"model":"deepfake_vs_real_image_detection"}
+        img = Image.open(io.BytesIO(image_data)).convert("RGB")
+        results = image_classifier(img)  # list of {label, score}
+        fake_score = 0.0
+        real_score = 0.0
+        for item in results:
+            label = str(item.get('label', '')).lower()
+            score = item.get('score', 0)
+            if label in ['fake', 'deepfake', 'label_1', 'artificial']:
+                fake_score = score
+            elif label in ['real', 'genuine', 'label_0', 'authentic']:
+                real_score = score
+        if fake_score == 0 and real_score == 0 and len(results) >= 2:
+            real_score = results[0]['score']
+            fake_score = results[1]['score']
+        deepfake_prob = round(fake_score * 100)
+        real_prob = round(real_score * 100)
+        return {"success":True,"deepfake_probability":deepfake_prob,"real_probability":real_prob,"verdict":"DEEPFAKE" if deepfake_prob > 50 else "AUTHENTIC","confidence":max(deepfake_prob, real_prob),"model":IMAGE_MODEL_NAME}
     except Exception as e:
         print(f"Image detection error: {type(e).__name__}: {e}")
-    return {"success":False,"error":"Model loading. Please try again in 30 seconds.","deepfake_probability":0,"real_probability":0,"verdict":"UNKNOWN","confidence":0,"model":"error"}
+        return {"success":False,"error":"Could not analyze this image. Try a clear JPG/PNG.","deepfake_probability":0,"real_probability":0,"verdict":"UNKNOWN","confidence":0,"model":"error"}
 
 
 def detect_voice_deepfake(audio_data, content_type):
@@ -490,7 +507,10 @@ def detect_image():
             return jsonify({"error": result.get("error", "Detection failed")}), 500
         return jsonify({
             "label": result["verdict"],
-            "confidence": result["confidence"] / 100
+            "confidence": result["confidence"] / 100,
+            "deepfake_probability": result["deepfake_probability"],
+            "real_probability": result["real_probability"],
+            "model": result["model"]
         })
     except Exception as e:
         print(f"Image route error: {e}")
@@ -773,13 +793,15 @@ def network_test():
     return jsonify(results)
 
 
-# ── MODEL STATUS (quick check that the text model loaded) ──
+# ── MODEL STATUS (quick check that the models loaded) ──
 
 @app.route('/model-status', methods=['GET'])
 def model_status():
     return jsonify({
         "text_model": TEXT_MODEL_NAME,
-        "loaded": _text_model is not None
+        "text_loaded": _text_model is not None,
+        "image_model": IMAGE_MODEL_NAME,
+        "image_loaded": image_classifier is not None
     })