fix(azure_ai): bill FLUX.2 by the megapixels Azure reports in request_meta

litellm/images/main.py

@@ -51,6 +51,7 @@
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import (
     LITELLM_IMAGE_VARIATION_PROVIDERS,
+    CustomPricingLiteLLMParams,
     LlmProviders,
     all_litellm_params,
 )
@@ -883,6 +884,10 @@ def image_edit(
             optional_params=dict(image_edit_request_params),
             litellm_params={
                 **image_edit_request_params,
+                **litellm_params.model_dump(
+                    include=frozenset(CustomPricingLiteLLMParams.model_fields),
+                    exclude_none=True,
+                ),
                 "litellm_call_id": litellm_call_id,
                 "model_info": model_info,
             },

litellm/llms/azure/azure.py

@@ -19,6 +19,7 @@
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing
 from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin
+from litellm.llms.azure_ai.image_generation import AzureFoundryFluxImageGenerationConfig
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
@@ -1323,7 +1324,9 @@ def image_generation(
                 deployment_name=model,
             )
             provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2"))
-            if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig):
+            if isinstance(
+                provider_config, (AzureFoundryMAIImageGenerationConfig, AzureFoundryFluxImageGenerationConfig)
+            ):
                 return provider_config.transform_image_generation_response(
                     model=data.get("model", "dall-e-2"),
                     raw_response=httpx_response,

litellm/llms/azure/image_generation/__init__.py

@@ -1,5 +1,8 @@
 from litellm._logging import verbose_logger
-from litellm.llms.azure_ai.image_generation import AzureFoundryMAIImageGenerationConfig
+from litellm.llms.azure_ai.image_generation import (
+    AzureFoundryFluxImageGenerationConfig,
+    AzureFoundryMAIImageGenerationConfig,
+)
 from litellm.llms.base_llm.image_generation.transformation import (
     BaseImageGenerationConfig,
 )
@@ -27,6 +30,8 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig:
         return AzureDallE3ImageGenerationConfig()
     elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model):
         return AzureFoundryMAIImageGenerationConfig()
+    elif AzureFoundryFluxImageGenerationConfig.is_flux2_model(model):
+        return AzureFoundryFluxImageGenerationConfig()
     else:
         verbose_logger.debug(
             "Using AzureGPTImageGenerationConfig for model: %s. This follows the gpt-image model format.", model

litellm/llms/azure_ai/image_edit/flux2_transformation.py

@@ -2,8 +2,9 @@
 from collections.abc import Mapping, Sequence
 from io import BufferedReader
 from types import MappingProxyType
-from typing import Any, Final
+from typing import TYPE_CHECKING, Any, Final
 
+import httpx
 from httpx._types import RequestFiles
 
 import litellm
@@ -13,12 +14,17 @@
 )
 from litellm.llms.azure_ai.image_generation.flux_transformation import (
     AzureFoundryFluxImageGenerationConfig,
+    with_flux2_billed_megapixels,
 )
 from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.images.main import ImageEditOptionalRequestParams
 from litellm.types.llms.openai import FileTypes
 from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import ImageResponse
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 
 
 class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
@@ -121,6 +127,17 @@ def transform_image_edit_request(
         }
         return request_body, []
 
+    def transform_image_edit_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        logging_obj: "LiteLLMLoggingObj",
+    ) -> ImageResponse:
+        return with_flux2_billed_megapixels(
+            super().transform_image_edit_response(model=model, raw_response=raw_response, logging_obj=logging_obj),
+            raw_response,
+        )
+
     def _convert_image_to_base64(self, image: Any) -> str:
         """Convert image file to base64 string"""
         if isinstance(image, BufferedReader):

litellm/llms/azure_ai/image_generation/cost_calculator.py

@@ -9,6 +9,8 @@
 )
 from litellm.types.utils import ImageResponse, ModelInfo
 
+_PIXELS_PER_MEGAPIXEL: Final = 1024 * 1024
+
 
 def _input_cost_per_pixel(resolved: ModelInfo) -> float:
     deployment_price: Final = _get_cost_per_unit(resolved, "input_cost_per_pixel", default_value=None)
@@ -53,7 +55,11 @@ def cost_calculator(
         if output_cost_per_image:
             return output_cost_per_image * num_images
 
-        if _input_cost_per_pixel(_model_info):
+        input_cost_per_pixel: Final = _input_cost_per_pixel(_model_info)
+        if input_cost_per_pixel and image_response.provider_billed_megapixels is not None:
+            return input_cost_per_pixel * _PIXELS_PER_MEGAPIXEL * image_response.provider_billed_megapixels
+
+        if input_cost_per_pixel:
             from litellm.cost_calculator import default_image_cost_calculator
 
             width: Final = optional_params.get("width") if optional_params else None

litellm/llms/azure_ai/image_generation/flux_transformation.py

@@ -1,10 +1,20 @@
+import math
 from collections.abc import Mapping
 from types import MappingProxyType
-from typing import Final
+from typing import TYPE_CHECKING, Final
+
+import httpx
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 from litellm.exceptions import BadRequestError, UnsupportedParamsError
 from litellm.llms.openai.image_generation import GPTImageGenerationConfig
 from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
+from litellm.types.utils import ImageResponse
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
+
+    from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
 
 FLUX2_DROPPED_OPENAI_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = (
     "background",
@@ -15,6 +25,31 @@
 )
 
 
+class _Flux2RequestMeta(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    input_mp: float = Field(ge=0)
+    output_mp: float = Field(ge=0)
+
+
+class _Flux2ResponseBody(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    request_meta: _Flux2RequestMeta
+
+
+def with_flux2_billed_megapixels(image_response: ImageResponse, raw_response: httpx.Response) -> ImageResponse:
+    try:
+        request_meta: Final = _Flux2ResponseBody.model_validate_json(raw_response.content).request_meta
+    except ValidationError:
+        return image_response
+    billed_megapixels: Final = request_meta.input_mp + request_meta.output_mp
+    if not math.isfinite(billed_megapixels):
+        return image_response
+    image_response.set_provider_billed_megapixels(billed_megapixels)
+    return image_response
+
+
 class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
     """Azure Foundry BFL API configuration for FLUX image generation."""
 
@@ -152,3 +187,32 @@ def map_openai_params(
             }
         )
         return {**optional_params, **mapped_params}  # mutable-ok: inherited config contract returns a dict
+
+    def transform_image_generation_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: ImageResponse,
+        logging_obj: "LiteLLMLoggingObj",
+        request_data: dict,
+        optional_params: dict,
+        litellm_params: dict,
+        encoding: "Tokenizer | None",
+        api_key: str | None = None,
+        json_mode: bool | None = None,
+    ) -> ImageResponse:
+        image_response: Final = super().transform_image_generation_response(
+            model=model,
+            raw_response=raw_response,
+            model_response=model_response,
+            logging_obj=logging_obj,
+            request_data=request_data,
+            optional_params=optional_params,
+            litellm_params=litellm_params,
+            encoding=encoding,
+            api_key=api_key,
+            json_mode=json_mode,
+        )
+        return (
+            with_flux2_billed_megapixels(image_response, raw_response) if self.is_flux2_model(model) else image_response
+        )

litellm/model_prices_and_context_window_backup.json

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

litellm/types/utils.py

@@ -2578,6 +2578,14 @@ class ImageUsage(BaseLiteLLMOpenAIResponseObject):
 
 class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject):
     _hidden_params: dict = {}
+    _provider_billed_megapixels: float | None = None
+
+    @property
+    def provider_billed_megapixels(self) -> float | None:
+        return self._provider_billed_megapixels
+
+    def set_provider_billed_megapixels(self, megapixels: float) -> None:
+        self._provider_billed_megapixels = megapixels
 
     usage: ImageUsage | None = None
     """

model_prices_and_context_window.json

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

tests/integration/pricing/test_image_pixel_pricing.py

@@ -0,0 +1,49 @@
+import json
+from typing import Final
+
+import pytest
+
+from tests.integration._support.client import Gateway, eventually
+from tests.integration._support.database import read_rows
+from tests.integration._support.wire import Reply, Request, wire_server
+
+_PIXELS_PER_MEGAPIXEL: Final = 1024 * 1024
+
+
+@pytest.mark.parametrize("input_cost_per_pixel", (1.5e-07, 0.0))
+def test_flux2_edit_bills_azure_reported_megapixels_at_the_deployment_rate(
+    gateway: Gateway, input_cost_per_pixel: float
+) -> None:
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        return Reply(
+            body=json.dumps(
+                {"data": [{"b64_json": "aW1n"}], "request_meta": {"input_mp": 0.39, "output_mp": 1.0}}
+            ).encode()
+        )
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model="azure_ai/FLUX.2-flex",
+            api_base=wire.url,
+            api_key="synthetic-azure-key",
+            input_cost_per_pixel=input_cost_per_pixel,
+        )
+        response: Final = gateway.request_multipart(
+            "/v1/images/edits",
+            {"model": model, "prompt": "add a hat", "size": "1024x1024"},
+            {"image": ("reference.png", b"reference", "image/png")},
+        )
+        assert response.status_code == 200, response.text
+        expected: Final = input_cost_per_pixel * _PIXELS_PER_MEGAPIXEL * (0.39 + 1.0)
+        if expected:
+            assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected), response.text
+        else:
+            assert response.headers.get("x-litellm-response-cost") in (None, "0", "0.0"), response.text
+        call_id: Final = response.headers["x-litellm-call-id"]
+        rows: Final = eventually(
+            lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)),
+            lambda values: len(values) == 1,
+            seconds=70,
+        )
+        assert float(rows[0]["spend"]) == pytest.approx(expected), response.text

tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py

@@ -28,6 +28,10 @@ def use_local_model_cost_map(monkeypatch):
     _invalidate_model_cost_lowercase_map()
 
 
+def _flex_rate() -> float:
+    return litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"]
+
+
 @pytest.mark.parametrize(
     ("model", "provider_path"),
     [
@@ -127,7 +131,8 @@ def test_flux2_flex_model_info():
     assert model_info["max_input_tokens"] == 32000
     assert model_info["max_tokens"] == 32000
     assert model_info["supported_endpoints"] == ["/v1/images/generations", "/v1/images/edits"]
-    assert catalog_info["input_cost_per_pixel"] == 5e-08
+    # Azure Retail Prices API "Flex Megapixel" is $0.05, and Cost Management metered a 1024x1024 image as 1.0 MP (2026-09-22)
+    assert catalog_info["input_cost_per_pixel"] * 1024 * 1024 == pytest.approx(0.05)
     assert catalog_info["supported_modalities"] == ["text", "image"]
     assert catalog_info["supported_output_modalities"] == ["image"]
 
@@ -148,7 +153,7 @@ def test_flux2_flex_cost_uses_generated_megapixels():
         call_type="image_generation",
     )
 
-    assert cost == pytest.approx(5e-08 * 2048 * 1024 * 2)
+    assert cost == pytest.approx(_flex_rate() * 2048 * 1024 * 2)
 
 
 @pytest.mark.parametrize("model", ("FLUX-1.1-pro", "FLUX.1-Kontext-pro"))
@@ -189,7 +194,40 @@ def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensi
         completion_response=response,
         optional_params=params,
         call_type="image_generation",
-    ) == pytest.approx(5e-08 * 2048 * 1024 * 2)
+    ) == pytest.approx(_flex_rate() * 2048 * 1024 * 2)
+
+
+@pytest.mark.parametrize(
+    ("request_meta", "expected_megapixels"),
+    (
+        ({"input_mp": 0.0, "output_mp": 0.39}, 0.39),
+        ({"input_mp": 45.78, "output_mp": 1.0}, 46.78),
+        ({"output_mp": 0.39}, 2048 * 1024 / (1024 * 1024)),
+        ({"input_mp": -1.0, "output_mp": 0.39}, 2048 * 1024 / (1024 * 1024)),
+        ({"input_mp": 1.7e308, "output_mp": 1.7e308}, 2048 * 1024 / (1024 * 1024)),
+    ),
+)
+def test_flux2_generation_bills_azure_reported_megapixels_and_falls_back_to_requested_size(
+    request_meta: Mapping[str, float], expected_megapixels: float
+):
+    params: Final = {"width": 2048, "height": 1024}
+    response: Final = get_azure_image_generation_config("FLUX.2-flex").transform_image_generation_response(
+        model="FLUX.2-flex",
+        raw_response=httpx.Response(200, json={"data": [{"b64_json": "aW1n"}], "request_meta": request_meta}),
+        model_response=ImageResponse(),
+        logging_obj=MagicMock(),
+        request_data={"prompt": "A red fox", **params},
+        optional_params=params,
+        litellm_params={},
+        encoding=None,
+    )
+
+    assert litellm.completion_cost(
+        model="azure_ai/FLUX.2-flex",
+        completion_response=response,
+        optional_params=params,
+        call_type="image_generation",
+    ) == pytest.approx(_flex_rate() * 1024 * 1024 * expected_megapixels)
 
 
 def test_flux2_flex_cost_accepts_lowercase_model_spelling():
@@ -202,7 +240,7 @@ def test_flux2_flex_cost_accepts_lowercase_model_spelling():
         call_type="image_generation",
     )
 
-    assert cost == pytest.approx(5e-08 * 1536 * 1024 * 2)
+    assert cost == pytest.approx(_flex_rate() * 1536 * 1024 * 2)
 
 
 def test_flux2_flex_cost_prefers_deployment_input_cost_per_pixel() -> None: