diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 00000000000..b1a62ff9ba2 --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1,4 @@ +# These are supported funding model platforms +# https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/displaying-a-sponsor-button-in-your-repository + +github: invoke-ai diff --git a/README.md b/README.md index afc0211bd91..6247b6ade6f 100644 --- a/README.md +++ b/README.md @@ -100,6 +100,15 @@ Get started with contributing by reading our [contribution documentation][contri We hope you enjoy using Invoke as much as we enjoy creating it, and we hope you will elect to become part of our community. +## Sponsors + +Invoke's open-source development is powered by our sponsors. If Invoke is valuable to you or your business, please consider [sponsoring us][sponsor link] — it directly funds maintenance, new features, and community support. + + + +[![Sponsor Invoke](https://img.shields.io/badge/Sponsor-Invoke-ea4aaa?logo=githubsponsors&logoColor=white)][sponsor link] + ## Thanks Invoke is a combined effort of [passionate and talented people from across the world][contributors]. We thank them for their time, hard work and effort. @@ -112,6 +121,7 @@ Original portions of the software are Copyright © 2024 by respective contributo [github issues]: https://github.com/invoke-ai/InvokeAI/issues [docs home]: https://invoke.ai [installation docs]: https://invoke.ai/start-here/installation/ +[sponsor link]: https://github.com/sponsors/invoke-ai [#dev-chat]: https://discord.com/channels/1020123559063990373/1049495067846524939 [contributing docs]: https://invoke.ai/contributing/ [CI checks on main badge]: https://flat.badgen.net/github/checks/invoke-ai/InvokeAI/main?label=CI%20status%20on%20main&cache=900&icon=github diff --git a/invokeai/app/api/extract_metadata_from_image.py b/invokeai/app/api/extract_metadata_from_image.py index 054b3cc38cc..642a4b1bce1 100644 --- a/invokeai/app/api/extract_metadata_from_image.py +++ b/invokeai/app/api/extract_metadata_from_image.py @@ -48,7 +48,11 @@ def extract_metadata_from_image( stringified_metadata: str | None = None # Use the metadata override if provided, else attempt to extract it from the image file. - metadata_raw = invokeai_metadata_override or pil_image.info.get("invokeai_metadata", None) + metadata_raw = ( + invokeai_metadata_override + if invokeai_metadata_override is not None + else pil_image.info.get("invokeai_metadata", None) + ) # If the metadata is present in the image file, we will attempt to parse it as JSON. When we create images, # we always store metadata as a stringified JSON dict. So, we expect it to be a string here. @@ -66,7 +70,11 @@ def extract_metadata_from_image( # We expect the workflow, if embedded in the image, to be a JSON-stringified WorkflowWithoutID. We will store it # as a string. - workflow_raw: str | None = invokeai_workflow_override or pil_image.info.get("invokeai_workflow", None) + workflow_raw: str | None = ( + invokeai_workflow_override + if invokeai_workflow_override is not None + else pil_image.info.get("invokeai_workflow", None) + ) # The fallback value for workflow is None. stringified_workflow: str | None = None @@ -85,7 +93,9 @@ def extract_metadata_from_image( # We expect the workflow, if embedded in the image, to be a JSON-stringified Graph. We will store it as a # string. - graph_raw: str | None = invokeai_graph_override or pil_image.info.get("invokeai_graph", None) + graph_raw: str | None = ( + invokeai_graph_override if invokeai_graph_override is not None else pil_image.info.get("invokeai_graph", None) + ) # The fallback value for graph is None. stringified_graph: str | None = None diff --git a/invokeai/app/invocations/z_image_denoise.py b/invokeai/app/invocations/z_image_denoise.py index c1e864ea179..576d10ac9a1 100644 --- a/invokeai/app/invocations/z_image_denoise.py +++ b/invokeai/app/invocations/z_image_denoise.py @@ -558,6 +558,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: transformer=transformer, regional_attn_mask=regional_extension.regional_attn_mask, img_seq_len=img_seq_len, + positive_cap_feats=pos_prompt_embeds, ) ) diff --git a/invokeai/backend/z_image/z_image_transformer_patch.py b/invokeai/backend/z_image/z_image_transformer_patch.py index 6e42e678558..a6707fb7f50 100644 --- a/invokeai/backend/z_image/z_image_transformer_patch.py +++ b/invokeai/backend/z_image/z_image_transformer_patch.py @@ -4,25 +4,38 @@ from typing import Callable, List, Optional, Tuple import torch -from torch.nn.utils.rnn import pad_sequence def create_regional_forward( original_forward: Callable, regional_attn_mask: torch.Tensor, img_seq_len: int, + positive_cap_feats: torch.Tensor, ) -> Callable: """Create a modified forward function that uses a regional attention mask. - The regional attention mask replaces the internally computed padding mask, - allowing for regional prompting where different image regions attend to - different text prompts. + The regional attention mask replaces the internally computed padding mask on the + main transformer layers (alternating with the plain padding mask), allowing for + regional prompting where different image regions attend to different text prompts. + + This delegates to the model's own helper methods (``patchify_and_embed``, + ``_prepare_sequence``, ``_build_unified_sequence``) so it stays in sync with the + upstream diffusers ``ZImageTransformer2DModel.forward`` implementation. Only the + main-layer attention mask is overridden. Args: - original_forward: The original forward method of ZImageTransformer2DModel. - regional_attn_mask: Attention mask of shape (seq_len, seq_len) where - seq_len = img_seq_len + txt_seq_len. - img_seq_len: Number of image tokens in the sequence. + original_forward: The original forward method of ZImageTransformer2DModel + (kept for signature compatibility; not used directly). + regional_attn_mask: Boolean attention mask of shape (seq_len, seq_len) where + seq_len = img_seq_len + txt_seq_len, ordered [img, txt]. + img_seq_len: Number of (unpadded) image tokens in the sequence. + positive_cap_feats: The exact caption-embedding tensor the regional mask was + built for (the conditioned/positive pass). The regional mask is applied only + to forward calls whose ``cap_feats`` is this same object; the negative/CFG + pass supplies a different tensor and is left to run with the plain padding + mask. Identity is used instead of a token-length heuristic so the positive + and negative passes can never be confused even when their padded lengths + coincide. Returns: A modified forward function with regional attention support. @@ -38,160 +51,146 @@ def regional_forward( ) -> Tuple[List[torch.Tensor], dict]: """Modified forward with regional attention mask injection. - This is based on the original ZImageTransformer2DModel.forward but - replaces the padding-based attention mask with a regional attention mask. + Mirrors the basic (non-omni) path of ZImageTransformer2DModel.forward but + injects a regional attention mask into the main transformer layers. """ assert patch_size in self.all_patch_size assert f_patch_size in self.all_f_patch_size - bsz = len(x) device = x[0].device - t_scaled = t * self.t_scale - t_emb = self.t_embedder(t_scaled) - SEQ_MULTI_OF = 32 # From diffusers transformer_z_image.py + # Identify which caption inputs belong to the conditioned (positive) pass the regional + # mask was built for. Capture this before patchify_and_embed reassigns ``cap_feats``. + # The negative/CFG pass supplies a different tensor, so object identity distinguishes the + # passes regardless of token length (avoids the positive mask leaking into the uncond + # prediction when prompt lengths happen to pad to the same multiple). + is_positive_pass = [ci is positive_cap_feats for ci in cap_feats] + + # Single adaLN embedding for all tokens (basic mode). + adaln_input = self.t_embedder(t * self.t_scale).type_as(x[0]) - # Patchify and embed (reusing the original method) + # Patchify & embed (basic mode: single image per batch item). ( x, cap_feats, x_size, x_pos_ids, cap_pos_ids, - x_inner_pad_mask, - cap_inner_pad_mask, + x_pad_mask, + cap_pad_mask, ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size) - # x embed & refine - x_item_seqlens = [len(_) for _ in x] - assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) - x_max_item_seqlen = max(x_item_seqlens) - - x_cat = torch.cat(x, dim=0) - x_cat = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x_cat) - - adaln_input = t_emb.type_as(x_cat) - x_cat[torch.cat(x_inner_pad_mask)] = self.x_pad_token - x_list = list(x_cat.split(x_item_seqlens, dim=0)) - x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) - - x_padded = pad_sequence(x_list, batch_first=True, padding_value=0.0) - x_freqs_cis_padded = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0) - x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(x_item_seqlens): - x_attn_mask[i, :seq_len] = 1 - - # Process through noise_refiner - if torch.is_grad_enabled() and self.gradient_checkpointing: - for layer in self.noise_refiner: - x_padded = self._gradient_checkpointing_func( - layer, x_padded, x_attn_mask, x_freqs_cis_padded, adaln_input - ) - else: - for layer in self.noise_refiner: - x_padded = layer(x_padded, x_attn_mask, x_freqs_cis_padded, adaln_input) - - # cap embed & refine - cap_item_seqlens = [len(_) for _ in cap_feats] - assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) - cap_max_item_seqlen = max(cap_item_seqlens) - - cap_cat = torch.cat(cap_feats, dim=0) - cap_cat = self.cap_embedder(cap_cat) - cap_cat[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token - cap_list = list(cap_cat.split(cap_item_seqlens, dim=0)) - cap_freqs_cis = list(self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0)) - - cap_padded = pad_sequence(cap_list, batch_first=True, padding_value=0.0) - cap_freqs_cis_padded = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0) - cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(cap_item_seqlens): - cap_attn_mask[i, :seq_len] = 1 - - # Process through context_refiner - if torch.is_grad_enabled() and self.gradient_checkpointing: - for layer in self.context_refiner: - cap_padded = self._gradient_checkpointing_func(layer, cap_padded, cap_attn_mask, cap_freqs_cis_padded) - else: - for layer in self.context_refiner: - cap_padded = layer(cap_padded, cap_attn_mask, cap_freqs_cis_padded) - - # Unified sequence: [img_tokens, txt_tokens] - unified = [] - unified_freqs_cis = [] - for i in range(bsz): - x_len = x_item_seqlens[i] - cap_len = cap_item_seqlens[i] - unified.append(torch.cat([x_padded[i][:x_len], cap_padded[i][:cap_len]])) - unified_freqs_cis.append(torch.cat([x_freqs_cis_padded[i][:x_len], cap_freqs_cis_padded[i][:cap_len]])) - - unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens, strict=False)] - assert unified_item_seqlens == [len(_) for _ in unified] - unified_max_item_seqlen = max(unified_item_seqlens) - - unified_padded = pad_sequence(unified, batch_first=True, padding_value=0.0) - unified_freqs_cis_padded = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0) + # X embed & refine. + x_seqlens = [len(xi) for xi in x] + x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](torch.cat(x, dim=0)) + x, x_freqs, x_mask, _, _ = self._prepare_sequence( + list(x.split(x_seqlens, dim=0)), x_pos_ids, x_pad_mask, self.x_pad_token, None, device + ) + for layer in self.noise_refiner: + x = layer(x, x_mask, x_freqs, adaln_input, None, None, None) + + # Cap embed & refine. + cap_seqlens = [len(ci) for ci in cap_feats] + cap_feats = self.cap_embedder(torch.cat(cap_feats, dim=0)) + cap_feats, cap_freqs, cap_mask, _, _ = self._prepare_sequence( + list(cap_feats.split(cap_seqlens, dim=0)), cap_pos_ids, cap_pad_mask, self.cap_pad_token, None, device + ) + for layer in self.context_refiner: + cap_feats = layer(cap_feats, cap_mask, cap_freqs) + + # Unified sequence: basic mode order [x, cap]. + unified, unified_freqs, unified_mask, _ = self._build_unified_sequence( + x, + x_freqs, + x_seqlens, + None, + cap_feats, + cap_freqs, + cap_seqlens, + None, + None, + None, + None, + None, + False, # omni_mode + device, + ) + + bsz = unified.shape[0] + unified_seqlen = unified.shape[1] # --- REGIONAL ATTENTION MASK INJECTION --- - # Instead of using the padding mask, we use the regional attention mask - # The regional mask is (seq_len, seq_len), we need to expand it to (batch, seq_len, seq_len) - # and then add the batch dimension for broadcasting: (batch, 1, seq_len, seq_len) - - # Expand regional mask to match the actual sequence length (may include padding) - if regional_attn_mask.shape[0] != unified_max_item_seqlen: - # Pad the regional mask to match unified sequence length - padded_regional_mask = torch.zeros( - (unified_max_item_seqlen, unified_max_item_seqlen), - dtype=regional_attn_mask.dtype, - device=device, + # The regional mask is (S, S) with S = img_seq_len + txt_seq_len, ordered [img, txt], + # using the *unpadded* image and text token counts. In the unified sequence, however, + # both the image block and the caption block are individually padded to a multiple of + # SEQ_MULTI_OF, so the real layout per item is: + # [ img_real | img_pad | txt_real | txt_pad ] + # We therefore scatter the four regional sub-blocks (img-img, img-txt, txt-img, txt-txt) + # into their padding-aware positions instead of assuming a contiguous top-left block. + # + # The patched forward also runs for the negative/CFG pass (a different prompt). The + # regional mask was built for the positive prompt only, so we apply it only to the + # conditioned items and fall back to the plain padding mask otherwise. + regional = regional_attn_mask.to(device=device, dtype=torch.bool) + txt_seq_len = regional.shape[0] - img_seq_len + + # Decide per item whether the regional mask applies, using only cheap scalar checks, so + # that on passes that never match (e.g. every negative/CFG pass) we avoid materializing + # the (bsz, 1, S, S) float mask at all. + applied_regional = [ + is_positive_pass[i] + and txt_seq_len > 0 + and img_seq_len <= x_seqlens[i] + and x_seqlens[i] + cap_seqlens[i] <= unified_seqlen + for i in range(bsz) + ] + + # Main transformer layers: alternate regional mask (even) with plain padding mask (odd). + # If no item matched the positive pass, skip regional injection entirely. + use_regional = any(applied_regional) + + float_mask = None + if use_regional: + # Build a per-item additive float mask. Start from the plain padding mask (0 where a + # token is valid, -inf where it is padding) so non-matching items behave normally. + neg_inf = torch.finfo(unified.dtype).min + zero = torch.zeros((), dtype=unified.dtype, device=device) + float_mask = ( + torch.where( + unified_mask.bool().unsqueeze(1).unsqueeze(1), # (bsz, 1, 1, S) + zero, + torch.full((), neg_inf, dtype=unified.dtype, device=device), + ) + .expand(bsz, 1, unified_seqlen, unified_seqlen) + .clone() ) - mask_size = min(regional_attn_mask.shape[0], unified_max_item_seqlen) - padded_regional_mask[:mask_size, :mask_size] = regional_attn_mask[:mask_size, :mask_size] - else: - padded_regional_mask = regional_attn_mask.to(device) - - # Convert boolean mask to additive float mask for attention - # True (attend) -> 0.0, False (block) -> -inf - # This is required because the attention backend expects additive masks for 4D inputs - # Use bfloat16 to match the transformer's query dtype - float_mask = torch.zeros_like(padded_regional_mask, dtype=torch.bfloat16) - float_mask[~padded_regional_mask] = float("-inf") - - # Expand to (batch, 1, seq_len, seq_len) for attention - unified_attn_mask = float_mask.unsqueeze(0).unsqueeze(0).expand(bsz, 1, -1, -1) - - # Process through main layers with regional attention mask - if torch.is_grad_enabled() and self.gradient_checkpointing: - for layer_idx, layer in enumerate(self.layers): - # Alternate between regional mask and full attention - if layer_idx % 2 == 0: - unified_padded = self._gradient_checkpointing_func( - layer, unified_padded, unified_attn_mask, unified_freqs_cis_padded, adaln_input - ) - else: - # Use padding mask only for odd layers (allows global coherence) - padding_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(unified_item_seqlens): - padding_mask[i, :seq_len] = 1 - unified_padded = self._gradient_checkpointing_func( - layer, unified_padded, padding_mask, unified_freqs_cis_padded, adaln_input - ) - else: - for layer_idx, layer in enumerate(self.layers): - # Alternate between regional mask and full attention - if layer_idx % 2 == 0: - unified_padded = layer(unified_padded, unified_attn_mask, unified_freqs_cis_padded, adaln_input) - else: - # Use padding mask only for odd layers (allows global coherence) - padding_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(unified_item_seqlens): - padding_mask[i, :seq_len] = 1 - unified_padded = layer(unified_padded, padding_mask, unified_freqs_cis_padded, adaln_input) - - # Final layer - unified_out = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified_padded, adaln_input) - unified_list = list(unified_out.unbind(dim=0)) - x_out = self.unpatchify(unified_list, x_size, patch_size, f_patch_size) + + for i in range(bsz): + if not applied_regional[i]: + continue + x_len = x_seqlens[i] + + ii, it = slice(0, img_seq_len), slice(img_seq_len, img_seq_len + txt_seq_len) + ui = slice(0, img_seq_len) # real image positions in unified item + ut = slice(x_len, x_len + txt_seq_len) # real text positions in unified item + + # Reset the masked region so only regional rules apply to real img/txt tokens; + # their rows start fully blocked and we open the allowed sub-blocks below. + float_mask[i, 0, ui, :] = neg_inf + float_mask[i, 0, ut, :] = neg_inf + + float_mask[i, 0, ui, ui] = torch.where(regional[ii, ii], zero, neg_inf) # img -> img + float_mask[i, 0, ui, ut] = torch.where(regional[ii, it], zero, neg_inf) # img -> txt + float_mask[i, 0, ut, ui] = torch.where(regional[it, ii], zero, neg_inf) # txt -> img + float_mask[i, 0, ut, ut] = torch.where(regional[it, it], zero, neg_inf) # txt -> txt + + for layer_idx, layer in enumerate(self.layers): + attn_mask = float_mask if (use_regional and layer_idx % 2 == 0) else unified_mask + unified = layer(unified, attn_mask, unified_freqs, adaln_input, None, None, None) + + # Final layer + unpatchify. + unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, c=adaln_input) + x_out = self.unpatchify(list(unified.unbind(dim=0)), x_size, patch_size, f_patch_size) return x_out, {} @@ -203,6 +202,7 @@ def patch_transformer_for_regional_prompting( transformer, regional_attn_mask: Optional[torch.Tensor], img_seq_len: int, + positive_cap_feats: Optional[torch.Tensor] = None, ): """Context manager to temporarily patch the transformer for regional prompting. @@ -211,6 +211,9 @@ def patch_transformer_for_regional_prompting( regional_attn_mask: Regional attention mask of shape (seq_len, seq_len). If None, the transformer is not patched. img_seq_len: Number of image tokens. + positive_cap_feats: The caption-embedding tensor the regional mask was built for. + Required when ``regional_attn_mask`` is provided; the mask is applied only to + forward calls whose ``cap_feats`` is this exact object (the conditioned pass). Yields: The (possibly patched) transformer. @@ -220,11 +223,14 @@ def patch_transformer_for_regional_prompting( yield transformer return + if positive_cap_feats is None: + raise ValueError("positive_cap_feats is required when regional_attn_mask is provided") + # Store original forward original_forward = transformer.forward # Create and bind the regional forward - regional_fwd = create_regional_forward(original_forward, regional_attn_mask, img_seq_len) + regional_fwd = create_regional_forward(original_forward, regional_attn_mask, img_seq_len, positive_cap_feats) transformer.forward = lambda *args, **kwargs: regional_fwd(transformer, *args, **kwargs) try: diff --git a/invokeai/frontend/webv2/src/workbench/auth/capabilities.test.ts b/invokeai/frontend/webv2/src/workbench/auth/capabilities.test.ts new file mode 100644 index 00000000000..fe95377119d --- /dev/null +++ b/invokeai/frontend/webv2/src/workbench/auth/capabilities.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, it } from 'vitest'; + +import type { AuthSession } from './session'; + +import { getCapabilities } from './capabilities'; + +const createSession = (overrides: Partial): AuthSession => ({ + multiuserEnabled: false, + phase: 'ready', + sessionExpired: false, + setupRequired: false, + strictPasswordChecking: false, + user: null, + ...overrides, +}); + +describe('getCapabilities', () => { + it('allows clearing intermediates in single-user mode', () => { + expect(getCapabilities(createSession({ multiuserEnabled: false })).canClearIntermediates).toBe(true); + }); + + it('allows clearing intermediates for multi-user admins only', () => { + expect( + getCapabilities( + createSession({ + multiuserEnabled: true, + user: { is_admin: true } as AuthSession['user'], + }) + ).canClearIntermediates + ).toBe(true); + + expect( + getCapabilities( + createSession({ + multiuserEnabled: true, + user: { is_admin: false } as AuthSession['user'], + }) + ).canClearIntermediates + ).toBe(false); + }); +}); diff --git a/invokeai/frontend/webv2/src/workbench/auth/capabilities.ts b/invokeai/frontend/webv2/src/workbench/auth/capabilities.ts index 92c61cfd257..d9bf11a511a 100644 --- a/invokeai/frontend/webv2/src/workbench/auth/capabilities.ts +++ b/invokeai/frontend/webv2/src/workbench/auth/capabilities.ts @@ -1,6 +1,7 @@ import { useAuthSession, type AuthSession } from './session'; export interface Capabilities { + canClearIntermediates: boolean; canManageModels: boolean; canManageNodes: boolean; canManageUsers: boolean; @@ -11,6 +12,7 @@ export const getCapabilities = (session: AuthSession): Capabilities => { const isAdmin = isSingleUser || session.user?.is_admin === true; return { + canClearIntermediates: isAdmin, canManageModels: isAdmin, canManageNodes: isAdmin, canManageUsers: session.multiuserEnabled && session.user?.is_admin === true, diff --git a/invokeai/frontend/webv2/src/workbench/images/intermediates.test.ts b/invokeai/frontend/webv2/src/workbench/images/intermediates.test.ts new file mode 100644 index 00000000000..2cdab856d7f --- /dev/null +++ b/invokeai/frontend/webv2/src/workbench/images/intermediates.test.ts @@ -0,0 +1,69 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const mocks = vi.hoisted(() => ({ + apiFetchJson: vi.fn(), +})); + +vi.mock('@workbench/backend/http', () => ({ + apiFetchJson: mocks.apiFetchJson, +})); + +describe('intermediates API helpers', () => { + beforeEach(() => { + mocks.apiFetchJson.mockReset(); + }); + + it('reads the intermediate image count from the images endpoint', async () => { + mocks.apiFetchJson.mockResolvedValueOnce(7); + const { getIntermediatesCount } = await import('./intermediates'); + + await expect(getIntermediatesCount()).resolves.toBe(7); + + expect(mocks.apiFetchJson).toHaveBeenCalledWith('/api/v1/images/intermediates'); + }); + + it('clears intermediate images with DELETE and returns the cleared count', async () => { + mocks.apiFetchJson.mockResolvedValueOnce(3); + const { clearIntermediates } = await import('./intermediates'); + + await expect(clearIntermediates()).resolves.toBe(3); + + expect(mocks.apiFetchJson).toHaveBeenCalledWith('/api/v1/images/intermediates', { method: 'DELETE' }); + }); +}); + +describe('clear intermediates UI state', () => { + it('requires permission, a loaded nonzero count, and no active queue work', async () => { + const { getClearIntermediatesState } = await import('./intermediates'); + + expect( + getClearIntermediatesState({ canClearIntermediates: false, hasActiveQueueWork: false, intermediatesCount: 4 }) + ).toEqual({ disabled: true, reason: 'Only administrators can clear intermediates.' }); + + expect( + getClearIntermediatesState({ canClearIntermediates: true, hasActiveQueueWork: true, intermediatesCount: 4 }) + ).toEqual({ disabled: true, reason: 'Wait for pending or running queue work to finish.' }); + + expect( + getClearIntermediatesState({ canClearIntermediates: true, hasActiveQueueWork: false, intermediatesCount: 0 }) + ).toEqual({ disabled: true, reason: 'There are no intermediates to clear.' }); + + expect( + getClearIntermediatesState({ canClearIntermediates: true, hasActiveQueueWork: false, intermediatesCount: null }) + ).toEqual({ disabled: true, reason: 'Loading intermediate count.' }); + + expect( + getClearIntermediatesState({ canClearIntermediates: true, hasActiveQueueWork: false, intermediatesCount: 4 }) + ).toEqual({ disabled: false }); + }); + + it('describes the destructive clear action before confirmation', async () => { + const { getClearIntermediatesConfirmation } = await import('./intermediates'); + + expect(getClearIntermediatesConfirmation(5)).toEqual({ + body: 'This permanently deletes 5 temporary intermediate images and clears canvas layers or staged images that may reference them. Final gallery images are not removed.', + confirmLabel: 'Clear intermediates', + title: 'Clear intermediate images?', + }); + }); +}); diff --git a/invokeai/frontend/webv2/src/workbench/images/intermediates.ts b/invokeai/frontend/webv2/src/workbench/images/intermediates.ts new file mode 100644 index 00000000000..b2562c9f237 --- /dev/null +++ b/invokeai/frontend/webv2/src/workbench/images/intermediates.ts @@ -0,0 +1,54 @@ +import { apiFetchJson } from '@workbench/backend/http'; + +const INTERMEDIATES_URL = '/api/v1/images/intermediates'; + +export interface ClearIntermediatesStateInput { + canClearIntermediates: boolean; + hasActiveQueueWork: boolean; + intermediatesCount: number | null; +} + +export interface ClearIntermediatesState { + disabled: boolean; + reason?: string; +} + +export interface ClearIntermediatesConfirmation { + body: string; + confirmLabel: string; + title: string; +} + +export const getIntermediatesCount = (): Promise => apiFetchJson(INTERMEDIATES_URL); + +export const clearIntermediates = (): Promise => apiFetchJson(INTERMEDIATES_URL, { method: 'DELETE' }); + +export const getClearIntermediatesState = ({ + canClearIntermediates, + hasActiveQueueWork, + intermediatesCount, +}: ClearIntermediatesStateInput): ClearIntermediatesState => { + if (!canClearIntermediates) { + return { disabled: true, reason: 'Only administrators can clear intermediates.' }; + } + + if (hasActiveQueueWork) { + return { disabled: true, reason: 'Wait for pending or running queue work to finish.' }; + } + + if (intermediatesCount === null) { + return { disabled: true, reason: 'Loading intermediate count.' }; + } + + if (intermediatesCount === 0) { + return { disabled: true, reason: 'There are no intermediates to clear.' }; + } + + return { disabled: false }; +}; + +export const getClearIntermediatesConfirmation = (count: number): ClearIntermediatesConfirmation => ({ + body: `This permanently deletes ${count} temporary intermediate image${count === 1 ? '' : 's'} and clears canvas layers or staged images that may reference them. Final gallery images are not removed.`, + confirmLabel: 'Clear intermediates', + title: 'Clear intermediate images?', +}); diff --git a/invokeai/frontend/webv2/src/workbench/settings/ClearIntermediatesSetting.tsx b/invokeai/frontend/webv2/src/workbench/settings/ClearIntermediatesSetting.tsx new file mode 100644 index 00000000000..8a5e70871d6 --- /dev/null +++ b/invokeai/frontend/webv2/src/workbench/settings/ClearIntermediatesSetting.tsx @@ -0,0 +1,166 @@ +import { HStack, Stack, Text } from '@chakra-ui/react'; +import { useCapabilities } from '@workbench/auth/capabilities'; +import { getApiErrorMessage } from '@workbench/backend/http'; +import { Button, ConfirmDialog } from '@workbench/components/ui'; +import { createExternalStore } from '@workbench/externalStore'; +import { + clearIntermediates, + getClearIntermediatesConfirmation, + getClearIntermediatesState, + getIntermediatesCount, +} from '@workbench/images/intermediates'; +import { useNotify } from '@workbench/useNotify'; +import { useQueueCounts } from '@workbench/widgets/queue/queueDataStore'; +import { getQueueStatus } from '@workbench/widgets/queue/queueServerApi'; +import { useOptionalWorkbenchDispatch } from '@workbench/WorkbenchContext'; +import { Trash2Icon } from 'lucide-react'; +import { useCallback, useEffect, useState } from 'react'; + +interface ClearIntermediatesSnapshot { + error: string | null; + hasGlobalQueueWork: boolean; + intermediatesCount: number | null; + loadState: 'idle' | 'loading' | 'loaded' | 'error'; +} + +const clearIntermediatesStore = createExternalStore({ + error: null, + hasGlobalQueueWork: false, + intermediatesCount: null, + loadState: 'idle', +}); + +let inflight: Promise | null = null; + +const hasQueueWork = ({ queue }: Awaited>): boolean => + queue.pending > 0 || queue.in_progress > 0; + +const refreshClearIntermediatesState = (): Promise => { + if (inflight) { + return inflight; + } + + clearIntermediatesStore.patchSnapshot({ error: null, loadState: 'loading' }); + + inflight = Promise.all([getIntermediatesCount(), getQueueStatus()]) + .then(([count, queueStatus]) => { + clearIntermediatesStore.patchSnapshot({ + error: null, + hasGlobalQueueWork: hasQueueWork(queueStatus), + intermediatesCount: count, + loadState: 'loaded', + }); + }) + .catch((error: unknown) => { + clearIntermediatesStore.patchSnapshot({ + error: getApiErrorMessage(error, 'Could not load intermediate images.'), + loadState: 'error', + }); + }) + .finally(() => { + inflight = null; + }); + + return inflight; +}; + +export const ClearIntermediatesSetting = () => { + const { canClearIntermediates } = useCapabilities(); + const dispatch = useOptionalWorkbenchDispatch(); + const notify = useNotify(); + const queueCounts = useQueueCounts(); + const { error, hasGlobalQueueWork, intermediatesCount, loadState } = clearIntermediatesStore.useSnapshot(); + const [busy, setBusy] = useState(false); + const [isConfirmOpen, setIsConfirmOpen] = useState(false); + + useEffect(() => { + void refreshClearIntermediatesState(); + }, []); + + const state = getClearIntermediatesState({ + canClearIntermediates, + hasActiveQueueWork: hasGlobalQueueWork || queueCounts.pending > 0 || queueCounts.in_progress > 0, + intermediatesCount, + }); + const disabledReason = loadState === 'error' ? error : state.reason; + const label = intermediatesCount === null ? 'Clear intermediates' : `Clear intermediates (${intermediatesCount})`; + const confirmation = getClearIntermediatesConfirmation(intermediatesCount ?? 0); + + const openConfirm = useCallback(() => { + if (!state.disabled) { + setIsConfirmOpen(true); + } + }, [state.disabled]); + + const closeConfirm = useCallback(() => setIsConfirmOpen(false), []); + + const handleClear = useCallback(async () => { + if (state.disabled || busy) { + return; + } + + setBusy(true); + + try { + const queueStatus = await getQueueStatus(); + const queueHasWork = hasQueueWork(queueStatus); + clearIntermediatesStore.patchSnapshot({ hasGlobalQueueWork: queueHasWork }); + + if (queueHasWork) { + notify.info('Intermediates not cleared', 'Wait for pending or running queue work to finish.'); + return; + } + + const clearedCount = await clearIntermediates(); + clearIntermediatesStore.patchSnapshot({ intermediatesCount: 0, loadState: 'loaded' }); + dispatch?.({ type: 'clearCanvasImageReferences' }); + dispatch?.({ type: 'refreshBackendData' }); + notify.info( + 'Intermediates cleared', + `Cleared ${clearedCount} intermediate image${clearedCount === 1 ? '' : 's'}.` + ); + } catch (error) { + notify.error('Failed to clear intermediates', getApiErrorMessage(error, 'Could not clear intermediate images.')); + } finally { + setBusy(false); + } + }, [busy, dispatch, notify, state.disabled]); + + return ( + + + + + Intermediate images + + + Clear temporary workflow outputs and progress artifacts. Final gallery images are not removed. + + + + + {disabledReason ? ( + + {disabledReason} + + ) : null} + + + ); +}; diff --git a/invokeai/frontend/webv2/src/workbench/settings/SettingsDialog.tsx b/invokeai/frontend/webv2/src/workbench/settings/SettingsDialog.tsx index 1f701e26b45..68f5c5d88c2 100644 --- a/invokeai/frontend/webv2/src/workbench/settings/SettingsDialog.tsx +++ b/invokeai/frontend/webv2/src/workbench/settings/SettingsDialog.tsx @@ -54,6 +54,7 @@ import { } from 'lucide-react'; import { useCallback, useEffect, useId, useMemo, useState, type ReactNode } from 'react'; +import { ClearIntermediatesSetting } from './ClearIntermediatesSetting'; import { HotkeysSettingsSection } from './HotkeysSettingsSection'; import { closeWorkbenchSettings, @@ -614,6 +615,7 @@ const QueueSection = () => { + ); }; diff --git a/invokeai/frontend/webv2/src/workbench/workbenchState.test.ts b/invokeai/frontend/webv2/src/workbench/workbenchState.test.ts index 91515cc4073..987b4973aba 100644 --- a/invokeai/frontend/webv2/src/workbench/workbenchState.test.ts +++ b/invokeai/frontend/webv2/src/workbench/workbenchState.test.ts @@ -471,6 +471,38 @@ describe('workbenchReducer Phase 5 generation flow', () => { expect(project.canvas.stagingArea.isVisible).toBe(false); }); + it('clears canvas image references when intermediates are cleared', () => { + let state = submitGenerate(primeGenerate()); + const queueItem = getActiveProject(state).queue.items[0]; + + state = workbenchReducer(state, { + images: [createImage('candidate-1.png', queueItem.id), createImage('candidate-2.png', queueItem.id)], + projectId: getActiveProject(state).id, + queueItemId: queueItem.id, + type: 'routeQueueItemResults', + }); + state = workbenchReducer(state, { type: 'acceptStagedImage' }); + state = workbenchReducer(state, { + images: [createImage('candidate-3.png', queueItem.id)], + projectId: getActiveProject(state).id, + queueItemId: queueItem.id, + type: 'routeQueueItemResults', + }); + + expect(getActiveProject(state).canvas.document.layers).toHaveLength(1); + expect(getActiveProject(state).canvas.stagingArea.pendingImageIds).toEqual(['candidate-3.png']); + expect(getActiveProject(state).undoRedo.past).toHaveLength(1); + + state = workbenchReducer(state, { type: 'clearCanvasImageReferences' }); + const project = getActiveProject(state); + + expect(project.canvas.document.layers).toEqual([]); + expect(project.canvas.stagingArea.pendingImages).toEqual([]); + expect(project.canvas.stagingArea.pendingImageIds).toEqual([]); + expect(project.canvas.stagingArea.isVisible).toBe(false); + expect(project.undoRedo).toEqual({ future: [], past: [] }); + }); + it('cycles staged canvas candidates and accepts the selected candidate placement into a raster layer', () => { let state = submitGenerate(primeGenerate()); const queueItem = getActiveProject(state).queue.items[0]; diff --git a/invokeai/frontend/webv2/src/workbench/workbenchState.ts b/invokeai/frontend/webv2/src/workbench/workbenchState.ts index a7fe5b31c71..5cd49d1eac5 100644 --- a/invokeai/frontend/webv2/src/workbench/workbenchState.ts +++ b/invokeai/frontend/webv2/src/workbench/workbenchState.ts @@ -192,6 +192,7 @@ type WorkbenchAction = | { type: 'touchGalleryImagesRefresh'; projectId?: string } | { type: 'removeGalleryImages'; imageNames: string[]; projectId?: string } | { type: 'setGalleryProjectBoardId'; boardId: string; projectId?: string } + | { type: 'clearCanvasImageReferences' } | { type: 'acceptStagedImage' } | { type: 'clearCanvasStaging' } | { type: 'cancelQueueItem'; queueItemId: string; projectId?: string } @@ -406,6 +407,15 @@ const clearStagingArea = (stagingArea: CanvasStateContract['stagingArea']): Canv sourceQueueItemId: undefined, }); +const clearCanvasImageReferences = (canvas: CanvasStateContract): CanvasStateContract => ({ + ...canvas, + document: { + ...canvas.document, + layers: [], + }, + stagingArea: clearStagingArea(canvas.stagingArea), +}); + const clampStagedImageIndex = (imageIndex: number, pendingImageCount: number): number => { const maxIndex = Math.max(0, pendingImageCount - 1); @@ -2320,6 +2330,16 @@ export const workbenchReducer = (state: WorkbenchState, action: WorkbenchAction) }; }); } + case 'clearCanvasImageReferences': { + return { + ...state, + projects: state.projects.map((project) => ({ + ...project, + canvas: clearCanvasImageReferences(project.canvas), + undoRedo: { future: [], past: [] }, + })), + }; + } case 'acceptStagedImage': { const activeProject = state.projects.find((project) => project.id === state.activeProjectId); const stagedImage = diff --git a/tests/app/test_extract_metadata_from_image.py b/tests/app/test_extract_metadata_from_image.py index 949d6df1656..7a299a57b3c 100644 --- a/tests/app/test_extract_metadata_from_image.py +++ b/tests/app/test_extract_metadata_from_image.py @@ -204,3 +204,20 @@ def test_with_no_metadata(mock_logger): assert result.invokeai_metadata is None assert result.invokeai_workflow is None assert result.invokeai_graph is None + + +def test_empty_string_overrides_do_not_fall_back_to_image_metadata( + mock_logger, valid_metadata, valid_workflow, valid_graph +): + mock_image = MagicMock(spec=Image.Image) + mock_image.info = { + "invokeai_metadata": valid_metadata, + "invokeai_workflow": valid_workflow, + "invokeai_graph": valid_graph, + } + + result = extract_metadata_from_image(mock_image, "", "", "", mock_logger) + + assert result.invokeai_metadata is None + assert result.invokeai_workflow is None + assert result.invokeai_graph is None