Update app.py

app.py

@@ -1,6 +1,6 @@
 import torch
 from torch import nn
-from transformers import AutoImageProcessor, Dinov2Model
+from transformers import pipeline
 from PIL import Image
 import base64
 from io import BytesIO
@@ -10,12 +10,16 @@ from pydantic import BaseModel
 # Nome do modelo no Hugging Face Hub
 MODEL_NAME = "facebook/dinov2-small"
 
-# Carregando processador e modelo
-processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
-model = Dinov2Model.from_pretrained(MODEL_NAME)
+# Usando um pipeline para carregar o modelo e o processador
+feature_extractor = pipeline(
+    "feature-extraction",
+    model=MODEL_NAME
+)
 
 # Projeção para 512D
-projection = nn.Linear(model.config.hidden_size, 512)
+# 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)
 
 # Inicializa o FastAPI
 app = FastAPI(
@@ -32,25 +36,19 @@ 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)
-        
-        # Abre a imagem com Pillow
         image = Image.open(BytesIO(image_data))
         
-        # Preprocessamento
-        inputs = processor(images=image, return_tensors="pt")
+        # Gera o embedding usando o pipeline
+        embedding_list = feature_extractor(images=[image])[0][0]
+        
+        # Converte a lista de embeddings para um tensor PyTorch para a projeção
+        embedding_tensor = torch.tensor(embedding_list)
         
         with torch.no_grad():
-            outputs = model(**inputs)
-            last_hidden_state = outputs.last_hidden_state
-            embedding = last_hidden_state[:, 0]
-            embedding_512 = projection(embedding)
+            embedding_512 = projection(embedding_tensor)
             
-        # Converte para lista Python e retorna
         return {"embedding": embedding_512.squeeze().tolist()}
         
     except Exception as e: