From b53f9ad65827eac2004b73d9f61eaeda5bb5981c Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 19:58:59 +0000 Subject: [PATCH] fix(fal_ai): validate returned image dimensions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/cost_calculator.py | 13 ++----- .../fal_ai/image_generation/transformation.py | 4 +- litellm/types/utils.py | 1 + litellm/utils.py | 1 + .../test_fal_ai_flux_dev_transformation.py | 38 +++++++++++++++++++ .../llms/fal_ai/test_cost_calculator.py | 27 +++++++++++++ 6 files changed, 73 insertions(+), 11 deletions(-) diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index 114bbabe4e6..fd7d82d314f 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -23,7 +23,6 @@ FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( ) _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: @@ -51,7 +50,7 @@ def _image_dimensions(image: object) -> tuple[int, int] | 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): + if type(width) is not int or width <= 0 or type(height) is not int or height <= 0: return None return width, height @@ -132,13 +131,9 @@ def cost_calculator( ) 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: Final = next( - ( - entry - for key in (f"{litellm.LlmProviders.FAL_AI.value}/{normalized_model}", normalized_model) - if (entry := _entry(key)) is not None - ), - _EMPTY_ENTRY, + model_info: Final = litellm.get_model_info( + model=normalized_model, + custom_llm_provider=litellm.LlmProviders.FAL_AI.value, ) raw_output_cost_per_image: Final = model_info.get("output_cost_per_image") output_cost_per_image: Final = ( diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index a5afdc41e9a..fd8e280da1c 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -47,8 +47,8 @@ def fal_images_to_image_objects(images: object) -> tuple[ImageObject, ...]: height: Final = image_map.get("height") content_type: Final = image_map.get("content_type") provider_specific_fields: Final[FalImageProviderSpecificFields] = { - **({"width": width} if isinstance(width, int) else {}), - **({"height": height} if isinstance(height, int) else {}), + **({"width": width} if isinstance(width, int) and type(width) is int and width > 0 else {}), + **({"height": height} if isinstance(height, int) and type(height) is int and height > 0 else {}), **({"content_type": content_type} if isinstance(content_type, str) else {}), } return ImageObject( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c661aa0839a..db9678eab43 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -315,6 +315,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_above_512k_tokens: float | None # MiniMax-M3: prompts >512K priced at 2x output 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] 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 9a80b115d4b..b54f700c67d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6087,6 +6087,7 @@ def _get_model_info_helper( output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), 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), 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/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py index 273b9151a0a..675d502240e 100644 --- a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py +++ b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py @@ -86,3 +86,41 @@ def test_flux_dev_response_omits_provider_specific_fields_when_fal_omits_metadat encoding=None, ) assert response.data[0].provider_specific_fields is None + + +@pytest.mark.parametrize( + "invalid_field, invalid_value, expected_fields", + ( + ("width", True, {"height": 768, "content_type": "image/png"}), + ("width", 0, {"height": 768, "content_type": "image/png"}), + ("width", -1, {"height": 768, "content_type": "image/png"}), + ("height", True, {"width": 1024, "content_type": "image/png"}), + ("height", 0, {"width": 1024, "content_type": "image/png"}), + ("height", -1, {"width": 1024, "content_type": "image/png"}), + ), +) +def test_flux_dev_response_drops_invalid_dimension_metadata(invalid_field, invalid_value, expected_fields): + metadata = {"width": 1024, "height": 768, "content_type": "image/png"} + metadata[invalid_field] = invalid_value + raw = httpx.Response( + 200, + json={ + "images": [ + { + "url": "https://fal.media/a.png", + **metadata, + } + ] + }, + ) + response = FalAIFluxDevConfig().transform_image_generation_response( + model="fal-ai/flux/dev", + raw_response=raw, + model_response=ImageResponse(), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert response.data[0].provider_specific_fields == expected_fields 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 e387586e417..6fb34d9f88e 100644 --- a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py +++ b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py @@ -123,6 +123,33 @@ def test_flux_dev_cost_uses_response_megapixels_per_image(): assert cost == pytest.approx(output_cost_per_pixel * 1_048_576 * (1 + 2 + 1)) +@pytest.mark.parametrize( + "dimensions", + ( + ((True, 1024),), + ((1024, 0),), + ((-1, 1024),), + ), +) +def test_flux_dev_invalid_response_dimensions_use_flat_price(dimensions): + model = "fal_ai/fal-ai/flux/dev" + cost = cost_calculator( + model=model, + image_response=_image_response_with_dimensions(dimensions), + optional_params={}, + ) + assert cost == litellm.model_cost[model]["output_cost_per_image"] * len(dimensions) + + +def test_unknown_fal_model_raises_when_flat_pricing_is_needed(): + with pytest.raises(Exception, match="isn't mapped yet"): + cost_calculator( + model="fal_ai/fal-ai/unknown-model", + image_response=_image_response(), + optional_params={}, + ) + + def test_image_edit_call_type_routes_to_fal_keyed_pricing(): model = "openai/gpt-image-2.5/flare/edit" cost = CostCalculatorUtils.route_image_generation_cost_calculator(