fix: encode_image and encode_text

- support bfloat16
- support image loader from url, PIL.image and data:image/ strings
- fix bug and optimize encode_image efficiency

modeling_clip.py

@@ -6,7 +6,10 @@
 
 from functools import partial
 from typing import List, Optional, Tuple, Union
-
+from PIL import Image
+from io import BytesIO
+import requests
+import base64
 import numpy as np
 import torch
 import torch.nn.functional as f
@@ -373,7 +376,7 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         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])
+            all_embeddings = np.asarray([emb.to(torch.float32).numpy() for emb in all_embeddings])
 
         if input_was_string:
             all_embeddings = all_embeddings[0]
@@ -381,10 +384,15 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         self.train(is_training)
         return all_embeddings
 
+    def decode_data_image(data_image_str):
+        header, data = data_image_str.split(',', 1)
+        image_data = base64.b64decode(data)
+        return Image.open(BytesIO(image_data))
+
     @torch.inference_mode()
     def encode_image(
         self,
-        images: Union[str, List[str]],
+        images: Union[str, List[Union[str, Image.Image]]],
         batch_size: int = 32,
         show_progress_bar: Optional[bool] = None,
         convert_to_numpy: bool = True,
@@ -394,10 +402,10 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
         Computes image embeddings.
-
+    
         Args:
-            images(`str` or `List[str]`):
-                image or images paths to be encoded
+            images(`str` or `List[Union[str, Image.Image]]`):
+                image paths, URLs, PIL images, or data:image/ strings to be encoded
             batch_size(`int`, *optional*, defaults to 32):
                 Batch size for the computation
             show_progress_bar(`bool`, *optional*, defaults to None):
@@ -421,35 +429,34 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
             If convert_to_tensor, a stacked tensor is returned.
             If convert_to_numpy, a numpy matrix is returned.
         """
-        from PIL import Image
-
+        
         is_training = self.training
         self.eval()
-
+    
         self.preprocess = self.get_preprocess()
         all_embeddings = []
-
+    
         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])
+    
+        permutation = np.argsort([-len(str(i)) for i in images])
         inverse_permutation = np.argsort(permutation)
         images = [images[idx] for idx in permutation]
-
+    
         if has_tqdm:
             range_iter = trange(
                 0,
@@ -460,26 +467,46 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
             )
         else:
             range_iter = range(0, len(images), batch_size)
-
-        for _ in range_iter:
-            processed_inputs = self.preprocess([Image.open(image) for image in images])
+    
+        for i in range_iter:
+            batch_images = images[i:i+batch_size]
+            processed_inputs = []
+    
+            for img in batch_images:
+                if isinstance(img, str):
+                    if img.startswith('http'):
+                        response = requests.get(img)
+                        image = Image.open(BytesIO(response.content)).convert('RGB')
+                    elif img.startswith('data:image/'):
+                        image = decode_data_image(img).convert('RGB')
+                    else:
+                        image = Image.open(img).convert('RGB')
+                elif isinstance(img, Image.Image):
+                    image = img.convert('RGB')
+                else:
+                    raise ValueError("Unsupported image format")
+    
+                processed_inputs.append(self.preprocess(image))
+    
+            processed_inputs = torch.stack(processed_inputs).to(device)
             embeddings = self.get_image_features(processed_inputs)
+            
             if normalize_embeddings:
                 embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
             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])
-
+            all_embeddings = np.asarray([emb.to(torch.float32).numpy() for emb in all_embeddings])
+    
         if input_was_single_img:
             all_embeddings = all_embeddings[0]
-
+    
         self.train(is_training)
         return all_embeddings