Update app.py

app.py

@@ -1,144 +1,44 @@
-# 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 torch import nn
 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
+import gradio as gr
+import numpy as np
+from PIL import Image
 
-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
+# Nome do modelo no Hugging Face Hub
+MODEL_NAME = "facebook/dinov2-small"
 
-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)
+# Carregando processador e modelo
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+model = AutoModel.from_pretrained(MODEL_NAME)
 
-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()
+# Projeção para 512D (caso a saída seja >512, reduzimos)
+projection = nn.Linear(model.config.hidden_size, 512)
 
-    inputs = preprocess_with_processor(pil_image)
+def get_embedding(image: Image.Image):
+    # Preprocessamento
+    inputs = processor(images=image, return_tensors="pt")
+    
     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()
+        # Usando o CLS token como embedding da imagem
+        last_hidden_state = outputs.last_hidden_state  # (batch, seq_len, hidden)
+        embedding = last_hidden_state[:, 0]  # pegando o [CLS] token
+
+        # Projeta para 512D
+        embedding_512 = projection(embedding)
+
+    # Converte para lista Python
+    return embedding_512.squeeze().tolist()
+
+# Cria API com Gradio (sem interface visual, apenas endpoint)
+iface = gr.Interface(
+    fn=get_embedding,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.JSON(),
+    live=False,
+    api_name="embed"  # endpoint em /embed
+)
+
+if __name__ == "__main__":
+    iface.launch()