Add opt-in PiD memory optimizations

docs/src/content/docs/configuration/low-vram-mode.mdx

@@ -32,13 +32,14 @@ enable_partial_loading: false
 
 ## Details and fine-tuning
 
-Low-VRAM mode involves 4 features, each of which can be configured or fine-tuned:
+Low-VRAM mode and related workload-specific optimizations include:
 
 - Partial model loading (`enable_partial_loading`)
 - PyTorch CUDA allocator config (`pytorch_cuda_alloc_conf`)
 - Dynamic RAM and VRAM cache sizes (`max_cache_ram_gb`, `max_cache_vram_gb`)
 - Working memory (`device_working_mem_gb`)
 - Keeping a RAM weight copy (`keep_ram_copy_of_weights`)
+- PiD decode activation chunking (`pid_memory_optimization`)
 
 Read on to learn about these features and understand how to fine-tune them for your system and use-cases.
 
@@ -146,6 +147,22 @@ Invoke has the option of keeping a RAM copy of all model weights, even when they
 keep_ram_copy_of_weights: false
 ```
 
+### PiD decode optimization
+
+PiD decodes directly into high-resolution pixels, so its activation and sampler memory can exceed the working memory needed by a normal VAE decode. Partial model loading reduces memory used by PiD's weights, but does not reduce these full-resolution intermediates.
+
+To reduce PiD's peak VRAM use, enable its experimental memory optimization in `invokeai.yaml` and restart InvokeAI:
+
+```yaml
+pid_memory_optimization: true
+```
+
+This setting processes parts of the PiD pixel pathway in chunks and uses float32 instead of float64 sampler intermediates. It applies to every supported PiD decoder. Measured on an RTX 4090, peak activation memory for a 2048px decode drops from ~3.7 GB to ~1.5 GB (1024px: ~0.9 GB to ~0.5 GB).
+
+It is disabled by default because it is **not output-preserving**: neither change is bit-exact, and the few-step sampler amplifies the difference into a slightly different image. The delta is small — around 43 dB PSNR, visually indistinguishable in side-by-side comparisons — but it is real, so the same seed and workflow will not reproduce an unoptimized decode exactly. Decoding speed is roughly unchanged: chunking costs about 4%, which the cheaper sampler math largely offsets.
+
+When the setting is active, each decode logs the resolution, the patch-token count and whether chunking actually engaged — worth checking, since the option is server-wide and never recorded in image metadata. See [PiD Super-Resolution Decode](/features/pid-decode/) for supported models and usage.
+
 ### Disabling Nvidia sysmem fallback (Windows only)
 
 On Windows, Nvidia GPUs are able to use system RAM when their VRAM fills up via **sysmem fallback**. While it sounds like a good idea on the surface, in practice it causes massive slowdowns during generation.

docs/src/content/docs/features/pid-decode.mdx

@@ -1,6 +1,6 @@
 ---
 title: PiD Super-Resolution Decode
-lastUpdated: 2026-07-01
+lastUpdated: 2026-08-04
 sidebar:
   order: 5
 ---
@@ -69,6 +69,16 @@ PiD is available in both the **Generate** tab (text-to-image) and on the **Canva
 
 ## Tips & limitations
 
+On GPUs with limited VRAM, enable the experimental PiD memory optimizations in `invokeai.yaml`, then restart InvokeAI:
+
+```yaml
+pid_memory_optimization: true
+```
+
+This setting applies to every supported PiD decoder. It reduces peak activation and sampler memory by processing parts of the pixel pathway in chunks and by running the sampler's intermediates in float32 instead of float64 — measured on an RTX 4090, a 2048px decode peaks at ~1.5 GB of activations instead of ~3.7 GB. Decoding speed is roughly unchanged.
+
+It is disabled by default because it **changes the decoded image**. Neither the chunked pathway nor the float32 sampler math is bit-exact with the default path, and the few-step sampler amplifies that into a small but real difference (~43 dB PSNR — visually indistinguishable, numerically not identical). A seed that reproduced an image with the setting off will not reproduce it exactly with the setting on, and because this is a server-wide `invokeai.yaml` option it is not recorded in image metadata — each decode logs that it ran optimized instead.
+
 - **Turn off "Scale Before Processing"** on the Canvas when using PiD — PiD already decodes at 4×, so pre-scaling would inflate the work and is blocked.
 - **Inpaint / Outpaint** are not supported with PiD yet; use text-to-image or image-to-image.
 - **SDXL Refiner** cannot be combined with PiD — disable one of them.

docs/src/generated/settings.json

@@ -561,6 +561,17 @@
       "type": "<class 'bool'>",
       "validation": {}
     },
+    {
+      "category": "GENERATION",
+      "default": false,
+      "description": "Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.",
+      "env_var": "INVOKEAI_PID_MEMORY_OPTIMIZATION",
+      "literal_values": [],
+      "name": "pid_memory_optimization",
+      "required": false,
+      "type": "<class 'bool'>",
+      "validation": {}
+    },
     {
       "category": "GENERATION",
       "default": "auto",

invokeai/app/invocations/flux2_pid_decode.py

@@ -203,7 +203,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         pid_info = context.models.load(self.pid_decoder.decoder)
         # The working-memory estimate scales with the OUTPUT pixel count, so it must see the PACKED latent
         # (spatial H/16), not the unpacked one - otherwise it over-reserves by 4x.
-        estimated_working_memory = estimate_pid_decode_working_memory(packed, BaseModelType.Flux2)
+        # Read once: the estimate and the decode must agree, or the cache reserves headroom for a
+        # peak that will not happen (or too little for one that will).
+        pid_memory_optimization = context.config.get().pid_memory_optimization
+        estimated_working_memory = estimate_pid_decode_working_memory(
+            packed, BaseModelType.Flux2, pid_memory_optimization
+        )
         with pid_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, pid_net):
             if not isinstance(pid_net, PidNet):
                 raise TypeError(f"Expected PidNet for PiD decoder, got {type(pid_net).__name__}.")
@@ -226,7 +231,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                 latent=denorm_latent,
                 caption_embs=caption_embs,
                 caption_mask=caption_mask,
-                config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed),
+                config=PiDDecodeConfig(
+                    num_inference_steps=self.num_inference_steps,
+                    seed=self.seed,
+                    pid_memory_optimization=pid_memory_optimization,
+                ),
             )
 
         TorchDevice.empty_cache()

invokeai/app/invocations/flux_pid_decode.py

@@ -129,7 +129,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 
         # 2) Run PiD decode (the loader already returns a live PidNet).
         pid_info = context.models.load(self.pid_decoder.decoder)
-        estimated_working_memory = estimate_pid_decode_working_memory(latents, BaseModelType.Flux)
+        # Read once: the estimate and the decode must agree, or the cache reserves headroom for a
+        # peak that will not happen (or too little for one that will).
+        pid_memory_optimization = context.config.get().pid_memory_optimization
+        estimated_working_memory = estimate_pid_decode_working_memory(
+            latents, BaseModelType.Flux, pid_memory_optimization
+        )
         with pid_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, pid_net):
             if not isinstance(pid_net, PidNet):
                 raise TypeError(f"Expected PidNet for PiD decoder, got {type(pid_net).__name__}.")
@@ -149,7 +154,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                 latent=denorm_latent,
                 caption_embs=caption_embs,
                 caption_mask=caption_mask,
-                config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed),
+                config=PiDDecodeConfig(
+                    num_inference_steps=self.num_inference_steps,
+                    seed=self.seed,
+                    pid_memory_optimization=pid_memory_optimization,
+                ),
             )
 
         TorchDevice.empty_cache()

invokeai/app/invocations/pid_upscale.py

@@ -173,7 +173,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 
         # 3) Run PiD decode (the loader already returns a live PidNet).
         pid_info = context.models.load(self.pid_decoder.decoder)
-        estimated_working_memory = estimate_pid_decode_working_memory(raw_latent, BaseModelType.Flux)
+        # Read once: the estimate and the decode must agree, or the cache reserves headroom for a
+        # peak that will not happen (or too little for one that will).
+        pid_memory_optimization = context.config.get().pid_memory_optimization
+        estimated_working_memory = estimate_pid_decode_working_memory(
+            raw_latent, BaseModelType.Flux, pid_memory_optimization
+        )
         with pid_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, pid_net):
             if not isinstance(pid_net, PidNet):
                 raise TypeError(f"Expected PidNet for PiD decoder, got {type(pid_net).__name__}.")
@@ -189,7 +194,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                 latent=latent_on_device,
                 caption_embs=caption_embs,
                 caption_mask=caption_mask,
-                config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed),
+                config=PiDDecodeConfig(
+                    num_inference_steps=self.num_inference_steps,
+                    seed=self.seed,
+                    pid_memory_optimization=pid_memory_optimization,
+                ),
             )
 
         TorchDevice.empty_cache()

invokeai/app/invocations/qwen_image_pid_decode.py

@@ -195,7 +195,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 
         # 4) Run PiD decode (the loader already returns a live PidNet).
         pid_info = context.models.load(self.pid_decoder.decoder)
-        estimated_working_memory = estimate_pid_decode_working_memory(latents, BaseModelType.QwenImage)
+        # Read once: the estimate and the decode must agree, or the cache reserves headroom for a
+        # peak that will not happen (or too little for one that will).
+        pid_memory_optimization = context.config.get().pid_memory_optimization
+        estimated_working_memory = estimate_pid_decode_working_memory(
+            latents, BaseModelType.QwenImage, pid_memory_optimization
+        )
         with pid_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, pid_net):
             if not isinstance(pid_net, PidNet):
                 raise TypeError(f"Expected PidNet for PiD decoder, got {type(pid_net).__name__}.")
@@ -215,7 +220,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                 latent=denorm_latent,
                 caption_embs=caption_embs,
                 caption_mask=caption_mask,
-                config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed),
+                config=PiDDecodeConfig(
+                    num_inference_steps=self.num_inference_steps,
+                    seed=self.seed,
+                    pid_memory_optimization=pid_memory_optimization,
+                ),
             )
 
         TorchDevice.empty_cache()

invokeai/app/invocations/sd3_pid_decode.py

@@ -126,7 +126,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         TorchDevice.empty_cache()
 
         pid_info = context.models.load(self.pid_decoder.decoder)
-        estimated_working_memory = estimate_pid_decode_working_memory(latents, BaseModelType.StableDiffusion3)
+        # Read once: the estimate and the decode must agree, or the cache reserves headroom for a
+        # peak that will not happen (or too little for one that will).
+        pid_memory_optimization = context.config.get().pid_memory_optimization
+        estimated_working_memory = estimate_pid_decode_working_memory(
+            latents, BaseModelType.StableDiffusion3, pid_memory_optimization
+        )
         with pid_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, pid_net):
             if not isinstance(pid_net, PidNet):
                 raise TypeError(f"Expected PidNet for PiD decoder, got {type(pid_net).__name__}.")
@@ -142,7 +147,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                 latent=denorm_latent,
                 caption_embs=caption_embs,
                 caption_mask=caption_mask,
-                config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed),
+                config=PiDDecodeConfig(
+                    num_inference_steps=self.num_inference_steps,
+                    seed=self.seed,
+                    pid_memory_optimization=pid_memory_optimization,
+                ),
             )
 
         TorchDevice.empty_cache()

invokeai/app/invocations/sdxl_pid_decode.py

@@ -171,7 +171,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 
         # 3) Run PiD decode (the loader already returns a live PidNet).
         pid_info = context.models.load(self.pid_decoder.decoder)
-        estimated_working_memory = estimate_pid_decode_working_memory(latents, BaseModelType.StableDiffusionXL)
+        # Read once: the estimate and the decode must agree, or the cache reserves headroom for a
+        # peak that will not happen (or too little for one that will).
+        pid_memory_optimization = context.config.get().pid_memory_optimization
+        estimated_working_memory = estimate_pid_decode_working_memory(
+            latents, BaseModelType.StableDiffusionXL, pid_memory_optimization
+        )
         with pid_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, pid_net):
             if not isinstance(pid_net, PidNet):
                 raise TypeError(f"Expected PidNet for PiD decoder, got {type(pid_net).__name__}.")
@@ -188,7 +193,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                 latent=denorm_latent,
                 caption_embs=caption_embs,
                 caption_mask=caption_mask,
-                config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed),
+                config=PiDDecodeConfig(
+                    num_inference_steps=self.num_inference_steps,
+                    seed=self.seed,
+                    pid_memory_optimization=pid_memory_optimization,
+                ),
             )
 
         TorchDevice.empty_cache()

invokeai/app/invocations/z_image_pid_decode.py

@@ -172,7 +172,12 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 
         # 2) Run PiD decode (the loader already returns a live PidNet).
         pid_info = context.models.load(self.pid_decoder.decoder)
-        estimated_working_memory = estimate_pid_decode_working_memory(latents, BaseModelType.Flux)
+        # Read once: the estimate and the decode must agree, or the cache reserves headroom for a
+        # peak that will not happen (or too little for one that will).
+        pid_memory_optimization = context.config.get().pid_memory_optimization
+        estimated_working_memory = estimate_pid_decode_working_memory(
+            latents, BaseModelType.Flux, pid_memory_optimization
+        )
         with pid_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, pid_net):
             if not isinstance(pid_net, PidNet):
                 raise TypeError(f"Expected PidNet for PiD decoder, got {type(pid_net).__name__}.")
@@ -198,7 +203,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                 latent=denorm_latent,
                 caption_embs=caption_embs,
                 caption_mask=caption_mask,
-                config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed),
+                config=PiDDecodeConfig(
+                    num_inference_steps=self.num_inference_steps,
+                    seed=self.seed,
+                    pid_memory_optimization=pid_memory_optimization,
+                ),
             )
             context.logger.info(
                 f"PiD output stats: shape={tuple(x0.shape)} dtype={x0.dtype} "

invokeai/app/services/config/config_default.py

@@ -113,6 +113,7 @@ class InvokeAIAppConfig(BaseSettings):
         device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)
         precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.<br>Valid values: `auto`, `float16`, `bfloat16`, `float32`
         sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.
+        pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.
         attention_type: Attention type.<br>Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`
         attention_slice_size: Slice size, valid when attention_type=="sliced".<br>Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`
         force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).
@@ -220,6 +221,7 @@ class InvokeAIAppConfig(BaseSettings):
 
     # GENERATION
     sequential_guidance:           bool = Field(default=False,              description="Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.")
+    pid_memory_optimization:       bool = Field(default=False,              description="Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.")
     attention_type:      ATTENTION_TYPE = Field(default="auto",             description="Attention type.")
     attention_slice_size: ATTENTION_SLICE_SIZE = Field(default="auto",      description='Slice size, valid when attention_type=="sliced".')
     force_tiled_decode:            bool = Field(default=False,              description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).")

invokeai/backend/pid/_src/networks/pid_net.py

@@ -276,6 +276,7 @@ def forward(
         # --- Feature extraction for GAN discriminator ---
         feature_indices=None,
         return_features_early: bool = False,
+        activation_chunk_size: Optional[int] = None,
     ):
         B, _, H, W = x.shape
         Hs = H // self.patch_size
@@ -449,7 +450,15 @@ def forward(
             x_pixels = split_inputs_cp(x_pixels, seq_dim=1, cp_group=cp_group)
             x_pixels = x_pixels.reshape(B * L_local, P2, self.pixel_hidden_size)
         for blk in self.pixel_blocks:
-            x_pixels = blk(x_pixels, s_cond, H, W, self.patch_size, mask)
+            x_pixels = blk(
+                x_pixels,
+                s_cond,
+                H,
+                W,
+                self.patch_size,
+                mask,
+                activation_chunk_size=activation_chunk_size,
+            )
 
         x_pixels = self.final_layer(x_pixels)  # [B*L_local, P², C_out]
         C_out = self.out_channels