fix: handle incomplete and non-base64 image data URLs

CHANGELOG.md

@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- fix: raise a clear error for data URLs missing a comma and decode non-base64 payloads by @kevin9327 in #2413
 - fix(ci): use supported Python versions on Ubuntu by @abetlen in #2407
 - feat: update llama.cpp to ggml-org/llama.cpp@7d7f02875
 

llama_cpp/_utils.py

@@ -97,3 +97,27 @@ class Singleton(object, metaclass=MetaSingleton):
 
     def __init__(self):
         super(Singleton, self).__init__()
+
+
+def _bytes_from_image_url(image_url: str) -> bytes:
+    """Load image bytes from a data URL or an http(s) URL."""
+    if image_url.startswith("data:"):
+        import base64
+
+        if "," not in image_url:
+            raise ValueError(
+                "Invalid data URL: expected 'data:<type>[;base64],<payload>'"
+            )
+        header, data = image_url.split(",", 1)
+        if ";base64" in header.lower():
+            return base64.b64decode(data)
+        # Non-base64 data URLs (e.g. data:image/svg+xml,<svg>...) carry the payload as text.
+        from urllib.parse import unquote_to_bytes
+
+        return unquote_to_bytes(data)
+
+    import urllib.request
+
+    with urllib.request.urlopen(image_url) as f:
+        return f.read()
+

llama_cpp/llama_chat_format.py

@@ -35,7 +35,7 @@
 import llama_cpp.llama_grammar as llama_grammar
 
 from ._logger import logger
-from ._utils import suppress_stdout_stderr, Singleton
+from ._utils import suppress_stdout_stderr, Singleton, _bytes_from_image_url
 
 ### Common Chat Templates and Special Tokens ###
 
@@ -3124,17 +3124,7 @@ def __call__(
     @staticmethod
     def _load_image(image_url: str) -> bytes:
         # TODO: Add Pillow support for other image formats beyond (jpg, png)
-        if image_url.startswith("data:"):
-            import base64
-
-            image_bytes = base64.b64decode(image_url.split(",")[1])
-            return image_bytes
-        else:
-            import urllib.request
-
-            with urllib.request.urlopen(image_url) as f:
-                image_bytes = f.read()
-                return image_bytes
+        return _bytes_from_image_url(image_url)
 
     @staticmethod
     def get_image_urls(messages: List[llama_types.ChatCompletionRequestMessage]):
@@ -3658,17 +3648,7 @@ def raise_exception(message: str):
 
     @staticmethod
     def _load_image(image_url: str) -> bytes:
-        if image_url.startswith("data:"):
-            import base64
-
-            image_bytes = base64.b64decode(image_url.split(",")[1])
-            return image_bytes
-        else:
-            import urllib.request
-
-            with urllib.request.urlopen(image_url) as f:
-                image_bytes = f.read()
-                return image_bytes
+        return _bytes_from_image_url(image_url)
 
     @staticmethod
     def get_image_urls(messages: List[llama_types.ChatCompletionRequestMessage]):

tests/test_image_url.py

@@ -0,0 +1,49 @@
+import base64
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+_UTILS = Path(__file__).resolve().parents[1] / "llama_cpp" / "_utils.py"
+_spec = importlib.util.spec_from_file_location("llama_cpp_utils_under_test", _UTILS)
+_utils = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_utils)
+_bytes_from_image_url = _utils._bytes_from_image_url
+
+
+def test_base64_data_url_round_trip():
+    payload = b"\x89PNG\r\n"
+    url = "data:image/png;base64," + base64.b64encode(payload).decode("ascii")
+    assert _bytes_from_image_url(url) == payload
+
+
+def test_base64_data_url_is_case_insensitive():
+    payload = b"hello"
+    url = "data:image/png;BASE64," + base64.b64encode(payload).decode("ascii")
+    assert _bytes_from_image_url(url) == payload
+
+
+def test_data_url_without_comma_raises():
+    with pytest.raises(ValueError, match="Invalid data URL"):
+        _bytes_from_image_url("data:image/png;base64")
+
+
+def test_plain_data_url_is_unquoted():
+    assert _bytes_from_image_url("data:image/svg+xml,%3Csvg%3E") == b"<svg>"
+
+
+def test_http_url(monkeypatch: pytest.MonkeyPatch):
+    import urllib.request
+
+    class FakeResp:
+        def read(self):
+            return b"img"
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *args):
+            return False
+
+    monkeypatch.setattr(urllib.request, "urlopen", lambda url: FakeResp())
+    assert _bytes_from_image_url("https://example.com/a.png") == b"img"