feat(azure_ai): support FLUX.2 flex images

[emerzon]

TLDR

Problem this solves:

• azure_ai/FLUX.2-flex requests are sent to the FLUX.2 Pro route, so Azure answers 404 for every generation and edit
• Flex controls (guidance, steps, width, height) and multiple reference images have no mapping
• FLUX.2 Flex has no pricing entry, so spend logs read zero

How it solves it:

• Pick the Flex or Pro Azure route from the model name; FLUX 1 requests are unchanged
• Map n and size to Azure's num_images, width, and height
• Pass numeric controls and up to ten reference images through on edits
• Bill FLUX.2 Flex per pixel from the generated image size and count, for either spelling of the model name, at $0.05 per megapixel (Azure's Global Standard Flex Megapixel meter in the Azure retail prices API, productName eq 'Azure BFL Flux Models')
• Keep FLUX.2 Pro traffic that carries OpenAI-only image fields working: user, quality, background, moderation, and output_compression are accepted and dropped, and size: "auto" sends no dimensions so Azure picks its default (1024x1024 in this run)
• Answer a size that is not WxH with a 400 naming the expected format instead of a 500

Decisions in this PR, each with the alternative not taken:

• The five OpenAI-only fields are dropped rather than forwarded as before (Azure ignored them, so forwarding was harmless but implied they did something) or rejected (a 400 on user, which most OpenAI SDK callers send, would have broken working Pro traffic)
• size: "auto" maps to no width/height rather than to a hardcoded 1024x1024, so Azure's own default applies; Flex is then billed at the cost calculator's 1024x1024 default, which matched what Azure generated
• A size that is not WxH is rejected with a 400 rather than dropped: before this PR Azure ignored the field and returned a default-size image, but silently ignoring a size is the bug this PR fixes and OpenAI rejects such values too
• Test expectation flipped: test_flux2_flex_rejects_invalid_size expected a bare ValueError for size: "large", which the proxy surfaced as a 500 APIConnectionError; it is now test_flux2_flex_rejects_invalid_size_as_bad_request and expects litellm.BadRequestError with status 400, driven through get_optional_params_image_gen the way the proxy reaches it

User Flow

Before: every FLUX.2 Flex request through Azure AI comes back 404, generation and edit alike, while FLUX.2 Pro works

1. They send POST https://litellm-domain/v1/images/generations with "model": "azure_ai/FLUX.2-flex", "prompt": "A red fox in a snowy forest at dawn", "n": 2, "size": "1536x1024", "guidance": 4.5, and "steps": 32
2. They receive HTTP 404 with litellm.NotFoundError: NotFoundError: Azure_aiException - NOT FOUND ... Received Model Group=azure_ai/FLUX.2-flex and no image
3. They send POST https://litellm-domain/v1/images/edits as multipart with the same model, two image files, a prompt, n=1, size=1024x1024, guidance=4.5, and steps=32
4. They receive the same HTTP 404 and x-litellm-response-cost: 0
5. They retry the generation as "model": "azure_ai/flux.2-flex" (lowercase) and receive the same HTTP 404
6. They send the same generation with "model": "azure_ai/FLUX.2-pro" and receive HTTP 200 with one image and x-litellm-response-cost: 0.04

After: the same requests generate and edit FLUX.2 Flex images, each priced per pixel, and FLUX.2 Pro is unchanged

1. They send POST https://litellm-domain/v1/images/generations with "model": "azure_ai/FLUX.2-flex", "prompt": "A red fox in a snowy forest at dawn", "n": 2, "size": "1536x1024", "guidance": 4.5, and "steps": 32
2. They receive HTTP 200 with one image object in data and x-litellm-response-cost: 0.0786432 (1536 x 1024 pixels at $0.05 per megapixel)
3. They send POST https://litellm-domain/v1/images/edits as multipart with the same model, two image files, a prompt, n=1, size=1024x1024, guidance=4.5, and steps=32
4. They receive HTTP 200 with the edited image in data and x-litellm-response-cost: 0.0524288
5. They retry the generation as "model": "azure_ai/flux.2-flex" (lowercase) and receive HTTP 200 with one image and x-litellm-response-cost: 0.0524288
6. They send the same generation with "model": "azure_ai/FLUX.2-pro" and receive HTTP 200 with one image and x-litellm-response-cost: 0.04

Relevant issues

Affected release

Linear ticket

Resolves LIT-5287

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• The handful of test files covering my change pass locally, e.g. uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*, make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
• My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

Live QA against Azure AI Foundry (eastus resource, GlobalStandard deployments FLUX.2-flex and FLUX.2-pro, real Azure calls). Both legs run the same script against a DB-less proxy booted with --num_workers 2 and a model_list carrying azure_ai/FLUX.2-flex, azure_ai/flux.2-flex, and azure_ai/FLUX.2-pro with api_base/api_key from the environment. Before is the merge base 3ad9a7f on port 43461, After is this PR's head c0b0ba2 on port 32230. Each response is shown as grep -iE '^HTTP|x-litellm-response-cost' over the saved headers plus a jq summary that replaces every b64_json with its character count. The same After leg and the P1, P4, P7, P9, P10, and P11 scenarios were re-run on c0b0ba2 merged into main at a262672 (port 22279) and returned the same statuses and costs

Before (3ad9a7f): Flex 404s on generation, edit, and the lowercase spelling; Pro works

$ curl -s -D gen-flex.headers -o gen-flex.json -X POST http://localhost:43461/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-flex","prompt":"A red fox in a snowy forest at dawn","n":2,"size":"1536x1024","guidance":4.5,"steps":32}'

HTTP/1.1 404 Not Found

{
  "error": {
    "message": "litellm.NotFoundError: NotFoundError: Azure_aiException - NOT FOUND\nPlease check this guide to understand why this error code might have been returned \nhttps://docs.microsoft.com/en-us/azure/machine-learning/how-to-troubleshoot-online-endpoints#http-status-codes\n. Received Model Group=azure_ai/FLUX.2-flex\nAvailable Model Group Fallbacks=None",
    "type": "invalid_request_error",
    "param": null,
    "code": "404"
  }
}

$ curl -s -D edit-flex.headers -o edit-flex.json -X POST http://localhost:43461/v1/images/edits -H 'Authorization: Bearer sk-lit5287-qa' -F model=azure_ai/FLUX.2-flex -F image=@ref-red-square.png -F image=@ref-blue-circle.png -F 'prompt=One poster combining the red square and the blue circle' -F n=1 -F size=1024x1024 -F guidance=4.5 -F steps=32

HTTP/1.1 404 Not Found
x-litellm-response-cost: 0

{
  "error": {
    "message": "litellm.NotFoundError: NotFoundError: Azure_aiException - NOT FOUND\nPlease check this guide to understand why this error code might have been returned \nhttps://docs.microsoft.com/en-us/azure/machine-learning/how-to-troubleshoot-online-endpoints#http-status-codes\n. Received Model Group=azure_ai/FLUX.2-flex\nAvailable Model Group Fallbacks=None",
    "type": "invalid_request_error",
    "param": null,
    "code": "404"
  }
}

$ curl -s -D gen-flex-lower.headers -o gen-flex-lower.json -X POST http://localhost:43461/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/flux.2-flex","prompt":"A lighthouse on a cliff at sunset","n":1,"size":"1024x1024"}'

HTTP/1.1 404 Not Found

{
  "error": {
    "message": "litellm.NotFoundError: NotFoundError: Azure_aiException - NOT FOUND\nPlease check this guide to understand why this error code might have been returned \nhttps://docs.microsoft.com/en-us/azure/machine-learning/how-to-troubleshoot-online-endpoints#http-status-codes\n. Received Model Group=azure_ai/flux.2-flex\nAvailable Model Group Fallbacks=None",
    "type": "invalid_request_error",
    "param": null,
    "code": "404"
  }
}

$ curl -s -D gen-pro.headers -o gen-pro.json -X POST http://localhost:43461/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-pro","prompt":"A lighthouse on a cliff at sunset","n":1,"size":"1024x1024"}'

HTTP/1.1 200 OK
x-litellm-response-cost: 0.04

{
  "created": 1789756502,
  "data": [
    {
      "b64_json_chars": 836936,
      "url": null
    }
  ]
}

After (c0b0ba2): Flex generates and edits at a per-pixel cost, lowercase spelling included; Pro unchanged

$ curl -s -D gen-flex.headers -o gen-flex.json -X POST http://localhost:32230/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-flex","prompt":"A red fox in a snowy forest at dawn","n":2,"size":"1536x1024","guidance":4.5,"steps":32}'

HTTP/1.1 200 OK
x-litellm-response-cost: 0.0786432

{
  "created": 1789759168,
  "data": [
    {
      "b64_json_chars": 1005824,
      "url": null
    }
  ]
}

$ curl -s -D edit-flex.headers -o edit-flex.json -X POST http://localhost:32230/v1/images/edits -H 'Authorization: Bearer sk-lit5287-qa' -F model=azure_ai/FLUX.2-flex -F image=@ref-red-square.png -F image=@ref-blue-circle.png -F 'prompt=One poster combining the red square and the blue circle' -F n=1 -F size=1024x1024 -F guidance=4.5 -F steps=32

HTTP/1.1 200 OK
x-litellm-response-cost: 0.0524288

{
  "created": 1789759199,
  "data": [
    {
      "b64_json_chars": 284096,
      "url": null
    }
  ]
}

$ curl -s -D gen-flex-lower.headers -o gen-flex-lower.json -X POST http://localhost:32230/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/flux.2-flex","prompt":"A lighthouse on a cliff at sunset","n":1,"size":"1024x1024"}'

HTTP/1.1 200 OK
x-litellm-response-cost: 0.0524288

{
  "created": 1789759221,
  "data": [
    {
      "b64_json_chars": 764248,
      "url": null
    }
  ]
}

$ curl -s -D gen-pro.headers -o gen-pro.json -X POST http://localhost:32230/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-pro","prompt":"A lighthouse on a cliff at sunset","n":1,"size":"1024x1024"}'

HTTP/1.1 200 OK
x-litellm-response-cost: 0.04

{
  "created": 1789759242,
  "data": [
    {
      "b64_json_chars": 1067260,
      "url": null
    }
  ]
}

Compatibility re-check on FLUX.2 Pro traffic: OpenAI-only fields, size auto, an invalid size, and edit form fields (Before 3ad9a7f on 43461, After c0b0ba2 on 32230)

Before

# P1: FLUX.2-pro generation with the OpenAI 'user' field
$ curl -s -D p1-pro-user.headers -o p1-pro-user.json -X POST http://localhost:43461/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-pro","prompt":"A paper boat on a pond","n":1,"size":"1024x1024","user":"qa-user-1"}'

HTTP/1.1 200 OK
x-litellm-response-cost: 0.04

{"created":1789758745,"data":[{"b64_json_chars":604792,"url":null}]}

# P4: FLUX.2-pro generation with 'size':'auto'
$ curl -s -D p4-pro-size-auto.headers -o p4-pro-size-auto.json -X POST http://localhost:43461/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-pro","prompt":"A paper boat on a pond","n":1,"size":"auto"}'

HTTP/1.1 200 OK
x-litellm-response-cost: 0.04

{"created":1789758780,"data":[{"b64_json_chars":603716,"url":null}]}

# P9: FLUX.2-pro generation with an unparseable 'size':'large'
$ curl -s -D p9-pro-size-large.headers -o p9-pro-size-large.json -X POST http://localhost:43461/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-pro","prompt":"A paper boat on a pond","n":1,"size":"large"}'

HTTP/1.1 200 OK
x-litellm-response-cost: 0.04

{"created":1789758816,"data":[{"b64_json_chars":580996,"url":null}]}

# P10: FLUX.2-pro image edit with 'user' and 'quality' form fields
$ curl -s -D p10-pro-edit-user-quality.headers -o p10-pro-edit-user-quality.json -X POST http://localhost:43461/v1/images/edits -H 'Authorization: Bearer sk-lit5287-qa' -F model=azure_ai/FLUX.2-pro -F image=@ref-red-square.png -F 'prompt=Turn the red square into a red balloon' -F n=1 -F size=1024x1024 -F user=qa-user-1 -F quality=high

HTTP/1.1 400 Bad Request
x-litellm-response-cost: 0

{"error":{"message":"litellm.UnsupportedParamsError: The following parameters are not supported for model FLUX.2-pro: quality, user. Received Model Group=azure_ai/FLUX.2-pro\nAvailable Model Group Fallbacks=None","type":"invalid_request_error","param":null,"code":"400"}}

After

# P1: FLUX.2-pro generation with the OpenAI 'user' field
$ curl -s -D p1-pro-user.headers -o p1-pro-user.json -X POST http://localhost:32230/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-pro","prompt":"A paper boat on a pond","n":1,"size":"1024x1024","user":"qa-user-1"}'

HTTP/1.1 200 OK
x-litellm-response-cost: 0.04

{"created":1789758914,"data":[{"b64_json_chars":592028,"url":null}]}

# P4: FLUX.2-pro generation with 'size':'auto'
$ curl -s -D p4-pro-size-auto.headers -o p4-pro-size-auto.json -X POST http://localhost:32230/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-pro","prompt":"A paper boat on a pond","n":1,"size":"auto"}'

HTTP/1.1 200 OK
x-litellm-response-cost: 0.04

{"created":1789758946,"data":[{"b64_json_chars":519396,"url":null}]}

# P9: FLUX.2-pro generation with an unparseable 'size':'large'
$ curl -s -D p9-pro-size-large.headers -o p9-pro-size-large.json -X POST http://localhost:32230/v1/images/generations -H 'Authorization: Bearer sk-lit5287-qa' -H 'Content-Type: application/json' -d '{"model":"azure_ai/FLUX.2-pro","prompt":"A paper boat on a pond","n":1,"size":"large"}'

HTTP/1.1 400 Bad Request

{"error":{"message":"litellm.BadRequestError: Invalid size format 'large'. Expected 'WxH', for example '1024x1024'.. Received Model Group=azure_ai/FLUX.2-pro\nAvailable Model Group Fallbacks=None","type":"invalid_request_error","param":null,"code":"400"}}

# P10: FLUX.2-pro image edit with 'user' and 'quality' form fields
$ curl -s -D p10-pro-edit-user-quality.headers -o p10-pro-edit-user-quality.json -X POST http://localhost:32230/v1/images/edits -H 'Authorization: Bearer sk-lit5287-qa' -F model=azure_ai/FLUX.2-pro -F image=@ref-red-square.png -F 'prompt=Turn the red square into a red balloon' -F n=1 -F size=1024x1024 -F user=qa-user-1 -F quality=high

HTTP/1.1 200 OK
x-litellm-response-cost: 0.04

{"created":1789759008,"data":[{"b64_json_chars":463340,"url":null}]}

Observations from the run:

• n: 2 returns one image; Azure outputs one per request (left alone)
• Generations send api-version=2025-02-01-preview, edits send preview; both work (left alone)
• A Flex 429 is retried three times before surfacing (pre-existing)
• One 404 cools the only deployment for seconds (pre-existing)
• Without dimensions Azure generated 1024x1024 on Pro and Flex (observed)
• FLUX.2 edits with user or quality were 400 before; 200 now (fixed)
• response_format on FLUX.2 is 400 on both sides (pre-existing)

Type

🆕 New Feature

Caveats (if any)

Low

• Azure returns one image per FLUX.2 request, so n above 1 yields one image and bills one image
• A FLUX.2 model name without flex in it takes the Pro route, so a future third FLUX.2 variant needs its own mapping before it works
• A FLUX.2 size that is not WxH (for example large) returned a default-size image with 200 before this PR because Azure ignored the field; it now returns 400 naming the expected format. Left as is: dropping the value instead would hide a caller's mistake the same way the old path did, and OpenAI rejects such sizes as well
• Without explicit dimensions (size: "auto" or no size), Flex is billed at 1024x1024, which is what Azure generated in this run; if Azure changes its default, billing would not follow it. Left as is: measuring the returned image would mean decoding every base64 payload in the logging path for a default Azure has not changed
• Flex and Pro edits are billed for the generated image only, while Azure also meters every reference image (Flex Ref Megapixel $0.05 per MP, Flux 2 Ref MP $0.015 per MP on Global Standard), so a Flex edit with one 1024x1024 reference logs about half of what Azure charges. Same gap as the pre-existing Pro and Kontext edit entries. Left for LIT-8140: the cost map has no key for reference-image pixels, and billing them means reading every uploaded reference's dimensions in the shared edit path
• Flex is billed at exact pixels times $0.05 per MP with no whole-megapixel rounding (1024x1024 bills $0.0524 where Azure meters whole megapixels), and the pre-existing Pro flat $0.04 per image does not follow Azure's $0.03 first MP plus $0.015 per additional MP either. Left as is: the cost map has no schema for tiered megapixel pricing, rolled into LIT-8140
• litellm.completion_cost(..., n=2) now bills n images when the caller passes n explicitly, where the Azure calculator ignored it before; the proxy never passes it and bills the returned image count (P6 above). Left as is: matches every other provider's image cost calculator
• A non-numeric string in a numeric field (n=abc, guidance=x) or more reference images than the model allows (11 on Flex, 9 on Pro) raises a plain ValueError, which the proxy surfaces as a 500, the same path the size fix moved to a 400. Left as is: no SDK client produces those inputs, and the fix is a behavioral commit that restarts the bot and QA cycle for an edge nobody hits

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
• c0b0ba2 passes /live-pr-risk

[devin-ai-integration]

PR #39424 labeled risk:medium (11 files, +367/−48, provider translation + model pricing; no sensitive area, below high thresholds). No enterprise label, so Linear routing skipped.

[codspeed]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing emerzon:litellm_azure_ai_flux_2_flex (c0b0ba2) with main (88e150b)}

[codecov]

Codecov Report

❌ Patch coverage is 94.80519% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/litellm_core_utils/llm_cost_calc/utils.py	0.00%	1 Missing ⚠️
...m/llms/azure_ai/image_edit/flux2_transformation.py	90.90%	1 Missing ⚠️
.../llms/azure_ai/image_generation/cost_calculator.py	92.30%	1 Missing ⚠️
...s/azure_ai/image_generation/flux_transformation.py	97.77%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge because no new actionable issue or outstanding prior finding remains.

Findings

1.  Model capabilities are hardcoded ▶

Summary

This PR adds Azure AI support for FLUX.2 Flex image generation and editing.

• Routes Flex and Pro requests to their model-specific Azure endpoints.
• Maps image dimensions, counts, numeric controls, and multiple edit reference images.
• Adds per-pixel FLUX.2 Flex cost calculation and model catalog entries.
• Preserves FLUX 1 behavior and compatibility with OpenAI-only FLUX.2 parameters.
• Adds focused generation, edit, routing, validation, and pricing tests.

_{Reviews (11) · Last reviewed commit: "fix(azure-ai): keep FLUX.2 tolerant of O..."}

[veria-ai]

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 1 · PR risk: 0/10

[emerzon]

Rebased onto main and verified the affected tests locally. @greptileai please review the current tip for remaining regressions

bugbot run

[emerzon]

@veria-ai please review the current tip after rebasing onto main and resolving the overlapping changes from upstream

[emerzon]

Resolved the CI compatibility findings and pushed the updated tip. @greptileai @veria-ai please review these latest changes

bugbot run

[emerzon]

@greptileai

[emerzon]

bugbot run

[emerzon]

@veria-ai

[cursor]

Stale Bugbot comment from a previous run.

[emerzon]

@greptileai

[emerzon]

bugbot run

[emerzon]

@veria-ai

[cursor]

Stale Bugbot comment from a previous run.

[mateo-berri]

@greptileai

[mateo-berri]

@greptileai

[mateo-berri]

bugbot run

[cursor]

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

^{Reviewed by Cursor Bugbot for commit c0b0ba2. Configure here.}

[greptile-apps]

Model capabilities are hardcoded

The FLUX.2 capability check and supported-parameter list are hardcoded in this transformation. This violates the repository directive that model-specific flags be stored in model_prices_and_context_window.json and read through get_model_info. As written, adding another FLUX.2 variant or changing its capabilities requires a LiteLLM code release rather than a catalog update. This repository requirement must be satisfied before merging. The same issue applies to the inline parameter list at lines 73–88.

Rule Used: What: Do not hardcode model-specific flags in the codebase. Instead, put them in model_prices_and_context_window.json and then read them in via get_model_info Why: Prevents need for users to upgrade litellm each time a new model supports this featu... (source)

[mateo-berri]

The name check predates this PR, every image provider config lists params in code, and custom deployment names have no cost map entry

[greptile-apps]

That's fair. I re-evaluated this and agree the finding is not actionable here.

is_flux2_model in litellm/llms/azure_ai/image_generation/flux_transformation.py:67-75 is endpoint routing, not a model capability flag, and custom deployment names should not need individual cost-map entries. The parameter list at :77-88 is the provider transformation contract; it also encodes mappings and behavior (n → num_images, size parsing, and dropping OpenAI-only fields) that model_prices_and_context_window.json cannot represent. Given the established provider-config pattern, this should not block the PR. No change is needed for this comment.

[mateo-berri]

@greptileai

[mateo-berri]

LGTM. Thanks for the contribution!