feat(ui): clear features if image used by them is deleted

This handles the case when an image is deleted but is still in use in as eg an init image on canvas, or a control image. If we just delete the image, canvas/controlnet/etc may break (the image would just fail to load).

When an image is deleted, the app checks to see if it is in use in:
- Image to Image
- ControlNet
- Unified Canvas
- Node Editor

The delete dialog will always open if the image is in use anywhere, and the user is advised that deleting the image will reset the feature(s).

Even if the user has ticked the box to not confirm on delete, the dialog will still show if the image is in use somewhere.

invokeai/frontend/web/src/app/components/App.tsx

@@ -21,7 +21,7 @@ import { ReactNode, memo, useCallback, useEffect, useState } from 'react';
 import { APP_HEIGHT, APP_WIDTH } from 'theme/util/constants';
 import GlobalHotkeys from './GlobalHotkeys';
 import Toaster from './Toaster';
-import DeleteModal from 'features/gallery/components/DeleteModal';
+import DeleteImageModal from 'features/gallery/components/DeleteImageModal';
 
 const DEFAULT_CONFIG = {};
 
@@ -134,7 +134,7 @@ const App = ({
           <FloatingGalleryButton />
         </Portal>
       </Grid>
-      <DeleteModal />
+      <DeleteImageModal />
       <Toaster />
       <GlobalHotkeys />
     </>

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

@@ -7,6 +7,12 @@ import { systemSelector } from 'features/system/store/systemSelectors';
 import { PropsWithChildren, createContext, useCallback, useState } from 'react';
 import { ImageDTO } from 'services/api';
 
+import { useImageUsage } from 'common/hooks/useImageUsage';
+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';
+
 type DeleteImageContextValue = {
   /**
    * Whether the delete image dialog is open.
@@ -16,16 +22,20 @@ type DeleteImageContextValue = {
    * Closes the delete image dialog.
    */
   onClose: () => void;
+  /**
+   * Opens the delete image dialog and handles all deletion-related checks.
+   */
+  onDelete: (image?: ImageDTO) => void;
+  /**
+   * The image pending deletion
+   */
+  image?: ImageDTO;
   /**
    * Immediately deletes an image.
    *
    * You probably don't want to use this - use `onDelete` instead.
    */
   onImmediatelyDelete: () => void;
-  /**
-   * Opens the delete image dialog and handles all deletion-related checks.
-   */
-  onDelete: (image?: ImageDTO) => void;
 };
 
 export const DeleteImageContext = createContext<DeleteImageContextValue>({
@@ -43,8 +53,6 @@ const selector = createSelector(
     return {
       canDeleteImage: isConnected && !isProcessing,
       shouldConfirmOnDelete,
-      isProcessing,
-      isConnected,
     };
   },
   defaultSelectorOptions
@@ -57,6 +65,35 @@ 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(
+    (image: ImageDTO) => {
+      dispatch(requestedImageDeletion(image));
+
+      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,
+    ]
+  );
 
   const closeAndClearImageToDelete = useCallback(() => {
     setImageToDelete(undefined);
@@ -65,20 +102,25 @@ export const DeleteImageContextProvider = (props: Props) => {
 
   const onImmediatelyDelete = useCallback(() => {
     if (canDeleteImage && imageToDelete) {
-      dispatch(requestedImageDeletion(imageToDelete));
+      handleActualDeletion(imageToDelete);
     }
     closeAndClearImageToDelete();
-  }, [canDeleteImage, imageToDelete, closeAndClearImageToDelete, dispatch]);
+  }, [
+    canDeleteImage,
+    imageToDelete,
+    closeAndClearImageToDelete,
+    handleActualDeletion,
+  ]);
 
-  const handleDelete = useCallback(
+  const handleGatedDeletion = useCallback(
     (image: ImageDTO) => {
-      if (shouldConfirmOnDelete) {
+      if (shouldConfirmOnDelete || imageUsage) {
         onOpen();
       } else {
-        dispatch(requestedImageDeletion(image));
+        handleActualDeletion(image);
       }
     },
-    [shouldConfirmOnDelete, onOpen, dispatch]
+    [shouldConfirmOnDelete, imageUsage, onOpen, handleActualDeletion]
   );
 
   const onDelete = useCallback(
@@ -87,15 +129,16 @@ export const DeleteImageContextProvider = (props: Props) => {
         return;
       }
       setImageToDelete(image);
-      handleDelete(image);
+      handleGatedDeletion(image);
     },
-    [handleDelete]
+    [handleGatedDeletion]
   );
 
   return (
     <DeleteImageContext.Provider
       value={{
         isOpen,
+        image: imageToDelete,
         onClose: closeAndClearImageToDelete,
         onDelete,
         onImmediatelyDelete,

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

@@ -1,4 +1,5 @@
 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';
@@ -7,13 +8,20 @@ import { nodesSelecter } from 'features/nodes/store/nodesSlice';
 import { generationSelector } from 'features/parameters/store/generationSelectors';
 import { some } from 'lodash-es';
 
-const selectIsImageInUse = createSelector(
+export type ImageUsage = {
+  isInitialImage: boolean;
+  isCanvasImage: boolean;
+  isNodesImage: boolean;
+  isControlNetImage: boolean;
+};
+
+const selectImageUsage = createSelector(
   [
     generationSelector,
     canvasSelector,
     nodesSelecter,
     controlNetSelector,
-    (state, image_name) => image_name,
+    (state: RootState, image_name?: string) => image_name,
   ],
   (generation, canvas, nodes, controlNet, image_name) => {
     const isInitialImage = generation.initialImage?.image_name === image_name;
@@ -37,18 +45,22 @@ const selectIsImageInUse = createSelector(
         c.processedControlImage?.image_name === image_name
     );
 
-    return {
+    const imageUsage: ImageUsage = {
       isInitialImage,
       isCanvasImage,
       isNodesImage,
       isControlNetImage,
     };
+
+    return imageUsage;
   },
   defaultSelectorOptions
 );
 
-export const useGetIsImageInUse = (image_name?: string) => {
-  const a = useAppSelector((state) => selectIsImageInUse(state, image_name));
+export const useImageUsage = (image_name?: string) => {
+  const imageUsage = useAppSelector((state) =>
+    selectImageUsage(state, image_name)
+  );
 
-  return a;
+  return imageUsage;
 };

invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts

@@ -187,6 +187,9 @@ export const controlNetSlice = createSlice({
         processorType
       ].default as RequiredControlNetProcessorNode;
     },
+    controlNetReset: () => {
+      return { ...initialControlNetState };
+    },
   },
   extraReducers: (builder) => {
     builder.addCase(controlNetImageProcessed, (state, action) => {
@@ -243,6 +246,7 @@ export const {
   controlNetEndStepPctChanged,
   controlNetProcessorParamsChanged,
   controlNetProcessorTypeChanged,
+  controlNetReset,
 } = controlNetSlice.actions;
 
 export default controlNetSlice.reducer;

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

@@ -50,7 +50,7 @@ import UpscaleSettings from 'features/parameters/components/Parameters/Upscale/U
 import { useAppToaster } from 'app/components/Toaster';
 import { setInitialCanvasImage } from 'features/canvas/store/canvasSlice';
 import { DeleteImageContext } from 'app/contexts/DeleteImageContext';
-import { DeleteImageButton } from './DeleteModal';
+import { DeleteImageButton } from './DeleteImageModal';
 
 const currentImageButtonsSelector = createSelector(
   [

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

@@ -15,7 +15,6 @@ import { imageSelected } from '../store/gallerySlice';
 import IAIDndImage from 'common/components/IAIDndImage';
 import { ImageDTO } from 'services/api';
 import { IAIImageFallback } from 'common/components/IAIImageFallback';
-import { useGetIsImageInUse } from 'common/hooks/useGetIsImageInUse';
 
 export const imagesSelector = createSelector(
   [uiSelector, gallerySelector, systemSelector],
@@ -55,8 +54,6 @@ const CurrentImagePreview = () => {
   const toaster = useAppToaster();
   const dispatch = useAppDispatch();
 
-  const isImageInUse = useGetIsImageInUse(image?.image_name);
-  console.log(isImageInUse);
   const handleError = useCallback(() => {
     dispatch(imageSelected());
     if (shouldFetchImages) {

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

@@ -5,19 +5,24 @@ import {
   AlertDialogFooter,
   AlertDialogHeader,
   AlertDialogOverlay,
+  Divider,
   Flex,
+  ListItem,
   Text,
+  UnorderedList,
 } from '@chakra-ui/react';
 import { createSelector } from '@reduxjs/toolkit';
 import { DeleteImageContext } 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';
-import { isEqual } from 'lodash-es';
+import { some } from 'lodash-es';
 
 import { ChangeEvent, memo, useCallback, useContext, useRef } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -28,31 +33,56 @@ const selector = createSelector(
   (system, config) => {
     const { shouldConfirmOnDelete } = system;
     const { canRestoreDeletedImagesFromBin } = config;
-    return { shouldConfirmOnDelete, canRestoreDeletedImagesFromBin };
+
+    return {
+      shouldConfirmOnDelete,
+      canRestoreDeletedImagesFromBin,
+    };
   },
-  {
-    memoizeOptions: {
-      resultEqualityCheck: isEqual,
-    },
-  }
+  defaultSelectorOptions
 );
 
+const ImageInUseMessage = (props: { imageUsage: ImageUsage }) => {
+  const { imageUsage } = props;
+
+  if (!some(imageUsage)) {
+    return null;
+  }
+
+  return (
+    <>
+      <Text>This image is currently in use in the following features:</Text>
+      <UnorderedList sx={{ paddingInlineStart: 6 }}>
+        {imageUsage.isInitialImage && <ListItem>Image to Image</ListItem>}
+        {imageUsage.isCanvasImage && <ListItem>Unified Canvas</ListItem>}
+        {imageUsage.isControlNetImage && <ListItem>ControlNet</ListItem>}
+        {imageUsage.isNodesImage && <ListItem>Node Editor</ListItem>}
+      </UnorderedList>
+      <Text>
+        If you delete this image, those features will immediately be reset.
+      </Text>
+    </>
+  );
+};
+
 const DeleteImageModal = () => {
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
 
+  const { isOpen, onClose, onImmediatelyDelete, image } =
+    useContext(DeleteImageContext);
+
   const { shouldConfirmOnDelete, canRestoreDeletedImagesFromBin } =
     useAppSelector(selector);
 
+  const imageUsage = useImageUsage(image?.image_name);
+
   const handleChangeShouldConfirmOnDelete = useCallback(
     (e: ChangeEvent<HTMLInputElement>) =>
       dispatch(setShouldConfirmOnDelete(!e.target.checked)),
     [dispatch]
   );
 
-  const { isOpen, onClose, onImmediatelyDelete } =
-    useContext(DeleteImageContext);
-
   const cancelRef = useRef<HTMLButtonElement>(null);
 
   return (
@@ -69,15 +99,15 @@ const DeleteImageModal = () => {
           </AlertDialogHeader>
 
           <AlertDialogBody>
-            <Flex direction="column" gap={5}>
-              <Flex direction="column" gap={2}>
-                <Text>{t('common.areYouSure')}</Text>
-                <Text>
-                  {canRestoreDeletedImagesFromBin
-                    ? t('gallery.deleteImageBin')
-                    : t('gallery.deleteImagePermanent')}
-                </Text>
-              </Flex>
+            <Flex direction="column" gap={3}>
+              <ImageInUseMessage imageUsage={imageUsage} />
+              <Divider />
+              <Text>
+                {canRestoreDeletedImagesFromBin
+                  ? t('gallery.deleteImageBin')
+                  : t('gallery.deleteImagePermanent')}
+              </Text>
+              <Text>{t('common.areYouSure')}</Text>
               <IAISwitch
                 label={t('common.dontAskMeAgain')}
                 isChecked={!shouldConfirmOnDelete}

invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts

@@ -93,6 +93,9 @@ const nodesSlice = createSlice({
         console.error(err);
       }
     },
+    nodeEditorReset: () => {
+      return { ...initialNodesState };
+    },
   },
   extraReducers(builder) {
     builder.addCase(receivedOpenAPISchema.fulfilled, (state, action) => {
@@ -127,6 +130,7 @@ export const {
   connectionEnded,
   shouldShowGraphOverlayChanged,
   parsedOpenAPISchema,
+  nodeEditorReset,
 } = nodesSlice.actions;
 
 export default nodesSlice.reducer;

invokeai/frontend/web/src/features/parameters/components/Parameters/ImageToImage/InitialImagePreview.tsx

@@ -14,7 +14,6 @@ import { useAppToaster } from 'app/components/Toaster';
 import IAIDndImage from 'common/components/IAIDndImage';
 import { ImageDTO } from 'services/api';
 import { IAIImageFallback } from 'common/components/IAIImageFallback';
-import { useGetIsImageInUse } from 'common/hooks/useGetIsImageInUse';
 
 const selector = createSelector(
   [generationSelector],