fix(azure_ai): bill FLUX.2-flex edit reference megapixels at Azure's rate

litellm-rust/crates/model-catalog/src/model_info.rs

@@ -347,6 +347,8 @@ pub struct ModelInfo {
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub input_cost_per_query: Option<f64>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub input_cost_per_reference_pixel: Option<f64>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
     pub input_cost_per_request: Option<f64>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub input_cost_per_second: Option<f64>,

litellm/images/utils.py

@@ -1,12 +1,157 @@
-from collections.abc import Collection, Mapping
+import os
+import struct
+from collections.abc import Collection, Mapping, Sequence
 from io import BufferedReader, BytesIO
-from typing import Any, Final, cast, get_type_hints
+from typing import IO, Any, Final, cast, get_type_hints
 
 import litellm
+from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.token_counter import get_image_type
 from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
 from litellm.types.files import FILE_MIME_TYPES, FileType
 from litellm.types.images.main import ImageEditOptionalRequestParams
+from litellm.types.llms.openai import FileTypes
+
+_JPEG_SOF_MARKERS: Final = frozenset(range(0xC0, 0xD0)) - {0xC4, 0xC8, 0xCC}
+_MAX_JPEG_SEGMENTS: Final = 1024
+_MAX_JPEG_HEADER_OFFSET: Final = 16 * 1024 * 1024
+_JPEG_FIRST_SEGMENT_OFFSET: Final = 2
+_JPEG_SOF_PAYLOAD_SIZE: Final = 5
+_HEADER_READ_SIZE: Final = 32
+_MIN_PNG_HEADER_SIZE: Final = 24
+_MIN_GIF_HEADER_SIZE: Final = 10
+_MIN_BMP_HEADER_SIZE: Final = 26
+_MIN_WEBP_HEADER_SIZE: Final = 30
+
+
+def _content_stream(image: FileTypes) -> IO[bytes] | None:
+    content: Final = image[1] if isinstance(image, tuple) else image
+    if isinstance(content, bytes):
+        return BytesIO(content)
+    if isinstance(content, (str, os.PathLike)):
+        return None
+    return content if content.seekable() else None
+
+
+def _png_dimensions(head: bytes) -> tuple[int, int] | None:
+    if len(head) < _MIN_PNG_HEADER_SIZE:
+        return None
+    width, height = struct.unpack(">II", head[16:24])
+    return width, height
+
+
+def _gif_dimensions(head: bytes) -> tuple[int, int] | None:
+    if len(head) < _MIN_GIF_HEADER_SIZE:
+        return None
+    width: Final = int.from_bytes(head[6:8], "little")
+    height: Final = int.from_bytes(head[8:10], "little")
+    return width, height
+
+
+def _bmp_dimensions(head: bytes) -> tuple[int, int] | None:
+    if len(head) < _MIN_BMP_HEADER_SIZE:
+        return None
+    if int.from_bytes(head[14:18], "little") == 12:
+        core_width: Final = int.from_bytes(head[18:20], "little")
+        core_height: Final = int.from_bytes(head[20:22], "little")
+        return core_width, core_height
+    width: Final = int.from_bytes(head[18:22], "little", signed=True)
+    height: Final = int.from_bytes(head[22:26], "little", signed=True)
+    return width, abs(height)
+
+
+def _webp_dimensions(head: bytes) -> tuple[int, int] | None:
+    if len(head) < _MIN_WEBP_HEADER_SIZE:
+        return None
+    match head[12:16]:  # pyright: ignore[reportMatchNotExhaustive]  # unknown fourccs fall through to the None below
+        case b"VP8X":
+            return int.from_bytes(head[24:27], "little") + 1, int.from_bytes(head[27:30], "little") + 1
+        case b"VP8 ":
+            width, height = struct.unpack("<HH", head[26:30])
+            return width & 0x3FFF, height & 0x3FFF
+        case b"VP8L":
+            bits: Final = int.from_bytes(head[21:25], "little")
+            return (bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1
+    return None
+
+
+def _jpeg_sof_dimensions(stream: IO[bytes], start: int, offset: int, segments_left: int) -> tuple[int, int] | None:
+    segment_offset = offset  # rebind-ok: advanced one JPEG segment per hop
+    for _ in range(segments_left):
+        if segment_offset > _MAX_JPEG_HEADER_OFFSET:
+            return None
+        stream.seek(start + segment_offset)
+        marker = stream.read(4)
+        if len(marker) < 4 or marker[0] != 0xFF:
+            return None
+        if marker[1] == 0xFF:
+            run = 1  # rebind-ok: consecutive fill bytes in this window
+            while run < 3 and marker[run + 1] == 0xFF:
+                run += 1
+            segment_offset += run
+            continue
+        if marker[1] in _JPEG_SOF_MARKERS:
+            sof = stream.read(_JPEG_SOF_PAYLOAD_SIZE)
+            if len(sof) < _JPEG_SOF_PAYLOAD_SIZE:
+                return None
+            _precision, height, width = struct.unpack(">BHH", sof)
+            return width, height
+        segment_length = int.from_bytes(marker[2:4], "big")
+        if segment_length < 2:
+            return None
+        segment_offset += 2 + segment_length
+    return None
+
+
+def _header_dimensions(stream: IO[bytes], position: int) -> tuple[int, int] | None:
+    stream.seek(position)
+    head: Final = stream.read(_HEADER_READ_SIZE)
+    match get_image_type(head):  # pyright: ignore[reportMatchNotExhaustive]  # unmatched types fall through to the None below
+        case "png":
+            return _png_dimensions(head)
+        case "webp":
+            return _webp_dimensions(head)
+        case "jpeg":
+            return _jpeg_sof_dimensions(stream, position, _JPEG_FIRST_SEGMENT_OFFSET, _MAX_JPEG_SEGMENTS)
+        case "gif":
+            return _gif_dimensions(head)
+    if head[:2] == b"BM":
+        return _bmp_dimensions(head)
+    return None
+
+
+def _measure_stream(stream: IO[bytes]) -> tuple[int, int] | None:
+    position: Final = stream.tell()
+    try:
+        dimensions: Final = _header_dimensions(stream, position)
+    finally:
+        stream.seek(position)
+    if dimensions is None or dimensions[0] <= 0 or dimensions[1] <= 0:
+        return None
+    return dimensions
+
+
+def measure_reference_image(image: FileTypes) -> tuple[int, int] | None:
+    try:
+        stream: Final = _content_stream(image)
+        if stream is None:
+            return None
+        return _measure_stream(stream)
+    except (OSError, ValueError, AttributeError, struct.error):
+        return None
+
+
+def measure_reference_pixels(images: Sequence[FileTypes]) -> int | None:
+    measured: Final = tuple(measure_reference_image(image) for image in images)
+    dimensions: Final = tuple(size for size in measured if size is not None)
+    if len(dimensions) != len(measured):
+        verbose_logger.debug(
+            "Reference image %d could not be measured (path, non-seekable stream, or no readable PNG, JPEG, WebP, "
+            "GIF or BMP header); billing generated pixels only",
+            measured.index(None),
+        )
+        return None
+    return sum(width * height for width, height in dimensions)
 
 
 class ImageEditRequestUtils:

litellm/llms/azure_ai/image_generation/cost_calculator.py

@@ -10,15 +10,21 @@
 from litellm.types.utils import ImageResponse, ModelInfo
 
 
-def _input_cost_per_pixel(resolved: ModelInfo) -> float:
-    deployment_price: Final = _get_cost_per_unit(resolved, "input_cost_per_pixel", default_value=None)
+def _pixel_rate(resolved: ModelInfo, cost_key: str) -> float:
+    deployment_price: Final = _get_cost_per_unit(resolved, cost_key, default_value=None)
     if deployment_price is not None:
         return deployment_price
     model_cost_key: Final = resolved.get("key")
     shared_entry: Final = litellm.model_cost.get(model_cost_key) if model_cost_key is not None else None
     if shared_entry is None:
         return 0.0
-    return shared_entry.get("input_cost_per_pixel") or 0.0
+    return shared_entry.get(cost_key) or 0.0
+
+
+def _reference_cost(resolved: ModelInfo, image_response: ImageResponse) -> float:
+    pixels: Final = image_response._hidden_params.get("reference_pixels")
+    rate: Final = _pixel_rate(resolved, "input_cost_per_reference_pixel")
+    return rate * pixels if isinstance(pixels, int) else 0.0
 
 
 def cost_calculator(
@@ -48,28 +54,31 @@ def cost_calculator(
         if token_based_cost is not None:
             return token_based_cost
 
+        from litellm.cost_calculator import default_image_cost_calculator
+
         num_images: Final = n if n is not None else len(image_response.data or ())
         output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0
-        if output_cost_per_image:
-            return output_cost_per_image * num_images
-
-        if _input_cost_per_pixel(_model_info):
-            from litellm.cost_calculator import default_image_cost_calculator
-
-            width: Final = optional_params.get("width") if optional_params else None
-            height: Final = optional_params.get("height") if optional_params else None
-            pixel_size: Final = (
-                f"{width}x{height}"
-                if type(width) is int and type(height) is int and width > 0 and height > 0
-                else size or image_response.size
-            )
-            return default_image_cost_calculator(
+        input_cost_per_pixel: Final[float] = _pixel_rate(_model_info, "input_cost_per_pixel")
+        width: Final = optional_params.get("width") if optional_params else None
+        height: Final = optional_params.get("height") if optional_params else None
+        pixel_size: Final = (
+            f"{width}x{height}"
+            if type(width) is int and type(height) is int and width > 0 and height > 0
+            else size or image_response.size
+        )
+        generated_cost: Final[float] = (
+            output_cost_per_image * num_images
+            if output_cost_per_image
+            else default_image_cost_calculator(
                 model=_model_info.get("key", model),
                 custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
                 size=pixel_size,
                 n=num_images,
                 model_info=model_info,
             )
-        return 0.0
+            if input_cost_per_pixel
+            else 0.0
+        )
+        return generated_cost + _reference_cost(_model_info, image_response)
 
     raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}")

litellm/llms/custom_httpx/llm_http_handler.py

@@ -33,6 +33,7 @@
 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.utils import measure_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._hidden_params["reference_pixels"] = reference_pixels  # rebind-ok: cost calculators read hidden params
+    return response
+
+
 class BaseLLMHTTPHandler:
     async def _make_common_async_call(
         self,
@@ -6881,6 +6888,7 @@ def image_edit_handler(
             litellm_params=dict(litellm_params),
         )
 
+        reference_pixels: Final = measure_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 @@ def image_edit_handler(
                 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 @@ async def async_image_edit_handler(
             litellm_params=dict(litellm_params),
         )
 
+        reference_pixels: Final = measure_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 @@ async def async_image_edit_handler(
                 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/model_prices_and_context_window_backup.json

@@ -10540,7 +10540,8 @@
         ]
     },
     "azure_ai/FLUX.2-flex": {
-        "input_cost_per_pixel": 5e-08,
+        "input_cost_per_pixel": 4.76837158203125e-08,
+        "input_cost_per_reference_pixel": 4.76837158203125e-08,
         "litellm_provider": "azure_ai",
         "max_input_tokens": 32000,
         "max_tokens": 32000,

litellm/types/utils.py

@@ -3702,6 +3702,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

model_prices_and_context_window.json

@@ -10540,7 +10540,8 @@
         ]
     },
     "azure_ai/FLUX.2-flex": {
-        "input_cost_per_pixel": 5e-08,
+        "input_cost_per_pixel": 4.76837158203125e-08,
+        "input_cost_per_reference_pixel": 4.76837158203125e-08,
         "litellm_provider": "azure_ai",
         "max_input_tokens": 32000,
         "max_tokens": 32000,

model_prices_and_context_window.schema.json

@@ -318,6 +318,10 @@
           "type": "number",
           "minimum": 0
         },
+        "input_cost_per_reference_pixel": {
+          "type": "number",
+          "minimum": 0
+        },
         "input_cost_per_request": {
           "type": "number",
           "minimum": 0

terraform/provider/docs/resources/model.md

@@ -157,6 +157,8 @@ The following arguments are supported:
 
 * `input_cost_per_pixel` - (Optional) float. Cost applied per input pixel for models that charge by image size.
 
+* `input_cost_per_reference_pixel` - (Optional) float. Cost applied per reference image pixel for image edit models that meter reference inputs.
+
 * `output_cost_per_pixel` - (Optional) float. Cost applied per output pixel for image-generation models.
 
 * `input_cost_per_second` - (Optional) float. Cost applied per input second for audio/transcription models.

terraform/provider/litellm/resource_model.go

@@ -119,6 +119,10 @@ func resourceLiteLLMModel() *schema.Resource {
 				Type:     schema.TypeFloat,
 				Optional: true,
 			},
+			"input_cost_per_reference_pixel": {
+				Type:     schema.TypeFloat,
+				Optional: true,
+			},
 			"output_cost_per_pixel": {
 				Type:     schema.TypeFloat,
 				Optional: true,

terraform/provider/litellm/resource_model_crud.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/hashicorp/go-cty/cty"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 )
 
@@ -124,6 +125,9 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e
 	if inputCostPerPixel := d.Get("input_cost_per_pixel").(float64); inputCostPerPixel > 0 {
 		litellmParams["input_cost_per_pixel"] = inputCostPerPixel
 	}
+	if raw, err := d.GetRawConfigAt(cty.GetAttrPath("input_cost_per_reference_pixel")); err == nil && !raw.IsNull() && raw.Type() == cty.Number {
+		litellmParams["input_cost_per_reference_pixel"] = d.Get("input_cost_per_reference_pixel").(float64)
+	}
 	if outputCostPerPixel := d.Get("output_cost_per_pixel").(float64); outputCostPerPixel > 0 {
 		litellmParams["output_cost_per_pixel"] = outputCostPerPixel
 	}

terraform/provider/litellm/types.go

@@ -93,6 +93,7 @@ type LiteLLMParams struct {
 	InputCostPerToken              float64                `json:"input_cost_per_token,omitempty"`
 	OutputCostPerToken             float64                `json:"output_cost_per_token,omitempty"`
 	InputCostPerPixel              float64                `json:"input_cost_per_pixel,omitempty"`
+	InputCostPerReferencePixel     float64                `json:"input_cost_per_reference_pixel,omitempty"`
 	OutputCostPerPixel             float64                `json:"output_cost_per_pixel,omitempty"`
 	InputCostPerSecond             float64                `json:"input_cost_per_second,omitempty"`
 	OutputCostPerSecond            float64                `json:"output_cost_per_second,omitempty"`