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

@@ -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 = (

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(

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

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),

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

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(