Fix FAL image generation routing and size mapping

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
 

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

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

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

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",
+        }

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"},