From 21ddd778f32fa4702538f534a8c4d02952a8b0ee Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 4 Aug 2026 16:50:48 -0500 Subject: [PATCH 1/7] perf(pid): chunk pixel activations --- .../pid/_src/networks/pixeldit_official.py | 89 +++++++++++++++++-- invokeai/backend/pid/decode.py | 6 +- tests/backend/pid/test_pid_decode.py | 29 +++++- tests/backend/pid/test_pixeldit_official.py | 79 ++++++++++++++++ 4 files changed, 193 insertions(+), 10 deletions(-) create mode 100644 tests/backend/pid/test_pixeldit_official.py diff --git a/invokeai/backend/pid/_src/networks/pixeldit_official.py b/invokeai/backend/pid/_src/networks/pixeldit_official.py index 6fdda4917db..6e319ca3b01 100644 --- a/invokeai/backend/pid/_src/networks/pixeldit_official.py +++ b/invokeai/backend/pid/_src/networks/pixeldit_official.py @@ -426,8 +426,11 @@ def __init__( rope_mode: str = "original", rope_ref_grid_h: int = 32, rope_ref_grid_w: int = 32, + activation_chunk_size: Optional[int] = 1024, ): super().__init__() + if activation_chunk_size is not None and activation_chunk_size <= 0: + raise ValueError("activation_chunk_size must be positive when set") self.pixel_dim = int(pixel_hidden_size) self.context_dim = int(patch_hidden_size) self.patch_size = int(patch_size) @@ -436,6 +439,7 @@ def __init__( self.rope_mode = rope_mode self.rope_ref_grid_h = rope_ref_grid_h self.rope_ref_grid_w = rope_ref_grid_w + self.activation_chunk_size = activation_chunk_size assert self.attn_dim % self.num_heads == 0, "pixel attention hidden size must be divisible by pixel num_heads" p2 = self.patch_size * self.patch_size self.compress_to_attn = nn.Linear(p2 * self.pixel_dim, self.attn_dim, bias=True) @@ -489,25 +493,98 @@ def forward( assert s_cond.shape[0] == BL, "s_cond batch must match x batch" assert BL % L_local == 0, "Total sequences must be a multiple of local patch count" B = BL // L_local + + # Chunking bounds full-resolution AdaLN and MLP intermediates during inference. Attention remains global: + # compressed tokens are assembled in original order before the unchanged attention call. Keep the original + # path while gradients are enabled because copy-based output assembly is inference-only. + chunk_size = self.activation_chunk_size + if chunk_size is not None and BL > chunk_size and not torch.is_grad_enabled(): + return self._forward_chunked(x, s_cond, B, L_local, Hs, Ws, mask, chunk_size) + + return self._forward_unchunked(x, s_cond, B, L_local, Hs, Ws, mask) + + def _forward_unchunked( + self, + x: torch.Tensor, + s_cond: torch.Tensor, + batch_size: int, + local_patch_count: int, + patch_grid_height: int, + patch_grid_width: int, + mask: Optional[torch.Tensor], + ) -> torch.Tensor: + patch_batch, pixels_per_patch, _ = x.shape # adaLN per pixel (within patch): params cond_params = self.adaLN_modulation(s_cond) # [BL, 6*pixel_dim*P2] - cond_params = cond_params.view(BL, P2, 6 * self.pixel_dim) + cond_params = cond_params.view(patch_batch, pixels_per_patch, 6 * self.pixel_dim) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(cond_params, 6, dim=-1) x_norm = apply_adaln(self.norm1(x), shift_msa, scale_msa) - x_flat = x_norm.view(BL, P2 * self.pixel_dim) - x_comp = self.compress_to_attn(x_flat).view(B, L_local, self.attn_dim) + x_flat = x_norm.view(patch_batch, pixels_per_patch * self.pixel_dim) + x_comp = self.compress_to_attn(x_flat).view(batch_size, local_patch_count, self.attn_dim) # attention across patch tokens (L) — pos is full-length; the CP-aware # RotaryAttention gathers k/v across CP ranks internally. - pos_comp = self._fetch_pos(Hs, Ws, x.device) + pos_comp = self._fetch_pos(patch_grid_height, patch_grid_width, x.device) attn_out = self.attn(x_comp, pos_comp, mask) # [B, L_local, attn_dim] - attn_flat = self.expand_from_attn(attn_out.view(B * L_local, self.attn_dim)) - attn_exp = attn_flat.view(BL, P2, self.pixel_dim) + attn_flat = self.expand_from_attn(attn_out.view(batch_size * local_patch_count, self.attn_dim)) + attn_exp = attn_flat.view(patch_batch, pixels_per_patch, self.pixel_dim) # residual & MLP locally x = x + gate_msa * attn_exp mlp_out = self.mlp(apply_adaln(self.norm2(x), shift_mlp, scale_mlp)) x = x + gate_mlp * mlp_out return x + def _compress_activation_chunk(self, x: torch.Tensor, s_cond: torch.Tensor) -> torch.Tensor: + patch_batch, pixels_per_patch, _ = x.shape + cond_params = self.adaLN_modulation(s_cond).view(patch_batch, pixels_per_patch, 6 * self.pixel_dim) + shift_msa, scale_msa = torch.chunk(cond_params, 6, dim=-1)[:2] + x_norm = apply_adaln(self.norm1(x), shift_msa, scale_msa) + return self.compress_to_attn(x_norm.view(patch_batch, pixels_per_patch * self.pixel_dim)) + + def _finish_activation_chunk( + self, x: torch.Tensor, s_cond: torch.Tensor, attention_output: torch.Tensor + ) -> torch.Tensor: + patch_batch, pixels_per_patch, _ = x.shape + cond_params = self.adaLN_modulation(s_cond).view(patch_batch, pixels_per_patch, 6 * self.pixel_dim) + _, _, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(cond_params, 6, dim=-1) + attention_expanded = self.expand_from_attn(attention_output).view(patch_batch, pixels_per_patch, self.pixel_dim) + x = x + gate_msa * attention_expanded + return x + gate_mlp * self.mlp(apply_adaln(self.norm2(x), shift_mlp, scale_mlp)) + + def _forward_chunked( + self, + x: torch.Tensor, + s_cond: torch.Tensor, + batch_size: int, + local_patch_count: int, + patch_grid_height: int, + patch_grid_width: int, + mask: Optional[torch.Tensor], + chunk_size: int, + ) -> torch.Tensor: + patch_batch = x.shape[0] + compressed: Optional[torch.Tensor] = None + for start in range(0, patch_batch, chunk_size): + end = min(start + chunk_size, patch_batch) + compressed_chunk = self._compress_activation_chunk(x[start:end], s_cond[start:end]) + if compressed is None: + compressed = compressed_chunk.new_empty((patch_batch, self.attn_dim)) + compressed[start:end].copy_(compressed_chunk) + assert compressed is not None + + compressed = compressed.view(batch_size, local_patch_count, self.attn_dim) + pos_comp = self._fetch_pos(patch_grid_height, patch_grid_width, x.device) + attention_output = self.attn(compressed, pos_comp, mask).view(patch_batch, self.attn_dim) + + output: Optional[torch.Tensor] = None + for start in range(0, patch_batch, chunk_size): + end = min(start + chunk_size, patch_batch) + output_chunk = self._finish_activation_chunk(x[start:end], s_cond[start:end], attention_output[start:end]) + if output is None: + output = output_chunk.new_empty(x.shape) + output[start:end].copy_(output_chunk) + assert output is not None + return output + # ============================================================================= # From pixdit_core/pixeldit_t2i.py diff --git a/invokeai/backend/pid/decode.py b/invokeai/backend/pid/decode.py index 8c51504d43a..dea038ca664 100644 --- a/invokeai/backend/pid/decode.py +++ b/invokeai/backend/pid/decode.py @@ -226,10 +226,10 @@ def _get_t_list(device: torch.device, *, num_steps: Optional[int] = None) -> Ten def _velocity_to_x0(x_t: Tensor, net_output: Tensor, t: Tensor) -> Tensor: - """Convert the network's velocity prediction back to x0 at time *t*.""" + """Convert the network's velocity prediction back to x0 at time *t* using float32 intermediates.""" s = [x_t.shape[0]] + [1] * (x_t.ndim - 1) - t_shaped = t.double().view(*s) - return (x_t.double() - t_shaped * net_output.double()).to(x_t.dtype) + t_shaped = t.float().view(*s) + return torch.addcmul(x_t.float(), net_output.float(), t_shaped, value=-1).to(x_t.dtype) @torch.no_grad() diff --git a/tests/backend/pid/test_pid_decode.py b/tests/backend/pid/test_pid_decode.py index 44e8f377617..26f0c1a29cf 100644 --- a/tests/backend/pid/test_pid_decode.py +++ b/tests/backend/pid/test_pid_decode.py @@ -1,10 +1,12 @@ """Regression tests for the PiD distill schedule and decoder/base validation.""" +from unittest.mock import patch + import pytest import torch from invokeai.backend.model_manager.taxonomy import BaseModelType -from invokeai.backend.pid.decode import _get_t_list, assert_pid_decoder_matches_base +from invokeai.backend.pid.decode import _get_t_list, _velocity_to_x0, assert_pid_decoder_matches_base _CPU = torch.device("cpu") @@ -35,6 +37,31 @@ def test_out_of_range_step_count_trips_the_safety_net() -> None: _get_t_list(_CPU, num_steps=5) +@pytest.mark.parametrize( + ("x_dtype", "net_output_dtype"), + [ + (torch.float32, torch.float32), + (torch.float32, torch.bfloat16), + (torch.bfloat16, torch.bfloat16), + ], +) +def test_velocity_to_x0_uses_float32_math_and_preserves_input_dtype( + x_dtype: torch.dtype, net_output_dtype: torch.dtype +) -> None: + x_t = torch.tensor([[[[1.0, -2.0], [3.0, -4.0]]]], dtype=x_dtype) + net_output = torch.tensor([[[[0.5, -0.25], [0.125, -0.0625]]]], dtype=net_output_dtype) + timestep = torch.tensor([0.634], dtype=torch.float32) + expected = (x_t.float() - timestep.view(1, 1, 1, 1) * net_output.float()).to(x_dtype) + + # A float64 intermediate doubles memory for each full-resolution sampler tensor. PiD already + # predicts under bf16 autocast, so perform this update in float32 without calling Tensor.double(). + with patch.object(torch.Tensor, "double", side_effect=AssertionError("unexpected float64 conversion")): + actual = _velocity_to_x0(x_t, net_output, timestep) + + assert actual.dtype == x_dtype + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_matching_decoder_base_is_accepted() -> None: assert_pid_decoder_matches_base(BaseModelType.Flux, BaseModelType.Flux, node_title="FLUX PiD Decode") diff --git a/tests/backend/pid/test_pixeldit_official.py b/tests/backend/pid/test_pixeldit_official.py new file mode 100644 index 00000000000..f5c7a254ea6 --- /dev/null +++ b/tests/backend/pid/test_pixeldit_official.py @@ -0,0 +1,79 @@ +import math + +import pytest +import torch + +from invokeai.backend.pid._src.networks.pixeldit_official import PiTBlock + + +def _build_pit_block(activation_chunk_size: int | None) -> PiTBlock: + return PiTBlock( + pixel_hidden_size=4, + patch_hidden_size=8, + patch_size=2, + num_heads=2, + mlp_ratio=2.0, + attn_hidden_size=8, + attn_num_heads=2, + activation_chunk_size=activation_chunk_size, + ).eval() + + +@pytest.mark.parametrize("use_autocast", [False, True]) +def test_pit_block_chunked_forward_matches_unchunked_and_bounds_adaln_batch(use_autocast: bool) -> None: + torch.manual_seed(0) + unchunked = _build_pit_block(activation_chunk_size=None) + chunk_size = 3 + chunked = _build_pit_block(activation_chunk_size=chunk_size) + chunked.load_state_dict(unchunked.state_dict()) + + batch_size = 2 + image_height = 4 + image_width = 4 + patch_size = 2 + patch_count = image_height * image_width // patch_size**2 + patch_batch = batch_size * patch_count + pixels = torch.randn(patch_batch, patch_size**2, 4) + condition = torch.randn(patch_batch, 8) + + adaln_batch_sizes: list[int] = [] + + def record_adaln_batch_size(_module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]) -> None: + adaln_batch_sizes.append(inputs[0].shape[0]) + + handle = chunked.adaLN_modulation.register_forward_pre_hook(record_adaln_batch_size) + try: + with torch.no_grad(), torch.autocast("cpu", dtype=torch.bfloat16, enabled=use_autocast): + expected = unchunked(pixels, condition, image_height, image_width, patch_size) + actual = chunked(pixels, condition, image_height, image_width, patch_size) + finally: + handle.remove() + + torch.testing.assert_close(actual, expected) + assert max(adaln_batch_sizes) <= chunk_size + assert len(adaln_batch_sizes) == 2 * math.ceil(patch_batch / chunk_size) + + +def test_pit_block_rejects_non_positive_activation_chunk_size() -> None: + with pytest.raises(ValueError, match="activation_chunk_size must be positive"): + _build_pit_block(activation_chunk_size=0) + + +def test_pit_block_uses_unchunked_path_when_gradients_are_enabled() -> None: + block = _build_pit_block(activation_chunk_size=1) + pixels = torch.randn(4, 4, 4, requires_grad=True) + condition = torch.randn(4, 8) + adaln_batch_sizes: list[int] = [] + + def record_adaln_batch_size(_module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]) -> None: + adaln_batch_sizes.append(inputs[0].shape[0]) + + handle = block.adaLN_modulation.register_forward_pre_hook(record_adaln_batch_size) + try: + output = block(pixels, condition, image_height=4, image_width=4, patch_size=2) + output.sum().backward() + finally: + handle.remove() + + assert adaln_batch_sizes == [pixels.shape[0]] + assert pixels.grad is not None From 9b711afd8853ae904c5f1f63670c53bc39be1a20 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 4 Aug 2026 17:14:13 -0500 Subject: [PATCH 2/7] perf(pid): gate memory optimizations --- .../docs/configuration/low-vram-mode.mdx | 15 ++++- docs/src/content/docs/features/pid-decode.mdx | 10 ++- docs/src/generated/settings.json | 11 ++++ invokeai/app/invocations/flux2_pid_decode.py | 6 +- invokeai/app/invocations/flux_pid_decode.py | 6 +- invokeai/app/invocations/pid_upscale.py | 6 +- .../app/invocations/qwen_image_pid_decode.py | 6 +- invokeai/app/invocations/sd3_pid_decode.py | 6 +- invokeai/app/invocations/sdxl_pid_decode.py | 6 +- .../app/invocations/z_image_pid_decode.py | 6 +- .../app/services/config/config_default.py | 2 + invokeai/backend/pid/_src/networks/pid_net.py | 11 +++- .../pid/_src/networks/pixeldit_official.py | 17 ++++-- invokeai/backend/pid/decode.py | 21 +++++-- .../frontend/web/src/services/api/schema.ts | 7 +++ tests/backend/pid/test_pid_decode.py | 50 ++++++++++++++- tests/backend/pid/test_pixeldit_official.py | 61 ++++++++++++++++--- tests/test_config.py | 9 +++ 18 files changed, 224 insertions(+), 32 deletions(-) diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index fa15cef8735..9511d5e9f87 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/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_optimization`) Read on to learn about these features and understand how to fine-tune them for your system and use-cases. @@ -146,6 +147,18 @@ 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_optimization: true +``` + +This setting uses lower-precision sampler intermediates and processes parts of the PiD pixel pathway in chunks. It applies to every supported PiD decoder and is disabled by default because chunking makes decoding slower. 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. diff --git a/docs/src/content/docs/features/pid-decode.mdx b/docs/src/content/docs/features/pid-decode.mdx index 3897844f75c..52aed18675e 100644 --- a/docs/src/content/docs/features/pid-decode.mdx +++ b/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,14 @@ 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_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, at the cost of slower decoding. It is disabled by default. + - **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. diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index 2c183f55400..fe9cd8a7040 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -561,6 +561,17 @@ "type": "", "validation": {} }, + { + "category": "GENERATION", + "default": false, + "description": "Enable experimental PiD decode memory optimizations at the cost of slower decoding.", + "env_var": "INVOKEAI_PID_OPTIMIZATION", + "literal_values": [], + "name": "pid_optimization", + "required": false, + "type": "", + "validation": {} + }, { "category": "GENERATION", "default": "auto", diff --git a/invokeai/app/invocations/flux2_pid_decode.py b/invokeai/app/invocations/flux2_pid_decode.py index 0d343952e42..4569f75a823 100644 --- a/invokeai/app/invocations/flux2_pid_decode.py +++ b/invokeai/app/invocations/flux2_pid_decode.py @@ -226,7 +226,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_optimization=context.config.get().pid_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/flux_pid_decode.py b/invokeai/app/invocations/flux_pid_decode.py index 182bcc9cdaf..ab0efb61c46 100644 --- a/invokeai/app/invocations/flux_pid_decode.py +++ b/invokeai/app/invocations/flux_pid_decode.py @@ -149,7 +149,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_optimization=context.config.get().pid_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/pid_upscale.py b/invokeai/app/invocations/pid_upscale.py index 3a50ffd4cf3..dba4d5ac59a 100644 --- a/invokeai/app/invocations/pid_upscale.py +++ b/invokeai/app/invocations/pid_upscale.py @@ -189,7 +189,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_optimization=context.config.get().pid_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/qwen_image_pid_decode.py b/invokeai/app/invocations/qwen_image_pid_decode.py index 3503da8bbd6..81f09c30df7 100644 --- a/invokeai/app/invocations/qwen_image_pid_decode.py +++ b/invokeai/app/invocations/qwen_image_pid_decode.py @@ -215,7 +215,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_optimization=context.config.get().pid_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/sd3_pid_decode.py b/invokeai/app/invocations/sd3_pid_decode.py index b22b0fa3bf1..b60e9b92fe5 100644 --- a/invokeai/app/invocations/sd3_pid_decode.py +++ b/invokeai/app/invocations/sd3_pid_decode.py @@ -142,7 +142,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_optimization=context.config.get().pid_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/sdxl_pid_decode.py b/invokeai/app/invocations/sdxl_pid_decode.py index 5a09d03fe0e..426678afd4e 100644 --- a/invokeai/app/invocations/sdxl_pid_decode.py +++ b/invokeai/app/invocations/sdxl_pid_decode.py @@ -188,7 +188,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_optimization=context.config.get().pid_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/z_image_pid_decode.py b/invokeai/app/invocations/z_image_pid_decode.py index 392f11e3639..06e29ada463 100644 --- a/invokeai/app/invocations/z_image_pid_decode.py +++ b/invokeai/app/invocations/z_image_pid_decode.py @@ -198,7 +198,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_optimization=context.config.get().pid_optimization, + ), ) context.logger.info( f"PiD output stats: shape={tuple(x0.shape)} dtype={x0.dtype} " diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index ae1e38e0ccc..46d7bfb3044 100644 --- a/invokeai/app/services/config/config_default.py +++ b/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.
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.
Valid values: `auto`, `float16`, `bfloat16`, `float32` sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements. + pid_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding. attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` attention_slice_size: Slice size, valid when attention_type=="sliced".
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_optimization: bool = Field(default=False, description="Enable experimental PiD decode memory optimizations at the cost of slower decoding.") 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).") diff --git a/invokeai/backend/pid/_src/networks/pid_net.py b/invokeai/backend/pid/_src/networks/pid_net.py index 290ccd50a3d..c8b24d77603 100644 --- a/invokeai/backend/pid/_src/networks/pid_net.py +++ b/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 diff --git a/invokeai/backend/pid/_src/networks/pixeldit_official.py b/invokeai/backend/pid/_src/networks/pixeldit_official.py index 6e319ca3b01..4318d836bca 100644 --- a/invokeai/backend/pid/_src/networks/pixeldit_official.py +++ b/invokeai/backend/pid/_src/networks/pixeldit_official.py @@ -426,11 +426,8 @@ def __init__( rope_mode: str = "original", rope_ref_grid_h: int = 32, rope_ref_grid_w: int = 32, - activation_chunk_size: Optional[int] = 1024, ): super().__init__() - if activation_chunk_size is not None and activation_chunk_size <= 0: - raise ValueError("activation_chunk_size must be positive when set") self.pixel_dim = int(pixel_hidden_size) self.context_dim = int(patch_hidden_size) self.patch_size = int(patch_size) @@ -439,7 +436,6 @@ def __init__( self.rope_mode = rope_mode self.rope_ref_grid_h = rope_ref_grid_h self.rope_ref_grid_w = rope_ref_grid_w - self.activation_chunk_size = activation_chunk_size assert self.attn_dim % self.num_heads == 0, "pixel attention hidden size must be divisible by pixel num_heads" p2 = self.patch_size * self.patch_size self.compress_to_attn = nn.Linear(p2 * self.pixel_dim, self.attn_dim, bias=True) @@ -472,7 +468,14 @@ def _fetch_pos(self, height: int, width: int, device): return pos def forward( - self, x: torch.Tensor, s_cond: torch.Tensor, image_height: int, image_width: int, patch_size: int, mask=None + self, + x: torch.Tensor, + s_cond: torch.Tensor, + image_height: int, + image_width: int, + patch_size: int, + mask=None, + activation_chunk_size: Optional[int] = None, ) -> torch.Tensor: # x: [B*L_local, P2, C]; under CP, L_local = (Hs*Ws)/cp_size. Without CP, # L_local == L_full. The reshape uses L_local for the (B, L_local, ...) @@ -497,7 +500,9 @@ def forward( # Chunking bounds full-resolution AdaLN and MLP intermediates during inference. Attention remains global: # compressed tokens are assembled in original order before the unchanged attention call. Keep the original # path while gradients are enabled because copy-based output assembly is inference-only. - chunk_size = self.activation_chunk_size + if activation_chunk_size is not None and activation_chunk_size <= 0: + raise ValueError("activation_chunk_size must be positive when set") + chunk_size = activation_chunk_size if chunk_size is not None and BL > chunk_size and not torch.is_grad_enabled(): return self._forward_chunked(x, s_cond, B, L_local, Hs, Ws, mask, chunk_size) diff --git a/invokeai/backend/pid/decode.py b/invokeai/backend/pid/decode.py index dea038ca664..4c854e3dcf2 100644 --- a/invokeai/backend/pid/decode.py +++ b/invokeai/backend/pid/decode.py @@ -28,6 +28,8 @@ from invokeai.backend.model_manager.taxonomy import BaseModelType from invokeai.backend.pid._src.networks.pid_net import PidNet +_PID_ACTIVATION_CHUNK_SIZE = 1024 + # --------------------------------------------------------------------------- # Network hyperparameters per backbone # --------------------------------------------------------------------------- @@ -225,11 +227,14 @@ def _get_t_list(device: torch.device, *, num_steps: Optional[int] = None) -> Ten return t -def _velocity_to_x0(x_t: Tensor, net_output: Tensor, t: Tensor) -> Tensor: - """Convert the network's velocity prediction back to x0 at time *t* using float32 intermediates.""" +def _velocity_to_x0(x_t: Tensor, net_output: Tensor, t: Tensor, *, pid_optimization: bool = False) -> Tensor: + """Convert the network's velocity prediction back to x0 at time *t*.""" s = [x_t.shape[0]] + [1] * (x_t.ndim - 1) - t_shaped = t.float().view(*s) - return torch.addcmul(x_t.float(), net_output.float(), t_shaped, value=-1).to(x_t.dtype) + if pid_optimization: + t_shaped = t.float().view(*s) + return torch.addcmul(x_t.float(), net_output.float(), t_shaped, value=-1).to(x_t.dtype) + t_shaped = t.double().view(*s) + return (x_t.double() - t_shaped * net_output.double()).to(x_t.dtype) @torch.no_grad() @@ -245,6 +250,7 @@ def _student_sample_loop( sample_type: str = "sde", autocast_dtype: Optional[torch.dtype] = None, generator: Optional[torch.Generator] = None, + pid_optimization: bool = False, ) -> Tensor: """Few-step distilled sampler. @@ -279,9 +285,10 @@ def _student_sample_loop( lq_video_or_image=None, lq_latent=lq_latent, degrade_sigma=degrade_sigma, + activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE if pid_optimization else None, ) if t_next.item() > 0: - x0_pred = _velocity_to_x0(x, v_pred, t_cur_batch) + x0_pred = _velocity_to_x0(x, v_pred, t_cur_batch, pid_optimization=pid_optimization) eps_infer = torch.randn( x0_pred.shape, device=x0_pred.device, @@ -297,7 +304,7 @@ def _student_sample_loop( else: x = (1.0 - t_next_b) * x0_pred + t_next_b * eps_infer else: - x = _velocity_to_x0(x, v_pred, t_cur_batch) + x = _velocity_to_x0(x, v_pred, t_cur_batch, pid_optimization=pid_optimization) return x @@ -322,6 +329,7 @@ class PiDDecodeConfig: # from_clean upscale path passes the LDM scheduler's per-step sigma here. degrade_sigma: float | list[float] | Tensor = 0.0 seed: int = 0 + pid_optimization: bool = False student_t_list: list[float] = field(default_factory=lambda: list(_STUDENT_T_LIST)) @@ -427,6 +435,7 @@ def decode( sample_type=cfg.sample_type, autocast_dtype=autocast_dtype, generator=gen, + pid_optimization=cfg.pid_optimization, ) return x0.clamp(-1, 1) diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 2928b7170d6..9cd8d2b3924 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -18918,6 +18918,7 @@ export type components = { * device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
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.
Valid values: `auto`, `float16`, `bfloat16`, `float32` * sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements. + * pid_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding. * attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` * attention_slice_size: Slice size, valid when attention_type=="sliced".
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). @@ -19259,6 +19260,12 @@ export type components = { * @default false */ sequential_guidance?: boolean; + /** + * Pid Optimization + * @description Enable experimental PiD decode memory optimizations at the cost of slower decoding. + * @default false + */ + pid_optimization?: boolean; /** * Attention Type * @description Attention type. diff --git a/tests/backend/pid/test_pid_decode.py b/tests/backend/pid/test_pid_decode.py index 26f0c1a29cf..3c51bf4b3eb 100644 --- a/tests/backend/pid/test_pid_decode.py +++ b/tests/backend/pid/test_pid_decode.py @@ -6,7 +6,13 @@ import torch from invokeai.backend.model_manager.taxonomy import BaseModelType -from invokeai.backend.pid.decode import _get_t_list, _velocity_to_x0, assert_pid_decoder_matches_base +from invokeai.backend.pid.decode import ( + PiDDecodeConfig, + _get_t_list, + _student_sample_loop, + _velocity_to_x0, + assert_pid_decoder_matches_base, +) _CPU = torch.device("cpu") @@ -56,12 +62,52 @@ def test_velocity_to_x0_uses_float32_math_and_preserves_input_dtype( # A float64 intermediate doubles memory for each full-resolution sampler tensor. PiD already # predicts under bf16 autocast, so perform this update in float32 without calling Tensor.double(). with patch.object(torch.Tensor, "double", side_effect=AssertionError("unexpected float64 conversion")): - actual = _velocity_to_x0(x_t, net_output, timestep) + actual = _velocity_to_x0(x_t, net_output, timestep, pid_optimization=True) assert actual.dtype == x_dtype torch.testing.assert_close(actual, expected, rtol=0, atol=0) +def test_velocity_to_x0_uses_original_float64_math_when_optimization_is_disabled() -> None: + x_t = torch.tensor([[[[1.0, -2.0], [3.0, -4.0]]]], dtype=torch.bfloat16) + net_output = torch.tensor([[[[0.5, -0.25], [0.125, -0.0625]]]], dtype=torch.bfloat16) + timestep = torch.tensor([0.634], dtype=torch.float32) + expected = (x_t.double() - timestep.double().view(1, 1, 1, 1) * net_output.double()).to(x_t.dtype) + + with patch("torch.addcmul", side_effect=AssertionError("unexpected optimized path")): + actual = _velocity_to_x0(x_t, net_output, timestep, pid_optimization=False) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize(("pid_optimization", "expected_chunk_size"), [(False, None), (True, 1024)]) +def test_student_sample_loop_passes_per_call_activation_chunk_size( + pid_optimization: bool, expected_chunk_size: int | None +) -> None: + activation_chunk_sizes: list[int | None] = [] + + def net(x: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor: + activation_chunk_sizes.append(kwargs.get("activation_chunk_size")) # type: ignore[arg-type] + return torch.zeros_like(x) + + _student_sample_loop( + net, # type: ignore[arg-type] + noise=torch.zeros(1, 3, 2, 2), + t_list=torch.tensor([0.999, 0.0]), + caption_embs=torch.zeros(1, 1, 1), + caption_mask=None, + lq_latent=None, + degrade_sigma=torch.zeros(1), + pid_optimization=pid_optimization, + ) + + assert activation_chunk_sizes == [expected_chunk_size] + + +def test_pid_optimization_defaults_to_disabled() -> None: + assert PiDDecodeConfig().pid_optimization is False + + def test_matching_decoder_base_is_accepted() -> None: assert_pid_decoder_matches_base(BaseModelType.Flux, BaseModelType.Flux, node_title="FLUX PiD Decode") diff --git a/tests/backend/pid/test_pixeldit_official.py b/tests/backend/pid/test_pixeldit_official.py index f5c7a254ea6..57cff37f337 100644 --- a/tests/backend/pid/test_pixeldit_official.py +++ b/tests/backend/pid/test_pixeldit_official.py @@ -6,7 +6,7 @@ from invokeai.backend.pid._src.networks.pixeldit_official import PiTBlock -def _build_pit_block(activation_chunk_size: int | None) -> PiTBlock: +def _build_pit_block() -> PiTBlock: return PiTBlock( pixel_hidden_size=4, patch_hidden_size=8, @@ -15,16 +15,15 @@ def _build_pit_block(activation_chunk_size: int | None) -> PiTBlock: mlp_ratio=2.0, attn_hidden_size=8, attn_num_heads=2, - activation_chunk_size=activation_chunk_size, ).eval() @pytest.mark.parametrize("use_autocast", [False, True]) def test_pit_block_chunked_forward_matches_unchunked_and_bounds_adaln_batch(use_autocast: bool) -> None: torch.manual_seed(0) - unchunked = _build_pit_block(activation_chunk_size=None) + unchunked = _build_pit_block() chunk_size = 3 - chunked = _build_pit_block(activation_chunk_size=chunk_size) + chunked = _build_pit_block() chunked.load_state_dict(unchunked.state_dict()) batch_size = 2 @@ -45,7 +44,14 @@ def record_adaln_batch_size(_module: torch.nn.Module, inputs: tuple[torch.Tensor try: with torch.no_grad(), torch.autocast("cpu", dtype=torch.bfloat16, enabled=use_autocast): expected = unchunked(pixels, condition, image_height, image_width, patch_size) - actual = chunked(pixels, condition, image_height, image_width, patch_size) + actual = chunked( + pixels, + condition, + image_height, + image_width, + patch_size, + activation_chunk_size=chunk_size, + ) finally: handle.remove() @@ -55,12 +61,15 @@ def record_adaln_batch_size(_module: torch.nn.Module, inputs: tuple[torch.Tensor def test_pit_block_rejects_non_positive_activation_chunk_size() -> None: + block = _build_pit_block() + pixels = torch.randn(4, 4, 4) + condition = torch.randn(4, 8) with pytest.raises(ValueError, match="activation_chunk_size must be positive"): - _build_pit_block(activation_chunk_size=0) + block(pixels, condition, image_height=4, image_width=4, patch_size=2, activation_chunk_size=0) def test_pit_block_uses_unchunked_path_when_gradients_are_enabled() -> None: - block = _build_pit_block(activation_chunk_size=1) + block = _build_pit_block() pixels = torch.randn(4, 4, 4, requires_grad=True) condition = torch.randn(4, 8) adaln_batch_sizes: list[int] = [] @@ -70,10 +79,46 @@ def record_adaln_batch_size(_module: torch.nn.Module, inputs: tuple[torch.Tensor handle = block.adaLN_modulation.register_forward_pre_hook(record_adaln_batch_size) try: - output = block(pixels, condition, image_height=4, image_width=4, patch_size=2) + output = block( + pixels, + condition, + image_height=4, + image_width=4, + patch_size=2, + activation_chunk_size=1, + ) output.sum().backward() finally: handle.remove() assert adaln_batch_sizes == [pixels.shape[0]] assert pixels.grad is not None + + +def test_pit_block_activation_chunking_is_not_sticky_between_calls() -> None: + block = _build_pit_block() + pixels = torch.randn(4, 4, 4) + condition = torch.randn(4, 8) + adaln_batch_sizes: list[int] = [] + + def record_adaln_batch_size(_module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]) -> None: + adaln_batch_sizes.append(inputs[0].shape[0]) + + handle = block.adaLN_modulation.register_forward_pre_hook(record_adaln_batch_size) + try: + with torch.no_grad(): + block( + pixels, + condition, + image_height=4, + image_width=4, + patch_size=2, + activation_chunk_size=1, + ) + optimized_call_count = len(adaln_batch_sizes) + block(pixels, condition, image_height=4, image_width=4, patch_size=2) + finally: + handle.remove() + + assert adaln_batch_sizes[:optimized_call_count] == [1] * 8 + assert adaln_batch_sizes[optimized_call_count:] == [pixels.shape[0]] diff --git a/tests/test_config.py b/tests/test_config.py index 85eaaef5328..f441e07072e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -90,6 +90,15 @@ def test_read_config_from_file(tmp_path: Path, patch_rootdir: None): assert config.port == 8080 +def test_pid_optimization_defaults_to_false_and_loads_from_yaml(tmp_path: Path, patch_rootdir: None) -> None: + assert InvokeAIAppConfig().pid_optimization is False + + temp_config_file = tmp_path / "temp_invokeai.yaml" + temp_config_file.write_text('schema_version: "4.0.3"\npid_optimization: true\n') + + assert load_and_migrate_config(temp_config_file).pid_optimization is True + + def test_migrate_v3_config_from_file(tmp_path: Path, patch_rootdir: None): """Test reading configuration from a file.""" temp_config_file = tmp_path / "temp_invokeai.yaml" From da164a866f2f045ce5f0b1c91baf9941c4cbb96e Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 4 Aug 2026 17:49:45 -0500 Subject: [PATCH 3/7] chore: openapi schema --- invokeai/frontend/web/openapi.json | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 0499e6f426d..8380125c584 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -49027,6 +49027,12 @@ "description": "Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.", "default": false }, + "pid_optimization": { + "type": "boolean", + "title": "Pid Optimization", + "description": "Enable experimental PiD decode memory optimizations at the cost of slower decoding.", + "default": false + }, "attention_type": { "type": "string", "enum": ["auto", "normal", "xformers", "sliced", "torch-sdp"], @@ -49291,7 +49297,7 @@ "additionalProperties": false, "type": "object", "title": "InvokeAIAppConfig", - "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)\n 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.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set." + "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)\n 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.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n pid_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set." }, "InvokeAIAppConfigWithSetFields": { "properties": { From 682b0d973f2f1501a2f7699755e80fe3c5d11a54 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 4 Aug 2026 19:22:08 -0500 Subject: [PATCH 4/7] refactor(pid): rename memory optimization setting --- .../content/docs/configuration/low-vram-mode.mdx | 4 ++-- docs/src/content/docs/features/pid-decode.mdx | 2 +- docs/src/generated/settings.json | 8 ++++---- invokeai/app/invocations/flux2_pid_decode.py | 2 +- invokeai/app/invocations/flux_pid_decode.py | 2 +- invokeai/app/invocations/pid_upscale.py | 2 +- .../app/invocations/qwen_image_pid_decode.py | 2 +- invokeai/app/invocations/sd3_pid_decode.py | 2 +- invokeai/app/invocations/sdxl_pid_decode.py | 2 +- invokeai/app/invocations/z_image_pid_decode.py | 2 +- invokeai/app/services/config/config_default.py | 4 ++-- invokeai/backend/pid/decode.py | 16 ++++++++-------- invokeai/frontend/web/openapi.json | 6 +++--- invokeai/frontend/web/src/services/api/schema.ts | 6 +++--- tests/backend/pid/test_pid_decode.py | 14 +++++++------- tests/test_config.py | 8 ++++---- 16 files changed, 41 insertions(+), 41 deletions(-) diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index 9511d5e9f87..705c07fc11f 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/low-vram-mode.mdx @@ -39,7 +39,7 @@ Low-VRAM mode and related workload-specific optimizations include: - 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_optimization`) +- 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. @@ -154,7 +154,7 @@ PiD decodes directly into high-resolution pixels, so its activation and sampler To reduce PiD's peak VRAM use, enable its experimental memory optimization in `invokeai.yaml` and restart InvokeAI: ```yaml -pid_optimization: true +pid_memory_optimization: true ``` This setting uses lower-precision sampler intermediates and processes parts of the PiD pixel pathway in chunks. It applies to every supported PiD decoder and is disabled by default because chunking makes decoding slower. See [PiD Super-Resolution Decode](/features/pid-decode/) for supported models and usage. diff --git a/docs/src/content/docs/features/pid-decode.mdx b/docs/src/content/docs/features/pid-decode.mdx index 52aed18675e..c39a3608d4d 100644 --- a/docs/src/content/docs/features/pid-decode.mdx +++ b/docs/src/content/docs/features/pid-decode.mdx @@ -72,7 +72,7 @@ PiD is available in both the **Generate** tab (text-to-image) and on the **Canva On GPUs with limited VRAM, enable the experimental PiD memory optimizations in `invokeai.yaml`, then restart InvokeAI: ```yaml -pid_optimization: true +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, at the cost of slower decoding. It is disabled by default. diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index fe9cd8a7040..88da0dfbb02 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -562,12 +562,12 @@ "validation": {} }, { - "category": "GENERATION", + "category": "OTHER", "default": false, - "description": "Enable experimental PiD decode memory optimizations at the cost of slower decoding.", - "env_var": "INVOKEAI_PID_OPTIMIZATION", + "description": "Enable experimental Wan memory optimizations at the cost of slower generation.", + "env_var": "INVOKEAI_WAN_MEMORY_OPTIMIZATION", "literal_values": [], - "name": "pid_optimization", + "name": "wan_memory_optimization", "required": false, "type": "", "validation": {} diff --git a/invokeai/app/invocations/flux2_pid_decode.py b/invokeai/app/invocations/flux2_pid_decode.py index 4569f75a823..68b52281810 100644 --- a/invokeai/app/invocations/flux2_pid_decode.py +++ b/invokeai/app/invocations/flux2_pid_decode.py @@ -229,7 +229,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_optimization=context.config.get().pid_optimization, + pid_memory_optimization=context.config.get().pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/flux_pid_decode.py b/invokeai/app/invocations/flux_pid_decode.py index ab0efb61c46..37461ec8141 100644 --- a/invokeai/app/invocations/flux_pid_decode.py +++ b/invokeai/app/invocations/flux_pid_decode.py @@ -152,7 +152,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_optimization=context.config.get().pid_optimization, + pid_memory_optimization=context.config.get().pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/pid_upscale.py b/invokeai/app/invocations/pid_upscale.py index dba4d5ac59a..00a94670204 100644 --- a/invokeai/app/invocations/pid_upscale.py +++ b/invokeai/app/invocations/pid_upscale.py @@ -192,7 +192,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_optimization=context.config.get().pid_optimization, + pid_memory_optimization=context.config.get().pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/qwen_image_pid_decode.py b/invokeai/app/invocations/qwen_image_pid_decode.py index 81f09c30df7..0bea9ac1b24 100644 --- a/invokeai/app/invocations/qwen_image_pid_decode.py +++ b/invokeai/app/invocations/qwen_image_pid_decode.py @@ -218,7 +218,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_optimization=context.config.get().pid_optimization, + pid_memory_optimization=context.config.get().pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/sd3_pid_decode.py b/invokeai/app/invocations/sd3_pid_decode.py index b60e9b92fe5..48926731068 100644 --- a/invokeai/app/invocations/sd3_pid_decode.py +++ b/invokeai/app/invocations/sd3_pid_decode.py @@ -145,7 +145,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_optimization=context.config.get().pid_optimization, + pid_memory_optimization=context.config.get().pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/sdxl_pid_decode.py b/invokeai/app/invocations/sdxl_pid_decode.py index 426678afd4e..b9dfbd7400b 100644 --- a/invokeai/app/invocations/sdxl_pid_decode.py +++ b/invokeai/app/invocations/sdxl_pid_decode.py @@ -191,7 +191,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_optimization=context.config.get().pid_optimization, + pid_memory_optimization=context.config.get().pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/z_image_pid_decode.py b/invokeai/app/invocations/z_image_pid_decode.py index 06e29ada463..0f18a5da341 100644 --- a/invokeai/app/invocations/z_image_pid_decode.py +++ b/invokeai/app/invocations/z_image_pid_decode.py @@ -201,7 +201,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_optimization=context.config.get().pid_optimization, + pid_memory_optimization=context.config.get().pid_memory_optimization, ), ) context.logger.info( diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 46d7bfb3044..72ea3f28f8c 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -113,7 +113,7 @@ class InvokeAIAppConfig(BaseSettings): device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
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.
Valid values: `auto`, `float16`, `bfloat16`, `float32` sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements. - pid_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding. + pid_memory_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding. attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` attention_slice_size: Slice size, valid when attention_type=="sliced".
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). @@ -221,7 +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_optimization: bool = Field(default=False, description="Enable experimental PiD decode memory optimizations at the cost of slower decoding.") + pid_memory_optimization: bool = Field(default=False, description="Enable experimental PiD decode memory optimizations at the cost of slower decoding.") 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).") diff --git a/invokeai/backend/pid/decode.py b/invokeai/backend/pid/decode.py index 4c854e3dcf2..2d675b92b27 100644 --- a/invokeai/backend/pid/decode.py +++ b/invokeai/backend/pid/decode.py @@ -227,10 +227,10 @@ def _get_t_list(device: torch.device, *, num_steps: Optional[int] = None) -> Ten return t -def _velocity_to_x0(x_t: Tensor, net_output: Tensor, t: Tensor, *, pid_optimization: bool = False) -> Tensor: +def _velocity_to_x0(x_t: Tensor, net_output: Tensor, t: Tensor, *, pid_memory_optimization: bool = False) -> Tensor: """Convert the network's velocity prediction back to x0 at time *t*.""" s = [x_t.shape[0]] + [1] * (x_t.ndim - 1) - if pid_optimization: + if pid_memory_optimization: t_shaped = t.float().view(*s) return torch.addcmul(x_t.float(), net_output.float(), t_shaped, value=-1).to(x_t.dtype) t_shaped = t.double().view(*s) @@ -250,7 +250,7 @@ def _student_sample_loop( sample_type: str = "sde", autocast_dtype: Optional[torch.dtype] = None, generator: Optional[torch.Generator] = None, - pid_optimization: bool = False, + pid_memory_optimization: bool = False, ) -> Tensor: """Few-step distilled sampler. @@ -285,10 +285,10 @@ def _student_sample_loop( lq_video_or_image=None, lq_latent=lq_latent, degrade_sigma=degrade_sigma, - activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE if pid_optimization else None, + activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE if pid_memory_optimization else None, ) if t_next.item() > 0: - x0_pred = _velocity_to_x0(x, v_pred, t_cur_batch, pid_optimization=pid_optimization) + x0_pred = _velocity_to_x0(x, v_pred, t_cur_batch, pid_memory_optimization=pid_memory_optimization) eps_infer = torch.randn( x0_pred.shape, device=x0_pred.device, @@ -304,7 +304,7 @@ def _student_sample_loop( else: x = (1.0 - t_next_b) * x0_pred + t_next_b * eps_infer else: - x = _velocity_to_x0(x, v_pred, t_cur_batch, pid_optimization=pid_optimization) + x = _velocity_to_x0(x, v_pred, t_cur_batch, pid_memory_optimization=pid_memory_optimization) return x @@ -329,7 +329,7 @@ class PiDDecodeConfig: # from_clean upscale path passes the LDM scheduler's per-step sigma here. degrade_sigma: float | list[float] | Tensor = 0.0 seed: int = 0 - pid_optimization: bool = False + pid_memory_optimization: bool = False student_t_list: list[float] = field(default_factory=lambda: list(_STUDENT_T_LIST)) @@ -435,7 +435,7 @@ def decode( sample_type=cfg.sample_type, autocast_dtype=autocast_dtype, generator=gen, - pid_optimization=cfg.pid_optimization, + pid_memory_optimization=cfg.pid_memory_optimization, ) return x0.clamp(-1, 1) diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 8380125c584..0f202f54cbf 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -49027,9 +49027,9 @@ "description": "Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.", "default": false }, - "pid_optimization": { + "pid_memory_optimization": { "type": "boolean", - "title": "Pid Optimization", + "title": "Pid Memory Optimization", "description": "Enable experimental PiD decode memory optimizations at the cost of slower decoding.", "default": false }, @@ -49297,7 +49297,7 @@ "additionalProperties": false, "type": "object", "title": "InvokeAIAppConfig", - "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)\n 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.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n pid_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set." + "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)\n 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.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set." }, "InvokeAIAppConfigWithSetFields": { "properties": { diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 9cd8d2b3924..9078ed1f536 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -18918,7 +18918,7 @@ export type components = { * device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
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.
Valid values: `auto`, `float16`, `bfloat16`, `float32` * sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements. - * pid_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding. + * pid_memory_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding. * attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` * attention_slice_size: Slice size, valid when attention_type=="sliced".
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). @@ -19261,11 +19261,11 @@ export type components = { */ sequential_guidance?: boolean; /** - * Pid Optimization + * Pid Memory Optimization * @description Enable experimental PiD decode memory optimizations at the cost of slower decoding. * @default false */ - pid_optimization?: boolean; + pid_memory_optimization?: boolean; /** * Attention Type * @description Attention type. diff --git a/tests/backend/pid/test_pid_decode.py b/tests/backend/pid/test_pid_decode.py index 3c51bf4b3eb..379ca5e0229 100644 --- a/tests/backend/pid/test_pid_decode.py +++ b/tests/backend/pid/test_pid_decode.py @@ -62,7 +62,7 @@ def test_velocity_to_x0_uses_float32_math_and_preserves_input_dtype( # A float64 intermediate doubles memory for each full-resolution sampler tensor. PiD already # predicts under bf16 autocast, so perform this update in float32 without calling Tensor.double(). with patch.object(torch.Tensor, "double", side_effect=AssertionError("unexpected float64 conversion")): - actual = _velocity_to_x0(x_t, net_output, timestep, pid_optimization=True) + actual = _velocity_to_x0(x_t, net_output, timestep, pid_memory_optimization=True) assert actual.dtype == x_dtype torch.testing.assert_close(actual, expected, rtol=0, atol=0) @@ -75,14 +75,14 @@ def test_velocity_to_x0_uses_original_float64_math_when_optimization_is_disabled expected = (x_t.double() - timestep.double().view(1, 1, 1, 1) * net_output.double()).to(x_t.dtype) with patch("torch.addcmul", side_effect=AssertionError("unexpected optimized path")): - actual = _velocity_to_x0(x_t, net_output, timestep, pid_optimization=False) + actual = _velocity_to_x0(x_t, net_output, timestep, pid_memory_optimization=False) torch.testing.assert_close(actual, expected, rtol=0, atol=0) -@pytest.mark.parametrize(("pid_optimization", "expected_chunk_size"), [(False, None), (True, 1024)]) +@pytest.mark.parametrize(("pid_memory_optimization", "expected_chunk_size"), [(False, None), (True, 1024)]) def test_student_sample_loop_passes_per_call_activation_chunk_size( - pid_optimization: bool, expected_chunk_size: int | None + pid_memory_optimization: bool, expected_chunk_size: int | None ) -> None: activation_chunk_sizes: list[int | None] = [] @@ -98,14 +98,14 @@ def net(x: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor: caption_mask=None, lq_latent=None, degrade_sigma=torch.zeros(1), - pid_optimization=pid_optimization, + pid_memory_optimization=pid_memory_optimization, ) assert activation_chunk_sizes == [expected_chunk_size] -def test_pid_optimization_defaults_to_disabled() -> None: - assert PiDDecodeConfig().pid_optimization is False +def test_pid_memory_optimization_defaults_to_disabled() -> None: + assert PiDDecodeConfig().pid_memory_optimization is False def test_matching_decoder_base_is_accepted() -> None: diff --git a/tests/test_config.py b/tests/test_config.py index f441e07072e..78d5dbe3466 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -90,13 +90,13 @@ def test_read_config_from_file(tmp_path: Path, patch_rootdir: None): assert config.port == 8080 -def test_pid_optimization_defaults_to_false_and_loads_from_yaml(tmp_path: Path, patch_rootdir: None) -> None: - assert InvokeAIAppConfig().pid_optimization is False +def test_pid_memory_optimization_defaults_to_false_and_loads_from_yaml(tmp_path: Path, patch_rootdir: None) -> None: + assert InvokeAIAppConfig().pid_memory_optimization is False temp_config_file = tmp_path / "temp_invokeai.yaml" - temp_config_file.write_text('schema_version: "4.0.3"\npid_optimization: true\n') + temp_config_file.write_text('schema_version: "4.0.3"\npid_memory_optimization: true\n') - assert load_and_migrate_config(temp_config_file).pid_optimization is True + assert load_and_migrate_config(temp_config_file).pid_memory_optimization is True def test_migrate_v3_config_from_file(tmp_path: Path, patch_rootdir: None): From bd7e7f072eb485121d0712f8e9bd3429bc5c8561 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 4 Aug 2026 21:57:14 -0500 Subject: [PATCH 5/7] Fix generated JSON settings file --- docs/src/generated/settings.json | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index 88da0dfbb02..5fbb3a61269 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -562,12 +562,12 @@ "validation": {} }, { - "category": "OTHER", + "category": "GENERATION", "default": false, - "description": "Enable experimental Wan memory optimizations at the cost of slower generation.", - "env_var": "INVOKEAI_WAN_MEMORY_OPTIMIZATION", + "description": "Enable experimental PiD decode memory optimizations at the cost of slower decoding.", + "env_var": "INVOKEAI_PID_MEMORY_OPTIMIZATION", "literal_values": [], - "name": "wan_memory_optimization", + "name": "pid_memory_optimization", "required": false, "type": "", "validation": {} From cfe6f5bd2d3170192fb18830ff08ec177211fa72 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 6 Aug 2026 04:51:24 +0200 Subject: [PATCH 6/7] fix(pid): make the memory optimization actually free VRAM, and say what it costs Addresses every point from the review of #9460. The setting freed activation memory that the cache then withheld anyway. `estimate_pid_decode_working_memory` was not flag-aware, so a decode that peaks at 1.5 GiB reserved the unoptimized 3.9 GiB; the cache takes max(working_mem_bytes, device_working_mem_gb) and subtracts that from the weight budget, so the saving never became weight residency - it only avoided a hard OOM, on precisely the low-VRAM machines this feature exists for. The estimate now takes the flag, and each node reads the setting once and feeds both the estimate and the decode from it, so the two cannot drift apart. Calibrated against measured peaks (RTX 4090, fp32 PidNet, bf16 autocast, 4 steps, B=1): 1024px 509 MiB 1536px 934 MiB 2048px 1533 MiB which is 85.3 * U + 167 MiB, not a pure multiple of the output area - chunking bounds the per-block activations to a fixed working set, so a single scaling constant would under-reserve at small sizes or over-reserve at large ones. The constants carry ~15% headroom. Below the chunk size the pixel blocks run unchunked, so the estimate clamps to the unoptimized one rather than charging for a working set that is never allocated. The documented cost was the wrong cost. "at the cost of slower decoding" is not observable - 2.78 s either way at 2048px, median of 3 with warmup - while the cost that is real went unmentioned: the option changes the decoded image. Both parts are non-bit-exact and the few-step sampler amplifies them (~43 dB PSNR end-to-end, visually indistinguishable, not reproducible against an unoptimized decode). The setting description and both docs pages now state that, with the measured VRAM numbers. The equivalence test could not fail. It asserted `assert_close` at pixel_hidden_size=4 / BL=8 / chunk=3 on the CPU, while the shipped path is CUDA under bf16 autocast with BL in the thousands. Measured at production dimensions: CPU fp32 is bit-identical (including B=2, where chunks straddle images), CUDA fp32 differs by 9.5e-07, CUDA bf16 by 1.57e-02 - systematic, not noise, both paths being internally deterministic. So chunking is exact as *mathematics*; the divergence is cuBLAS picking different kernels for 1024-row slices. `test_pid_chunked_equivalence.py` now pins both halves: exactness on the CPU, and an absolute tolerance contract on CUDA with 3x headroom (relative tolerances are meaningless here - activations cross zero, so max|rel| reaches 1e3 on elements whose absolute error is one bf16 ULP). Two review points did not survive measurement, and are documented rather than "fixed": - The fp32 `_velocity_to_x0` branch is kept. Fusing the multiply-subtract in fp64 is bit-identical to the default expression but frees nothing (288 MiB either way at 2048px), so the 192 MiB is bought entirely with precision - max|diff| 4.8e-07 per call, 8.6% of the flag's total saving. That is a fair trade for someone who opted into trading quality for VRAM, but it is now written down, in the function, in the setting description and in the docs. - The doubled `adaLN_modulation` per chunk stays. Global attention sits between the two halves, so reusing the slices means holding them for every chunk - the full-resolution tensor the path exists to avoid (536 MiB in bf16 at 2048px). Projecting only the needed slices is no better: they are interleaved per pixel position, so selecting them gathers rows of a 1536x24576 weight, ~50 MiB per call. The recompute is ~9.9 TFLOP per 2048px decode, about 4% wall clock - and the honest source of the "slower decoding" the setting advertises. Observability: a decode with the flag on now logs the resolution, the patch-token count and whether chunking actually engaged. The setting is server-wide and never enters image metadata, so this is the only record that a given decode ran optimized, and the only feedback that a yaml-only, restart-required knob took effect at all. Tests: AST sweeps assert that every module building a `PiDDecodeConfig` forwards the flag, estimates working memory for the same mode, and reads the setting exactly once - so an eighth PiD node is covered the day it lands. Plus estimate behaviour (shrinks when enabled, keeps the fixed term, never exceeds the unoptimized estimate, still returns 0 for unsupported backbones), batch > 1 across chunk boundaries, and the two paths the flag must not reach: the discriminator feature extraction returns before the pixel blocks, and context parallelism is unreachable in this codebase (its only caller lives in a vendored class InvokeAI never instantiates) - both pinned so a refactor cannot quietly put them under the flag. All new tests mutation-verified: breaking the chunked assembly, making the estimate ignore the flag, dropping the fixed term, and omitting the flag from a node's decode or estimate each fail at least one test. tests/app + tests/backend/pid: 2202 passed. The 9 failures are the pre-existing network-dependent ones in test_model_install / test_load_api / test_download_queue. --- .../docs/configuration/low-vram-mode.mdx | 6 +- docs/src/content/docs/features/pid-decode.mdx | 4 +- docs/src/generated/settings.json | 2 +- invokeai/app/invocations/flux2_pid_decode.py | 9 +- invokeai/app/invocations/flux_pid_decode.py | 9 +- invokeai/app/invocations/pid_upscale.py | 9 +- .../app/invocations/qwen_image_pid_decode.py | 9 +- invokeai/app/invocations/sd3_pid_decode.py | 9 +- invokeai/app/invocations/sdxl_pid_decode.py | 9 +- .../app/invocations/z_image_pid_decode.py | 9 +- .../app/services/config/config_default.py | 4 +- .../pid/_src/networks/pixeldit_official.py | 13 ++ invokeai/backend/pid/decode.py | 71 +++++++- invokeai/frontend/web/openapi.json | 4 +- .../frontend/web/src/services/api/schema.ts | 4 +- .../test_pid_memory_optimization_wiring.py | 82 +++++++++ .../pid/test_pid_chunked_equivalence.py | 170 ++++++++++++++++++ tests/backend/pid/test_pid_decode.py | 55 ++++++ 18 files changed, 452 insertions(+), 26 deletions(-) create mode 100644 tests/app/invocations/test_pid_memory_optimization_wiring.py create mode 100644 tests/backend/pid/test_pid_chunked_equivalence.py diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index 705c07fc11f..5e15e78cb9a 100644 --- a/docs/src/content/docs/configuration/low-vram-mode.mdx +++ b/docs/src/content/docs/configuration/low-vram-mode.mdx @@ -157,7 +157,11 @@ To reduce PiD's peak VRAM use, enable its experimental memory optimization in `i pid_memory_optimization: true ``` -This setting uses lower-precision sampler intermediates and processes parts of the PiD pixel pathway in chunks. It applies to every supported PiD decoder and is disabled by default because chunking makes decoding slower. See [PiD Super-Resolution Decode](/features/pid-decode/) for supported models and usage. +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) diff --git a/docs/src/content/docs/features/pid-decode.mdx b/docs/src/content/docs/features/pid-decode.mdx index c39a3608d4d..7f6745d6d03 100644 --- a/docs/src/content/docs/features/pid-decode.mdx +++ b/docs/src/content/docs/features/pid-decode.mdx @@ -75,7 +75,9 @@ On GPUs with limited VRAM, enable the experimental PiD memory optimizations in ` 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, at the cost of slower decoding. It is disabled by default. +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. diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index 5fbb3a61269..76b464eef34 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -564,7 +564,7 @@ { "category": "GENERATION", "default": false, - "description": "Enable experimental PiD decode memory optimizations at the cost of slower decoding.", + "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", diff --git a/invokeai/app/invocations/flux2_pid_decode.py b/invokeai/app/invocations/flux2_pid_decode.py index 68b52281810..907aabde180 100644 --- a/invokeai/app/invocations/flux2_pid_decode.py +++ b/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__}.") @@ -229,7 +234,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_memory_optimization=context.config.get().pid_memory_optimization, + pid_memory_optimization=pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/flux_pid_decode.py b/invokeai/app/invocations/flux_pid_decode.py index 37461ec8141..d5cf45a08c9 100644 --- a/invokeai/app/invocations/flux_pid_decode.py +++ b/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__}.") @@ -152,7 +157,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_memory_optimization=context.config.get().pid_memory_optimization, + pid_memory_optimization=pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/pid_upscale.py b/invokeai/app/invocations/pid_upscale.py index 00a94670204..348158008e3 100644 --- a/invokeai/app/invocations/pid_upscale.py +++ b/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__}.") @@ -192,7 +197,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_memory_optimization=context.config.get().pid_memory_optimization, + pid_memory_optimization=pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/qwen_image_pid_decode.py b/invokeai/app/invocations/qwen_image_pid_decode.py index 0bea9ac1b24..58d32a84e40 100644 --- a/invokeai/app/invocations/qwen_image_pid_decode.py +++ b/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__}.") @@ -218,7 +223,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_memory_optimization=context.config.get().pid_memory_optimization, + pid_memory_optimization=pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/sd3_pid_decode.py b/invokeai/app/invocations/sd3_pid_decode.py index 48926731068..237e4d0aadc 100644 --- a/invokeai/app/invocations/sd3_pid_decode.py +++ b/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__}.") @@ -145,7 +150,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_memory_optimization=context.config.get().pid_memory_optimization, + pid_memory_optimization=pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/sdxl_pid_decode.py b/invokeai/app/invocations/sdxl_pid_decode.py index b9dfbd7400b..bcaabe9925f 100644 --- a/invokeai/app/invocations/sdxl_pid_decode.py +++ b/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__}.") @@ -191,7 +196,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_memory_optimization=context.config.get().pid_memory_optimization, + pid_memory_optimization=pid_memory_optimization, ), ) diff --git a/invokeai/app/invocations/z_image_pid_decode.py b/invokeai/app/invocations/z_image_pid_decode.py index 0f18a5da341..bdd957cc5dc 100644 --- a/invokeai/app/invocations/z_image_pid_decode.py +++ b/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__}.") @@ -201,7 +206,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: config=PiDDecodeConfig( num_inference_steps=self.num_inference_steps, seed=self.seed, - pid_memory_optimization=context.config.get().pid_memory_optimization, + pid_memory_optimization=pid_memory_optimization, ), ) context.logger.info( diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 72ea3f28f8c..1dc8b560326 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -113,7 +113,7 @@ class InvokeAIAppConfig(BaseSettings): device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
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.
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 at the cost of slower decoding. + 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.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` attention_slice_size: Slice size, valid when attention_type=="sliced".
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). @@ -221,7 +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 at the cost of slower decoding.") + 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).") diff --git a/invokeai/backend/pid/_src/networks/pixeldit_official.py b/invokeai/backend/pid/_src/networks/pixeldit_official.py index 4318d836bca..f71b15176f1 100644 --- a/invokeai/backend/pid/_src/networks/pixeldit_official.py +++ b/invokeai/backend/pid/_src/networks/pixeldit_official.py @@ -539,6 +539,19 @@ def _forward_unchunked( return x def _compress_activation_chunk(self, x: torch.Tensor, s_cond: torch.Tensor) -> torch.Tensor: + """First half of a chunk: adaLN + norm1, compressed to one attention token per patch. + + This recomputes ``adaLN_modulation`` that ``_finish_activation_chunk`` computes again, and that + is deliberate. Global attention sits between the two halves, so the four slices the second half + needs would have to be held for *every* chunk across the attention call - i.e. the full-resolution + ``[BL, P2, 4*pixel_dim]`` tensor this whole path exists to avoid (536 MiB in bf16 at 2048px). + Recomputing trades roughly 9.9 TFLOP per 2048px decode for that, measured at ~4% wall clock, + which is the honest source of the "slower decoding" the setting advertises. + + Computing only the needed slices instead is not the way out either: the six slices are interleaved + per pixel position in the projection's output, so selecting two of them means gathering rows of a + 1536x24576 weight - a ~50 MiB copy per call, spending memory to save the compute. + """ patch_batch, pixels_per_patch, _ = x.shape cond_params = self.adaLN_modulation(s_cond).view(patch_batch, pixels_per_patch, 6 * self.pixel_dim) shift_msa, scale_msa = torch.chunk(cond_params, 6, dim=-1)[:2] diff --git a/invokeai/backend/pid/decode.py b/invokeai/backend/pid/decode.py index 2d675b92b27..9ecfec3ad32 100644 --- a/invokeai/backend/pid/decode.py +++ b/invokeai/backend/pid/decode.py @@ -27,6 +27,7 @@ from invokeai.backend.model_manager.taxonomy import BaseModelType from invokeai.backend.pid._src.networks.pid_net import PidNet +from invokeai.backend.util.logging import InvokeAILogger _PID_ACTIVATION_CHUNK_SIZE = 1024 @@ -133,13 +134,33 @@ # ~4GB at a 2048px output is a small headroom above the 3GB default. Experimentally-tunable; calibrate to peak. _PID_DECODE_WORKING_MEMORY_SCALING_CONSTANT = 250 +# The same estimate for `pid_memory_optimization=True`. Chunking bounds the per-block activations to a fixed +# working set, so the peak stops being a pure multiple of the output size: it is a smaller per-pixel term plus a +# constant for the chunk working set (constant because `_PID_ACTIVATION_CHUNK_SIZE` is fixed). +# +# Measured peaks on an RTX 4090 (fp32 PidNet, bf16 autocast, 4 steps, B=1), against U = out_h * out_w * 4 bytes: +# 1024px 509 MiB (127.2 * U) 1536px 934 MiB (103.8 * U) 2048px 1533 MiB (95.8 * U) +# Least-squares fit: 85.3 * U + 167 MiB. The constants below carry ~15% headroom over that fit. +# +# Keeping the unoptimized constant here would be the bug the flag is supposed to avoid: the cache takes +# max(this_estimate, device_working_mem_gb) and subtracts it from the weight budget, so reserving 4GB for a decode +# that peaks at 1.5GB withholds VRAM that PidNet could have stayed resident in - exactly the partial-load-to-CPU +# outcome the comment above warns about, on the low-VRAM systems this feature exists for. +_PID_DECODE_CHUNKED_SCALING_CONSTANT = 95 +_PID_DECODE_CHUNKED_FIXED_BYTES = 224 * 2**20 -def estimate_pid_decode_working_memory(latent: Tensor, backbone: BaseModelType) -> int: + +def estimate_pid_decode_working_memory( + latent: Tensor, backbone: BaseModelType, pid_memory_optimization: bool = False +) -> int: """Estimate the working (activation) memory in bytes for a PiD decode of *latent*. The decoded image is ``latent_spatial * sr_scale * latent_spatial_down_factor`` pixels per side. PidNet runs in float32 (see ``model_loaders/pid_decoder.py``), so the element size is 4 bytes. Returns 0 for unsupported backbones so callers fall back to the cache's default working-memory reservation. + + ``pid_memory_optimization`` must mirror the flag passed to :class:`PiDDecodeConfig` for the same decode - + otherwise the cache reserves headroom for a peak that will not happen. """ per_backbone = _PER_BACKBONE.get(backbone) if per_backbone is None: @@ -148,7 +169,16 @@ def estimate_pid_decode_working_memory(latent: Tensor, backbone: BaseModelType) out_h = int(latent.shape[-2]) * total_up out_w = int(latent.shape[-1]) * total_up element_size = 4 # PidNet runs in float32 (see model_loaders/pid_decoder.py) - return int(out_h * out_w * element_size * _PID_DECODE_WORKING_MEMORY_SCALING_CONSTANT) + output_bytes = out_h * out_w * element_size + unoptimized = int(output_bytes * _PID_DECODE_WORKING_MEMORY_SCALING_CONSTANT) + if not pid_memory_optimization: + return unoptimized + chunked = int(output_bytes * _PID_DECODE_CHUNKED_SCALING_CONSTANT + _PID_DECODE_CHUNKED_FIXED_BYTES) + # The fixed term is the chunk working set, so it only exists once chunking actually engages - + # below `_PID_ACTIVATION_CHUNK_SIZE` patch tokens the pixel blocks run unchunked and the peak is + # the unoptimized one. Without this clamp a small output would be *charged* for a working set it + # never allocates, and the optimized estimate could exceed the unoptimized one. + return min(chunked, unoptimized) def build_pid_net(backbone: BaseModelType) -> PidNet: @@ -228,7 +258,27 @@ def _get_t_list(device: torch.device, *, num_steps: Optional[int] = None) -> Ten def _velocity_to_x0(x_t: Tensor, net_output: Tensor, t: Tensor, *, pid_memory_optimization: bool = False) -> Tensor: - """Convert the network's velocity prediction back to x0 at time *t*.""" + """Convert the network's velocity prediction back to x0 at time *t*. + + The optimized branch is a genuine precision reduction, not just a cheaper spelling, so it is worth + being explicit about what it buys. Measured on an RTX 4090 (B=1, 3xHxW, ``x_t`` fp32 / ``net_output`` + bf16), transient peak for this call alone: + + ====== ========== ============== ============== + size fp64 (dflt) fused fp64 fused fp32 + ====== ========== ============== ============== + 1024px 72 MiB 72 MiB 24 MiB + 2048px 288 MiB 288 MiB 96 MiB + ====== ========== ============== ============== + + Fusing the multiply-subtract in fp64 is bit-identical to the default expression but frees nothing, + so the 192 MiB at 2048px is bought entirely with precision: ``max|diff| = 4.8e-07`` per call against + the fp64 result. That is ~8.6% of the 2.2 GiB the flag saves overall, and the 4-step SDE sampler + amplifies the per-call error into a visible-in-numbers-only image delta (see the tolerance contract + in ``tests/backend/pid/test_pid_chunked_equivalence.py``). It stays under the same flag because a + user who opted into "trade quality for VRAM" wants both parts; it is documented here, in the setting + description and in the docs so nobody has to rediscover that this option is not output-preserving. + """ s = [x_t.shape[0]] + [1] * (x_t.ndim - 1) if pid_memory_optimization: t_shaped = t.float().view(*s) @@ -423,6 +473,21 @@ def decode( t_list = _get_t_list(device, num_steps=cfg.num_inference_steps) + if cfg.pid_memory_optimization: + # The setting is server-level and never reaches image metadata, so this log line is the + # only record that a decode ran optimized - and the only feedback the user gets that a + # yaml-only, restart-required knob took effect. It also reports whether chunking really + # engaged: below the chunk size the pixel blocks run unchunked and only the sampler-math + # change applies. + patch_tokens = batch_size * (img_h // self.net.patch_size) * (img_w // self.net.patch_size) + engaged = patch_tokens > _PID_ACTIVATION_CHUNK_SIZE + InvokeAILogger.get_logger(__name__).info( + f"PiD memory optimization enabled for a {img_w}x{img_h} decode: " + f"{patch_tokens} patch tokens vs chunk size {_PID_ACTIVATION_CHUNK_SIZE} " + f"({'chunked' if engaged else 'below the chunk size, activations run unchunked'}), " + "float32 sampler intermediates. Output differs slightly from an unoptimized decode." + ) + self.net.eval() x0 = _student_sample_loop( self.net, diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 0f202f54cbf..6ae841bac9c 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -49030,7 +49030,7 @@ "pid_memory_optimization": { "type": "boolean", "title": "Pid Memory Optimization", - "description": "Enable experimental PiD decode memory optimizations at the cost of slower decoding.", + "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.", "default": false }, "attention_type": { @@ -49297,7 +49297,7 @@ "additionalProperties": false, "type": "object", "title": "InvokeAIAppConfig", - "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)\n 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.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations at the cost of slower decoding.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set." + "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)\n 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.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n 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.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set." }, "InvokeAIAppConfigWithSetFields": { "properties": { diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 9078ed1f536..46badaf2f6e 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -18918,7 +18918,7 @@ export type components = { * device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
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.
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 at the cost of slower decoding. + * 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.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp` * attention_slice_size: Slice size, valid when attention_type=="sliced".
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). @@ -19262,7 +19262,7 @@ export type components = { sequential_guidance?: boolean; /** * Pid Memory Optimization - * @description Enable experimental PiD decode memory optimizations at the cost of slower decoding. + * @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. * @default false */ pid_memory_optimization?: boolean; diff --git a/tests/app/invocations/test_pid_memory_optimization_wiring.py b/tests/app/invocations/test_pid_memory_optimization_wiring.py new file mode 100644 index 00000000000..8084ee5210b --- /dev/null +++ b/tests/app/invocations/test_pid_memory_optimization_wiring.py @@ -0,0 +1,82 @@ +"""Every PiD node must forward `pid_memory_optimization` — to the decode *and* to the memory estimate. + +Two things go wrong silently here, so both are pinned structurally rather than per-node: + +1. A new PiD node that forgets `pid_memory_optimization=` on its `PiDDecodeConfig` decodes unoptimized + while the user believes the setting applies. There is no error and no log; the only symptom is VRAM + use nobody is measuring. +2. Passing the flag to the decode but not to `estimate_pid_decode_working_memory` is worse than not + wiring it at all. The cache takes `max(working_mem_bytes, device_working_mem_gb)` and subtracts it + from the weight budget, so an estimate calibrated for the unoptimized peak withholds VRAM that the + optimization just freed — PidNet then partial-loads to CPU on exactly the low-VRAM machines the + feature exists for. + +An AST sweep is used instead of invoking each node because it needs no model fixtures and, unlike a +per-node test, it automatically covers the eighth PiD node the day someone adds it. +""" + +import ast +from pathlib import Path + +import pytest + +_INVOCATIONS_DIR = Path(__file__).parents[3] / "invokeai" / "app" / "invocations" +_FLAG = "pid_memory_optimization" +_ESTIMATE = "estimate_pid_decode_working_memory" + + +def _modules_constructing_a_decode_config() -> list[Path]: + modules = sorted(p for p in _INVOCATIONS_DIR.glob("*.py") if "PiDDecodeConfig(" in p.read_text(encoding="utf-8")) + assert modules, "no PiD nodes found - has the invocations directory moved?" + return modules + + +def _calls(tree: ast.AST, func_name: str) -> list[ast.Call]: + return [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and ( + (isinstance(node.func, ast.Name) and node.func.id == func_name) + or (isinstance(node.func, ast.Attribute) and node.func.attr == func_name) + ) + ] + + +@pytest.mark.parametrize("module_path", _modules_constructing_a_decode_config(), ids=lambda p: p.stem) +def test_pid_node_forwards_the_setting_to_the_decode(module_path: Path) -> None: + tree = ast.parse(module_path.read_text(encoding="utf-8")) + configs = _calls(tree, "PiDDecodeConfig") + assert configs, f"{module_path.name}: expected at least one PiDDecodeConfig(...)" + + for call in configs: + keywords = {kw.arg for kw in call.keywords} + assert _FLAG in keywords, ( + f"{module_path.name}:{call.lineno} builds a PiDDecodeConfig without {_FLAG}=; " + "the app config setting would silently not apply to this node." + ) + + +@pytest.mark.parametrize("module_path", _modules_constructing_a_decode_config(), ids=lambda p: p.stem) +def test_pid_node_estimates_working_memory_for_the_same_mode_it_decodes_in(module_path: Path) -> None: + tree = ast.parse(module_path.read_text(encoding="utf-8")) + estimates = _calls(tree, _ESTIMATE) + assert estimates, f"{module_path.name}: expected a call to {_ESTIMATE}(...)" + + for call in estimates: + passes_flag = len(call.args) >= 3 or any(kw.arg == _FLAG for kw in call.keywords) + assert passes_flag, ( + f"{module_path.name}:{call.lineno} estimates working memory without {_FLAG}; " + "the cache would reserve the unoptimized peak and withhold the VRAM the flag frees." + ) + + +@pytest.mark.parametrize("module_path", _modules_constructing_a_decode_config(), ids=lambda p: p.stem) +def test_pid_node_reads_the_setting_exactly_once(module_path: Path) -> None: + """The estimate and the decode must be fed from one read, so they cannot disagree.""" + source = module_path.read_text(encoding="utf-8") + reads = source.count(f"context.config.get().{_FLAG}") + assert reads == 1, ( + f"{module_path.name}: reads the setting {reads} times; a single read keeps the working-memory " + "estimate and the decode in lockstep." + ) diff --git a/tests/backend/pid/test_pid_chunked_equivalence.py b/tests/backend/pid/test_pid_chunked_equivalence.py new file mode 100644 index 00000000000..78c115dfb6b --- /dev/null +++ b/tests/backend/pid/test_pid_chunked_equivalence.py @@ -0,0 +1,170 @@ +"""What `pid_memory_optimization` guarantees about the decoded image, at production dimensions. + +`test_pixeldit_official.py` pins the chunked `PiTBlock` against the unchunked one at toy size on the +CPU. That is a real assertion about the *math* - and it holds exactly - but it cannot see the shape +of the problem the setting actually has, because the shipped path runs on CUDA under bf16 autocast +(`PiDDecoder.decode` sets `autocast_dtype = torch.bfloat16` for every CUDA decode) with `BL` in the +thousands rather than 8. There, splitting a GEMM into 1024-row slices makes cuBLAS pick different +kernels and reduction orders, and the results differ by bf16 ULPs. + +Measured on an RTX 4090 (torch 2.7.1+cu128), production `PiTBlock` dimensions: + + CPU fp32, BL=1024 / 2048 -> max|diff| = 0 (bit-identical) + CUDA fp32, BL=4096 -> max|diff| = 9.5e-07 + CUDA bf16 autocast, BL=4096..32768 -> max|diff| = 1.57e-02, mean|diff| = 2.5e-05 + +Both paths are internally deterministic, so those numbers are systematic, not run noise. The +consequence downstream is an image that differs slightly: ~43 dB PSNR end-to-end, visually +indistinguishable but not reproducible against an unoptimized decode. + +This module therefore states the contract the setting really offers, in two halves: chunking is +*exact* as mathematics, and on the accelerated path it is bounded rather than exact. Relative +tolerances are useless here - activations pass through zero, so `max|rel|` reaches 1e3 on elements +whose absolute error is a bf16 ULP - hence an absolute bound. +""" + +import pytest +import torch + +from invokeai.backend.pid._src.networks.pixeldit_official import PiTBlock +from invokeai.backend.pid.decode import _PID_ACTIVATION_CHUNK_SIZE + +# Production PiTBlock geometry, from `_PID_SR4X_BASE` in invokeai/backend/pid/decode.py. +_PIXEL_HIDDEN_SIZE = 16 +_PATCH_HIDDEN_SIZE = 1536 +_PATCH_SIZE = 16 +_NUM_GROUPS = 24 +_ATTN_HIDDEN_SIZE = 1152 +_ATTN_NUM_GROUPS = 16 + +# The contract. Measured worst case under bf16 autocast is 1.57e-02 and is stable across image sizes +# and batch sizes; this leaves ~3x headroom so kernel-selection differences on other GPUs stay inside +# it, while a genuine breakage of the chunked path (wrong slice, misassembled output) blows past it +# by orders of magnitude. +_BF16_ABSOLUTE_TOLERANCE = 5e-2 +_BF16_MEAN_ABSOLUTE_TOLERANCE = 1e-3 + + +def _build_block(device: str) -> PiTBlock: + torch.manual_seed(0) + block = PiTBlock( + pixel_hidden_size=_PIXEL_HIDDEN_SIZE, + patch_hidden_size=_PATCH_HIDDEN_SIZE, + patch_size=_PATCH_SIZE, + num_heads=_NUM_GROUPS, + mlp_ratio=4.0, + attn_hidden_size=_ATTN_HIDDEN_SIZE, + attn_num_heads=_ATTN_NUM_GROUPS, + rope_mode="ntk_aware", + rope_ref_grid_h=64, + rope_ref_grid_w=64, + ).eval() + return block.to(device) + + +def _inputs(image_px: int, batch_size: int, device: str) -> tuple[torch.Tensor, torch.Tensor]: + patch_grid = image_px // _PATCH_SIZE + patch_batch = batch_size * patch_grid * patch_grid + torch.manual_seed(1) + x = torch.randn(patch_batch, _PATCH_SIZE**2, _PIXEL_HIDDEN_SIZE, device=device) + s_cond = torch.randn(patch_batch, _PATCH_HIDDEN_SIZE, device=device) + return x, s_cond + + +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_chunking_is_exact_mathematics_on_cpu(batch_size: int) -> None: + """Chunking only reorders work, so on a path with no kernel-selection freedom it is bit-exact. + + Batch > 1 is covered here because that is where chunk boundaries stop lining up with image + boundaries: with `BL = B * Hs * Ws`, a chunk can straddle two images, and nothing in + `_forward_chunked` may depend on an image staying within one chunk. + """ + image_px = 512 # BL = 1024 (B=1) / 2048 (B=2), i.e. at and above the chunk size + block = _build_block("cpu") + x, s_cond = _inputs(image_px, batch_size, "cpu") + + with torch.no_grad(): + unchunked = block(x, s_cond, image_px, image_px, _PATCH_SIZE) + chunked = block(x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE) + + assert torch.equal(chunked, unchunked) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the divergence only exists on the accelerated path") +@pytest.mark.parametrize(("image_px", "batch_size"), [(1024, 1), (2048, 1), (2048, 2)]) +def test_chunking_stays_within_the_documented_tolerance_on_cuda_bf16(image_px: int, batch_size: int) -> None: + """The shipped path: CUDA + bf16 autocast, BL well above the chunk size. + + Asserting equality here would be asserting something false. Asserting a bound is the honest + contract, and it is what the docs promise users who enable the setting. + """ + block = _build_block("cuda") + x, s_cond = _inputs(image_px, batch_size, "cuda") + patch_batch = x.shape[0] + assert patch_batch >= 2 * _PID_ACTIVATION_CHUNK_SIZE, "the chunked path must actually be exercised" + + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + unchunked = block(x, s_cond, image_px, image_px, _PATCH_SIZE) + chunked = block(x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE) + chunked_again = block( + x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE + ) + + # Determinism first: without it, a passing tolerance would prove nothing about the chunking. + assert torch.equal(chunked, chunked_again) + + difference = (chunked.float() - unchunked.float()).abs() + assert difference.max().item() <= _BF16_ABSOLUTE_TOLERANCE + assert difference.mean().item() <= _BF16_MEAN_ABSOLUTE_TOLERANCE + # And it is genuinely not exact - if this ever starts failing, the divergence is gone and the + # tolerance contract (plus the "output changes" wording in the docs) should be revisited. + assert difference.max().item() > 0.0 + + +def test_discriminator_feature_extraction_never_reaches_the_chunked_pixel_blocks() -> None: + """`PidNet.forward` returns from the `return_features_early` branch before the pixel pathway. + + The chunk size is only handed to `self.pixel_blocks`, so the discriminator/feature path is + structurally unaffected by the setting. Pinning it here keeps a future refactor from moving the + early exit below the pixel loop and quietly putting an untested path under the flag. + """ + import inspect + + from invokeai.backend.pid._src.networks.pid_net import PidNet + + source = inspect.getsource(PidNet.forward) + early_exit = source.index("return self._unpatchify_features(") + pixel_loop = source.index("for blk in self.pixel_blocks:") + chunk_handoff = source.index("activation_chunk_size=activation_chunk_size") + + assert early_exit < pixel_loop < chunk_handoff + + +def test_context_parallelism_is_unreachable_so_chunking_cannot_interact_with_it() -> None: + """Chunking assembles compressed tokens in original order before the *unchanged* attention call, + but under context parallelism attention gathers k/v across ranks — an interaction nothing here + tests, and which cannot be tested single-GPU. + + It is also unreachable: `enable_context_parallel` is only ever called from + `pid/_src/models/pixeldit_model.py`, a vendored upstream class InvokeAI never instantiates (the + loader goes `load_pid_decoder` -> `build_pid_net` -> `PidNet` directly). So `_cp_group` stays + None on every decode this application performs. Pinned here so that wiring CP up later cannot + quietly put chunking into an untested regime. + """ + from pathlib import Path + + import invokeai + + invokeai_root = Path(invokeai.__file__).parent + vendored = invokeai_root / "backend" / "pid" / "_src" + # Scan the Python trees only: `invokeai/frontend/web` carries node_modules, which is huge and + # contains symlinks that break a naive walk. + callers = [ + path + for tree in (invokeai_root / "app", invokeai_root / "backend") + for path in tree.rglob("*.py") + if vendored not in path.parents and "enable_context_parallel(" in path.read_text(encoding="utf-8") + ] + assert callers == [], f"context parallelism is now reachable from {callers}; revisit the chunked attention path" + + assert _build_block("cpu")._cp_group is None diff --git a/tests/backend/pid/test_pid_decode.py b/tests/backend/pid/test_pid_decode.py index 379ca5e0229..97230d7482b 100644 --- a/tests/backend/pid/test_pid_decode.py +++ b/tests/backend/pid/test_pid_decode.py @@ -12,6 +12,7 @@ _student_sample_loop, _velocity_to_x0, assert_pid_decoder_matches_base, + estimate_pid_decode_working_memory, ) _CPU = torch.device("cpu") @@ -108,6 +109,60 @@ def test_pid_memory_optimization_defaults_to_disabled() -> None: assert PiDDecodeConfig().pid_memory_optimization is False +def test_working_memory_estimate_shrinks_when_the_optimization_is_enabled() -> None: + """Wiring the flag into the decode but not into the estimate is worse than not wiring it at all. + + The cache takes `max(working_mem_bytes, device_working_mem_gb)` and subtracts it from the weight + budget, so an estimate calibrated for the unoptimized peak withholds precisely the VRAM the + optimization just freed, and PidNet partial-loads to CPU on the machines this feature targets. + """ + latent = torch.zeros(1, 16, 64, 64) # FLUX: 64 * 4 * 8 = 2048px output + + unoptimized = estimate_pid_decode_working_memory(latent, BaseModelType.Flux) + optimized = estimate_pid_decode_working_memory(latent, BaseModelType.Flux, True) + + assert optimized < unoptimized + # Measured peaks at 2048px on an RTX 4090: 3.68 GiB unoptimized, 1.50 GiB optimized. The estimates + # must sit above their own peak (headroom) and below the other mode's (or the flag buys nothing). + gib = 1024**3 + assert 1.50 * gib < optimized < 2.50 * gib + assert 3.68 * gib < unoptimized < 4.50 * gib + + +def test_working_memory_estimate_keeps_a_fixed_term_for_the_chunk_working_set() -> None: + """The optimized peak is not a pure multiple of the output size. + + Chunking bounds the per-block activations to a fixed working set, so halving the output area does + not halve the peak (measured: 509 MiB at 1024px vs 1533 MiB at 2048px — a factor of 3.0, not 4). + A pure scaling constant would therefore under-reserve at small sizes or over-reserve at large ones. + """ + small = estimate_pid_decode_working_memory(torch.zeros(1, 16, 32, 32), BaseModelType.Flux, True) + large = estimate_pid_decode_working_memory(torch.zeros(1, 16, 64, 64), BaseModelType.Flux, True) + + assert large < 4 * small, "the estimate scales purely with area — the fixed chunk term is missing" + assert large > 2 * small, "the per-pixel term has been lost" + + +def test_working_memory_estimate_never_exceeds_the_unoptimized_one() -> None: + """Below the chunk size the pixel blocks run unchunked, so the fixed chunk-working-set term must + not be charged: a small output would otherwise reserve *more* with the optimization enabled than + without it. + """ + # 8 * 4 * 8 = 256px output -> 256 patch tokens, well under the 1024-token chunk size. + small_latent = torch.zeros(1, 16, 8, 8) + + optimized = estimate_pid_decode_working_memory(small_latent, BaseModelType.Flux, True) + unoptimized = estimate_pid_decode_working_memory(small_latent, BaseModelType.Flux) + + assert optimized <= unoptimized + + +def test_working_memory_estimate_still_returns_zero_for_unsupported_backbones() -> None: + latent = torch.zeros(1, 16, 64, 64) + for optimized in (False, True): + assert estimate_pid_decode_working_memory(latent, BaseModelType.StableDiffusion1, optimized) == 0 + + def test_matching_decoder_base_is_accepted() -> None: assert_pid_decoder_matches_base(BaseModelType.Flux, BaseModelType.Flux, node_title="FLUX PiD Decode") From 8dff4eb0ad6782ba38d821b1e2a8e32e4f2781b3 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 6 Aug 2026 06:15:29 +0200 Subject: [PATCH 7/7] test(pid): fix a chunking test that was both unportable and partly vacuous CI caught this on macos-default py3.11; every other job in the matrix was cancelled by fail-fast. Test-only change, no production code touched. Two separate mistakes, both mine: 1. The CPU comparison asserted `torch.equal`. That held on x86-64 with MKL and failed on macOS/Accelerate. Splitting a GEMM along its row dimension can select a micro-kernel with different K-blocking, so bit-exactness there is a property of the BLAS, not of the chunking. Only reassociation-closeness is portable. 2. Worse, and only found while investigating the first: the `batch_size=1` parametrization never entered the chunked path at all. The dispatch guard is `BL > chunk_size`, and 512px with B=1 puts BL at exactly 1024 - so it compared the unchunked path against itself and passed for the wrong reason. Verified by spying on `_forward_chunked`: zero calls. Both cases now demonstrably chunk - 768px/B=1 (BL 2304, boundaries inside one image) and 512px/B=2 (BL 2048, boundaries straddling images) - and a context manager fails the test if `_forward_chunked` is not entered, so the comparison cannot silently empty out again. Bit-equality is replaced by a signal-relative bound, calibrated rather than guessed. At these dimensions the signal is ~5.7, so one fp32 ULP is ~6.8e-07: correct code, x86-64/MKL max|diff| = 0 attention contribution off by 1e-6 max|diff| = 7.2e-07 (1.3e-07 relative, sub-ULP) attention contribution off by 1e-4 max|diff| = 1.0e-05 (1.9e-06 relative, ~15 ULP) 1e-5 relative is ~84 ULP: above any BLAS reassociation, four orders of magnitude below a structural break. The docstring states what that gives up - a uniform scaling error below ~2e-06 relative is indistinguishable from legitimate reassociation and no portable test can claim it - and what it still guards, which is the bug class that matters. Mutation-verified against realistic breakage: an off-by-one on the last chunk, a wrong `s_cond` slice, and skipping the chunked path each fail 5 of the 7 tests. 98 passed locally, ruff clean. --- .../pid/test_pid_chunked_equivalence.py | 106 ++++++++++++++---- 1 file changed, 85 insertions(+), 21 deletions(-) diff --git a/tests/backend/pid/test_pid_chunked_equivalence.py b/tests/backend/pid/test_pid_chunked_equivalence.py index 78c115dfb6b..de2c08f14d4 100644 --- a/tests/backend/pid/test_pid_chunked_equivalence.py +++ b/tests/backend/pid/test_pid_chunked_equivalence.py @@ -9,20 +9,26 @@ Measured on an RTX 4090 (torch 2.7.1+cu128), production `PiTBlock` dimensions: - CPU fp32, BL=1024 / 2048 -> max|diff| = 0 (bit-identical) - CUDA fp32, BL=4096 -> max|diff| = 9.5e-07 - CUDA bf16 autocast, BL=4096..32768 -> max|diff| = 1.57e-02, mean|diff| = 2.5e-05 + CPU fp32 (x86-64, MKL), BL=2048 -> max|diff| = 0 + CUDA fp32, BL=4096 -> max|diff| = 9.5e-07 + CUDA bf16 autocast, BL=4096..32768 -> max|diff| = 1.57e-02, mean|diff| = 2.5e-05 Both paths are internally deterministic, so those numbers are systematic, not run noise. The consequence downstream is an image that differs slightly: ~43 dB PSNR end-to-end, visually indistinguishable but not reproducible against an unoptimized decode. -This module therefore states the contract the setting really offers, in two halves: chunking is -*exact* as mathematics, and on the accelerated path it is bounded rather than exact. Relative -tolerances are useless here - activations pass through zero, so `max|rel|` reaches 1e3 on elements -whose absolute error is a bf16 ULP - hence an absolute bound. +The CPU column is *not* a portable guarantee, and this module originally claimed it was. Asserting +`torch.equal` there passed on x86-64/MKL and failed on macOS/Accelerate in CI: splitting a GEMM along +its row dimension can select a micro-kernel with different K-blocking, so bit-exactness is a property +of the BLAS, not of the chunking. What is portable is that chunking only *reassociates* work, so the +contract here is a distance bound on both paths - tight and scaled to the signal on fp32, wider and +absolute under bf16 autocast (relative tolerances are useless there: activations pass through zero, +so `max|rel|` reaches 1e3 on elements whose absolute error is a single bf16 ULP). """ +from collections.abc import Iterator +from contextlib import contextmanager + import pytest import torch @@ -44,6 +50,47 @@ _BF16_ABSOLUTE_TOLERANCE = 5e-2 _BF16_MEAN_ABSOLUTE_TOLERANCE = 1e-3 +# The CPU bound, relative to the signal. This started out as `torch.equal`, which held on x86-64 with +# MKL and failed on macOS/Accelerate in CI: splitting a GEMM along M can select a different +# micro-kernel whose K-blocking differs, so even fp32 is only reassociation-close, not bit-equal. +# +# Calibration, measured at these dimensions (signal ~5.7, so one fp32 ULP is ~6.8e-07 absolute): +# correct code, x86-64/MKL max|diff| = 0 +# attention contribution off by 1e-6 max|diff| = 7.2e-07 (1.3e-07 relative) - one ULP, invisible +# attention contribution off by 1e-4 max|diff| = 1.0e-05 (1.9e-06 relative) - ~15 ULP +# +# 1e-5 relative is ~84 ULP: comfortably above whatever reassociation any BLAS produces, and four +# orders of magnitude below a structurally broken chunk, which lands at O(signal). What it does *not* +# do is catch a uniform scaling error smaller than ~2e-06 relative - that band is indistinguishable +# from legitimate backend reassociation, which is the lesson macOS taught here, so no portable test +# can claim it. Structural mistakes (wrong slice, dropped row, unassembled output) are the bug class +# this guards, and they are nowhere near that band. +_CPU_RELATIVE_TOLERANCE = 1e-5 + + +@contextmanager +def _assert_chunked_path_is_taken() -> Iterator[None]: + """Fail if `_forward_chunked` is not reached. + + Not paranoia: the dispatch guard is `BL > chunk_size`, so the original `batch_size=1` case here + sat exactly *at* 1024 and silently ran unchunked - the test compared the unchunked path against + itself and passed for the wrong reason. Asserting the path is taken keeps a future change to the + guard, the chunk size or the test dimensions from quietly emptying this file out again. + """ + calls: list[int] = [] + original = PiTBlock._forward_chunked + + def _spy(self: PiTBlock, *args: object, **kwargs: object) -> torch.Tensor: + calls.append(1) + return original(self, *args, **kwargs) + + PiTBlock._forward_chunked = _spy # type: ignore[method-assign] + try: + yield + finally: + PiTBlock._forward_chunked = original # type: ignore[method-assign] + assert calls, "the chunked path was never entered - this comparison proves nothing" + def _build_block(device: str) -> PiTBlock: torch.manual_seed(0) @@ -71,23 +118,37 @@ def _inputs(image_px: int, batch_size: int, device: str) -> tuple[torch.Tensor, return x, s_cond -@pytest.mark.parametrize("batch_size", [1, 2]) -def test_chunking_is_exact_mathematics_on_cpu(batch_size: int) -> None: - """Chunking only reorders work, so on a path with no kernel-selection freedom it is bit-exact. - - Batch > 1 is covered here because that is where chunk boundaries stop lining up with image - boundaries: with `BL = B * Hs * Ws`, a chunk can straddle two images, and nothing in - `_forward_chunked` may depend on an image staying within one chunk. +@pytest.mark.parametrize( + ("image_px", "batch_size"), + [ + (768, 1), # BL = 2304: chunk boundaries fall inside a single image + (512, 2), # BL = 2048: chunk boundaries straddle two images + ], +) +def test_chunking_matches_the_unchunked_path_on_cpu(image_px: int, batch_size: int) -> None: + """Chunking splits the adaLN/MLP GEMMs along their row dimension only, so the result must stay + at floating-point-reassociation distance from the unchunked path - orders of magnitude below the + signal, where a real breakage (wrong slice, misassembled output) lands. + + Batch > 1 matters because that is where chunk boundaries stop lining up with image boundaries: + with `BL = B * Hs * Ws` a chunk can straddle two images, and nothing in `_forward_chunked` may + depend on an image staying inside one chunk. """ - image_px = 512 # BL = 1024 (B=1) / 2048 (B=2), i.e. at and above the chunk size block = _build_block("cpu") x, s_cond = _inputs(image_px, batch_size, "cpu") with torch.no_grad(): unchunked = block(x, s_cond, image_px, image_px, _PATCH_SIZE) - chunked = block(x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE) + with _assert_chunked_path_is_taken(): + chunked = block( + x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE + ) - assert torch.equal(chunked, unchunked) + signal = unchunked.abs().max().item() + difference = (chunked - unchunked).abs().max().item() + assert difference <= _CPU_RELATIVE_TOLERANCE * signal, ( + f"chunked output drifted by {difference:.3e} against a signal of {signal:.3e}" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="the divergence only exists on the accelerated path") @@ -105,10 +166,13 @@ def test_chunking_stays_within_the_documented_tolerance_on_cuda_bf16(image_px: i with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): unchunked = block(x, s_cond, image_px, image_px, _PATCH_SIZE) - chunked = block(x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE) - chunked_again = block( - x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE - ) + with _assert_chunked_path_is_taken(): + chunked = block( + x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE + ) + chunked_again = block( + x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE + ) # Determinism first: without it, a passing tolerance would prove nothing about the chunking. assert torch.equal(chunked, chunked_again)