diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index f91a87ba0f4..942fb64dde6 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -1,12 +1,14 @@ import base64 from collections.abc import Mapping, Sequence -from io import BufferedReader +from io import BufferedReader, IOBase from types import MappingProxyType -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final +import httpx from httpx._types import RequestFiles import litellm +from litellm.litellm_core_utils.token_counter import image_dimensions_from_bytes from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, get_azure_ai_auth_headers, @@ -14,11 +16,18 @@ from litellm.llms.azure_ai.common_utils import ( from litellm.llms.azure_ai.image_generation.flux_transformation import ( AzureFoundryFluxImageGenerationConfig, ) +from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.openai import FileTypes from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +REFERENCE_IMAGE_PIXELS_HIDDEN_PARAM: Final = "reference_image_pixels" class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): @@ -30,6 +39,10 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): with the image passed as base64 in JSON body. """ + def __init__(self) -> None: + super().__init__() + self.reference_image_pixels: int = 0 + def get_supported_openai_params(self, model: str) -> list: return AzureFoundryFluxImageGenerationConfig().get_supported_openai_params(model) @@ -107,10 +120,12 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): if len(images) > max_reference_images: 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) reference_images: Final[Mapping[str, str]] = MappingProxyType( { - "input_image" if index == 1 else f"input_image_{index}": self._convert_image_to_base64(reference_image) - for index, reference_image in enumerate(images, start=1) + "input_image" if index == 1 else f"input_image_{index}": base64.b64encode(image_bytes).decode("utf-8") + for index, image_bytes in enumerate(reference_bytes, start=1) } ) request_body: Final[dict[str, Any]] = { @@ -121,19 +136,30 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): } return request_body, [] - def _convert_image_to_base64(self, image: Any) -> str: - """Convert image file to base64 string""" + def _read_image_bytes(self, image: FileTypes | Sequence[FileTypes]) -> bytes: if isinstance(image, BufferedReader): - image_bytes = image.read() - image.seek(0) # Reset file pointer for potential reuse - elif isinstance(image, bytes): - image_bytes = image - elif hasattr(image, "read"): - image_bytes = image.read() - else: - raise ValueError(f"Unsupported image type: {type(image)}") + image_bytes: Final = image.read() + image.seek(0) + return image_bytes + if isinstance(image, bytes): + return image + if isinstance(image, IOBase): + return image.read() + raise ValueError(f"Unsupported image type: {type(image)}") - return base64.b64encode(image_bytes).decode("utf-8") + def transform_image_edit_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> ImageResponse: + try: + raw_response_json: Final = raw_response.json() + except Exception: + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) + pixels: Final = self.reference_image_pixels + hidden_params: Final = {REFERENCE_IMAGE_PIXELS_HIDDEN_PARAM: pixels} # mutable-ok: the cost path writes into it + return ImageResponse(**raw_response_json, hidden_params=hidden_params) def get_complete_url( self, @@ -165,3 +191,11 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): model=model, api_version=api_version, ) + + +def _pixel_count(image_bytes: bytes) -> int: + dimensions: Final = image_dimensions_from_bytes(image_bytes) + if dimensions is None: + return 0 + width, height = dimensions + return width * height diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 08b732197f3..106676a6e14 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -7,18 +7,26 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( calculate_image_response_cost_from_usage, resolve_image_model_info, ) +from litellm.llms.azure_ai.image_edit.flux2_transformation import REFERENCE_IMAGE_PIXELS_HIDDEN_PARAM from litellm.types.utils import ImageResponse, ModelInfo -def _input_cost_per_pixel(resolved: ModelInfo) -> float: - deployment_price: Final = _get_cost_per_unit(resolved, "input_cost_per_pixel", default_value=None) +def _pixel_rate(resolved: ModelInfo, cost_key: str) -> float: + 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("input_cost_per_pixel") or 0.0 + return shared_entry.get(cost_key) or 0.0 + + +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: + return 0.0 + return _pixel_rate(resolved, "input_cost_per_reference_pixel") * pixels def cost_calculator( @@ -48,28 +56,48 @@ def cost_calculator( if token_based_cost is not None: return token_based_cost - num_images: Final = n if n is not None else len(image_response.data or ()) - output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0 - if output_cost_per_image: - return output_cost_per_image * num_images - - if _input_cost_per_pixel(_model_info): - 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=_model_info.get("key", model), - custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, - size=pixel_size, - n=num_images, - model_info=model_info, - ) - return 0.0 + return _generated_cost( + model=model, + resolved=_model_info, + image_response=image_response, + size=size, + 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)}") + + +def _generated_cost( + model: str, + resolved: ModelInfo, + image_response: ImageResponse, + size: str | None, + n: int | None, + optional_params: Mapping[str, object] | None, + model_info: ModelInfo | None, +) -> float: + num_images: Final = n if n is not None else len(image_response.data or ()) + 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 + if not _pixel_rate(resolved, "input_cost_per_pixel"): + return 0.0 + + 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, + size=pixel_size, + n=num_images, + model_info=model_info, + ) 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 54713e1a5eb..29da0e782dc 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 @@ -14,7 +14,7 @@ from litellm.llms.azure.image_generation.http_utils import azure_deployment_imag from litellm.llms.azure_ai.image_generation.flux_transformation import ( AzureFoundryFluxImageGenerationConfig, ) -from litellm.types.utils import ImageObject, ImageResponse +from litellm.types.utils import ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params_image_gen @@ -32,6 +32,29 @@ def _flex_pixel_rate() -> float: return litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"] +def _flex_reference_rate() -> float: + return litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_reference_pixel"] + + +def _flex_edit_cost(response: ImageResponse, model_info: dict | None = None) -> float: + return CostCalculatorUtils.route_image_generation_cost_calculator( + model="FLUX.2-flex", + completion_response=response, + custom_llm_provider="azure_ai", + size="1024x1024", + call_type="image_edit", + model_info=model_info, + ) + + +def _edit_response(reference_image_pixels: object) -> ImageResponse: + return ImageResponse( + data=[ImageObject(b64_json="aW1n")], + size="1024x1024", + hidden_params={"reference_image_pixels": reference_image_pixels}, + ) + + @pytest.mark.parametrize( ("model", "provider_path"), [ @@ -131,7 +154,7 @@ 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"] * 1024 * 1024 == pytest.approx(0.05) + 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"] @@ -209,6 +232,138 @@ def test_flux2_flex_cost_accepts_lowercase_model_spelling(): assert cost == pytest.approx(_flex_pixel_rate() * 1536 * 1024 * 2) +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")) + + assert generated_only == pytest.approx(_flex_pixel_rate() * 1024 * 1024) + 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( + _flex_reference_rate() * 3 * 1024 * 1024 + ) + + +def test_flux2_flex_edit_reference_cost_scales_with_reference_pixels_not_image_count(): + two_outputs: Final = ImageResponse( + data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")], + size="1024x1024", + hidden_params={"reference_image_pixels": 1536 * 1024}, + ) + + assert _flex_edit_cost(two_outputs) == pytest.approx( + _flex_pixel_rate() * 1024 * 1024 * 2 + _flex_reference_rate() * 1536 * 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): + 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 + + +def test_flux2_flex_edit_ignores_reference_pixels_when_provider_reports_token_usage(): + response: Final = ImageResponse( + data=[ImageObject(b64_json="aW1n")], + size="1024x1024", + hidden_params={"reference_image_pixels": 1024 * 1024}, + usage=ImageUsage( + input_tokens=150, + input_tokens_details=ImageUsageInputTokensDetails(image_tokens=100, text_tokens=50), + output_tokens=1000, + total_tokens=1150, + ), + ) + token_rates: Final = { + "input_cost_per_token": 1e-05, + "input_cost_per_image_token": 2e-05, + "output_cost_per_image_token": 4e-05, + "output_cost_per_token": 4e-05, + } + + assert _flex_edit_cost(response, model_info=token_rates) == pytest.approx(50 * 1e-05 + 100 * 2e-05 + 1000 * 4e-05) + + +def test_flux2_flex_edit_reads_the_reference_rate_from_its_own_catalog_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setitem(litellm.model_cost["azure_ai/FLUX.2-flex"], "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( + _flex_pixel_rate() * 1024 * 1024 + 3e-07 * 1024 * 1024 + ) + + +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}, + ) + + assert _flex_edit_cost(large_reference_small_output) == pytest.approx( + _flex_pixel_rate() * 1024 * 1024 + _flex_reference_rate() * 2048 * 2048 + ) + + +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}) + + 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), + model_info={"input_cost_per_pixel": 2e-07, "input_cost_per_reference_pixel": 3e-07}, + ) + + assert cost == pytest.approx(2e-07 * 1024 * 1024 + 3e-07 * 1024 * 1024) + + +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}) + + 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}) + + assert cost == pytest.approx(_flex_pixel_rate() * 1024 * 1024) + + +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), + custom_llm_provider="azure_ai", + size="1024x1024", + call_type="image_edit", + model_info={"input_cost_per_pixel": 1e-07, "input_cost_per_reference_pixel": 1e-07}, + ) + + assert cost == pytest.approx(1e-07 * 1024 * 1024 * 2) + + +@pytest.mark.parametrize("model", ("FLUX-1.1-pro", "FLUX.1-Kontext-pro")) +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), + custom_llm_provider="azure_ai", + size="1024x1024", + call_type="image_edit", + ) + + assert cost == pytest.approx(litellm.model_cost[f"azure_ai/{model}"]["output_cost_per_image"]) + + def test_flux2_flex_cost_prefers_deployment_input_cost_per_pixel() -> None: response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")]) 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 5efdf1232a0..ce39513d1b2 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,5 +1,7 @@ import base64 +import io import json +import struct from collections.abc import Mapping from typing import Final @@ -14,7 +16,7 @@ from litellm.llms.azure_ai.image_edit.flux2_transformation import ( from litellm.llms.azure_ai.image_edit.transformation import ( AzureFoundryFluxImageEditConfig, ) -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler def test_azure_ai_validate_environment(): @@ -191,3 +193,110 @@ def test_flux2_image_edit_accepts_and_drops_openai_only_parameters(): ) assert optional_params == {"num_images": 1} + + +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" + + +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" + + +def _webp(width: int, height: int) -> bytes: + payload: Final = b"\x00\x00\x00\x9d\x01\x2a" + struct.pack(" tuple[float, float]: + row: Final = litellm.model_cost["azure_ai/FLUX.2-flex"] + return row["input_cost_per_pixel"], row["input_cost_per_reference_pixel"] + + +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(): + sent: Final[dict[str, object]] = {} + + def respond(request: httpx.Request) -> httpx.Response: + sent.update(json.loads(request.content)) + return _edit_ok(request) + + references: Final = (_png(1024, 1024), _jpeg(800, 600), _webp(640, 480)) + response: Final = litellm.image_edit( + model="azure_ai/FLUX.2-flex", + image=list(references), + prompt="Blend every reference", + 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() + 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["response_cost"] == pytest.approx( + generated_rate * 1024 * 1024 + reference_rate * reference_pixels + ) + + +def test_flux2_image_edit_reads_streams_once_and_still_measures_them(): + uploaded: Final = io.BytesIO(_png(2048, 2048)) + response: Final = 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(_edit_ok))), + size="1024x1024", + ) + generated_rate, reference_rate = _flex_rates() + + 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_only_the_measurable_references(): + 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"], + prompt="Blend every reference", + 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"] == 1024 * 1024 + assert response._hidden_params["response_cost"] == pytest.approx( + generated_rate * 1024 * 1024 + reference_rate * 1024 * 1024 + ) + + +async def test_flux2_aimage_edit_bills_references_like_image_edit(): + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_edit_ok)) + response: Final = await litellm.aimage_edit( + model="azure_ai/FLUX.2-flex", + image=[_png(1024, 1024), _png(1024, 1024)], + prompt="Blend both references", + api_key="test-key", + api_base="https://example.services.ai.azure.com", + client=client, + size="1024x1024", + ) + generated_rate, reference_rate = _flex_rates() + + assert response._hidden_params["reference_image_pixels"] == 2 * 1024 * 1024 + assert response._hidden_params["response_cost"] == pytest.approx( + generated_rate * 1024 * 1024 + reference_rate * 2 * 1024 * 1024 + )