diff --git a/litellm/llms/black_forest_labs/common_utils.py b/litellm/llms/black_forest_labs/common_utils.py index 4469a9df404..507ef17c500 100644 --- a/litellm/llms/black_forest_labs/common_utils.py +++ b/litellm/llms/black_forest_labs/common_utils.py @@ -36,4 +36,7 @@ IMAGE_GENERATION_MODELS: Dict[str, str] = { "flux-pro-1.1-ultra": "/v1/flux-pro-1.1-ultra", "flux-dev": "/v1/flux-dev", "flux-pro": "/v1/flux-pro", + # Kontext models support both text-to-image and image editing + "flux-kontext-pro": "/v1/flux-kontext-pro", + "flux-kontext-max": "/v1/flux-kontext-max", } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bb27acff72b..ce173f69aa3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7792,7 +7792,8 @@ "output_cost_per_image": 0.04, "source": "https://bfl.ai/pricing", "supported_endpoints": [ - "/v1/images/edits" + "/v1/images/edits", + "/v1/images/generations" ] }, "black_forest_labs/flux-kontext-max": { @@ -7801,7 +7802,8 @@ "output_cost_per_image": 0.08, "source": "https://bfl.ai/pricing", "supported_endpoints": [ - "/v1/images/edits" + "/v1/images/edits", + "/v1/images/generations" ] }, "black_forest_labs/flux-pro-1.0-fill": { diff --git a/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py b/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py index 95884beb2ae..a839983f8e4 100644 --- a/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py +++ b/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py @@ -178,6 +178,16 @@ class TestBlackForestLabsImageGenerationTransformation: endpoint = self.config._get_model_endpoint("flux-pro") assert endpoint == "/v1/flux-pro" + def test_get_model_endpoint_flux_kontext_pro(self): + """Test endpoint for flux-kontext-pro model (supports both generation and editing).""" + endpoint = self.config._get_model_endpoint("flux-kontext-pro") + assert endpoint == "/v1/flux-kontext-pro" + + def test_get_model_endpoint_flux_kontext_max(self): + """Test endpoint for flux-kontext-max model (supports both generation and editing).""" + endpoint = self.config._get_model_endpoint("flux-kontext-max") + assert endpoint == "/v1/flux-kontext-max" + def test_get_model_endpoint_unknown_raises(self): """Test that unknown models raise ValueError.""" with pytest.raises(ValueError, match="Unknown BFL image generation model"):