diff --git a/litellm/images/main.py b/litellm/images/main.py index 1f722eb752a..4794785b866 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -51,6 +51,7 @@ 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, ) @@ -883,6 +884,7 @@ 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/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 942fb64dde6..83497f6b624 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -1,6 +1,6 @@ import base64 from collections.abc import Mapping, Sequence -from io import BufferedReader, IOBase +from io import IOBase from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -8,6 +8,7 @@ import httpx from httpx._types import RequestFiles import litellm +from litellm._logging import verbose_logger from litellm.litellm_core_utils.token_counter import image_dimensions_from_bytes from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, @@ -28,6 +29,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj REFERENCE_IMAGE_PIXELS_HIDDEN_PARAM: Final = "reference_image_pixels" +UNMEASURED_REFERENCE_IMAGE_PIXELS: Final = 1024 * 1024 class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): @@ -137,15 +139,16 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): return request_body, [] def _read_image_bytes(self, image: FileTypes | Sequence[FileTypes]) -> bytes: - if isinstance(image, BufferedReader): - image_bytes: Final = image.read() - image.seek(0) - return image_bytes if isinstance(image, bytes): return image - if isinstance(image, IOBase): + if not isinstance(image, IOBase): + raise ValueError(f"Unsupported image type: {type(image)}") + if not image.seekable(): return image.read() - raise ValueError(f"Unsupported image type: {type(image)}") + image.seek(0) + image_bytes: Final = image.read() + image.seek(0) + return image_bytes def transform_image_edit_response( self, @@ -196,6 +199,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): def _pixel_count(image_bytes: bytes) -> int: dimensions: Final = image_dimensions_from_bytes(image_bytes) if dimensions is None: - return 0 + 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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8db2a91e81b..a3fd2bd4c81 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11371,6 +11371,7 @@ ] }, "azure_ai/flux.2-pro": { + "input_cost_per_reference_pixel": 1.430511474609375e-08, "litellm_provider": "azure_ai", "mode": "image_generation", "output_cost_per_image": 0.04, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8db2a91e81b..a3fd2bd4c81 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11371,6 +11371,7 @@ ] }, "azure_ai/flux.2-pro": { + "input_cost_per_reference_pixel": 1.430511474609375e-08, "litellm_provider": "azure_ai", "mode": "image_generation", "output_cost_per_image": 0.04, 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 ce39513d1b2..8fbdcb12ef6 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,7 +1,9 @@ import base64 +import datetime import io import json import struct +import uuid from collections.abc import Mapping from typing import Final @@ -11,11 +13,13 @@ import pytest import litellm from litellm.images.utils import ImageEditRequestUtils from litellm.llms.azure_ai.image_edit.flux2_transformation import ( + UNMEASURED_REFERENCE_IMAGE_PIXELS, AzureFoundryFlux2ImageEditConfig, ) from litellm.llms.azure_ai.image_edit.transformation import ( AzureFoundryFluxImageEditConfig, ) +from litellm.llms.custom_httpx import llm_http_handler as llm_http_handler_module from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -156,7 +160,7 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[ assert body == { "model": "FLUX.2-flex", "prompt": "Add a hat", - "input_image": base64.b64encode(b"image").decode(), + "input_image": base64.b64encode(_png(512, 512)).decode(), "num_images": 2, "width": 2048, "height": 1024, @@ -168,7 +172,7 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[ client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) response: Final = litellm.image_edit( model="azure_ai/FLUX.2-flex", - image=b"image", + image=_png(512, 512), prompt="Add a hat", api_key="test-key", api_base="https://example.services.ai.azure.com", @@ -178,9 +182,10 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[ steps="32", **dimensions, ) + generated_rate, reference_rate = _flex_rates() assert response._hidden_params["response_cost"] == pytest.approx( - litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"] * 2048 * 1024 * 2 + generated_rate * 2048 * 1024 * 2 + reference_rate * 512 * 512 ) @@ -196,16 +201,26 @@ def test_flux2_image_edit_accepts_and_drops_openai_only_parameters(): def _png(width: int, height: int) -> bytes: - return 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 ( + b"\x89PNG\r\n\x1a\n" + + (13).to_bytes(4, "big") + + b"IHDR" + + struct.pack(">II", width, height) + + b"\x08\x02\x00\x00\x00" + ) def _jpeg(width: int, height: int) -> bytes: - return b"\xff\xd8\xff\xc0" + struct.pack(">HBHHB", 17, 8, height, width, 3) + b"\x01\x22\x00\x02\x11\x01\x03\x11\x01" + return ( + b"\xff\xd8\xff\xc0" + struct.pack(">HBHHB", 17, 8, height, width, 3) + b"\x01\x22\x00\x02\x11\x01\x03\x11\x01" + ) def _webp(width: int, height: int) -> bytes: payload: Final = b"\x00\x00\x00\x9d\x01\x2a" + struct.pack(" tuple[float, float]: @@ -264,10 +279,10 @@ def test_flux2_image_edit_reads_streams_once_and_still_measures_them(): ) -def test_flux2_image_edit_bills_only_the_measurable_references(): +def test_flux2_image_edit_bills_unmeasurable_references_as_one_megapixel_each(): response: Final = litellm.image_edit( model="azure_ai/FLUX.2-flex", - image=[_png(1024, 1024), b"not an image", b"\x89PNG\r\n\x1a\n\x00\x00"], + image=[_png(640, 640), b"BM not a parseable header", b"\x89PNG\r\n\x1a\n\x00\x00"], prompt="Blend every reference", api_key="test-key", api_base="https://example.services.ai.azure.com", @@ -275,8 +290,98 @@ def test_flux2_image_edit_bills_only_the_measurable_references(): 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 + ) + + +@pytest.mark.parametrize("stream_position", ("start", "end")) +def test_flux2_image_edit_resends_and_rebills_a_reused_stream(stream_position: str): + 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) + + reference: Final = _png(2048, 1024) + uploaded: Final = io.BytesIO(reference) + if stream_position == "end": + uploaded.read() + responses: Final = tuple( + litellm.image_edit( + model="azure_ai/FLUX.2-flex", + image=[uploaded], + 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", + ) + for _attempt in range(2) + ) + + assert sent_images == [base64.b64encode(reference).decode()] * 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(): + pro_row: Final = litellm.model_cost["azure_ai/flux.2-pro"] + response: Final = litellm.image_edit( + model="azure_ai/flux.2-pro", + image=[_png(1024, 1024), _jpeg(4032, 3024)], + 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", + ) + 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 + ) + + +async def test_flux2_image_edit_bills_the_deployment_rates_when_logging_starts_before_routing( + 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) + generated_rate: Final = 1e-07 + reference_rate: Final = 2e-07 + router: Final = litellm.Router( + model_list=[ + { + "model_name": "flux2-flex-deployment", + "litellm_params": { + "model": "azure_ai/FLUX.2-flex", + "api_base": "https://example.services.ai.azure.com", + "api_key": "test-key", + "input_cost_per_pixel": generated_rate, + "input_cost_per_reference_pixel": reference_rate, + }, + } + ] + ) + 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 + ) - assert response._hidden_params["reference_image_pixels"] == 1024 * 1024 assert response._hidden_params["response_cost"] == pytest.approx( generated_rate * 1024 * 1024 + reference_rate * 1024 * 1024 )