ok

app/services/embeddings_service.py

@@ -5,8 +5,11 @@ import os
 from typing import Dict, List, Optional
 
 import numpy as np
+import torch
+import torch.nn.functional as F
 from PIL import Image
 from sentence_transformers import SentenceTransformer
+from transformers import AutoImageProcessor, AutoModel
 
 _logger = logging.getLogger(__name__)
 
@@ -33,7 +36,8 @@ class EmbeddingService:
             self._device = "cpu"
 
         self._loaded_dimensions: List[int] = []
-        self._vision_model: Optional[SentenceTransformer] = None
+        self._vision_processor: Optional[AutoImageProcessor] = None
+        self._vision_model: Optional[AutoModel] = None
         self._vision_loaded = False
 
     def load_model(self, dimension: int) -> None:
@@ -77,12 +81,14 @@ class EmbeddingService:
                 _logger.info("Patched vision model config: n_inner float -> int")
 
         _logger.info("Loading vision embedding model from %s", source)
-        self._vision_model = SentenceTransformer(
+        self._vision_processor = AutoImageProcessor.from_pretrained(source)
+        self._vision_model = AutoModel.from_pretrained(
             source,
-            device=self._device,
             trust_remote_code=True,
+            _fast_init=False,
         )
         self._vision_model.eval()
+        self._vision_model.to(self._device)
         self._vision_loaded = True
         _logger.info("Loaded vision embedding model (device=%s)", self._device)
 
@@ -99,14 +105,18 @@ class EmbeddingService:
         return result.tolist()
 
     def generate_image_embedding(self, images: List[Image.Image]) -> List[List[float]]:
-        if not self._vision_loaded or self._vision_model is None:
+        if not self._vision_loaded or self._vision_model is None or self._vision_processor is None:
             raise ValueError("Vision model not loaded")
-        result: np.ndarray = self._vision_model.encode(
-            images,
-            convert_to_numpy=True,
-            show_progress_bar=False,
-        )
-        return result.tolist()
+        all_embeddings: List[List[float]] = []
+        with torch.no_grad():
+            for image in images:
+                inputs = self._vision_processor(image, return_tensors="pt")
+                inputs = {k: v.to(self._device) for k, v in inputs.items()}
+                outputs = self._vision_model(**inputs)
+                emb = outputs.last_hidden_state[:, 0]
+                emb = F.normalize(emb, p=2, dim=1)
+                all_embeddings.append(emb.cpu().numpy().flatten().tolist())
+        return all_embeddings
 
     @property
     def loaded_dimensions(self) -> List[int]: