feat(images): stamp measured reference pixels on image edit responses

litellm/llms/custom_httpx/llm_http_handler.py

@@ -33,6 +33,7 @@ from litellm._logging import _redact_string, verbose_logger
 from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta
 from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
 from litellm.files.types import FileContentStreamingResult
+from litellm.images.dimensions import total_reference_pixels
 from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs
 from litellm.litellm_core_utils.agentic_loop_settings import (
     DEFAULT_MAX_AGENTIC_LOOPS,
@@ -382,6 +383,12 @@ def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]:
     )
 
 
+def _with_reference_pixels(response: ImageResponse, reference_pixels: int | None) -> ImageResponse:
+    if reference_pixels is not None:
+        response._reference_pixels = reference_pixels  # mutable-ok: stamp the measured reference pixels
+    return response
+
+
 class BaseLLMHTTPHandler:
     async def _make_common_async_call(
         self,
@@ -6881,6 +6888,7 @@ class BaseLLMHTTPHandler:
             litellm_params=dict(litellm_params),
         )
 
+        reference_pixels: Final = total_reference_pixels(image if isinstance(image, list) else [image])
         data, files = image_edit_provider_config.transform_image_edit_request(
             model=model,
             image=image,
@@ -6928,10 +6936,13 @@ class BaseLLMHTTPHandler:
                 provider_config=image_edit_provider_config,
             )
 
-        return image_edit_provider_config.transform_image_edit_response(
-            model=model,
-            raw_response=response,
-            logging_obj=logging_obj,
+        return _with_reference_pixels(
+            image_edit_provider_config.transform_image_edit_response(
+                model=model,
+                raw_response=response,
+                logging_obj=logging_obj,
+            ),
+            reference_pixels,
         )
 
     async def async_image_edit_handler(
@@ -6980,6 +6991,7 @@ class BaseLLMHTTPHandler:
             litellm_params=dict(litellm_params),
         )
 
+        reference_pixels: Final = total_reference_pixels(image if isinstance(image, list) else [image])
         data, files = await image_edit_provider_config.async_transform_image_edit_request(
             model=model,
             image=image,
@@ -7027,10 +7039,13 @@ class BaseLLMHTTPHandler:
                 provider_config=image_edit_provider_config,
             )
 
-        return image_edit_provider_config.transform_image_edit_response(
-            model=model,
-            raw_response=response,
-            logging_obj=logging_obj,
+        return _with_reference_pixels(
+            image_edit_provider_config.transform_image_edit_response(
+                model=model,
+                raw_response=response,
+                logging_obj=logging_obj,
+            ),
+            reference_pixels,
         )
 
     def image_generation_handler(

litellm/types/utils.py

@@ -318,6 +318,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     input_cost_per_audio_token_batches: ReadOnly[float | None]
     input_cost_per_image_token_batches: ReadOnly[float | None]
     input_cost_per_second: float | None  # for OpenAI Speech models
+    input_cost_per_pixel: ReadOnly[float | None]
+    input_cost_per_reference_pixel: ReadOnly[float | None]
     input_cost_per_token_batches: float | None
     input_cost_per_video_token_batches: ReadOnly[float | None]
     input_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None]
@@ -2582,6 +2584,7 @@ from openai.types.images_response import ImagesResponse as OpenAIImageResponse
 
 class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject):
     _hidden_params: dict = {}
+    _reference_pixels: int | None = None
 
     usage: ImageUsage | None = None
     """
@@ -3702,6 +3705,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
     output_cost_per_image_1024: float | None = None
     output_cost_per_image_1536: float | None = None
     input_cost_per_pixel: float | None = None
+    input_cost_per_reference_pixel: float | None = None
     output_cost_per_pixel: float | None = None
 
     # Include all ModelInfoBase fields as optional

tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py

@@ -2,6 +2,7 @@ import asyncio
 import base64
 import json
 import logging
+import struct
 import threading
 import time
 from typing import Final
@@ -3898,6 +3899,77 @@ def test_image_edit_handler_keeps_the_sync_transform():
     assert response.data[0].b64_json == "sync"
 
 
+def _tiny_png(width: int, height: int) -> bytes:
+    return (
+        b"\x89PNG\r\n\x1a\n"
+        + struct.pack(">I", 13)
+        + b"IHDR"
+        + struct.pack(">II", width, height)
+        + b"\x08\x06\x00\x00\x00"
+        + bytes(4)
+    )
+
+
+def test_image_edit_handler_stamps_measured_reference_pixels():
+    client = HTTPHandler()
+    client.client = httpx.Client(transport=_echo_json_transport({}))
+
+    response = BaseLLMHTTPHandler().image_edit_handler(
+        model="edit-model",
+        image=[_tiny_png(4, 2), _tiny_png(1, 1)],
+        prompt="add a hat",
+        image_edit_provider_config=_ImageEditRecordingConfig(),
+        image_edit_optional_request_params={},
+        custom_llm_provider="openai",
+        litellm_params=GenericLiteLLMParams(),
+        logging_obj=Mock(),
+        timeout=10.0,
+        client=client,
+    )
+
+    assert response._reference_pixels == 4 * 2 + 1 * 1
+
+
+async def test_async_image_edit_handler_stamps_measured_reference_pixels():
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=_echo_json_transport({}))
+
+    response = await BaseLLMHTTPHandler().async_image_edit_handler(
+        model="edit-model",
+        image=_tiny_png(4, 2),
+        prompt="add a hat",
+        image_edit_provider_config=_ImageEditRecordingConfig(),
+        image_edit_optional_request_params={},
+        custom_llm_provider="openai",
+        litellm_params=GenericLiteLLMParams(),
+        logging_obj=Mock(),
+        timeout=10.0,
+        client=client,
+    )
+
+    assert response._reference_pixels == 8
+
+
+def test_image_edit_handler_leaves_reference_pixels_unset_when_unmeasurable():
+    client = HTTPHandler()
+    client.client = httpx.Client(transport=_echo_json_transport({}))
+
+    response = BaseLLMHTTPHandler().image_edit_handler(
+        model="edit-model",
+        image=b"not-an-image",
+        prompt="add a hat",
+        image_edit_provider_config=_ImageEditRecordingConfig(),
+        image_edit_optional_request_params={},
+        custom_llm_provider="openai",
+        litellm_params=GenericLiteLLMParams(),
+        logging_obj=Mock(),
+        timeout=10.0,
+        client=client,
+    )
+
+    assert response._reference_pixels is None
+
+
 class _ScriptedClientWebSocket(_FakeClientWebSocket):
     def __init__(self, messages: list[str], last_event_type: str) -> None:
         super().__init__()