From 32e25cdb11a6825e17bbdd8eb877708358218f98 Mon Sep 17 00:00:00 2001 From: "hugo.lima" Date: Fri, 17 May 2024 16:46:56 +0100 Subject: [PATCH] added decoding of base64 image data in _gemini_vision_convert_messages function --- litellm/llms/prompt_templates/factory.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index cf593369c45..77c511358f2 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -1328,6 +1328,7 @@ def _gemini_vision_convert_messages(messages: list): # Case 1: Image from URL image = _load_image_from_url(img) processed_images.append(image) + else: try: from PIL import Image @@ -1335,8 +1336,22 @@ def _gemini_vision_convert_messages(messages: list): raise Exception( "gemini image conversion failed please run `pip install Pillow`" ) - # Case 2: Image filepath (e.g. temp.jpeg) given - image = Image.open(img) + + if "base64" in img: + # Case 2: Base64 image data + import base64 + import io + # Extract the base64 image data + base64_data = img.split("base64,")[1] + + # Decode the base64 image data + image_data = base64.b64decode(base64_data) + + # Load the image from the decoded data + image = Image.open(io.BytesIO(image_data)) + else: + # Case 3: Image filepath (e.g. temp.jpeg) given + image = Image.open(img) processed_images.append(image) content = [prompt] + processed_images return content