From d6c08ba46921cdadccbba767bc39de8f0337faf7 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Sat, 3 Jun 2023 15:05:49 +1000 Subject: [PATCH] feat(ui): add mini/advanced controlnet ui --- .../middleware/listenerMiddleware/index.ts | 2 +- ...amsChanged.ts => controlNetAutoProcess.ts} | 52 ++--- .../web/src/common/components/IAISwitch.tsx | 8 +- .../controlNet/components/ControlNet.tsx | 194 +++++++++++++++--- .../components/ControlNetImagePreview.tsx | 95 +++++---- .../controlNet/components/ControlNetMini.tsx | 153 -------------- .../components/ControlNetPreprocessButton.tsx | 4 +- .../ParamControlNetBeginStepPct.tsx | 58 ------ .../parameters/ParamControlNetEndStepPct.tsx | 42 ---- .../parameters/ParamControlNetIsEnabled.tsx | 2 +- .../ParamControlNetIsPreprocessed.tsx | 4 +- .../parameters/ParamControlNetModel.tsx | 4 +- .../features/controlNet/store/constants.ts | 26 ++- .../controlNet/store/controlNetSlice.ts | 71 +++---- .../nodes/util/addControlNetToLinearGraph.ts | 100 +++++++++ .../graphBuilders/buildImageToImageGraph.ts | 3 + .../graphBuilders/buildTextToImageGraph.ts | 93 +-------- .../ControlNet/ParamControlNetCollapse.tsx | 67 +----- 18 files changed, 430 insertions(+), 548 deletions(-) rename invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/{controlNetProcessorParamsChanged.ts => controlNetAutoProcess.ts} (54%) delete mode 100644 invokeai/frontend/web/src/features/controlNet/components/ControlNetMini.tsx delete mode 100644 invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetBeginStepPct.tsx delete mode 100644 invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetEndStepPct.tsx create mode 100644 invokeai/frontend/web/src/features/nodes/util/addControlNetToLinearGraph.ts 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 8c6503521c..a9349dc863 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts @@ -71,7 +71,7 @@ import { addStagingAreaImageSavedListener } from './listeners/stagingAreaImageSa import { addCommitStagingAreaImageListener } from './listeners/addCommitStagingAreaImageListener'; import { addImageCategoriesChangedListener } from './listeners/imageCategoriesChanged'; import { addControlNetImageProcessedListener } from './listeners/controlNetImageProcessed'; -import { addControlNetAutoProcessListener } from './listeners/controlNetProcessorParamsChanged'; +import { addControlNetAutoProcessListener } from './listeners/controlNetAutoProcess'; export const listenerMiddleware = createListenerMiddleware(); diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetProcessorParamsChanged.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetAutoProcess.ts similarity index 54% rename from invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetProcessorParamsChanged.ts rename to invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetAutoProcess.ts index 11237c9d27..d53907e673 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetProcessorParamsChanged.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetAutoProcess.ts @@ -1,3 +1,4 @@ +import { AnyAction } from '@reduxjs/toolkit'; import { startAppListening } from '..'; import { log } from 'app/logging/useLogger'; import { controlNetImageProcessed } from 'features/controlNet/store/actions'; @@ -5,10 +6,37 @@ import { controlNetImageChanged, controlNetProcessorParamsChanged, controlNetProcessorTypeChanged, + isControlNetImagePreprocessedToggled, } from 'features/controlNet/store/controlNetSlice'; +import { RootState } from 'app/store/store'; const moduleLog = log.child({ namespace: 'controlNet' }); +const predicate = (action: AnyAction, state: RootState) => { + const isActionMatched = + controlNetProcessorParamsChanged.match(action) || + controlNetImageChanged.match(action) || + controlNetProcessorTypeChanged.match(action) || + isControlNetImagePreprocessedToggled.match(action); + + if (!isActionMatched) { + return false; + } + + const { controlNetId } = action.payload; + + const shouldAutoProcess = + !state.controlNet.controlNets[controlNetId].isPreprocessed; + + const isBusy = state.system.isProcessing; + + const hasControlImage = Boolean( + state.controlNet.controlNets[controlNetId].controlImage + ); + + return shouldAutoProcess && !isBusy && hasControlImage; +}; + /** * Listener that automatically processes a ControlNet image when its processor parameters are changed. * @@ -16,35 +44,13 @@ const moduleLog = log.child({ namespace: 'controlNet' }); */ export const addControlNetAutoProcessListener = () => { startAppListening({ - predicate: (action) => - controlNetProcessorParamsChanged.match(action) || - controlNetImageChanged.match(action) || - controlNetProcessorTypeChanged.match(action), + predicate, effect: async ( action, { dispatch, getState, cancelActiveListeners, delay } ) => { - const state = getState(); - if (!state.controlNet.shouldAutoProcess) { - // silently skip - return; - } - - if (state.system.isProcessing) { - moduleLog.trace('System busy, skipping ControlNet auto-processing'); - return; - } - const { controlNetId } = action.payload; - if (!state.controlNet.controlNets[controlNetId].controlImage) { - moduleLog.trace( - { data: { controlNetId } }, - 'No ControlNet image to auto-process' - ); - return; - } - // Cancel any in-progress instances of this listener cancelActiveListeners(); diff --git a/invokeai/frontend/web/src/common/components/IAISwitch.tsx b/invokeai/frontend/web/src/common/components/IAISwitch.tsx index e1bddb9f43..9a7ba7eb76 100644 --- a/invokeai/frontend/web/src/common/components/IAISwitch.tsx +++ b/invokeai/frontend/web/src/common/components/IAISwitch.tsx @@ -36,9 +36,11 @@ const IAISwitch = (props: Props) => { alignItems="center" {...formControlProps} > - - {label} - + {label && ( + + {label} + + )} ); diff --git a/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx b/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx index 2b86ca0e4d..fe5f07f89f 100644 --- a/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx +++ b/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx @@ -1,32 +1,42 @@ import { memo, useCallback } from 'react'; import { - ControlNet, - controlNetProcessedImageChanged, + ControlNetConfig, + controlNetAdded, controlNetRemoved, + controlNetToggled, + isControlNetImagePreprocessedToggled, } from '../store/controlNetSlice'; import { useAppDispatch } from 'app/store/storeHooks'; import ParamControlNetModel from './parameters/ParamControlNetModel'; import ParamControlNetWeight from './parameters/ParamControlNetWeight'; import { - Box, + Checkbox, Flex, - Tab, + FormControl, + FormLabel, + HStack, TabList, - TabPanel, TabPanels, Tabs, - Text, + Tab, + TabPanel, + Box, } from '@chakra-ui/react'; -import IAIButton from 'common/components/IAIButton'; -import { FaUndo } from 'react-icons/fa'; +import { FaCopy, FaTrash } from 'react-icons/fa'; + +import ParamControlNetBeginEnd from './parameters/ParamControlNetBeginEnd'; +import ControlNetImagePreview from './ControlNetImagePreview'; +import IAIIconButton from 'common/components/IAIIconButton'; +import { v4 as uuidv4 } from 'uuid'; +import { useToggle } from 'react-use'; import ParamControlNetProcessorSelect from './parameters/ParamControlNetProcessorSelect'; import ControlNetProcessorComponent from './ControlNetProcessorComponent'; import ControlNetPreprocessButton from './ControlNetPreprocessButton'; -import ParamControlNetBeginEnd from './parameters/ParamControlNetBeginEnd'; -import ControlNetImagePreview from './ControlNetImagePreview'; +import IAIButton from 'common/components/IAIButton'; +import IAISwitch from 'common/components/IAISwitch'; type ControlNetProps = { - controlNet: ControlNet; + controlNet: ControlNetConfig; }; const ControlNet = (props: ControlNetProps) => { @@ -38,24 +48,160 @@ const ControlNet = (props: ControlNetProps) => { beginStepPct, endStepPct, controlImage, - isControlImageProcessed, + isPreprocessed: isControlImageProcessed, processedControlImage, processorNode, } = props.controlNet; const dispatch = useAppDispatch(); - const handleReset = useCallback(() => { - dispatch( - controlNetProcessedImageChanged({ - controlNetId, - processedControlImage: null, - }) - ); + const [shouldShowAdvanced, onToggleAdvanced] = useToggle(true); + + const handleDelete = useCallback(() => { + dispatch(controlNetRemoved({ controlNetId })); }, [controlNetId, dispatch]); - const handleControlNetRemoved = useCallback(() => { - dispatch(controlNetRemoved(controlNetId)); + const handleDuplicate = useCallback(() => { + dispatch( + controlNetAdded({ controlNetId: uuidv4(), controlNet: props.controlNet }) + ); + }, [dispatch, props.controlNet]); + + const handleToggleIsEnabled = useCallback(() => { + dispatch(controlNetToggled({ controlNetId })); }, [controlNetId, dispatch]); + const handleToggleIsPreprocessed = useCallback(() => { + dispatch(isControlNetImagePreprocessedToggled({ controlNetId })); + }, [controlNetId, dispatch]); + + return ( + + + + + + + } + /> + } + /> + + {isEnabled && ( + <> + + {!shouldShowAdvanced && ( + + + + )} + + + + + + Preprocessed + + + + + + Advanced + + + + + + + + {shouldShowAdvanced && ( + <> + + + + {!isControlImageProcessed && ( + <> + + + + )} + + )} + + )} + + ); + return ( @@ -101,18 +247,18 @@ const ControlNet = (props: ControlNetProps) => { processorNode={processorNode} /> - } onClick={handleReset} isDisabled={Boolean(!processedControlImage)} > Reset Processing - + */} - Remove ControlNet + Remove ControlNet ); }; diff --git a/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx b/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx index 86a5a06569..099e58ce80 100644 --- a/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx +++ b/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx @@ -1,7 +1,7 @@ import { memo, useCallback, useRef, useState } from 'react'; import { ImageDTO } from 'services/api'; import { - ControlNet, + ControlNetConfig, controlNetImageChanged, controlNetSelector, } from '../store/controlNetSlice'; @@ -24,7 +24,7 @@ const selector = createSelector( ); type Props = { - controlNet: ControlNet; + controlNet: ControlNetConfig; }; const ControlNetImagePreview = (props: Props) => { @@ -32,7 +32,7 @@ const ControlNetImagePreview = (props: Props) => { controlNetId, controlImage, processedControlImage, - isControlImageProcessed, + isPreprocessed: isControlImageProcessed, } = props.controlNet; const dispatch = useAppDispatch(); const { isProcessingControlImage } = useAppSelector(selector); @@ -63,63 +63,62 @@ const ControlNetImagePreview = (props: Props) => { - {controlImage && - processedControlImage && - shouldShowProcessedImage && - !isProcessingControlImage && ( - + + {shouldShowProcessedImageBackdrop && ( + + )} - {shouldShowProcessedImageBackdrop && ( - - )} - - - + - - )} + + + )} {isProcessingControlImage && ( { - const { - controlNetId, - isEnabled, - model, - weight, - beginStepPct, - endStepPct, - controlImage, - isControlImageProcessed, - processedControlImage, - processorNode, - } = props.controlNet; - const dispatch = useAppDispatch(); - - const handleDelete = useCallback(() => { - dispatch(controlNetRemoved(controlNetId)); - }, [controlNetId, dispatch]); - - const handleDuplicate = useCallback(() => { - dispatch( - controlNetAdded({ controlNetId: uuidv4(), controlNet: props.controlNet }) - ); - }, [dispatch, props.controlNet]); - - const handleToggleIsEnabled = useCallback(() => { - dispatch(controlNetToggled(controlNetId)); - }, [controlNetId, dispatch]); - - const handleToggleIsPreprocessed = useCallback(() => { - dispatch(isControlNetImageProcessedToggled(controlNetId)); - }, [controlNetId, dispatch]); - - return ( - - - - } - /> - } - /> - - - - - - - - - - - - - Enabled - - - - - - Preprocessed - - - - - - - ); -}; - -export default memo(ControlNet); diff --git a/invokeai/frontend/web/src/features/controlNet/components/ControlNetPreprocessButton.tsx b/invokeai/frontend/web/src/features/controlNet/components/ControlNetPreprocessButton.tsx index 94b1f86501..95a4f968e5 100644 --- a/invokeai/frontend/web/src/features/controlNet/components/ControlNetPreprocessButton.tsx +++ b/invokeai/frontend/web/src/features/controlNet/components/ControlNetPreprocessButton.tsx @@ -1,12 +1,12 @@ import IAIButton from 'common/components/IAIButton'; import { memo, useCallback } from 'react'; -import { ControlNet } from '../store/controlNetSlice'; +import { ControlNetConfig } from '../store/controlNetSlice'; import { useAppDispatch } from 'app/store/storeHooks'; import { controlNetImageProcessed } from '../store/actions'; import { useIsReadyToInvoke } from 'common/hooks/useIsReadyToInvoke'; type Props = { - controlNet: ControlNet; + controlNet: ControlNetConfig; }; const ControlNetPreprocessButton = (props: Props) => { diff --git a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetBeginStepPct.tsx b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetBeginStepPct.tsx deleted file mode 100644 index d94db5e272..0000000000 --- a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetBeginStepPct.tsx +++ /dev/null @@ -1,58 +0,0 @@ -import { useAppDispatch } from 'app/store/storeHooks'; -import IAISlider from 'common/components/IAISlider'; -import { - controlNetBeginStepPctChanged, - controlNetEndStepPctChanged, -} from 'features/controlNet/store/controlNetSlice'; -import { memo, useCallback } from 'react'; - -type ParamControlNetBeginStepPctProps = { - controlNetId: string; - beginStepPct: number; -}; - -const ParamControlNetBeginStepPct = ( - props: ParamControlNetBeginStepPctProps -) => { - const { controlNetId, beginStepPct } = props; - const dispatch = useAppDispatch(); - - const handleBeginStepPctChanged = useCallback( - (beginStepPct: number) => { - dispatch(controlNetBeginStepPctChanged({ controlNetId, beginStepPct })); - }, - [controlNetId, dispatch] - ); - - const handleBeginStepPctReset = useCallback(() => { - dispatch(controlNetBeginStepPctChanged({ controlNetId, beginStepPct: 0 })); - }, [controlNetId, dispatch]); - - const handleEndStepPctChanged = useCallback( - (endStepPct: number) => { - dispatch(controlNetEndStepPctChanged({ controlNetId, endStepPct })); - }, - [controlNetId, dispatch] - ); - - const handleEndStepPctReset = useCallback(() => { - dispatch(controlNetEndStepPctChanged({ controlNetId, endStepPct: 0 })); - }, [controlNetId, dispatch]); - - return ( - - ); -}; - -export default memo(ParamControlNetBeginStepPct); diff --git a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetEndStepPct.tsx b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetEndStepPct.tsx deleted file mode 100644 index d3d831cf31..0000000000 --- a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetEndStepPct.tsx +++ /dev/null @@ -1,42 +0,0 @@ -import { useAppDispatch } from 'app/store/storeHooks'; -import IAISlider from 'common/components/IAISlider'; -import { controlNetEndStepPctChanged } from 'features/controlNet/store/controlNetSlice'; -import { memo, useCallback } from 'react'; - -type ParamControlNetEndStepPctProps = { - controlNetId: string; - endStepPct: number; -}; - -const ParamControlNetEndStepPct = (props: ParamControlNetEndStepPctProps) => { - const { controlNetId, endStepPct } = props; - const dispatch = useAppDispatch(); - - const handleEndStepPctChanged = useCallback( - (endStepPct: number) => { - dispatch(controlNetEndStepPctChanged({ controlNetId, endStepPct })); - }, - [controlNetId, dispatch] - ); - - const handleEndStepPctReset = () => { - dispatch(controlNetEndStepPctChanged({ controlNetId, endStepPct: 0 })); - }; - - return ( - - ); -}; - -export default memo(ParamControlNetEndStepPct); diff --git a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsEnabled.tsx b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsEnabled.tsx index f42265cb22..d7f519a7b6 100644 --- a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsEnabled.tsx +++ b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsEnabled.tsx @@ -13,7 +13,7 @@ const ParamControlNetIsEnabled = (props: ParamControlNetIsEnabledProps) => { const dispatch = useAppDispatch(); const handleIsEnabledChanged = useCallback(() => { - dispatch(controlNetToggled(controlNetId)); + dispatch(controlNetToggled({ controlNetId })); }, [dispatch, controlNetId]); return ( diff --git a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsPreprocessed.tsx b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsPreprocessed.tsx index 9e2658964d..6db61a0d15 100644 --- a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsPreprocessed.tsx +++ b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsPreprocessed.tsx @@ -3,7 +3,7 @@ import IAIFullCheckbox from 'common/components/IAIFullCheckbox'; import IAISwitch from 'common/components/IAISwitch'; import { controlNetToggled, - isControlNetImageProcessedToggled, + isControlNetImagePreprocessedToggled, } from 'features/controlNet/store/controlNetSlice'; import { memo, useCallback } from 'react'; @@ -18,7 +18,7 @@ const ParamControlNetIsEnabled = (props: ParamControlNetIsEnabledProps) => { const handleIsControlImageProcessedToggled = useCallback(() => { dispatch( - isControlNetImageProcessedToggled({ + isControlNetImagePreprocessedToggled({ controlNetId, }) ); diff --git a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx index e5a1ca1c39..beb34e9d24 100644 --- a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx +++ b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx @@ -3,8 +3,8 @@ import IAICustomSelect from 'common/components/IAICustomSelect'; import { CONTROLNET_MODELS, ControlNetModel, - controlNetModelChanged, -} from 'features/controlNet/store/controlNetSlice'; +} from 'features/controlNet/store/constants'; +import { controlNetModelChanged } from 'features/controlNet/store/controlNetSlice'; import { memo, useCallback } from 'react'; type ParamIsControlNetModelProps = { diff --git a/invokeai/frontend/web/src/features/controlNet/store/constants.ts b/invokeai/frontend/web/src/features/controlNet/store/constants.ts index a7e20a78d7..da3a9c57b5 100644 --- a/invokeai/frontend/web/src/features/controlNet/store/constants.ts +++ b/invokeai/frontend/web/src/features/controlNet/store/constants.ts @@ -22,7 +22,7 @@ type ControlNetProcessorsDict = Record< * * TODO: Generate from the OpenAPI schema */ -export const CONTROLNET_PROCESSORS = { +export const CONTROLNET_PROCESSORS: ControlNetProcessorsDict = { canny_image_processor: { type: 'canny_image_processor', label: 'Canny', @@ -164,3 +164,27 @@ export const CONTROLNET_PROCESSORS = { }, }, }; + +export const CONTROLNET_MODELS = [ + 'lllyasviel/sd-controlnet-canny', + 'lllyasviel/sd-controlnet-depth', + 'lllyasviel/sd-controlnet-hed', + 'lllyasviel/sd-controlnet-seg', + 'lllyasviel/sd-controlnet-openpose', + 'lllyasviel/sd-controlnet-scribble', + 'lllyasviel/sd-controlnet-normal', + 'lllyasviel/sd-controlnet-mlsd', +]; + +export type ControlNetModel = (typeof CONTROLNET_MODELS)[number]; + +export const CONTROLNET_MODEL_MAP: Record< + ControlNetModel, + ControlNetProcessorType +> = { + 'lllyasviel/sd-controlnet-canny': 'canny_image_processor', + 'lllyasviel/sd-controlnet-depth': 'midas_depth_image_processor', + 'lllyasviel/sd-controlnet-hed': 'hed_image_processor', + 'lllyasviel/sd-controlnet-openpose': 'openpose_image_processor', + 'lllyasviel/sd-controlnet-mlsd': 'mlsd_image_processor', +}; diff --git a/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts b/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts index 1155567d73..4847e3c1a5 100644 --- a/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts +++ b/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts @@ -7,36 +7,27 @@ import { RequiredCannyImageProcessorInvocation, RequiredControlNetProcessorNode, } from './types'; -import { CONTROLNET_PROCESSORS } from './constants'; +import { + CONTROLNET_MODELS, + CONTROLNET_PROCESSORS, + ControlNetModel, +} from './constants'; import { controlNetImageProcessed } from './actions'; -export const CONTROLNET_MODELS = [ - 'lllyasviel/sd-controlnet-canny', - 'lllyasviel/sd-controlnet-depth', - 'lllyasviel/sd-controlnet-hed', - 'lllyasviel/sd-controlnet-seg', - 'lllyasviel/sd-controlnet-openpose', - 'lllyasviel/sd-controlnet-scribble', - 'lllyasviel/sd-controlnet-normal', - 'lllyasviel/sd-controlnet-mlsd', -]; - -export type ControlNetModel = (typeof CONTROLNET_MODELS)[number]; - -export const initialControlNet: Omit = { +export const initialControlNet: Omit = { isEnabled: true, model: CONTROLNET_MODELS[0], weight: 1, beginStepPct: 0, endStepPct: 1, controlImage: null, - isControlImageProcessed: false, + isPreprocessed: false, processedControlImage: null, processorNode: CONTROLNET_PROCESSORS.canny_image_processor .default as RequiredCannyImageProcessorInvocation, }; -export type ControlNet = { +export type ControlNetConfig = { controlNetId: string; isEnabled: boolean; model: ControlNetModel; @@ -44,22 +35,20 @@ export type ControlNet = { beginStepPct: number; endStepPct: number; controlImage: ImageDTO | null; - isControlImageProcessed: boolean; + isPreprocessed: boolean; processedControlImage: ImageDTO | null; processorNode: RequiredControlNetProcessorNode; }; export type ControlNetState = { - controlNets: Record; + controlNets: Record; isEnabled: boolean; - shouldAutoProcess: boolean; isProcessingControlImage: boolean; }; export const initialControlNetState: ControlNetState = { controlNets: {}, isEnabled: false, - shouldAutoProcess: true, isProcessingControlImage: false, }; @@ -72,7 +61,10 @@ export const controlNetSlice = createSlice({ }, controlNetAdded: ( state, - action: PayloadAction<{ controlNetId: string; controlNet?: ControlNet }> + action: PayloadAction<{ + controlNetId: string; + controlNet?: ControlNetConfig; + }> ) => { const { controlNetId, controlNet } = action.payload; state.controlNets[controlNetId] = { @@ -91,12 +83,18 @@ export const controlNetSlice = createSlice({ controlImage, }; }, - controlNetRemoved: (state, action: PayloadAction) => { - const controlNetId = action.payload; + controlNetRemoved: ( + state, + action: PayloadAction<{ controlNetId: string }> + ) => { + const { controlNetId } = action.payload; delete state.controlNets[controlNetId]; }, - controlNetToggled: (state, action: PayloadAction) => { - const controlNetId = action.payload; + controlNetToggled: ( + state, + action: PayloadAction<{ controlNetId: string }> + ) => { + const { controlNetId } = action.payload; state.controlNets[controlNetId].isEnabled = !state.controlNets[controlNetId].isEnabled; }, @@ -110,17 +108,20 @@ export const controlNetSlice = createSlice({ const { controlNetId, controlImage } = action.payload; state.controlNets[controlNetId].controlImage = controlImage; state.controlNets[controlNetId].processedControlImage = null; - if (state.shouldAutoProcess && controlImage !== null) { + if ( + controlImage !== null && + !state.controlNets[controlNetId].isPreprocessed + ) { state.isProcessingControlImage = true; } }, - isControlNetImageProcessedToggled: ( + isControlNetImagePreprocessedToggled: ( state, - action: PayloadAction + action: PayloadAction<{ controlNetId: string }> ) => { - const controlNetId = action.payload; - state.controlNets[controlNetId].isControlImageProcessed = - !state.controlNets[controlNetId].isControlImageProcessed; + const { controlNetId } = action.payload; + state.controlNets[controlNetId].isPreprocessed = + !state.controlNets[controlNetId].isPreprocessed; }, controlNetProcessedImageChanged: ( state, @@ -191,9 +192,6 @@ export const controlNetSlice = createSlice({ processorType ].default as RequiredControlNetProcessorNode; }, - shouldAutoProcessToggled: (state) => { - state.shouldAutoProcess = !state.shouldAutoProcess; - }, }, extraReducers: (builder) => { builder.addCase(controlNetImageProcessed, (state, action) => { @@ -212,7 +210,7 @@ export const { controlNetAddedFromImage, controlNetRemoved, controlNetImageChanged, - isControlNetImageProcessedToggled, + isControlNetImagePreprocessedToggled, controlNetProcessedImageChanged, controlNetToggled, controlNetModelChanged, @@ -221,7 +219,6 @@ export const { controlNetEndStepPctChanged, controlNetProcessorParamsChanged, controlNetProcessorTypeChanged, - shouldAutoProcessToggled, } = controlNetSlice.actions; export default controlNetSlice.reducer; diff --git a/invokeai/frontend/web/src/features/nodes/util/addControlNetToLinearGraph.ts b/invokeai/frontend/web/src/features/nodes/util/addControlNetToLinearGraph.ts new file mode 100644 index 0000000000..b386b41dc7 --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/addControlNetToLinearGraph.ts @@ -0,0 +1,100 @@ +import { RootState } from 'app/store/store'; +import { forEach, size } from 'lodash-es'; +import { CollectInvocation, ControlNetInvocation } from 'services/api'; +import { NonNullableGraph } from '../types/types'; + +const CONTROL_NET_COLLECT = 'control_net_collect'; + +export const addControlNetToLinearGraph = ( + graph: NonNullableGraph, + baseNodeId: string, + state: RootState +): void => { + const { isEnabled: isControlNetEnabled, controlNets } = state.controlNet; + + // Add ControlNet + if (isControlNetEnabled) { + if (size(controlNets) > 1) { + const controlNetIterateNode: CollectInvocation = { + id: CONTROL_NET_COLLECT, + type: 'collect', + }; + graph.nodes[controlNetIterateNode.id] = controlNetIterateNode; + graph.edges.push({ + source: { node_id: controlNetIterateNode.id, field: 'collection' }, + destination: { + node_id: baseNodeId, + field: 'control', + }, + }); + } + + forEach(controlNets, (controlNet, index) => { + const { + controlNetId, + isEnabled, + isPreprocessed: isControlImageProcessed, + controlImage, + processedControlImage, + beginStepPct, + endStepPct, + model, + processorNode, + weight, + } = controlNet; + + if (!isEnabled) { + // Skip disabled ControlNets + return; + } + + const controlNetNode: ControlNetInvocation = { + id: `control_net_${controlNetId}`, + type: 'controlnet', + begin_step_percent: beginStepPct, + end_step_percent: endStepPct, + control_model: model as ControlNetInvocation['control_model'], + control_weight: weight, + }; + + if (processedControlImage && !isControlImageProcessed) { + // We've already processed the image in the app, so we can just use the processed image + const { image_name, image_origin } = processedControlImage; + controlNetNode.image = { + image_name, + image_origin, + }; + } else if (controlImage && isControlImageProcessed) { + // The control image is preprocessed + const { image_name, image_origin } = controlImage; + controlNetNode.image = { + image_name, + image_origin, + }; + } else { + // Skip ControlNets without an unprocessed image - should never happen if everything is working correctly + return; + } + + graph.nodes[controlNetNode.id] = controlNetNode; + + if (size(controlNets) > 1) { + graph.edges.push({ + source: { node_id: controlNetNode.id, field: 'control' }, + destination: { + node_id: CONTROL_NET_COLLECT, + field: 'item', + }, + }); + } else { + graph.edges.push({ + source: { node_id: controlNetNode.id, field: 'control' }, + destination: { + node_id: baseNodeId, + field: 'control', + }, + }); + } + }); + } +}; diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildImageToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildImageToImageGraph.ts index fe4f6c63b5..a1dc5d48ab 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildImageToImageGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildImageToImageGraph.ts @@ -14,6 +14,7 @@ import { import { NonNullableGraph } from 'features/nodes/types/types'; import { log } from 'app/logging/useLogger'; import { set } from 'lodash-es'; +import { addControlNetToLinearGraph } from '../addControlNetToLinearGraph'; const moduleLog = log.child({ namespace: 'nodes' }); @@ -408,5 +409,7 @@ export const buildImageToImageGraph = (state: RootState): Graph => { }); } + addControlNetToLinearGraph(graph, LATENTS_TO_LATENTS, state); + return graph; }; diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildTextToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildTextToImageGraph.ts index 161f857bd7..ae71f569b6 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildTextToImageGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildTextToImageGraph.ts @@ -1,8 +1,6 @@ import { RootState } from 'app/store/store'; import { - CollectInvocation, CompelInvocation, - ControlNetInvocation, Graph, IterateInvocation, LatentsToImageInvocation, @@ -12,7 +10,7 @@ import { TextToLatentsInvocation, } from 'services/api'; import { NonNullableGraph } from 'features/nodes/types/types'; -import { forEach, size } from 'lodash-es'; +import { addControlNetToLinearGraph } from '../addControlNetToLinearGraph'; const POSITIVE_CONDITIONING = 'positive_conditioning'; const NEGATIVE_CONDITIONING = 'negative_conditioning'; @@ -22,7 +20,6 @@ const NOISE = 'noise'; const RANDOM_INT = 'rand_int'; const RANGE_OF_SIZE = 'range_of_size'; const ITERATE = 'iterate'; -const CONTROL_NET_COLLECT = 'control_net_collect'; /** * Builds the Text to Image tab graph. @@ -42,8 +39,6 @@ export const buildTextToImageGraph = (state: RootState): Graph => { shouldRandomizeSeed, } = state.generation; - const { isEnabled: isControlNetEnabled, controlNets } = state.controlNet; - const graph: NonNullableGraph = { nodes: {}, edges: [], @@ -315,91 +310,7 @@ export const buildTextToImageGraph = (state: RootState): Graph => { }); } - // Add ControlNet - if (isControlNetEnabled) { - if (size(controlNets) > 1) { - const controlNetIterateNode: CollectInvocation = { - id: CONTROL_NET_COLLECT, - type: 'collect', - }; - graph.nodes[controlNetIterateNode.id] = controlNetIterateNode; - graph.edges.push({ - source: { node_id: controlNetIterateNode.id, field: 'collection' }, - destination: { - node_id: TEXT_TO_LATENTS, - field: 'control', - }, - }); - } - - forEach(controlNets, (controlNet, index) => { - const { - controlNetId, - isEnabled, - isControlImageProcessed, - controlImage, - processedControlImage, - beginStepPct, - endStepPct, - model, - processorNode, - weight, - } = controlNet; - - if (!isEnabled) { - // Skip disabled ControlNets - return; - } - - const controlNetNode: ControlNetInvocation = { - id: `control_net_${controlNetId}`, - type: 'controlnet', - begin_step_percent: beginStepPct, - end_step_percent: endStepPct, - control_model: model as ControlNetInvocation['control_model'], - control_weight: weight, - }; - - if (processedControlImage && !isControlImageProcessed) { - // We've already processed the image in the app, so we can just use the processed image - const { image_name, image_origin } = processedControlImage; - controlNetNode.image = { - image_name, - image_origin, - }; - } else if (controlImage && isControlImageProcessed) { - // The control image is preprocessed - const { image_name, image_origin } = controlImage; - controlNetNode.image = { - image_name, - image_origin, - }; - } else { - // Skip ControlNets without an unprocessed image - should never happen if everything is working correctly - return; - } - - graph.nodes[controlNetNode.id] = controlNetNode; - - if (size(controlNets) > 1) { - graph.edges.push({ - source: { node_id: controlNetNode.id, field: 'control' }, - destination: { - node_id: CONTROL_NET_COLLECT, - field: 'item', - }, - }); - } else { - graph.edges.push({ - source: { node_id: controlNetNode.id, field: 'control' }, - destination: { - node_id: TEXT_TO_LATENTS, - field: 'control', - }, - }); - } - }); - } + addControlNetToLinearGraph(graph, TEXT_TO_LATENTS, state); return graph; }; diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/ControlNet/ParamControlNetCollapse.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/ControlNet/ParamControlNetCollapse.tsx index 6369ab56de..06c6108dcb 100644 --- a/invokeai/frontend/web/src/features/parameters/components/Parameters/ControlNet/ParamControlNetCollapse.tsx +++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/ControlNet/ParamControlNetCollapse.tsx @@ -1,18 +1,7 @@ -import { - Divider, - Flex, - Tab, - TabList, - TabPanel, - TabPanels, - Tabs, -} from '@chakra-ui/react'; +import { Divider, Flex } from '@chakra-ui/react'; import { useTranslation } from 'react-i18next'; import IAICollapse from 'common/components/IAICollapse'; import { Fragment, memo, useCallback } from 'react'; -import IAIIconButton from 'common/components/IAIIconButton'; -import { FaPlus } from 'react-icons/fa'; -import ControlNet from 'features/controlNet/components/ControlNet'; import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; import { createSelector } from '@reduxjs/toolkit'; import { @@ -23,9 +12,9 @@ import { import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions'; import { map } from 'lodash-es'; import { v4 as uuidv4 } from 'uuid'; -import ControlNetMini from 'features/controlNet/components/ControlNetMini'; import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus'; import IAIButton from 'common/components/IAIButton'; +import ControlNet from 'features/controlNet/components/ControlNet'; const selector = createSelector( controlNetSelector, @@ -62,61 +51,19 @@ const ParamControlNetCollapse = () => { onToggle={handleClickControlNetToggle} withSwitch > - {controlNetsArray.length === 0 && ( - - Add ControlNet - - )} - + {controlNetsArray.map((c, i) => ( {i > 0 && } - + ))} + + Add ControlNet + ); - - return ( - - - - {controlNetsArray.map((c, i) => ( - - {i + 1} - - ))} - } - /> - - - {controlNetsArray.map((c) => ( - - - {/* */} - - ))} - - - - ); }; export default memo(ParamControlNetCollapse);