From 531d54097f7390e753d1f8d3d7febbc8950f2dee Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 19:08:16 +0000 Subject: [PATCH] refactor(fal_ai): bill flux dev per 1024x1024 megapixel and drop ImageResponse retyping Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../amazon_nova_canvas_transformation.py | 2 +- .../amazon_stability1_transformation.py | 2 +- .../amazon_stability3_transformation.py | 2 +- .../amazon_titan_transformation.py | 2 +- litellm/llms/fal_ai/cost_calculator.py | 99 +++++++++---------- .../llms/gemini/image_edit/transformation.py | 2 +- .../vertex_gemini_transformation.py | 2 +- .../vertex_imagen_transformation.py | 2 +- .../image_generation_handler.py | 2 +- ...odel_prices_and_context_window_backup.json | 2 +- litellm/types/utils.py | 5 +- model_prices_and_context_window.json | 2 +- .../providers/test_fal_ai_image_wire.py | 2 +- .../llms/fal_ai/test_cost_calculator.py | 2 +- 14 files changed, 63 insertions(+), 65 deletions(-) diff --git a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py index 32f069be6c3..ce61a6253f6 100644 --- a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py @@ -181,7 +181,7 @@ class AmazonNovaCanvasConfig: for _img in nova_response.get("images", []): openai_images.append(Image(b64_json=_img)) - model_response.data = openai_images # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = openai_images return model_response @classmethod diff --git a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py index a8ae5f6e1eb..df91e736239 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py @@ -125,7 +125,7 @@ class AmazonStabilityConfig: _image = Image(b64_json=artifact["base64"]) image_list.append(_image) - model_response.data = image_list # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = image_list return model_response diff --git a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py index 489440a7430..98e4cbbfd4d 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py @@ -101,7 +101,7 @@ class AmazonStability3Config: for _img in stability_3_response.get("images", []): openai_images.append(Image(b64_json=_img)) - model_response.data = openai_images # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = openai_images return model_response @classmethod diff --git a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py index 93b23672377..e1b06791c9d 100644 --- a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py @@ -134,7 +134,7 @@ class AmazonTitanImageGenerationConfig: _image = Image(b64_json=image) image_list.append(_image) - model_response.data = image_list # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = image_list return model_response diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index 497a792fd8a..114bbabe4e6 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -10,6 +10,7 @@ from litellm.types.utils import ImageObject, ImageResponse FAL_KEYED_PRICING_DEFAULT_QUALITY: Final[str] = "high" FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = "1024-x-768" +FAL_PIXELS_PER_MEGAPIXEL: Final[int] = 1_048_576 FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( { "square_hd": "1024-x-1024", @@ -21,10 +22,8 @@ FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( } ) -_MODEL_COST_MAP: Final[TypeAdapter[Mapping[str, Mapping[str, object]]]] = TypeAdapter( - Mapping[str, Mapping[str, object]] -) _OBJECT_MAP: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_EMPTY_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) def _keyed_size(optional_params: Mapping[str, object]) -> str | None: @@ -43,46 +42,6 @@ def _keyed_size(optional_params: Mapping[str, object]) -> str | None: return None -def _response_size(image: object) -> str | None: - if not isinstance(image, ImageObject): - return None - raw_provider_specific_fields: Final = image.provider_specific_fields - if not isinstance(raw_provider_specific_fields, Mapping): - return None - provider_specific_fields: Final = _OBJECT_MAP.validate_python(raw_provider_specific_fields) - width: Final = provider_specific_fields.get("width") - height: Final = provider_specific_fields.get("height") - if not isinstance(width, int) or not isinstance(height, int): - return None - return f"{width}-x-{height}" - - -def _keyed_quality(optional_params: Mapping[str, object]) -> str: - raw_quality: Final = optional_params.get("quality") - return raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY - - -def _keyed_cost_per_image( - model: str, - image: object, - optional_params: Mapping[str, object], - model_cost_map: Mapping[str, Mapping[str, object]], -) -> float | None: - quality: Final = _keyed_quality(optional_params) - request_size: Final = _keyed_size(optional_params) or FAL_TEXT_TO_IMAGE_DEFAULT_SIZE - sizes: Final = (_response_size(image), request_size, FAL_TEXT_TO_IMAGE_DEFAULT_SIZE) - for size in sizes: - if size is None: - continue - keyed_entry = model_cost_map.get(f"fal_ai/{quality}/{size}/{model}") - if keyed_entry is None: - continue - keyed_cost = keyed_entry.get("output_cost_per_image") - if isinstance(keyed_cost, (int, float)): - return float(keyed_cost) - return None - - def _image_dimensions(image: object) -> tuple[int, int] | None: if not isinstance(image, ImageObject): return None @@ -97,6 +56,39 @@ def _image_dimensions(image: object) -> tuple[int, int] | None: return width, height +def _response_size(image: object) -> str | None: + dimensions: Final = _image_dimensions(image) + if dimensions is None: + return None + width, height = dimensions + return f"{width}-x-{height}" + + +def _keyed_quality(optional_params: Mapping[str, object]) -> str: + raw_quality: Final = optional_params.get("quality") + return raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY + + +def _keyed_cost_per_image( + model: str, + image: object, + optional_params: Mapping[str, object], +) -> float | None: + quality: Final = _keyed_quality(optional_params) + request_size: Final = _keyed_size(optional_params) or FAL_TEXT_TO_IMAGE_DEFAULT_SIZE + sizes: Final = (_response_size(image), request_size, FAL_TEXT_TO_IMAGE_DEFAULT_SIZE) + for size in sizes: + if size is None: + continue + keyed_entry = _entry(f"fal_ai/{quality}/{size}/{model}") + if keyed_entry is None: + continue + keyed_cost = keyed_entry.get("output_cost_per_image") + if isinstance(keyed_cost, (int, float)): + return float(keyed_cost) + return None + + def _flat_cost_per_image( image: object, output_cost_per_image: float, @@ -106,8 +98,15 @@ def _flat_cost_per_image( if dimensions is None or output_cost_per_pixel is None: return output_cost_per_image width, height = dimensions - megapixels: Final = 1 if (width, height) == (1024, 1024) else ceil(width * height / 1_000_000) - return output_cost_per_pixel * 1_000_000 * megapixels + megapixels: Final = ceil(width * height / FAL_PIXELS_PER_MEGAPIXEL) + return output_cost_per_pixel * FAL_PIXELS_PER_MEGAPIXEL * megapixels + + +def _entry(key: str) -> Mapping[str, object] | None: + raw_entry: Final[object] = litellm.model_cost.get(key) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # global catalog is untyped + if not isinstance(raw_entry, Mapping): + return None + return _OBJECT_MAP.validate_python(raw_entry) def cost_calculator( @@ -123,28 +122,24 @@ def cost_calculator( normalized_model: Final = model.removeprefix(f"{litellm.LlmProviders.FAL_AI.value}/") params: Final[Mapping[str, object]] = optional_params or MappingProxyType({}) images: Final = tuple(image_response.data or ()) - raw_model_cost: Final[object] = litellm.model_cost # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # global catalog is untyped - model_cost_map: Final = _MODEL_COST_MAP.validate_python(raw_model_cost) keyed_costs: Final = tuple( _keyed_cost_per_image( model=normalized_model, image=image, optional_params=params, - model_cost_map=model_cost_map, ) for image in images ) if all(cost is not None for cost in keyed_costs): return sum(cost for cost in keyed_costs if cost is not None) - model_info_entry: Final = next( + model_info: Final = next( ( entry for key in (f"{litellm.LlmProviders.FAL_AI.value}/{normalized_model}", normalized_model) - if (entry := model_cost_map.get(key)) is not None + if (entry := _entry(key)) is not None ), - None, + _EMPTY_ENTRY, ) - model_info: Final = _OBJECT_MAP.validate_python(model_info_entry or MappingProxyType({})) raw_output_cost_per_image: Final = model_info.get("output_cost_per_image") output_cost_per_image: Final = ( float(raw_output_cost_per_image) if isinstance(raw_output_cost_per_image, (int, float)) else 0.0 diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 594b4dfd4ce..e6c22dc60b4 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -148,7 +148,7 @@ class GeminiImageEditConfig(BaseImageEditConfig): ) ) - model_response.data = cast(list[OpenAIImage], data_list) # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = cast(list[OpenAIImage], data_list) if "usageMetadata" in response_json: model_response.usage = transform_gemini_image_usage(response_json["usageMetadata"]) return model_response diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 6cf32733550..725a7f39917 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -219,7 +219,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): if (inline_data := part.get("inlineData")) and (b64_json := inline_data.get("data")) ] - model_response.data = cast(list[OpenAIImage], data_list) # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = cast(list[OpenAIImage], data_list) return model_response def _map_size_to_aspect_ratio(self, size: str) -> str: diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index d0485275bf3..c6ad5928b74 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -220,7 +220,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): ) ) - model_response.data = cast(list[OpenAIImage], data_list) # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = cast(list[OpenAIImage], data_list) return model_response def _map_size_to_aspect_ratio(self, size: str) -> str: diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index d0fb6a2f937..6a5bb484540 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -40,7 +40,7 @@ class VertexImageGeneration(VertexLLM): image_object = Image(b64_json=bytes_base64_encoded) response_data.append(image_object) - model_response.data = response_data # pyright: ignore[reportAttributeAccessIssue] # legacy OpenAI image response type + model_response.data = response_data return model_response def transform_optional_params(self, optional_params: dict | None) -> dict: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 73912668a2a..b54bb72c9df 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24886,7 +24886,7 @@ }, "mode": "image_generation", "output_cost_per_image": 0.025, - "output_cost_per_pixel": 2.5e-08, + "output_cost_per_pixel": 2.384185791015625e-08, "source": "https://fal.ai/models/fal-ai/flux/dev", "supported_endpoints": [ "/v1/images/generations" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index b22f482df72..0eae96d066d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2542,7 +2542,6 @@ from openai.types.images_response import ImagesResponse as OpenAIImageResponse class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): _hidden_params: dict = {} - data: list[ImageObject] usage: ImageUsage | None = None """ Users might use litellm with older python versions, we don't want this to break for them. @@ -2551,6 +2550,10 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): model_config = ConfigDict(extra="allow", protected_namespaces=()) + @field_serializer("data") + def _serialize_image_data(self, data: list[OpenAIImage]) -> list[dict[str, object]]: + return [image.model_dump() for image in data] + def __init__( self, created: int | None = None, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 73912668a2a..b54bb72c9df 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24886,7 +24886,7 @@ }, "mode": "image_generation", "output_cost_per_image": 0.025, - "output_cost_per_pixel": 2.5e-08, + "output_cost_per_pixel": 2.384185791015625e-08, "source": "https://fal.ai/models/fal-ai/flux/dev", "supported_endpoints": [ "/v1/images/generations" diff --git a/tests/integration/providers/test_fal_ai_image_wire.py b/tests/integration/providers/test_fal_ai_image_wire.py index e0c05e82b33..f9ceac0b037 100644 --- a/tests/integration/providers/test_fal_ai_image_wire.py +++ b/tests/integration/providers/test_fal_ai_image_wire.py @@ -163,7 +163,7 @@ def test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image(gate }, ] cost: Final = _response_cost(response) - assert cost == _approx(4 * _catalog_cost("fal_ai/fal-ai/flux/dev", "output_cost_per_pixel") * 1_000_000) + assert cost == _approx(3 * _catalog_cost("fal_ai/fal-ai/flux/dev", "output_cost_per_pixel") * 1_048_576) assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/fal-ai/flux/dev")] diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py index 05b6014c8e2..e387586e417 100644 --- a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py +++ b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py @@ -120,7 +120,7 @@ def test_flux_dev_cost_uses_response_megapixels_per_image(): optional_params={}, ) output_cost_per_pixel = litellm.model_cost[model]["output_cost_per_pixel"] - assert cost == pytest.approx(output_cost_per_pixel * 1_000_000 * (1 + 3 + 1)) + assert cost == pytest.approx(output_cost_per_pixel * 1_048_576 * (1 + 2 + 1)) def test_image_edit_call_type_routes_to_fal_keyed_pricing():