From ee0fd200953ea14f402424c27486b3763aec4d49 Mon Sep 17 00:00:00 2001
From: shrey kharbanda <shreshth@berri.ai>
Date: Thu, 24 Sep 2026 23:38:53 +0000
Subject: [PATCH] fix(azure_ai): bill FLUX.2 Pro edit references on the Pro ref
 meter

---
 ...odel_prices_and_context_window_backup.json |  1 +
 model_prices_and_context_window.json          |  1 +
 .../image_edit/test_flux2_transformation.py   | 20 +++++++++++++++++++
 .../test_azure_ai_flux2_image_generation.py   | 18 +++++++++++++++++
 4 files changed, 40 insertions(+)

diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 990033da4ed..71643813044 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -10954,6 +10954,7 @@
         ]
     },
     "azure_ai/flux.2-pro": {
+        "input_cost_per_pixel": 1.430511474609375e-08,
         "litellm_provider": "azure_ai",
         "mode": "image_generation",
         "output_cost_per_image": 0.04,
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 990033da4ed..71643813044 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -10954,6 +10954,7 @@
         ]
     },
     "azure_ai/flux.2-pro": {
+        "input_cost_per_pixel": 1.430511474609375e-08,
         "litellm_provider": "azure_ai",
         "mode": "image_generation",
         "output_cost_per_image": 0.04,
diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_flux2_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_flux2_transformation.py
index e9e807cfc97..eab1111d8b7 100644
--- a/tests/test_litellm/llms/azure_ai/image_edit/test_flux2_transformation.py
+++ b/tests/test_litellm/llms/azure_ai/image_edit/test_flux2_transformation.py
@@ -1,5 +1,6 @@
 import base64
 import io
+import tempfile
 from typing import Final
 
 from litellm.llms.azure_ai.image_edit.flux2_transformation import AzureFoundryFlux2ImageEditConfig
@@ -37,3 +38,22 @@ def test_convert_image_to_base64_rewinds_seekable_streams_before_encoding() -> N
     encoded: Final = AzureFoundryFlux2ImageEditConfig()._convert_image_to_base64(stream)
 
     assert base64.b64decode(encoded) == b"png-bytes"
+
+
+def test_convert_image_to_base64_sends_the_full_image_on_every_retry() -> None:
+    config: Final = AzureFoundryFlux2ImageEditConfig()
+    stream: Final = io.BytesIO(b"png-bytes")
+
+    first: Final = config._convert_image_to_base64(stream)
+    retry: Final = config._convert_image_to_base64(stream)
+
+    assert base64.b64decode(retry) == base64.b64decode(first) == b"png-bytes"
+
+
+def test_convert_image_to_base64_accepts_named_temporary_files() -> None:
+    with tempfile.NamedTemporaryFile() as upload:
+        upload.write(b"png-bytes")
+
+        encoded: Final = AzureFoundryFlux2ImageEditConfig()._convert_image_to_base64(upload)
+
+    assert base64.b64decode(encoded) == b"png-bytes"
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 f05bb87f91d..247bc117d4e 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
@@ -520,3 +520,21 @@ def test_flux2_cost_optional_params_dimensions_beat_response_size() -> None:
 def test_flux2_cost_rejects_non_image_response() -> None:
     with pytest.raises(TypeError, match="must be of type ImageResponse"):
         cost_calculator(model="FLUX.2-flex", image_response=object())
+
+
+def test_flux2_pro_edit_bills_references_on_the_pro_ref_meter() -> None:
+    pro_pixel_rate: Final = litellm.model_cost["azure_ai/flux.2-pro"]["input_cost_per_pixel"]
+    assert pro_pixel_rate * _ONE_MEGAPIXEL == pytest.approx(0.015), (
+        "Azure Retail Prices API, product 'Azure BFL Flux Models', meter 'Flux 2 Ref MP Megapixel' is $0.015 "
+        "per MP, checked 2026-09-24"
+    )
+
+    cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator(
+        model="flux.2-pro",
+        completion_response=_edit_response(),
+        custom_llm_provider="azure_ai",
+        size="1024x1024",
+        call_type="image_edit",
+    )
+
+    assert cost == pytest.approx(litellm.model_cost["azure_ai/flux.2-pro"]["output_cost_per_image"] + 0.015 * 2)