Update app.py

app.py

@@ -80,20 +80,23 @@ def _hf_predict_proba(pil_img: Image.Image) -> float:
     Returns P(Deepfake) in [0,1] using the ViT classifier.
     """
     inputs = _hf_processor(images=pil_img.convert("RGB"), return_tensors="pt")
-    logits = _hf_model(**inputs).logits  # (1, C)
+    with torch.no_grad():
+        logits = _hf_model(**inputs).logits  # (1, C)
+
     if logits.shape[-1] == 1:
-        # Unlikely for this model, but handle binary-sigmoid heads
+        # Binary sigmoid head
         return torch.sigmoid(logits.squeeze(0))[0].item()
 
-    probs = torch.softmax(logits.squeeze(0), dim=-1).cpu().numpy()
+    # Softmax head
+    probs = torch.softmax(logits.squeeze(0), dim=-1).detach().cpu().numpy()
     if _DEEP_IDX is not None and 0 <= _DEEP_IDX < probs.shape[0]:
         return float(probs[_DEEP_IDX])
 
-    # Binary softmax fallback: assume index 1 = deepfake
+    # Binary fallback
     if probs.shape[0] == 2:
         return float(probs[1])
 
-    # Last resort: take the highest class prob (not ideal, but safe)
+    # Last resort: take max
     return float(probs.max())
 
 # -------------------- Output card --------------------