Fix NaN on CPU inference

app.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import time
+import math
 from contextlib import nullcontext
 from pathlib import Path
 
@@ -38,7 +39,8 @@ inference_device = None
 def select_torch_dtype() -> torch.dtype:
     if torch.cuda.is_available():
         return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
-    return torch.float16
+    # float16 on CPU is numerically unstable for this model and can produce NaNs.
+    return torch.float32
 
 
 def load_model_and_processor():
@@ -168,7 +170,13 @@ def infer_image(image: Image.Image):
     elapsed_ms = (time.perf_counter() - start_time) * 1000.0
 
     logit = float(outputs["logit"].squeeze().detach().float().cpu().item())
+    if not math.isfinite(logit):
+        raise gr.Error("Model produced a non-finite logit (NaN/Inf). Please retry.")
+
     prob_fake = float(torch.sigmoid(torch.tensor(logit)).item())
+    if not math.isfinite(prob_fake):
+        raise gr.Error("Model produced a non-finite probability (NaN/Inf). Please retry.")
+
     label = "Fake" if prob_fake >= THRESHOLD else "Real"
     confidence = prob_fake if label == "Fake" else 1.0 - prob_fake