Update app.py

app.py

@@ -1,95 +1,158 @@
-import os
-import torch
+import sys
+import types
+import importlib.util
+
+# ─── Mock flash_attn ───────────────────────────────────────────────
+flash_mock = types.ModuleType("flash_attn")
+flash_mock.__version__ = "2.0.0"
+flash_mock.__spec__ = importlib.util.spec_from_loader("flash_attn", loader=None)
+sys.modules["flash_attn"] = flash_mock
+sys.modules["flash_attn.flash_attn_interface"] = types.ModuleType("flash_attn.flash_attn_interface")
+sys.modules["flash_attn.bert_padding"] = types.ModuleType("flash_attn.bert_padding")
+# ───────────────────────────────────────────────────────────────────
+
 import io
+import base64
+import time
+import torch
 from PIL import Image
-from fastapi import FastAPI, File, UploadFile
+from transformers import AutoProcessor, AutoModelForCausalLM
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
-from transformers import AutoProcessor, AutoModelForCausalLM, AutoConfig
+from pydantic import BaseModel
+from contextlib import asynccontextmanager
+import uvicorn
+
+# ─── المهام المتاحة ────────────────────────────────────────────────
+AVAILABLE_TASKS = [
+    "<CAPTION>",
+    "<DETAILED_CAPTION>",
+    "<MORE_DETAILED_CAPTION>",
+    "<OD>",
+    "<OCR>",
+    "<OCR_WITH_REGION>",
+]
+
+# ─── تحميل النموذج عند بدء التطبيق ───────────────────────────────
+MODEL = {}
 
-# إعداد المسارات لضمان الصلاحيات داخل Hugging Face
-os.environ["HF_HOME"] = "/app/cache"
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    print("📥 جاري تحميل Florence-2-large...")
+    start = time.time()
 
-app = FastAPI(title="Florence-2 Image Captioning API")
+    processor = AutoProcessor.from_pretrained(
+        "microsoft/Florence-2-large",
+        trust_remote_code=True
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        "microsoft/Florence-2-large",
+        torch_dtype=torch.float32,
+        trust_remote_code=True,
+        attn_implementation="eager"
+    )
+    model.eval()
 
-# تفعيل CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
+    MODEL["processor"] = processor
+    MODEL["model"] = model
+    MODEL["device"] = "cpu"
+
+    print(f"✅ النموذج جاهز في {time.time()-start:.1f}s")
+    yield
+    MODEL.clear()
+
+# ─── FastAPI ───────────────────────────────────────────────────────
+app = FastAPI(
+    title="Florence-2-large API",
+    description="Microsoft Florence-2-large - Image Analysis API",
+    version="1.0.0",
+    lifespan=lifespan
 )
 
-model_id = "microsoft/Florence-2-large"
-device = "cpu"
+# ─── Schemas ───────────────────────────────────────────────────────
+class PredictRequest(BaseModel):
+    image: str          # base64 encoded image
+    task: str = "<CAPTION>"
+    text_input: str = None
 
-print(f"⏳ Loading {model_id} (CPU Optimized)...")
+class PredictResponse(BaseModel):
+    result: dict
+    task: str
+    elapsed_seconds: float
 
-# 1. تحميل الإعدادات وإصلاح مشكلة forced_bos_token_id
-config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
-if not hasattr(config, 'forced_bos_token_id'):
-    config.forced_bos_token_id = None
+# ─── Endpoints ─────────────────────────────────────────────────────
+@app.get("/")
+def root():
+    return {
+        "status": "running",
+        "model": "microsoft/Florence-2-large",
+        "device": MODEL.get("device", "loading..."),
+        "available_tasks": AVAILABLE_TASKS
+    }
 
-# 2. تحميل المعالج
-processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+@app.get("/health")
+def health():
+    return {"status": "ok", "model_loaded": "model" in MODEL}
 
-# 3. تحميل النموذج مع تعطيل Flash Attention لاستخدامه على الـ CPU
-model = AutoModelForCausalLM.from_pretrained(
-    model_id, 
-    config=config,
-    trust_remote_code=True,
-    torch_dtype=torch.float32,
-    attn_implementation="sdpa"  # 👈 هذا السطر يحل مشكلة ImportError: flash_attn
-).to(device).eval()
+@app.post("/predict", response_model=PredictResponse)
+def predict(request: PredictRequest):
+    # التحقق من المهمة
+    if request.task not in AVAILABLE_TASKS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"مهمة غير صالحة. المهام المتاحة: {AVAILABLE_TASKS}"
+        )
 
-@app.post("/describe")
-async def describe_image(file: UploadFile = File(...)):
+    # فك تشفير الصورة
     try:
-        # قراءة الصورة
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-        
-        # مهمة الوصف التفصيلي
-        task_prompt = "<DETAILED_CAPTION>"
-        
-        # تجهيز المدخلات
-        inputs = processor(text=task_prompt, images=image, return_tensors="pt").to(device)
-
-        # توليد النص
+        image_bytes = base64.b64decode(request.image)
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"خطأ في الصورة: {str(e)}")
+
+    # تحضير الـ prompt
+    prompt = request.task
+    if request.text_input:
+        prompt += request.text_input
+
+    # الاستدلال
+    try:
+        processor = MODEL["processor"]
+        model = MODEL["model"]
+
+        inputs = processor(
+            text=prompt,
+            images=image,
+            return_tensors="pt"
+        )
+
+        start = time.time()
         with torch.no_grad():
             generated_ids = model.generate(
                 input_ids=inputs["input_ids"],
                 pixel_values=inputs["pixel_values"],
-                max_new_tokens=1024,
-                num_beams=3,
-                do_sample=False
-            )
-            
-            generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-            
-            # معالجة النتائج
-            parsed_answer = processor.post_process_generation(
-                generated_text, 
-                task=task_prompt, 
-                image_size=(image.width, image.height)
+                max_new_tokens=512,
+                do_sample=False,
+                num_beams=1
             )
 
-        return {
-            "description": parsed_answer[task_prompt],
-            "status": "success"
-        }
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+        parsed = processor.post_process_generation(
+            generated_text,
+            task=request.task,
+            image_size=(image.width, image.height)
+        )
+        elapsed = time.time() - start
 
-    except Exception as e:
-        print(f"❌ Error: {str(e)}")
-        return JSONResponse(
-            status_code=500, 
-            content={"error": "Processing failed", "details": str(e)}
+        return PredictResponse(
+            result=parsed,
+            task=request.task,
+            elapsed_seconds=round(elapsed, 2)
         )
 
-@app.get("/")
-def health_check():
-    return {"status": "online", "device": device, "endpoint": "/describe"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"خطأ في الاستدلال: {str(e)}")
+
 
 if __name__ == "__main__":
-    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)