feat(ui): more tweaking controlnet ui

2024-08-30 20:32:17 +00:00 · 2023-06-01 22:52:25 +10:00 · 2023-06-01 22:52:25 +10:00 · b17f4c1650
commit b17f4c1650
parent 98493ed9e2
14 changed files with 271 additions and 148 deletions
--- a/invokeai/frontend/web/src/common/components/IAICustomSelect.tsx
+++ b/invokeai/frontend/web/src/common/components/IAICustomSelect.tsx
@ -104,7 +104,8 @@ const IAICustomSelect = (props: IAICustomSelectProps) => {
            ref={refs.setFloating}
            sx={{
              ...floatingStyles,
-              width: 'max-content',
+              width: 'full',
+              // width: 'max-content',
              top: 0,
              left: 0,
              flexDirection: 'column',
@ -118,6 +119,7 @@ const IAICustomSelect = (props: IAICustomSelectProps) => {
              px: 0,
              h: 'fit-content',
              maxH: 64,
+              minW: 48,
            }}
          >
            <OverlayScrollbarsComponent>
--- a/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/ControlNet.tsx
@ -3,8 +3,10 @@ import { ControlNetProcessorNode } from '../store/types';
 import { ImageDTO } from 'services/api';
 import CannyProcessor from './processors/CannyProcessor';
 import {
+  CONTROLNET_PROCESSORS,
  ControlNet,
  ControlNetModel,
+  ControlNetProcessor,
  controlNetBeginStepPctChanged,
  controlNetEndStepPctChanged,
  controlNetImageChanged,
@ -23,7 +25,18 @@ import ParamControlNetModel from './parameters/ParamControlNetModel';
 import ParamControlNetWeight from './parameters/ParamControlNetWeight';
 import ParamControlNetBeginStepPct from './parameters/ParamControlNetBeginStepPct';
 import ParamControlNetEndStepPct from './parameters/ParamControlNetEndStepPct';
-import { Flex, HStack, VStack } from '@chakra-ui/react';
+import {
+  Box,
+  Flex,
+  HStack,
+  Tab,
+  TabList,
+  TabPanel,
+  TabPanels,
+  Tabs,
+  VStack,
+  useDisclosure,
+} from '@chakra-ui/react';
 import IAISelectableImage from './parameters/IAISelectableImage';
 import IAIButton from 'common/components/IAIButton';
 import IAIIconButton from 'common/components/IAIIconButton';
@ -31,6 +44,7 @@ import IAISwitch from 'common/components/IAISwitch';
 import ParamControlNetIsPreprocessed from './parameters/ParamControlNetIsPreprocessed';
 import IAICollapse from 'common/components/IAICollapse';
 import ControlNetProcessorCollapse from './ControlNetProcessorCollapse';
+import IAICustomSelect from 'common/components/IAICustomSelect';

 type ControlNetProps = {
  controlNet: ControlNet;
@ -50,9 +64,14 @@ const ControlNet = (props: ControlNetProps) => {
  } = props.controlNet;
  const dispatch = useAppDispatch();

-  const [processorType, setProcessorType] = useState<
-    ControlNetProcessorNode['type']
-  >('canny_image_processor');
+  const [processorType, setProcessorType] =
+    useState<ControlNetProcessor>('canny');
+
+  const handleProcessorTypeChanged = (type: string | null | undefined) => {
+    setProcessorType(type as ControlNetProcessor);
+  };
+
+  const { isOpen, onToggle } = useDisclosure();

  const handleControlImageChanged = useCallback(
    (controlImage: ImageDTO) => {
@ -82,34 +101,82 @@ const ControlNet = (props: ControlNetProps) => {
  );

  return (
-    <Flex sx={{ flexDir: 'column', gap: 3, pb: 4 }}>
-      <IAIButton onClick={handleControlNetRemoved}>Remove ControlNet</IAIButton>
+    <Flex sx={{ flexDir: 'column', gap: 3 }}>
      <IAISelectableImage
        image={processedControlImage || controlImage}
        onChange={handleControlImageChanged}
        onReset={handleControlImageReset}
        resetIconSize="sm"
      />
-      <ControlNetProcessorCollapse
-        controlNetId={controlNetId}
-        image={controlImage}
-      />
      <ParamControlNetModel controlNetId={controlNetId} model={model} />
-      <ParamControlNetIsEnabled
-        controlNetId={controlNetId}
-        isEnabled={isEnabled}
-      />
-      <ParamControlNetWeight controlNetId={controlNetId} weight={weight} />
-      <ParamControlNetBeginStepPct
-        controlNetId={controlNetId}
-        beginStepPct={beginStepPct}
-      />
-      <ParamControlNetEndStepPct
-        controlNetId={controlNetId}
-        endStepPct={endStepPct}
-      />
+      <Tabs
+        isFitted
+        orientation="horizontal"
+        variant="enclosed"
+        size="sm"
+        colorScheme="accent"
+      >
+        <TabList>
+          <Tab
+            sx={{ 'button&': { _selected: { borderBottomColor: 'base.800' } } }}
+          >
+            Model Config
+          </Tab>
+          <Tab
+            sx={{ 'button&': { _selected: { borderBottomColor: 'base.800' } } }}
+          >
+            Preprocess
+          </Tab>
+        </TabList>
+        <TabPanels sx={{ pt: 2 }}>
+          <TabPanel sx={{ p: 0 }}>
+            <ParamControlNetWeight
+              controlNetId={controlNetId}
+              weight={weight}
+            />
+            <ParamControlNetBeginStepPct
+              controlNetId={controlNetId}
+              beginStepPct={beginStepPct}
+            />
+            <ParamControlNetEndStepPct
+              controlNetId={controlNetId}
+              endStepPct={endStepPct}
+            />
+          </TabPanel>
+          <TabPanel sx={{ p: 0 }}>
+            <IAICustomSelect
+              label="Processor"
+              items={CONTROLNET_PROCESSORS}
+              selectedItem={processorType}
+              setSelectedItem={handleProcessorTypeChanged}
+            />
+            <ProcessorComponent
+              controlNetId={controlNetId}
+              controlImage={controlImage}
+              processedControlImage={processedControlImage}
+              type={processorType}
+            />
+          </TabPanel>
+        </TabPanels>
+      </Tabs>
+      <IAIButton onClick={handleControlNetRemoved}>Remove ControlNet</IAIButton>
    </Flex>
  );
 };

 export default memo(ControlNet);
+
+export type ControlNetProcessorProps = {
+  controlNetId: string;
+  controlImage: ImageDTO | null;
+  processedControlImage: ImageDTO | null;
+  type: ControlNetProcessor;
+};
+
+const ProcessorComponent = (props: ControlNetProcessorProps) => {
+  const { type } = props;
+  if (type === 'canny') {
+    return <CannyProcessor {...props} />;
+  }
+  return null;
+};
--- a/invokeai/frontend/web/src/features/controlNet/components/ControlNetProcessorCollapse.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/ControlNetProcessorCollapse.tsx
@ -1,67 +1,76 @@
-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';
+// import { Collapse, Flex, useDisclosure } from '@chakra-ui/react';
+// 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';
+// import IAISwitch from 'common/components/IAISwitch';

-export type ControlNetProcessorProps = {
-  controlNetId: string;
-  image: ImageDTO;
-  type: ControlNetProcessor;
-};
+// export type ControlNetProcessorProps = {
+//   controlNetId: string;
+//   controlImage: ImageDTO | null;
+//   processedControlImage: ImageDTO | null;
+//   type: ControlNetProcessor;
+// };

-const ProcessorComponent = (props: ControlNetProcessorProps) => {
-  const { type } = props;
-  if (type === 'canny') {
-    return <CannyProcessor {...props} />;
-  }
-  return null;
-};
+// const ProcessorComponent = (props: ControlNetProcessorProps) => {
+//   const { type } = props;
+//   if (type === 'canny') {
+//     return <CannyProcessor {...props} />;
+//   }
+//   return null;
+// };

-type ControlNetProcessorCollapseProps = {
-  controlNetId: string;
-  image: ImageDTO | null;
-};
+// type ControlNetProcessorCollapseProps = {
+//   isOpen: boolean;
+//   controlNetId: string;
+//   controlImage: ImageDTO | null;
+//   processedControlImage: ImageDTO | null;
+// };
+// const ControlNetProcessorCollapse = (
+//   props: ControlNetProcessorCollapseProps
+// ) => {
+//   const { isOpen, controlImage, controlNetId, processedControlImage } = props;

-const ControlNetProcessorCollapse = (
-  props: ControlNetProcessorCollapseProps
-) => {
-  const { image, controlNetId } = props;
-  const { isOpen, onToggle } = useDisclosure();
+//   const [processorType, setProcessorType] =
+//     useState<ControlNetProcessor>('canny');

-  const [processorType, setProcessorType] =
-    useState<ControlNetProcessor>('canny');
+//   const handleProcessorTypeChanged = (type: string | null | undefined) => {
+//     setProcessorType(type as ControlNetProcessor);
+//   };

-  const handleProcessorTypeChanged = (type: string | null | undefined) => {
-    setProcessorType(type as ControlNetProcessor);
-  };
+//   return (
+//     <Flex
+//       sx={{
+//         gap: 2,
+//         p: 4,
+//         mt: 2,
+//         bg: 'base.850',
+//         borderRadius: 'base',
+//         flexDirection: 'column',
+//       }}
+//     >
+//       <IAICustomSelect
+//         label="Processor"
+//         items={CONTROLNET_PROCESSORS}
+//         selectedItem={processorType}
+//         setSelectedItem={handleProcessorTypeChanged}
+//       />
+//       {controlImage && (
+//         <ProcessorComponent
+//           controlNetId={controlNetId}
+//           controlImage={controlImage}
+//           processedControlImage={processedControlImage}
+//           type={processorType}
+//         />
+//       )}
+//     </Flex>
+//   );
+// };

-  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);

-export default memo(ControlNetProcessorCollapse);
+export default {};
--- a/invokeai/frontend/web/src/features/controlNet/components/parameters/IAISelectableImage.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/parameters/IAISelectableImage.tsx
@ -64,6 +64,7 @@ const IAISelectableImage = (props: IAISelectableImageProps) => {
            fallbackStrategy="beforeLoadOrError"
            fallback={<ImageFallback />}
            onError={onError}
+            draggable={false}
            sx={{
              borderRadius: 'base',
            }}
@ -80,6 +81,7 @@ const IAISelectableImage = (props: IAISelectableImageProps) => {
            >
              <IAIIconButton
                size={resetIconSize}
+                tooltip="Reset Image"
                aria-label="Reset Image"
                icon={<FaTimes />}
                onClick={onReset}
--- a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsEnabled.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetIsEnabled.tsx
@ -1,5 +1,5 @@
 import { useAppDispatch } from 'app/store/storeHooks';
-import IAIFullCheckbox from 'common/components/IAIFullCheckbox';
+import IAISwitch from 'common/components/IAISwitch';
 import { controlNetToggled } from 'features/controlNet/store/controlNetSlice';
 import { memo, useCallback } from 'react';

@ -17,7 +17,7 @@ const ParamControlNetIsEnabled = (props: ParamControlNetIsEnabledProps) => {
  }, [dispatch, controlNetId]);

  return (
-    <IAIFullCheckbox
+    <IAISwitch
      label="Enabled"
      isChecked={isEnabled}
      onChange={handleIsEnabledChanged}
--- a/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/parameters/ParamControlNetModel.tsx
@ -27,7 +27,6 @@ const ParamIsControlNetModel = (props: ParamIsControlNetModelProps) => {

  return (
    <IAICustomSelect
-      label="Model"
      items={CONTROLNET_MODELS}
      selectedItem={model}
      setSelectedItem={handleModelChanged}
--- a/invokeai/frontend/web/src/features/controlNet/components/processors/CannyProcessor.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/processors/CannyProcessor.tsx
@ -1,23 +1,22 @@
 import { Flex } from '@chakra-ui/react';
 import IAISlider from 'common/components/IAISlider';
-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 ControlNetResetProcessedImageButton from './common/ControlNetResetProcessedImageButton';
-import { ControlNetProcessorProps } from '../ControlNetProcessorCollapse';
 import { controlNetProcessedImageChanged } from 'features/controlNet/store/controlNetSlice';
+import ControlNetProcessorButtons from './common/ControlNetProcessorButtons';
+import { memo, useCallback, useState } from 'react';
+import { ControlNetProcessorProps } from '../ControlNet';

 export const CANNY_PROCESSOR = 'canny_image_processor';

 const CannyProcessor = (props: ControlNetProcessorProps) => {
-  const { controlNetId, image, type } = props;
+  const { controlNetId, controlImage, processedControlImage, type } = props;
  const dispatch = useAppDispatch();
  const [lowThreshold, setLowThreshold] = useState(100);
  const [highThreshold, setHighThreshold] = useState(200);

  const handleProcess = useCallback(() => {
-    if (!image) {
+    if (!controlImage) {
      return;
    }

@ -28,15 +27,15 @@ const CannyProcessor = (props: ControlNetProcessorProps) => {
          id: CANNY_PROCESSOR,
          type: 'canny_image_processor',
          image: {
-            image_name: image.image_name,
-            image_origin: image.image_origin,
+            image_name: controlImage.image_name,
+            image_origin: controlImage.image_origin,
          },
          low_threshold: lowThreshold,
          high_threshold: highThreshold,
        },
      })
    );
-  }, [controlNetId, dispatch, highThreshold, image, lowThreshold]);
+  }, [controlNetId, dispatch, highThreshold, controlImage, lowThreshold]);

  const handleReset = useCallback(() => {
    dispatch(
@ -65,10 +64,12 @@ const CannyProcessor = (props: ControlNetProcessorProps) => {
        max={255}
        withInput
      />
-      <Flex sx={{ gap: 4 }}>
-        <ControlNetProcessButton onClick={handleProcess} />
-        <ControlNetResetProcessedImageButton onClick={handleReset} />
-      </Flex>
+      <ControlNetProcessorButtons
+        handleProcess={handleProcess}
+        isProcessDisabled={Boolean(!controlImage)}
+        handleReset={handleReset}
+        isResetDisabled={Boolean(!processedControlImage)}
+      />
    </Flex>
  );
 };
--- a/invokeai/frontend/web/src/features/controlNet/components/processors/common/ControlNetProcessButton.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/processors/common/ControlNetProcessButton.tsx
@ -1,13 +0,0 @@
-import IAIButton from 'common/components/IAIButton';
-import { memo } from 'react';
-
-type ControlNetProcessButtonProps = {
-  onClick: () => void;
-};
-
-const ControlNetProcessButton = (props: ControlNetProcessButtonProps) => {
-  const { onClick } = props;
-  return <IAIButton onClick={onClick}>Process Control Image</IAIButton>;
-};
-
-export default memo(ControlNetProcessButton);
--- a/invokeai/frontend/web/src/features/controlNet/components/processors/common/ControlNetProcessorButtons.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/processors/common/ControlNetProcessorButtons.tsx
@ -0,0 +1,44 @@
+import { Flex } from '@chakra-ui/react';
+import { memo } from 'react';
+import { FaUndo } from 'react-icons/fa';
+import IAIButton from 'common/components/IAIButton';
+
+type ControlNetProcessorButtonsProps = {
+  handleProcess: () => void;
+  isProcessDisabled: boolean;
+  handleReset: () => void;
+  isResetDisabled: boolean;
+};
+
+const ControlNetProcessorButtons = (props: ControlNetProcessorButtonsProps) => {
+  const { handleProcess, isProcessDisabled, handleReset, isResetDisabled } =
+    props;
+  return (
+    <Flex
+      sx={{
+        gap: 4,
+        w: 'full',
+        alignItems: 'center',
+        justifyContent: 'stretch',
+      }}
+    >
+      <IAIButton
+        size="sm"
+        onClick={handleProcess}
+        isDisabled={isProcessDisabled}
+      >
+        Preprocess
+      </IAIButton>
+      <IAIButton
+        size="sm"
+        leftIcon={<FaUndo />}
+        onClick={handleReset}
+        isDisabled={isResetDisabled}
+      >
+        Reset Processing
+      </IAIButton>
+    </Flex>
+  );
+};
+
+export default memo(ControlNetProcessorButtons);
--- a/invokeai/frontend/web/src/features/controlNet/components/processors/common/ControlNetResetProcessedImageButton.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/processors/common/ControlNetResetProcessedImageButton.tsx
@ -1,20 +0,0 @@
-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);
--- a/invokeai/frontend/web/src/features/gallery/store/gallerySlice.ts
+++ b/invokeai/frontend/web/src/features/gallery/store/gallerySlice.ts
@ -50,7 +50,10 @@ export const gallerySlice = createSlice({
  },
  extraReducers: (builder) => {
    builder.addCase(imageUpserted, (state, action) => {
-      if (state.shouldAutoSwitchToNewImages) {
+      if (
+        state.shouldAutoSwitchToNewImages &&
+        action.payload.image_category === 'general'
+      ) {
        state.selectedImage = action.payload;
      }
    });
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildTextToImageGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildTextToImageGraph.ts
@ -19,6 +19,7 @@ const NOISE = 'noise';
 const RANDOM_INT = 'rand_int';
 const RANGE_OF_SIZE = 'range_of_size';
 const ITERATE = 'iterate';
+const CONTROL_NET = 'control_net';

 /**
 * Builds the Text to Image tab graph.
--- 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,4 +1,12 @@
-import { Flex, useDisclosure } from '@chakra-ui/react';
+import {
+  Flex,
+  Spacer,
+  Tab,
+  TabList,
+  TabPanel,
+  TabPanels,
+  Tabs,
+} from '@chakra-ui/react';
 import { useTranslation } from 'react-i18next';
 import IAICollapse from 'common/components/IAICollapse';
 import { memo, useCallback } from 'react';
@ -13,22 +21,23 @@ import {
  isControlNetEnabledToggled,
 } from 'features/controlNet/store/controlNetSlice';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
-import { map } from 'lodash-es';
+import { map, startCase } from 'lodash-es';
 import { v4 as uuidv4 } from 'uuid';
+import { CloseIcon } from '@chakra-ui/icons';

 const selector = createSelector(
  controlNetSelector,
  (controlNet) => {
    const { controlNets, isEnabled } = controlNet;

-    return { controlNets, isEnabled };
+    return { controlNetsArray: map(controlNets), isEnabled };
  },
  defaultSelectorOptions
 );

 const ParamControlNetCollapse = () => {
  const { t } = useTranslation();
-  const { controlNets, isEnabled } = useAppSelector(selector);
+  const { controlNetsArray, isEnabled } = useAppSelector(selector);
  const dispatch = useAppDispatch();

  const handleClickControlNetToggle = useCallback(() => {
@ -46,17 +55,35 @@ const ParamControlNetCollapse = () => {
      onToggle={handleClickControlNetToggle}
      withSwitch
    >
-      <Flex sx={{ alignItems: 'flex-end' }}>
-        <IAIIconButton
-          size="sm"
-          aria-label="Add ControlNet"
-          onClick={handleClickedAddControlNet}
-          icon={<FaPlus />}
-        />
-      </Flex>
-      {map(controlNets, (c) => (
-        <ControlNet key={c.controlNetId} controlNet={c} />
-      ))}
+      <Tabs
+        isFitted
+        orientation="horizontal"
+        variant="line"
+        size="sm"
+        colorScheme="accent"
+      >
+        <TabList alignItems="center" borderBottomColor="base.800" pb={4}>
+          {controlNetsArray.map((c, i) => (
+            <Tab key={`tab_${c.controlNetId}`} borderTopRadius="base">
+              {i + 1}
+            </Tab>
+          ))}
+          <IAIIconButton
+            marginInlineStart={2}
+            size="sm"
+            aria-label="Add ControlNet"
+            onClick={handleClickedAddControlNet}
+            icon={<FaPlus />}
+          />
+        </TabList>
+        <TabPanels>
+          {controlNetsArray.map((c) => (
+            <TabPanel key={`tabPanel_${c.controlNetId}`} sx={{ p: 0 }}>
+              <ControlNet controlNet={c} />
+            </TabPanel>
+          ))}
+        </TabPanels>
+      </Tabs>
    </IAICollapse>
  );
 };
--- a/invokeai/frontend/web/src/theme/components/tabs.ts
+++ b/invokeai/frontend/web/src/theme/components/tabs.ts
@ -26,6 +26,7 @@ const invokeAITablist = defineStyle((_props) => ({
    padding: 2,
    borderRadius: 'base',
    _selected: {
+      borderBottomColor: 'base.800',
      bg: 'accent.700',
      color: 'accent.100',
      _hover: {