From e0e0629c4ca269f735bef904831aa757b60da84c Mon Sep 17 00:00:00 2001 From: Shreshth Kharbanda Date: Fri, 25 Sep 2026 13:01:36 -0700 Subject: [PATCH] fix(azure_ai): bill FLUX.2 images in whole megapixels the way Azure meters them Azure's request_meta for FLUX.2-pro edits showed that it rounds every generated image up to whole megapixels, bills a lone reference at its rounded-up megapixels capped at 4, and bills each reference of a multi-reference edit as exactly 1 MP. The per-pixel math under-billed small and fractional references and over-billed photos, for example a 4032x3024 reference logged 11.63 MP where Azure charged 4 The FLUX.2 edit now reports one pixel count per reference, and the Azure AI cost calculator applies the megapixel rule to FLUX.2 models only. A reference whose header reports zero pixels is billed as one unmeasured megapixel so it no longer voids the whole edit's reference cost azure_ai/flux.2-pro now bills $0.03 for the first generated megapixel and $0.015 for each extra one, matching Azure's price list, instead of a flat $0.04 per image --- .../image_edit/flux2_transformation.py | 14 +-- .../image_generation/cost_calculator.py | 82 +++++++++++--- ...odel_prices_and_context_window_backup.json | 5 +- model_prices_and_context_window.json | 5 +- .../test_azure_ai_flux2_image_generation.py | 103 +++++++++++++++--- ...test_azure_ai_image_edit_transformation.py | 79 +++++++++++--- 6 files changed, 230 insertions(+), 58 deletions(-) diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 83497f6b624..98edddc9261 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -43,7 +43,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): def __init__(self) -> None: super().__init__() - self.reference_image_pixels: int = 0 + self.reference_image_pixels: tuple[int, ...] = () def get_supported_openai_params(self, model: str) -> list: return AzureFoundryFluxImageGenerationConfig().get_supported_openai_params(model) @@ -123,7 +123,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): raise ValueError(f"{model} supports at most {max_reference_images} reference images.") reference_bytes: Final = tuple(self._read_image_bytes(reference_image) for reference_image in images) - self.reference_image_pixels = sum(_pixel_count(image_bytes) for image_bytes in reference_bytes) + self.reference_image_pixels = tuple(_pixel_count(image_bytes) for image_bytes in reference_bytes) reference_images: Final[Mapping[str, str]] = MappingProxyType( { "input_image" if index == 1 else f"input_image_{index}": base64.b64encode(image_bytes).decode("utf-8") @@ -198,8 +198,8 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): def _pixel_count(image_bytes: bytes) -> int: dimensions: Final = image_dimensions_from_bytes(image_bytes) - if dimensions is None: - verbose_logger.warning("Could not read the dimensions of a FLUX.2 reference image, billing it as one megapixel") - return UNMEASURED_REFERENCE_IMAGE_PIXELS - width, height = dimensions - return width * height + pixels: Final = dimensions[0] * dimensions[1] if dimensions is not None else 0 + if pixels > 0: + return pixels + verbose_logger.warning("Could not read the dimensions of a FLUX.2 reference image, billing it as one megapixel") + return UNMEASURED_REFERENCE_IMAGE_PIXELS diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 106676a6e14..d05fc43f9e8 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -1,5 +1,8 @@ +import math from collections.abc import Mapping -from typing import Any, Final +from typing import Annotated, Any, Final + +from pydantic import Field, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import ( @@ -8,25 +11,58 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( resolve_image_model_info, ) from litellm.llms.azure_ai.image_edit.flux2_transformation import REFERENCE_IMAGE_PIXELS_HIDDEN_PARAM +from litellm.llms.azure_ai.image_generation.flux_transformation import AzureFoundryFluxImageGenerationConfig from litellm.types.utils import ImageResponse, ModelInfo +MEGAPIXEL: Final = 1024 * 1024 +MAX_LONE_REFERENCE_MEGAPIXELS: Final = 4 +_REFERENCE_PIXELS: Final = TypeAdapter(tuple[Annotated[int, Field(strict=True, gt=0)], ...]) -def _pixel_rate(resolved: ModelInfo, cost_key: str) -> float: + +def _price(resolved: ModelInfo, cost_key: str) -> float | None: deployment_price: Final = _get_cost_per_unit(resolved, cost_key, default_value=None) if deployment_price is not None: return deployment_price model_cost_key: Final = resolved.get("key") shared_entry: Final = litellm.model_cost.get(model_cost_key) if model_cost_key is not None else None if shared_entry is None: - return 0.0 - return shared_entry.get(cost_key) or 0.0 + return None + return shared_entry.get(cost_key) + + +def _pixel_rate(resolved: ModelInfo, cost_key: str) -> float: + return _price(resolved, cost_key) or 0.0 + + +def _billable_megapixels(pixels: int) -> int: + return math.ceil(pixels / MEGAPIXEL) + + +def _billable_reference_megapixels(reference_pixels: tuple[int, ...]) -> int: + # Azure's FLUX.2 request_meta (2026-09-25) bills a lone reference at no more than 4 MP, and each reference of a + # multi-reference edit as exactly 1 MP whatever its size + match reference_pixels: + case (): + return 0 + case (lone_reference,): + return min(_billable_megapixels(lone_reference), MAX_LONE_REFERENCE_MEGAPIXELS) + case _: + return len(reference_pixels) def _reference_cost(resolved: ModelInfo, image_response: ImageResponse) -> float: - pixels: Final = image_response._hidden_params.get(REFERENCE_IMAGE_PIXELS_HIDDEN_PARAM) - if type(pixels) is not int or pixels <= 0: + reported_pixels: Final = image_response._hidden_params.get(REFERENCE_IMAGE_PIXELS_HIDDEN_PARAM) + if reported_pixels is None: return 0.0 - return _pixel_rate(resolved, "input_cost_per_reference_pixel") * pixels + try: + reference_pixels: Final = _REFERENCE_PIXELS.validate_python(reported_pixels) + except ValidationError: + return 0.0 + return ( + _pixel_rate(resolved, "input_cost_per_reference_pixel") + * MEGAPIXEL + * _billable_reference_megapixels(reference_pixels) + ) def cost_calculator( @@ -79,6 +115,9 @@ def _generated_cost( model_info: ModelInfo | None, ) -> float: num_images: Final = n if n is not None else len(image_response.data or ()) + pixel_size: Final = _output_size(size, optional_params, image_response) + if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model): + return num_images * _flux2_image_cost(resolved, _size_pixels(pixel_size)) output_cost_per_image: Final[float] = resolved.get("output_cost_per_image") or 0.0 if output_cost_per_image: return output_cost_per_image * num_images @@ -87,13 +126,6 @@ def _generated_cost( from litellm.cost_calculator import default_image_cost_calculator - 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 = ( - f"{width}x{height}" - 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( model=resolved.get("key", model), custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, @@ -101,3 +133,25 @@ def _generated_cost( n=num_images, model_info=model_info, ) + + +def _flux2_image_cost(resolved: ModelInfo, pixels: int) -> float: + megapixel_rate: Final = _pixel_rate(resolved, "input_cost_per_pixel") * MEGAPIXEL + first_megapixel_price: Final = _price(resolved, "output_cost_per_image") + first_megapixel: Final = megapixel_rate if first_megapixel_price is None else first_megapixel_price + return first_megapixel + megapixel_rate * (_billable_megapixels(pixels) - 1) + + +def _output_size( + size: str | None, optional_params: Mapping[str, object] | None, image_response: ImageResponse +) -> str | None: + width: Final = optional_params.get("width") if optional_params else None + height: Final = optional_params.get("height") if optional_params else None + if type(width) is int and type(height) is int and width > 0 and height > 0: + return f"{width}x{height}" + return size or image_response.size + + +def _size_pixels(size: str | None) -> int: + width, height = (int(dimension) for dimension in (size or "1024x1024").replace("-x-", "x").split("x")) + return width * height diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a3fd2bd4c81..bade2687d52 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11371,11 +11371,12 @@ ] }, "azure_ai/flux.2-pro": { + "input_cost_per_pixel": 1.430511474609375e-08, "input_cost_per_reference_pixel": 1.430511474609375e-08, "litellm_provider": "azure_ai", "mode": "image_generation", - "output_cost_per_image": 0.04, - "source": "https://ai.azure.com/explore/models/flux.2-pro/version/1/registry/azureml-blackforestlabs", + "output_cost_per_image": 0.03, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/", "supported_endpoints": [ "/v1/images/generations" ] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a3fd2bd4c81..bade2687d52 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11371,11 +11371,12 @@ ] }, "azure_ai/flux.2-pro": { + "input_cost_per_pixel": 1.430511474609375e-08, "input_cost_per_reference_pixel": 1.430511474609375e-08, "litellm_provider": "azure_ai", "mode": "image_generation", - "output_cost_per_image": 0.04, - "source": "https://ai.azure.com/explore/models/flux.2-pro/version/1/registry/azureml-blackforestlabs", + "output_cost_per_image": 0.03, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/", "supported_endpoints": [ "/v1/images/generations" ] 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 29da0e782dc..c29b538e664 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 @@ -55,6 +55,15 @@ def _edit_response(reference_image_pixels: object) -> ImageResponse: ) +def _pro_megapixel_rates() -> tuple[float, float, float]: + row: Final = litellm.model_cost["azure_ai/flux.2-pro"] + return ( + row["output_cost_per_image"], + row["input_cost_per_pixel"] * 1024 * 1024, + row["input_cost_per_reference_pixel"] * 1024 * 1024, + ) + + @pytest.mark.parametrize( ("model", "provider_path"), [ @@ -225,11 +234,68 @@ def test_flux2_flex_cost_accepts_lowercase_model_spelling(): cost: Final = litellm.completion_cost( model="azure_ai/flux.2-flex", completion_response=response, + optional_params={"width": 2048, "height": 1024, "num_images": 2}, + call_type="image_generation", + ) + + assert cost == pytest.approx(_flex_pixel_rate() * 2048 * 1024 * 2) + + +def test_flux2_flex_generation_rounds_each_image_up_to_whole_megapixels(): + 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(_flex_pixel_rate() * 1536 * 1024 * 2) + assert cost == pytest.approx(_flex_pixel_rate() * 2 * 1024 * 1024 * 2) + + +@pytest.mark.parametrize(("size", "megapixels"), (("256x256", 1), ("1024x1280", 2), ("2048x2048", 4))) +def test_flux2_pro_generation_bills_the_first_megapixel_then_each_additional_one(size: str, megapixels: int): + first, additional, _reference = _pro_megapixel_rates() + + cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( + model="flux.2-pro", + completion_response=ImageResponse(data=[ImageObject(b64_json="aW1n")]), + custom_llm_provider="azure_ai", + size=size, + call_type="image_generation", + ) + + assert cost == pytest.approx(first + additional * (megapixels - 1)) + + +# Billable megapixels as Azure's request_meta reported them for FLUX.2-pro edits on 2026-09-25 +@pytest.mark.parametrize( + ("size", "reference_pixels", "output_megapixels", "reference_megapixels"), + ( + pytest.param("1024x1024", (1024 * 1024,), 1, 1, id="one-whole-megapixel-reference"), + pytest.param("1024x1024", (1024 * 1280,), 1, 2, id="fractional-reference-rounds-up"), + pytest.param("1024x1280", (1024 * 1280,), 2, 2, id="fractional-output-rounds-up"), + pytest.param("1024x1024", (4032 * 3024,), 1, 4, id="lone-reference-caps-at-four-megapixels"), + pytest.param("1024x1024", (1024 * 1280,) * 2, 1, 2, id="each-of-several-references-is-one-megapixel"), + pytest.param("1024x1024", (1024 * 1024, 4032 * 3024), 1, 2, id="large-reference-among-several"), + pytest.param("1024x1024", (1024 * 1280,) * 3, 1, 3, id="three-references"), + ), +) +def test_flux2_pro_edit_bills_the_megapixels_azure_meters( + size: str, reference_pixels: tuple[int, ...], output_megapixels: int, reference_megapixels: int +): + first, additional, reference = _pro_megapixel_rates() + + cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( + model="flux.2-pro", + completion_response=_edit_response(reference_pixels), + custom_llm_provider="azure_ai", + size=size, + call_type="image_edit", + ) + + assert cost == pytest.approx(first + additional * (output_megapixels - 1) + reference * reference_megapixels) def test_flux2_flex_reference_rate_matches_generated_rate(): @@ -240,28 +306,31 @@ def test_flux2_flex_edit_bills_reference_pixels_on_top_of_generated_pixels(): generated_only: Final = _flex_edit_cost(ImageResponse(data=[ImageObject(b64_json="aW1n")], size="1024x1024")) assert generated_only == pytest.approx(_flex_pixel_rate() * 1024 * 1024) - assert _flex_edit_cost(_edit_response(1024 * 1024)) - generated_only == pytest.approx( + assert _flex_edit_cost(_edit_response((1024 * 1024,))) - generated_only == pytest.approx( _flex_reference_rate() * 1024 * 1024 ) - assert _flex_edit_cost(_edit_response(3 * 1024 * 1024)) - generated_only == pytest.approx( + assert _flex_edit_cost(_edit_response((3 * 1024 * 1024,))) - generated_only == pytest.approx( _flex_reference_rate() * 3 * 1024 * 1024 ) -def test_flux2_flex_edit_reference_cost_scales_with_reference_pixels_not_image_count(): +def test_flux2_flex_edit_reference_cost_does_not_scale_with_image_count(): two_outputs: Final = ImageResponse( data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")], size="1024x1024", - hidden_params={"reference_image_pixels": 1536 * 1024}, + hidden_params={"reference_image_pixels": (2048 * 1024,)}, ) assert _flex_edit_cost(two_outputs) == pytest.approx( - _flex_pixel_rate() * 1024 * 1024 * 2 + _flex_reference_rate() * 1536 * 1024 + _flex_pixel_rate() * 1024 * 1024 * 2 + _flex_reference_rate() * 2048 * 1024 ) -@pytest.mark.parametrize("reference_image_pixels", (0, -1, True, None, 1048576.0, "1048576")) -def test_flux2_flex_edit_bills_only_a_positive_integer_reference_count(reference_image_pixels: object): +@pytest.mark.parametrize( + "reference_image_pixels", + ((), (0,), (-1,), (True,), None, (1048576.0,), ("1048576",), 1048576, (1048576, "1048576"), (1048576, 0)), +) +def test_flux2_flex_edit_bills_only_positive_integer_reference_counts(reference_image_pixels: object): generated_only: Final = _flex_edit_cost(ImageResponse(data=[ImageObject(b64_json="aW1n")], size="1024x1024")) assert _flex_edit_cost(_edit_response(reference_image_pixels)) == generated_only @@ -271,7 +340,7 @@ def test_flux2_flex_edit_ignores_reference_pixels_when_provider_reports_token_us response: Final = ImageResponse( data=[ImageObject(b64_json="aW1n")], size="1024x1024", - hidden_params={"reference_image_pixels": 1024 * 1024}, + hidden_params={"reference_image_pixels": (1024 * 1024,)}, usage=ImageUsage( input_tokens=150, input_tokens_details=ImageUsageInputTokensDetails(image_tokens=100, text_tokens=50), @@ -294,7 +363,7 @@ def test_flux2_flex_edit_reads_the_reference_rate_from_its_own_catalog_key(monke litellm.get_model_info.cache_clear() _invalidate_model_cost_lowercase_map() - assert _flex_edit_cost(_edit_response(1024 * 1024)) == pytest.approx( + assert _flex_edit_cost(_edit_response((1024 * 1024,))) == pytest.approx( _flex_pixel_rate() * 1024 * 1024 + 3e-07 * 1024 * 1024 ) @@ -303,7 +372,7 @@ def test_flux2_flex_edit_prices_output_from_the_request_size_not_the_reference() large_reference_small_output: Final = ImageResponse( data=[ImageObject(b64_json="aW1n")], size="1024x1024", - hidden_params={"reference_image_pixels": 2048 * 2048}, + hidden_params={"reference_image_pixels": (2048 * 2048,)}, ) assert _flex_edit_cost(large_reference_small_output) == pytest.approx( @@ -312,14 +381,14 @@ def test_flux2_flex_edit_prices_output_from_the_request_size_not_the_reference() def test_flux2_flex_edit_honors_explicit_zero_deployment_generated_rate(): - cost: Final = _flex_edit_cost(_edit_response(1024 * 1024), model_info={"input_cost_per_pixel": 0.0}) + cost: Final = _flex_edit_cost(_edit_response((1024 * 1024,)), model_info={"input_cost_per_pixel": 0.0}) assert cost == pytest.approx(_flex_reference_rate() * 1024 * 1024) def test_flux2_flex_edit_prefers_deployment_reference_rate(): cost: Final = _flex_edit_cost( - _edit_response(1024 * 1024), + _edit_response((1024 * 1024,)), model_info={"input_cost_per_pixel": 2e-07, "input_cost_per_reference_pixel": 3e-07}, ) @@ -327,13 +396,13 @@ def test_flux2_flex_edit_prefers_deployment_reference_rate(): def test_flux2_flex_edit_deployment_reference_rate_alone_keeps_catalog_generated_rate(): - cost: Final = _flex_edit_cost(_edit_response(1024 * 1024), model_info={"input_cost_per_reference_pixel": 3e-07}) + cost: Final = _flex_edit_cost(_edit_response((1024 * 1024,)), model_info={"input_cost_per_reference_pixel": 3e-07}) assert cost == pytest.approx(_flex_pixel_rate() * 1024 * 1024 + 3e-07 * 1024 * 1024) def test_flux2_flex_edit_honors_explicit_zero_deployment_reference_rate(): - cost: Final = _flex_edit_cost(_edit_response(1024 * 1024), model_info={"input_cost_per_reference_pixel": 0.0}) + cost: Final = _flex_edit_cost(_edit_response((1024 * 1024,)), model_info={"input_cost_per_reference_pixel": 0.0}) assert cost == pytest.approx(_flex_pixel_rate() * 1024 * 1024) @@ -341,7 +410,7 @@ def test_flux2_flex_edit_honors_explicit_zero_deployment_reference_rate(): def test_unlisted_azure_ai_model_bills_deployment_reference_rate() -> None: cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( model="unlisted-flux-deployment", - completion_response=_edit_response(1024 * 1024), + completion_response=_edit_response((1024 * 1024,)), custom_llm_provider="azure_ai", size="1024x1024", call_type="image_edit", @@ -355,7 +424,7 @@ def test_unlisted_azure_ai_model_bills_deployment_reference_rate() -> None: def test_flat_priced_flux_edit_ignores_reference_pixels(model: str): cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( model=model, - completion_response=_edit_response(4 * 1024 * 1024), + completion_response=_edit_response((4 * 1024 * 1024,)), custom_llm_provider="azure_ai", size="1024x1024", call_type="image_edit", 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 8fbdcb12ef6..056eae205e3 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 @@ -185,7 +185,7 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[ generated_rate, reference_rate = _flex_rates() assert response._hidden_params["response_cost"] == pytest.approx( - generated_rate * 2048 * 1024 * 2 + reference_rate * 512 * 512 + generated_rate * 2048 * 1024 * 2 + reference_rate * 1024 * 1024 ) @@ -232,7 +232,7 @@ def _edit_ok(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"data": [{"b64_json": "aW1n"}]}) -def test_flux2_image_edit_bills_every_reference_by_its_header_dimensions(): +def test_flux2_image_edit_measures_every_reference_but_bills_each_of_several_as_one_megapixel(): sent: Final[dict[str, object]] = {} def respond(request: httpx.Request) -> httpx.Response: @@ -250,13 +250,12 @@ def test_flux2_image_edit_bills_every_reference_by_its_header_dimensions(): size="1024x1024", ) generated_rate, reference_rate = _flex_rates() - reference_pixels: Final = 1024 * 1024 + 800 * 600 + 640 * 480 assert sent["input_image"] == base64.b64encode(references[0]).decode() assert sent["input_image_3"] == base64.b64encode(references[2]).decode() - assert response._hidden_params["reference_image_pixels"] == reference_pixels + assert response._hidden_params["reference_image_pixels"] == (1024 * 1024, 800 * 600, 640 * 480) assert response._hidden_params["response_cost"] == pytest.approx( - generated_rate * 1024 * 1024 + reference_rate * reference_pixels + generated_rate * 1024 * 1024 + reference_rate * 3 * 1024 * 1024 ) @@ -273,28 +272,77 @@ def test_flux2_image_edit_reads_streams_once_and_still_measures_them(): ) generated_rate, reference_rate = _flex_rates() - assert response._hidden_params["reference_image_pixels"] == 2048 * 2048 + assert response._hidden_params["reference_image_pixels"] == (2048 * 2048,) assert response._hidden_params["response_cost"] == pytest.approx( generated_rate * 1024 * 1024 + reference_rate * 2048 * 2048 ) -def test_flux2_image_edit_bills_unmeasurable_references_as_one_megapixel_each(): +@pytest.mark.parametrize( + ("reference", "billed_megapixels"), + ( + pytest.param(_png(640, 640), 1, id="small-reference-rounds-up"), + pytest.param(_png(1024, 1280), 2, id="fractional-reference-rounds-up"), + pytest.param(_jpeg(4032, 3024), 4, id="photo-reference-caps-at-four-megapixels"), + ), +) +def test_flux2_image_edit_bills_a_lone_reference_in_whole_megapixels(reference: bytes, billed_megapixels: int): response: Final = litellm.image_edit( model="azure_ai/FLUX.2-flex", - image=[_png(640, 640), b"BM not a parseable header", b"\x89PNG\r\n\x1a\n\x00\x00"], - prompt="Blend every reference", + image=reference, + prompt="Make it a watercolor", api_key="test-key", api_base="https://example.services.ai.azure.com", client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_edit_ok))), size="1024x1024", ) generated_rate, reference_rate = _flex_rates() - reference_pixels: Final = 640 * 640 + 2 * UNMEASURED_REFERENCE_IMAGE_PIXELS - assert response._hidden_params["reference_image_pixels"] == reference_pixels assert response._hidden_params["response_cost"] == pytest.approx( - generated_rate * 1024 * 1024 + reference_rate * reference_pixels + generated_rate * 1024 * 1024 + reference_rate * billed_megapixels * 1024 * 1024 + ) + + +@pytest.mark.parametrize( + "reference", + ( + pytest.param(b"BM not a parseable header", id="unsupported-format"), + pytest.param(b"\x89PNG\r\n\x1a\n\x00\x00", id="truncated-header"), + pytest.param(_png(0, 640), id="zero-width-header"), + ), +) +def test_flux2_image_edit_bills_a_lone_unmeasurable_reference_as_one_megapixel(reference: bytes): + response: Final = litellm.image_edit( + model="azure_ai/FLUX.2-flex", + image=[reference], + prompt="Make it a watercolor", + api_key="test-key", + api_base="https://example.services.ai.azure.com", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_edit_ok))), + size="1024x1024", + ) + generated_rate, reference_rate = _flex_rates() + + assert response._hidden_params["reference_image_pixels"] == (UNMEASURED_REFERENCE_IMAGE_PIXELS,) + assert response._hidden_params["response_cost"] == pytest.approx( + generated_rate * 1024 * 1024 + reference_rate * 1024 * 1024 + ) + + +def test_flux2_image_edit_still_bills_every_reference_when_one_header_reports_zero_pixels(): + response: Final = litellm.image_edit( + model="azure_ai/FLUX.2-flex", + image=[_png(1024, 1024), _jpeg(640, 0)], + prompt="Blend both references", + api_key="test-key", + api_base="https://example.services.ai.azure.com", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_edit_ok))), + size="1024x1024", + ) + generated_rate, reference_rate = _flex_rates() + + assert response._hidden_params["response_cost"] == pytest.approx( + generated_rate * 1024 * 1024 + reference_rate * 2 * 1024 * 1024 ) @@ -324,7 +372,7 @@ def test_flux2_image_edit_resends_and_rebills_a_reused_stream(stream_position: s ) assert sent_images == [base64.b64encode(reference).decode()] * 2 - assert [response._hidden_params["reference_image_pixels"] for response in responses] == [2048 * 1024] * 2 + assert [response._hidden_params["reference_image_pixels"] for response in responses] == [(2048 * 1024,)] * 2 def test_flux2_pro_image_edit_bills_references_on_the_pro_reference_rate(): @@ -338,11 +386,10 @@ def test_flux2_pro_image_edit_bills_references_on_the_pro_reference_rate(): client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_edit_ok))), size="1024x1024", ) - reference_pixels: Final = 1024 * 1024 + 4032 * 3024 assert pro_row["input_cost_per_reference_pixel"] > 0 assert response._hidden_params["response_cost"] == pytest.approx( - pro_row["output_cost_per_image"] + pro_row["input_cost_per_reference_pixel"] * reference_pixels + pro_row["output_cost_per_image"] + pro_row["input_cost_per_reference_pixel"] * 2 * 1024 * 1024 ) @@ -401,7 +448,7 @@ async def test_flux2_aimage_edit_bills_references_like_image_edit(): ) generated_rate, reference_rate = _flex_rates() - assert response._hidden_params["reference_image_pixels"] == 2 * 1024 * 1024 + assert response._hidden_params["reference_image_pixels"] == (1024 * 1024, 1024 * 1024) assert response._hidden_params["response_cost"] == pytest.approx( generated_rate * 1024 * 1024 + reference_rate * 2 * 1024 * 1024 )