Create app.py

app.py

@@ -0,0 +1,144 @@
+# app.py
+import os
+import io
+import asyncio
+import numpy as np
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from PIL import Image
+import torch
+from transformers import AutoImageProcessor, AutoModel
+from typing import Optional
+
+app = FastAPI(title="DINOv2 Image Embedding API")
+
+# Configurações — altere MODEL_REPO se quiser outra variante
+MODEL_REPO = os.environ.get("MODEL_REPO", "facebook/dinov2-small")
+HF_TOKEN = os.environ.get("HF_TOKEN", None)  # se repo privado
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Limite de inferências concorrentes (ajuste via var de ambiente se quiser)
+MAX_CONCURRENT = int(os.environ.get("MAX_CONCURRENT", "4"))
+inference_semaphore = asyncio.Semaphore(MAX_CONCURRENT)
+
+# Globals
+model = None
+processor = None
+
+def load_model_and_processor():
+    """
+    Carrega AutoImageProcessor e AutoModel do Hugging Face Hub.
+    Usa AutoModel (gera embeddings via pooler_output ou CLS token).
+    """
+    global model, processor
+    if model is not None and processor is not None:
+        return
+
+    # Opções de auth
+    use_auth = True if HF_TOKEN else False
+    auth = HF_TOKEN if HF_TOKEN else None
+
+    print(f"Loading processor and model from: {MODEL_REPO} (device={DEVICE})")
+    # Carrega processor (pré-processamento oficial do repositório)
+    processor = AutoImageProcessor.from_pretrained(MODEL_REPO, use_auth_token=auth) if use_auth else AutoImageProcessor.from_pretrained(MODEL_REPO)
+    # Carrega o modelo base (sem cabeça de classificação explicita)
+    model = AutoModel.from_pretrained(MODEL_REPO, use_auth_token=auth) if use_auth else AutoModel.from_pretrained(MODEL_REPO)
+    model.to(DEVICE)
+    model.eval()
+    # imprimir dimensão de saída (útil para debug)
+    try:
+        hidden_size = model.config.hidden_size
+        print(f"Model loaded. hidden_size = {hidden_size}")
+    except Exception:
+        print("Model loaded. (no hidden_size in config)")
+
+def extract_embedding_from_outputs(outputs):
+    """
+    Tenta extrair um embedding a partir da saída do AutoModel:
+    - usa pooler_output se disponível
+    - senão usa last_hidden_state[:, 0, :] (token CLS)
+    """
+    if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
+        emb = outputs.pooler_output
+    elif hasattr(outputs, "last_hidden_state"):
+        emb = outputs.last_hidden_state[:, 0, :]  # CLS token
+    else:
+        # fallback: pegar o primeiro tensor qualquer
+        if isinstance(outputs, (tuple, list)):
+            out = outputs[0]
+            emb = out[:, 0, :]
+        else:
+            raise RuntimeError("Não foi possível extrair embedding das saídas do modelo.")
+    return emb
+
+def preprocess_with_processor(pil_image: Image.Image):
+    """
+    Usa o AutoImageProcessor para pré-processar a PIL image em tensores PyTorch.
+    Retorna dict com tensores enviados ao device.
+    """
+    # processor aceita uma lista de imagens
+    inputs = processor(images=pil_image, return_tensors="pt")
+    # mover tensores para device
+    for k, v in inputs.items():
+        inputs[k] = v.to(DEVICE)
+    return inputs
+
+async def run_inference(pil_image: Image.Image) -> np.ndarray:
+    """
+    Executa inferência em thread pool (para não bloquear o loop do FastAPI).
+    Retorna um vetor numpy 1D (embedding L2-normalizado).
+    """
+    loop = asyncio.get_running_loop()
+    return await loop.run_in_executor(None, _sync_inference, pil_image)
+
+def _sync_inference(pil_image: Image.Image) -> np.ndarray:
+    """
+    Função síncrona que faz preprocess, forward e extrai embedding.
+    """
+    global model, processor
+    if model is None or processor is None:
+        load_model_and_processor()
+
+    inputs = preprocess_with_processor(pil_image)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        emb_tensor = extract_embedding_from_outputs(outputs)  # shape (1, dim)
+        emb = emb_tensor.cpu().numpy().reshape(-1)
+
+    # Normalizar L2
+    norm = np.linalg.norm(emb)
+    if norm > 0:
+        emb = emb / norm
+    return emb.astype(float)
+
+@app.post("/embed")
+async def embed_image(file: UploadFile = File(...)):
+    # Proteção básica: tipo e tamanho máximo (ex: 6 MB)
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Envie um arquivo de imagem.")
+    content = await file.read()
+    if len(content) > (6 * 1024 * 1024):
+        raise HTTPException(status_code=413, detail="Arquivo muito grande (max 6MB).")
+
+    try:
+        pil_img = Image.open(io.BytesIO(content)).convert("RGB")
+    except Exception:
+        raise HTTPException(status_code=400, detail="Imagem inválida.")
+
+    # Controle de concorrência
+    async with inference_semaphore:
+        try:
+            emb = await run_inference(pil_img)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Erro durante inferência: {e}")
+
+    # Retorna embedding (lista de floats)
+    return JSONResponse({"embedding": emb.tolist(), "dim": len(emb)})
+
+@app.get("/healthz")
+async def health():
+    loaded = model is not None and processor is not None
+    return {"status": "ok", "model_loaded": loaded, "model_repo": MODEL_REPO}
+
+# Ao iniciar em runtime, podemos preparar o model (opcional)
+# load_model_and_processor()