From d5c1ae1cb2c11b011e4a0b46da986ed51813eaf7 Mon Sep 17 00:00:00 2001 From: frob Date: Sun, 7 Apr 2024 13:05:39 +0200 Subject: [PATCH 1/6] Update ollama.py for image handling Some clients (eg librechat) send images in datauri format, not plain base64. Strip off the prerix when passing images to ollama. --- litellm/llms/ollama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 779896abfdd..2eb1ce26a61 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -157,7 +157,7 @@ def get_ollama_response( if format is not None: data["format"] = format if images is not None: - data["images"] = images + data["images"] = [image.split(",")[-1] if image.startswith("data:") else image for image in images] ## LOGGING logging_obj.pre_call( From 59ed4fb51e89c8e0f18972eb82f0a4da389bcf60 Mon Sep 17 00:00:00 2001 From: frob Date: Mon, 8 Apr 2024 03:28:24 +0200 Subject: [PATCH 2/6] Update ollama.py for image handling ollama wants plain base64 jpeg images, and some clients send dataURI and/or webp. Remove prefixes and convert all non-jpeg images to jpeg. --- litellm/llms/ollama.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 2eb1ce26a61..65a87e514aa 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -120,6 +120,31 @@ class OllamaConfig: and v is not None } +# ollama wants plain base64 jpeg files as images. strip any leading dataURI +# and convert to jpeg. +def _convert_image(image): + import base64, io + try: + from PIL import Image + except: + raise Exception( + "ollama image conversion failed please run `pip install Pillow`" + ) + + orig = image + if image.startswith("data:"): + image = image.split(",")[-1] + try: + image_data = Image.open(io.BytesIO(base64.b64decode(image))) + if image_data.format == "JPEG": + return image + except: + return orig + jpeg_image = io.BytesIO() + image_data.convert("RGB").save(jpeg_image, "JPEG") + jpeg_image.seek(0) + return base64.b64encode(jpeg_image.getvalue()).decode("utf-8") + # ollama implementation def get_ollama_response( @@ -157,7 +182,7 @@ def get_ollama_response( if format is not None: data["format"] = format if images is not None: - data["images"] = [image.split(",")[-1] if image.startswith("data:") else image for image in images] + data["images"] = [_convert_image(image) for image in images] ## LOGGING logging_obj.pre_call( From 82a4232dce80e841f59c06cb291d5fe653b9d1ba Mon Sep 17 00:00:00 2001 From: frob Date: Mon, 8 Apr 2024 03:35:02 +0200 Subject: [PATCH 3/6] ollama also accepts PNG --- litellm/llms/ollama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 65a87e514aa..400bf480be2 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -136,7 +136,7 @@ def _convert_image(image): image = image.split(",")[-1] try: image_data = Image.open(io.BytesIO(base64.b64decode(image))) - if image_data.format == "JPEG": + if image_data.format in ["JPEG", "PNG"]: return image except: return orig From 2492fade3a1373f756d7290165afe6edbb0c65f1 Mon Sep 17 00:00:00 2001 From: frob Date: Tue, 16 Apr 2024 01:12:24 +0200 Subject: [PATCH 4/6] Update comment. --- litellm/llms/ollama.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 860875a46be..670c565904e 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -120,8 +120,8 @@ class OllamaConfig: and v is not None } -# ollama wants plain base64 jpeg files as images. strip any leading dataURI -# and convert to jpeg. +# ollama wants plain base64 jpeg/png files as images. strip any leading dataURI +# and convert to jpeg if necessary. def _convert_image(image): import base64, io try: From 19194ece783a03115e173ac1de7f6030184c8ce2 Mon Sep 17 00:00:00 2001 From: frob Date: Tue, 28 May 2024 16:48:15 +0200 Subject: [PATCH 5/6] Add tests for ollama image handling --- litellm/tests/test_completion.py | 77 ++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index c1cf907ae19..b8b445bcceb 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -1397,6 +1397,83 @@ def test_hf_classifier_task(): pytest.fail(f"Error occurred: {str(e)}") +def test_ollama_image(): + """ + Test that datauri prefixes are removed, JPEG/PNG images are passed + through, and other image formats are converted to JPEG. Non-image + data is untouched. + """ + + import io, base64 + from PIL import Image + + def mock_post(url, **kwargs): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Type": "application/json"} + mock_response.json.return_value = { + # return the image in the response so that it can be tested + # against the original + "response": kwargs["json"]["images"] + } + return mock_response + + def make_b64image(format): + image = Image.new(mode='RGB', size=(1, 1)) + image_buffer = io.BytesIO() + image.save(image_buffer, format) + return base64.b64encode(image_buffer.getvalue()).decode("utf-8") + + jpeg_image = make_b64image("JPEG") + webp_image = make_b64image("WEBP") + png_image = make_b64image("PNG") + + base64_data = base64.b64encode(b"some random data") + datauri_base64_data = f"data:text/plain;base64,{base64_data}" + + tests = [ + # input expected + [ jpeg_image, jpeg_image ], + [ webp_image, None ], + [ png_image, png_image ], + [ f"data:image/jpeg;base64,{jpeg_image}", jpeg_image ], + [ f"data:image/webpb;base64,{webp_image}", None ], + [ f"data:image/png;base64,{png_image}", png_image ], + [ datauri_base64_data, datauri_base64_data ] + ] + + for test in tests: + try: + with patch("requests.post", side_effect=mock_post): + response = completion( + model="ollama/llava", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Whats in this image?"}, + { + "type": "image_url", + "image_url": { + "url": test[0] + }, + }, + ], + } + ], + ) + if not test[1]: + # the conversion process may not always generate the same image, + # so just check for a JPEG image when a conversion was done. + image_data = response["choices"][0]["message"]["content"][0] + image = Image.open(io.BytesIO(base64.b64decode(image_data))) + assert image.format == "JPEG" + else: + assert response["choices"][0]["message"]["content"][0] == test[1] + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + ########################### End of Hugging Face Tests ############################################## # def test_completion_hf_api(): # # failing on circle-ci commenting out From c6af8977f6874080b9be76f38fd9818e2b37a78b Mon Sep 17 00:00:00 2001 From: frob Date: Tue, 28 May 2024 17:00:08 +0200 Subject: [PATCH 6/6] Fix typo --- litellm/tests/test_completion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index b8b445bcceb..a17cff4b43f 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -1437,7 +1437,7 @@ def test_ollama_image(): [ webp_image, None ], [ png_image, png_image ], [ f"data:image/jpeg;base64,{jpeg_image}", jpeg_image ], - [ f"data:image/webpb;base64,{webp_image}", None ], + [ f"data:image/webp;base64,{webp_image}", None ], [ f"data:image/png;base64,{png_image}", png_image ], [ datauri_base64_data, datauri_base64_data ] ]