feat(working-memory): ship VAE working-memory coefficients as data

[Pfannkuchensack]

Summary

VRAM working-memory estimates are hand-fitted constants scattered across the estimators, and most recent working-memory fixes were re-fits of those constants. This PR makes estimates data-backed, starting with the AutoencoderKL family; reservations are unchanged.

• New invokeai/backend/working_memory/:
  • Estimators return a WorkingMemoryEstimate, made of fitted terms (coefficient × scale) and analytic terms (score matrix, image buffers).
  • Each estimator declares an EstimatorSpec.
  • Fitted coefficients ship as JSON under working_memory/data/, together with the measured points that bound them.
  • A data test holds every estimate at or above its measured points, and within 1.3× of them (or the point's tighter max_ratio).
• The five AutoencoderKL-family VAE estimators (SD1.5/SDXL, SD3, CogView4, FLUX.1, FLUX.2) move to this contract.
  • Their shared cuDNN/MIOpen table becomes data: one carrying entry, which the other four borrow via values_from.
  • MPS, XPU and CPU are priced with the cuDNN values through one declared rule (UNMEASURED_PRICED_AS).
• model_on_device(working_memory=estimate) takes an estimate; working_mem_bytes= stays for plain reservations and custom nodes.

Follow-ups build on this: an in-server probe that measures operations in admin calibration runs, and per-dtype VAE coefficients.

Related Issues / Discussions

None.

QA Instructions

• Equivalence with upstream/main: a one-off scratch grid (not committed) compared old and new reservations over 1,427,600 cases, and every one was identical.
  • Inputs: the five estimators × encode/decode × sizes 64–8200 (incl. non-square, non-multiples of 8) × tiles (None/0/negative/8–8192) × fp16/bf16/fp32 × batch 1–3.
  • Environments: 8 devices including device=None × 5 SDPA regimes × HIP on/off, KL vs Tiny.
• Tests:
  • uv run --no-sync pytest tests/backend/working_memory tests/app/invocations/test_flux2_working_memory.py tests/app/invocations/test_vae_working_memory_backend.py tests/app/invocations/test_z_image_tiled_decode.py tests/app/invocations/test_flux_vae_decode_tiling.py tests/backend/util/test_rocm_sdpa_head_dim_guard.py tests/backend/model_manager/load/test_loaded_model.py tests/test_package_data.py: 313 passed, 4 skipped. Each new test file also passes alone.
  • Mutation checks (coefficients changed in memory, each caught):
    • cuDNN k_encode=700 fails 7 points;
    • MIOpen k_decode=3450 fails 2 points;
    • a borrower with its own wrong values fails;
    • a duplicate estimator id fails at import.
• Server check: VAE encode→decode graphs at 1024² bf16 on an RTX 4090 and an RX 9060 XT (ROCm). The new debug line logs the same reservations the old formula gives:
  • RTX 4090 (shipped:conv_backend=cudnn): vae.flux and vae.sd, 2200 / 4400 MB.
  • RX 9060 XT (miopen): 5500 / 7200 MB.
• Not run: docs build. Full suite left to CI.

Review

Resolved findings

• The device_working_mem_gb floor was applied to every point, which made the lower bound vacuous below 3 GiB. The floor now applies only where an analytic term decides.
• Copied coefficients were unguarded. Sharing is now structural (values_from).
• The platform identity now reports MPS/XPU/CPU truthfully instead of calling them cuDNN.
• One resolution owner (SPEC.estimate).
• Registry completeness through ESTIMATOR_MODULES plus a source scan.
• A duplicate-id guard that distinguishes a reload from a collision.
• A WorkingMemoryTerm.fitted() constructor, so a fitted term's bytes always equal coefficient × scale.
• Points that kept a tighter bound in their original tests keep it (max_ratio).

Remaining limits

• Under jurigged --dev_reload, editing a register(...) statement raises; restart the server.

Compatibility / Rollout

• The five estimate_vae_working_memory_* functions (sd15_sdxl, sd3, cogview4, flux, flux2) now return a WorkingMemoryEstimate instead of an int. External callers read .reservation_bytes or pass it as working_memory=.
  • Passing an estimate as working_mem_bytes raises a TypeError that says so.
• working_mem_bytes=<int> is unchanged.
• No persisted state, no API change.

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.}