app.py

import io
import time
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageTextToText
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager

MODEL_ID = "HuggingFaceTB/SmolVLM-500M-Instruct"

VQA_QUESTION = (
    "Is there a human being or any part of a human body in the picture? Answer yes or no"
)

MODEL_DATA = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    print(f"📥 Loading {MODEL_ID}...")
    start = time.time()

    MODEL_DATA["processor"] = AutoProcessor.from_pretrained(
        MODEL_ID,
        size={"longest_edge": 1 * 512}  # تقليل الدقة لتسريع CPU
    )
    MODEL_DATA["model"] = AutoModelForImageTextToText.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,
        _attn_implementation="eager"
    ).eval()

    print(f"✅ Model ready in {time.time()-start:.1f}s")
    yield
    MODEL_DATA.clear()

app = FastAPI(
    title="Female Detection API - SmolVLM-500M",
    description="HuggingFaceTB/SmolVLM-500M-Instruct | VQA",
    version="1.0.0",
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": "model" in MODEL_DATA}

def decide(answer: str) -> tuple[str, str]:
    a = answer.strip().lower()
    if a == "no" or a.startswith("no"):
        return "allow", "model_answered_no"
    elif "yes" in a:
        return "block", "model_answered_yes"
    else:
        return "block", "unexpected_answer_blocked_for_safety"

@app.post("/analyze")
async def analyze_image(file: UploadFile = File(...)):

    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="الملف ليس صورة")

    try:
        image = Image.open(io.BytesIO(await file.read())).convert("RGB")
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"خطأ في قراءة الصورة: {str(e)}")

    try:
        processor = MODEL_DATA["processor"]
        model     = MODEL_DATA["model"]

        # ─── SmolVLM طريقة الرسمية ────────────────────────────────
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image"},        # ← بدون image= هنا
                    {"type": "text", "text": VQA_QUESTION}
                ]
            }
        ]

        # ─── apply_chat_template ──────────────────────────────────
        prompt = processor.apply_chat_template(
            messages,
            add_generation_prompt=True
        )

        # ─── processor يستقبل الصورة هنا ─────────────────────────
        inputs = processor(
            text=prompt,
            images=[image],
            return_tensors="pt"
        )

        start_time = time.time()
        with torch.no_grad():
            generated_ids = model.generate(
                **inputs,
                max_new_tokens=20,
                do_sample=False,
            )

        # ─── استخراج الإجابة الجديدة فقط ─────────────────────────
        generated_ids_trimmed = [
            out_ids[len(in_ids):]
            for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
        ]
        answer = processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False
        )[0].strip()

        elapsed          = round(time.time() - start_time, 2)
        decision, reason = decide(answer)

        return {
            "decision":       decision,
            "reason":         reason,
            "vqa_answer":     answer,
            "question":       VQA_QUESTION,
            "execution_time": elapsed,
            "status":         "success"
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)