[Pfannkuchensack]

Stacked on #9899 (and through it #9894, #9896, #9897, #9898). Until those merge, the diff also shows their commits. This PR's own changes are 87079a62d9, cb91aab513 and 7714be86b6; the rest are merges.

Summary

The shipped VAE working-memory estimates come from a few reference GPUs. A user's card, driver or allocator can need less, or more. Admins can now calibrate a model's VAE decode on their own GPU from the Model Manager, and estimates on that GPU then use the measurement.

Model detail, "Working memory" section (actions are admin-only)

• Calibrate: a dialog shows the planned runs (sizes per dtype, 1 warm-up + 2 measured each). The runs then go through the queue as short generations, with progress and Stop.
• Reset: removes this model's measurements on this GPU.
• Per line (operation · dtype): a table with Measured, Built-in and Calibrated estimate per size. Each column header has an info popover. Sizes that did not count are muted with their reason.
• States: never calibrated, in use, may be outdated, not used, too few sizes, running, incomplete, failed, unsupported, no GPU, more than one generation device, store unreadable, and measurement method not verified on this platform.
• Shared lines: models whose decode prices on the same line share the calibration. The section lists who measured it.

Estimates

• Within the measured sizes: the user line, an upper envelope over the counted measurements plus 5 % (the shipped margin).
• Above the largest measured size: the larger of the user line and built-in.
• No usable fit: built-in exactly as before.
• Environment changes: a hard change (backend, allocator, torch, settings) stops using the fit. A soft change (driver) keeps it and marks it "may be outdated".
• Below the working-memory floor (device_working_mem_gb), nothing changes.

Robustness found in hardware runs

• Finished attempts only: a calibration attempt's windows join no fit until every run of that attempt has finished. Larger cases are never priced from the same attempt's smaller ones.
• Fragmentation: allocator fragmentation in a warm server is detected per attempt and excluded. This happens on ROCm with expandable segments: reserved 1.3–1.47× allocated while allocated stays flat. The shipped-data builder uses the same rule (invokeai/backend/working_memory/fragmentation.py).
• Full offload before the measured lock: before a measured calibration lock, other unlocked models are moved off the GPU whole. The cache's partial offload otherwise leaves the measured model partly resident ([bug]: model cache leaves the requested model partly resident after splitting the last offload candidate #9907). Normal locks are unchanged.
• Plan capacity: the plan sizes cases for the capacity the cache will grant, including driver memory kept after earlier runs (HIP's retained pool).

Before/after on an RX 9060 XT (SDXL, fp32 decode)

	1024² reservation
Built-in	~6.1 GiB
User fit without the fragmentation rule (first hardware run)	8.3 GiB, which would evict the UNet again
This PR	follows the counted line up to 1152², below built-in (1152²: measured 6.95, calibrated 7.3, built-in 7.48 GB)

Related Issues / Discussions

• Builds on feat(e2e): add per-family end-to-end regression runner #9894, feat(working-memory): ship VAE working-memory coefficients as data #9896, feat(working-memory): measure working memory in admin calibration runs #9897, fix(working-memory): reserve VAE working memory per dtype and allocator #9898, feat(e2e): calibrate working memory and build shipped data from probe runs #9899.
• Found while testing: [bug]: model cache leaves the requested model partly resident after splitting the last offload candidate #9907 (the model cache leaves the requested model partly resident after a partial offload) and [bug]: first load after an update shows raw translation keys (service worker serves stale locales) #9908 (the first load after an update shows raw translation keys).

QA Instructions

Tests

• uv run --no-sync pytest -n 4: 4818 passed, 158 skipped, 5 xfailed. Covered: working-memory backend, calibration service, router and runner, session processor, tests/backend/model_manager/load, tests/app/invocations, the builder, tests/family_e2e.
• uv tool run ruff@0.11.2 check and format --check.
• API regenerated (make frontend-openapi/frontend-typegen equivalent); pnpm lint and pnpm test in invokeai/frontend/api.
• Frontend:
  • pnpm build (format, oxlint, tsc, architecture);
  • vitest src/features/models 255;
  • browser tests for the section and the graph compiler 24;
  • node --test scripts/mock-backend-fixtures.test.mjs 28;
  • pnpm test:performance:architecture passes.
• Full pnpm test here: 3 failures from the de-DE locale of this machine (number formatting) and 2 load timeouts that pass alone. All unrelated.

Mutation checks, each caught:

• the margin applied to the whole reservation;
• a user fit outside its range;
• live-attempt windows joining a fit;
• attempts pooled in the fragmentation rule;
• the calibration offload disabled, or applied outside calibration;
• splitting instead of whole moves;
• the plan's outside-usage floor removed;
• reset scope, stale-failure display, uncounted-row display.

Hardware, SDXL from the UI (18 runs per attempt) on an RTX 4090 (CUDA, cudaMallocAsync) and an RX 9060 XT (ROCm 10, Windows, expandable segments):

• Final attempt: 36 of 36 windows fully resident on both cards.
• 9060 XT: fp32 1152² counts; the earlier fragmented 1216² is excluded.
• 4090: calibrated below built-in at every size (0.91–0.96×).

Browser

• Every state on the mock backend (pnpm dev:mock; fault scenarios through POST /__faults) and on both real servers.
• Keyboard flows: Calibrate, Stop and Reset confirmations; focus lands on Stop after start and returns after each dialog; popovers open with Enter and close with Escape.
• Screenshots: to be attached.

Linux NVIDIA: the NVML counter reader and this calibration round ran on 2× RTX PRO 6000 Blackwell (#9910 (comment)); with #9897's edba0dd882 it counts as verified.

Not verified on hardware: the Linux AMD, MPS and XPU counter readers. The UI marks them as unverified.

Review

Resolved material findings

• Shared lines: they showed as "built-in". Reset copy now says other models' measurements still apply.
• Stale state: it never cleared after a recalibration. The header now ignores other environments and superseded estimator versions.
• Reset scope: Reset removed less than the summary attributed to the model.
• Fit refresh during an attempt: fits were refreshed mid-attempt and the plan was built from stale fits. Now held back per attempt.
• Fragmentation: fragmented windows set the user line, and windows from different attempts were pooled.
• Margin test: it could not catch a margin on analytic terms.
• UI states: a stale failure was shown after a lost connection, and the table columns were misaligned.
• Plan sizes: they ran partly resident on 16 GB cards.

Remaining limits

• First calibration on ROCm: the plan reads capacity live, so its largest size may not count until the retained pool has been recorded once. Partly resident windows never count.
• Growth rule margin: a clean attempt's first-to-largest growth was 1.095×, close to the 1.10 threshold. A false flag falls back towards built-in.
• Plan capacity formula: it re-implements the cache's budget rule in the app layer (without the reclaimable-block credit, so it is conservative).
• Encodes: image-to-image and inpainting encodes keep built-in estimates.

Compatibility / Rollout

• Store: the store file per device class under <db_dir>/working_memory/ comes from feat(working-memory): measure working memory in admin calibration runs #9897. New optional fields (attempt_id, run soft) are left out when absent, so older files load. No database migration.
• API: new GET /api/v2/models/i/{key}/working_memory/plan and /summary (404 for unknown models). attempt_id on calibrate. Line fields measured_by, own_windows, and size fields counted, excluded_reason. Generated openapi.json and schema.ts are updated.
• Without calibration: estimates are unchanged; the resolver falls back to the shipped values.
• Locale: +3.8 KB of strings. browser-baseline.json records the measured otherRawBytes (215114 B on every route, main 211323 B); nothing else in the baseline changed.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[Pfannkuchensack]

@joshistoast Thanks, fixed in e00e817. Your case reproduced exactly. On a 24 GiB card with the shipped cuDNN/native lines, bf16 plans [896, 1216, 1792] and fp32 plans [640, 896, 1344].

• One shared dtype rule. The plan now uses the dtype the decode actually runs in. The loaders' fp16-avoidance rule moved into TorchDevice.dtype_avoiding_float16, which both the loaders and the calibration service call. A standalone FLUX.1/FLUX.2 VAE kept in float32 is therefore planned on decode_fp32, and the UI's precision label follows.
• Pipelines with their own VAE.
  • Those that decode with their own VAE, as Generate leaves the slot empty, plan that VAE's dtype. That covers FLUX.2 and Z-Image Diffusers pipelines and self-contained SDNQ pipelines.
  • A non-bundled FLUX.2 main plans the VAE of the component source Generate picks for it.
  • The loader dtypes the plan restates are pinned by tests against the real loaders.
• No silent wrong measurement. A standalone VAE is no longer calibrated through a FLUX.1 SDNQ pipeline, which ignores the VAE pick. The plan also no longer pins a standalone VAE for a pipeline that decodes with its own.
• Follow-up. While checking this I found that SDNQ pipelines load their VAE in fp32, and a standalone SDNQ FLUX VAE in fp16 under auto. That is filed as [bug]: SDNQ pipelines load their VAE in fp32, and a standalone SDNQ FLUX VAE in fp16 under auto #9996, since changing the loaders is an inference change outside this PR.

[joshistoast]

[P2] Refresh the memory-dependent plan after its inputs change

The plan depends on current free GPU memory and fitted estimates, but reopening the panel does not refresh it. If another process occupies GPU memory when the plan loads, Calibrate can remain disabled after memory recovers. Completion events refresh summaries/status only, so recalibration can also reuse sizes calculated before previous measurements. Refresh the plan when reopening or preparing calibration, and invalidate it after completed measurements.

[Pfannkuchensack]

Fixed in a1d9318. The section now re-reads the plan every time it opens, not only the first time, and Calibrate re-reads it right before it starts: the button shows a loading state meanwhile, and the dialog and the run use the plan just read. If that plan has become unsupported, no dialog opens and Calibrate is disabled with the new reason; a failed read shows the error with Retry. Calibration status events now mark every plan stale (they share fits and recorded memory), so the next opening or start reads it fresh. I mark rather than refetch on the event itself, because free memory during or right after a calibration is not what a normal generation sees. Covered by unit and browser tests.

[JPPhoto]

Findings/issues:

1. invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:1101: What: _live_runs() checks only queue item IDs already stored on a run. Problem: a retry is a new same-origin item, absent from the store until admitted, so a failed run's valid points can affect estimates while its retry waits. Effect: queued generations can use a partial calibration fit and under-reserve VRAM. Likelihood: plausible when an admin retries an item behind queued work. Recovery: a later refresh may replace the fit. Fix: include pending retry items when deriving live runs. If skipped: reservations can stay wrong while the retry is pending. Test: fail a run after recording a valid point, queue a retry without dequeuing it, refresh fits, and assert the run remains excluded.

2. invokeai/app/services/working_memory_calibration/working_memory_calibration_plan.py:224: What: vae_source() falls back to the first FLUX.2 pipeline when none matches the target variant. Problem: with only an incompatible source and no standalone VAE or encoder, the plan says calibration is supported, but graph compilation rejects the source and queues no cases. Effect: the Model Manager offers calibration that cannot run. Likelihood: limited to installs with only an incompatible component source. Recovery: install a compatible source or standalone components. Fix: mark the plan unsupported unless its graph can compile. If skipped: users get a compile error instead of calibration. Test: compile a Klein 9B plan with only a Klein 4B source and no standalone components; assert it is unsupported or compiles with compatible components.

3. invokeai/frontend/webv2/src/features/models/ui/detail/WorkingMemorySection.tsx:279: What: the confirmation dialog starts the plan fetched before it opened. Problem: if free GPU memory or fitted measurements change while the dialog is open, confirmation still queues the old cases. Effect: runs may be tiled, partly resident, or fail, wasting calibration work. Likelihood: plausible if the dialog stays open while GPU use or another calibration changes. Recovery: manually retry after memory stabilizes. Fix: refresh the plan on confirmation and stop if it is no longer supported. If skipped: users can confirm stale sizes. Test: change the plan while the dialog is open; confirm and assert it refreshes or queues nothing if unsupported.

4. invokeai/frontend/webv2/src/features/models/data/workingMemoryCalibration.ts:108: What: postStep() stops retrying on every ApiError. Problem: the API documents that a keyed request can return 503 after its queue items were committed; retrying the same key can recover that run, but the client stops and skips later steps. Effect: the attempt is partial, and a manual retry can duplicate already queued work. Likelihood: low, during transient store-lock or write failures. Recovery: manually retry the calibration. Fix: retry 503 with the same idempotency key within the existing bound. If skipped: transient storage errors leave incomplete attempts. Test: return 503 after enqueue, then success for the same key; assert the step is not duplicated and later steps are posted.

[Pfannkuchensack]

@JPPhoto Thanks. All four are fixed in 623a5e3, and every new test fails without its fix. The branch also has main and the stack below merged in (4339e78). Its only conflict was the generated byte fields in browser-baseline.json.

1. A pending retry and live runs

• _live_runs now asks the queue for the origins of its unfinished calibration items (pending, in progress or waiting) through the new SessionQueue.get_unfinished_origins.
• A retry carries its run's origin from the moment it is queued, so a failed run holds its windows back again as soon as the retry is queued. It releases them once the retry finishes, is canceled or is removed.
• The by-id reads and the settled cache are gone.
• Cost: the read finds rows through the status index and tests the origin on those rows. A query-plan test pins this. On SQLite it costs about 0.05 ms with a handful of unfinished items, whatever the history size, and 34 ms with 1000 pending 64 KB sessions. It runs once per refresh of a store that has windows.
• Test: a run fails after recording a valid point and a retry is queued but not dequeued. The fit stays shipped. Once the retry is canceled, the run's points count.

2. A FLUX.2 plan with only an incompatible source

• vae_source returns the first pipeline whose encoder fits. It falls back to another pipeline only when a standalone encoder exists, because only then does Generate keep that pipeline, and then only for its VAE.
• Without a usable text encoder the plan is supported: false, and the reason names the encoder to install.
• A standalone FLUX.2 VAE goes through the first main model whose graph compiles.
• I checked the rule against the real graph compiler on 20 model sets. The tests cover:
  • a Klein 9B with only a Klein 4B source, also with a standalone VAE or with an encoder of the wrong size (all unsupported);
  • a standalone FLUX.2 VAE as the target;
  • a Qwen3 8B encoder beside the unfitting pipeline (supported).
• Beyond FLUX.2: the section now compiles the plan's cases with the same compiler before it opens the dialog, and again on confirm. A case that does not compile is shown as the section's reason and queues nothing. This also covers other bases that need a standalone encoder, such as FLUX.1 without T5.

3. A stale plan at confirmation

• Confirm reads the plan again and starts the plan it read.
• If the plan no longer holds, or the read fails, the dialog closes and the section shows the reason, or the error with a retry.
• If the lines or the number of runs changed, the dialog stays open and shows the new plan with an alert, and you confirm again.
• A change in size alone, such as free memory moving during a generation, starts the new sizes directly, so the prompt does not repeat while a generation runs.
• While the read runs, the button shows a spinner and a second click is ignored. Cancel or Escape takes the confirmation back.

4. A 503 after the enqueue

• postStep retries a 503 with the same request and key, within the existing bound. That replays the run the first request created, so nothing is queued twice.
• 429, 404 and 409 still fail at once.
• Tests:
  • 503 after the enqueue, then success: no step is duplicated and the later steps are posted;
  • 503 every time: three posts, then stop;
  • a non-503 error is not retried.

Checks: backend tests for working-memory calibration, session queue and processor, ruff, frontend lint, unit and browser tests, a production-build browser check of the new flows, and the frontend performance gate. MySQL/MariaDB was not run locally; the external-database CI job covers the new query.