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>
This commit is contained in:
kerry 2026-09-21 19:08:16 +00:00
parent adc4e6a132
commit 9b54c4b077
14 changed files with 63 additions and 65 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -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:

View file

@ -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:

View file

@ -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:

View file

@ -24921,7 +24921,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"

View file

@ -2543,7 +2543,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.
@ -2552,6 +2551,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,

View file

@ -24921,7 +24921,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"

View file

@ -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")]

View file

@ -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():