fix(azure_ai): bill FLUX.2-flex edit reference images as one megapixel each

litellm/llms/azure_ai/image_generation/cost_calculator.py

@@ -9,6 +9,8 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
 )
 from litellm.types.utils import ImageResponse, ModelInfo
 
+_BILLED_PIXELS_PER_INPUT_IMAGE: 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)
@@ -63,13 +65,15 @@ def cost_calculator(
                 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(
+            output_cost: Final = 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,
             )
+            input_image_pixels: Final = image_response.input_image_count * _BILLED_PIXELS_PER_INPUT_IMAGE
+            return output_cost + _input_cost_per_pixel(_model_info) * input_image_pixels
         return 0.0
 
     raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}")

litellm/llms/custom_httpx/llm_http_handler.py

@@ -6968,6 +6968,7 @@ class BaseLLMHTTPHandler:
             logging_obj=logging_obj,
         )
         set_provider_response_headers_in_hidden_params(image_edit_response, response.headers)
+        image_edit_response.set_input_image_count(len(image) if isinstance(image, list) else 1)
         return image_edit_response
 
     async def async_image_edit_handler(
@@ -7069,6 +7070,7 @@ class BaseLLMHTTPHandler:
             logging_obj=logging_obj,
         )
         set_provider_response_headers_in_hidden_params(image_edit_response, response.headers)
+        image_edit_response.set_input_image_count(len(image) if isinstance(image, list) else 1)
         return image_edit_response
 
     def image_generation_handler(

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 @@ from openai.types.images_response import ImagesResponse as OpenAIImageResponse
 
 class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject):
     _hidden_params: dict = {}
+    _input_image_count: int = 0
+
+    @property
+    def input_image_count(self) -> int:
+        return self._input_image_count
+
+    def set_input_image_count(self, count: int) -> None:
+        self._input_image_count = count
 
     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/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,7 @@ 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)
 
 
 def test_flux2_flex_cost_accepts_lowercase_model_spelling():
@@ -202,7 +207,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:

tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py

@@ -4249,6 +4249,28 @@ async def test_async_image_edit_handler_records_upstream_response_headers():
     _assert_upstream_headers_recorded(response)
 
 
+def test_image_edit_handler_records_input_image_count():
+    client = HTTPHandler(client=httpx.Client(transport=_json_with_upstream_headers({"transformed_by": "sync"})))
+
+    response = BaseLLMHTTPHandler().image_edit_handler(
+        client=client, **{**_image_edit_call_kwargs(), "image": [b"one", b"two", b"three"]}
+    )
+
+    assert response.input_image_count == 3
+
+
+@pytest.mark.asyncio
+async def test_async_image_edit_handler_records_input_image_count():
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=_json_with_upstream_headers({"transformed_by": "async"}))
+
+    response = await BaseLLMHTTPHandler().async_image_edit_handler(
+        client=client, **{**_image_edit_call_kwargs(), "image": [b"one", b"two"]}
+    )
+
+    assert response.input_image_count == 2
+
+
 class _HeaderImageGenerationConfig(BaseImageGenerationConfig):
     def get_supported_openai_params(self, model):
         return []

tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py

@@ -175,7 +175,31 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[
         **dimensions,
     )
 
-    assert response._hidden_params["response_cost"] == pytest.approx(5e-08 * 2048 * 1024 * 2)
+    rate: Final = litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"]
+    assert response._hidden_params["response_cost"] == pytest.approx(rate * (2048 * 1024 * 2 + 1024 * 1024))
+
+
+@pytest.mark.parametrize("reference_images", (1, 3, 10))
+@pytest.mark.usefixtures("local_model_cost_map")
+def test_flux2_image_edit_bills_each_reference_image_as_one_megapixel(reference_images: int):
+    client: Final = HTTPHandler(
+        client=httpx.Client(
+            transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"data": [{"b64_json": "aW1n"}]}))
+        )
+    )
+
+    response: Final = litellm.image_edit(
+        model="azure_ai/FLUX.2-flex",
+        image=[b"image"] * reference_images,
+        prompt="Blend every reference",
+        api_key="test-key",
+        api_base="https://example.services.ai.azure.com",
+        client=client,
+        size="1024x1024",
+    )
+
+    rate: Final = litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"]
+    assert response._hidden_params["response_cost"] == pytest.approx(rate * 1024 * 1024 * (1 + reference_images))
 
 
 def test_flux2_image_edit_accepts_and_drops_openai_only_parameters():