Update app.py

app.py

@@ -1,7 +1,6 @@
 import torch
 from torch import nn
 from transformers import AutoImageProcessor, Dinov2Model
-from huggingface_hub import hf_hub_download
 from PIL import Image
 import base64
 from io import BytesIO
@@ -10,24 +9,11 @@ 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 o AutoImageProcessor, que funciona, e carregamos o modelo manualmente
+# Usamos a classe específica Dinov2Model para garantir que o modelo seja carregado corretamente
 processor = AutoImageProcessor.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)
-
+model = Dinov2Model.from_pretrained(MODEL_NAME)
 
 # Projeção para 512D
 projection = nn.Linear(model.config.hidden_size, 512)
@@ -51,6 +37,7 @@ async def get_embedding(request: ImageRequest):
         image_data = base64.b64decode(img_base64)
         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():