From 61040b8983d2d15c1ee716024ce5feafa3308609 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Fri, 25 Sep 2026 00:25:59 +0000 Subject: [PATCH] fix(azure_ai): bill FLUX.2-flex edit reference images as one megapixel each --- .../image_generation/cost_calculator.py | 6 ++++- litellm/llms/custom_httpx/llm_http_handler.py | 2 ++ ...odel_prices_and_context_window_backup.json | 2 +- litellm/types/utils.py | 8 ++++++ model_prices_and_context_window.json | 2 +- .../test_azure_ai_flux2_image_generation.py | 13 +++++++--- .../custom_httpx/test_llm_http_handler.py | 22 ++++++++++++++++ ...test_azure_ai_image_edit_transformation.py | 26 ++++++++++++++++++- 8 files changed, 73 insertions(+), 8 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 08b732197f3..d0888ab89ff 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/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)}") diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4f31742aaaa..04f61366a9a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3e2aa011164..f490beb6b1d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index caf88e5d517..aba602708a3 100644 --- a/litellm/types/utils.py +++ b/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 """ diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3e2aa011164..f490beb6b1d 100644 --- a/model_prices_and_context_window.json +++ b/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, diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index fdd21c87732..634f65fcacc 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/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: diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 75b6ce4c626..41fbf84033e 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/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 [] diff --git a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 51ba2c34cd7..a4b702b7877 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/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():