fix(fal_ai): price images from the dimensions fal returns

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/fal_ai/cost_calculator.py

@@ -3,7 +3,7 @@ from types import MappingProxyType
 from typing import Final
 
 import litellm
-from litellm.types.utils import ImageResponse
+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"
@@ -34,16 +34,22 @@ def _keyed_size(optional_params: Mapping[str, object]) -> str | None:
     return None
 
 
-def _keyed_cost_per_image(model: str, optional_params: Mapping[str, object] | None) -> float | None:
-    if optional_params is None:
-        return None
-    size: Final = _keyed_size(optional_params)
+def _response_size(image: ImageObject) -> str | None:
+    fields = image.provider_specific_fields or {}  # mutable-ok: empty fallback is never mutated
+    width, height = fields.get("width"), fields.get("height")
+    if (
+        isinstance(width, int)
+        and isinstance(height, int)
+        and not isinstance(width, bool)
+        and not isinstance(height, bool)
+    ):
+        return f"{width}-x-{height}"
+    return None
+
+
+def _keyed_cost_per_image(model: str, quality: str, size: str | None) -> float | None:
     if size is None:
         return None
-    raw_quality: Final = optional_params.get("quality")
-    quality: Final = (
-        raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY
-    )
     keyed_entry: Final = litellm.model_cost.get(f"fal_ai/{quality}/{size}/{model}")
     if keyed_entry is None:
         return None
@@ -64,9 +70,22 @@ def cost_calculator(
     # the proxy cost path passes the provider-prefixed model name
     model = model.removeprefix(f"{litellm.LlmProviders.FAL_AI.value}/")
     num_images: Final[int] = len(image_response.data) if image_response.data else 0
-    keyed_cost_per_image: Final = _keyed_cost_per_image(model=model, optional_params=optional_params)
-    if keyed_cost_per_image is not None:
-        return keyed_cost_per_image * num_images
+    params: Final[Mapping[str, object]] = optional_params or {}  # mutable-ok: empty fallback is never mutated
+    raw_quality: Final = params.get("quality")
+    quality: Final = (
+        raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY
+    )
+    request_size: Final = _keyed_size(params)
+    keyed_costs: Final = tuple(
+        _keyed_cost_per_image(
+            model=model,
+            quality=quality,
+            size=(_response_size(image) if isinstance(image, ImageObject) else None) or request_size,
+        )
+        for image in image_response.data or ()
+    )
+    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 = litellm.get_model_info(
         model=model,
         custom_llm_provider=litellm.LlmProviders.FAL_AI.value,

litellm/llms/fal_ai/image_generation/transformation.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Final
 
 import httpx
@@ -22,15 +23,38 @@ else:
     LiteLLMLoggingObj = Any
 
 
+def _fal_image_to_image_object(image_data: Mapping[str, object] | str) -> ImageObject:
+    if isinstance(image_data, str):
+        return ImageObject(url=image_data, b64_json=None)
+    url: Final = image_data.get("url")
+    b64_json: Final = image_data.get("b64_json")
+    width: Final = image_data.get("width")
+    height: Final = image_data.get("height")
+    if (
+        isinstance(width, int)
+        and not isinstance(width, bool)
+        and isinstance(height, int)
+        and not isinstance(height, bool)
+    ):
+        return ImageObject(
+            url=url if isinstance(url, str) else None,
+            b64_json=b64_json if isinstance(b64_json, str) else None,
+            provider_specific_fields={  # mutable-ok: provider metadata is constructed once
+                "width": width,
+                "height": height,
+            },
+        )
+    return ImageObject(
+        url=url if isinstance(url, str) else None,
+        b64_json=b64_json if isinstance(b64_json, str) else None,
+    )
+
+
 def fal_images_to_image_objects(images: object) -> tuple[ImageObject, ...]:
     if not isinstance(images, list):
         return ()
     return tuple(
-        ImageObject(url=image_data.get("url", None), b64_json=image_data.get("b64_json", None))
-        if isinstance(image_data, dict)
-        else ImageObject(url=image_data, b64_json=None)
-        for image_data in images
-        if isinstance(image_data, (dict, str))
+        _fal_image_to_image_object(image_data) for image_data in images if isinstance(image_data, (Mapping, str))
     )
 
 

tests/integration/contracts.json

@@ -169,6 +169,9 @@
     "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_row": [
       "other.provider_wire.fal_ai.gpt_image_generation_quality_size_wire_and_keyed_pricing"
     ],
+    "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_charges_the_size_fal_returned_not_the_requested_size": [
+      "other.provider_wire.fal_ai.image_pricing_uses_response_dimensions"
+    ],
     "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image": [
       "other.provider_wire.fal_ai.flux_dev_endpoint_and_per_image_pricing"
     ],

tests/integration/providers/test_fal_ai_image_wire.py

@@ -30,7 +30,7 @@ def _catalog_cost(key: str) -> float:
     return float(cost_value)
 
 
-def _image_response(urls: tuple[str, ...], prompt: str) -> bytes:
+def _image_response(urls: tuple[str, ...], prompt: str, width: int = 1024, height: int = 768) -> bytes:
     return json.dumps(
         {
             "images": [
@@ -39,8 +39,8 @@ def _image_response(urls: tuple[str, ...], prompt: str) -> bytes:
                     "content_type": "image/png",
                     "file_name": url.rsplit("/", 1)[-1],
                     "file_size": 123456,
-                    "width": 1024,
-                    "height": 768,
+                    "width": width,
+                    "height": height,
                 }
                 for url in urls
             ],
@@ -69,7 +69,7 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro
         body: Final = _JSON_OBJECT.validate_json(request.body)
         if body.get("quality") == "high":
             assert body == {"prompt": _PROMPT, "quality": "high", "image_size": {"width": 1024, "height": 1536}}
-            return Reply(body=_image_response((f"{wire_url}/files/high.png",), _PROMPT))
+            return Reply(body=_image_response((f"{wire_url}/files/high.png",), _PROMPT, width=1024, height=1536))
         assert body == {"prompt": _PROMPT, "quality": "low"}
         return Reply(body=_image_response((f"{wire_url}/files/low.png",), _PROMPT))
 
@@ -106,6 +106,53 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro
         ]
 
 
+@pytest.mark.covers("other.provider_wire.fal_ai.image_pricing_uses_response_dimensions")
+def test_fal_gpt_image_25_charges_the_size_fal_returned_not_the_requested_size(gateway: Gateway) -> None:
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        assert request.headers["authorization"] == "Key synthetic-fal-key"
+        assert request.target == "/openai/gpt-image-2.5/flare/text-to-image"
+        body: Final = _JSON_OBJECT.validate_json(request.body)
+        if body.get("image_size") is not None:
+            assert body == {"prompt": _PROMPT, "quality": "low", "image_size": {"width": 1024, "height": 1024}}
+            return Reply(body=_image_response((f"{wire_url}/files/rounded.png",), _PROMPT, width=1024, height=1536))
+        assert body == {"prompt": _PROMPT, "quality": "low"}
+        return Reply(body=_image_response((f"{wire_url}/files/default.png",), _PROMPT, width=1920, height=1080))
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        wire_url: Final = wire.url
+        model: Final = scenario.model(
+            model=f"fal_ai/{_GPT_IMAGE_MODEL}", api_base=wire.url, api_key="synthetic-fal-key"
+        )
+        requested_size_response: Final = gateway.request(
+            "POST",
+            "/v1/images/generations",
+            {"model": model, "prompt": _PROMPT, "quality": "low", "size": "1024x1024"},
+        )
+        assert requested_size_response.status_code == 200, requested_size_response.text
+        requested_size_cost: Final = _response_cost(requested_size_response)
+        assert requested_size_cost == _approx(
+            _catalog_cost("fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image")
+        )
+        assert requested_size_cost != _catalog_cost("fal_ai/low/1024-x-1024/openai/gpt-image-2.5/flare/text-to-image")
+
+        default_size_response: Final = gateway.request(
+            "POST",
+            "/v1/images/generations",
+            {"model": model, "prompt": _PROMPT, "quality": "low"},
+        )
+        assert default_size_response.status_code == 200, default_size_response.text
+        default_size_cost: Final = _response_cost(default_size_response)
+        assert default_size_cost == _approx(
+            _catalog_cost("fal_ai/low/1920-x-1080/openai/gpt-image-2.5/flare/text-to-image")
+        )
+        assert default_size_cost != _catalog_cost("fal_ai/low/1024-x-768/openai/gpt-image-2.5/flare/text-to-image")
+        assert [(request.method, request.target) for request in wire.drain()] == [
+            ("POST", "/openai/gpt-image-2.5/flare/text-to-image"),
+            ("POST", "/openai/gpt-image-2.5/flare/text-to-image"),
+        ]
+
+
 @pytest.mark.covers("other.provider_wire.fal_ai.flux_dev_endpoint_and_per_image_pricing")
 def test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image(gateway: Gateway) -> None:
     def respond(request: Request) -> Reply:

tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py

@@ -1,7 +1,5 @@
 import pytest
 
-import litellm
-from litellm.llms.fal_ai.cost_calculator import cost_calculator
 from litellm.llms.fal_ai.image_generation import (
     FalAIGPTImage2Config,
     FalAINanoBananaConfig,
@@ -11,7 +9,7 @@ from litellm.llms.fal_ai.image_generation.gpt_image_2_transformation import (
     map_gpt_image_quality,
     supported_gpt_image_qualities,
 )
-from litellm.types.utils import ImageObject, ImageResponse
+from litellm.llms.fal_ai.image_generation.transformation import fal_images_to_image_objects
 
 
 @pytest.mark.parametrize(
@@ -185,3 +183,17 @@ def test_supported_qualities_derived_from_pricing_rows(model):
 
 def test_map_gpt_image_quality_passes_through_when_no_pricing_rows():
     assert map_gpt_image_quality("xhigh", "some-new-model", {}) == "xhigh"
+
+
+def test_fal_images_to_image_objects_keeps_response_dimensions():
+    image_objects = fal_images_to_image_objects(
+        [
+            {"url": "https://example.com/dimensions.png", "width": 1024, "height": 1536},
+            "https://example.com/url.png",
+            {"url": "https://example.com/no-dimensions.png"},
+        ]
+    )
+
+    assert image_objects[0].provider_specific_fields == {"width": 1024, "height": 1536}
+    assert image_objects[1].provider_specific_fields is None
+    assert image_objects[2].provider_specific_fields is None

tests/test_litellm/llms/fal_ai/test_cost_calculator.py

@@ -1,3 +1,5 @@
+from typing import Final
+
 import pytest
 
 import litellm
@@ -15,8 +17,20 @@ def _use_local_model_cost_map(monkeypatch):
     litellm.get_model_info.cache_clear()
 
 
-def _image_response(num_images: int = 1) -> ImageResponse:
-    return ImageResponse(data=[ImageObject(url="https://example.com/img.png") for _ in range(num_images)])
+def _image_response(
+    num_images: int = 1,
+    width: int | None = None,
+    height: int | None = None,
+) -> ImageResponse:
+    provider_specific_fields: Final[dict[str, int] | None] = (
+        {"width": width, "height": height} if width is not None and height is not None else None
+    )
+    return ImageResponse(
+        data=[
+            ImageObject(url="https://example.com/img.png", provider_specific_fields=provider_specific_fields)
+            for _ in range(num_images)
+        ]
+    )
 
 
 GPT_IMAGE_25_MODELS = (
@@ -44,6 +58,89 @@ def test_gpt_image_25_quality_and_size_pick_keyed_row(model):
     assert cost == 2 * litellm.model_cost[f"fal_ai/max/3840-x-2160/{model}"]["output_cost_per_image"] > 0
 
 
+def test_gpt_image_25_response_size_wins_over_request_size():
+    model = "fal_ai/openai/gpt-image-2.5/flare/text-to-image"
+    cost = cost_calculator(
+        model=model,
+        image_response=_image_response(width=1024, height=1536),
+        optional_params={"quality": "low", "image_size": {"width": 1024, "height": 1024}},
+    )
+    response_size_cost = litellm.model_cost["fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image"][
+        "output_cost_per_image"
+    ]
+    request_size_cost = litellm.model_cost["fal_ai/low/1024-x-1024/openai/gpt-image-2.5/flare/text-to-image"][
+        "output_cost_per_image"
+    ]
+    assert cost == response_size_cost > 0
+    assert cost != request_size_cost
+
+
+def test_gpt_image_25_auto_size_uses_response_size():
+    model = "fal_ai/openai/gpt-image-2.5/flare/text-to-image"
+    cost = cost_calculator(
+        model=model,
+        image_response=_image_response(width=1920, height=1080),
+        optional_params={"quality": "low", "image_size": "auto"},
+    )
+    response_size_cost = litellm.model_cost["fal_ai/low/1920-x-1080/openai/gpt-image-2.5/flare/text-to-image"][
+        "output_cost_per_image"
+    ]
+    default_size_cost = litellm.model_cost["fal_ai/low/1024-x-768/openai/gpt-image-2.5/flare/text-to-image"][
+        "output_cost_per_image"
+    ]
+    assert cost == response_size_cost > 0
+    assert cost != default_size_cost
+
+
+def test_gpt_image_25_images_sum_response_size_keyed_rows():
+    model = "fal_ai/openai/gpt-image-2.5/flare/text-to-image"
+    cost = cost_calculator(
+        model=model,
+        image_response=ImageResponse(
+            data=[
+                ImageObject(
+                    url="https://example.com/one.png", provider_specific_fields={"width": 1024, "height": 1024}
+                ),
+                ImageObject(
+                    url="https://example.com/two.png", provider_specific_fields={"width": 1024, "height": 1536}
+                ),
+            ]
+        ),
+        optional_params={"quality": "medium"},
+    )
+    expected = sum(
+        litellm.model_cost[f"fal_ai/medium/{size}/{model.removeprefix('fal_ai/')}"]["output_cost_per_image"]
+        for size in ("1024-x-1024", "1024-x-1536")
+    )
+    assert cost == expected
+
+
+def test_gpt_image_25_missing_response_size_uses_request_size_and_none_params_uses_response_size():
+    model = "fal_ai/openai/gpt-image-2.5/flare/text-to-image"
+    request_size_cost = cost_calculator(
+        model=model,
+        image_response=_image_response(),
+        optional_params={"quality": "low", "image_size": {"width": 1024, "height": 1024}},
+    )
+    response_size_cost = cost_calculator(
+        model=model,
+        image_response=_image_response(width=1024, height=1536),
+        optional_params=None,
+    )
+    assert (
+        request_size_cost
+        == litellm.model_cost["fal_ai/low/1024-x-1024/openai/gpt-image-2.5/flare/text-to-image"][
+            "output_cost_per_image"
+        ]
+    )
+    assert (
+        response_size_cost
+        == litellm.model_cost["fal_ai/high/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image"][
+            "output_cost_per_image"
+        ]
+    )
+
+
 def test_gpt_image_25_edit_auto_size_still_honors_quality():
     model = "fal_ai/openai/gpt-image-2.5/flare/edit"
     low = cost_calculator(