mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
feat(azure_ai): bill FLUX.2 edit reference images by megapixel
Azure meters a FLUX.2 edit for the pixels it generates and again for every pixel of the reference images sent with it. The proxy priced the generated image only, so an edit cost the same with one reference or three The Flux2 edit config already reads each reference's bytes to base64 them into the request. It now measures those same bytes with image_dimensions_from_bytes, sums the areas, and hands the total to the response as the reference_image_pixels hidden param. The azure_ai image cost calculator adds that count times input_cost_per_reference_pixel, resolving the rate the same way it resolves the generated-pixel rate: deployment pricing first, shared catalog second. A reference whose header cannot be read contributes zero pixels and the request goes ahead, so the fallback only ever under-bills. Provider-reported token usage still short-circuits everything. Flat-priced rows carry no reference rate, so their edits are unchanged
This commit is contained in:
Shreshth Kharbanda
parent
10323380cf
commit
a598888ab0
4 changed files with 370 additions and 44 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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")])
|
||||
|
||||
|
|
|
|||
|
|
@ -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("<HH", width, height)
|
||||
return b"RIFF" + struct.pack("<I", 12 + len(payload)) + b"WEBP" + b"VP8 " + struct.pack("<I", len(payload)) + payload
|
||||
|
||||
|
||||
def _flex_rates() -> 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
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue