fix(qwen): estimate Qwen Image VAE working memory so the cache frees room before decode/encode

[lstein]

Summary

The Qwen Image qwen_image_l2i (decode) and qwen_image_i2l (encode) invocations called model_on_device() without a working_mem_bytes estimate — unlike the SD/SDXL l2i path, which calls estimate_vae_working_memory_sd15_sdxl(...). As a result, the model cache only reserved the default device_working_mem_gb and never evicted the resident transformer / text encoder before the VAE decode.

On a near-full card this OOMs. Reproduced with Qwen Image Edit 2511 (Q8_0) + the standard Qwen Image VAE on a 48 GB AMD W7900: with the transformer (~20.7 GB) and text encoder (~15.8 GB) resident, the autoencoder decode tried to allocate ~5 GiB into the fragmented ~8 GiB remainder and failed:

CUDA out of memory. Tried to allocate 5.01 GiB. GPU 0 has a total capacity of 44.98 GiB
of which 3.69 GiB is free. ... 2.48 GiB is reserved by PyTorch but unallocated.

Root cause

ModelCache._load_locked_model() computes vram_available = free_vram − working_mem and only evicts other models when that drops below what the locked model needs. The VAE is tiny (~242 MB) and already resident, so model_vram_needed ≈ 0 and nothing is ever evicted — the big transformer/text encoder stay put and the decode is squeezed into whatever fragmented VRAM is left.

Passing a realistic working_mem_bytes lets the cache make room (evicting other models) before the operation runs, which is exactly what the SD/SDXL path already does.

Fix

• Add estimate_vae_working_memory_qwen_image() in vae_working_memory.py.
• Pass the estimate into model_on_device(working_mem_bytes=...) in both the decode and encode invocations.

Calibration

The estimate is calibrated against a measured decode on a W7900. At 1248×832 the decode grew CUDA reserved memory by ~10.06 GiB (implied constant ~5082); rounded up to 5500 for headroom. The current SD constant (2200) under-modeled this heavier video-style VAE by ~2.4×.

The constant intentionally tracks peak reserved (not just allocated) memory. The cache's guarantee is "if it doesn't evict, then free ≥ estimate," so the estimate must be ≥ the decode's true reserved footprint. This closes the danger zone where the cache would skip eviction yet the decode would still reserve more than the free VRAM:

• Tight card (transformer + text encoder resident, ~8 GiB free): estimate ~10.6 GiB > free → cache evicts the text encoder → ~24 GiB clean headroom → decode succeeds.
• Roomy card (free ≥ estimate): no eviction, but free already ≥ the decode's need → fits.

Testing

• Reproduced the OOM, then confirmed a clean run at 1248×832 with the calibrated constant and device_working_mem_gb back at its default — the text encoder is offloaded just before the VAE decode and the generation completes.
• ruff check / ruff format / compile all clean.

Notes / open question

• The encode constant (2750) follows the SD-style "half of decode" convention and is not independently measured — a conservative default. Worth a follow-up measurement if encode-side OOMs surface (relevant for Qwen Image Edit, which encodes an input image).
• Calibration was done on ROCm/W7900; expandable_segments:True did not resolve the fragmentation on that stack — eviction is what reliably works.

🤖 Generated with Claude Code

[Pfannkuchensack]

Findings

Medium - Missing test coverage for the actual fix

Paths: invokeai/app/invocations/qwen_image_latents_to_image.py:45-52, invokeai/app/invocations/qwen_image_image_to_latents.py:48-54

Both the decode and encode paths now compute and pass working_mem_bytes, but the PR adds zero tests. The repo already has a directly-applicable pattern: tests/app/invocations/test_z_image_working_memory.py mocks the VAE/context and asserts model_on_device.assert_called_once_with(working_mem_bytes=expected_memory) plus estimate.assert_called_once().

Without an equivalent, a future refactor that drops the working_mem_bytes= argument, reverts to a bare model_on_device(), or wires the wrong operation/tensor into the estimator would silently reintroduce the original OOM regression with no CI signal. The entire value of this PR is "the estimate is actually passed to the cache," and nothing proves that automatically.

To expose this issue, add a test that constructs QwenImageLatentsToImageInvocation and QwenImageImageToLatentsInvocation with a mocked AutoencoderKLQwenImage, patches estimate_vae_working_memory_qwen_image, invokes them, and asserts model_on_device was called once with working_mem_bytes=<estimate> for both operation="decode" and operation="encode" (mirroring test_z_image_working_memory.py).

Low - Encode constant is unmeasured and the docstring is inconsistent

Path: invokeai/backend/util/vae_working_memory.py:117-118

The encode constant is hardcoded as 2750, which the PR description itself flags as "not independently measured" and merely "half of decode." The new docstring says "Encoding uses ~half the working memory of decoding," but the sibling estimators it claims to match document encode as "~45%" / "~50%" of decode (invokeai/backend/util/vae_working_memory.py:25, :65, :86, :138).

For Qwen Image Edit, which encodes a real input image, an under-modeled encode constant reproduces the exact failure mode this PR fixes (cache declines to evict because free >= estimate, yet the encode reserves more than free and OOMs). This is a residual correctness gap, not just a style nit, because the encode path is the one Qwen Image Edit actually exercises.

To expose this issue, add a test that calls estimate_vae_working_memory_qwen_image(operation="encode", ...) and asserts the returned value is at least the measured encode reserved footprint once an encode measurement exists; until then, treat the encode constant as unverified.

Open Questions

• Single-point linear calibration for a video-style VAE. The decode constant 5500 (invokeai/backend/util/vae_working_memory.py:118) is extrapolated linearly in h*w from one calibration point (1248x832, ~10.06 GiB reserved). At that point the formula yields 1038336 * 2 * 5500 = ~10.64 GiB, only ~5.7% above the measured value. The estimator assumes peak working memory is strictly linear in spatial area. If the VAE has any super-linear (attention) memory term at higher resolutions, a large decode (e.g. 2048x2048) could under-estimate, the cache would skip eviction (free >= estimate), and the decode would still OOM. A second calibration point at a larger resolution would settle whether the linear model and the thin margin hold.

[lstein]

Thanks @Pfannkuchensack — addressed in dd285a1.

Medium (missing test coverage): Added tests/app/invocations/test_qwen_image_working_memory.py, mirroring test_z_image_working_memory.py. It constructs QwenImageLatentsToImageInvocation (decode) and exercises QwenImageImageToLatentsInvocation.vae_encode (encode) with a mocked AutoencoderKLQwenImage, patches estimate_vae_working_memory_qwen_image, and asserts:

• estimate is called once with the correct operation (decode/encode), and
• model_on_device.assert_called_once_with(working_mem_bytes=<estimate>) for both paths.

So dropping the working_mem_bytes= argument or reverting to a bare model_on_device() now fails CI.

Low (encode constant unmeasured / inconsistent docstring): Reworded the comment to stop claiming the encode constant "matches the other estimators." It now states plainly that the encode constant is not independently measured, is set to half the decode constant (consistent with the siblings' ~45–50% ratio), and should be recalibrated against a measured encode — the path Qwen Image Edit actually exercises. I left the value at 2750 since there's no measurement to justify a different number yet.

Open question (single-point linear calibration): Agreed this is a real limitation. The decode constant is extrapolated linearly in h*w from one calibration point (1248x832), and a super-linear attention term at very high resolution could under-estimate and let the cache skip eviction. I don't have a second high-res measurement on hand to settle it; the decode comment now documents the single calibration point and the headroom rounding so the assumption is explicit. I'll treat a second calibration point (e.g. 2048x2048) as a follow-up rather than block this OOM fix on it.

[lstein]

Multi-point calibration — follow-up on the single-point concern

Added scripts/calibrate_qwen_vae_working_memory.py (backend-portable; uses only torch.cuda.*, which works on both CUDA and ROCm) and ran the full grid in fp16 on an AMD W7900. Each (operation, resolution) is measured in a fresh subprocess to keep the caching allocator's fragmentation history from contaminating the reserved-memory reading. We measure peak reserved (not allocated) delta, since that's what the cache's free >= estimate check compares against.

Implied constant = reserved / (h·w·element_size) (fp16, W7900):

op	512²	768²	1024²	1536²	1792²	2048²
decode	5132	4596	4570	3273	3735	4813
encode	5864	5858	5858	3532	4364	OOM

Three takeaways:

1. The encode constant was a real bug — fixed. 2750 ("half of decode") under-estimated by ~2× at every measured resolution (512² encode needs 2.86 GiB; 2750 budgeted 1.34 GiB). This is precisely the cache-skips-eviction OOM you flagged for Qwen Image Edit. Encoding actually reserves ≥ decoding at matched area, so I raised encode 2750 → 6300, which now covers all 9 measured encode points with headroom.

2. Decode 5500 holds up. It covers all 10 decode points including 2048² (37.6 GiB) with margin (max implied k = 5132). Left unchanged, now backed by 10 points instead of 1.

3. The linear model breaks down at the top end. Above ~1792² memory grows super-linearly in area (an attention term — 1792²→2048² grows ~quadratically), and the implied constant is non-monotonic (a drop in the mid-range, consistent with an SDPA-backend crossover on ROCm). So a single linear constant under-estimates for very large decodes on big-VRAM cards — though such resolutions OOM on a 48 GB card regardless. Documented in the estimator as a known limitation rather than over-engineering an area² term now.

On CUDA vs ROCm: these are the conservative (ROCm) numbers. The curve shape is architectural and should reproduce, but the absolute constant and the mid-range crossover are backend-specific (MIOpen vs cuDNN conv workspaces, flash-attention availability). I'll run the same harness on an NVIDIA card and ship the per-backend max — the CUDA numbers can only push the constants up, never down, so the encode fix above is safe to land now either way.

To reproduce on NVIDIA:

python scripts/calibrate_qwen_vae_working_memory.py --csv qwen_vae_cuda.csv

(auto-discovers the VAE under $INVOKEAI_ROOT/models, or pass --vae <dir>; defaults to fp16 to match InvokeAI's default precision).

[lstein]

CUDA results in — constants now branched per backend

Ran the same fp16 grid on an NVIDIA card. Implied constant = reserved / (h·w·element_size):

	512²	768²	1024²	1536²	1792²	2048²
ROCm decode	5132	4596	4570	3273	3735	4813
ROCm encode	5864	5858	5858	3532	4364	OOM
CUDA decode	2660	2519	2690	2671	2281	OOM¹
CUDA encode	1456	1451	1458	1456	1455	1455

¹ card-size OOM, not a super-linear blowup.

This decisively answers the open question and explains the earlier weirdness:

• The super-linear / non-monotonic behavior was a ROCm artifact, not architectural. On CUDA the curve is flat-to-sub-linear (decode even drops to 2281 at 1792²; encode is dead flat at ~1455). The Qwen VAE is attention-heavy: CUDA uses Flash/efficient attention (O(area)), the ROCm build falls back to math attention (O(area²)) → ~2× (decode) to ~4× (encode) more memory, and the super-linear tail above ~1792².
• "Encode is half of decode" is a CUDA truth (1456/2680 ≈ 0.55) that's simply false on ROCm, where encode reserves ≥ decode. That's why the original unmeasured ROCm encode constant OOM'd.

Because the two backends diverge far beyond any headroom, a single constant would either under-estimate on ROCm (OOM) or over-budget CUDA ~2–4× (needless eviction on tight NVIDIA cards). So the estimator now branches on torch.version.hip:

decode: ROCm 5500 / CUDA 2900
encode: ROCm 6300 / CUDA 1600

Each verified to cover its full measured grid (19 points/backend) with ~8% headroom, and there's a parametrized test asserting the constant selected for each (operation, backend) so a refactor can't silently swap them. Pushed in ee560c0.

[Pfannkuchensack]

Qwen-Image VAE — Working-Memory-Kalibrierung

Setup: torch 2.7.1+cu128 · NVIDIA GeForce RTX 4090 · hip=None · dtype=float16
VAE: S:/model_library/qwenImage_qwenImageVAE.safetensors

Decode

HxW	area	reserved (GiB)	alloc (GiB)	implied_k
512×512	262 144	1.311	1.274	2684.0
768×768	589 824	2.992	2.563	2723.6
832×1248	1 038 336	5.211	4.331	2694.3
1024×1024	1 048 576	5.254	4.369	2690.0
1088×1920	2 088 960	10.408	8.468	2674.9
1280×1280	1 638 400	8.188	6.693	2682.9
1536×1024	1 572 864	7.857	6.434	2682.0
1536×1536	2 359 296	11.736	9.532	2670.7
1792×1792	3 211 264	15.949	12.884	2666.4
2048×2048	4 194 304	20.799	16.754	2662.2

Encode

HxW	area	reserved (GiB)	alloc (GiB)	implied_k
512×512	262 144	0.711	0.900	1456.0
768×768	589 824	1.594	1.724	1450.7
832×1248	1 038 336	2.824	2.854	1460.3
1024×1024	1 048 576	2.848	2.879	1458.0
1088×1920	2 088 960	5.672	5.498	1457.7
1280×1280	1 638 400	4.459	4.364	1461.1
1536×1024	1 572 864	4.281	4.199	1461.3
1536×1536	2 359 296	6.396	6.179	1455.6
1792×1792	3 211 264	8.705	8.320	1455.3
2048×2048	4 194 304	11.365	10.793	1454.8

Summary (max implied constant = candidate scaling_constant, vor Headroom)

op	n	min_k	max_k	empfohlen
decode	10	2662.2	2723.6	>= 2724 (+ Headroom)
encode	10	1450.7	1461.3	>= 1461 (+ Headroom)