From 6c53abc034131055f22f2551546378345f8e77ca Mon Sep 17 00:00:00 2001 From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com> Date: Wed, 14 Jun 2023 20:01:17 +1200 Subject: [PATCH] feat: Add ControlMode to Linear UI --- invokeai/frontend/web/public/locales/en.json | 3 +- .../controlNet/components/ControlNet.tsx | 96 ++++++++++--------- .../parameters/ParamControlNetControlMode.tsx | 45 +++++++++ .../features/controlNet/store/constants.ts | 12 ++- .../controlNet/store/controlNetSlice.ts | 88 ++++++++++------- .../nodes/util/addControlNetToLinearGraph.ts | 2 + 6 files changed, 163 insertions(+), 83 deletions(-) create mode 100644 invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetControlMode.tsx diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 7a73bae411..fbd41e877d 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -524,7 +524,8 @@ "initialImage": "Initial Image", "showOptionsPanel": "Show Options Panel", "hidePreview": "Hide Preview", - "showPreview": "Show Preview" + "showPreview": "Show Preview", + "controlNetControlMode": "Control Mode" }, "settings": { "models": "Models", diff --git a/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx b/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx index 927daa9dc0..17cd5693e8 100644 --- a/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx +++ b/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx @@ -1,26 +1,27 @@ +import { Box, ChakraProps, Flex } from '@chakra-ui/react'; +import { useAppDispatch } from 'app/store/storeHooks'; import { memo, useCallback } from 'react'; +import { FaCopy, FaTrash } from 'react-icons/fa'; import { ControlNetConfig, controlNetAdded, controlNetRemoved, controlNetToggled, } from '../store/controlNetSlice'; -import { useAppDispatch } from 'app/store/storeHooks'; import ParamControlNetModel from './parameters/ParamControlNetModel'; import ParamControlNetWeight from './parameters/ParamControlNetWeight'; -import { Flex, Box, ChakraProps } from '@chakra-ui/react'; -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 IAISwitch from 'common/components/IAISwitch'; import { ChevronUpIcon } from '@chakra-ui/icons'; +import IAIIconButton from 'common/components/IAIIconButton'; +import IAISwitch from 'common/components/IAISwitch'; +import { useToggle } from 'react-use'; +import { v4 as uuidv4 } from 'uuid'; +import ControlNetImagePreview from './ControlNetImagePreview'; +import ControlNetProcessorComponent from './ControlNetProcessorComponent'; import ParamControlNetShouldAutoConfig from './ParamControlNetShouldAutoConfig'; +import ParamControlNetBeginEnd from './parameters/ParamControlNetBeginEnd'; +import ParamControlNetControlMode from './parameters/ParamControlNetControlMode'; +import ParamControlNetProcessorSelect from './parameters/ParamControlNetProcessorSelect'; const expandedControlImageSx: ChakraProps['sx'] = { maxH: 96 }; @@ -36,6 +37,7 @@ const ControlNet = (props: ControlNetProps) => { weight, beginStepPct, endStepPct, + controlMode, controlImage, processedControlImage, processorNode, @@ -137,45 +139,51 @@ const ControlNet = (props: ControlNetProps) => { {isEnabled && ( <> - - - - - - {!isExpanded && ( + + - + + - )} + {!isExpanded && ( + + + + )} + + + {isExpanded && ( <> diff --git a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetControlMode.tsx b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetControlMode.tsx new file mode 100644 index 0000000000..b8737004fd --- /dev/null +++ b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetControlMode.tsx @@ -0,0 +1,45 @@ +import { useAppDispatch } from 'app/store/storeHooks'; +import IAIMantineSelect from 'common/components/IAIMantineSelect'; +import { + ControlModes, + controlNetControlModeChanged, +} from 'features/controlNet/store/controlNetSlice'; +import { useCallback } from 'react'; +import { useTranslation } from 'react-i18next'; + +type ParamControlNetControlModeProps = { + controlNetId: string; + controlMode: string; +}; + +const CONTROL_MODE_DATA = [ + { label: 'Balanced', value: 'balanced' }, + { label: 'Prompt', value: 'more_prompt' }, + { label: 'Control', value: 'more_control' }, + { label: 'Mega Control', value: 'unbalanced' }, +]; + +export default function ParamControlNetControlMode( + props: ParamControlNetControlModeProps +) { + const { controlNetId, controlMode = false } = props; + const dispatch = useAppDispatch(); + + const { t } = useTranslation(); + + const handleControlModeChange = useCallback( + (controlMode: ControlModes) => { + dispatch(controlNetControlModeChanged({ controlNetId, controlMode })); + }, + [controlNetId, dispatch] + ); + + return ( + + ); +} diff --git a/invokeai/frontend/web/src/features/controlNet/store/constants.ts b/invokeai/frontend/web/src/features/controlNet/store/constants.ts index 7df0fda8e6..df6ee653dc 100644 --- a/invokeai/frontend/web/src/features/controlNet/store/constants.ts +++ b/invokeai/frontend/web/src/features/controlNet/store/constants.ts @@ -1,6 +1,5 @@ import { ControlNetProcessorType, - RequiredCannyImageProcessorInvocation, RequiredControlNetProcessorNode, } from './types'; @@ -23,7 +22,7 @@ type ControlNetProcessorsDict = Record< * * TODO: Generate from the OpenAPI schema */ -export const CONTROLNET_PROCESSORS = { +export const CONTROLNET_PROCESSORS: ControlNetProcessorsDict = { none: { type: 'none', label: 'none', @@ -174,6 +173,8 @@ export const CONTROLNET_PROCESSORS = { }, }; +type ControlNetModelsDict = Record; + type ControlNetModel = { type: string; label: string; @@ -181,7 +182,7 @@ type ControlNetModel = { defaultProcessor?: ControlNetProcessorType; }; -export const CONTROLNET_MODELS = { +export const CONTROLNET_MODELS: ControlNetModelsDict = { 'lllyasviel/control_v11p_sd15_canny': { type: 'lllyasviel/control_v11p_sd15_canny', label: 'Canny', @@ -190,6 +191,7 @@ export const CONTROLNET_MODELS = { 'lllyasviel/control_v11p_sd15_inpaint': { type: 'lllyasviel/control_v11p_sd15_inpaint', label: 'Inpaint', + defaultProcessor: 'none', }, 'lllyasviel/control_v11p_sd15_mlsd': { type: 'lllyasviel/control_v11p_sd15_mlsd', @@ -209,6 +211,7 @@ export const CONTROLNET_MODELS = { 'lllyasviel/control_v11p_sd15_seg': { type: 'lllyasviel/control_v11p_sd15_seg', label: 'Segmentation', + defaultProcessor: 'none', }, 'lllyasviel/control_v11p_sd15_lineart': { type: 'lllyasviel/control_v11p_sd15_lineart', @@ -223,6 +226,7 @@ export const CONTROLNET_MODELS = { 'lllyasviel/control_v11p_sd15_scribble': { type: 'lllyasviel/control_v11p_sd15_scribble', label: 'Scribble', + defaultProcessor: 'none', }, 'lllyasviel/control_v11p_sd15_softedge': { type: 'lllyasviel/control_v11p_sd15_softedge', @@ -242,10 +246,12 @@ export const CONTROLNET_MODELS = { 'lllyasviel/control_v11f1e_sd15_tile': { type: 'lllyasviel/control_v11f1e_sd15_tile', label: 'Tile (experimental)', + defaultProcessor: 'none', }, 'lllyasviel/control_v11e_sd15_ip2p': { type: 'lllyasviel/control_v11e_sd15_ip2p', label: 'Pix2Pix (experimental)', + defaultProcessor: 'none', }, 'CrucibleAI/ControlNetMediaPipeFace': { type: 'CrucibleAI/ControlNetMediaPipeFace', diff --git a/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts b/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts index d71ff4da68..a90c69028b 100644 --- a/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts +++ b/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts @@ -1,36 +1,27 @@ -import { PayloadAction } from '@reduxjs/toolkit'; -import { createSlice } from '@reduxjs/toolkit'; +import { PayloadAction, createSlice } from '@reduxjs/toolkit'; import { RootState } from 'app/store/store'; +import { forEach } from 'lodash-es'; import { ImageDTO } from 'services/api'; -import { - ControlNetProcessorType, - RequiredCannyImageProcessorInvocation, - RequiredControlNetProcessorNode, -} from './types'; +import { appSocketInvocationError } from 'services/events/actions'; +import { imageDeleted, imageUrlsReceived } from 'services/thunks/image'; +import { isAnySessionRejected } from 'services/thunks/session'; +import { controlNetImageProcessed } from './actions'; import { CONTROLNET_MODELS, CONTROLNET_PROCESSORS, ControlNetModelName, } from './constants'; -import { controlNetImageProcessed } from './actions'; -import { imageDeleted, imageUrlsReceived } from 'services/thunks/image'; -import { forEach } from 'lodash-es'; -import { isAnySessionRejected } from 'services/thunks/session'; -import { appSocketInvocationError } from 'services/events/actions'; +import { + ControlNetProcessorType, + RequiredCannyImageProcessorInvocation, + RequiredControlNetProcessorNode, +} from './types'; -export const initialControlNet: Omit = { - isEnabled: true, - model: CONTROLNET_MODELS['lllyasviel/control_v11p_sd15_canny'].type, - weight: 1, - beginStepPct: 0, - endStepPct: 1, - controlImage: null, - processedControlImage: null, - processorType: 'canny_image_processor', - processorNode: CONTROLNET_PROCESSORS.canny_image_processor - .default as RequiredCannyImageProcessorInvocation, - shouldAutoConfig: true, -}; +export type ControlModes = + | 'balanced' + | 'more_prompt' + | 'more_control' + | 'unbalanced'; export type ControlNetConfig = { controlNetId: string; @@ -39,6 +30,7 @@ export type ControlNetConfig = { weight: number; beginStepPct: number; endStepPct: number; + controlMode: ControlModes; controlImage: ImageDTO | null; processedControlImage: ImageDTO | null; processorType: ControlNetProcessorType; @@ -46,6 +38,21 @@ export type ControlNetConfig = { shouldAutoConfig: boolean; }; +export const initialControlNet: Omit = { + isEnabled: true, + model: CONTROLNET_MODELS['lllyasviel/control_v11p_sd15_canny'].type, + weight: 1, + beginStepPct: 0, + endStepPct: 1, + controlMode: 'balanced', + controlImage: null, + processedControlImage: null, + processorType: 'canny_image_processor', + processorNode: CONTROLNET_PROCESSORS.canny_image_processor + .default as RequiredCannyImageProcessorInvocation, + shouldAutoConfig: true, +}; + export type ControlNetState = { controlNets: Record; isEnabled: boolean; @@ -147,11 +154,13 @@ export const controlNetSlice = createSlice({ state.controlNets[controlNetId].processedControlImage = null; if (state.controlNets[controlNetId].shouldAutoConfig) { - const processorType = CONTROLNET_MODELS[model].defaultProcessor; + const processorType = + CONTROLNET_MODELS[model as keyof typeof CONTROLNET_MODELS] + .defaultProcessor; if (processorType) { state.controlNets[controlNetId].processorType = processorType; state.controlNets[controlNetId].processorNode = CONTROLNET_PROCESSORS[ - processorType + processorType as keyof typeof CONTROLNET_PROCESSORS ].default as RequiredControlNetProcessorNode; } else { state.controlNets[controlNetId].processorType = 'none'; @@ -181,6 +190,13 @@ export const controlNetSlice = createSlice({ const { controlNetId, endStepPct } = action.payload; state.controlNets[controlNetId].endStepPct = endStepPct; }, + controlNetControlModeChanged: ( + state, + action: PayloadAction<{ controlNetId: string; controlMode: ControlModes }> + ) => { + const { controlNetId, controlMode } = action.payload; + state.controlNets[controlNetId].controlMode = controlMode; + }, controlNetProcessorParamsChanged: ( state, action: PayloadAction<{ @@ -210,7 +226,7 @@ export const controlNetSlice = createSlice({ state.controlNets[controlNetId].processedControlImage = null; state.controlNets[controlNetId].processorType = processorType; state.controlNets[controlNetId].processorNode = CONTROLNET_PROCESSORS[ - processorType + processorType as keyof typeof CONTROLNET_PROCESSORS ].default as RequiredControlNetProcessorNode; state.controlNets[controlNetId].shouldAutoConfig = false; }, @@ -227,12 +243,14 @@ export const controlNetSlice = createSlice({ if (newShouldAutoConfig) { // manage the processor for the user const processorType = - CONTROLNET_MODELS[state.controlNets[controlNetId].model] - .defaultProcessor; + CONTROLNET_MODELS[ + state.controlNets[controlNetId] + .model as keyof typeof CONTROLNET_MODELS + ].defaultProcessor; if (processorType) { state.controlNets[controlNetId].processorType = processorType; state.controlNets[controlNetId].processorNode = CONTROLNET_PROCESSORS[ - processorType + processorType as keyof typeof CONTROLNET_PROCESSORS ].default as RequiredControlNetProcessorNode; } else { state.controlNets[controlNetId].processorType = 'none'; @@ -271,8 +289,7 @@ export const controlNetSlice = createSlice({ }); builder.addCase(imageUrlsReceived.fulfilled, (state, action) => { - const { image_name, image_origin, image_url, thumbnail_url } = - action.payload; + const { image_name, image_url, thumbnail_url } = action.payload; forEach(state.controlNets, (c) => { if (c.controlImage?.image_name === image_name) { @@ -286,11 +303,11 @@ export const controlNetSlice = createSlice({ }); }); - builder.addCase(appSocketInvocationError, (state, action) => { + builder.addCase(appSocketInvocationError, (state) => { state.pendingControlImages = []; }); - builder.addMatcher(isAnySessionRejected, (state, action) => { + builder.addMatcher(isAnySessionRejected, (state) => { state.pendingControlImages = []; }); }, @@ -308,6 +325,7 @@ export const { controlNetWeightChanged, controlNetBeginStepPctChanged, controlNetEndStepPctChanged, + controlNetControlModeChanged, controlNetProcessorParamsChanged, controlNetProcessorTypeChanged, controlNetReset, diff --git a/invokeai/frontend/web/src/features/nodes/util/addControlNetToLinearGraph.ts b/invokeai/frontend/web/src/features/nodes/util/addControlNetToLinearGraph.ts index 9aab13784f..a488961b71 100644 --- a/invokeai/frontend/web/src/features/nodes/util/addControlNetToLinearGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/addControlNetToLinearGraph.ts @@ -45,6 +45,7 @@ export const addControlNetToLinearGraph = ( processedControlImage, beginStepPct, endStepPct, + controlMode, model, processorType, weight, @@ -60,6 +61,7 @@ export const addControlNetToLinearGraph = ( type: 'controlnet', begin_step_percent: beginStepPct, end_step_percent: endStepPct, + control_mode: controlMode, control_model: model as ControlNetInvocation['control_model'], control_weight: weight, };