Updated app.py with new model

app.py

@@ -1,4 +1,3 @@
-# app.py
 import gradio as gr
 from transformers import pipeline
 from PIL import Image
@@ -8,7 +7,6 @@ import traceback
 MODEL_ID = "umm-maybe/AI-image-detector"
 
 # Load the model pipeline (image-classification)
-# device = -1 forces CPU; Spaces manages hardware, so omit device for default
 pipe = pipeline("image-classification", model=MODEL_ID)
 
 def predict_image(image: Image.Image):
@@ -21,37 +19,36 @@ def predict_image(image: Image.Image):
         res = pipe(image)
 
         # typical pipeline returns list of dicts: [{"label": "...", "score": ...}, ...]
-        # But some wrappers return richer JSON. Try to detect structure:
-        # If res is a list of simple label/score dicts, use the top result.
         if isinstance(res, list) and isinstance(res[0], dict) and "label" in res[0]:
             top = res[0]
             label = top.get("label", "")
             score = float(top.get("score", 0.0))
-            # Map label to friendly text if needed
+
             if label.lower().startswith("artificial") or "ai" in label.lower():
                 verdict = "AI-generated"
             elif label.lower().startswith("human") or "real" in label.lower():
                 verdict = "Human-made"
             else:
-                # fallback: if score high and label looks like 'artificial'
                 verdict = label
+
             return f"{verdict} — {score:.2f}"
-        # If pipeline returns nested structure (some wrappers do), attempt to use it:
+
+        # Handle nested response structures
         if isinstance(res, dict):
-            # check nested fields used by composite wrappers
             imr = res.get("individual_model_results") or res.get("models") or {}
-            ai_det = imr.get("ai_detector") or imr.get("umm-maybe") or None
+            ai_det = imr.get("ai_detector") or None
             if ai_det:
                 ai_prob = ai_det.get("ai_probability") or ai_det.get("score") or ai_det.get("confidence")
                 ai_prob = float(ai_prob)
                 if ai_prob >= 0.6:
                     return f"AI-generated — {ai_prob:.2f}"
                 elif ai_prob <= 0.4:
-                    return f"Human-made — {(1-ai_prob):.2f}"
+                    return f"Human-made — {(1 - ai_prob):.2f}"
                 else:
                     return f"Uncertain — {ai_prob:.2f}"
-        # Last fallback: stringify result
+
         return str(res)
+
     except Exception as e:
         traceback.print_exc()
         return f"Error: {e}"
@@ -61,8 +58,8 @@ demo = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="AI Image Detector (umm-maybe)",
-    description="Upload an image. The model returns whether it's AI-generated or human-made."
+    title="AI Image Detector",
+    description="Upload an image to detect whether it is AI-generated or human-made."
 )
 
 if __name__ == "__main__":