feat: Add ControlMode to Linear UI

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",

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) => {
       </Flex>
       {isEnabled && (
         <>
-          <Flex sx={{ gap: 4, w: 'full' }}>
-            <Flex
-              sx={{
-                flexDir: 'column',
-                gap: 2,
-                w: 'full',
-                h: isExpanded ? 28 : 24,
-                paddingInlineStart: 1,
-                paddingInlineEnd: isExpanded ? 1 : 0,
-                pb: 2,
-                justifyContent: 'space-between',
-              }}
-            >
-              <ParamControlNetWeight
-                controlNetId={controlNetId}
-                weight={weight}
-                mini={!isExpanded}
-              />
-              <ParamControlNetBeginEnd
-                controlNetId={controlNetId}
-                beginStepPct={beginStepPct}
-                endStepPct={endStepPct}
-                mini={!isExpanded}
-              />
-            </Flex>
-            {!isExpanded && (
+          <Flex sx={{ w: 'full', flexDirection: 'column' }}>
+            <Flex sx={{ gap: 4, w: 'full' }}>
               <Flex
                 sx={{
-                  alignItems: 'center',
-                  justifyContent: 'center',
-                  h: 24,
-                  w: 24,
-                  aspectRatio: '1/1',
+                  flexDir: 'column',
+                  gap: 3,
+                  w: 'full',
+                  paddingInlineStart: 1,
+                  paddingInlineEnd: isExpanded ? 1 : 0,
+                  pb: 2,
+                  justifyContent: 'space-between',
                 }}
               >
-                <ControlNetImagePreview controlNet={props.controlNet} />
+                <ParamControlNetWeight
+                  controlNetId={controlNetId}
+                  weight={weight}
+                  mini={!isExpanded}
+                />
+                <ParamControlNetBeginEnd
+                  controlNetId={controlNetId}
+                  beginStepPct={beginStepPct}
+                  endStepPct={endStepPct}
+                  mini={!isExpanded}
+                />
               </Flex>
-            )}
+              {!isExpanded && (
+                <Flex
+                  sx={{
+                    alignItems: 'center',
+                    justifyContent: 'center',
+                    h: 24,
+                    w: 24,
+                    aspectRatio: '1/1',
+                  }}
+                >
+                  <ControlNetImagePreview controlNet={props.controlNet} />
+                </Flex>
+              )}
+            </Flex>
+            <ParamControlNetControlMode
+              controlNetId={controlNetId}
+              controlMode={controlMode}
+            />
           </Flex>
+
           {isExpanded && (
             <>
               <Box mt={2}>

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 (
+    <IAIMantineSelect
+      label={t('parameters.controlNetControlMode')}
+      data={CONTROL_MODE_DATA}
+      value={String(controlMode)}
+      onChange={handleControlModeChange}
+    />
+  );
+}

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<string, ControlNetModel>;
+
 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',

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<ControlNetConfig, 'controlNetId'> = {
-  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<ControlNetConfig, 'controlNetId'> = {
+  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<string, ControlNetConfig>;
   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,

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,
       };