fix(azure_ai): bill FLUX.2-flex edit reference megapixels at Azure's rate

[devin-ai-integration]

TLDR

Problem this solves:

• Azure meters FLUX.2 Flex edits per reference megapixel too
• The proxy billed only the generated image, so spend ran low
• A two reference 1024x1024 edit was billed at one third of its Azure cost
• The Flex per pixel rate assumed 1 MP = 1,000,000 pixels; Azure and BFL use 1024x1024

How it solves it:

• Read each reference image header (PNG, JPEG, WebP, GIF, BMP) before the upload stream is consumed
• JPEG headers are walked with a loop bounded at 1024 segments and a 16 MiB offset, never recursion, so ICC profiles split across hundreds of APP2 chunks still measure
• Measurement never raises: closed or read only streams and headers with non positive dimensions (a negative width BMP) count as unmeasured instead of leaking an exception into every provider's edit path or a negative pixel total
• Sum reference pixels into one hidden integer on the image response
• Azure AI image cost adds pixels times input_cost_per_reference_pixel, read from the resolved deployment model info first and the shared catalog second, the same lookup the generated pixels use, so custom priced and unlisted deployments bill their own reference rate
• azure_ai/FLUX.2-flex gets that rate, and both rates become 0.05 / 1048576 per pixel
• Provider usage blocks still win, so token-reporting models are unchanged
• The Terraform provider's litellm_model resource gains input_cost_per_reference_pixel next to input_cost_per_pixel, forwarded whenever it is set in config so an explicit 0 override reaches the deployment, and the dashboard schema.d.ts is regenerated
• ModelGroupInfo carries input_cost_per_reference_pixel next to input_cost_per_pixel, so /model_group/info reports a configured reference rate instead of dropping it

User Flow

Before: a developer editing an image through FLUX.2 Flex sees the same spend whether they send one reference or two, and Azure bills them two or three times what the proxy shows

1. They send POST https://litellm-domain/v1/images/edits as multipart with model=azure_ai/FLUX.2-flex, prompt=make it a watercolor, size=1024x1024 and one 1024x1024 PNG as image
2. The response is 200 with the edited image and the header x-litellm-response-cost: 0.0524288
3. They repeat the request with the same PNG twice as image[] and get x-litellm-response-cost: 0.0524288 again
4. GET https://litellm-domain/spend/logs?request_id= shows spend: 0.0524288 for both calls, while Azure Cost Management shows 1.0 output MP plus 1.0 and 2.0 reference MP at $0.05 each ($0.10 and $0.15)

After: the same requests come back with the reference megapixels priced in at Azure's megapixel definition, so spend matches the Azure meter to the cent

1. They send POST https://litellm-domain/v1/images/edits as multipart with model=azure_ai/FLUX.2-flex, prompt=make it a watercolor, size=1024x1024 and one 1024x1024 PNG as image
2. The response is 200 with the edited image and the header x-litellm-response-cost: 0.1
3. They repeat the request with the same PNG twice as image[] and get x-litellm-response-cost: 0.15
4. GET https://litellm-domain/spend/logs?request_id= shows spend: 0.1 and spend: 0.15, matching the Azure meters

Relevant issues

Affected release

Linear ticket

Resolves LIT-8140

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.). The one red job, rust-wheel, fails on tests/test_litellm_rust/ocr/test_callbacks.py::test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks with the identical DID NOT RAISE on the merge base b0ac23d3 (main run 35807522715); nothing here touches OCR
• 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

Two proxies ran side by side against the real Azure AI Foundry FLUX.2-flex deployment with LITELLM_LOCAL_MODEL_COST_MAP=True: the merge base on port 4100 and this branch on port 4200. Every scenario is one curl, for example

curl -s -D - -o body.json http://localhost:4200/v1/images/edits -H "Authorization: Bearer $KEY" \
  -F model=azure_ai/FLUX.2-flex -F "prompt=make it a watercolor" -F size=1024x1024 \
  -F 'image[]=@img/png_1024.png' -F 'image[]=@img/png_1024.png' -F 'image[]=@img/png_1024.png'

The x-litellm-response-cost header from each call, base then branch. Expected is 0.05 * (output MP + reference MP) with 1 MP = 1024x1024

scenario	base	branch	expected
generation, no reference, 1024x1024	0.0524288	0.05	0.05
1 ref PNG 1024x1024	0.0524288	0.1	0.10
2 refs PNG 1024x1024	0.0524288	0.15000000000000002	0.15
3 refs PNG 1024x1024	0.0524288	0.2	0.20
1 ref JPEG baseline	400 from Azure on that call, 0	0.1	0.10
1 ref JPEG progressive (SOF2)	0.0524288	0.1	0.10
1 ref JPEG with EXIF APP1	0.0524288	0.1	0.10
1 ref WebP lossy (VP8)	0.0524288	0.1	0.10
1 ref WebP lossless (VP8L)	0.0524288	0.1	0.10
1 ref WebP alpha (VP8X)	0.0524288	0.1	0.10
1 ref PNG 768x1024	0.0524288	0.08750000000000001	0.0875
1 ref PNG 1536x1536, output 1024x768	0.0393216	0.15000000000000002	0.15
3 refs PNG + JPEG + WebP	0.0524288	0.2	0.20
1 ref GIF 1024x1024	0.0524288	0.1	0.10
1 ref BMP 1024x1024	0.0524288	0.1	0.10
1 ref PNG with truncated header	400 from Azure, 0	400 from Azure, 0	0
1 ref JPEG with 300 APP2 segments before SOF	0.0524288	0.1	0.10
1 ref BMP with negative width	400 from Azure, 0	400 from Azure, 0	0
2 refs PNG + negative width BMP	500 from Azure, 0	500 from Azure, 0	0

The whole matrix was rerun on the final head 1eaddc47e9 against the same merge base with byte identical fixtures; every branch value matched the expected column and every base value reproduced the under billing. Azure itself rejects the malformed BMP, and on the branch that request no longer produced a negative reference total on the way out

The base column shows both bugs at once: every edit costs the same as a generation because references were never counted, and 1024x1024 is priced at 0.0524288 because the old rate divided by 1,000,000 instead of 1,048,576. The base JPEG baseline 400 was Azure rejecting that one upload ("Unable to extract dimensions from input image"); the identical file was accepted on the branch run and on a retry, so it is provider side flakiness, not a difference in the proxy. The trailing digits in 0.15000000000000002 and 0.08750000000000001 are the usual float rendering of 0.05 * 3 and 0.05 * 1.75; spend is already stored as a float

Azure Cost Management for these calls posts the next day. The earlier run on 2026-09-22 (4 calls, all 1024x1024) reconciled exactly: Flex Megapixel 4.0 MP at $0.20 and Flex Ref Megapixel 6.0 MP at $0.30, for $0.50 against $0.21 of proxy spend on the base

Pricing source

Both input_cost_per_pixel and input_cost_per_reference_pixel on azure_ai/FLUX.2-flex are 4.76837158203125e-08, which is 0.05 / 1048576. The Azure retail prices API (filter productName eq 'Azure BFL Flux Models', checked 2026-09-23) lists meterName: Flex Megapixel and meterName: Flex Ref Megapixel at retailPrice: 0.05 per megapixel. Azure Cost Management usage for the QA Foundry resources on 2026-09-22 records whole 1024x1024 images as whole megapixels (UsageQuantity 4.0 on Flex Megapixel and 6.0 on Flex Ref Megapixel), which is only consistent with 1 MP = 1,048,576 pixels, the definition BFL also publishes. The previous 5e-08 treated a megapixel as 1,000,000 pixels and over-billed every Flex generation by 4.86%

Regression tests and mutation

On the merge base the extended edit tests fail with KeyError: 'reference_pixels' because the hidden key does not exist, and the schema invariant fails because azure_ai/FLUX.2-flex has no input_cost_per_reference_pixel

Mutations at the final head, each run against the mapped test files and then restored and verified with cmp -s: removing the outer except in measure_reference_image turns closed_and_read_only_streams and restores_position_on_read_error red; dropping the <= 0 dimension guard turns bmp-negative-width, bmp-zero-width and invalid_dimensions red; putting _MAX_JPEG_SEGMENTS back to 64 turns jpeg-300-app2-segments red; counting fill bytes as hops turns jpeg-1100-fill-bytes-before-sof red; skipping past a marker byte after a lone fill byte (the old 1 + int(marker[2] == 0xFF) + int(marker[3] == 0xFF) skip) turns jpeg-fill-before-large-segment red; skipping at most 3 fill bytes per hop instead of the whole run turns jpeg-200k-fill-bytes-before-sof red; consuming a whole 64 KiB fill chunk instead of leaving the last byte as the marker prefix turns jpeg-65535-fill-bytes-before-sof red; removing the 16 MiB offset cap turns jpeg-sof-beyond-16mib red; narrowing _JPEG_SOF_MARKERS turns 11 of the 13 every_jpeg_sof_marker cases red; dropping the C8 and CC exclusions turns both skips_non_sof_c_range_markers cases red; changing is not None to truthiness in _pixel_rate turns explicit_zero_deployment_rate red; weakening isinstance(pixels, int) turns both non integer reference_pixels cases red; replacing or 0.0 with .get(key, 0.0) turns catalog_rate_is_unset red

Earlier mutations, still killed: removing the _with_reference_pixels(...) wrapper from both handler return sites turns the edit tests red; removing the JPEG 0xFF padding skip turns jpeg-sof-after-0xff-padding red; lowering _MAX_JPEG_SEGMENTS to 63 turns jpeg-sof-at-segment-64-boundary red; dropping abs() on the BMP height turns bmp-negative-height red; swapping the GIF width and height offsets turns gif89a and gif87a red; dropping the 12 byte BITMAPCOREHEADER branch turns bmp-coreheader red. Reverting _reference_cost to read the shared catalog instead of the resolved model info turns test_flux2_flex_reference_cost_prefers_deployment_rate and test_unlisted_azure_ai_model_bills_deployment_reference_rate red. Restoring each gives a green run. The segment_length < 2 guard cannot be killed by any input because a sub 2 length can never land the next read on a valid 0xFF marker; it stays as a loop safety net

Adversarial review and formal check

A fresh context review council (six investigation lenses, three taxonomy auditors, four independent reviewers, and reproduce, falsify and artifact check verifiers on every verdict flipping finding) ran against the pre push head. It refuted the universal safety claim with three executable counterexamples that the final head fixes and pins with the tests above: a closed BytesIO raised ValueError out of the measurement path, a JPEG with 64 or more pre SOF segments fell to the fallback, and a BMP with a negative width produced a negative pixel total. It also flagged the Terraform provider drift and the tests that pinned the reference term to input_cost_per_pixel, both fixed here

The billing model is also checked in Lean 4 (v4.15.0, core only, 34 theorems, no sorry): the 2^20 megapixel constant and the 48576 pixel gap to 10^6, the Azure invoice reconciliation (4 MP plus 6 MP at 5 cents is 50 cents), additivity of the reference charge over any list of images, pro rata billing of non square and fractional sizes, the fallback only ever under billing by exactly the reference term, provider usage taking precedence over any pixel value, the deployment then catalog then zero rate chain honouring an explicit zero, rejection of non positive dimensions, and the bounds of the JPEG walk

Type

🐛 Bug Fix

Caveats (if any)

Medium

• A reference in a format the parser does not know (TIFF, HEIC) or with an unreadable header logs at debug and bills generated pixels only, for every reference on that request
• A path string or non-seekable stream as reference is not measured, same fallback
• A JPEG whose SOF sits past 16 MiB of metadata or more than 1024 segments is not measured, same fallback
• Deployments on the hosted cost map see the new rates only after this lands on main
• Existing Flex generation spend drops 4.86% because the megapixel constant was wrong

Low

• Every success callback sees one new integer reference_pixels in hidden params for image edits on the shared handler; a typed usage field is the follow-up
• The Rust model catalog ModelInfo and the inline cost map schema in tests/test_litellm/test_utils.py gained the same optional key so both schema checks accept it
• Bedrock Stability, Nova Canvas and Black Forest Labs edits bypass the shared handler and are not measured; none of their cost rows carries a reference meter today
• azure_ai/flux.2-pro is metered by Azure as first MP, additional MP and reference MP tiers; that pricing shape is left for a follow-up
• The sibling Terraform cost fields (input_cost_per_pixel, output_cost_per_pixel, per second and per million token rates) keep their pre-existing > 0 guard, so an explicit 0 for those is still dropped; that is a follow-up
• The Azure retail prices rows show no reference meter for FLUX.1-Kontext-pro, so it is left as per image

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

---

Note

Medium Risk
Changes spend calculation for Azure AI image edits and adjusts FLUX.2-flex per-megapixel rates; incorrect header parsing could under-bill reference images.

Overview
Azure AI FLUX.2 Flex image edit billing now counts reference image megapixels separately from generated output, and uses Azure’s 1 MP = 1024×1024 definition for per-pixel rates.

Before upload streams are consumed, the shared image-edit path measures reference dimensions from binary headers (PNG, JPEG, WebP, GIF, BMP) and stores the summed pixel count on the response as reference_pixels. The Azure AI cost calculator adds reference_pixels × input_cost_per_reference_pixel to existing generated-pixel or per-image charges; provider-reported token usage still overrides pixel math when present.

Catalog and types gain input_cost_per_reference_pixel; azure_ai/FLUX.2-flex sets both generated and reference rates to 0.05 / 1,048,576 (replacing the old 1,000,000-pixel megapixel assumption). Unreadable references fall back to billing generated pixels only.

^{Reviewed by Cursor Bugbot for commit 6db3435. Bugbot is set up for automated code reviews on this repo. Configure here.}

Link to Devin session: https://app.devin.ai/sessions/70b30e9f289d49e08561526f9110a933
Open in Devin Desktop: https://app.devin.ai/desktop/session/70b30e9f289d49e08561526f9110a933?variant=devin
Requested by: @shreshth-berri

[devin-ai-integration]

I'll fix CI failures and address comments from users with write access. I'll skip comments containing "(aside)".

• Disable automatic comment, CI, and merge conflict monitoring

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[codspeed]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_lit8140_flux2_reference_pixels (4347e20) with main (6593566)¹}

Footnotes

1. No successful run was found on main (eb7eeb5) during the generation of this report, so 6593566 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the previous billing, Terraform zero-override, and JPEG fill-run findings are resolved, with no new actionable issue in the latest changes.

Summary

This PR corrects Azure AI FLUX.2-flex billing by measuring reference-image pixels before upload streams are consumed and adding the corresponding reference-pixel charge to generated-image cost.

• Parses dimensions from PNG, JPEG, WebP, GIF, and BMP reference headers while restoring stream positions and safely falling back when measurement fails.
• Resolves generated and reference pixel rates from deployment pricing before the shared catalog, preserving explicit zero overrides.
• Updates FLUX.2-flex pricing to Azure's 1024×1024 megapixel definition.
• Extends model schemas, Rust catalog types, Terraform configuration, and generated dashboard types with input_cost_per_reference_pixel.
• Adds regression coverage for parser boundaries, deployment pricing, provider-reported usage, and synchronous and asynchronous edit paths.

_{Reviews (9) · Last reviewed commit: "fix(images): keep the marker prefix at a..."}

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous 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 6a1bf0e. Configure here.}

[devin-ai-integration]

Superseded by #42829