diff --git a/invokeai/frontend/web/src/app/contexts/DeleteImageContext.tsx b/invokeai/frontend/web/src/app/contexts/DeleteImageContext.tsx index 2f2bc4625b..8263b48114 100644 --- a/invokeai/frontend/web/src/app/contexts/DeleteImageContext.tsx +++ b/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'; -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'; +export type ImageUsage = { + isInitialImage: boolean; + isCanvasImage: boolean; + isNodesImage: boolean; + 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(); 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, - ] + // Check where the image to be deleted is used (eg init image, controlnet, etc.) + const imageUsage = useAppSelector((state) => + selectImageUsage(state, imageToDelete?.image_name) ); + // 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( - (image?: ImageDTO) => { - if (!image) { - return; - } - setImageToDelete(image); - handleGatedDeletion(image); - }, - [handleGatedDeletion] - ); + // Consumers of the context call this to delete an image + const onDelete = useCallback((image?: ImageDTO) => { + if (!image) { + return; + } + // Set the image to delete, then let the effect call the actual deletion + setImageToDelete(image); + }, []); + + 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 ( { onClose: closeAndClearImageToDelete, onDelete, onImmediatelyDelete, + imageUsage, }} > {props.children} diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDeleted.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDeleted.ts index b527b5d00b..f4376a4959 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDeleted.ts +++ b/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; - if (!image) { - moduleLog.warn('No image provided'); - return; - } + const { image, imageUsage } = action.payload; 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)); diff --git a/invokeai/frontend/web/src/common/hooks/useImageUsage.ts b/invokeai/frontend/web/src/common/hooks/useImageUsage.ts deleted file mode 100644 index cf762f7880..0000000000 --- a/invokeai/frontend/web/src/common/hooks/useImageUsage.ts +++ /dev/null @@ -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; -}; diff --git a/invokeai/frontend/web/src/features/gallery/components/DeleteImageModal.tsx b/invokeai/frontend/web/src/features/gallery/components/DeleteImageModal.tsx index 335944df43..0ce7bb3666 100644 --- a/invokeai/frontend/web/src/features/gallery/components/DeleteImageModal.tsx +++ b/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) => dispatch(setShouldConfirmOnDelete(!e.target.checked)), diff --git a/invokeai/frontend/web/src/features/gallery/store/actions.ts b/invokeai/frontend/web/src/features/gallery/store/actions.ts index 7c00201da9..8b2beb9c13 100644 --- a/invokeai/frontend/web/src/features/gallery/store/actions.ts +++ b/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< - ImageDTO | ImageNameAndOrigin | undefined ->('gallery/requestedImageDeletion'); +export type RequestedImageDeletionArg = { + image: ImageDTO; + imageUsage: ImageUsage; +}; + +export const requestedImageDeletion = createAction( + 'gallery/requestedImageDeletion' +); export const sentImageToCanvas = createAction('gallery/sentImageToCanvas');