Update app.py

app.py

@@ -11,7 +11,8 @@ import imagehash
 MODEL_NAME = "facebook/dinov2-small"
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
 model = Dinov2Model.from_pretrained(MODEL_NAME)
-projection = nn.Linear(model.config.hidden_size, 512)
+# A camada de projeção para 512 dimensões agora é criada dentro da função,
+# para permitir a escolha entre 384 e 512.
 
 app = FastAPI(
     title="API de Embedding de Imagem",
@@ -21,7 +22,8 @@ app = FastAPI(
 
 class ImageRequest(BaseModel):
     image: str
-    use_float16: bool = False # <-- NOVO: Parâmetro opcional com valor padrão False
+    target_dim: int = 512  # <-- NOVO: Parâmetro opcional para a dimensão do embedding
+    use_float16: bool = False # <-- NOVO: Parâmetro opcional para usar float16
 
 @app.post("/embed")
 async def get_embedding(request: ImageRequest):
@@ -30,19 +32,28 @@ async def get_embedding(request: ImageRequest):
         image_data = base64.b64decode(img_base64)
         image = Image.open(BytesIO(image_data)).convert("RGB")
         inputs = processor(images=image, return_tensors="pt")
+
         with torch.no_grad():
             outputs = model(**inputs)
             last_hidden_state = outputs.last_hidden_state
             embedding = last_hidden_state[:, 0]
-            embedding_512 = projection(embedding)
-
-        # <-- NOVA LÓGICA: Conversão condicional para float16
-        if request.use_float16:
-            embedding_512 = embedding_512.half()
+            
+            # Lógica para definir a dimensão
+            if request.target_dim == 384:
+                final_embedding = embedding
+            elif request.target_dim == 512:
+                projection = nn.Linear(model.config.hidden_size, 512)
+                final_embedding = projection(embedding)
+            else:
+                raise HTTPException(status_code=400, detail="Dimensão inválida. Escolha entre 384 ou 512.")
+            
+            # Lógica para a conversão para float16
+            if request.use_float16:
+                final_embedding = final_embedding.half()
 
         phash = str(imagehash.phash(image))
         return {
-            "embedding": embedding_512.squeeze().tolist(),
+            "embedding": final_embedding.squeeze().tolist(),
             "phash": phash
         }
     except Exception as e: