From 723a9fe542cbc4819b6312565e7f70d89f2395ac Mon Sep 17 00:00:00 2001
From: shrey kharbanda <shrey@berri.ai>
Date: Wed, 28 Oct 2026 21:49:11 -0700
Subject: [PATCH] fix(azure_ai): preserve fallback for empty FLUX.2 image
 responses

---
 .../azure_ai/image_generation/cost_calculator.py   |  4 +++-
 .../test_azure_ai_flux2_image_generation.py        | 14 ++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py
index 92d9cb7a8ce..b3d844f17cd 100644
--- a/litellm/llms/azure_ai/image_generation/cost_calculator.py
+++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py
@@ -50,13 +50,15 @@ def record_megapixel_usage(response: object, response_body: bytes | Mapping[str,
 
 
 def _megapixel_cost(model_info: ModelInfo, image_response: ImageResponse) -> float | None:
+    images: Final = len(image_response.data or ())
+    if not images:
+        return None
     try:
         prices: Final = _MegapixelPrices.model_validate(model_info)
         usage: Final = _MegapixelUsage.model_validate(image_response.hidden_params.get(_MEGAPIXEL_USAGE_HIDDEN_PARAM))
     except ValidationError:
         return None
     first_rate: Final = prices.output_cost_per_first_megapixel
-    images: Final = max(len(image_response.data or ()), 1)
     output_mp_per_image: Final = usage.output_mp / images
     return images * (
         (prices.output_cost_per_megapixel if first_rate is None else first_rate)
diff --git a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py
index 5a010f81556..4b85081f098 100644
--- a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py
+++ b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py
@@ -334,6 +334,20 @@ def test_flux2_pro_generation_bills_first_megapixel_per_returned_image(
     )
 
 
+@respx.mock
+def test_flux2_pro_generation_without_returned_images_has_no_cost() -> None:
+    respx.post(url__regex=_flux2_route("flux-2-pro")).mock(
+        return_value=httpx.Response(200, json=_metered_body(3.0, 1.5, num_images=0))
+    )
+
+    response: Final = litellm.image_generation(
+        model="azure_ai/flux.2-pro", prompt="A red fox", n=4, api_base=FOUNDRY_BASE, api_key="test-key"
+    )
+
+    assert response.data == []
+    assert response._hidden_params["response_cost"] == 0.0
+
+
 @pytest.mark.asyncio
 async def test_async_flux2_pro_generation_bills_the_reported_megapixels(
     respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch