Update main.py

main.py

@@ -1,60 +1,60 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from PIL import Image
-import io
-from transformers import pipeline
-
-app = FastAPI()
-
-# Allow CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Load AI Detection Model
-# We use a pre-trained model from Hugging Face
-# 'umm-maybe/AI-image-detector' is a good general purpose detector
-print("Loading AI Model...")
-classifier = pipeline("image-classification", model="umm-maybe/AI-image-detector")
-print("Model Loaded!")
-
-@app.post("/analyze")
-async def analyze_media(file: UploadFile = File(...)):
-    try:
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents))
-        
-        # Run Inference
-        results = classifier(image)
-        # results is a list like [{'label': 'artificial', 'score': 0.99}, {'label': 'human', 'score': 0.01}]
-        
-        # Find the 'artificial' or 'AI' score
-        ai_score = 0.0
-        for r in results:
-            label = r['label'].lower()
-            if 'artificial' in label or 'ai' in label:
-                ai_score = r['score']
-                break
-            if 'human' in label or 'real' in label:
-                # If we found human score, AI score is 1 - human
-                ai_score = 1.0 - r['score']
-        
-        is_ai = ai_score > 0.5
-        
-        return {
-            "filename": file.filename,
-            "is_ai": is_ai,
-            "confidence": round(ai_score * 100, 2), # Return 0-100
-            "details": results
-        }
-    except Exception as e:
-        print(f"Error: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/")
-def read_root():
-    return {"status": "AI Detector Neural Network is Running"}
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image
+import io
+from transformers import pipeline
+
+app = FastAPI()
+
+# Allow CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load AI Detection Model
+# We use a pre-trained model from Hugging Face
+# 'Organika/sdxl-detector' is specialized for Stable Diffusion detection
+print("Loading AI Model...")
+classifier = pipeline("image-classification", model="Organika/sdxl-detector")
+print("Model Loaded!")
+
+@app.post("/analyze")
+async def analyze_media(file: UploadFile = File(...)):
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents))
+        
+        # Run Inference
+        results = classifier(image)
+        # results is a list like [{'label': 'artificial', 'score': 0.99}, {'label': 'human', 'score': 0.01}]
+        
+        # Find the 'artificial' or 'AI' score
+        ai_score = 0.0
+        for r in results:
+            label = r['label'].lower()
+            if 'artificial' in label or 'ai' in label:
+                ai_score = r['score']
+                break
+            if 'human' in label or 'real' in label:
+                # If we found human score, AI score is 1 - human
+                ai_score = 1.0 - r['score']
+        
+        is_ai = ai_score > 0.5
+        
+        return {
+            "filename": file.filename,
+            "is_ai": is_ai,
+            "confidence": round(ai_score * 100, 2), # Return 0-100
+            "details": results
+        }
+    except Exception as e:
+        print(f"Error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/")
+def read_root():
+    return {"status": "AI Detector Neural Network is Running"}