From 3608ba392e8627d8590eaac74936a6b5fa54ef79 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Fri, 25 Sep 2026 02:15:37 +0000 Subject: [PATCH] fix(azure_ai): ignore non-finite FLUX.2 request_meta megapixels --- .../llms/azure_ai/image_generation/flux_transformation.py | 6 +++++- .../test_azure_ai_flux2_image_generation.py | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index 79df59eae5b..93362a57f92 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -1,3 +1,4 @@ +import math from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -42,7 +43,10 @@ def with_flux2_billed_megapixels(image_response: ImageResponse, raw_response: ht request_meta: Final = _Flux2ResponseBody.model_validate_json(raw_response.content).request_meta except ValidationError: return image_response - image_response.set_provider_billed_megapixels(request_meta.input_mp + request_meta.output_mp) + billed_megapixels: Final = request_meta.input_mp + request_meta.output_mp + if not math.isfinite(billed_megapixels): + return image_response + image_response.set_provider_billed_megapixels(billed_megapixels) return image_response 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 5a527ec5983..8b698bea638 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 @@ -204,6 +204,7 @@ def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensi ({"input_mp": 45.78, "output_mp": 1.0}, 46.78), ({"output_mp": 0.39}, 2048 * 1024 / (1024 * 1024)), ({"input_mp": -1.0, "output_mp": 0.39}, 2048 * 1024 / (1024 * 1024)), + ({"input_mp": 1.7e308, "output_mp": 1.7e308}, 2048 * 1024 / (1024 * 1024)), ), ) def test_flux2_generation_bills_azure_reported_megapixels_and_falls_back_to_requested_size(