From 243e76dd806bc8909211004668db5862812615db Mon Sep 17 00:00:00 2001 From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com> Date: Tue, 29 Aug 2023 23:48:28 +1200 Subject: [PATCH] feat: Send Canvas Image & Mask To ControlNet --- .../middleware/listenerMiddleware/index.ts | 8 ++- .../listeners/canvasImageToControlNet.ts | 58 +++++++++++++++ .../listeners/canvasMaskToControlNet.ts | 70 +++++++++++++++++++ .../web/src/features/canvas/store/actions.ts | 9 +++ .../controlNet/components/ControlNet.tsx | 8 +++ .../imports/ControlNetCanvasImageImports.tsx | 54 ++++++++++++++ 6 files changed, 205 insertions(+), 2 deletions(-) create mode 100644 invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasImageToControlNet.ts create mode 100644 invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMaskToControlNet.ts create mode 100644 invokeai/frontend/web/src/features/controlNet/components/imports/ControlNetCanvasImageImports.tsx diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts index abb17d1eec..4afe023fbb 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts @@ -15,7 +15,9 @@ import { addDeleteBoardAndImagesFulfilledListener } from './listeners/boardAndIm import { addBoardIdSelectedListener } from './listeners/boardIdSelected'; import { addCanvasCopiedToClipboardListener } from './listeners/canvasCopiedToClipboard'; import { addCanvasDownloadedAsImageListener } from './listeners/canvasDownloadedAsImage'; +import { addCanvasImageToControlNetListener } from './listeners/canvasImageToControlNet'; import { addCanvasMaskSavedToGalleryListener } from './listeners/canvasMaskSavedToGallery'; +import { addCanvasMaskToControlNetListener } from './listeners/canvasMaskToControlNet'; import { addCanvasMergedListener } from './listeners/canvasMerged'; import { addCanvasSavedToGalleryListener } from './listeners/canvasSavedToGallery'; import { addControlNetAutoProcessListener } from './listeners/controlNetAutoProcess'; @@ -41,6 +43,8 @@ import { addImageUploadedFulfilledListener, addImageUploadedRejectedListener, } from './listeners/imageUploaded'; +import { addImagesStarredListener } from './listeners/imagesStarred'; +import { addImagesUnstarredListener } from './listeners/imagesUnstarred'; import { addInitialImageSelectedListener } from './listeners/initialImageSelected'; import { addModelSelectedListener } from './listeners/modelSelected'; import { addModelsLoadedListener } from './listeners/modelsLoaded'; @@ -80,8 +84,6 @@ import { addUserInvokedCanvasListener } from './listeners/userInvokedCanvas'; import { addUserInvokedImageToImageListener } from './listeners/userInvokedImageToImage'; import { addUserInvokedNodesListener } from './listeners/userInvokedNodes'; import { addUserInvokedTextToImageListener } from './listeners/userInvokedTextToImage'; -import { addImagesStarredListener } from './listeners/imagesStarred'; -import { addImagesUnstarredListener } from './listeners/imagesUnstarred'; export const listenerMiddleware = createListenerMiddleware(); @@ -137,6 +139,8 @@ addSessionReadyToInvokeListener(); // Canvas actions addCanvasSavedToGalleryListener(); addCanvasMaskSavedToGalleryListener(); +addCanvasImageToControlNetListener(); +addCanvasMaskToControlNetListener(); addCanvasDownloadedAsImageListener(); addCanvasCopiedToClipboardListener(); addCanvasMergedListener(); diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasImageToControlNet.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasImageToControlNet.ts new file mode 100644 index 0000000000..fb411a6e25 --- /dev/null +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasImageToControlNet.ts @@ -0,0 +1,58 @@ +import { logger } from 'app/logging/logger'; +import { canvasImageToControlNet } from 'features/canvas/store/actions'; +import { getBaseLayerBlob } from 'features/canvas/util/getBaseLayerBlob'; +import { controlNetImageChanged } from 'features/controlNet/store/controlNetSlice'; +import { addToast } from 'features/system/store/systemSlice'; +import { imagesApi } from 'services/api/endpoints/images'; +import { startAppListening } from '..'; + +export const addCanvasImageToControlNetListener = () => { + startAppListening({ + actionCreator: canvasImageToControlNet, + effect: async (action, { dispatch, getState }) => { + const log = logger('canvas'); + const state = getState(); + + const blob = await getBaseLayerBlob(state); + + if (!blob) { + log.error('Problem getting base layer blob'); + dispatch( + addToast({ + title: 'Problem Saving Canvas', + description: 'Unable to export base layer', + status: 'error', + }) + ); + return; + } + + const { autoAddBoardId } = state.gallery; + + const imageDTO = await dispatch( + imagesApi.endpoints.uploadImage.initiate({ + file: new File([blob], 'savedCanvas.png', { + type: 'image/png', + }), + image_category: 'mask', + is_intermediate: false, + board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId, + crop_visible: true, + postUploadAction: { + type: 'TOAST', + toastOptions: { title: 'Canvas Sent to ControlNet & Assets' }, + }, + }) + ).unwrap(); + + const { image_name } = imageDTO; + + dispatch( + controlNetImageChanged({ + controlNetId: action.payload.controlNet.controlNetId, + controlImage: image_name, + }) + ); + }, + }); +}; diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMaskToControlNet.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMaskToControlNet.ts new file mode 100644 index 0000000000..6c97259f02 --- /dev/null +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMaskToControlNet.ts @@ -0,0 +1,70 @@ +import { logger } from 'app/logging/logger'; +import { canvasMaskToControlNet } from 'features/canvas/store/actions'; +import { getCanvasData } from 'features/canvas/util/getCanvasData'; +import { controlNetImageChanged } from 'features/controlNet/store/controlNetSlice'; +import { addToast } from 'features/system/store/systemSlice'; +import { imagesApi } from 'services/api/endpoints/images'; +import { startAppListening } from '..'; + +export const addCanvasMaskToControlNetListener = () => { + startAppListening({ + actionCreator: canvasMaskToControlNet, + effect: async (action, { dispatch, getState }) => { + const log = logger('canvas'); + const state = getState(); + + const canvasBlobsAndImageData = await getCanvasData( + state.canvas.layerState, + state.canvas.boundingBoxCoordinates, + state.canvas.boundingBoxDimensions, + state.canvas.isMaskEnabled, + state.canvas.shouldPreserveMaskedArea + ); + + if (!canvasBlobsAndImageData) { + return; + } + + const { maskBlob } = canvasBlobsAndImageData; + + if (!maskBlob) { + log.error('Problem getting mask layer blob'); + dispatch( + addToast({ + title: 'Problem Importing Mask', + description: 'Unable to export mask', + status: 'error', + }) + ); + return; + } + + const { autoAddBoardId } = state.gallery; + + const imageDTO = await dispatch( + imagesApi.endpoints.uploadImage.initiate({ + file: new File([maskBlob], 'canvasMaskImage.png', { + type: 'image/png', + }), + image_category: 'mask', + is_intermediate: false, + board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId, + crop_visible: true, + postUploadAction: { + type: 'TOAST', + toastOptions: { title: 'Mask Sent to ControlNet & Assets' }, + }, + }) + ).unwrap(); + + const { image_name } = imageDTO; + + dispatch( + controlNetImageChanged({ + controlNetId: action.payload.controlNet.controlNetId, + controlImage: image_name, + }) + ); + }, + }); +}; diff --git a/invokeai/frontend/web/src/features/canvas/store/actions.ts b/invokeai/frontend/web/src/features/canvas/store/actions.ts index b4efa76e42..85e0a7b406 100644 --- a/invokeai/frontend/web/src/features/canvas/store/actions.ts +++ b/invokeai/frontend/web/src/features/canvas/store/actions.ts @@ -1,4 +1,5 @@ import { createAction } from '@reduxjs/toolkit'; +import { ControlNetConfig } from 'features/controlNet/store/controlNetSlice'; import { ImageDTO } from 'services/api/types'; export const canvasSavedToGallery = createAction('canvas/canvasSavedToGallery'); @@ -20,3 +21,11 @@ export const canvasMerged = createAction('canvas/canvasMerged'); export const stagingAreaImageSaved = createAction<{ imageDTO: ImageDTO }>( 'canvas/stagingAreaImageSaved' ); + +export const canvasMaskToControlNet = createAction<{ + controlNet: ControlNetConfig; +}>('canvas/canvasMaskToControlNet'); + +export const canvasImageToControlNet = createAction<{ + controlNet: ControlNetConfig; +}>('canvas/canvasImageToControlNet'); diff --git a/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx b/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx index de9995c577..1f70542494 100644 --- a/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx +++ b/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx @@ -17,11 +17,13 @@ import { stateSelector } from 'app/store/store'; import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions'; import IAIIconButton from 'common/components/IAIIconButton'; import IAISwitch from 'common/components/IAISwitch'; +import { activeTabNameSelector } from 'features/ui/store/uiSelectors'; import { useToggle } from 'react-use'; import { v4 as uuidv4 } from 'uuid'; import ControlNetImagePreview from './ControlNetImagePreview'; import ControlNetProcessorComponent from './ControlNetProcessorComponent'; import ParamControlNetShouldAutoConfig from './ParamControlNetShouldAutoConfig'; +import ControlNetCanvasImageImports from './imports/ControlNetCanvasImageImports'; import ParamControlNetBeginEnd from './parameters/ParamControlNetBeginEnd'; import ParamControlNetControlMode from './parameters/ParamControlNetControlMode'; import ParamControlNetProcessorSelect from './parameters/ParamControlNetProcessorSelect'; @@ -36,6 +38,8 @@ const ControlNet = (props: ControlNetProps) => { const { controlNetId } = controlNet; const dispatch = useAppDispatch(); + const activeTabName = useAppSelector(activeTabNameSelector); + const selector = createSelector( stateSelector, ({ controlNet }) => { @@ -108,6 +112,9 @@ const ControlNet = (props: ControlNetProps) => { > + {activeTabName === 'unifiedCanvas' && ( + + )} { /> )} + { + const { controlNet } = props; + const dispatch = useAppDispatch(); + + const handleImportImageFromCanvas = useCallback(() => { + dispatch(canvasImageToControlNet({ controlNet })); + }, [controlNet, dispatch]); + + const handleImportMaskFromCanvas = useCallback(() => { + dispatch(canvasMaskToControlNet({ controlNet })); + }, [controlNet, dispatch]); + + return ( + + } + tooltip="Import Image From Canvas" + aria-label="Import Image From Canvas" + onClick={handleImportImageFromCanvas} + /> + } + tooltip="Import Mask From Canvas" + aria-label="Import Mask From Canvas" + onClick={handleImportMaskFromCanvas} + /> + + ); +}; + +export default memo(ControlNetCanvasImageImports);