Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 from torch import nn
-from transformers import AutoImageProcessor, AutoModel
+from transformers import AutoImageProcessor, Dinov2Model
+from huggingface_hub import hf_hub_download
 from PIL import Image
 import base64
 from io import BytesIO
@@ -9,11 +10,24 @@ from pydantic import BaseModel
 
 # Nome do modelo no Hugging Face Hub
 MODEL_NAME = "facebook/dinov2-small"
+MODEL_FILE = "pytorch_model.bin" # O nome do arquivo do modelo
 
 # Carregando processador e modelo
-# Usamos a mesma lógica do seu script original que funciona
+# Usamos o AutoImageProcessor, que funciona, e carregamos o modelo manualmente
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
-model = AutoModel.from_pretrained(MODEL_NAME)
+
+# Baixa o arquivo do modelo para o cache local
+try:
+    model_path = hf_hub_download(repo_id=MODEL_NAME, filename=MODEL_FILE)
+    
+    # Carrega o modelo com o estado do arquivo baixado
+    model = Dinov2Model.from_pretrained(model_path)
+except Exception as e:
+    # Caso a primeira tentativa falhe (por ex., com Safetensors)
+    MODEL_FILE = "model.safetensors"
+    model_path = hf_hub_download(repo_id=MODEL_NAME, filename=MODEL_FILE)
+    model = Dinov2Model.from_pretrained(MODEL_NAME, state_dict=torch.load(model_path), use_safetensors=True)
+
 
 # Projeção para 512D
 projection = nn.Linear(model.config.hidden_size, 512)
@@ -33,16 +47,10 @@ 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))
         
-        # --- Lógica de Inferência do Seu Script Original ---
         inputs = processor(images=image, return_tensors="pt")
         
         with torch.no_grad():
@@ -51,7 +59,6 @@ async def get_embedding(request: ImageRequest):
             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: