refactor(token_counter): split image_dimensions_from_bytes out of get_image_dimensions

get_image_dimensions fetched or decoded its input, parsed the header, and
substituted a 300x300 default when the header was unreadable, all in one
function. Billing code needs the middle step on its own: given bytes, return
the dimensions or None, never a default. The parsers are unchanged and the
string entry point still applies the default, so its callers see no
difference. Unreadable headers used to raise TypeError or struct.error out of
the parser; the bytes entry point turns those into None

litellm/litellm_core_utils/token_counter.py

@@ -248,15 +248,29 @@ def get_image_dimensions(
         _header, encoded = data.split(",", 1)
         img_data = base64.b64decode(encoded)
 
+    dimensions: Final = image_dimensions_from_bytes(img_data)
+    if dimensions is None:
+        return DEFAULT_IMAGE_WIDTH, DEFAULT_IMAGE_HEIGHT
+    return dimensions
+
+
+def image_dimensions_from_bytes(img_data: bytes) -> tuple[int, int] | None:
+    try:
+        return _header_dimensions(img_data)
+    except (struct.error, TypeError):
+        return None
+
+
+def _header_dimensions(img_data: bytes) -> tuple[int, int] | None:
     img_type: Final = get_image_type(img_data)
 
     if img_type == "png":
         w, h = _unpack_ints(">LL", img_data[16:24])
         return w, h
-    elif img_type == "gif":
+    if img_type == "gif":
         w, h = _unpack_ints("<HH", img_data[6:10])
         return w, h
-    elif img_type == "jpeg":
+    if img_type == "jpeg":
         with io.BytesIO(img_data) as fhandle:
             fhandle.seek(0)
             size = 2
@@ -271,27 +285,21 @@ def get_image_dimensions(
             fhandle.seek(1, 1)
             h, w = _unpack_ints(">HH", fhandle.read(4))
         return w, h
-    elif img_type == "webp":
-        # For WebP, the dimensions are stored at different offsets depending on the format
-        # Check for VP8X (extended format)
+    if img_type == "webp":
         if img_data[12:16] == b"VP8X":
             w = _unpack_ints("<I", img_data[24:27] + b"\x00")[0] + 1
             h = _unpack_ints("<I", img_data[27:30] + b"\x00")[0] + 1
             return w, h
-        # Check for VP8 (lossy format)
-        elif img_data[12:16] == b"VP8 ":
+        if img_data[12:16] == b"VP8 ":
             w = _unpack_ints("<H", img_data[26:28])[0] & 0x3FFF
             h = _unpack_ints("<H", img_data[28:30])[0] & 0x3FFF
             return w, h
-        # Check for VP8L (lossless format)
-        elif img_data[12:16] == b"VP8L":
+        if img_data[12:16] == b"VP8L":
             bits: Final = _unpack_ints("<I", img_data[21:25])[0]
             w = (bits & 0x3FFF) + 1
             h = ((bits >> 14) & 0x3FFF) + 1
             return w, h
-
-    # return sensible default image dimensions if unable to get dimensions
-    return DEFAULT_IMAGE_WIDTH, DEFAULT_IMAGE_HEIGHT
+    return None
 
 
 def calculate_img_tokens(

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -3,6 +3,7 @@
 import asyncio
 import base64
 import importlib
+import struct
 import threading
 import time
 import traceback
@@ -28,6 +29,7 @@ from litellm.litellm_core_utils.token_counter import (
     _get_tiktoken_count_function,
     calculate_img_tokens,
     high_detail_image_token_upper_bound,
+    image_dimensions_from_bytes,
     offload_token_count,
 )
 from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
@@ -1599,3 +1601,75 @@ def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int,
 def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None:
     assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound()
     assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound()
+
+
+def _png_bytes(width: int, height: int) -> bytes:
+    return (
+        b"\x89PNG\r\n\x1a\n"
+        + (13).to_bytes(4, "big")
+        + b"IHDR"
+        + struct.pack(">II", width, height)
+        + b"\x08\x02\x00\x00\x00"
+    )
+
+
+def _gif_bytes(width: int, height: int) -> bytes:
+    return b"GIF89a" + struct.pack("<HH", width, height) + b"\x00\x00\x00"
+
+
+def _jpeg_bytes(width: int, height: int, sof_marker: bytes, app_segments: int) -> bytes:
+    app: Final = b"".join(
+        b"\xff\xe0" + struct.pack(">H", 16) + b"JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00"
+        for _ in range(app_segments)
+    )
+    sof: Final = (
+        b"\xff" + sof_marker + struct.pack(">HBHHB", 17, 8, height, width, 3) + b"\x01\x22\x00\x02\x11\x01\x03\x11\x01"
+    )
+    return b"\xff\xd8" + app + sof
+
+
+def _webp_bytes(chunk: bytes, payload: bytes) -> bytes:
+    body: Final = chunk + struct.pack("<I", len(payload)) + payload
+    return b"RIFF" + struct.pack("<I", 4 + len(body)) + b"WEBP" + body
+
+
+def _webp_vp8_bytes(width: int, height: int) -> bytes:
+    return _webp_bytes(b"VP8 ", b"\x00\x00\x00\x9d\x01\x2a" + struct.pack("<HH", width, height))
+
+
+def _webp_vp8l_bytes(width: int, height: int) -> bytes:
+    return _webp_bytes(b"VP8L", b"\x2f" + struct.pack("<I", (width - 1) | ((height - 1) << 14)))
+
+
+def _webp_vp8x_bytes(width: int, height: int) -> bytes:
+    return _webp_bytes(b"VP8X", b"\x00" * 4 + (width - 1).to_bytes(3, "little") + (height - 1).to_bytes(3, "little"))
+
+
+@pytest.mark.parametrize(
+    ("image", "expected"),
+    [
+        pytest.param(_png_bytes(1024, 768), (1024, 768), id="png"),
+        pytest.param(_gif_bytes(100, 50), (100, 50), id="gif"),
+        pytest.param(_jpeg_bytes(800, 600, b"\xc0", 1), (800, 600), id="jpeg-baseline"),
+        pytest.param(_jpeg_bytes(640, 480, b"\xc2", 3), (640, 480), id="jpeg-progressive-after-app-segments"),
+        pytest.param(_webp_vp8_bytes(640, 480), (640, 480), id="webp-vp8"),
+        pytest.param(_webp_vp8l_bytes(320, 240), (320, 240), id="webp-vp8l"),
+        pytest.param(_webp_vp8x_bytes(1920, 1080), (1920, 1080), id="webp-vp8x"),
+    ],
+)
+def test_image_dimensions_from_bytes_reads_each_header_format(image: bytes, expected: tuple[int, int]) -> None:
+    assert image_dimensions_from_bytes(image) == expected
+
+
+@pytest.mark.parametrize(
+    "image",
+    [
+        pytest.param(b"", id="empty"),
+        pytest.param(b"BM" + b"\x00" * 30, id="unknown-format"),
+        pytest.param(b"\x89PNG\r\n\x1a\n\x00\x00", id="png-truncated-before-ihdr"),
+        pytest.param(b"\xff\xd8\xff\xe0\x00\x10JFIF", id="jpeg-truncated-inside-app0"),
+        pytest.param(b"\xff\xd8\xff\xe0\x00\x04\x00\x00", id="jpeg-ends-before-sof"),
+    ],
+)
+def test_image_dimensions_from_bytes_returns_none_for_unreadable_headers(image: bytes) -> None:
+    assert image_dimensions_from_bytes(image) is None