ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add fal-ai/flux/schnell support (#16580)

Browse Source
This commit is contained in:
Sameer Kankute 2025-11-14 12:01:31 +05:30 committed by GitHub
parent 266744a5bd
commit 13993d6ea3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 112 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/my-website/docs/providers/fal_ai.md
View File

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

2
litellm/llms/fal_ai/__init__.py
View File

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

4
litellm/llms/fal_ai/image_generation/__init__.py
View File

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

88
litellm/llms/fal_ai/image_generation/flux_schnell_transformation.py Normal file
View File

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

8
litellm/model_prices_and_context_window_backup.json
View File

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

8
model_prices_and_context_window.json
View File

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

1
tests/image_gen_tests/test_fal_ai_image_generation.py
View File

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