Update app.py

app.py

@@ -1,6 +1,6 @@
 import torch
 from torch import nn
-from transformers import pipeline
+from transformers import AutoImageProcessor, AutoModel
 from PIL import Image
 import base64
 from io import BytesIO
@@ -10,16 +10,13 @@ from pydantic import BaseModel
 # Nome do modelo no Hugging Face Hub
 MODEL_NAME = "facebook/dinov2-small"
 
-# Usando um pipeline para carregar o modelo e o processador
-feature_extractor = pipeline(
-    "feature-extraction",
-    model=MODEL_NAME
-)
+# Carregando processador e modelo
+# Usamos a mesma lógica do seu script original que funciona
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+model = AutoModel.from_pretrained(MODEL_NAME)
 
 # Projeção para 512D
-# O pipeline retorna um tensor, então a projeção ainda é necessária
-# Você pode remover isso se o embedding de 768D for suficiente
-projection = nn.Linear(768, 512)
+projection = nn.Linear(model.config.hidden_size, 512)
 
 # Inicializa o FastAPI
 app = FastAPI(
@@ -36,19 +33,25 @@ class ImageRequest(BaseModel):
 @app.post("/embed")
 async def get_embedding(request: ImageRequest):
     try:
+        # Extrai a string Base64 do formato "data:image/png;base64,..."
         header, img_base64 = request.image.split(",", 1)
+        
+        # Decodifica a string Base64
         image_data = base64.b64decode(img_base64)
-        image = Image.open(BytesIO(image_data))
         
-        # Gera o embedding usando o pipeline
-        embedding_list = feature_extractor(images=[image])[0][0]
+        # Abre a imagem com Pillow
+        image = Image.open(BytesIO(image_data))
         
-        # Converte a lista de embeddings para um tensor PyTorch para a projeção
-        embedding_tensor = torch.tensor(embedding_list)
+        # --- Lógica de Inferência do Seu Script Original ---
+        inputs = processor(images=image, return_tensors="pt")
         
         with torch.no_grad():
-            embedding_512 = projection(embedding_tensor)
+            outputs = model(**inputs)
+            last_hidden_state = outputs.last_hidden_state
+            embedding = last_hidden_state[:, 0]
+            embedding_512 = projection(embedding)
             
+        # Converte para lista Python e retorna
         return {"embedding": embedding_512.squeeze().tolist()}
         
     except Exception as e: