[Pfannkuchensack]

Stacked on #9910 (and through it #9894, #9897, #9898, #9899). Until those merge, the diff also shows their commits. This PR's own changes are f11f6029f7 and f0349da68d.

Summary

The Z-Image denoise reserved a hand-set 0.25 MiB per attended token plus a fixed 1 GiB on every platform and dtype. Fresh-process measurements on an RTX 4090 and an RX 9060 XT show that this is twice what bf16 needs at 1024², and less than fp32 needs at large sizes. The denoise now runs through the working-memory estimator framework, and its per-token lines are fitted from those measurements.

Z-Image denoise reservation, 512-token caption	1024² bf16	2048² bf16	2048² fp32
Before (every platform)	2176 MiB	5248 MiB	5248 MiB
RTX 4090, native allocator	1172	4297	7406
RTX 4090, cudaMallocAsync	841	3058	5506
RX 9060 XT, expandable segments (Windows)	993	2954	5256
RX 9060 XT, native (Windows)	1400	4538	8021
Unmeasured: XPU, MPS, CPU, Linux ROCm, other allocators	at least the previous values

Estimator

• denoise.z_image prices a fitted line per dtype over the attended tokens.
• Regional bias, the SDPA score matrix, ControlNet hints (n + 4), LoRA patches, the dequantization copy and a new fp8 cast transient are analytic terms. The fp8 term is the largest layer's widened weights; FP8 Storage hooks and kept fp8 checkpoint weights are both covered.
• The node passes model_on_device(working_memory=…), so calibration runs now record denoise windows.
• User calibration from the Model Manager stays VAE-only.

Framework

• Fitted lines are declared by each EstimatorSpec (fit.Line), so they can run over tokens as well as output pixels.
• Points with an analytic term never join a fit.
• unmeasured_floor keeps every unmeasured platform at or above the previous values.
• Entries can ship as legacy, i.e. hand-set; only Linux ROCm native does.
• scripts/working_memory/build_shipped.py takes fresh-process point files. It replaces a legacy entry only when capping points measure every line.

Measurement

• scripts/calibrate_z_image_working_memory.py measures one forward per fresh process, with demand = max(driver, reserved, allocated) growth. The cold and warm runs both count.
• It uses the server's runtime: run_app.apply_server_runtime shares the steps with run_app(), and a test pins that both apply the same sequence.
• It uses the probe's window definition, now shared in invokeai/backend/working_memory/measurement.py.

Related Issues / Discussions

• Builds on 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, feat(models): calibrate VAE working memory per model from the Model Manager #9910.
• The fp8 cast term replaces a share of the old fixed base, which used to stand for it.

QA Instructions

Measurements

• RTX 4090 (torch 2.13+cu130, native and cudaMallocAsync) and RX 9060 XT (torch 2.13+rocm10.0, Windows, expandable and native, AOTriton fused).
• bf16, fp16 and fp32 at 512–2048 px, with a 512-token caption plus 77-token captions at two sizes.
• Random-init model at the real geometry with 2 blocks; 4 blocks measured the same.
• 84 points ship in working_memory/data/denoise.json. Each sits at 1.05–1.22× its estimate.

Checks of the analytic terms

• ControlNet with 15 blocks: covered at 1.05–1.08×.
• Forced math kernel: the 14.2 bytes per score element measured is below the 17 charged.

In-server check with the real Z-Image Turbo, through the calibrate route: every window is below its new estimate.

• RTX 4090, native allocator, bf16: ratio 1.06–1.07 at 1024², 1536² and 2048².
• RTX 4090, cudaMallocAsync, bf16.
• RX 9060 XT, expandable segments, int8 build.

A/B, Z-Image Turbo bf16, 8 steps, family E2E tool, interleaved runs before → after → before, 1 warm-up + 3 measured per size:

• Times are within the spread of the two "before" runs on both cards: RTX 4090 at 1024²/1536²/2048², RX 9060 XT at 1024²/1536².
• No paging is caused by this change.
• On these roots device_working_mem_gb: 3 decides at ≤ 1536², so the gain shows above that, or with a lower floor.

Tests

• uv run --no-sync pytest -n 4: working memory, calibration services, router and runner, tests/app/invocations, tests/backend/model_manager/load, fp8, builder, calibration script and run_app environment. 4848 passed.
• Ruff is clean. The VAE data, fits and reservations are byte-identical to the base.
• Mutation checks, each caught:
  • the legacy pinning;
  • the term-derived line eligibility;
  • the floors;
  • legacy replacement by probe windows;
  • the builder's larger-reading rule;
  • the server/script runtime steps;
  • demand without the driver figure.

Review

Resolved findings

• Unmeasured platforms (XPU, MPS, CUDA with expandable segments, ROCm hipMallocAsync) had dropped to a measured card's line. They are floored at the previous values now.
• The measurement script had copied the server startup and the probe's window. Both are shared now.
• Line eligibility is derived from the terms instead of hand lists.
• The builder reports a lower remeasurement and never replaces legacy values from probe windows alone.
• The calibration summary's excluded_windows description is updated, and the VAE plan's capacity input ignores denoise windows.

Remaining limits

• Linux ROCm stays on the previous values: unmeasured, with fp32 below what Windows measured. A Z-Image fp32 transformer (24 GB) does not occur on ROCm cards that lack bf16.
• Only the transformer forward is measured. Sage attention and partial-load weight streaming are not measured.

Compatibility / Rollout

• No persisted state. The API changes only in a field description; openapi.json and schema.ts are regenerated.
• EstimatorSpec.line is now a required argument; the five VAE specs pass convolution_line.

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)

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

[Pfannkuchensack]

@JPPhoto Thanks, confirmed and fixed in 5b1c79d.

Fit-grid point marked in_fit by kind alone (scripts/calibrate_z_image_working_memory.py). point_line now sets in_fit only if the point is a fit-grid point and Z_IMAGE_DENOISE_SPEC.line(features).fits is true, matching fit.py. A fit-grid forward with a nonzero analytic term (e.g. score_matrix_bytes) ships as an out-of-fit point that bounds its estimate from below only.

Test: test_a_fit_grid_forward_its_line_does_not_admit_ships_out_of_the_fit_and_an_admitted_one_stays_bounded in tests/test_calibrate_z_image_working_memory.py.

Shipped data: all 84 denoise.z_image points are line-eligible, so in_fit is unchanged and no regeneration is needed (build_shipped check mode reports no changes).

[JPPhoto]

Good!

[joshistoast]

lgtm