add image usage for board images and listener to handle actual deletion

Author: Mary Hipp

invokeai/frontend/web/src/app/contexts/DeleteBoardImagesContext.tsx

@@ -2,13 +2,76 @@ import { useDisclosure } from '@chakra-ui/react';
 import { PropsWithChildren, createContext, useCallback, useState } from 'react';
 import { BoardDTO } from 'services/api/types';
 import { useDeleteBoardMutation } from '../../services/api/endpoints/boards';
+import { defaultSelectorOptions } from '../store/util/defaultMemoizeOptions';
+import { createSelector } from '@reduxjs/toolkit';
+import { some } from 'lodash-es';
+import { canvasSelector } from '../../features/canvas/store/canvasSelectors';
+import { controlNetSelector } from '../../features/controlNet/store/controlNetSlice';
+import { selectImagesById } from '../../features/gallery/store/imagesSlice';
+import { nodesSelector } from '../../features/nodes/store/nodesSlice';
+import { generationSelector } from '../../features/parameters/store/generationSelectors';
+import { RootState } from '../store/store';
+import { useAppDispatch, useAppSelector } from '../store/storeHooks';
+import { ImageUsage } from './DeleteImageContext';
+import { requestedBoardImagesDeletion } from '../../features/gallery/store/actions';
 
-export type ImageUsage = {
-  isInitialImage: boolean;
-  isCanvasImage: boolean;
-  isNodesImage: boolean;
-  isControlNetImage: boolean;
-};
+export const selectBoardImagesUsage = createSelector(
+  [
+    (state: RootState) => state,
+    generationSelector,
+    canvasSelector,
+    nodesSelector,
+    controlNetSelector,
+    (state: RootState, board_id?: string) => board_id,
+  ],
+  (state, generation, canvas, nodes, controlNet, board_id) => {
+    const initialImage = generation.initialImage
+      ? selectImagesById(state, generation.initialImage.imageName)
+      : undefined;
+    const isInitialImage = initialImage?.board_id === board_id;
+
+    const isCanvasImage = canvas.layerState.objects.some((obj) => {
+      if (obj.kind === 'image') {
+        const image = selectImagesById(state, obj.imageName);
+        return image?.board_id === board_id;
+      }
+      return false;
+    });
+
+    const isNodesImage = nodes.nodes.some((node) => {
+      return some(node.data.inputs, (input) => {
+        if (input.type === 'image' && input.value) {
+          const image = selectImagesById(state, input.value.image_name);
+          return image?.board_id === board_id;
+        }
+        return false;
+      });
+    });
+
+    const isControlNetImage = some(controlNet.controlNets, (c) => {
+      const controlImage = c.controlImage
+        ? selectImagesById(state, c.controlImage)
+        : undefined;
+      const processedControlImage = c.processedControlImage
+        ? selectImagesById(state, c.processedControlImage)
+        : undefined;
+      return (
+        controlImage?.board_id === board_id ||
+        processedControlImage?.board_id === board_id
+      );
+    });
+
+    const imageUsage: ImageUsage = {
+      isInitialImage,
+      isCanvasImage,
+      isNodesImage,
+      isControlNetImage,
+    };
+
+    return imageUsage;
+  },
+  defaultSelectorOptions
+);
 
 type DeleteBoardImagesContextValue = {
   /**
@@ -19,9 +82,7 @@ type DeleteBoardImagesContextValue = {
    * Closes the move image dialog.
    */
   onClose: () => void;
-  /**
-   * The image pending movement
-   */
+  imagesUsage?: ImageUsage;
   board?: BoardDTO;
   onClickDeleteBoardImages: (board: BoardDTO) => void;
   handleDeleteBoardImages: (boardId: string) => void;
@@ -42,8 +103,13 @@ type Props = PropsWithChildren;
 export const DeleteBoardImagesContextProvider = (props: Props) => {
   const [boardToDelete, setBoardToDelete] = useState<BoardDTO>();
   const { isOpen, onOpen, onClose } = useDisclosure();
+  const dispatch = useAppDispatch();
+
+  // Check where the board images to be deleted are used (eg init image, controlnet, etc.)
+  const imagesUsage = useAppSelector((state) =>
+    selectBoardImagesUsage(state, boardToDelete?.board_id)
+  );
 
-  const [deleteBoardAndImages] = useDeleteBoardAndImagesMutation();
   const [deleteBoard] = useDeleteBoardMutation();
 
   // Clean up after deleting or dismissing the modal
@@ -67,11 +133,13 @@ export const DeleteBoardImagesContextProvider = (props: Props) => {
   const handleDeleteBoardImages = useCallback(
     (boardId: string) => {
       if (boardToDelete) {
-        deleteBoardAndImages(boardId);
+        dispatch(
+          requestedBoardImagesDeletion({ board: boardToDelete, imagesUsage })
+        );
         closeAndClearBoardToDelete();
       }
     },
-    [deleteBoardAndImages, closeAndClearBoardToDelete, boardToDelete]
+    [dispatch, closeAndClearBoardToDelete, boardToDelete, imagesUsage]
   );
 
   const handleDeleteBoardOnly = useCallback(
@@ -93,6 +161,7 @@ export const DeleteBoardImagesContextProvider = (props: Props) => {
         onClickDeleteBoardImages,
         handleDeleteBoardImages,
         handleDeleteBoardOnly,
+        imagesUsage,
       }}
     >
       {props.children}

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts

@@ -83,6 +83,7 @@ import {
   addImageRemovedFromBoardRejectedListener,
 } from './listeners/imageRemovedFromBoard';
 import { addReceivedOpenAPISchemaListener } from './listeners/receivedOpenAPISchema';
+import { addRequestedBoardImageDeletionListener } from './listeners/boardImagesDeleted';
 
 export const listenerMiddleware = createListenerMiddleware();
 
@@ -124,6 +125,7 @@ addRequestedImageDeletionListener();
 addImageDeletedPendingListener();
 addImageDeletedFulfilledListener();
 addImageDeletedRejectedListener();
+addRequestedBoardImageDeletionListener();
 
 // Image metadata
 addImageMetadataReceivedFulfilledListener();

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardImagesDeleted.ts

@@ -0,0 +1,79 @@
+import { requestedBoardImagesDeletion } from 'features/gallery/store/actions';
+import { startAppListening } from '..';
+import { imageSelected } from 'features/gallery/store/gallerySlice';
+import {
+  imagesRemoved,
+  selectImagesAll,
+  selectImagesById,
+} from 'features/gallery/store/imagesSlice';
+import { resetCanvas } from 'features/canvas/store/canvasSlice';
+import { controlNetReset } from 'features/controlNet/store/controlNetSlice';
+import { clearInitialImage } from 'features/parameters/store/generationSlice';
+import { nodeEditorReset } from 'features/nodes/store/nodesSlice';
+import { LIST_TAG, api } from 'services/api';
+import { boardsApi } from '../../../../../services/api/endpoints/boards';
+
+export const addRequestedBoardImageDeletionListener = () => {
+  startAppListening({
+    actionCreator: requestedBoardImagesDeletion,
+    effect: async (action, { dispatch, getState, condition }) => {
+      const { board, imagesUsage } = action.payload;
+
+      const { board_id } = board;
+
+      const state = getState();
+      const selectedImage = state.gallery.selectedImage
+        ? selectImagesById(state, state.gallery.selectedImage)
+        : undefined;
+
+      if (selectedImage && selectedImage.board_id === board_id) {
+        dispatch(imageSelected());
+      }
+
+      // We need to reset the features where the board images are in use - none of these work if their image(s) don't exist
+
+      if (imagesUsage.isCanvasImage) {
+        dispatch(resetCanvas());
+      }
+
+      if (imagesUsage.isControlNetImage) {
+        dispatch(controlNetReset());
+      }
+
+      if (imagesUsage.isInitialImage) {
+        dispatch(clearInitialImage());
+      }
+
+      if (imagesUsage.isNodesImage) {
+        dispatch(nodeEditorReset());
+      }
+
+      // Preemptively remove from gallery
+      const images = selectImagesAll(state).reduce((acc: string[], img) => {
+        if (img.board_id === board_id) {
+          acc.push(img.image_name);
+        }
+        return acc;
+      }, []);
+      dispatch(imagesRemoved(images));
+
+      // Delete from server
+      dispatch(boardsApi.endpoints.deleteBoardAndImages.initiate(board_id));
+      const result =
+        boardsApi.endpoints.deleteBoardAndImages.select(board_id)(state);
+      const { isSuccess } = result;
+
+      // Wait for successful deletion, then trigger boards to re-fetch
+      const wasBoardDeleted = await condition(() => !!isSuccess, 30000);
+
+      if (wasBoardDeleted) {
+        dispatch(
+          api.util.invalidateTags([
+            { type: 'Board', id: board_id },
+            { type: 'Image', id: LIST_TAG },
+          ])
+        );
+      }
+    },
+  });
+};

invokeai/frontend/web/src/features/gallery/components/Boards/DeleteBoardImagesModal.tsx

@@ -7,12 +7,46 @@ import {
   AlertDialogOverlay,
   Divider,
   Flex,
+  ListItem,
   Text,
+  UnorderedList,
 } from '@chakra-ui/react';
 import IAIButton from 'common/components/IAIButton';
 import { memo, useContext, useRef } from 'react';
 import { useTranslation } from 'react-i18next';
 import { DeleteBoardImagesContext } from '../../../../app/contexts/DeleteBoardImagesContext';
+import { some } from 'lodash-es';
+import { ImageUsage } from '../../../../app/contexts/DeleteImageContext';
+
+const BoardImageInUseMessage = (props: { imagesUsage?: ImageUsage }) => {
+  const { imagesUsage } = props;
+
+  if (!imagesUsage) {
+    return null;
+  }
+
+  if (!some(imagesUsage)) {
+    return null;
+  }
+
+  return (
+    <>
+      <Text>
+        An image from this board is currently in use in the following features:
+      </Text>
+      <UnorderedList sx={{ paddingInlineStart: 6 }}>
+        {imagesUsage.isInitialImage && <ListItem>Image to Image</ListItem>}
+        {imagesUsage.isCanvasImage && <ListItem>Unified Canvas</ListItem>}
+        {imagesUsage.isControlNetImage && <ListItem>ControlNet</ListItem>}
+        {imagesUsage.isNodesImage && <ListItem>Node Editor</ListItem>}
+      </UnorderedList>
+      <Text>
+        If you delete images from this board, those features will immediately be
+        reset.
+      </Text>
+    </>
+  );
+};
 
 const DeleteBoardImagesModal = () => {
   const { t } = useTranslation();
@@ -23,6 +57,7 @@ const DeleteBoardImagesModal = () => {
     board,
     handleDeleteBoardImages,
     handleDeleteBoardOnly,
+    imagesUsage,
   } = useContext(DeleteBoardImagesContext);
 
   const cancelRef = useRef<HTMLButtonElement>(null);
@@ -43,6 +78,7 @@ const DeleteBoardImagesModal = () => {
 
             <AlertDialogBody>
               <Flex direction="column" gap={3}>
+                <BoardImageInUseMessage imagesUsage={imagesUsage} />
                 <Divider />
                 <Text>{t('common.areYouSure')}</Text>
                 <Text fontWeight="bold">

invokeai/frontend/web/src/features/gallery/store/actions.ts

@@ -1,6 +1,6 @@
 import { createAction } from '@reduxjs/toolkit';
 import { ImageUsage } from 'app/contexts/DeleteImageContext';
-import { ImageDTO } from 'services/api/types';
+import { ImageDTO, BoardDTO } from 'services/api/types';
 
 export type RequestedImageDeletionArg = {
   image: ImageDTO;
@@ -11,6 +11,16 @@ export const requestedImageDeletion = createAction<RequestedImageDeletionArg>(
   'gallery/requestedImageDeletion'
 );
 
+export type RequestedBoardImagesDeletionArg = {
+  board: BoardDTO;
+  imagesUsage: ImageUsage;
+};
+
+export const requestedBoardImagesDeletion =
+  createAction<RequestedBoardImagesDeletionArg>(
+    'gallery/requestedBoardImagesDeletion'
+  );
+
 export const sentImageToCanvas = createAction('gallery/sentImageToCanvas');
 
 export const sentImageToImg2Img = createAction('gallery/sentImageToImg2Img');

invokeai/frontend/web/src/features/gallery/store/imagesSlice.ts

@@ -60,6 +60,9 @@ const imagesSlice = createSlice({
     imageRemoved: (state, action: PayloadAction<string>) => {
       imagesAdapter.removeOne(state, action.payload);
     },
+    imagesRemoved: (state, action: PayloadAction<string[]>) => {
+      imagesAdapter.removeMany(state, action.payload);
+    },
     imageCategoriesChanged: (state, action: PayloadAction<ImageCategory[]>) => {
       state.categories = action.payload;
     },
@@ -117,6 +120,7 @@ export const {
   imageUpserted,
   imageUpdatedOne,
   imageRemoved,
+  imagesRemoved,
   imageCategoriesChanged,
 } = imagesSlice.actions;