Receive base64 encoded images.

server/text_generation_server/models/custom_modeling/idefics_image_processing.py

@@ -35,6 +35,7 @@ from transformers.image_utils import (
     valid_images,
 )
 from io import BytesIO
+import base64
 import requests
 from transformers import TensorType, is_torch_available
 
@@ -194,10 +195,17 @@ class IdeficsImageProcessor(BaseImageProcessor):
         if isinstance(image_url_or_urls, list):
             return [self.fetch_images(x) for x in image_url_or_urls]
         elif isinstance(image_url_or_urls, str):
+            image = image_url_or_urls
+
+            if image.startswith("http://") or image.startswith("https://"):
                 response = requests.get(image_url_or_urls, stream=True, headers=headers, timeout=(1, 5))
                 response.raise_for_status()
+                content = response.content
+            else:
+                content = base64.b64decode(image)
+
             try:
-                image = Image.open(BytesIO(response.content))
+                image = Image.open(BytesIO(content))
                 # image.verify()
             except Exception:
                 raise ValueError(f"Could not load image from url {image_url_or_urls}")