fix(flux2): estimate working memory for denoise and both VAE directions

[Pfannkuchensack]

Summary

Fix. The FLUX.2 path called model_on_device() without a working_mem_bytes estimate in every one of its load sites — the denoise node, the VAE decode and encode nodes, and the reference-image encode inside Flux2RefImageExtension. Every other base (SD1/SDXL, FLUX.1, SD3, CogView4, Qwen-Image, Wan, Anima, Krea-2) passes one. Without it the model cache reserves only the default device_working_mem_gb (3 GB) and fills the remainder of the card with the model, so it has no idea that the operation about to run needs far more than that.

Reference images are what turns this from tight into fatal. FLUX.2 concatenates reference latents onto the image stream, so attaching three 1024×1024 references to a 1024×1024 generation takes the attended sequence from 4 608 to 16 896 tokens — and the activation footprint from ~1.7 GB to ~6.5 GB, against a 3 GB reservation. Tile-based refiner workflows do exactly this, once per tile.

How. Two estimators, both calibrated against measured peak reserved memory (the conservative quantity, including allocator overhead), passed at every load site:

• Flux2DenoiseInvocation._estimate_working_memory() — FLUX.2 attention runs through SDPA and never materializes the O(seq²) score matrix, so activations scale linearly with the total attended sequence. Measured slope on the Klein 9B geometry in bf16: ~0.39 MB per token, flat from 1.5k to 28k tokens. It is also independent of block count (a no-grad forward frees each block's intermediates as it goes), so the same constant covers Klein 4B and 9B. Image, reference and text tokens all count; LoRA sidecar patches and the regional-prompting additive bias get their own terms.
• estimate_vae_working_memory_flux2() — the FLUX.2 VAE scales linearly in pixel area at ~2170 (decode) / ~1070 (encode) bytes per pixel per element byte, which rounds to the same 2200/1100 constants the FLUX.1 VAE already uses. A 1024×1024 decode peaks at ~4.3 GB, a 1536×1536 decode at ~9.6 GB. The tiled branch bounds the estimate by one tile, matching the 512px tiling the reference-image encode already forces.

Measurements (RTX 4090, bf16, peak reserved, each point in a fresh subprocess so allocator history cannot contaminate it):

image tokens	reference tokens	total seq	measured peak
4 096 (1024px)	0	4 608	1.66 GB
4 096	4 096 (1 ref)	8 704	3.25 GB
4 096	12 288 (3 refs)	16 896	6.38 GB
6 889 (1328px)	12 288	19 689	7.35 GB
6 889	20 667 (3×1328px refs)	28 068	10.70 GB
16 384 (2048px)	0	16 896	6.38 GB

The last row is worth noting: 2048px with no references and 1024px with three references produce the same sequence length and the same peak, which is what makes a single per-token constant the right model.

Related Issues / Discussions

Closes #9500

Note on the report: the reporter states the same workflow ran fine on 6.13. I could not confirm a 6.13 → 6.14 regression. _get_vram_available, _load_locked_model and _offload_unlocked_models are logically identical between v6.13.8 and v6.14.0-rc1 for single-GPU (the large model_cache.py diff is multi-GPU plumbing), and the only FLUX.2 memory change in 6.14 — tiling the reference-image VAE encode — reduces peak usage. The defect described here is present in both versions; given how narrow the margin is (see QA below), 6.13 most likely just got lucky more often. That also matches the reporter's own "sometimes it goes through, other times it OOMs at the Nth tile".

QA Instructions

Unit tests: pytest tests/app/invocations/test_flux2_working_memory.py (35 tests). The measured peaks above are pinned as both lower and upper bounds on the estimate, so a future constant change that would reintroduce the OOM — or one that over-reserves so hard the cache pushes the transformer to RAM — fails the suite. The wiring tests were mutation-checked: removing working_mem_bytes= from the denoise call site fails two of them.

End-to-end on a 24 GB card (RTX 4090), device_working_mem_gb: 3, enable_partial_loading: true, Klein 9B fp8 (17.35 GB resident as bf16), three 1024×1024 reference images, 4 steps, Qwen3 encoder on CPU.

Roomy card, 1024×1024 output — both pass, but look at the residency:

	transformer residency	free VRAM for the forward	result
before	17 350 MB (100 %)	~7 GB against a 6.5 GB need	completes
after	15 046 MB (86.7 %)	~9.5 GB	completes

That ~0.5 GB of margin in the "before" row is the whole bug. It is not a comfortable pass; it is a coin flip that lands differently depending on what else the cache happens to be holding — precisely the "sometimes it goes through, other times it OOMs at the Nth tile" the issue describes. With the estimate the cache deliberately holds 2.3 GB of the transformer back in RAM.

Tight card, 1328×1328 output — the difference stops being theoretical. A second process pinned 4.5 GB (scripts/allocate_vram.py) to emulate a card with a desktop and other apps on it, and PYTORCH_CUDA_ALLOC_CONF left at the stock native allocator:

	transformer residency	result
before	14 182 MB (81.7 %)	still not finished after 12 minutes; cancelled
after	13 894 MB (80.1 %)	completes in 58 s (denoise 35 s)

Near-identical residency, wildly different outcomes: without the reservation the forward has to claw its working set out of a card the cache believes is fine, and the allocator spends its time synchronizing and releasing cached blocks instead of computing.

Reproducing the reporter's hard OOM. I was not able to make the baseline OOM outright on this hardware — with partial loading enabled it degrades into the PCIe-thrashing case above instead of raising. A card that cannot fall back that way (partial loading disabled, or a model that must be fully resident) is where the same shortfall surfaces as torch.cuda.OutOfMemoryError.

On the choice of constant. Both estimators target peak reserved memory, not peak allocated, consistent with every other estimator in vae_working_memory.py. At 19 689 tokens the denoise allocates ~3.3 GB but reserves ~7.4 GB; targeting the allocated figure would be enough on backend:cudaMallocAsync and would reintroduce the bug for everyone on the default allocator. The cost is that cudaMallocAsync users reserve more than they strictly need.

Merge Plan

Nothing special — backend only, no schema, DB or redux changes. Node versions are unbumped, matching the precedent set by #9305 (the equivalent Qwen-Image working-memory fix), since no node interface changed.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[JPPhoto]

Please fix:

• invokeai/backend/util/vae_working_memory.py:110: CUDA-linear estimate underestimates FLUX.2 VAE attention on ROCm, where code notes materialized attention. High-resolution encode/decode can still OOM. Effect: fix fails on supported ROCm. Likelihood: Normal ROCm high-resolution use. Recovery: lower resolution; nodes expose no tiling control. Test: Run ROCm Torch 2.10 at 1024/1536px; compare peak reserved memory with estimate.

• invokeai/app/invocations/flux2_denoise.py:468: Regional prompting passes full float S x S mask via invokeai/backend/flux2/extensions/regional_prompting_extension.py:41; if SDPA selects math fallback, score workspace is quadratic, but estimate adds only mask storage. PyTorch documents backend-dependent SDPA dispatch and math intermediates here. Effect: high-resolution regional prompts can still OOM. Likelihood: Plausible backend/resolution edge. Recovery: disable regional prompting or lower resolution. Test: Measure torch.cuda.max_memory_reserved() for masked 1024/2048px FLUX.2 forwards on Torch 2.7/2.10.

Suggestions:

• Consider backend-specific peak-memory calibration for ROCm and regional-mask attention.

[Pfannkuchensack]

Cant really test the ROCm stuff maybe @lstein can take a look

[JPPhoto]

Merge blockers:

• invokeai/backend/util/attention.py:71-74 treats MPS as fused, but Torch 2.7.1 routes MPS SDPA through math, materializing Q @ K^T (dispatch, math path). vae_working_memory.py:158-164 therefore omits about 3.5 GB at 1024px. Effect: MPS Flux2 VAE decode can OOM after cache admission. Likelihood: Normal MPS 1024px+ decode. Recovery: Lower resolution or CPU fallback. Test: Run Torch 2.7.1 MPS 1024px decode; assert nonzero score reservation.

Other findings/issues:

• invokeai/backend/util/attention.py:93-94 treats every probe exception as fused. Probe OOM or backend failure then returns zero score bytes. Effect: Under-reservation and forward OOM. Likelihood: Low free VRAM or backend mismatch. Recovery: Clear cache or lower resolution. Test: Inject torch.cuda.OutOfMemoryError from torch.empty; assert conservative budgeting.
• invokeai/backend/util/attention.py:76-92 probes native Torch eligibility, but Flux2 can use another Diffusers backend through dispatch_attention_fn (Diffusers dispatch). Forced native-math bypasses the probe. Effect: CUDA math attention can materialize O(S^2) while estimate adds zero. Likelihood: Custom attention backend users. Recovery: Restore fused backend or lower resolution. Test: Force _native_math; assert score bytes are included.

Suggestions:

• Instead of assuming non-CUDA is fused, probe the active backend or conservatively charge math on MPS.
• Instead of returning fused on exceptions, budget unknown probe results as math.