feat(ui): get processed images back into controlnet ui

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetImageProcessed.ts

@@ -1,15 +1,13 @@
 import { startAppListening } from '..';
-import { imageMetadataReceived, imageUploaded } from 'services/thunks/image';
+import { imageMetadataReceived } from 'services/thunks/image';
-import { addToast } from 'features/system/store/systemSlice';
 import { log } from 'app/logging/useLogger';
 import { controlNetImageProcessed } from 'features/controlNet/store/actions';
 import { Graph } from 'services/api';
 import { sessionCreated } from 'services/thunks/session';
 import { sessionReadyToInvoke } from 'features/system/store/actions';
-import { appSocketInvocationComplete } from 'services/events/actions';
+import { socketInvocationComplete } from 'services/events/actions';
 import { isImageOutput } from 'services/types/guards';
 import { controlNetProcessedImageChanged } from 'features/controlNet/store/controlNetSlice';
-import { selectImagesById } from 'features/gallery/store/imagesSlice';
 
 const moduleLog = log.child({ namespace: 'controlNet' });
 
@@ -18,27 +16,36 @@ export const addControlNetImageProcessedListener = () => {
     actionCreator: controlNetImageProcessed,
     effect: async (action, { dispatch, getState, take }) => {
       const { controlNetId, processorNode } = action.payload;
-      const { id } = processorNode;
+
+      // ControlNet one-off procressing graph is just he processor node, no edges
       const graph: Graph = {
-        nodes: { [id]: processorNode },
+        nodes: { [processorNode.id]: processorNode },
       };
+
+      // Create a session to run the graph & wait til it's ready to invoke
       const sessionCreatedAction = dispatch(sessionCreated({ graph }));
       const [sessionCreatedFulfilledAction] = await take(
         (action): action is ReturnType<typeof sessionCreated.fulfilled> =>
           sessionCreated.fulfilled.match(action) &&
           action.meta.requestId === sessionCreatedAction.requestId
       );
+
       const sessionId = sessionCreatedFulfilledAction.payload.id;
+
+      // Invoke the session & wait til it's complete
       dispatch(sessionReadyToInvoke());
-      const [processorAction] = await take(
+      const [invocationCompleteAction] = await take(
-        (action): action is ReturnType<typeof appSocketInvocationComplete> =>
+        (action): action is ReturnType<typeof socketInvocationComplete> =>
-          appSocketInvocationComplete.match(action) &&
+          socketInvocationComplete.match(action) &&
           action.payload.data.graph_execution_state_id === sessionId
       );
 
-      if (isImageOutput(processorAction.payload.data.result)) {
+      // We still have to check the output type
-        const { image_name } = processorAction.payload.data.result.image;
+      if (isImageOutput(invocationCompleteAction.payload.data.result)) {
+        const { image_name } =
+          invocationCompleteAction.payload.data.result.image;
 
+        // Wait for the ImageDTO to be received
         const [imageMetadataReceivedAction] = await take(
           (
             action
@@ -46,8 +53,9 @@ export const addControlNetImageProcessedListener = () => {
             imageMetadataReceived.fulfilled.match(action) &&
             action.payload.image_name === image_name
         );
-
         const processedControlImage = imageMetadataReceivedAction.payload;
+
+        // Update the processed image in the store
         dispatch(
           controlNetProcessedImageChanged({
             controlNetId,

invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx

@@ -1,4 +1,4 @@
-import { memo, useCallback } from 'react';
+import { memo, useCallback, useState } from 'react';
 import { ControlNetProcessorNode } from '../store/types';
 import { ImageDTO } from 'services/api';
 import CannyProcessor from './processors/CannyProcessor';
@@ -27,19 +27,10 @@ import { Flex, HStack, VStack } from '@chakra-ui/react';
 import IAISelectableImage from './parameters/IAISelectableImage';
 import IAIButton from 'common/components/IAIButton';
 import IAIIconButton from 'common/components/IAIIconButton';
-
+import IAISwitch from 'common/components/IAISwitch';
-export type ControlNetProcessorProps = {
+import ParamControlNetIsPreprocessed from './parameters/ParamControlNetIsPreprocessed';
-  controlNetId: string;
+import IAICollapse from 'common/components/IAICollapse';
-  image: ImageDTO;
+import ControlNetProcessorCollapse from './ControlNetProcessorCollapse';
-  type: ControlNetProcessorNode['type'];
-};
-
-const renderProcessorComponent = (props: ControlNetProcessorProps) => {
-  const { type } = props;
-  if (type === 'canny_image_processor') {
-    return <CannyProcessor {...props} />;
-  }
-};
 
 type ControlNetProps = {
   controlNet: ControlNet;
@@ -59,6 +50,10 @@ const ControlNet = (props: ControlNetProps) => {
   } = props.controlNet;
   const dispatch = useAppDispatch();
 
+  const [processorType, setProcessorType] = useState<
+    ControlNetProcessorNode['type']
+  >('canny_image_processor');
+
   const handleControlImageChanged = useCallback(
     (controlImage: ImageDTO) => {
       dispatch(controlNetImageChanged({ controlNetId, controlImage }));
@@ -74,14 +69,6 @@ const ControlNet = (props: ControlNetProps) => {
     dispatch(controlNetRemoved(controlNetId));
   }, [controlNetId, dispatch]);
 
-  const handleIsControlImageProcessedToggled = useCallback(() => {
-    dispatch(
-      isControlNetImageProcessedToggled({
-        controlNetId,
-      })
-    );
-  }, [controlNetId, dispatch]);
-
   const handleProcessedControlImageChanged = useCallback(
     (processedControlImage: ImageDTO | null) => {
       dispatch(
@@ -98,11 +85,15 @@ const ControlNet = (props: ControlNetProps) => {
     <Flex sx={{ flexDir: 'column', gap: 3, pb: 4 }}>
       <IAIButton onClick={handleControlNetRemoved}>Remove ControlNet</IAIButton>
       <IAISelectableImage
-        image={controlImage}
+        image={processedControlImage || controlImage}
         onChange={handleControlImageChanged}
         onReset={handleControlImageReset}
         resetIconSize="sm"
       />
+      <ControlNetProcessorCollapse
+        controlNetId={controlNetId}
+        image={controlImage}
+      />
       <ParamControlNetModel controlNetId={controlNetId} model={model} />
       <ParamControlNetIsEnabled
         controlNetId={controlNetId}

invokeai/frontend/web/src/features/controlNet/components/ControlNetProcessorCollapse.tsx

@@ -0,0 +1,67 @@
+import { useDisclosure } from '@chakra-ui/react';
+import IAICollapse from 'common/components/IAICollapse';
+import { memo, useState } from 'react';
+import CannyProcessor from './processors/CannyProcessor';
+import { ImageDTO } from 'services/api';
+import IAICustomSelect from 'common/components/IAICustomSelect';
+import {
+  CONTROLNET_PROCESSORS,
+  ControlNetProcessor,
+} from '../store/controlNetSlice';
+
+export type ControlNetProcessorProps = {
+  controlNetId: string;
+  image: ImageDTO;
+  type: ControlNetProcessor;
+};
+
+const ProcessorComponent = (props: ControlNetProcessorProps) => {
+  const { type } = props;
+  if (type === 'canny') {
+    return <CannyProcessor {...props} />;
+  }
+  return null;
+};
+
+type ControlNetProcessorCollapseProps = {
+  controlNetId: string;
+  image: ImageDTO | null;
+};
+
+const ControlNetProcessorCollapse = (
+  props: ControlNetProcessorCollapseProps
+) => {
+  const { image, controlNetId } = props;
+  const { isOpen, onToggle } = useDisclosure();
+
+  const [processorType, setProcessorType] =
+    useState<ControlNetProcessor>('canny');
+
+  const handleProcessorTypeChanged = (type: string | null | undefined) => {
+    setProcessorType(type as ControlNetProcessor);
+  };
+
+  return (
+    <IAICollapse
+      isOpen={Boolean(isOpen && image)}
+      onToggle={onToggle}
+      label="Process Image"
+      withSwitch
+    >
+      <IAICustomSelect
+        items={CONTROLNET_PROCESSORS}
+        selectedItem={processorType}
+        setSelectedItem={handleProcessorTypeChanged}
+      />
+      {image && (
+        <ProcessorComponent
+          controlNetId={controlNetId}
+          image={image}
+          type={processorType}
+        />
+      )}
+    </IAICollapse>
+  );
+};
+
+export default memo(ControlNetProcessorCollapse);

invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsPreprocessed.tsx

@@ -0,0 +1,36 @@
+import { useAppDispatch } from 'app/store/storeHooks';
+import IAIFullCheckbox from 'common/components/IAIFullCheckbox';
+import IAISwitch from 'common/components/IAISwitch';
+import {
+  controlNetToggled,
+  isControlNetImageProcessedToggled,
+} from 'features/controlNet/store/controlNetSlice';
+import { memo, useCallback } from 'react';
+
+type ParamControlNetIsEnabledProps = {
+  controlNetId: string;
+  isControlImageProcessed: boolean;
+};
+
+const ParamControlNetIsEnabled = (props: ParamControlNetIsEnabledProps) => {
+  const { controlNetId, isControlImageProcessed } = props;
+  const dispatch = useAppDispatch();
+
+  const handleIsControlImageProcessedToggled = useCallback(() => {
+    dispatch(
+      isControlNetImageProcessedToggled({
+        controlNetId,
+      })
+    );
+  }, [controlNetId, dispatch]);
+
+  return (
+    <IAISwitch
+      label="Preprocess"
+      isChecked={isControlImageProcessed}
+      onChange={handleIsControlImageProcessedToggled}
+    />
+  );
+};
+
+export default memo(ParamControlNetIsEnabled);

invokeai/frontend/web/src/features/controlNet/components/processors/CannyProcessor.tsx

@@ -4,9 +4,11 @@ import { memo, useCallback, useState } from 'react';
 import ControlNetProcessButton from './common/ControlNetProcessButton';
 import { useAppDispatch } from 'app/store/storeHooks';
 import { controlNetImageProcessed } from 'features/controlNet/store/actions';
-import { ControlNetProcessorProps } from '../ControlNet';
+import ControlNetResetProcessedImageButton from './common/ControlNetResetProcessedImageButton';
+import { ControlNetProcessorProps } from '../ControlNetProcessorCollapse';
+import { controlNetProcessedImageChanged } from 'features/controlNet/store/controlNetSlice';
 
-export const CANNY_PROCESSOR = 'canny_processor';
+export const CANNY_PROCESSOR = 'canny_image_processor';
 
 const CannyProcessor = (props: ControlNetProcessorProps) => {
   const { controlNetId, image, type } = props;
@@ -36,6 +38,15 @@ const CannyProcessor = (props: ControlNetProcessorProps) => {
     );
   }, [controlNetId, dispatch, highThreshold, image, lowThreshold]);
 
+  const handleReset = useCallback(() => {
+    dispatch(
+      controlNetProcessedImageChanged({
+        controlNetId,
+        processedControlImage: null,
+      })
+    );
+  }, [controlNetId, dispatch]);
+
   return (
     <Flex sx={{ flexDirection: 'column', gap: 2 }}>
       <IAISlider
@@ -54,7 +65,10 @@ const CannyProcessor = (props: ControlNetProcessorProps) => {
         max={255}
         withInput
       />
-      <ControlNetProcessButton onClick={handleProcess} />
+      <Flex sx={{ gap: 4 }}>
+        <ControlNetProcessButton onClick={handleProcess} />
+        <ControlNetResetProcessedImageButton onClick={handleReset} />
+      </Flex>
     </Flex>
   );
 };

invokeai/frontend/web/src/features/controlNet/components/processors/common/ControlNetResetProcessedImageButton.tsx

@@ -0,0 +1,20 @@
+import IAIButton from 'common/components/IAIButton';
+import { memo } from 'react';
+import { FaUnderline, FaUndo } from 'react-icons/fa';
+
+type ControlNetResetProcessedImageButtonProps = {
+  onClick: () => void;
+};
+
+const ControlNetResetProcessedImageButton = (
+  props: ControlNetResetProcessedImageButtonProps
+) => {
+  const { onClick } = props;
+  return (
+    <IAIButton leftIcon={<FaUndo />} onClick={onClick}>
+      Reset Processing
+    </IAIButton>
+  );
+};
+
+export default memo(ControlNetResetProcessedImageButton);

invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts

@@ -18,20 +18,22 @@ export const CONTROLNET_MODELS = [
   'lllyasviel/sd-controlnet-mlsd',
 ];
 
-// export const CONTROLNET_PROCESSORS = [
+export const CONTROLNET_PROCESSORS = [
-//   'canny',
+  'canny',
-//   'contentShuffle',
+  'contentShuffle',
-//   'hed',
+  'hed',
-//   'lineart',
+  'lineart',
-//   'lineartAnime',
+  'lineartAnime',
-//   'mediapipeFace',
+  'mediapipeFace',
-//   'midasDepth',
+  'midasDepth',
-//   'mlsd',
+  'mlsd',
-//   'normalBae',
+  'normalBae',
-//   'openpose',
+  'openpose',
-//   'pidi',
+  'pidi',
-//   'zoeDepth',
+  'zoeDepth',
-// ] as const;
+];
+
+export type ControlNetProcessor = (typeof CONTROLNET_PROCESSORS)[number];
 
 export type ControlNetModel = (typeof CONTROLNET_MODELS)[number];
 
@@ -109,6 +111,7 @@ export const controlNetSlice = createSlice({
     ) => {
       const { controlNetId, controlImage } = action.payload;
       state.controlNets[controlNetId].controlImage = controlImage;
+      state.controlNets[controlNetId].processedControlImage = null;
     },
     isControlNetImageProcessedToggled: (
       state,