Skip to content

Commit 66f7e4a

Browse files
krrishdholakiakrrishdholakia
authored
Merge pull request #14790 from eycjur/support_flux_image_edit
[Feat] Support flux image edit
2 parents a5ed344 + 9cb4099 commit 66f7e4a

File tree

6 files changed

+413
-0
lines changed

6 files changed

+413
-0
lines changed

‎docs/my-website/docs/providers/azure_ai_img_edit.md‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
import Tabs from '@theme/Tabs';
2+
import TabItem from '@theme/TabItem';
3+
4+
# Azure AI Image Editing
5+
6+
Azure AI provides powerful image editing capabilities using FLUX models from Black Forest Labs to modify existing images based on text descriptions.
7+
8+
## Overview
9+
10+
| Property | Details |
11+
|----------|---------|
12+
| Description | Azure AI Image Editing uses FLUX models to modify existing images based on text prompts. |
13+
| Provider Route on LiteLLM | `azure_ai/` |
14+
| Provider Doc | [Azure AI FLUX Models ↗](https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/black-forest-labs-flux-1-kontext-pro-and-flux1-1-pro-now-available-in-azure-ai-f/4434659) |
15+
| Supported Operations | [`/images/edits`](#image-editing) |
16+
17+
## Setup
18+
19+
### API Key & Base URL & API Version
20+
21+
```python showLineNumbers
22+
# Set your Azure AI API credentials
23+
import os
24+
os.environ["AZURE_AI_API_KEY"] = "your-api-key-here"
25+
os.environ["AZURE_AI_API_BASE"] = "your-azure-ai-endpoint" # e.g., https://your-endpoint.eastus2.inference.ai.azure.com/
26+
os.environ["AZURE_AI_API_VERSION"] = "2025-04-01-preview" # Example API version
27+
```
28+
29+
Get your API key and endpoint from [Azure AI Studio](https://ai.azure.com/).
30+
31+
## Supported Models
32+
33+
| Model Name | Description | Cost per Image |
34+
|------------|-------------|----------------|
35+
| `azure_ai/FLUX.1-Kontext-pro` | FLUX 1 Kontext Pro model with enhanced context understanding for editing | $0.04 |
36+
37+
## Image Editing
38+
39+
### Usage - LiteLLM Python SDK
40+
41+
<Tabs>
42+
<TabItem value="basic-edit" label="Basic Usage">
43+
44+
```python showLineNumbers title="Basic Image Editing"
45+
import os
46+
import base64
47+
from pathlib import Path
48+
49+
import litellm
50+
51+
# Set your API credentials
52+
os.environ["AZURE_AI_API_KEY"] = "your-api-key-here"
53+
os.environ["AZURE_AI_API_BASE"] = "your-azure-ai-endpoint"
54+
os.environ["AZURE_AI_API_VERSION"] = "2025-04-01-preview"
55+
56+
# Edit an image with a prompt
57+
response = litellm.image_edit(
58+
model="azure_ai/FLUX.1-Kontext-pro",
59+
image=open("path/to/your/image.png", "rb"),
60+
prompt="Add a winter theme with snow and cold colors",
61+
api_base=os.environ["AZURE_AI_API_BASE"],
62+
api_key=os.environ["AZURE_AI_API_KEY"],
63+
api_version=os.environ["AZURE_AI_API_VERSION"]
64+
)
65+
66+
img_base64 = response.data[0].get("b64_json")
67+
img_bytes = base64.b64decode(img_base64)
68+
path = Path("edited_image.png")
69+
path.write_bytes(img_bytes)
70+
```
71+
72+
</TabItem>
73+
74+
<TabItem value="async-edit" label="Async Usage">
75+
76+
```python showLineNumbers title="Async Image Editing"
77+
import os
78+
import base64
79+
from pathlib import Path
80+
81+
import litellm
82+
import asyncio
83+
84+
# Set your API credentials
85+
os.environ["AZURE_AI_API_KEY"] = "your-api-key-here"
86+
os.environ["AZURE_AI_API_BASE"] = "your-azure-ai-endpoint"
87+
os.environ["AZURE_AI_API_VERSION"] = "2025-04-01-preview"
88+
89+
async def edit_image():
90+
# Edit image asynchronously
91+
response = await litellm.aimage_edit(
92+
model="azure_ai/FLUX.1-Kontext-pro",
93+
image=open("path/to/your/image.png", "rb"),
94+
prompt="Make this image look like a watercolor painting",
95+
api_base=os.environ["AZURE_AI_API_BASE"],
96+
api_key=os.environ["AZURE_AI_API_KEY"],
97+
api_version=os.environ["AZURE_AI_API_VERSION"]
98+
)
99+
img_base64 = response.data[0].get("b64_json")
100+
img_bytes = base64.b64decode(img_base64)
101+
path = Path("async_edited_image.png")
102+
path.write_bytes(img_bytes)
103+
104+
# Run the async function
105+
asyncio.run(edit_image())
106+
```
107+
108+
</TabItem>
109+
110+
<TabItem value="advanced-edit" label="Advanced Parameters">
111+
112+
```python showLineNumbers title="Advanced Image Editing with Parameters"
113+
import os
114+
import base64
115+
from pathlib import Path
116+
117+
import litellm
118+
119+
# Set your API credentials
120+
os.environ["AZURE_AI_API_KEY"] = "your-api-key-here"
121+
os.environ["AZURE_AI_API_BASE"] = "your-azure-ai-endpoint"
122+
os.environ["AZURE_AI_API_VERSION"] = "2025-04-01-preview"
123+
124+
# Edit image with additional parameters
125+
response = litellm.image_edit(
126+
model="azure_ai/FLUX.1-Kontext-pro",
127+
image=open("path/to/your/image.png", "rb"),
128+
prompt="Add magical elements like floating crystals and mystical lighting",
129+
api_base=os.environ["AZURE_AI_API_BASE"],
130+
api_key=os.environ["AZURE_AI_API_KEY"],
131+
api_version=os.environ["AZURE_AI_API_VERSION"],
132+
n=1
133+
)
134+
img_base64 = response.data[0].get("b64_json")
135+
img_bytes = base64.b64decode(img_base64)
136+
path = Path("advanced_edited_image.png")
137+
path.write_bytes(img_bytes)
138+
```
139+
140+
</TabItem>
141+
</Tabs>
142+
143+
### Usage - LiteLLM Proxy Server
144+
145+
#### 1. Configure your config.yaml
146+
147+
```yaml showLineNumbers title="Azure AI Image Editing Configuration"
148+
model_list:
149+
- model_name: azure-flux-kontext-edit
150+
litellm_params:
151+
model: azure_ai/FLUX.1-Kontext-pro
152+
api_key: os.environ/AZURE_AI_API_KEY
153+
api_base: os.environ/AZURE_AI_API_BASE
154+
api_version: os.environ/AZURE_AI_API_VERSION
155+
model_info:
156+
mode: image_edit
157+
158+
general_settings:
159+
master_key: sk-1234
160+
```
161+
162+
#### 2. Start LiteLLM Proxy Server
163+
164+
```bash showLineNumbers title="Start LiteLLM Proxy Server"
165+
litellm --config /path/to/config.yaml
166+
167+
# RUNNING on http://0.0.0.0:4000
168+
```
169+
170+
#### 3. Make image editing requests with OpenAI Python SDK
171+
172+
<Tabs>
173+
<TabItem value="openai-edit-sdk" label="OpenAI SDK">
174+
175+
```python showLineNumbers title="Azure AI Image Editing via Proxy - OpenAI SDK"
176+
from openai import OpenAI
177+
178+
# Initialize client with your proxy URL
179+
client = OpenAI(
180+
base_url="http://localhost:4000", # Your proxy URL
181+
api_key="sk-1234" # Your proxy API key
182+
)
183+
184+
# Edit image with FLUX Kontext Pro
185+
response = client.images.edit(
186+
model="azure-flux-kontext-edit",
187+
image=open("path/to/your/image.png", "rb"),
188+
prompt="Transform this image into a beautiful oil painting style",
189+
)
190+
191+
img_base64 = response.data[0].b64_json
192+
img_bytes = base64.b64decode(img_base64)
193+
path = Path("proxy_edited_image.png")
194+
path.write_bytes(img_bytes)
195+
```
196+
197+
</TabItem>
198+
199+
<TabItem value="litellm-edit-sdk" label="LiteLLM SDK">
200+
201+
```python showLineNumbers title="Azure AI Image Editing via Proxy - LiteLLM SDK"
202+
import litellm
203+
204+
# Edit image through proxy
205+
response = litellm.image_edit(
206+
model="litellm_proxy/azure-flux-kontext-edit",
207+
image=open("path/to/your/image.png", "rb"),
208+
prompt="Add a mystical forest background with magical creatures",
209+
api_base="http://localhost:4000",
210+
api_key="sk-1234"
211+
)
212+
213+
img_base64 = response.data[0].b64_json
214+
img_bytes = base64.b64decode(img_base64)
215+
path = Path("proxy_edited_image.png")
216+
path.write_bytes(img_bytes)
217+
```
218+
219+
</TabItem>
220+
221+
<TabItem value="curl-edit" label="cURL">
222+
223+
```bash showLineNumbers title="Azure AI Image Editing via Proxy - cURL"
224+
curl --location 'http://localhost:4000/v1/images/edits' \
225+
--header 'Authorization: Bearer sk-1234' \
226+
--form 'model="azure-flux-kontext-edit"' \
227+
--form 'prompt="Convert this image to a vintage sepia tone with old-fashioned effects"' \
228+
--form 'image=@"path/to/your/image.png"'
229+
```
230+
231+
</TabItem>
232+
</Tabs>
233+
234+
## Supported Parameters
235+
236+
Azure AI Image Editing supports the following OpenAI-compatible parameters:
237+
238+
| Parameter | Type | Description | Default | Example |
239+
|-----------|------|-------------|---------|---------|
240+
| `image` | file | The image file to edit | Required | File object or binary data |
241+
| `prompt` | string | Text description of the desired changes | Required | `"Add snow and winter elements"` |
242+
| `model` | string | The FLUX model to use for editing | Required | `"azure_ai/FLUX.1-Kontext-pro"` |
243+
| `n` | integer | Number of edited images to generate (You can specify only 1) | `1` | `1` |
244+
| `api_base` | string | Your Azure AI endpoint URL | Required | `"https://your-endpoint.eastus2.inference.ai.azure.com/"` |
245+
| `api_key` | string | Your Azure AI API key | Required | Environment variable or direct value |
246+
| `api_version` | string | API version for Azure AI | Required | `"2025-04-01-preview"` |
247+
248+
## Getting Started
249+
250+
1. Create an account at [Azure AI Studio](https://ai.azure.com/)
251+
2. Deploy a FLUX model in your Azure AI Studio workspace
252+
3. Get your API key and endpoint from the deployment details
253+
4. Set your `AZURE_AI_API_KEY`, `AZURE_AI_API_BASE` and `AZURE_AI_API_VERSION` environment variables
254+
5. Prepare your source image
255+
6. Use `litellm.image_edit()` to modify your images with text instructions
256+
257+
## Additional Resources
258+
259+
- [Azure AI Studio Documentation](https://docs.microsoft.com/en-us/azure/ai-services/)
260+
- [FLUX Models Announcement](https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/black-forest-labs-flux-1-kontext-pro-and-flux1-1-pro-now-available-in-azure-ai-f/4434659)

‎docs/my-website/sidebars.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -383,6 +383,7 @@ const sidebars = {
383383
items: [
384384
"providers/azure_ai",
385385
"providers/azure_ai_img",
386+
"providers/azure_ai_img_edit",
386387
]
387388
},
388389
{

‎litellm/llms/azure_ai/image_edit/__init__.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
2+
3+
from .transformation import AzureFoundryFluxImageEditConfig
4+
5+
__all__ = ["AzureFoundryFluxImageEditConfig"]
6+
7+
8+
def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig:
9+
model = model.lower()
10+
model = model.replace("-", "")
11+
model = model.replace("_", "")
12+
if model == "" or "flux" in model: # empty model is flux
13+
return AzureFoundryFluxImageEditConfig()
14+
else:
15+
raise ValueError(f"Model {model} is not supported for Azure AI image editing.")

‎litellm/llms/azure_ai/image_edit/transformation.py‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
from typing import Optional
2+
3+
import httpx
4+
5+
import litellm
6+
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
7+
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
8+
from litellm.secret_managers.main import get_secret_str
9+
from litellm.utils import _add_path_to_api_base
10+
11+
12+
class AzureFoundryFluxImageEditConfig(OpenAIImageEditConfig):
13+
"""
14+
Azure AI Foundry FLUX image edit config
15+
16+
Supports FLUX models including FLUX-1-kontext-pro for image editing.
17+
18+
Azure AI Foundry FLUX models handle image editing through the /images/edits endpoint,
19+
same as standard Azure OpenAI models. The request format uses multipart/form-data
20+
with image files and prompt.
21+
"""
22+
23+
def validate_environment(
24+
self,
25+
headers: dict,
26+
model: str,
27+
api_key: Optional[str] = None,
28+
) -> dict:
29+
"""
30+
Validate Azure AI Foundry environment and set up authentication
31+
Uses Api-Key header format
32+
"""
33+
api_key = AzureFoundryModelInfo.get_api_key(api_key)
34+
35+
if not api_key:
36+
raise ValueError(
37+
f"Azure AI API key is required for model {model}. Set AZURE_AI_API_KEY environment variable or pass api_key parameter."
38+
)
39+
40+
headers.update(
41+
{
42+
"Api-Key": api_key, # Azure AI Foundry uses Api-Key header format
43+
}
44+
)
45+
return headers
46+
47+
def get_complete_url(
48+
self,
49+
model: str,
50+
api_base: Optional[str],
51+
litellm_params: dict,
52+
) -> str:
53+
"""
54+
Constructs a complete URL for Azure AI Foundry image edits API request.
55+
56+
Azure AI Foundry FLUX models handle image editing through the /images/edits
57+
endpoint.
58+
59+
Args:
60+
- model: Model name (deployment name for Azure AI Foundry)
61+
- api_base: Base URL for Azure AI endpoint
62+
- litellm_params: Additional parameters including api_version
63+
64+
Returns:
65+
- Complete URL for the image edits endpoint
66+
"""
67+
api_base = AzureFoundryModelInfo.get_api_base(api_base)
68+
69+
if api_base is None:
70+
raise ValueError(
71+
"Azure AI API base is required. Set AZURE_AI_API_BASE environment variable or pass api_base parameter."
72+
)
73+
74+
api_version = (litellm_params.get("api_version") or litellm.api_version
75+
or get_secret_str("AZURE_AI_API_VERSION")
76+
)
77+
if api_version is None:
78+
# API version is mandatory for Azure AI Foundry
79+
raise ValueError(
80+
"Azure API version is required. Set AZURE_AI_API_VERSION environment variable or pass api_version parameter."
81+
)
82+
83+
# Add the path to the base URL using the model as deployment name
84+
# Azure AI Foundry FLUX models use /images/edits for editing
85+
if "/openai/deployments/" in api_base:
86+
new_url = _add_path_to_api_base(
87+
api_base=api_base,
88+
ending_path="/images/edits",
89+
)
90+
else:
91+
new_url = _add_path_to_api_base(
92+
api_base=api_base,
93+
ending_path=f"/openai/deployments/{model}/images/edits",
94+
)
95+
96+
# Use the new query_params dictionary
97+
final_url = httpx.URL(new_url).copy_with(params={"api-version": api_version})
98+
99+
return str(final_url)

‎litellm/utils.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7532,6 +7532,12 @@ def get_provider_image_edit_config(
75327532
)
75337533

75347534
return RecraftImageEditConfig()
7535+
elif LlmProviders.AZURE_AI == provider:
7536+
from litellm.llms.azure_ai.image_edit import (
7537+
get_azure_ai_image_edit_config,
7538+
)
7539+
7540+
return get_azure_ai_image_edit_config(model)
75357541
elif LlmProviders.LITELLM_PROXY == provider:
75367542
from litellm.llms.litellm_proxy.image_edit.transformation import (
75377543
LiteLLMProxyImageEditConfig,

0 commit comments

Comments
 (0)