feat(ui): fix bugs with image deletion (#3506)

- `imageUsage` object was always stale due to react component lifecycle,
fixed this
- cleaned up the deletion listener and context

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

@@ -4,14 +4,69 @@ import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import { requestedImageDeletion } from 'features/gallery/store/actions';
 import { systemSelector } from 'features/system/store/systemSelectors';
-import { PropsWithChildren, createContext, useCallback, useState } from 'react';
+import {
+  PropsWithChildren,
+  createContext,
+  useCallback,
+  useEffect,
+  useState,
+} from 'react';
 import { ImageDTO } from 'services/api';
+import { RootState } from 'app/store/store';
+import { canvasSelector } from 'features/canvas/store/canvasSelectors';
+import { controlNetSelector } from 'features/controlNet/store/controlNetSlice';
+import { nodesSelecter } from 'features/nodes/store/nodesSlice';
+import { generationSelector } from 'features/parameters/store/generationSelectors';
+import { some } from 'lodash-es';
 
-import { useImageUsage } from 'common/hooks/useImageUsage';
+export type ImageUsage = {
-import { resetCanvas } from 'features/canvas/store/canvasSlice';
+  isInitialImage: boolean;
-import { controlNetReset } from 'features/controlNet/store/controlNetSlice';
+  isCanvasImage: boolean;
-import { clearInitialImage } from 'features/parameters/store/generationSlice';
+  isNodesImage: boolean;
-import { nodeEditorReset } from 'features/nodes/store/nodesSlice';
+  isControlNetImage: boolean;
+};
+
+export const selectImageUsage = createSelector(
+  [
+    generationSelector,
+    canvasSelector,
+    nodesSelecter,
+    controlNetSelector,
+    (state: RootState, image_name?: string) => image_name,
+  ],
+  (generation, canvas, nodes, controlNet, image_name) => {
+    const isInitialImage = generation.initialImage?.image_name === image_name;
+
+    const isCanvasImage = canvas.layerState.objects.some(
+      (obj) => obj.kind === 'image' && obj.image.image_name === image_name
+    );
+
+    const isNodesImage = nodes.nodes.some((node) => {
+      return some(
+        node.data.inputs,
+        (input) =>
+          input.type === 'image' && input.value?.image_name === image_name
+      );
+    });
+
+    const isControlNetImage = some(
+      controlNet.controlNets,
+      (c) =>
+        c.controlImage?.image_name === image_name ||
+        c.processedControlImage?.image_name === image_name
+    );
+
+    const imageUsage: ImageUsage = {
+      isInitialImage,
+      isCanvasImage,
+      isNodesImage,
+      isControlNetImage,
+    };
+
+    return imageUsage;
+  },
+  defaultSelectorOptions
+);
 
 type DeleteImageContextValue = {
   /**
@@ -30,6 +85,10 @@ type DeleteImageContextValue = {
    * The image pending deletion
    */
   image?: ImageDTO;
+  /**
+   * The features in which this image is used
+   */
+  imageUsage?: ImageUsage;
   /**
    * Immediately deletes an image.
    *
@@ -65,41 +124,28 @@ export const DeleteImageContextProvider = (props: Props) => {
   const [imageToDelete, setImageToDelete] = useState<ImageDTO>();
   const dispatch = useAppDispatch();
   const { isOpen, onOpen, onClose } = useDisclosure();
-  const imageUsage = useImageUsage(imageToDelete?.image_name);
 
-  const handleActualDeletion = useCallback(
+  // Check where the image to be deleted is used (eg init image, controlnet, etc.)
-    (image: ImageDTO) => {
+  const imageUsage = useAppSelector((state) =>
-      dispatch(requestedImageDeletion(image));
+    selectImageUsage(state, imageToDelete?.image_name)
-
-      if (imageUsage.isCanvasImage) {
-        dispatch(resetCanvas());
-      }
-
-      if (imageUsage.isControlNetImage) {
-        dispatch(controlNetReset());
-      }
-
-      if (imageUsage.isInitialImage) {
-        dispatch(clearInitialImage());
-      }
-
-      if (imageUsage.isControlNetImage) {
-        dispatch(nodeEditorReset());
-      }
-    },
-    [
-      dispatch,
-      imageUsage.isCanvasImage,
-      imageUsage.isControlNetImage,
-      imageUsage.isInitialImage,
-    ]
   );
 
+  // Clean up after deleting or dismissing the modal
   const closeAndClearImageToDelete = useCallback(() => {
     setImageToDelete(undefined);
     onClose();
   }, [onClose]);
 
+  // Dispatch the actual deletion action, to be handled by listener middleware
+  const handleActualDeletion = useCallback(
+    (image: ImageDTO) => {
+      dispatch(requestedImageDeletion({ image, imageUsage }));
+      closeAndClearImageToDelete();
+    },
+    [closeAndClearImageToDelete, dispatch, imageUsage]
+  );
+
+  // This is intended to be called by the delete button in the dialog
   const onImmediatelyDelete = useCallback(() => {
     if (canDeleteImage && imageToDelete) {
       handleActualDeletion(imageToDelete);
@@ -114,25 +160,31 @@ export const DeleteImageContextProvider = (props: Props) => {
 
   const handleGatedDeletion = useCallback(
     (image: ImageDTO) => {
-      if (shouldConfirmOnDelete || imageUsage) {
+      if (shouldConfirmOnDelete || some(imageUsage)) {
+        // If we should confirm on delete, or if the image is in use, open the dialog
         onOpen();
       } else {
         handleActualDeletion(image);
       }
     },
-    [shouldConfirmOnDelete, imageUsage, onOpen, handleActualDeletion]
+    [imageUsage, shouldConfirmOnDelete, onOpen, handleActualDeletion]
   );
 
-  const onDelete = useCallback(
+  // Consumers of the context call this to delete an image
-    (image?: ImageDTO) => {
+  const onDelete = useCallback((image?: ImageDTO) => {
-      if (!image) {
+    if (!image) {
-        return;
+      return;
-      }
+    }
-      setImageToDelete(image);
+    // Set the image to delete, then let the effect call the actual deletion
-      handleGatedDeletion(image);
+    setImageToDelete(image);
-    },
+  }, []);
-    [handleGatedDeletion]
+
-  );
+  useEffect(() => {
+    // We need to use an effect here to trigger the image usage selector, else we get a stale value
+    if (imageToDelete) {
+      handleGatedDeletion(imageToDelete);
+    }
+  }, [handleGatedDeletion, imageToDelete]);
 
   return (
     <DeleteImageContext.Provider
@@ -142,6 +194,7 @@ export const DeleteImageContextProvider = (props: Props) => {
         onClose: closeAndClearImageToDelete,
         onDelete,
         onImmediatelyDelete,
+        imageUsage,
       }}
     >
       {props.children}

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

@@ -9,6 +9,10 @@ import {
   selectImagesEntities,
   selectImagesIds,
 } 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';
 
 const moduleLog = log.child({ namespace: 'addRequestedImageDeletionListener' });
 
@@ -19,11 +23,7 @@ export const addRequestedImageDeletionListener = () => {
   startAppListening({
     actionCreator: requestedImageDeletion,
     effect: (action, { dispatch, getState }) => {
-      const image = action.payload;
+      const { image, imageUsage } = action.payload;
-      if (!image) {
-        moduleLog.warn('No image provided');
-        return;
-      }
 
       const { image_name, image_origin } = image;
 
@@ -57,6 +57,24 @@ export const addRequestedImageDeletionListener = () => {
         }
       }
 
+      // We need to reset the features where the image is in use - none of these work if their image(s) don't exist
+
+      if (imageUsage.isCanvasImage) {
+        dispatch(resetCanvas());
+      }
+
+      if (imageUsage.isControlNetImage) {
+        dispatch(controlNetReset());
+      }
+
+      if (imageUsage.isInitialImage) {
+        dispatch(clearInitialImage());
+      }
+
+      if (imageUsage.isNodesImage) {
+        dispatch(nodeEditorReset());
+      }
+
       // Preemptively remove from gallery
       dispatch(imageRemoved(image_name));
 

invokeai/frontend/web/src/common/hooks/useImageUsage.ts

@@ -1,66 +0,0 @@
-import { createSelector } from '@reduxjs/toolkit';
-import { RootState } from 'app/store/store';
-import { useAppSelector } from 'app/store/storeHooks';
-import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
-import { canvasSelector } from 'features/canvas/store/canvasSelectors';
-import { controlNetSelector } from 'features/controlNet/store/controlNetSlice';
-import { nodesSelecter } from 'features/nodes/store/nodesSlice';
-import { generationSelector } from 'features/parameters/store/generationSelectors';
-import { some } from 'lodash-es';
-
-export type ImageUsage = {
-  isInitialImage: boolean;
-  isCanvasImage: boolean;
-  isNodesImage: boolean;
-  isControlNetImage: boolean;
-};
-
-const selectImageUsage = createSelector(
-  [
-    generationSelector,
-    canvasSelector,
-    nodesSelecter,
-    controlNetSelector,
-    (state: RootState, image_name?: string) => image_name,
-  ],
-  (generation, canvas, nodes, controlNet, image_name) => {
-    const isInitialImage = generation.initialImage?.image_name === image_name;
-
-    const isCanvasImage = canvas.layerState.objects.some(
-      (obj) => obj.kind === 'image' && obj.image.image_name === image_name
-    );
-
-    const isNodesImage = nodes.nodes.some((node) => {
-      return some(
-        node.data.inputs,
-        (input) =>
-          input.type === 'image' && input.value?.image_name === image_name
-      );
-    });
-
-    const isControlNetImage = some(
-      controlNet.controlNets,
-      (c) =>
-        c.controlImage?.image_name === image_name ||
-        c.processedControlImage?.image_name === image_name
-    );
-
-    const imageUsage: ImageUsage = {
-      isInitialImage,
-      isCanvasImage,
-      isNodesImage,
-      isControlNetImage,
-    };
-
-    return imageUsage;
-  },
-  defaultSelectorOptions
-);
-
-export const useImageUsage = (image_name?: string) => {
-  const imageUsage = useAppSelector((state) =>
-    selectImageUsage(state, image_name)
-  );
-
-  return imageUsage;
-};

invokeai/frontend/web/src/features/gallery/components/DeleteImageModal.tsx

@@ -12,13 +12,15 @@ import {
   UnorderedList,
 } from '@chakra-ui/react';
 import { createSelector } from '@reduxjs/toolkit';
-import { DeleteImageContext } from 'app/contexts/DeleteImageContext';
+import {
+  DeleteImageContext,
+  ImageUsage,
+} from 'app/contexts/DeleteImageContext';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import IAIButton from 'common/components/IAIButton';
 import IAIIconButton from 'common/components/IAIIconButton';
 import IAISwitch from 'common/components/IAISwitch';
-import { ImageUsage, useImageUsage } from 'common/hooks/useImageUsage';
 import { configSelector } from 'features/system/store/configSelectors';
 import { systemSelector } from 'features/system/store/systemSelectors';
 import { setShouldConfirmOnDelete } from 'features/system/store/systemSlice';
@@ -42,9 +44,13 @@ const selector = createSelector(
   defaultSelectorOptions
 );
 
-const ImageInUseMessage = (props: { imageUsage: ImageUsage }) => {
+const ImageInUseMessage = (props: { imageUsage?: ImageUsage }) => {
   const { imageUsage } = props;
 
+  if (!imageUsage) {
+    return null;
+  }
+
   if (!some(imageUsage)) {
     return null;
   }
@@ -69,14 +75,12 @@ const DeleteImageModal = () => {
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
 
-  const { isOpen, onClose, onImmediatelyDelete, image } =
+  const { isOpen, onClose, onImmediatelyDelete, image, imageUsage } =
     useContext(DeleteImageContext);
 
   const { shouldConfirmOnDelete, canRestoreDeletedImagesFromBin } =
     useAppSelector(selector);
 
-  const imageUsage = useImageUsage(image?.image_name);
-
   const handleChangeShouldConfirmOnDelete = useCallback(
     (e: ChangeEvent<HTMLInputElement>) =>
       dispatch(setShouldConfirmOnDelete(!e.target.checked)),

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

@@ -1,10 +1,15 @@
 import { createAction } from '@reduxjs/toolkit';
-import { ImageNameAndOrigin } from 'features/parameters/store/actions';
+import { ImageUsage } from 'app/contexts/DeleteImageContext';
 import { ImageDTO } from 'services/api';
 
-export const requestedImageDeletion = createAction<
+export type RequestedImageDeletionArg = {
-  ImageDTO | ImageNameAndOrigin | undefined
+  image: ImageDTO;
->('gallery/requestedImageDeletion');
+  imageUsage: ImageUsage;
+};
+
+export const requestedImageDeletion = createAction<RequestedImageDeletionArg>(
+  'gallery/requestedImageDeletion'
+);
 
 export const sentImageToCanvas = createAction('gallery/sentImageToCanvas');