diff --git a/docs/src/content/docs/configuration/low-vram-mode.mdx b/docs/src/content/docs/configuration/low-vram-mode.mdx index fa15cef8735..5e15e78cb9a 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_memory_optimization`) Read on to learn about these features and understand how to fine-tune them for your system and use-cases. @@ -146,6 +147,22 @@ Invoke has the option of keeping a RAM copy of all model weights, even when they keep_ram_copy_of_weights: false ``` +### PiD decode optimization + +PiD decodes directly into high-resolution pixels, so its activation and sampler memory can exceed the working memory needed by a normal VAE decode. Partial model loading reduces memory used by PiD's weights, but does not reduce these full-resolution intermediates. + +To reduce PiD's peak VRAM use, enable its experimental memory optimization in `invokeai.yaml` and restart InvokeAI: + +```yaml +pid_memory_optimization: true +``` + +This setting processes parts of the PiD pixel pathway in chunks and uses float32 instead of float64 sampler intermediates. It applies to every supported PiD decoder. Measured on an RTX 4090, peak activation memory for a 2048px decode drops from ~3.7 GB to ~1.5 GB (1024px: ~0.9 GB to ~0.5 GB). + +It is disabled by default because it is **not output-preserving**: neither change is bit-exact, and the few-step sampler amplifies the difference into a slightly different image. The delta is small — around 43 dB PSNR, visually indistinguishable in side-by-side comparisons — but it is real, so the same seed and workflow will not reproduce an unoptimized decode exactly. Decoding speed is roughly unchanged: chunking costs about 4%, which the cheaper sampler math largely offsets. + +When the setting is active, each decode logs the resolution, the patch-token count and whether chunking actually engaged — worth checking, since the option is server-wide and never recorded in image metadata. See [PiD Super-Resolution Decode](/features/pid-decode/) for supported models and usage. + ### Disabling Nvidia sysmem fallback (Windows only) On Windows, Nvidia GPUs are able to use system RAM when their VRAM fills up via **sysmem fallback**. While it sounds like a good idea on the surface, in practice it causes massive slowdowns during generation. diff --git a/docs/src/content/docs/features/pid-decode.mdx b/docs/src/content/docs/features/pid-decode.mdx index 3897844f75c..7f6745d6d03 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,16 @@ PiD is available in both the **Generate** tab (text-to-image) and on the **Canva ## Tips & limitations +On GPUs with limited VRAM, enable the experimental PiD memory optimizations in `invokeai.yaml`, then restart InvokeAI: + +```yaml +pid_memory_optimization: true +``` + +This setting applies to every supported PiD decoder. It reduces peak activation and sampler memory by processing parts of the pixel pathway in chunks and by running the sampler's intermediates in float32 instead of float64 — measured on an RTX 4090, a 2048px decode peaks at ~1.5 GB of activations instead of ~3.7 GB. Decoding speed is roughly unchanged. + +It is disabled by default because it **changes the decoded image**. Neither the chunked pathway nor the float32 sampler math is bit-exact with the default path, and the few-step sampler amplifies that into a small but real difference (~43 dB PSNR — visually indistinguishable, numerically not identical). A seed that reproduced an image with the setting off will not reproduce it exactly with the setting on, and because this is a server-wide `invokeai.yaml` option it is not recorded in image metadata — each decode logs that it ran optimized instead. + - **Turn off "Scale Before Processing"** on the Canvas when using PiD — PiD already decodes at 4×, so pre-scaling would inflate the work and is blocked. - **Inpaint / Outpaint** are not supported with PiD yet; use text-to-image or image-to-image. - **SDXL Refiner** cannot be combined with PiD — disable one of them. diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index 2c183f55400..76b464eef34 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. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.", + "env_var": "INVOKEAI_PID_MEMORY_OPTIMIZATION", + "literal_values": [], + "name": "pid_memory_optimization", + "required": false, + "type": "", + "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..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__}.") @@ -226,7 +231,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latent=denorm_latent, caption_embs=caption_embs, caption_mask=caption_mask, - config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed), + config=PiDDecodeConfig( + num_inference_steps=self.num_inference_steps, + seed=self.seed, + pid_memory_optimization=pid_memory_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/flux_pid_decode.py b/invokeai/app/invocations/flux_pid_decode.py index 182bcc9cdaf..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__}.") @@ -149,7 +154,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latent=denorm_latent, caption_embs=caption_embs, caption_mask=caption_mask, - config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed), + config=PiDDecodeConfig( + num_inference_steps=self.num_inference_steps, + seed=self.seed, + pid_memory_optimization=pid_memory_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/pid_upscale.py b/invokeai/app/invocations/pid_upscale.py index 3a50ffd4cf3..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__}.") @@ -189,7 +194,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latent=latent_on_device, caption_embs=caption_embs, caption_mask=caption_mask, - config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed), + config=PiDDecodeConfig( + num_inference_steps=self.num_inference_steps, + seed=self.seed, + pid_memory_optimization=pid_memory_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/qwen_image_pid_decode.py b/invokeai/app/invocations/qwen_image_pid_decode.py index 3503da8bbd6..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__}.") @@ -215,7 +220,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latent=denorm_latent, caption_embs=caption_embs, caption_mask=caption_mask, - config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed), + config=PiDDecodeConfig( + num_inference_steps=self.num_inference_steps, + seed=self.seed, + pid_memory_optimization=pid_memory_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/sd3_pid_decode.py b/invokeai/app/invocations/sd3_pid_decode.py index b22b0fa3bf1..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__}.") @@ -142,7 +147,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latent=denorm_latent, caption_embs=caption_embs, caption_mask=caption_mask, - config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed), + config=PiDDecodeConfig( + num_inference_steps=self.num_inference_steps, + seed=self.seed, + pid_memory_optimization=pid_memory_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/sdxl_pid_decode.py b/invokeai/app/invocations/sdxl_pid_decode.py index 5a09d03fe0e..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__}.") @@ -188,7 +193,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latent=denorm_latent, caption_embs=caption_embs, caption_mask=caption_mask, - config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed), + config=PiDDecodeConfig( + num_inference_steps=self.num_inference_steps, + seed=self.seed, + pid_memory_optimization=pid_memory_optimization, + ), ) TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/z_image_pid_decode.py b/invokeai/app/invocations/z_image_pid_decode.py index 392f11e3639..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__}.") @@ -198,7 +203,11 @@ def invoke(self, context: InvocationContext) -> ImageOutput: latent=denorm_latent, caption_embs=caption_embs, caption_mask=caption_mask, - config=PiDDecodeConfig(num_inference_steps=self.num_inference_steps, seed=self.seed), + config=PiDDecodeConfig( + num_inference_steps=self.num_inference_steps, + seed=self.seed, + pid_memory_optimization=pid_memory_optimization, + ), ) context.logger.info( f"PiD output stats: shape={tuple(x0.shape)} dtype={x0.dtype} " diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index ae1e38e0ccc..1dc8b560326 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_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). @@ -220,6 +221,7 @@ class InvokeAIAppConfig(BaseSettings): # GENERATION sequential_guidance: bool = Field(default=False, description="Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.") + pid_memory_optimization: bool = Field(default=False, description="Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.") attention_type: ATTENTION_TYPE = Field(default="auto", description="Attention type.") attention_slice_size: ATTENTION_SLICE_SIZE = Field(default="auto", description='Slice size, valid when attention_type=="sliced".') force_tiled_decode: bool = Field(default=False, description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).") 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 6fdda4917db..f71b15176f1 100644 --- a/invokeai/backend/pid/_src/networks/pixeldit_official.py +++ b/invokeai/backend/pid/_src/networks/pixeldit_official.py @@ -468,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, ...) @@ -489,25 +496,113 @@ 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. + 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) + + 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: + """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] + 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..9ecfec3ad32 100644 --- a/invokeai/backend/pid/decode.py +++ b/invokeai/backend/pid/decode.py @@ -27,6 +27,9 @@ 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 # --------------------------------------------------------------------------- # Network hyperparameters per backbone @@ -131,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: @@ -146,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: @@ -225,9 +257,32 @@ 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*.""" +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*. + + 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) + 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) @@ -245,6 +300,7 @@ def _student_sample_loop( sample_type: str = "sde", autocast_dtype: Optional[torch.dtype] = None, generator: Optional[torch.Generator] = None, + pid_memory_optimization: bool = False, ) -> Tensor: """Few-step distilled sampler. @@ -279,9 +335,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_memory_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_memory_optimization=pid_memory_optimization) eps_infer = torch.randn( x0_pred.shape, device=x0_pred.device, @@ -297,7 +354,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_memory_optimization=pid_memory_optimization) return x @@ -322,6 +379,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_memory_optimization: bool = False student_t_list: list[float] = field(default_factory=lambda: list(_STUDENT_T_LIST)) @@ -415,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, @@ -427,6 +500,7 @@ def decode( sample_type=cfg.sample_type, autocast_dtype=autocast_dtype, generator=gen, + 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 e19d0163e31..f91163dbeb2 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -48320,6 +48320,12 @@ "description": "Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.", "default": false }, + "pid_memory_optimization": { + "type": "boolean", + "title": "Pid Memory Optimization", + "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": { "type": "string", "enum": ["auto", "normal", "xformers", "sliced", "torch-sdp"], @@ -48584,7 +48590,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_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 596472a0d43..7aeac79574a 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -18630,6 +18630,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. 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). @@ -18971,6 +18972,12 @@ export type components = { * @default false */ sequential_guidance?: boolean; + /** + * Pid Memory Optimization + * @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; /** * Attention Type * @description Attention type. 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..de2c08f14d4 --- /dev/null +++ b/tests/backend/pid/test_pid_chunked_equivalence.py @@ -0,0 +1,234 @@ +"""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 (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. + +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 + +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 + +# 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) + 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( + ("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. + """ + 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) + with _assert_chunked_path_is_taken(): + chunked = block( + x, s_cond, image_px, image_px, _PATCH_SIZE, activation_chunk_size=_PID_ACTIVATION_CHUNK_SIZE + ) + + 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") +@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) + 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) + + 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 44e8f377617..97230d7482b 100644 --- a/tests/backend/pid/test_pid_decode.py +++ b/tests/backend/pid/test_pid_decode.py @@ -1,10 +1,19 @@ """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 ( + PiDDecodeConfig, + _get_t_list, + _student_sample_loop, + _velocity_to_x0, + assert_pid_decoder_matches_base, + estimate_pid_decode_working_memory, +) _CPU = torch.device("cpu") @@ -35,6 +44,125 @@ 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, pid_memory_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_memory_optimization=False) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@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_memory_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_memory_optimization=pid_memory_optimization, + ) + + assert activation_chunk_sizes == [expected_chunk_size] + + +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") diff --git a/tests/backend/pid/test_pixeldit_official.py b/tests/backend/pid/test_pixeldit_official.py new file mode 100644 index 00000000000..57cff37f337 --- /dev/null +++ b/tests/backend/pid/test_pixeldit_official.py @@ -0,0 +1,124 @@ +import math + +import pytest +import torch + +from invokeai.backend.pid._src.networks.pixeldit_official import PiTBlock + + +def _build_pit_block() -> 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, + ).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() + chunk_size = 3 + chunked = _build_pit_block() + 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, + activation_chunk_size=chunk_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: + 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"): + 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() + 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, + 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..78d5dbe3466 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_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_memory_optimization: true\n') + + 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): """Test reading configuration from a file.""" temp_config_file = tmp_path / "temp_invokeai.yaml"