From 99a61c8c5538659cab20c8d94767c314daa68c2a Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 7 Jan 2026 17:31:39 +0530
Subject: [PATCH] get_flux2_image_generation_url

---
 .../image_generation/flux_transformation.py   | 56 ++++++++++++++++++-
 1 file changed, 55 insertions(+), 1 deletion(-)

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