diff --git a/docs/my-website/docs/providers/azure_ai_img.md b/docs/my-website/docs/providers/azure_ai_img.md
index 8e2f5226866..6434856fd2a 100644
--- a/docs/my-website/docs/providers/azure_ai_img.md
+++ b/docs/my-website/docs/providers/azure_ai_img.md
@@ -1,7 +1,7 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
-# Azure AI Image Generation
+# Azure AI Image Generation (Black Forest Labs - Flux)
Azure AI provides powerful image generation capabilities using FLUX models from Black Forest Labs to create high-quality images from text descriptions.
@@ -33,6 +33,7 @@ Get your API key and endpoint from [Azure AI Studio](https://ai.azure.com/).
|------------|-------------|----------------|
| `azure_ai/FLUX-1.1-pro` | Latest FLUX 1.1 Pro model for high-quality image generation | $0.04 |
| `azure_ai/FLUX.1-Kontext-pro` | FLUX 1 Kontext Pro model with enhanced context understanding | $0.04 |
+| `azure_ai/flux.2-pro` | FLUX 2 Pro model for next-generation image generation | $0.04 |
## Image Generation
@@ -85,6 +86,32 @@ print(response.data[0].url)
+
+
+```python showLineNumbers title="FLUX 2 Pro Image Generation"
+import litellm
+import os
+
+# Set your API credentials
+os.environ["AZURE_AI_API_KEY"] = "your-api-key-here"
+os.environ["AZURE_AI_API_BASE"] = "your-azure-ai-endpoint" # e.g., https://litellm-ci-cd-prod.services.ai.azure.com
+
+# Generate image with FLUX 2 Pro
+response = litellm.image_generation(
+ model="azure_ai/flux.2-pro",
+ prompt="A photograph of a red fox in an autumn forest",
+ api_base=os.environ["AZURE_AI_API_BASE"],
+ api_key=os.environ["AZURE_AI_API_KEY"],
+ api_version="preview",
+ size="1024x1024",
+ n=1
+)
+
+print(response.data[0].b64_json) # FLUX 2 returns base64 encoded images
+```
+
+
+
```python showLineNumbers title="Async Image Generation"
@@ -165,6 +192,15 @@ model_list:
model_info:
mode: image_generation
+ - model_name: azure-flux-2-pro
+ litellm_params:
+ model: azure_ai/flux.2-pro
+ api_key: os.environ/AZURE_AI_API_KEY
+ api_base: os.environ/AZURE_AI_API_BASE
+ api_version: preview
+ model_info:
+ mode: image_generation
+
general_settings:
master_key: sk-1234
```
diff --git a/flux2_test_image.png b/flux2_test_image.png
new file mode 100644
index 00000000000..d40fa1a65f2
Binary files /dev/null and b/flux2_test_image.png differ
diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py
index 994afa26e9c..ec4553fac4f 100644
--- a/litellm/llms/azure/azure.py
+++ b/litellm/llms/azure/azure.py
@@ -990,6 +990,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
def create_azure_base_url(
self, azure_client_params: dict, model: Optional[str]
) -> str:
+ from litellm.llms.azure_ai.image_generation import (
+ AzureFoundryFluxImageGenerationConfig,
+ )
+
api_base: str = azure_client_params.get(
"azure_endpoint", ""
) # "https://example-endpoint.openai.azure.com"
@@ -999,6 +1003,15 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
if model is None:
model = ""
+ # Handle FLUX 2 models on Azure AI which use a different URL pattern
+ # e.g., /providers/blackforestlabs/v1/flux-2-pro instead of /openai/deployments/{model}/images/generations
+ if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model):
+ return AzureFoundryFluxImageGenerationConfig.get_flux2_image_generation_url(
+ api_base=api_base,
+ model=model,
+ api_version=api_version,
+ )
+
if "/openai/deployments/" in api_base:
base_url_with_deployment = api_base
else:
diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py
index 5325f32ef63..6a1868d94cc 100644
--- a/litellm/llms/azure_ai/image_generation/flux_transformation.py
+++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
from litellm.llms.openai.image_generation import GPTImageGenerationConfig
@@ -11,4 +13,56 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
From our test suite - following GPTImageGenerationConfig is working for this model
"""
- pass
+
+ @staticmethod
+ def get_flux2_image_generation_url(
+ api_base: Optional[str],
+ model: str,
+ api_version: Optional[str],
+ ) -> str:
+ """
+ Constructs the complete URL for Azure AI FLUX 2 image generation.
+
+ FLUX 2 models on Azure AI use a different URL pattern than standard Azure OpenAI:
+ - Standard: /openai/deployments/{model}/images/generations
+ - FLUX 2: /providers/blackforestlabs/v1/flux-2-pro
+
+ Args:
+ api_base: Base URL (e.g., https://litellm-ci-cd-prod.services.ai.azure.com)
+ model: Model name (e.g., flux.2-pro)
+ api_version: API version (e.g., preview)
+
+ Returns:
+ Complete URL for the FLUX 2 image generation endpoint
+ """
+ if api_base is None:
+ raise ValueError(
+ "api_base is required for Azure AI FLUX 2 image generation"
+ )
+
+ api_base = api_base.rstrip("/")
+ api_version = api_version or "preview"
+
+ # If the api_base already contains /providers/, it's already a complete path
+ if "/providers/" in api_base:
+ if "?" in api_base:
+ return api_base
+ return f"{api_base}?api-version={api_version}"
+
+ # Construct the FLUX 2 provider path
+ # Model name flux.2-pro maps to endpoint flux-2-pro
+ return f"{api_base}/providers/blackforestlabs/v1/flux-2-pro?api-version={api_version}"
+
+ @staticmethod
+ def is_flux2_model(model: str) -> bool:
+ """
+ Check if the model is an Azure AI FLUX 2 model.
+
+ Args:
+ model: Model name (e.g., flux.2-pro, azure_ai/flux.2-pro)
+
+ Returns:
+ True if the model is a FLUX 2 model
+ """
+ model_lower = model.lower().replace(".", "-").replace("_", "-")
+ return "flux-2" in model_lower or "flux2" in model_lower
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 90b73e4709c..fb00f636409 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -4909,6 +4909,15 @@
"/v1/images/generations"
]
},
+ "azure_ai/flux.2-pro": {
+ "litellm_provider": "azure_ai",
+ "mode": "image_generation",
+ "output_cost_per_image": 0.04,
+ "source": "https://ai.azure.com/explore/models/flux.2-pro/version/1/registry/azureml-blackforestlabs",
+ "supported_endpoints": [
+ "/v1/images/generations"
+ ]
+ },
"azure_ai/Llama-3.2-11B-Vision-Instruct": {
"input_cost_per_token": 3.7e-07,
"litellm_provider": "azure_ai",
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 90b73e4709c..fb00f636409 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -4909,6 +4909,15 @@
"/v1/images/generations"
]
},
+ "azure_ai/flux.2-pro": {
+ "litellm_provider": "azure_ai",
+ "mode": "image_generation",
+ "output_cost_per_image": 0.04,
+ "source": "https://ai.azure.com/explore/models/flux.2-pro/version/1/registry/azureml-blackforestlabs",
+ "supported_endpoints": [
+ "/v1/images/generations"
+ ]
+ },
"azure_ai/Llama-3.2-11B-Vision-Instruct": {
"input_cost_per_token": 3.7e-07,
"litellm_provider": "azure_ai",
diff --git a/test_image_edit.png b/test_image_edit.png
new file mode 100644
index 00000000000..0f2de3749df
Binary files /dev/null and b/test_image_edit.png differ