diff --git a/docs/my-website/docs/providers/fal_ai.md b/docs/my-website/docs/providers/fal_ai.md index d42182b57a..e50ef919da 100644 --- a/docs/my-website/docs/providers/fal_ai.md +++ b/docs/my-website/docs/providers/fal_ai.md @@ -31,6 +31,7 @@ Get your API key from [fal.ai](https://fal.ai/). | Model Name | Description | Documentation | |------------|-------------|---------------| +| `fal_ai/flux/schnell` | Flux Schnell - Low-latency generation with `image_size` support | [Docs ↗](https://fal.ai/models/fal-ai/flux/schnell) | | `fal_ai/fal-ai/flux-pro/v1.1-ultra` | FLUX Pro v1.1 Ultra - High-quality image generation | [Docs ↗](https://fal.ai/models/fal-ai/flux-pro/v1.1-ultra) | | `fal_ai/fal-ai/imagen4/preview` | Google's Imagen 4 - Highest quality model | [Docs ↗](https://fal.ai/models/fal-ai/imagen4/preview) | | `fal_ai/fal-ai/recraft/v3/text-to-image` | Recraft v3 - Multiple style options | [Docs ↗](https://fal.ai/models/fal-ai/recraft/v3/text-to-image) | diff --git a/litellm/llms/fal_ai/__init__.py b/litellm/llms/fal_ai/__init__.py index 492197951e..1f4cbe0e9c 100644 --- a/litellm/llms/fal_ai/__init__.py +++ b/litellm/llms/fal_ai/__init__.py @@ -3,6 +3,7 @@ from .image_generation import ( FalAIBaseConfig, FalAIBriaConfig, FalAIFluxProV11UltraConfig, + FalAIFluxSchnellConfig, FalAIImageGenerationConfig, FalAIImagen4Config, FalAIRecraftV3Config, @@ -18,6 +19,7 @@ __all__ = [ "FalAIRecraftV3Config", "FalAIBriaConfig", "FalAIFluxProV11UltraConfig", + "FalAIFluxSchnellConfig", "FalAIStableDiffusionConfig", "get_fal_ai_image_generation_config", ] diff --git a/litellm/llms/fal_ai/image_generation/__init__.py b/litellm/llms/fal_ai/image_generation/__init__.py index 74d3b434b8..b4ae6734c6 100644 --- a/litellm/llms/fal_ai/image_generation/__init__.py +++ b/litellm/llms/fal_ai/image_generation/__init__.py @@ -4,6 +4,7 @@ from litellm.llms.base_llm.image_generation.transformation import ( from .bria_transformation import FalAIBriaConfig from .flux_pro_v11_ultra_transformation import FalAIFluxProV11UltraConfig +from .flux_schnell_transformation import FalAIFluxSchnellConfig from .imagen4_transformation import FalAIImagen4Config from .recraft_v3_transformation import FalAIRecraftV3Config from .stable_diffusion_transformation import FalAIStableDiffusionConfig @@ -16,6 +17,7 @@ __all__ = [ "FalAIRecraftV3Config", "FalAIBriaConfig", "FalAIFluxProV11UltraConfig", + "FalAIFluxSchnellConfig", "FalAIStableDiffusionConfig", ] @@ -41,6 +43,8 @@ def get_fal_ai_image_generation_config(model: str) -> BaseImageGenerationConfig: return FalAIBriaConfig() elif "flux-pro" in model_lower and "ultra" in model_lower: return FalAIFluxProV11UltraConfig() + elif "flux/schnell" in model_lower or "flux-schnell" in model_lower or "schnell" in model_lower: + return FalAIFluxSchnellConfig() elif "stable-diffusion" in model_lower: return FalAIStableDiffusionConfig() diff --git a/litellm/llms/fal_ai/image_generation/flux_schnell_transformation.py b/litellm/llms/fal_ai/image_generation/flux_schnell_transformation.py new file mode 100644 index 0000000000..ed6ed37fb4 --- /dev/null +++ b/litellm/llms/fal_ai/image_generation/flux_schnell_transformation.py @@ -0,0 +1,88 @@ +from typing import Any + +from .flux_pro_v11_ultra_transformation import FalAIFluxProV11UltraConfig + + +class FalAIFluxSchnellConfig(FalAIFluxProV11UltraConfig): + """ + Configuration for Fal AI Flux Schnell model. + + Flux Schnell shares the same response format as Flux Pro models but expects + the OpenAI `size` parameter to be translated into Fal AI's `image_size` + enum/object. + + Model endpoint: fal-ai/flux/schnell + Documentation: https://fal.ai/models/fal-ai/flux/schnell + """ + + IMAGE_GENERATION_ENDPOINT: str = "fal-ai/flux/schnell" + + _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", + } + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + supported_params = self.get_supported_openai_params(model) + + param_mapping = { + "n": "num_images", + "response_format": "output_format", + "size": "image_size", + } + + for k in non_default_params.keys(): + if k not in optional_params.keys(): + if k in supported_params: + mapped_key = param_mapping.get(k, k) + mapped_value = non_default_params[k] + + if k == "response_format": + if mapped_value in ["b64_json", "url"]: + mapped_value = "jpeg" + elif k == "size": + mapped_value = self._map_image_size(mapped_value) + + optional_params[mapped_key] = mapped_value + elif drop_params: + continue + else: + raise ValueError( + f"Parameter {k} is not supported for model {model}. " + f"Supported parameters are {supported_params}. " + "Set drop_params=True to drop unsupported parameters." + ) + + return optional_params + + def _map_image_size(self, size: Any) -> Any: + if isinstance(size, dict): + return size + + if not isinstance(size, str): + return size + + if size in self._OPENAI_SIZE_TO_IMAGE_SIZE: + return self._OPENAI_SIZE_TO_IMAGE_SIZE[size] + + if "x" in size: + try: + width_str, height_str = size.split("x") + width = int(width_str) + height = int(height_str) + return {"width": width, "height": height} + except (ValueError, AttributeError, ZeroDivisionError): + pass + + return "landscape_4_3" + diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 41e36994e0..c16980211e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -8523,6 +8523,14 @@ "/v1/images/generations" ] }, + "fal_ai/fal-ai/flux/schnell": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.003, + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "fal_ai/fal-ai/imagen4/preview": { "litellm_provider": "fal_ai", "mode": "image_generation", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 41e36994e0..c16980211e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -8523,6 +8523,14 @@ "/v1/images/generations" ] }, + "fal_ai/fal-ai/flux/schnell": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.003, + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "fal_ai/fal-ai/imagen4/preview": { "litellm_provider": "fal_ai", "mode": "image_generation", 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 949606a58a..7b5415e1d7 100644 --- a/tests/image_gen_tests/test_fal_ai_image_generation.py +++ b/tests/image_gen_tests/test_fal_ai_image_generation.py @@ -14,6 +14,7 @@ from litellm import aimage_generation "model", [ "fal_ai/fal-ai/flux-pro/v1.1-ultra", + "fal_ai/fal-ai/flux/schnell", "fal_ai/fal-ai/recraft/v3/text-to-image", "fal_ai/bria/text-to-image/3.2", "fal_ai/fal-ai/stable-diffusion-v35-medium"