Update app.py

app.py

@@ -1,107 +1,73 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
+import torch
 from PIL import Image
 import io
-from threading import Lock
-from datetime import datetime
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+from transformers import AutoProcessor, AutoModelForCausalLM
 
-print("Starting API...")
+app = FastAPI(title="Florence-2-Logit-Shield")
 
-app = FastAPI()
+# إعداد النموذج
+model_id = "microsoft/Florence-2-large"
+device = "cpu"
 
-MODEL_ID = "microsoft/Florence-2-large"
-processor = None
-model = None
-torch = None
-AutoProcessor = None
-AutoModelForCausalLM = None
-model_lock = Lock()
-model_load_error = None
-model_loaded_at = None
+print(f"⏳ Loading {model_id}...")
+# Florence-2 يتطلب trust_remote_code=True
+model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to(device).eval()
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
 
-def ensure_model_loaded():
-    global processor, model, torch, AutoProcessor, AutoModelForCausalLM, model_load_error, model_loaded_at
-    if processor is not None and model is not None:
-        return
-    with model_lock:
-        if processor is not None and model is not None:
-            return
-        print("Loading model...")
-        try:
-            if torch is None or AutoProcessor is None or AutoModelForCausalLM is None:
-                import torch as _torch
-                from transformers import AutoProcessor as _AutoProcessor, AutoModelForCausalLM as _AutoModelForCausalLM
-                torch = _torch
-                AutoProcessor = _AutoProcessor
-                AutoModelForCausalLM = _AutoModelForCausalLM
+# استخراج معرفات التوكنات (Florence يستخدم توكنات مختلفة قليلاً)
+# سنبحث عن توكنات "Yes" و "No" في قاموس النموذج
+YES_TOKEN_ID = processor.tokenizer.convert_tokens_to_ids("Yes")
+NO_TOKEN_ID = processor.tokenizer.convert_tokens_to_ids("No")
 
-            processor = AutoProcessor.from_pretrained(MODEL_ID)
-            model = AutoModelForCausalLM.from_pretrained(
-                MODEL_ID,
-                torch_dtype=torch.float32,
-                trust_remote_code=True
-            )
-            model.to("cpu")
-            model.eval()
-            model_loaded_at = datetime.utcnow().isoformat() + "Z"
-            model_load_error = None
-            print("Model loaded successfully")
-        except Exception as exc:
-            model_load_error = str(exc)
-            print(f"Model loading failed: {model_load_error}")
-            raise
-
-@app.get("/")
-def home():
-    return {
-        "status": "API running",
-        "model_id": MODEL_ID,
-        "model_loaded": processor is not None and model is not None,
-        "model_loaded_at": model_loaded_at,
-        "model_load_error": model_load_error
-    }
-
-@app.get("/healthz")
-def healthz():
-    return {"ok": True, "time": datetime.utcnow().isoformat() + "Z"}
-
-@app.post("/caption")
-async def caption(file: UploadFile = File(...)):
+@app.post("/test-image")
+async def test_image(file: UploadFile = File(...)):
     try:
-        ensure_model_loaded()
-    except Exception as exc:
-        raise HTTPException(status_code=500, detail=f"Model load failed: {exc}")
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        
+        # Florence-2 يعمل بنظام المهام (Task-based)
+        # نستخدم مهمة الإجابة على الأسئلة الوصفية
+        prompt = "<VQA> Is there an adult woman or any adult female body part? Answer No if it is a child."
+        
+        # تجهيز المدخلات
+        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
 
-    image_bytes = await file.read()
-    try:
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    except Exception as exc:
-        raise HTTPException(status_code=400, detail=f"Invalid image file: {exc}")
-
-    prompt = "<CAPTION>"
+        with torch.no_grad():
+            # استخراج المخرجات للحصول على اللوجيت
+            outputs = model(**inputs)
+            
+            # اللوجيت لآخر توكن تم إنتاجه
+            logits = outputs.logits[:, -1, :]
+            probs = torch.softmax(logits, dim=-1)
+            
+            prob_yes = probs[0, YES_TOKEN_ID].item()
+            prob_no = probs[0, NO_TOKEN_ID].item()
 
-    inputs = processor(
-        text=prompt,
-        images=image,
-        return_tensors="pt"
-    )
+        # العتبة (Threshold) لـ Florence-2
+        # هذا النموذج دقيق جداً، لذا نبدأ بـ 0.05 (5%)
+        THRESHOLD = 0.05
+        is_unsafe = prob_yes > THRESHOLD
 
-    with torch.inference_mode():
-        generated_ids = model.generate(
-            input_ids=inputs["input_ids"],
-            pixel_values=inputs["pixel_values"],
-            max_new_tokens=50
-        )
+        return {
+            "is_blur_needed": is_unsafe,
+            "decision": "BLOCK" if is_unsafe else "ALLOW",
+            "analysis": {
+                "female_score": round(prob_yes * 100, 4),
+                "no_score": round(prob_no * 100, 4),
+                "threshold": THRESHOLD
+            },
+            "model": "Florence-2-large"
+        }
 
-    generated_text = processor.batch_decode(
-        generated_ids,
-        skip_special_tokens=True
-    )[0]
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
-    parsed = processor.post_process_generation(
-        generated_text,
-        task=prompt,
-        image_size=(image.width, image.height)
-    )
-    caption = parsed.get(prompt, generated_text) if isinstance(parsed, dict) else str(parsed)
+@app.get("/")
+def home():
+    return {"message": "Florence-2 Shield is Running"}
 
-    return {"caption": caption}
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)