From cf08440557bd0a377f1ce082497d037294c6aca0 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 18 Sep 2026 11:42:52 -0700
Subject: [PATCH] fix(azure-ai): price FLUX.2 Flex by its resolved cost map key

---
 .../azure_ai/image_generation/cost_calculator.py    |  5 +----
 .../test_azure_ai_flux2_image_generation.py         | 13 +++++++++++++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py
index 086293f26c0..35d0f4fb6c3 100644
--- a/litellm/llms/azure_ai/image_generation/cost_calculator.py
+++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py
@@ -42,9 +42,6 @@ def cost_calculator(
         if input_cost_per_pixel:
             from litellm.cost_calculator import default_image_cost_calculator
 
-            cost_model: Final = (
-                model if model.startswith(f"{litellm.LlmProviders.AZURE_AI.value}/") else f"azure_ai/{model}"
-            )
             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 = (
@@ -53,7 +50,7 @@ def cost_calculator(
                 else size or image_response.size
             )
             return default_image_cost_calculator(
-                model=cost_model,
+                model=_model_info["key"],
                 custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
                 size=pixel_size,
                 n=num_images,
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 c1d46b7e919..388cca9f054 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
@@ -172,6 +172,19 @@ def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensi
     ) == pytest.approx(5e-08 * 2048 * 1024 * 2)
 
 
+def test_flux2_flex_cost_accepts_lowercase_model_spelling():
+    response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")])
+
+    cost: Final = litellm.completion_cost(
+        model="azure_ai/flux.2-flex",
+        completion_response=response,
+        optional_params={"width": 1536, "height": 1024, "num_images": 2},
+        call_type="image_generation",
+    )
+
+    assert cost == pytest.approx(5e-08 * 1536 * 1024 * 2)
+
+
 def test_flux2_response_preserves_mapped_dimensions():
     config = AzureFoundryFluxImageGenerationConfig()
     params = config.map_openai_params(