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

docs/src/content/docs/contributing/new-model-integration.mdx

@@ -601,11 +601,11 @@ Packages are discovered automatically. A directory of node modules without an `_
            ...
    ```
 
-   Report progress with `context.util.sd_step_callback(state, BaseModelType.NewModel)`; it renders previews through the architecture's `LatentSpaceFacet`. The denoise node is also where these are checked: interrupts (`context.util.is_canceled()`), inpainting, img2img through `denoising_start`, and a working-memory estimate passed to `model_on_device(working_mem_bytes=...)`.
+   Report progress with `context.util.sd_step_callback(state, BaseModelType.NewModel)`; it renders previews through the architecture's `LatentSpaceFacet`. The denoise node is also where these are checked: interrupts (`context.util.is_canceled()`), inpainting, img2img through `denoising_start`, and a working-memory estimate passed to `model_on_device(working_mem_bytes=<int>)`, or `working_memory=` for an estimator that declares an `EstimatorSpec`.
 
 4. **VAE encode and decode** (`vae/new_model_image_to_latents.py`, `vae/new_model_latents_to_image.py`)
 
-   Apply the VAE's latent normalization, scaling and shift or per-channel mean and std, exactly as the reference pipeline does. Handle tiling and give a working-memory estimate. The estimate helpers live in `invokeai/backend/util/vae_working_memory.py`, with calibration scripts in `scripts/calibrate_*_working_memory.py`. Node types conventionally end in `_i2l` and `_l2i`.
+   Apply the VAE's latent normalization, scaling and shift or per-channel mean and std, exactly as the reference pipeline does. Handle tiling and give a working-memory estimate. The estimate helpers live in `invokeai/backend/util/vae_working_memory.py`, with calibration scripts in `scripts/calibrate_*_working_memory.py`. A new estimator declares an `EstimatorSpec`, returns a `WorkingMemoryEstimate` and takes its fitted coefficients from `invokeai/backend/working_memory/data/`, next to the measured points they were fitted on: the estimate must be at least each point and, where a fitted term decides it, at most 1.3 times it (or the point's `max_ratio`). The [model manager](/development/architecture/model-manager/) lists what such an estimator needs. `working_mem_bytes=<int>` remains for a plain reservation. Node types conventionally end in `_i2l` and `_l2i`.
 
 5. **LoRA loaders** (if LoRA is supported, see [Optional features](#12-optional-features))
 

docs/src/content/docs/development/Architecture/model-manager.mdx

@@ -972,6 +972,26 @@ with loaded_model.model_on_device() as (state_dict, vae):
 
 Since not all models have state dicts, the `state_dict` return value can be None.
 
+An operation that needs more device memory than `device_working_mem_gb` reserves it when it locks the model, so the cache offloads other models instead of running out of memory. Pass the operation's working-memory estimate:
+
+```py
+estimate = estimate_vae_working_memory_flux2(
+    operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device
+)
+with vae_info.model_on_device(working_memory=estimate) as (_, vae):
+    image = vae.decode(latents)[0]
+```
+
+An estimator that declares an `EstimatorSpec` returns a `WorkingMemoryEstimate` (`invokeai/backend/working_memory/estimate.py`): the reservation, the terms it adds up to, the features it was evaluated from, and the source of its coefficients, which `model_on_device()` logs at debug level. A plain byte count goes in as `working_mem_bytes=<int>` instead; never both. The cache never reserves less than `device_working_mem_gb`.
+
+To write such an estimator (`invokeai/backend/util/vae_working_memory.py` has five):
+
+- Register an `EstimatorSpec` at module import, and list the module in `ESTIMATOR_MODULES` (`working_memory/spec.py`); the data test fails on a module that declares a spec without being listed.
+- Split the estimator in two. The impure part reads the model and the device and records every runtime decision (tile, element size, score-matrix bytes) as a feature. The pure `terms(features, values)` turns features and coefficient values into terms, building fitted terms with `WorkingMemoryTerm.fitted(...)`. The public function returns `SPEC.estimate(features, device)`.
+- Bump `version` whenever `terms` changes the formula.
+- Ship one coefficient entry per measured platform value in `invokeai/backend/working_memory/data/`, either with its own `values` or with `values_from` another estimator's entry. Unmeasured platform values (MPS, XPU, CPU) are priced as `UNMEASURED_PRICED_AS` in `working_memory/platform.py` says.
+- Ship the measured points next to them: peak reserved bytes (`reserved_peak`), one fresh process per point. `tests/backend/working_memory/` holds each estimate to at least its point, and where a fitted term decides it, to at most 1.3 times the point, or the point's tighter `max_ratio`. The `device_working_mem_gb` floor counts only where an analytic term decides. Every value an entry carries must be decided by at least one point.
+
 ### Emitting model loading events
 
 When the `context` argument is passed to `load_model_*()`, it will retrieve the invocation event bus from the passed `InvocationContext` object to emit events on the invocation bus. The two events are "model_load_started" and "model_load_completed". Both carry the following payload:

invokeai/app/invocations/vae/cogview4_image_to_latents.py

@@ -43,7 +43,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso
         estimated_working_memory = estimate_vae_working_memory_cogview4(
             operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_info.compute_device
         )
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             assert isinstance(vae, AutoencoderKL)
 
             vae.disable_tiling()

invokeai/app/invocations/vae/cogview4_latents_to_image.py

@@ -50,7 +50,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         )
         with (
             SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes),
-            vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae),
+            vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae),
         ):
             context.util.signal_progress("Running VAE")
             assert isinstance(vae, (AutoencoderKL))

invokeai/app/invocations/vae/flux2_vae_decode.py

@@ -24,6 +24,7 @@
 from invokeai.backend.util.oom import is_oom_error
 from invokeai.backend.util.vae_tiling_scope import MIN_TILE_SAMPLE_SIZE, scoped_vae_tiling
 from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2, should_pretile_vae_decode
+from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate
 
 
 @invocation(
@@ -67,7 +68,7 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents
         # Decoding at FLUX.2 resolutions costs multiple GB of activations (~4.3GB at 1024x1024),
         # far above the default working memory the cache would otherwise reserve. Tell it up front so
         # it offloads enough of the (possibly still resident) transformer to leave room.
-        def estimate(tiled: bool) -> int:
+        def estimate(tiled: bool) -> WorkingMemoryEstimate:
             return estimate_vae_working_memory_flux2(
                 operation="decode",
                 image_tensor=latents,
@@ -82,12 +83,12 @@ def estimate(tiled: bool) -> int:
         if (
             not use_tiling
             and config.auto_tiled_decode
-            and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory)
+            and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory.reservation_bytes)
         ):
             use_tiling = True
             estimated_working_memory = estimate(use_tiling)
 
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             vae_dtype = next(iter(vae.parameters())).dtype
             # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from
             # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would

invokeai/app/invocations/vae/flux2_vae_encode.py

@@ -75,7 +75,7 @@ def _vae_encode(
             device=device,
         )
 
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             vae_dtype = next(iter(vae.parameters())).dtype
             image_tensor = image_tensor.to(device=device, dtype=vae_dtype)
 

invokeai/app/invocations/vae/flux_vae_decode.py

@@ -80,7 +80,7 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents
         if (
             not use_tiling
             and config.auto_tiled_decode
-            and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory)
+            and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory.reservation_bytes)
         ):
             use_tiling = True
             tile_size = self.tile_size
@@ -92,7 +92,7 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents
                 device=vae_info.compute_device,
             )
 
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             assert isinstance(vae, AutoencoderKL)
             vae_dtype = next(iter(vae.parameters())).dtype
             # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from

invokeai/app/invocations/vae/flux_vae_encode.py

@@ -90,7 +90,7 @@ def vae_encode(
         # generator held outside would already have been drawn from by the attempt that failed, so
         # the retry would produce a different latent than a directly tiled run.
         generator = torch.Generator(device=vae_info.compute_device).manual_seed(0)
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             assert isinstance(vae, AutoencoderKL)
             vae_dtype = next(iter(vae.parameters())).dtype
             # The VAE's own device, as in the decode nodes: a cpu_only VAE (or another worker's GPU)

invokeai/app/invocations/vae/ideogram4_latents_to_image.py

@@ -81,7 +81,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             operation="decode", image_tensor=z, vae=vae_info.model, device=device
         )
 
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             context.util.signal_progress("Running VAE")
             vae_dtype = next(vae.parameters()).dtype
             TorchDevice.empty_cache()

invokeai/app/invocations/vae/image_to_latents.py

@@ -83,7 +83,7 @@ def vae_encode(
             fp32=upcast,
             device=vae_info.compute_device,
         )
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)), "VAE must be of type SD-1.5 or SDXL"
             orig_dtype = vae.dtype
             if upcast:

invokeai/app/invocations/vae/latents_to_image.py

@@ -65,7 +65,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         )
         with (
             SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes),
-            vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae),
+            vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae),
         ):
             context.util.signal_progress("Running VAE decoder")
             assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))

invokeai/app/invocations/vae/sd3_image_to_latents.py

@@ -39,7 +39,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso
         estimated_working_memory = estimate_vae_working_memory_sd3(
             operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_info.compute_device
         )
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             assert isinstance(vae, AutoencoderKL)
 
             vae.disable_tiling()

invokeai/app/invocations/vae/sd3_latents_to_image.py

@@ -52,7 +52,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         )
         with (
             SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes),
-            vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae),
+            vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae),
         ):
             context.util.signal_progress("Running VAE")
             assert isinstance(vae, (AutoencoderKL))

invokeai/app/invocations/vae/z_image_image_to_latents.py

@@ -72,7 +72,7 @@ def vae_encode(
             device=vae_info.compute_device,
         )
 
-        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             if not isinstance(vae, AutoencoderKL):
                 raise TypeError(
                     f"Expected AutoencoderKL, got {type(vae).__name__}. VAE model type changed unexpectedly after loading."

invokeai/app/invocations/vae/z_image_latents_to_image.py

@@ -74,7 +74,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         if (
             not use_tiling
             and config.auto_tiled_decode
-            and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory)
+            and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory.reservation_bytes)
         ):
             use_tiling = True
             estimated_working_memory = estimate_vae_working_memory_flux(
@@ -90,7 +90,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         # FLUX VAE now honours the axes the UI has always offered for it.
         seamless_context = SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes)
 
-        with seamless_context, vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+        with seamless_context, vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
             context.util.signal_progress("Running VAE")
             if not isinstance(vae, AutoencoderKL):
                 raise TypeError(

invokeai/backend/AGENTS.md

@@ -6,6 +6,7 @@
 - Use existing device/model-management abstractions, `util/devices.py` included. Keep CPU/MPS/ROCm/XPU/multi-device support. No hardcoded CUDA, GPU zero, or dtype assumptions.
 - Respect session/thread device ownership. Restore temporary model patches, device state, hooks, allocations on success, cancellation, and exceptions. No request-specific mutable globals.
 - Define model/cache lifetimes and working-memory estimates. Bound queues, caches, downloads, batches, buffers. Consider peak memory and throughput.
+- New or refitted working-memory coefficients go in `working_memory/data/` with the measured points that bound them; their estimators declare an `EstimatorSpec` and return a `WorkingMemoryEstimate`.
 - Keep model validation/loading safeguards, vendored exclusions, upstream provenance. No unrelated vendored rewrites or formatting.
 
 ## Performance

invokeai/backend/flux/extensions/kontext_extension.py

@@ -138,7 +138,7 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
                 operation="encode", image_tensor=image_tensor, vae=vae_info.model
             )
 
-            with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+            with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
                 assert isinstance(vae, AutoencoderKL)
                 vae_dtype = next(iter(vae.parameters())).dtype
                 image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)

invokeai/backend/flux2/ref_image_extension.py

@@ -218,7 +218,7 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]:
                 device=TorchDevice.choose_torch_device(),
             )
 
-            with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
+            with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
                 vae_dtype = next(iter(vae.parameters())).dtype
                 image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
 

invokeai/backend/model_manager/load/load_base.py

@@ -25,6 +25,8 @@
     ModelCache,
 )
 from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType
+from invokeai.backend.util.logging import InvokeAILogger
+from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate
 
 
 class LoadedModelWithoutConfig:
@@ -194,13 +196,27 @@ def __exit__(self, *args: Any, **kwargs: Any) -> None:
 
     @contextmanager
     def model_on_device(
-        self, working_mem_bytes: Optional[int] = None
+        self, working_mem_bytes: Optional[int] = None, *, working_memory: Optional[WorkingMemoryEstimate] = None
     ) -> Generator[Tuple[Optional[Dict[str, torch.Tensor]], AnyModel], None, None]:
         """Return a tuple consisting of the model's state dict (if it exists) and the locked model on execution device.
 
         :param working_mem_bytes: The amount of working memory to keep available on the compute device when loading the
             model.
+        :param working_memory: A working-memory estimate whose `reservation_bytes` is kept available instead; see
+            `invokeai.backend.working_memory`. Pass this or `working_mem_bytes`, not both.
         """
+        if isinstance(working_mem_bytes, WorkingMemoryEstimate):
+            raise TypeError(
+                "model_on_device() got a WorkingMemoryEstimate as working_mem_bytes; pass it as working_memory=."
+            )
+        if working_memory is not None:
+            if working_mem_bytes is not None:
+                raise ValueError("Pass either working_mem_bytes or working_memory to model_on_device(), not both.")
+            working_mem_bytes = working_memory.reservation_bytes
+            InvokeAILogger.get_logger(__name__).debug(
+                f"Working memory for {self._cache_record.key}: {working_memory.estimator} reserves "
+                f"{working_mem_bytes / 2**20:.0f} MB ({working_memory.source})"
+            )
         self._lock_paced(working_mem_bytes)
         self._end_first_use_window()
         try: