fix(ui): handle deleting images in use in generation tab

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

@@ -1,6 +1,7 @@
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { resetCanvas } from 'features/canvas/store/canvasSlice';
 import { controlAdaptersReset } from 'features/controlAdapters/store/controlAdaptersSlice';
+import { allLayersDeleted } from 'features/controlLayers/store/controlLayersSlice';
 import { getImageUsage } from 'features/deleteImageModal/store/selectors';
 import { nodeEditorReset } from 'features/nodes/store/nodesSlice';
 import { imagesApi } from 'services/api/endpoints/images';
@@ -16,10 +17,11 @@ export const addDeleteBoardAndImagesFulfilledListener = (startAppListening: AppS
       let wasCanvasReset = false;
       let wasNodeEditorReset = false;
       let wereControlAdaptersReset = false;
+      let wereControlLayersReset = false;
 
-      const { generation, canvas, nodes, controlAdapters } = getState();
+      const { canvas, nodes, controlAdapters, controlLayers } = getState();
       deleted_images.forEach((image_name) => {
-        const imageUsage = getImageUsage(generation, canvas, nodes, controlAdapters, image_name);
+        const imageUsage = getImageUsage(canvas, nodes, controlAdapters, controlLayers.present, image_name);
 
         if (imageUsage.isCanvasImage && !wasCanvasReset) {
           dispatch(resetCanvas());
@@ -35,6 +37,11 @@ export const addDeleteBoardAndImagesFulfilledListener = (startAppListening: AppS
           dispatch(controlAdaptersReset());
           wereControlAdaptersReset = true;
         }
+
+        if (imageUsage.isControlLayerImage && !wereControlLayersReset) {
+          dispatch(allLayersDeleted());
+          wereControlLayersReset = true;
+        }
       });
     },
   });

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

@@ -1,5 +1,6 @@
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
+import type { AppDispatch, RootState } from 'app/store/store';
 import { resetCanvas } from 'features/canvas/store/canvasSlice';
 import {
   controlAdapterImageChanged,
@@ -7,6 +8,13 @@ import {
   selectControlAdapterAll,
 } from 'features/controlAdapters/store/controlAdaptersSlice';
 import { isControlNetOrT2IAdapter } from 'features/controlAdapters/store/types';
+import {
+  isControlAdapterLayer,
+  isInitialImageLayer,
+  isIPAdapterLayer,
+  isRegionalGuidanceLayer,
+  layerDeleted,
+} from 'features/controlLayers/store/controlLayersSlice';
 import { imageDeletionConfirmed } from 'features/deleteImageModal/store/actions';
 import { isModalOpenChanged } from 'features/deleteImageModal/store/slice';
 import { selectListImagesQueryArgs } from 'features/gallery/store/gallerySelectors';
@@ -17,8 +25,79 @@ import { isInvocationNode } from 'features/nodes/types/invocation';
 import { clamp, forEach } from 'lodash-es';
 import { api } from 'services/api';
 import { imagesApi } from 'services/api/endpoints/images';
+import type { ImageDTO } from 'services/api/types';
 import { imagesSelectors } from 'services/api/util';
 
+const deleteNodesImages = (state: RootState, dispatch: AppDispatch, imageDTO: ImageDTO) => {
+  state.nodes.nodes.forEach((node) => {
+    if (!isInvocationNode(node)) {
+      return;
+    }
+
+    forEach(node.data.inputs, (input) => {
+      if (isImageFieldInputInstance(input) && input.value?.image_name === imageDTO.image_name) {
+        dispatch(
+          fieldImageValueChanged({
+            nodeId: node.data.id,
+            fieldName: input.name,
+            value: undefined,
+          })
+        );
+      }
+    });
+  });
+};
+
+const deleteControlAdapterImages = (state: RootState, dispatch: AppDispatch, imageDTO: ImageDTO) => {
+  forEach(selectControlAdapterAll(state.controlAdapters), (ca) => {
+    if (
+      ca.controlImage === imageDTO.image_name ||
+      (isControlNetOrT2IAdapter(ca) && ca.processedControlImage === imageDTO.image_name)
+    ) {
+      dispatch(
+        controlAdapterImageChanged({
+          id: ca.id,
+          controlImage: null,
+        })
+      );
+      dispatch(
+        controlAdapterProcessedImageChanged({
+          id: ca.id,
+          processedControlImage: null,
+        })
+      );
+    }
+  });
+};
+
+const deleteControlLayerImages = (state: RootState, dispatch: AppDispatch, imageDTO: ImageDTO) => {
+  state.controlLayers.present.layers.forEach((l) => {
+    if (isRegionalGuidanceLayer(l)) {
+      if (l.ipAdapters.some((ipa) => ipa.image?.imageName === imageDTO.image_name)) {
+        dispatch(layerDeleted(l.id));
+      }
+    }
+    if (isControlAdapterLayer(l)) {
+      if (
+        l.controlAdapter.image?.imageName === imageDTO.image_name ||
+        l.controlAdapter.processedImage?.imageName === imageDTO.image_name
+      ) {
+        dispatch(layerDeleted(l.id));
+      }
+    }
+    if (isIPAdapterLayer(l)) {
+      if (l.ipAdapter.image?.imageName === imageDTO.image_name) {
+        dispatch(layerDeleted(l.id));
+      }
+    }
+    if (isInitialImageLayer(l)) {
+      if (l.image?.imageName === imageDTO.image_name) {
+        dispatch(layerDeleted(l.id));
+      }
+    }
+  });
+};
+
 export const addRequestedSingleImageDeletionListener = (startAppListening: AppStartListening) => {
   startAppListening({
     actionCreator: imageDeletionConfirmed,
@@ -72,45 +151,9 @@ export const addRequestedSingleImageDeletionListener = (startAppListening: AppSt
       }
 
       imageDTOs.forEach((imageDTO) => {
-        // reset control adapters that use the deleted images
+        deleteControlAdapterImages(state, dispatch, imageDTO);
-        forEach(selectControlAdapterAll(getState().controlAdapters), (ca) => {
+        deleteNodesImages(state, dispatch, imageDTO);
-          if (
+        deleteControlLayerImages(state, dispatch, imageDTO);
-            ca.controlImage === imageDTO.image_name ||
-            (isControlNetOrT2IAdapter(ca) && ca.processedControlImage === imageDTO.image_name)
-          ) {
-            dispatch(
-              controlAdapterImageChanged({
-                id: ca.id,
-                controlImage: null,
-              })
-            );
-            dispatch(
-              controlAdapterProcessedImageChanged({
-                id: ca.id,
-                processedControlImage: null,
-              })
-            );
-          }
-        });
-
-        // reset nodes that use the deleted images
-        getState().nodes.nodes.forEach((node) => {
-          if (!isInvocationNode(node)) {
-            return;
-          }
-
-          forEach(node.data.inputs, (input) => {
-            if (isImageFieldInputInstance(input) && input.value?.image_name === imageDTO.image_name) {
-              dispatch(
-                fieldImageValueChanged({
-                  nodeId: node.data.id,
-                  fieldName: input.name,
-                  value: undefined,
-                })
-              );
-            }
-          });
-        });
       });
 
       // Delete from server
@@ -162,45 +205,9 @@ export const addRequestedSingleImageDeletionListener = (startAppListening: AppSt
         }
 
         imageDTOs.forEach((imageDTO) => {
-          // reset control adapters that use the deleted images
+          deleteControlAdapterImages(state, dispatch, imageDTO);
-          forEach(selectControlAdapterAll(getState().controlAdapters), (ca) => {
+          deleteNodesImages(state, dispatch, imageDTO);
-            if (
+          deleteControlLayerImages(state, dispatch, imageDTO);
-              ca.controlImage === imageDTO.image_name ||
-              (isControlNetOrT2IAdapter(ca) && ca.processedControlImage === imageDTO.image_name)
-            ) {
-              dispatch(
-                controlAdapterImageChanged({
-                  id: ca.id,
-                  controlImage: null,
-                })
-              );
-              dispatch(
-                controlAdapterProcessedImageChanged({
-                  id: ca.id,
-                  processedControlImage: null,
-                })
-              );
-            }
-          });
-
-          // reset nodes that use the deleted images
-          getState().nodes.nodes.forEach((node) => {
-            if (!isInvocationNode(node)) {
-              return;
-            }
-
-            forEach(node.data.inputs, (input) => {
-              if (isImageFieldInputInstance(input) && input.value?.image_name === imageDTO.image_name) {
-                dispatch(
-                  fieldImageValueChanged({
-                    nodeId: node.data.id,
-                    fieldName: input.name,
-                    value: undefined,
-                  })
-                );
-              }
-            });
-          });
         });
       } catch {
         // no-op

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

@@ -3,6 +3,7 @@ import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { selectCanvasSlice } from 'features/canvas/store/canvasSlice';
 import { selectControlAdaptersSlice } from 'features/controlAdapters/store/controlAdaptersSlice';
+import { selectControlLayersSlice } from 'features/controlLayers/store/controlLayersSlice';
 import { imageDeletionConfirmed } from 'features/deleteImageModal/store/actions';
 import { getImageUsage, selectImageUsage } from 'features/deleteImageModal/store/selectors';
 import {
@@ -12,7 +13,6 @@ import {
 } from 'features/deleteImageModal/store/slice';
 import type { ImageUsage } from 'features/deleteImageModal/store/types';
 import { selectNodesSlice } from 'features/nodes/store/nodesSlice';
-import { selectGenerationSlice } from 'features/parameters/store/generationSlice';
 import { setShouldConfirmOnDelete } from 'features/system/store/systemSlice';
 import { some } from 'lodash-es';
 import type { ChangeEvent } from 'react';
@@ -24,23 +24,24 @@ import ImageUsageMessage from './ImageUsageMessage';
 const selectImageUsages = createMemoizedSelector(
   [
     selectDeleteImageModalSlice,
-    selectGenerationSlice,
     selectCanvasSlice,
     selectNodesSlice,
     selectControlAdaptersSlice,
+    selectControlLayersSlice,
     selectImageUsage,
   ],
-  (deleteImageModal, generation, canvas, nodes, controlAdapters, imagesUsage) => {
+  (deleteImageModal, canvas, nodes, controlAdapters, controlLayers, imagesUsage) => {
     const { imagesToDelete } = deleteImageModal;
 
     const allImageUsage = (imagesToDelete ?? []).map(({ image_name }) =>
-      getImageUsage(generation, canvas, nodes, controlAdapters, image_name)
+      getImageUsage(canvas, nodes, controlAdapters, controlLayers.present, image_name)
     );
 
     const imageUsageSummary: ImageUsage = {
       isCanvasImage: some(allImageUsage, (i) => i.isCanvasImage),
       isNodesImage: some(allImageUsage, (i) => i.isNodesImage),
       isControlImage: some(allImageUsage, (i) => i.isControlImage),
+      isControlLayerImage: some(allImageUsage, (i) => i.isControlLayerImage),
     };
 
     return {

invokeai/frontend/web/src/features/deleteImageModal/components/ImageUsageMessage.tsx

@@ -29,9 +29,10 @@ const ImageUsageMessage = (props: Props) => {
     <>
       <Text>{topMessage}</Text>
       <UnorderedList paddingInlineStart={6}>
-        {imageUsage.isCanvasImage && <ListItem>{t('common.unifiedCanvas')}</ListItem>}
+        {imageUsage.isCanvasImage && <ListItem>{t('ui.tabs.canvasTab')}</ListItem>}
         {imageUsage.isControlImage && <ListItem>{t('common.controlNet')}</ListItem>}
-        {imageUsage.isNodesImage && <ListItem>{t('common.nodeEditor')}</ListItem>}
+        {imageUsage.isNodesImage && <ListItem>{t('ui.tabs.workflowsTab')}</ListItem>}
+        {imageUsage.isControlLayerImage && <ListItem>{t('ui.tabs.generationTab')}</ListItem>}
       </UnorderedList>
       <Text>{bottomMessage}</Text>
     </>

invokeai/frontend/web/src/features/deleteImageModal/store/selectors.ts

@@ -7,22 +7,28 @@ import {
 } from 'features/controlAdapters/store/controlAdaptersSlice';
 import type { ControlAdaptersState } from 'features/controlAdapters/store/types';
 import { isControlNetOrT2IAdapter } from 'features/controlAdapters/store/types';
+import {
+  isControlAdapterLayer,
+  isInitialImageLayer,
+  isIPAdapterLayer,
+  isRegionalGuidanceLayer,
+  selectControlLayersSlice,
+} from 'features/controlLayers/store/controlLayersSlice';
+import type { ControlLayersState } from 'features/controlLayers/store/types';
 import { selectDeleteImageModalSlice } from 'features/deleteImageModal/store/slice';
 import { selectNodesSlice } from 'features/nodes/store/nodesSlice';
 import type { NodesState } from 'features/nodes/store/types';
 import { isImageFieldInputInstance } from 'features/nodes/types/field';
 import { isInvocationNode } from 'features/nodes/types/invocation';
-import { selectGenerationSlice } from 'features/parameters/store/generationSlice';
-import type { GenerationState } from 'features/parameters/store/types';
 import { some } from 'lodash-es';
 
 import type { ImageUsage } from './types';
 
 export const getImageUsage = (
-  generation: GenerationState,
   canvas: CanvasState,
   nodes: NodesState,
   controlAdapters: ControlAdaptersState,
+  controlLayers: ControlLayersState,
   image_name: string
 ) => {
   const isCanvasImage = canvas.layerState.objects.some((obj) => obj.kind === 'image' && obj.imageName === image_name);
@@ -38,10 +44,29 @@ export const getImageUsage = (
     (ca) => ca.controlImage === image_name || (isControlNetOrT2IAdapter(ca) && ca.processedControlImage === image_name)
   );
 
+  const isControlLayerImage = controlLayers.layers.some((l) => {
+    if (isRegionalGuidanceLayer(l)) {
+      return l.ipAdapters.some((ipa) => ipa.image?.imageName === image_name);
+    }
+    if (isControlAdapterLayer(l)) {
+      return (
+        l.controlAdapter.image?.imageName === image_name || l.controlAdapter.processedImage?.imageName === image_name
+      );
+    }
+    if (isIPAdapterLayer(l)) {
+      return l.ipAdapter.image?.imageName === image_name;
+    }
+    if (isInitialImageLayer(l)) {
+      return l.image?.imageName === image_name;
+    }
+    return false;
+  });
+
   const imageUsage: ImageUsage = {
     isCanvasImage,
     isNodesImage,
     isControlImage,
+    isControlLayerImage,
   };
 
   return imageUsage;
@@ -49,11 +74,11 @@ export const getImageUsage = (
 
 export const selectImageUsage = createMemoizedSelector(
   selectDeleteImageModalSlice,
-  selectGenerationSlice,
   selectCanvasSlice,
   selectNodesSlice,
   selectControlAdaptersSlice,
-  (deleteImageModal, generation, canvas, nodes, controlAdapters) => {
+  selectControlLayersSlice,
+  (deleteImageModal, canvas, nodes, controlAdapters, controlLayers) => {
     const { imagesToDelete } = deleteImageModal;
 
     if (!imagesToDelete.length) {
@@ -61,7 +86,7 @@ export const selectImageUsage = createMemoizedSelector(
     }
 
     const imagesUsage = imagesToDelete.map((i) =>
-      getImageUsage(generation, canvas, nodes, controlAdapters, i.image_name)
+      getImageUsage(canvas, nodes, controlAdapters, controlLayers.present, i.image_name)
     );
 
     return imagesUsage;

invokeai/frontend/web/src/features/deleteImageModal/store/types.ts

@@ -9,4 +9,5 @@ export type ImageUsage = {
   isCanvasImage: boolean;
   isNodesImage: boolean;
   isControlImage: boolean;
+  isControlLayerImage: boolean;
 };

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

@@ -15,11 +15,11 @@ import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
 import { useAppSelector } from 'app/store/storeHooks';
 import { selectCanvasSlice } from 'features/canvas/store/canvasSlice';
 import { selectControlAdaptersSlice } from 'features/controlAdapters/store/controlAdaptersSlice';
+import { selectControlLayersSlice } from 'features/controlLayers/store/controlLayersSlice';
 import ImageUsageMessage from 'features/deleteImageModal/components/ImageUsageMessage';
 import { getImageUsage } from 'features/deleteImageModal/store/selectors';
 import type { ImageUsage } from 'features/deleteImageModal/store/types';
 import { selectNodesSlice } from 'features/nodes/store/nodesSlice';
-import { selectGenerationSlice } from 'features/parameters/store/generationSlice';
 import { some } from 'lodash-es';
 import { memo, useCallback, useMemo, useRef } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -43,16 +43,17 @@ const DeleteBoardModal = (props: Props) => {
   const selectImageUsageSummary = useMemo(
     () =>
       createMemoizedSelector(
-        [selectGenerationSlice, selectCanvasSlice, selectNodesSlice, selectControlAdaptersSlice],
+        [selectCanvasSlice, selectNodesSlice, selectControlAdaptersSlice, selectControlLayersSlice],
-        (generation, canvas, nodes, controlAdapters) => {
+        (canvas, nodes, controlAdapters, controlLayers) => {
           const allImageUsage = (boardImageNames ?? []).map((imageName) =>
-            getImageUsage(generation, canvas, nodes, controlAdapters, imageName)
+            getImageUsage(canvas, nodes, controlAdapters, controlLayers.present, imageName)
           );
 
           const imageUsageSummary: ImageUsage = {
             isCanvasImage: some(allImageUsage, (i) => i.isCanvasImage),
             isNodesImage: some(allImageUsage, (i) => i.isNodesImage),
             isControlImage: some(allImageUsage, (i) => i.isControlImage),
+            isControlLayerImage: some(allImageUsage, (i) => i.isControlLayerImage),
           };
 
           return imageUsageSummary;