[Feat] Add Azure BFL - Flux 2 models (#18764)

* add azure_ai/flux.2-pro

* get_flux2_image_generation_url

* azure_client_params

* docs

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)
 
 </TabItem>
 
+<TabItem value="flux2" label="FLUX 2 Pro">
+
+```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
+```
+
+</TabItem>
+
 <TabItem value="async" label="Async Usage">
 
 ```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
 ```

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:

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

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

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