From f5d62116aef3f43786e25bc9c2e68187de371930 Mon Sep 17 00:00:00 2001 From: axAilotl Date: Sun, 29 Mar 2026 14:15:25 -0400 Subject: [PATCH] Fix FAL image generation routing and size mapping --- .../fal_ai/image_generation/transformation.py | 125 ++++++++++++++++-- ...odel_prices_and_context_window_backup.json | 2 + model_prices_and_context_window.json | 2 + .../test_fal_ai_image_generation.py | 57 +++++++- ..._fal_ai_image_generation_transformation.py | 69 ++++++++++ tests/test_litellm/test_utils.py | 1 + 6 files changed, 235 insertions(+), 21 deletions(-) create mode 100644 tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index 4a0dea48a10..0ec3763dc4e 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -20,6 +20,91 @@ else: LiteLLMLoggingObj = Any +def _normalize_model_endpoint(model: str) -> str: + normalized = model.strip() + if normalized.startswith("fal_ai/"): + return normalized[len("fal_ai/") :] + return normalized + + +def _get_model_info_for_image_size_mapping(model: str) -> Optional[dict]: + import litellm + from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + + normalized_model = _normalize_model_endpoint(model) + candidate_models = [normalized_model] + + if normalized_model.startswith("fal-ai/"): + candidate_models.append(normalized_model[len("fal-ai/") :]) + + for candidate_model in candidate_models: + model_info = litellm.model_cost.get(candidate_model) + if ( + model_info is not None + and "supports_raw_image_size_dimensions" in model_info + ): + return model_info + + model_cost_map = GetModelCostMap.load_local_model_cost_map() + for candidate_model in candidate_models: + if candidate_model in model_cost_map: + return model_cost_map[candidate_model] + + return None + + +def _supports_raw_image_size_dimensions(model: str) -> bool: + model_info = _get_model_info_for_image_size_mapping(model) + if model_info is None: + return False + return bool(model_info.get("supports_raw_image_size_dimensions")) + + +def _map_openai_size_to_model_image_size(model: str, size: Any) -> Any: + if _supports_raw_image_size_dimensions(model): + if isinstance(size, dict): + width = size.get("width") + height = size.get("height") + if isinstance(width, int) and isinstance(height, int): + return f"{width}x{height}" + return size + + return _map_openai_size_to_image_size(size) + + +def _map_openai_size_to_image_size(size: Any) -> Any: + if isinstance(size, dict): + return size + + if not isinstance(size, str): + return size + + openai_size_to_image_size = { + "1024x1024": "square_hd", + "512x512": "square", + "1792x1024": "landscape_16_9", + "1024x1792": "portrait_16_9", + "1024x768": "landscape_4_3", + "768x1024": "portrait_4_3", + "1536x1024": "landscape_4_3", + "1024x1536": "portrait_4_3", + } + if size in openai_size_to_image_size: + return openai_size_to_image_size[size] + + if "x" in size: + try: + width_str, height_str = size.split("x") + return { + "width": int(width_str), + "height": int(height_str), + } + except (AttributeError, ValueError): + return size + + return size + + class FalAIBaseConfig(BaseImageGenerationConfig): """ Base configuration for Fal AI image generation models. @@ -39,17 +124,19 @@ class FalAIBaseConfig(BaseImageGenerationConfig): stream: Optional[bool] = None, ) -> str: """ - Get the complete url for the request + Get the complete url for the request. - Some providers need `model` in `api_base` + Newer Fal AI image generation models are addressed by their model path + directly under the base URL, e.g. `https://fal.run/fal-ai/flux-2`. """ complete_url: str = ( api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL ) complete_url = complete_url.rstrip("/") - if self.IMAGE_GENERATION_ENDPOINT: - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" + endpoint = self.IMAGE_GENERATION_ENDPOINT or _normalize_model_endpoint(model) + if endpoint: + complete_url = f"{complete_url}/{endpoint.lstrip('/')}" return complete_url def validate_environment( @@ -144,16 +231,26 @@ class FalAIImageGenerationConfig(FalAIBaseConfig): drop_params: bool, ) -> dict: supported_params = self.get_supported_openai_params(model) - for k in non_default_params.keys(): - if k not in optional_params.keys(): - if k in supported_params: - optional_params[k] = non_default_params[k] - elif drop_params: - pass - else: - raise ValueError( - f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." - ) + for key, value in non_default_params.items(): + if key in optional_params: + continue + + if key not in supported_params: + if drop_params: + continue + raise ValueError( + f"Parameter {key} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." + ) + + if key == "n": + optional_params["num_images"] = value + elif key == "response_format": + output_format = "png" if value in ["url", "b64_json"] else value + optional_params["output_format"] = output_format + elif key == "size": + optional_params["image_size"] = _map_openai_size_to_model_image_size( + model=model, size=value + ) return optional_params diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d781c91992d..2cd33f4cfba 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -18154,6 +18154,7 @@ "supported_endpoints": [ "/v1/images/generations" ], + "supports_raw_image_size_dimensions": true, "supports_vision": true, "supports_pdf_input": true }, @@ -18169,6 +18170,7 @@ "supported_endpoints": [ "/v1/images/generations" ], + "supports_raw_image_size_dimensions": true, "supports_vision": true, "supports_pdf_input": true }, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cfdb2911fdf..6f5868ae7b2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -18154,6 +18154,7 @@ "supported_endpoints": [ "/v1/images/generations" ], + "supports_raw_image_size_dimensions": true, "supports_vision": true, "supports_pdf_input": true }, @@ -18169,6 +18170,7 @@ "supported_endpoints": [ "/v1/images/generations" ], + "supports_raw_image_size_dimensions": true, "supports_vision": true, "supports_pdf_input": true }, diff --git a/tests/image_gen_tests/test_fal_ai_image_generation.py b/tests/image_gen_tests/test_fal_ai_image_generation.py index 44e6c34cea0..9641c9a2754 100644 --- a/tests/image_gen_tests/test_fal_ai_image_generation.py +++ b/tests/image_gen_tests/test_fal_ai_image_generation.py @@ -1,4 +1,3 @@ -import asyncio import os import sys from unittest.mock import MagicMock, patch @@ -7,13 +6,13 @@ import pytest sys.path.insert(0, os.path.abspath("../..")) -import litellm from litellm import aimage_generation @pytest.mark.parametrize( "model,expected_endpoint", [ + ("fal_ai/fal-ai/flux-2", "fal-ai/flux-2"), ("fal_ai/fal-ai/flux-pro/v1.1-ultra", "fal-ai/flux-pro/v1.1-ultra"), ("fal_ai/fal-ai/stable-diffusion-v35-medium", "fal-ai/stable-diffusion-v35-medium"), ], @@ -79,16 +78,60 @@ async def test_fal_ai_image_generation_basic(model, expected_endpoint): assert captured_url is not None assert "fal.run" in captured_url assert expected_endpoint in captured_url - print(f"Validated URL: {captured_url}") - # Validate headers assert captured_headers is not None assert "Authorization" in captured_headers assert captured_headers["Authorization"] == f"Key {test_api_key}" - print(f"Validated headers: {captured_headers}") - + # Validate request body assert captured_json_data is not None assert captured_json_data["prompt"] == test_prompt - print(f"Validated request body: {captured_json_data}") + +@pytest.mark.asyncio +async def test_fal_ai_gpt_image_uses_literal_dimension_image_size(): + """ + GPT Image 1.5 expects literal dimension strings like `1024x1024` for + `image_size`, not generic Fal image-size enums like `square_hd`. + """ + captured_json_data = None + + def capture_post_call(*args, **kwargs): + nonlocal captured_json_data + + captured_json_data = kwargs.get("json") + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = { + "images": [ + { + "url": "https://example.com/generated-image.png", + } + ], + } + + return mock_response + + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: + mock_post.side_effect = capture_post_call + + response = await aimage_generation( + model="fal_ai/fal-ai/gpt-image-1.5", + prompt="A cute baby sea otter", + api_key="test-fal-ai-key-12345", + size="1024x1024", + n=2, + response_format="url", + ) + + assert response is not None + assert captured_json_data is not None + assert captured_json_data["prompt"] == "A cute baby sea otter" + assert captured_json_data["image_size"] == "1024x1024" + assert captured_json_data["num_images"] == 2 + assert captured_json_data["output_format"] == "png" + assert "size" not in captured_json_data + assert "n" not in captured_json_data + assert "response_format" not in captured_json_data diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py new file mode 100644 index 00000000000..57408dc3ba3 --- /dev/null +++ b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py @@ -0,0 +1,69 @@ +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.fal_ai.image_generation.transformation import ( + FalAIImageGenerationConfig, +) + + +class TestFalAIImageGenerationTransformation: + def setup_method(self): + self.config = FalAIImageGenerationConfig() + + def test_get_complete_url_uses_model_path_for_generic_models(self): + result = self.config.get_complete_url( + api_base=None, + api_key="test_key", + model="fal_ai/fal-ai/flux-2", + optional_params={}, + litellm_params={}, + ) + + assert result == "https://fal.run/fal-ai/flux-2" + + def test_get_complete_url_uses_custom_base_for_generic_models(self): + result = self.config.get_complete_url( + api_base="https://custom.fal.run/", + api_key="test_key", + model="fal_ai/fal-ai/flux-2", + optional_params={}, + litellm_params={}, + ) + + assert result == "https://custom.fal.run/fal-ai/flux-2" + + def test_map_openai_params_generic_model_maps_supported_fields(self): + result = self.config.map_openai_params( + non_default_params={ + "n": 2, + "response_format": "url", + "size": "1024x1024", + }, + optional_params={}, + model="fal_ai/fal-ai/flux-2", + drop_params=False, + ) + + assert result == { + "num_images": 2, + "output_format": "png", + "image_size": "square_hd", + } + + def test_map_openai_params_gpt_image_keeps_raw_dimension_size(self): + result = self.config.map_openai_params( + non_default_params={ + "size": "1024x1024", + }, + optional_params={}, + model="fal_ai/fal-ai/gpt-image-1.5", + drop_params=False, + ) + + assert result == { + "image_size": "1024x1024", + } diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 50dc3c6c6ec..2a962c9f0cd 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -769,6 +769,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_minimal_reasoning_effort": {"type": "boolean"}, "supports_none_reasoning_effort": {"type": "boolean"}, "supports_xhigh_reasoning_effort": {"type": "boolean"}, + "supports_raw_image_size_dimensions": {"type": "boolean"}, "supports_service_tier": {"type": "boolean"}, "supports_preset": {"type": "boolean"}, "tool_use_system_prompt_tokens": {"type": "number"},