Update app.py

app.py

@@ -3,45 +3,62 @@ from transformers import AutoModelForImageClassification, ViTImageProcessor
 from PIL import Image
 import torch
 
+# -------------------------------
 # Load model once (global)
+# -------------------------------
 model_id = "jacoballessio/ai-image-detect-distilled"
 
 processor = ViTImageProcessor.from_pretrained(model_id)
 model = AutoModelForImageClassification.from_pretrained(
     model_id,
-    torch_dtype=torch.float32
+    dtype=torch.float32,
+    low_cpu_mem_usage=True
 )
-model.eval()
 
+model.eval()
 device = "cpu"
 model.to(device)
 
 
+# -------------------------------
+# Prediction function
+# -------------------------------
 def predict(image: Image.Image):
     if image is None:
-        return "Please upload an image", ""
+        return "Please upload an image", None
 
     # Preprocess
     inputs = processor(image, return_tensors="pt").to(device)
 
+    # Inference
     with torch.no_grad():
         outputs = model(**inputs)
 
-    # Softmax
+    # Probabilities
     probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
     confidence = probs.max().item()
     predicted_label = model.config.id2label[probs.argmax().item()]
 
+    # Convert to dict for Gradio Label
+    labels = model.config.id2label
+    scores = probs.squeeze().tolist()
+
+    confidence_dict = {
+        labels[i]: float(scores[i]) for i in range(len(scores))
+    }
+
     # Result text
     if predicted_label.lower() == "fake":
         result = f"⚠️ AI-GENERATED\nConfidence: {confidence:.3f}"
     else:
         result = f"✅ REAL IMAGE\nConfidence: {confidence:.3f}"
 
-    return result, probs.squeeze().tolist()
+    return result, confidence_dict
 
 
+# -------------------------------
 # UI
+# -------------------------------
 app = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload Image"),
@@ -53,5 +70,9 @@ app = gr.Interface(
     description="Upload an image to check if it's AI-generated or real."
 )
 
+
+# -------------------------------
+# Run app
+# -------------------------------
 if __name__ == "__main__":
     app.launch(server_name="0.0.0.0", server_port=7860)