Skip to content

Commit

Permalink
add image usage for board images and listener to handle actual deletion
Browse files Browse the repository at this point in the history
  • Loading branch information
Mary Hipp authored and psychedelicious committed Jun 29, 2023
1 parent ba67e57 commit 723d68e
Show file tree
Hide file tree
Showing 6 changed files with 213 additions and 13 deletions.
93 changes: 81 additions & 12 deletions invokeai/frontend/web/src/app/contexts/DeleteBoardImagesContext.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
/**
Expand All @@ -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;
Expand All @@ -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
Expand All @@ -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(
Expand All @@ -93,6 +161,7 @@ export const DeleteBoardImagesContextProvider = (props: Props) => {
onClickDeleteBoardImages,
handleDeleteBoardImages,
handleDeleteBoardOnly,
imagesUsage,
}}
>
{props.children}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ import {
addImageRemovedFromBoardRejectedListener,
} from './listeners/imageRemovedFromBoard';
import { addReceivedOpenAPISchemaListener } from './listeners/receivedOpenAPISchema';
import { addRequestedBoardImageDeletionListener } from './listeners/boardImagesDeleted';

export const listenerMiddleware = createListenerMiddleware();

Expand Down Expand Up @@ -124,6 +125,7 @@ addRequestedImageDeletionListener();
addImageDeletedPendingListener();
addImageDeletedFulfilledListener();
addImageDeletedRejectedListener();
addRequestedBoardImageDeletionListener();

// Image metadata
addImageMetadataReceivedFulfilledListener();
Expand Down
Original file line number Diff line number Diff line change
@@ -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 },
])
);
}
},
});
};
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -23,6 +57,7 @@ const DeleteBoardImagesModal = () => {
board,
handleDeleteBoardImages,
handleDeleteBoardOnly,
imagesUsage,
} = useContext(DeleteBoardImagesContext);

const cancelRef = useRef<HTMLButtonElement>(null);
Expand All @@ -43,6 +78,7 @@ const DeleteBoardImagesModal = () => {

<AlertDialogBody>
<Flex direction="column" gap={3}>
<BoardImageInUseMessage imagesUsage={imagesUsage} />
<Divider />
<Text>{t('common.areYouSure')}</Text>
<Text fontWeight="bold">
Expand Down
12 changes: 11 additions & 1 deletion invokeai/frontend/web/src/features/gallery/store/actions.ts
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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');
Original file line number Diff line number Diff line change
Expand Up @@ -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;
},
Expand Down Expand Up @@ -117,6 +120,7 @@ export const {
imageUpserted,
imageUpdatedOne,
imageRemoved,
imagesRemoved,
imageCategoriesChanged,
} = imagesSlice.actions;

Expand Down

0 comments on commit 723d68e

Please sign in to comment.