diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 6cef9d7ab68..7b9e6f9d62a 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -222,7 +222,7 @@ COST_DESCRIPTIONS: dict[str, str] = { "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", - "input_cost_per_reference_pixel": "USD per pixel of each reference image sent with an image edit.", + "input_cost_per_reference_pixel": "USD per billable pixel of the reference images sent with an image edit (Azure AI FLUX.2 counts them in whole 1024x1024 megapixels).", } diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index 54a7ebfc5b4..4924a7b37d3 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -158,9 +158,10 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_pixel: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub input_cost_per_reference_pixel: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_query: Option, + /// USD per billable pixel of the reference images sent with an image edit (Azure AI FLUX.2 counts them in whole 1024x1024 megapixels). + #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_reference_pixel: Option, #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_request: Option, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/litellm/images/main.py b/litellm/images/main.py index 4794785b866..1f722eb752a 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -51,7 +51,6 @@ from litellm.types.llms.openai import ImageGenerationRequestQuality from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( LITELLM_IMAGE_VARIATION_PROVIDERS, - CustomPricingLiteLLMParams, LlmProviders, all_litellm_params, ) @@ -884,7 +883,6 @@ def image_edit( optional_params=dict(image_edit_request_params), litellm_params={ **image_edit_request_params, - **litellm_params.model_dump(include=set(CustomPricingLiteLLMParams.model_fields), exclude_none=True), "litellm_call_id": litellm_call_id, "model_info": model_info, }, diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index bf0673ec816..e9531427f48 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -54,6 +54,9 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import Message, SelectTokenizerResponse +MAX_JPEG_HEADER_SEGMENTS: Final = 1024 +MAX_JPEG_FILL_BYTES: Final = 1024 + def get_modified_max_tokens( model: str, @@ -248,7 +251,7 @@ def get_image_dimensions( _header, encoded = data.split(",", 1) img_data = base64.b64decode(encoded) - dimensions: Final = image_dimensions_from_bytes(img_data) + dimensions: Final = _header_dimensions(img_data) if dimensions is None: return DEFAULT_IMAGE_WIDTH, DEFAULT_IMAGE_HEIGHT return dimensions @@ -271,20 +274,7 @@ def _header_dimensions(img_data: bytes) -> tuple[int, int] | None: w, h = _unpack_ints("H", fhandle.read(2))[0] - 2 - fhandle.seek(1, 1) - h, w = _unpack_ints(">HH", fhandle.read(4)) - return w, h + return _jpeg_dimensions(img_data) if img_type == "webp": if img_data[12:16] == b"VP8X": w = _unpack_ints(" tuple[int, int] | None: return None +def _jpeg_dimensions(img_data: bytes) -> tuple[int, int] | None: + with io.BytesIO(img_data) as fhandle: + fhandle.seek(2) + for _ in range(MAX_JPEG_HEADER_SEGMENTS): + marker = _next_jpeg_marker(fhandle) + if marker is None: + return None + segment_length = _unpack_ints(">H", fhandle.read(2))[0] + if 0xC0 <= marker <= 0xCF and marker not in (0xC4, 0xC8, 0xCC): + fhandle.seek(1, 1) + h, w = _unpack_ints(">HH", fhandle.read(4)) + return w, h + if segment_length < 2: + return None + fhandle.seek(segment_length - 2, 1) + return None + + +def _next_jpeg_marker(fhandle: io.BytesIO) -> int | None: + for _ in range(MAX_JPEG_FILL_BYTES): + byte = fhandle.read(1) + if byte != b"\xff": + return ord(byte) + return None + + def calculate_img_tokens( data, mode: Literal["low", "high", "auto"] = "auto", diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 98edddc9261..11554d1d94b 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -1,6 +1,5 @@ import base64 from collections.abc import Mapping, Sequence -from io import IOBase from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -141,13 +140,15 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): def _read_image_bytes(self, image: FileTypes | Sequence[FileTypes]) -> bytes: if isinstance(image, bytes): return image - if not isinstance(image, IOBase): + read: Final[object] = getattr(image, "read", None) + if not callable(read): raise ValueError(f"Unsupported image type: {type(image)}") - if not image.seekable(): - return image.read() - image.seek(0) - image_bytes: Final = image.read() - image.seek(0) + rewound: Final = _rewind(image) + image_bytes: Final = read() + if rewound: + _rewind(image) + if not isinstance(image_bytes, bytes): + raise TypeError("FLUX.2 reference images must be opened in binary mode") return image_bytes def transform_image_edit_response( @@ -196,6 +197,20 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): ) +def _rewind(image: object) -> bool: + seekable: Final[object] = getattr(image, "seekable", None) + if callable(seekable) and not seekable(): + return False + seek: Final[object] = getattr(image, "seek", None) + if not callable(seek): + return False + try: + seek(0) + except OSError: + return False + return True + + def _pixel_count(image_bytes: bytes) -> int: dimensions: Final = image_dimensions_from_bytes(image_bytes) pixels: Final = dimensions[0] * dimensions[1] if dimensions is not None else 0 diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index d05fc43f9e8..40324f68e84 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -1,24 +1,35 @@ +import base64 import math from collections.abc import Mapping +from dataclasses import dataclass from typing import Annotated, Any, Final from pydantic import Field, TypeAdapter, ValidationError import litellm +from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.utils import ( _get_cost_per_unit, calculate_image_response_cost_from_usage, resolve_image_model_info, ) +from litellm.litellm_core_utils.token_counter import image_dimensions_from_bytes 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 +from litellm.types.utils import ImageObject, ImageResponse, ModelInfo MEGAPIXEL: Final = 1024 * 1024 MAX_LONE_REFERENCE_MEGAPIXELS: Final = 4 _REFERENCE_PIXELS: Final = TypeAdapter(tuple[Annotated[int, Field(strict=True, gt=0)], ...]) +@dataclass(frozen=True, slots=True, kw_only=True) +class _Flux2MegapixelPrices: + first: float + additional: float + reference: 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: @@ -27,13 +38,47 @@ def _price(resolved: ModelInfo, cost_key: str) -> float | None: 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 None - return shared_entry.get(cost_key) + return _get_cost_per_unit(shared_entry, cost_key, default_value=None) def _pixel_rate(resolved: ModelInfo, cost_key: str) -> float: return _price(resolved, cost_key) or 0.0 +def _deployment_price(deployment: ModelInfo | None, cost_key: str) -> float | None: + if deployment is None: + return None + return _get_cost_per_unit(deployment, cost_key, default_value=None) + + +def _flux2_prices(resolved: ModelInfo, deployment: ModelInfo | None) -> _Flux2MegapixelPrices: + # A deployment that prices the generated image itself also prices its references, unless it sets a reference + # rate: a flat per-image price covers them, and a pixel rate bills them at that rate + deployment_reference_rate: Final = _deployment_price(deployment, "input_cost_per_reference_pixel") + deployment_image_price: Final = _deployment_price(deployment, "output_cost_per_image") + if deployment_image_price is not None: + return _Flux2MegapixelPrices( + first=deployment_image_price, + additional=0.0, + reference=(deployment_reference_rate or 0.0) * MEGAPIXEL, + ) + deployment_pixel_rate: Final = _deployment_price(deployment, "input_cost_per_pixel") + if deployment_pixel_rate is not None: + return _Flux2MegapixelPrices( + first=deployment_pixel_rate * MEGAPIXEL, + additional=deployment_pixel_rate * MEGAPIXEL, + reference=(deployment_pixel_rate if deployment_reference_rate is None else deployment_reference_rate) + * MEGAPIXEL, + ) + catalog_megapixel_rate: Final = _pixel_rate(resolved, "input_cost_per_pixel") * MEGAPIXEL + catalog_first_megapixel: Final = _price(resolved, "output_cost_per_image") + return _Flux2MegapixelPrices( + first=catalog_megapixel_rate if catalog_first_megapixel is None else catalog_first_megapixel, + additional=catalog_megapixel_rate, + reference=_pixel_rate(resolved, "input_cost_per_reference_pixel") * MEGAPIXEL, + ) + + def _billable_megapixels(pixels: int) -> int: return math.ceil(pixels / MEGAPIXEL) @@ -50,21 +95,47 @@ def _billable_reference_megapixels(reference_pixels: tuple[int, ...]) -> int: return len(reference_pixels) -def _reference_cost(resolved: ModelInfo, image_response: ImageResponse) -> float: +def _reference_cost(prices: _Flux2MegapixelPrices, image_response: ImageResponse) -> float: reported_pixels: Final = image_response._hidden_params.get(REFERENCE_IMAGE_PIXELS_HIDDEN_PARAM) if reported_pixels is None: return 0.0 try: reference_pixels: Final = _REFERENCE_PIXELS.validate_python(reported_pixels) except ValidationError: + verbose_logger.warning("Ignoring malformed FLUX.2 reference pixel counts: %r", reported_pixels) return 0.0 - return ( - _pixel_rate(resolved, "input_cost_per_reference_pixel") - * MEGAPIXEL - * _billable_reference_megapixels(reference_pixels) + return prices.reference * _billable_reference_megapixels(reference_pixels) + + +def _flux2_generated_cost( + prices: _Flux2MegapixelPrices, image_response: ImageResponse, requested_pixels: int, n: int | None +) -> float: + return sum( + prices.first + prices.additional * (_billable_megapixels(pixels) - 1) + for pixels in _generated_pixels(image_response, requested_pixels, n) ) +def _generated_pixels(image_response: ImageResponse, requested_pixels: int, n: int | None) -> tuple[int, ...]: + images: Final = image_response.data or () + if not images: + return (requested_pixels,) * (n or 0) + return tuple(_measured_pixels(image) or requested_pixels for image in images) + + +def _measured_pixels(image: ImageObject) -> int | None: + if not image.b64_json: + return None + try: + image_bytes: Final = base64.b64decode(image.b64_json) + except ValueError: + return None + dimensions: Final = image_dimensions_from_bytes(image_bytes) + if dimensions is None: + return None + return dimensions[0] * dimensions[1] or None + + def cost_calculator( model: str, image_response: Any, @@ -74,7 +145,7 @@ def cost_calculator( model_info: ModelInfo | None = None, ) -> float: """ - Azure AI image generation cost calculator + Azure AI image generation and image edit cost calculator """ _model_info: Final = resolve_image_model_info( model=model, @@ -92,6 +163,13 @@ def cost_calculator( if token_based_cost is not None: return token_based_cost + if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model): + prices: Final = _flux2_prices(_model_info, model_info) + requested_pixels: Final = _size_pixels(_output_size(size, optional_params, image_response)) or MEGAPIXEL + return _flux2_generated_cost(prices, image_response, requested_pixels, n) + _reference_cost( + prices, image_response + ) + return _generated_cost( model=model, resolved=_model_info, @@ -100,7 +178,7 @@ def cost_calculator( n=n, optional_params=optional_params, model_info=model_info, - ) + _reference_cost(_model_info, image_response) + ) raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") @@ -115,9 +193,6 @@ 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 @@ -129,19 +204,12 @@ def _generated_cost( return default_image_cost_calculator( model=resolved.get("key", model), custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, - size=pixel_size, + size=_output_size(size, optional_params, image_response), 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: @@ -152,6 +220,8 @@ def _output_size( 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 +def _size_pixels(size: str | None) -> int | None: + dimensions: Final = (size or "").lower().replace("-x-", "x").split("x") + if len(dimensions) != 2 or not all(dimension.isdigit() for dimension in dimensions): + return None + return int(dimensions[0]) * int(dimensions[1]) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bade2687d52..c252fe68a7b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11378,7 +11378,15 @@ "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" + "/v1/images/generations", + "/v1/images/edits" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "image" ] }, "azure_ai/FLUX.2-flex": { diff --git a/litellm/types/router.py b/litellm/types/router.py index b72809f625f..0d89c9b5081 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -788,6 +788,7 @@ class ModelGroupInfo(BaseModel): input_cost_per_token: float | None = None output_cost_per_token: float | None = None input_cost_per_pixel: float | None = None + input_cost_per_reference_pixel: float | None = None mode: ( str | Literal["chat", "embedding", "completion", "image_generation", "audio_transcription", "rerank", "moderations"] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 746576594f8..bb2899fb0a6 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -355,6 +355,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_character_above_128k_tokens: float | None # only for vertex ai models output_cost_per_image: float | None output_cost_per_pixel: ReadOnly[float | None] + input_cost_per_reference_pixel: ReadOnly[float | None] output_cost_per_image_token: float | None output_cost_per_video_token: float | None # for gemini omni models with video output output_vector_size: int | None diff --git a/litellm/utils.py b/litellm/utils.py index 4ea0769ea11..4551ffcedc6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6256,6 +6256,7 @@ def _get_model_info_helper( output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), output_cost_per_pixel=_model_info.get("output_cost_per_pixel", None), + input_cost_per_reference_pixel=_model_info.get("input_cost_per_reference_pixel", None), output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None), output_vector_size=_model_info.get("output_vector_size", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bade2687d52..c252fe68a7b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11378,7 +11378,15 @@ "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" + "/v1/images/generations", + "/v1/images/edits" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "image" ] }, "azure_ai/FLUX.2-flex": { diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 09a1831660f..c38adc685ed 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -331,7 +331,7 @@ "input_cost_per_reference_pixel": { "type": "number", "minimum": 0, - "description": "USD per pixel of each reference image sent with an image edit." + "description": "USD per billable pixel of the reference images sent with an image edit (Azure AI FLUX.2 counts them in whole 1024x1024 megapixels)." }, "input_cost_per_request": { "type": "number", diff --git a/terraform/provider/docs/resources/model.md b/terraform/provider/docs/resources/model.md index 620fd963091..8068cfad5ec 100644 --- a/terraform/provider/docs/resources/model.md +++ b/terraform/provider/docs/resources/model.md @@ -159,7 +159,7 @@ The following arguments are supported: * `input_cost_per_pixel` - (Optional) float. Cost applied per input pixel for models that charge by image size. -* `input_cost_per_reference_pixel` - (Optional) float. Cost applied per pixel of each reference image sent with an image edit, for models that meter reference images separately. +* `input_cost_per_reference_pixel` - (Optional) float. Cost applied per billable pixel of the reference images sent with an image edit, for models that meter reference images separately. Azure AI FLUX.2 counts reference images in whole 1024x1024 megapixels. * `output_cost_per_pixel` - (Optional) float. Cost applied per output pixel for image-generation models. diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 53a09d8feed..9f1f13bc5a8 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.token_counter import ( _get_extrapolating_count_function, _get_tiktoken_count_function, calculate_img_tokens, + get_image_dimensions, high_detail_image_token_upper_bound, image_dimensions_from_bytes, offload_token_count, @@ -1673,3 +1674,37 @@ def test_image_dimensions_from_bytes_reads_each_header_format(image: bytes, expe ) def test_image_dimensions_from_bytes_returns_none_for_unreadable_headers(image: bytes) -> None: assert image_dimensions_from_bytes(image) is None + + +def _jpeg_sof(width: int, height: int) -> bytes: + return b"\xff\xc0" + struct.pack(">HBHHB", 17, 8, height, width, 3) + b"\x01\x22\x00\x02\x11\x01\x03\x11\x01" + + +@pytest.mark.parametrize( + "image", + [ + pytest.param(b"\xff\xd8" + b"\xff\xe0\x00\x02" * 1025 + _jpeg_sof(800, 600), id="too-many-segments"), + pytest.param(b"\xff\xd8" + b"\xff" * 2000 + _jpeg_sof(800, 600)[1:], id="too-many-fill-bytes"), + pytest.param(b"\xff\xd8\xff\xe0\x00\x00\x02" + _jpeg_sof(800, 600), id="segment-length-below-two"), + ], +) +def test_image_dimensions_from_bytes_gives_up_on_pathological_jpeg_headers(image: bytes) -> None: + assert image_dimensions_from_bytes(image) is None + + +def test_image_dimensions_from_bytes_still_reads_a_jpeg_with_many_real_segments() -> None: + image: Final = b"\xff\xd8" + b"\xff\xe0\x00\x02" * 1000 + b"\xff" * 64 + _jpeg_sof(800, 600)[1:] + + assert image_dimensions_from_bytes(image) == (800, 600) + + +@pytest.mark.parametrize( + "header", + [ + pytest.param(b"\x89PNG\r\n\x1a\n\x00\x00", id="png-truncated"), + pytest.param(b"\xff\xd8\xff\xe0\x00\x10JFIF", id="jpeg-truncated"), + ], +) +def test_get_image_dimensions_still_raises_for_a_truncated_header(header: bytes) -> None: + with pytest.raises((struct.error, TypeError)): + get_image_dimensions(data="data:image/png;base64," + base64.b64encode(header).decode()) 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 056eae205e3..50bbeefb520 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 @@ -1,10 +1,11 @@ import base64 -import datetime +import contextlib import io import json +import pathlib import struct -import uuid -from collections.abc import Mapping +import tempfile +from collections.abc import Callable, Iterator, Mapping from typing import Final import httpx @@ -259,7 +260,143 @@ def test_flux2_image_edit_measures_every_reference_but_bills_each_of_several_as_ ) -def test_flux2_image_edit_reads_streams_once_and_still_measures_them(): +class _ReadOnlyUpload: + def __init__(self, data: bytes) -> None: + self._data: Final = data + + def read(self) -> bytes: + return self._data + + +class _SeekableUploadWithoutSeekable: + def __init__(self, data: bytes) -> None: + self._stream: Final = io.BytesIO(data) + self._stream.read() + + def read(self) -> bytes: + return self._stream.read() + + def seek(self, offset: int, whence: int = 0) -> int: + return self._stream.seek(offset, whence) + + +class _UploadWithBrokenSeek: + def __init__(self, data: bytes) -> None: + self._data: Final = data + + def read(self) -> bytes: + return self._data + + def seek(self, offset: int, whence: int = 0) -> int: + raise io.UnsupportedOperation("seek") + + +class _NonSeekableUpload: + def __init__(self, data: bytes) -> None: + self._data: Final = data + + def read(self) -> bytes: + return self._data + + def seekable(self) -> bool: + return False + + def seek(self, offset: int, whence: int = 0) -> int: + raise AssertionError("a non-seekable upload must not be seeked") + + +def _written_temporary_file(stack: contextlib.ExitStack, data: bytes, spooled: bool) -> object: + upload: Final = stack.enter_context(tempfile.SpooledTemporaryFile() if spooled else tempfile.NamedTemporaryFile()) + upload.write(data) + upload.flush() + return upload + + +UPLOADS: Final[Mapping[str, Callable[[contextlib.ExitStack, pathlib.Path, bytes], object]]] = { + "bytesio": lambda _stack, _path, data: io.BytesIO(data), + "buffered-reader": lambda stack, path, _data: stack.enter_context(path.open("rb")), + "named-temporary-file-at-eof": lambda stack, _path, data: _written_temporary_file(stack, data, spooled=False), + "spooled-temporary-file-at-eof": lambda stack, _path, data: _written_temporary_file(stack, data, spooled=True), + "duck-typed-read-only": lambda _stack, _path, data: _ReadOnlyUpload(data), + "duck-typed-seek-without-seekable-at-eof": lambda _stack, _path, data: _SeekableUploadWithoutSeekable(data), + "duck-typed-seek-raises": lambda _stack, _path, data: _UploadWithBrokenSeek(data), + "non-seekable-stream": lambda _stack, _path, data: _NonSeekableUpload(data), +} + + +@pytest.fixture +def exit_stack() -> Iterator[contextlib.ExitStack]: + with contextlib.ExitStack() as stack: + yield stack + + +@pytest.mark.parametrize("upload_kind", tuple(UPLOADS)) +def test_flux2_image_edit_sends_and_bills_every_readable_upload( + upload_kind: str, tmp_path: pathlib.Path, exit_stack: contextlib.ExitStack +): + reference: Final = _png(2048, 1024) + path: Final = tmp_path / "reference.png" + path.write_bytes(reference) + sent_images: Final[list[str]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + sent_images.append(json.loads(request.content)["input_image"]) + return _edit_ok(request) + + response: Final = litellm.image_edit( + model="azure_ai/FLUX.2-flex", + image=UPLOADS[upload_kind](exit_stack, path, 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(respond))), + size="1024x1024", + ) + generated_rate, reference_rate = _flex_rates() + + assert sent_images == [base64.b64encode(reference).decode()] + assert response._hidden_params["response_cost"] == pytest.approx( + generated_rate * 1024 * 1024 + reference_rate * 2 * 1024 * 1024 + ) + + +@pytest.mark.parametrize("upload_kind", ("bytesio", "buffered-reader", "named-temporary-file-at-eof")) +def test_flux2_image_edit_leaves_a_seekable_upload_rewound_for_reuse( + upload_kind: str, tmp_path: pathlib.Path, exit_stack: contextlib.ExitStack +): + path: Final = tmp_path / "reference.png" + path.write_bytes(_png(1024, 1024)) + upload: Final = UPLOADS[upload_kind](exit_stack, path, _png(1024, 1024)) + + litellm.image_edit( + model="azure_ai/FLUX.2-flex", + image=upload, + 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", + ) + + assert upload.tell() == 0 + + +def test_flux2_image_edit_rejects_a_text_mode_upload_with_a_clear_error(tmp_path: pathlib.Path): + path: Final = tmp_path / "reference.png" + path.write_bytes(_png(1024, 1024)) + + with path.open("r", encoding="latin-1") as text_upload, pytest.raises(TypeError, match="binary mode"): + AzureFoundryFlux2ImageEditConfig().transform_image_edit_request( + model="FLUX.2-flex", + prompt="Make it a watercolor", + image=text_upload, + image_edit_optional_request_params={}, + litellm_params={}, + headers={}, + ) + + +def test_flux2_image_edit_measures_a_stream_reference(): uploaded: Final = io.BytesIO(_png(2048, 2048)) response: Final = litellm.image_edit( model="azure_ai/FLUX.2-flex", @@ -278,6 +415,7 @@ def test_flux2_image_edit_reads_streams_once_and_still_measures_them(): ) +# Billable megapixels for a lone reference as Azure's FLUX.2-pro request_meta reported them on 2026-09-25 @pytest.mark.parametrize( ("reference", "billed_megapixels"), ( @@ -393,9 +531,7 @@ def test_flux2_pro_image_edit_bills_references_on_the_pro_reference_rate(): ) -async def test_flux2_image_edit_bills_the_deployment_rates_when_logging_starts_before_routing( - monkeypatch: pytest.MonkeyPatch, -): +async def test_flux2_router_image_edit_bills_the_deployment_rates(monkeypatch: pytest.MonkeyPatch): mock_client: Final = AsyncHTTPHandler() mock_client.client = httpx.AsyncClient(transport=httpx.MockTransport(_edit_ok)) monkeypatch.setattr(llm_http_handler_module, "get_async_httpx_client", lambda **_kwargs: mock_client) @@ -415,22 +551,43 @@ async def test_flux2_image_edit_bills_the_deployment_rates_when_logging_starts_b } ] ) - logging_obj, request_data = litellm.utils.function_setup( - original_function="aimage_edit", - rules_obj=litellm.utils.Rules(), - start_time=datetime.datetime.now(), - model="flux2-flex-deployment", - prompt="Make it a watercolor", - size="1024x1024", - litellm_call_id=str(uuid.uuid4()), - ) response: Final = await router.aimage_edit( - **request_data, image=[_png(1024, 1024)], litellm_logging_obj=logging_obj + model="flux2-flex-deployment", + prompt="Make it a watercolor", + image=[_png(1024, 1280)], + size="1024x1280", ) assert response._hidden_params["response_cost"] == pytest.approx( - generated_rate * 1024 * 1024 + reference_rate * 1024 * 1024 + generated_rate * 2 * 1024 * 1024 + reference_rate * 2 * 1024 * 1024 + ) + + +def _edit_returning(image: bytes) -> Callable[[httpx.Request], httpx.Response]: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"b64_json": base64.b64encode(image).decode()}]}) + + return respond + + +@pytest.mark.parametrize("size", (None, "1024x1024"), ids=("no-size", "size-azure-did-not-use")) +def test_flux2_image_edit_bills_the_generated_image_azure_returned(size: str | None): + response: Final = litellm.image_edit( + model="azure_ai/flux.2-pro", + image=_png(1024, 1280), + 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_returning(_png(1024, 1280))))), + **({} if size is None else {"size": size}), + ) + pro_row: Final = litellm.model_cost["azure_ai/flux.2-pro"] + + assert response._hidden_params["response_cost"] == pytest.approx( + pro_row["output_cost_per_image"] + + pro_row["input_cost_per_pixel"] * 1024 * 1024 + + pro_row["input_cost_per_reference_pixel"] * 2 * 1024 * 1024 ) 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 c29b538e664..16f94ee6fd3 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 @@ -1,4 +1,6 @@ -from collections.abc import Mapping +import base64 +import struct +from collections.abc import Callable, Mapping from typing import Final from unittest.mock import MagicMock @@ -163,11 +165,21 @@ def test_flux2_flex_model_info(): assert model_info["max_input_tokens"] == 32000 assert model_info["max_tokens"] == 32000 assert model_info["supported_endpoints"] == ["/v1/images/generations", "/v1/images/edits"] - assert catalog_info["input_cost_per_pixel"] == catalog_info["input_cost_per_reference_pixel"] assert catalog_info["supported_modalities"] == ["text", "image"] assert catalog_info["supported_output_modalities"] == ["image"] +@pytest.mark.parametrize("model", ("flux.2-pro", "FLUX.2-flex")) +def test_flux2_model_info_reports_the_reference_rate_and_the_edit_endpoint(model: str): + model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") + + assert ( + model_info["input_cost_per_reference_pixel"] + == litellm.model_cost[f"azure_ai/{model}"]["input_cost_per_reference_pixel"] + ) + assert "/v1/images/edits" in model_info["supported_endpoints"] + + def test_flux2_flex_cost_uses_generated_megapixels(): response = ImageResponse( data=[ @@ -269,6 +281,203 @@ def test_flux2_pro_generation_bills_the_first_megapixel_then_each_additional_one assert cost == pytest.approx(first + additional * (megapixels - 1)) +def _catalog_image_cost(model: str, megapixels: int) -> float: + row: Final = litellm.model_cost[f"azure_ai/{model}"] + megapixel_rate: Final = row["input_cost_per_pixel"] * 1024 * 1024 + return row.get("output_cost_per_image", megapixel_rate) + megapixel_rate * (megapixels - 1) + + +@pytest.mark.parametrize("model", ("flux.2-pro", "FLUX.2-flex")) +@pytest.mark.parametrize(("size", "megapixels"), (("1024x1024", 1), ("2048x2048", 4))) +@pytest.mark.parametrize( + ("deployment_prices", "expected"), + ( + pytest.param(None, None, id="catalog-prices"), + pytest.param({"output_cost_per_image": 0.07}, lambda megapixels: 0.07, id="deployment-flat-price"), + pytest.param( + {"output_cost_per_image": 0.07, "input_cost_per_pixel": 1e-07}, + lambda megapixels: 0.07, + id="deployment-flat-price-wins-over-its-pixel-rate", + ), + pytest.param( + {"input_cost_per_pixel": 1e-07}, + lambda megapixels: 1e-07 * 1024 * 1024 * megapixels, + id="deployment-pixel-rate-prices-every-megapixel", + ), + pytest.param( + {"input_cost_per_reference_pixel": 3e-07}, None, id="deployment-reference-rate-keeps-catalog-output" + ), + ), +) +def test_flux2_generation_bills_each_price_source_as_its_owner_set_it( + model: str, + size: str, + megapixels: int, + deployment_prices: Mapping[str, float] | None, + expected: Callable[[int], float] | None, +): + cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( + model=model, + completion_response=ImageResponse(data=[ImageObject(b64_json="aW1n")]), + custom_llm_provider="azure_ai", + size=size, + call_type="image_generation", + model_info=None if deployment_prices is None else dict(deployment_prices), + ) + + assert cost == pytest.approx(_catalog_image_cost(model, megapixels) if expected is None else expected(megapixels)) + + +@pytest.mark.parametrize("model", ("flux.2-pro", "FLUX.2-flex")) +@pytest.mark.parametrize( + ("deployment_prices", "reference_megapixel_price"), + ( + pytest.param(None, None, id="catalog-reference-rate"), + pytest.param({"output_cost_per_image": 0.07}, 0.0, id="deployment-flat-price-covers-references"), + pytest.param( + {"output_cost_per_image": 0.07, "input_cost_per_reference_pixel": 3e-07}, + 3e-07 * 1024 * 1024, + id="deployment-flat-price-with-its-reference-rate", + ), + pytest.param( + {"input_cost_per_pixel": 1e-07}, 1e-07 * 1024 * 1024, id="deployment-pixel-rate-prices-references" + ), + pytest.param( + {"input_cost_per_pixel": 1e-07, "input_cost_per_reference_pixel": 3e-07}, + 3e-07 * 1024 * 1024, + id="deployment-reference-rate-wins-over-its-pixel-rate", + ), + pytest.param( + {"input_cost_per_reference_pixel": 3e-07}, 3e-07 * 1024 * 1024, id="deployment-reference-rate-alone" + ), + ), +) +def test_flux2_edit_prices_references_from_the_source_that_priced_the_image( + model: str, deployment_prices: Mapping[str, float] | None, reference_megapixel_price: float | None +): + def edit_cost(reference_pixels: tuple[int, ...]) -> float: + return CostCalculatorUtils.route_image_generation_cost_calculator( + model=model, + completion_response=_edit_response(reference_pixels), + custom_llm_provider="azure_ai", + size="1024x1024", + call_type="image_edit", + model_info=None if deployment_prices is None else dict(deployment_prices), + ) + + expected_per_megapixel: Final = ( + litellm.model_cost[f"azure_ai/{model}"]["input_cost_per_reference_pixel"] * 1024 * 1024 + if reference_megapixel_price is None + else reference_megapixel_price + ) + + assert edit_cost((1024 * 1280,)) - edit_cost(()) == pytest.approx(expected_per_megapixel * 2) + + +def _png_b64(width: int, height: int) -> str: + png: Final = ( + b"\x89PNG\r\n\x1a\n" + + (13).to_bytes(4, "big") + + b"IHDR" + + struct.pack(">II", width, height) + + b"\x08\x02\x00\x00\x00" + ) + return base64.b64encode(png).decode() + + +@pytest.mark.parametrize( + ("returned_images", "requested_size", "expected_megapixels"), + ( + pytest.param(((1024, 1280),), "1024x1024", (2,), id="returned-size-wins-over-requested"), + pytest.param(((512, 512), (2048, 1024)), "1024x1024", (1, 2), id="each-returned-image-measured"), + pytest.param((None,), "1024x1280", (2,), id="unmeasurable-image-falls-back-to-requested"), + ), +) +def test_flux2_pro_bills_each_generated_image_by_its_returned_size( + returned_images: tuple[tuple[int, int] | None, ...], requested_size: str, expected_megapixels: tuple[int, ...] +): + first, additional, _reference = _pro_megapixel_rates() + response: Final = ImageResponse( + data=[ImageObject(b64_json="aW1n" if image is None else _png_b64(*image)) for image in returned_images] + ) + + cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( + model="flux.2-pro", + completion_response=response, + custom_llm_provider="azure_ai", + size=requested_size, + call_type="image_generation", + ) + + assert cost == pytest.approx(sum(first + additional * (megapixels - 1) for megapixels in expected_megapixels)) + + +@pytest.mark.parametrize("size", ("auto", "large", "0x1024", "1024x")) +def test_flux2_pro_bills_one_megapixel_for_a_size_it_cannot_measure(size: str): + 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) + + +@pytest.mark.parametrize(("n", "billed_images"), ((2, 2), (None, 0))) +def test_flux2_pro_bills_the_requested_image_count_when_the_response_lists_none(n: int | None, billed_images: int): + first, _additional, _reference = _pro_megapixel_rates() + + cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( + model="flux.2-pro", + completion_response=ImageResponse(data=[]), + custom_llm_provider="azure_ai", + size="1024x1024", + n=n, + call_type="image_generation", + ) + + assert cost == pytest.approx(first * billed_images) + + +def test_flux2_pro_reads_an_uppercase_size(): + 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="2048X1024", + call_type="image_generation", + ) + + assert cost == pytest.approx(first + additional) + + +def test_flux2_edit_parses_string_prices_registered_in_the_cost_map(monkeypatch: pytest.MonkeyPatch): + row: Final = litellm.model_cost["azure_ai/FLUX.2-flex"] + monkeypatch.setitem(row, "input_cost_per_pixel", "2e-07") + monkeypatch.setitem(row, "input_cost_per_reference_pixel", "3e-07") + litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() + + assert _flex_edit_cost(_edit_response((1024 * 1024,))) == pytest.approx(2e-07 * 1024 * 1024 + 3e-07 * 1024 * 1024) + + +@pytest.fixture +def distinct_pro_megapixel_prices(monkeypatch: pytest.MonkeyPatch) -> tuple[float, float, float]: + row: Final = litellm.model_cost["azure_ai/flux.2-pro"] + monkeypatch.setitem(row, "output_cost_per_image", 0.05) + monkeypatch.setitem(row, "input_cost_per_pixel", 2e-08) + monkeypatch.setitem(row, "input_cost_per_reference_pixel", 3e-08) + litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() + return 0.05, 2e-08 * 1024 * 1024, 3e-08 * 1024 * 1024 + + # 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"), @@ -283,9 +492,13 @@ def test_flux2_pro_generation_bills_the_first_megapixel_then_each_additional_one ), ) def test_flux2_pro_edit_bills_the_megapixels_azure_meters( - size: str, reference_pixels: tuple[int, ...], output_megapixels: int, reference_megapixels: int + size: str, + reference_pixels: tuple[int, ...], + output_megapixels: int, + reference_megapixels: int, + distinct_pro_megapixel_prices: tuple[float, float, float], ): - first, additional, reference = _pro_megapixel_rates() + first, additional, reference = distinct_pro_megapixel_prices cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( model="flux.2-pro", @@ -298,10 +511,6 @@ def test_flux2_pro_edit_bills_the_megapixels_azure_meters( assert cost == pytest.approx(first + additional * (output_megapixels - 1) + reference * reference_megapixels) -def test_flux2_flex_reference_rate_matches_generated_rate(): - assert _flex_reference_rate() == _flex_pixel_rate() - - 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")) @@ -380,10 +589,10 @@ 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(): +def test_flux2_flex_edit_with_a_free_deployment_pixel_rate_bills_nothing(): 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) + assert cost == 0.0 def test_flux2_flex_edit_prefers_deployment_reference_rate(): @@ -407,17 +616,17 @@ def test_flux2_flex_edit_honors_explicit_zero_deployment_reference_rate(): assert cost == pytest.approx(_flex_pixel_rate() * 1024 * 1024) -def test_unlisted_azure_ai_model_bills_deployment_reference_rate() -> None: +def test_custom_named_flux2_deployment_bills_its_own_megapixel_and_reference_rates() -> None: cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( - model="unlisted-flux-deployment", - completion_response=_edit_response((1024 * 1024,)), + model="my-flux2-prod", + completion_response=_edit_response((1024 * 1280,)), custom_llm_provider="azure_ai", - size="1024x1024", + size="1024x1280", call_type="image_edit", - model_info={"input_cost_per_pixel": 1e-07, "input_cost_per_reference_pixel": 1e-07}, + model_info={"input_cost_per_pixel": 1e-07, "input_cost_per_reference_pixel": 2e-07}, ) - assert cost == pytest.approx(1e-07 * 1024 * 1024 * 2) + assert cost == pytest.approx(1e-07 * 1024 * 1024 * 2 + 2e-07 * 1024 * 1024 * 2) @pytest.mark.parametrize("model", ("FLUX-1.1-pro", "FLUX.1-Kontext-pro")) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b1b151ebc8d..edcc9ed1e67 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35934,6 +35934,8 @@ export interface components { health_status?: string | null; /** Input Cost Per Pixel */ input_cost_per_pixel?: number | null; + /** Input Cost Per Reference Pixel */ + input_cost_per_reference_pixel?: number | null; /** Input Cost Per Token */ input_cost_per_token?: number | null; /**