refactor: refine load images

modeling_clip.py

@@ -223,6 +223,7 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
             self.text_projection = nn.Identity()
 
         self.tokenizer = None
+        self.preprocess = None
         self.post_init()
 
     def get_text_features(
@@ -249,7 +250,7 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
 
     def get_tokenizer(self):
         if not self.tokenizer:
-            self.tokenizer = AutoTokenizer.from_pretrained(config._name_or_path)
+            self.tokenizer = AutoTokenizer.from_pretrained(config._name_or_path, trust_remote_code=True)
         return self.tokenizer
 
     @torch.inference_mode()
@@ -264,7 +265,7 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = False,
         **tokenizer_kwargs,
-    ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]::
+    ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
        Computes sentence embeddings
         Args:
@@ -373,19 +374,81 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         self.train(is_training)
         return all_embeddings
 
+    def get_preprocess(self):
+        if not self.preprocess:
+            self.preprocess = AutoImageProcessor.from_pretrained(config._name_or_path, trust_remote_code=True)
+        return self.preprocess
+
+    
+    @torch.inference_mode()
     def encode_image(
         self,
-        pixel_values: Union[None, torch.FloatTensor, BatchFeature] = None,
-        return_dict: Optional[bool] = None,
-        *_,
-        **__,
-    ) -> Union[Tuple[Optional[torch.FloatTensor], ...], CLIPVisionModelOutput]:
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
-        feats = self.get_image_features(pixel_values=pixel_values)
-        out = CLIPVisionModelOutput(image_embeds=feats)
-        return out if return_dict else out.to_tuple()
+        images: Union[str, List[str]],
+        batch_size: int = 32,
+        show_progress_bar: Optional[bool] = None,
+        convert_to_numpy: bool = True,
+        convert_to_tensor: bool = False,
+        device: Optional[torch.device] = None,
+        normalize_embeddings: bool = False,
+    ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
+        from PIL.Image import  Image
+        
+        is_training = self.training
+        self.eval()
+
+        self.preprocess = self.get_preprocess()
+
+        if show_progress_bar is None:
+            show_progress_bar = (
+                logger.getEffectiveLevel() == logging.INFO
+                or logger.getEffectiveLevel() == logging.DEBUG
+            )
+
+        if convert_to_tensor:
+            convert_to_numpy = False
+
+        input_was_single_img = False
+        if isinstance(images, str) or not hasattr(images, '__len__'):
+            images = [images]
+            input_was_single_img = True
+
+        if device is not None:
+            self.to(device)
+
+        permutation = np.argsort([-len(i) for i in images])
+        inverse_permutation = np.argsort(permutation)
+        images = [images[idx] for idx in permutation]
+
+        if has_tqdm:
+            range_iter = trange(
+                0,
+                len(sentences),
+                batch_size,
+                desc="Encoding",
+                disable=not show_progress_bar,
+            )
+        else:
+            range_iter = range(0, len(sentences), batch_size)
+
+        for i in range_iter:
+            processed_inputs = self.process([Image.open(image) for image in images])
+            embeddings = self.get_image_features(processed_inputs)
+            if convert_to_numpy:
+                embeddings = embeddings.cpu()
+            all_embeddings.extend(embeddings)
+
+        all_embeddings = [all_embeddings[idx] for idx in inverse_permutation]
+
+        if convert_to_tensor:
+            all_embeddings = torch.stack(all_embeddings)
+        elif convert_to_numpy:
+            all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
+
+        if input_was_single_img:
+            all_embeddings = all_embeddings[0]
+
+        self.train(is_training)
+        return all_embeddings
 
     def forward(
         self,