Merge pull request #37729 from BerriAI/devin_ai_fal_gpt_image_2

feat(fal_ai): add gpt-image-2 image generation support

Author: mateo-berri

litellm/llms/fal_ai/image_generation/__init__.py

@@ -12,6 +12,7 @@
 from .flux_pro_v11_transformation import FalAIFluxProV11Config
 from .flux_pro_v11_ultra_transformation import FalAIFluxProV11UltraConfig
 from .flux_schnell_transformation import FalAIFluxSchnellConfig
+from .gpt_image_2_transformation import FalAIGPTImage2Config
 from .ideogram_v3_transformation import FalAIIdeogramV3Config
 from .imagen4_transformation import FalAIImagen4Config
 from .nano_banana_transformation import FalAINanoBananaConfig
@@ -27,6 +28,7 @@
     "FalAIFluxProV11Config",
     "FalAIFluxProV11UltraConfig",
     "FalAIFluxSchnellConfig",
+    "FalAIGPTImage2Config",
     "FalAIIdeogramV3Config",
     "FalAIImageGenerationConfig",
     "FalAIImagen4Config",
@@ -49,7 +51,9 @@ def get_fal_ai_image_generation_config(model: str) -> BaseImageGenerationConfig:
     model_lower: Final = model.lower()
 
     # Map model names to their corresponding configuration classes
-    if "nano-banana" in model_lower or "gemini-25-flash-image" in model_lower:
+    if "gpt-image-2" in model_lower:
+        return FalAIGPTImage2Config()
+    elif "nano-banana" in model_lower or "gemini-25-flash-image" in model_lower:
         return FalAINanoBananaConfig()
     elif "imagen4" in model_lower or "imagen-4" in model_lower:
         return FalAIImagen4Config()

litellm/llms/fal_ai/image_generation/gpt_image_2_transformation.py

@@ -0,0 +1,124 @@
+from collections.abc import Mapping
+from types import MappingProxyType
+from typing import Final
+
+from typing_extensions import ReadOnly, TypedDict
+
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
+
+from .transformation import FalAIBaseConfig
+
+
+class FalAIImageSize(TypedDict):
+    width: ReadOnly[int]
+    height: ReadOnly[int]
+
+
+SUPPORTED_OPENAI_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = (
+    "n",
+    "output_format",
+    "quality",
+    "response_format",
+    "size",
+)
+
+
+class FalAIGPTImage2Config(FalAIBaseConfig):
+    """
+    Configuration for OpenAI's GPT Image 2 served through Fal AI.
+
+    Model endpoints:
+    - openai/gpt-image-2 (text-to-image)
+    - openai/gpt-image-2/edit (editing, with optional mask)
+
+    Documentation: https://fal.ai/models/openai/gpt-image-2/api
+    """
+
+    MODEL_PREFIX: Final[str] = "openai/"
+    SUPPORTED_QUALITIES: Final[frozenset[str]] = frozenset({"auto", "low", "medium", "high"})
+    OPENAI_QUALITY_ALIASES: Final[Mapping[str, str]] = MappingProxyType({"hd": "high", "standard": "medium"})
+    PARAM_TRANSLATION: Final[Mapping[str, str]] = MappingProxyType(
+        {
+            "n": "num_images",
+            "size": "image_size",
+            "quality": "quality",
+            "output_format": "output_format",
+        }
+    )
+
+    def get_complete_url(
+        self,
+        api_base: str | None,
+        api_key: str | None,
+        model: str,
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        stream: bool | None = None,
+    ) -> str:
+        base_url: Final[str] = (api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL).rstrip("/")
+        endpoint: Final[str] = model if model.startswith(self.MODEL_PREFIX) else f"{self.MODEL_PREFIX}{model}"
+        return f"{base_url}/{endpoint}"
+
+    def get_supported_openai_params(  # mutable-ok: base class contract returns a list
+        self, model: str
+    ) -> list[OpenAIImageGenerationOptionalParams]:
+        return list(SUPPORTED_OPENAI_PARAMS)  # mutable-ok: base class contract returns a list
+
+    def map_openai_params(  # mutable-ok: base class contract returns a dict
+        self,
+        non_default_params: Mapping[str, object],
+        optional_params: Mapping[str, object],
+        model: str,
+        drop_params: bool,
+    ) -> dict:
+        unsupported_params: Final = tuple(
+            key for key in non_default_params if key not in SUPPORTED_OPENAI_PARAMS and key not in optional_params
+        )
+        if unsupported_params and not drop_params:
+            raise ValueError(
+                f"Parameters {unsupported_params} are not supported for model {model}. "
+                f"Supported parameters are {SUPPORTED_OPENAI_PARAMS}. "
+                "Set drop_params=True to drop unsupported parameters."
+            )
+        translated_params: Final[Mapping[str, object]] = MappingProxyType(
+            {
+                self.PARAM_TRANSLATION[key]: self._translate_value(key, value)
+                for key, value in non_default_params.items()
+                if key in self.PARAM_TRANSLATION and self.PARAM_TRANSLATION[key] not in optional_params
+            }
+        )
+        return {**optional_params, **translated_params}  # mutable-ok: base class contract returns a dict
+
+    def _translate_value(self, key: str, value: object) -> object:
+        if key == "size":
+            return self._map_image_size(value)
+        if key == "quality":
+            return self._map_quality(value)
+        return value
+
+    def _map_image_size(self, size: object) -> object:
+        if not isinstance(size, str) or size == "auto":
+            return size
+        try:
+            width, height = (int(part) for part in size.lower().split("x"))
+        except ValueError:
+            return size
+        image_size: Final[FalAIImageSize] = {"width": width, "height": height}
+        return image_size
+
+    def _map_quality(self, quality: object) -> object:
+        if not isinstance(quality, str):
+            return quality
+        normalized: Final[str] = self.OPENAI_QUALITY_ALIASES.get(quality, quality)
+        return normalized if normalized in self.SUPPORTED_QUALITIES else "auto"
+
+    def transform_image_generation_request(  # mutable-ok: base class contract returns a dict
+        self,
+        model: str,
+        prompt: str,
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        headers: Mapping[str, str],
+    ) -> dict:
+        return {"prompt": prompt, **optional_params}  # mutable-ok: base class contract returns a dict

litellm/model_prices_and_context_window_backup.json

@@ -17397,6 +17397,45 @@
             "/v1/images/generations"
         ]
     },
+    "fal_ai/openai/gpt-image-2": {
+        "litellm_provider": "fal_ai",
+        "metadata": {
+            "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token, so the flat output_cost_per_image here is fal's published per-image rate for a default request (quality=high, image_size=landscape_4_3 at 1024x768). Other canonical sizes at high quality: 1024x1024 $0.211, 1024x1536 $0.165, 1920x1080 $0.158, 2560x1440 $0.222, 3840x2160 $0.401"
+        },
+        "mode": "image_generation",
+        "output_cost_per_image": 0.145,
+        "source": "https://fal.ai/models/openai/gpt-image-2",
+        "supported_endpoints": [
+            "/v1/images/generations"
+        ],
+        "supports_vision": true
+    },
+    "fal_ai/gpt-image-2": {
+        "litellm_provider": "fal_ai",
+        "metadata": {
+            "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rate, see that entry for the size and quality caveat"
+        },
+        "mode": "image_generation",
+        "output_cost_per_image": 0.145,
+        "source": "https://fal.ai/models/openai/gpt-image-2",
+        "supported_endpoints": [
+            "/v1/images/generations"
+        ],
+        "supports_vision": true
+    },
+    "fal_ai/openai/gpt-image-2/edit": {
+        "litellm_provider": "fal_ai",
+        "metadata": {
+            "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Same rate as fal_ai/openai/gpt-image-2, see that entry for the size and quality caveat"
+        },
+        "mode": "image_generation",
+        "output_cost_per_image": 0.145,
+        "source": "https://fal.ai/models/openai/gpt-image-2/edit",
+        "supported_endpoints": [
+            "/v1/images/generations"
+        ],
+        "supports_vision": true
+    },
     "featherless_ai/featherless-ai/Qwerky-72B": {
         "litellm_provider": "featherless_ai",
         "max_input_tokens": 32768,

model_prices_and_context_window.json

@@ -17397,6 +17397,45 @@
             "/v1/images/generations"
         ]
     },
+    "fal_ai/openai/gpt-image-2": {
+        "litellm_provider": "fal_ai",
+        "metadata": {
+            "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token, so the flat output_cost_per_image here is fal's published per-image rate for a default request (quality=high, image_size=landscape_4_3 at 1024x768). Other canonical sizes at high quality: 1024x1024 $0.211, 1024x1536 $0.165, 1920x1080 $0.158, 2560x1440 $0.222, 3840x2160 $0.401"
+        },
+        "mode": "image_generation",
+        "output_cost_per_image": 0.145,
+        "source": "https://fal.ai/models/openai/gpt-image-2",
+        "supported_endpoints": [
+            "/v1/images/generations"
+        ],
+        "supports_vision": true
+    },
+    "fal_ai/gpt-image-2": {
+        "litellm_provider": "fal_ai",
+        "metadata": {
+            "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rate, see that entry for the size and quality caveat"
+        },
+        "mode": "image_generation",
+        "output_cost_per_image": 0.145,
+        "source": "https://fal.ai/models/openai/gpt-image-2",
+        "supported_endpoints": [
+            "/v1/images/generations"
+        ],
+        "supports_vision": true
+    },
+    "fal_ai/openai/gpt-image-2/edit": {
+        "litellm_provider": "fal_ai",
+        "metadata": {
+            "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Same rate as fal_ai/openai/gpt-image-2, see that entry for the size and quality caveat"
+        },
+        "mode": "image_generation",
+        "output_cost_per_image": 0.145,
+        "source": "https://fal.ai/models/openai/gpt-image-2/edit",
+        "supported_endpoints": [
+            "/v1/images/generations"
+        ],
+        "supports_vision": true
+    },
     "featherless_ai/featherless-ai/Qwerky-72B": {
         "litellm_provider": "featherless_ai",
         "max_input_tokens": 32768,

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

@@ -0,0 +1,146 @@
+import pytest
+
+import litellm
+from litellm.llms.fal_ai.cost_calculator import cost_calculator
+from litellm.llms.fal_ai.image_generation import (
+    FalAIGPTImage2Config,
+    FalAINanoBananaConfig,
+    get_fal_ai_image_generation_config,
+)
+from litellm.types.utils import ImageObject, ImageResponse
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        "openai/gpt-image-2",
+        "gpt-image-2",
+        "openai/gpt-image-2/edit",
+    ],
+)
+def test_gpt_image_2_config_selected(model):
+    assert isinstance(get_fal_ai_image_generation_config(model), FalAIGPTImage2Config)
+
+
+def test_nano_banana_still_routes_to_nano_banana_config():
+    assert isinstance(
+        get_fal_ai_image_generation_config("fal-ai/nano-banana"),
+        FalAINanoBananaConfig,
+    )
+
+
+@pytest.mark.parametrize(
+    "model,expected_url",
+    [
+        ("openai/gpt-image-2", "https://fal.run/openai/gpt-image-2"),
+        ("gpt-image-2", "https://fal.run/openai/gpt-image-2"),
+        ("openai/gpt-image-2/edit", "https://fal.run/openai/gpt-image-2/edit"),
+    ],
+)
+def test_get_complete_url_derives_endpoint_from_model(model, expected_url):
+    url = FalAIGPTImage2Config().get_complete_url(
+        api_base=None,
+        api_key="test-key",
+        model=model,
+        optional_params={},
+        litellm_params={},
+    )
+    assert url == expected_url
+
+
+def test_get_complete_url_respects_api_base_override():
+    url = FalAIGPTImage2Config().get_complete_url(
+        api_base="https://proxy.internal/",
+        api_key="test-key",
+        model="openai/gpt-image-2",
+        optional_params={},
+        litellm_params={},
+    )
+    assert url == "https://proxy.internal/openai/gpt-image-2"
+
+
+@pytest.mark.parametrize(
+    "non_default_params,expected",
+    [
+        ({"n": 3}, {"num_images": 3}),
+        ({"size": "1024x1536"}, {"image_size": {"width": 1024, "height": 1536}}),
+        ({"size": "auto"}, {"image_size": "auto"}),
+        ({"quality": "medium"}, {"quality": "medium"}),
+        ({"quality": "hd"}, {"quality": "high"}),
+        ({"quality": "standard"}, {"quality": "medium"}),
+        ({"quality": "nonsense"}, {"quality": "auto"}),
+        ({"output_format": "webp"}, {"output_format": "webp"}),
+        ({"response_format": "url"}, {}),
+    ],
+)
+def test_map_openai_params(non_default_params, expected):
+    assert (
+        FalAIGPTImage2Config().map_openai_params(
+            non_default_params=non_default_params,
+            optional_params={},
+            model="openai/gpt-image-2",
+            drop_params=False,
+        )
+        == expected
+    )
+
+
+def test_map_openai_params_keeps_explicit_provider_params():
+    mapped = FalAIGPTImage2Config().map_openai_params(
+        non_default_params={"n": 4, "size": "1024x1024"},
+        optional_params={"num_images": 1, "image_size": "square_hd"},
+        model="openai/gpt-image-2",
+        drop_params=False,
+    )
+    assert mapped == {"num_images": 1, "image_size": "square_hd"}
+
+
+def test_map_openai_params_raises_on_unsupported_param():
+    with pytest.raises(ValueError, match="style"):
+        FalAIGPTImage2Config().map_openai_params(
+            non_default_params={"style": "vivid"},
+            optional_params={},
+            model="openai/gpt-image-2",
+            drop_params=False,
+        )
+
+
+def test_map_openai_params_drops_unsupported_param():
+    assert (
+        FalAIGPTImage2Config().map_openai_params(
+            non_default_params={"style": "vivid"},
+            optional_params={},
+            model="openai/gpt-image-2",
+            drop_params=True,
+        )
+        == {}
+    )
+
+
+def test_transform_image_generation_request():
+    assert FalAIGPTImage2Config().transform_image_generation_request(
+        model="openai/gpt-image-2",
+        prompt="a red bicycle",
+        optional_params={"quality": "high", "num_images": 2},
+        litellm_params={},
+        headers={},
+    ) == {"prompt": "a red bicycle", "quality": "high", "num_images": 2}
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        "openai/gpt-image-2",
+        "gpt-image-2",
+        "openai/gpt-image-2/edit",
+    ],
+)
+def test_cost_calculator_uses_registry_price(model, monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
+    response = ImageResponse(
+        data=[
+            ImageObject(url="https://v3b.fal.media/files/b/one.png"),
+            ImageObject(url="https://v3b.fal.media/files/b/two.png"),
+        ]
+    )
+    assert cost_calculator(model=model, image_response=response) == pytest.approx(0.29)