feat(ui): refactor embedding ui; now is autocomplete

invokeai/frontend/web/src/common/components/IAIMantineMultiSelect.tsx

@@ -1,15 +1,16 @@
 import { Tooltip, useColorMode, useToken } from '@chakra-ui/react';
 import { MultiSelect, MultiSelectProps } from '@mantine/core';
 import { useChakraThemeTokens } from 'common/hooks/useChakraThemeTokens';
-import { memo } from 'react';
+import { RefObject, memo } from 'react';
 import { mode } from 'theme/util/mode';
 
 type IAIMultiSelectProps = MultiSelectProps & {
   tooltip?: string;
+  inputRef?: RefObject<HTMLInputElement>;
 };
 
 const IAIMantineMultiSelect = (props: IAIMultiSelectProps) => {
-  const { searchable = true, tooltip, ...rest } = props;
+  const { searchable = true, tooltip, inputRef, ...rest } = props;
   const {
     base50,
     base100,
@@ -33,6 +34,7 @@ const IAIMantineMultiSelect = (props: IAIMultiSelectProps) => {
   return (
     <Tooltip label={tooltip} placement="top" hasArrow>
       <MultiSelect
+        ref={inputRef}
         searchable={searchable}
         styles={() => ({
           label: {
@@ -49,8 +51,6 @@ const IAIMantineMultiSelect = (props: IAIMultiSelectProps) => {
             borderWidth: '2px',
             borderColor: mode(base200, base800)(colorMode),
             color: mode(base900, base100)(colorMode),
-            paddingTop: 6,
-            paddingBottom: 6,
             paddingRight: 24,
             fontWeight: 600,
             '&:hover': { borderColor: mode(base300, base600)(colorMode) },

invokeai/frontend/web/src/features/embedding/components/AddEmbeddingButton.tsx

@@ -0,0 +1,33 @@
+import IAIIconButton from 'common/components/IAIIconButton';
+import { memo } from 'react';
+import { BiCode } from 'react-icons/bi';
+
+type Props = {
+  onClick: () => void;
+};
+
+const AddEmbeddingButton = (props: Props) => {
+  const { onClick } = props;
+  return (
+    <IAIIconButton
+      size="sm"
+      aria-label="Add Embedding"
+      tooltip="Add Embedding"
+      icon={<BiCode />}
+      sx={{
+        p: 2,
+        color: 'base.700',
+        _hover: {
+          color: 'base.550',
+        },
+        _active: {
+          color: 'base.500',
+        },
+      }}
+      variant="link"
+      onClick={onClick}
+    />
+  );
+};
+
+export default memo(AddEmbeddingButton);

invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingCollapse.tsx

@@ -1,18 +0,0 @@
-import { Flex } from '@chakra-ui/react';
-import { RootState } from 'app/store/store';
-import { useAppSelector } from 'app/store/storeHooks';
-import ParamEmbeddingSelect from './ParamEmbeddingSelect';
-
-export default function ParamEmbeddingCollapse() {
-  const shouldShowEmbeddingPicker = useAppSelector(
-    (state: RootState) => state.ui.shouldShowEmbeddingPicker
-  );
-
-  return (
-    shouldShowEmbeddingPicker && (
-      <Flex sx={{ flexDir: 'column', gap: 2 }}>
-        <ParamEmbeddingSelect />
-      </Flex>
-    )
-  );
-}

invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingPopover.tsx

@@ -0,0 +1,129 @@
+import {
+  Popover,
+  PopoverBody,
+  PopoverContent,
+  PopoverTrigger,
+  Text,
+} from '@chakra-ui/react';
+import IAIMantineMultiSelect from 'common/components/IAIMantineMultiSelect';
+import { forEach } from 'lodash-es';
+import {
+  PropsWithChildren,
+  forwardRef,
+  useCallback,
+  useMemo,
+  useRef,
+} from 'react';
+import { useGetTextualInversionModelsQuery } from 'services/api/endpoints/models';
+import { PARAMETERS_PANEL_WIDTH } from 'theme/util/constants';
+
+type EmbeddingSelectItem = {
+  label: string;
+  value: string;
+  description?: string;
+};
+
+type Props = PropsWithChildren & {
+  onSelect: (v: string) => void;
+  isOpen: boolean;
+  onClose: () => void;
+};
+
+const ParamEmbeddingPopover = (props: Props) => {
+  const { onSelect, isOpen, onClose, children } = props;
+  const { data: embeddingQueryData } = useGetTextualInversionModelsQuery();
+  const inputRef = useRef<HTMLInputElement>(null);
+
+  const data = useMemo(() => {
+    if (!embeddingQueryData) {
+      return [];
+    }
+
+    const data: EmbeddingSelectItem[] = [];
+
+    forEach(embeddingQueryData.entities, (embedding, _) => {
+      if (!embedding) return;
+
+      data.push({
+        value: embedding.name,
+        label: embedding.name,
+        description: embedding.description,
+      });
+    });
+
+    return data;
+  }, [embeddingQueryData]);
+
+  const handleChange = useCallback(
+    (v: string[]) => {
+      if (v.length === 0) {
+        return;
+      }
+
+      onSelect(v[0]);
+    },
+    [onSelect]
+  );
+
+  return (
+    <Popover
+      initialFocusRef={inputRef}
+      returnFocusOnClose={true}
+      isOpen={isOpen}
+      onClose={onClose}
+      placement="bottom"
+      openDelay={0}
+      closeDelay={0}
+    >
+      <PopoverTrigger>{children}</PopoverTrigger>
+      <PopoverContent sx={{ p: 0, top: -1, shadow: 'dark-lg' }}>
+        <PopoverBody
+          sx={{ p: 1, w: `calc(${PARAMETERS_PANEL_WIDTH} - 2rem )` }}
+        >
+          <IAIMantineMultiSelect
+            inputRef={inputRef}
+            placeholder={data.length === 0 ? 'No Embeddings' : 'Add Embedding'}
+            value={[]}
+            data={data}
+            maxDropdownHeight={400}
+            nothingFound="No matching Embeddings"
+            itemComponent={SelectItem}
+            disabled={data.length === 0}
+            filter={(value, selected, item: EmbeddingSelectItem) =>
+              item.label.toLowerCase().includes(value.toLowerCase().trim()) ||
+              item.value.toLowerCase().includes(value.toLowerCase().trim())
+            }
+            onChange={handleChange}
+          />
+        </PopoverBody>
+      </PopoverContent>
+    </Popover>
+  );
+};
+
+export default ParamEmbeddingPopover;
+
+interface ItemProps extends React.ComponentPropsWithoutRef<'div'> {
+  value: string;
+  label: string;
+  description?: string;
+}
+
+const SelectItem = forwardRef<HTMLDivElement, ItemProps>(
+  ({ label, description, ...others }: ItemProps, ref) => {
+    return (
+      <div ref={ref} {...others}>
+        <div>
+          <Text>{label}</Text>
+          {description && (
+            <Text size="xs" color="base.600">
+              {description}
+            </Text>
+          )}
+        </div>
+      </div>
+    );
+  }
+);
+
+SelectItem.displayName = 'SelectItem';

invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingSelect.tsx

@@ -1,128 +0,0 @@
-import { Flex, Text } from '@chakra-ui/react';
-import { RootState } from 'app/store/store';
-import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import IAIButton from 'common/components/IAIButton';
-import IAIMantineMultiSelect from 'common/components/IAIMantineMultiSelect';
-import {
-  setNegativePrompt,
-  setPositivePrompt,
-} from 'features/parameters/store/generationSlice';
-import { forEach, join, map } from 'lodash-es';
-import { forwardRef, useMemo, useState } from 'react';
-import { useGetTextualInversionModelsQuery } from 'services/api/endpoints/models';
-
-type EmbeddingSelectItem = {
-  label: string;
-  value: string;
-  description?: string;
-};
-
-export default function ParamEmbeddingSelect() {
-  const { data: embeddingQueryData } = useGetTextualInversionModelsQuery();
-  const [selectedEmbeddings, setSelectedEmbeddings] = useState<
-    string[] | undefined
-  >(undefined);
-
-  const dispatch = useAppDispatch();
-
-  const positivePrompt = useAppSelector(
-    (state: RootState) => state.generation.positivePrompt
-  );
-
-  const negativePrompt = useAppSelector(
-    (state: RootState) => state.generation.negativePrompt
-  );
-
-  const data = useMemo(() => {
-    if (!embeddingQueryData) {
-      return [];
-    }
-
-    const data: EmbeddingSelectItem[] = [];
-
-    forEach(embeddingQueryData.entities, (embedding, _) => {
-      if (!embedding) return;
-
-      data.push({
-        value: embedding.name,
-        label: embedding.name,
-        description: embedding.description,
-      });
-    });
-
-    return data;
-  }, [embeddingQueryData]);
-
-  const handlePositiveAdd = () => {
-    if (!selectedEmbeddings) return;
-    const parsedEmbeddings = join(
-      map(selectedEmbeddings, (embedding) => `<${embedding}>`),
-      ' '
-    );
-    dispatch(setPositivePrompt(`${positivePrompt} ${parsedEmbeddings}`));
-    setSelectedEmbeddings([]);
-  };
-
-  const handleNegativeAdd = () => {
-    if (!selectedEmbeddings) return;
-    const parsedEmbeddings = join(
-      map(selectedEmbeddings, (embedding) => `<${embedding}>`),
-      ' '
-    );
-    dispatch(setNegativePrompt(`${negativePrompt} ${parsedEmbeddings}`));
-    setSelectedEmbeddings([]);
-  };
-
-  return (
-    <Flex gap={2} flexDirection="column">
-      <IAIMantineMultiSelect
-        placeholder="Pick Embedding"
-        value={selectedEmbeddings}
-        onChange={(v) => setSelectedEmbeddings(v)}
-        data={data}
-        maxDropdownHeight={400}
-        nothingFound="No matching Embeddings"
-        itemComponent={SelectItem}
-        disabled={data.length === 0}
-        filter={(value, selected, item: EmbeddingSelectItem) =>
-          item.label.toLowerCase().includes(value.toLowerCase().trim()) ||
-          item.value.toLowerCase().includes(value.toLowerCase().trim())
-        }
-        clearable
-      />
-      <Flex gap={2}>
-        <IAIButton size="sm" w="100%" onClick={handlePositiveAdd}>
-          Add To Positive
-        </IAIButton>
-        <IAIButton size="sm" w="100%" onClick={handleNegativeAdd}>
-          Add To Negative
-        </IAIButton>
-      </Flex>
-    </Flex>
-  );
-}
-
-interface ItemProps extends React.ComponentPropsWithoutRef<'div'> {
-  value: string;
-  label: string;
-  description?: string;
-}
-
-const SelectItem = forwardRef<HTMLDivElement, ItemProps>(
-  ({ label, description, ...others }: ItemProps, ref) => {
-    return (
-      <div ref={ref} {...others}>
-        <div>
-          <Text>{label}</Text>
-          {description && (
-            <Text size="xs" color="base.600">
-              {description}
-            </Text>
-          )}
-        </div>
-      </div>
-    );
-  }
-);
-
-SelectItem.displayName = 'SelectItem';

invokeai/frontend/web/src/features/lora/components/ParamLora.tsx

@@ -4,7 +4,12 @@ import IAIIconButton from 'common/components/IAIIconButton';
 import IAISlider from 'common/components/IAISlider';
 import { memo, useCallback } from 'react';
 import { FaTrash } from 'react-icons/fa';
-import { Lora, loraRemoved, loraWeightChanged } from '../store/loraSlice';
+import {
+  Lora,
+  loraRemoved,
+  loraWeightChanged,
+  loraWeightReset,
+} from '../store/loraSlice';
 
 type Props = {
   lora: Lora;
@@ -22,7 +27,7 @@ const ParamLora = (props: Props) => {
   );
 
   const handleReset = useCallback(() => {
-    dispatch(loraWeightChanged({ id: lora.id, weight: 0.75 }));
+    dispatch(loraWeightReset(lora.id));
   }, [dispatch, lora.id]);
 
   const handleRemoveLora = useCallback(() => {

invokeai/frontend/web/src/features/lora/store/loraSlice.ts

@@ -38,9 +38,14 @@ export const loraSlice = createSlice({
       const { id, weight } = action.payload;
       state.loras[id].weight = weight;
     },
+    loraWeightReset: (state, action: PayloadAction<string>) => {
+      const id = action.payload;
+      state.loras[id].weight = defaultLoRAConfig.weight;
+    },
   },
 });
 
-export const { loraAdded, loraRemoved, loraWeightChanged } = loraSlice.actions;
+export const { loraAdded, loraRemoved, loraWeightChanged, loraWeightReset } =
+  loraSlice.actions;
 
 export default loraSlice.reducer;

invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamNegativeConditioning.tsx

@@ -1,29 +1,90 @@
-import { FormControl } from '@chakra-ui/react';
+import { Box, FormControl, useDisclosure } from '@chakra-ui/react';
 import type { RootState } from 'app/store/store';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import IAITextarea from 'common/components/IAITextarea';
+import AddEmbeddingButton from 'features/embedding/components/AddEmbeddingButton';
+import ParamEmbeddingPopover from 'features/embedding/components/ParamEmbeddingPopover';
 import { setNegativePrompt } from 'features/parameters/store/generationSlice';
+import { ChangeEvent, KeyboardEvent, useCallback, useRef } from 'react';
 import { useTranslation } from 'react-i18next';
 
 const ParamNegativeConditioning = () => {
   const negativePrompt = useAppSelector(
     (state: RootState) => state.generation.negativePrompt
   );
-
+  const promptRef = useRef<HTMLTextAreaElement>(null);
+  const { isOpen, onClose, onOpen } = useDisclosure();
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
 
+  const handleChangePrompt = useCallback(
+    (e: ChangeEvent<HTMLTextAreaElement>) => {
+      dispatch(setNegativePrompt(e.target.value));
+    },
+    [dispatch]
+  );
+  const handleKeyDown = useCallback(
+    (e: KeyboardEvent<HTMLTextAreaElement>) => {
+      if (e.key === '<') {
+        onOpen();
+      }
+    },
+    [onOpen]
+  );
+
+  const handleSelect = useCallback(
+    (v: string) => {
+      const caret = promptRef.current?.selectionStart;
+
+      if (caret === undefined) {
+        return;
+      }
+
+      let newPrompt = negativePrompt.slice(0, caret);
+
+      if (newPrompt[newPrompt.length - 1] !== '<') {
+        newPrompt += '<';
+      }
+
+      newPrompt += `${v}>`;
+      newPrompt += negativePrompt.slice(caret);
+
+      dispatch(setNegativePrompt(newPrompt));
+    },
+    [dispatch, negativePrompt]
+  );
+
   return (
     <FormControl>
+      <ParamEmbeddingPopover
+        isOpen={isOpen}
+        onClose={onClose}
+        onSelect={handleSelect}
+      >
         <IAITextarea
           id="negativePrompt"
           name="negativePrompt"
+          ref={promptRef}
           value={negativePrompt}
-        onChange={(e) => dispatch(setNegativePrompt(e.target.value))}
           placeholder={t('parameters.negativePromptPlaceholder')}
+          onChange={handleChangePrompt}
+          onKeyDown={handleKeyDown}
+          resize="vertical"
           fontSize="sm"
           minH={16}
         />
+      </ParamEmbeddingPopover>
+      {!isOpen && (
+        <Box
+          sx={{
+            position: 'absolute',
+            top: 0,
+            insetInlineEnd: 0,
+          }}
+        >
+          <AddEmbeddingButton onClick={onOpen} />
+        </Box>
+      )}
     </FormControl>
   );
 };

invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamPositiveConditioning.tsx

@@ -1,4 +1,4 @@
-import { Box, FormControl } from '@chakra-ui/react';
+import { Box, FormControl, useDisclosure } from '@chakra-ui/react';
 import { RootState } from 'app/store/store';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { ChangeEvent, KeyboardEvent, useCallback, useRef } from 'react';
@@ -12,14 +12,13 @@ import {
 import { activeTabNameSelector } from 'features/ui/store/uiSelectors';
 
 import { userInvoked } from 'app/store/actions';
-import IAIIconButton from 'common/components/IAIIconButton';
 import IAITextarea from 'common/components/IAITextarea';
 import { useIsReadyToInvoke } from 'common/hooks/useIsReadyToInvoke';
-import { toggleEmbeddingPicker } from 'features/ui/store/uiSlice';
+import AddEmbeddingButton from 'features/embedding/components/AddEmbeddingButton';
+import ParamEmbeddingPopover from 'features/embedding/components/ParamEmbeddingPopover';
 import { isEqual } from 'lodash-es';
 import { useHotkeys } from 'react-hotkeys-hook';
 import { useTranslation } from 'react-i18next';
-import { BiCode } from 'react-icons/bi';
 
 const promptInputSelector = createSelector(
   [(state: RootState) => state.generation, activeTabNameSelector],
@@ -43,14 +42,16 @@ const ParamPositiveConditioning = () => {
   const dispatch = useAppDispatch();
   const { prompt, activeTabName } = useAppSelector(promptInputSelector);
   const isReady = useIsReadyToInvoke();
-
   const promptRef = useRef<HTMLTextAreaElement>(null);
-
+  const { isOpen, onClose, onOpen } = useDisclosure();
   const { t } = useTranslation();
 
-  const handleChangePrompt = (e: ChangeEvent<HTMLTextAreaElement>) => {
+  const handleChangePrompt = useCallback(
+    (e: ChangeEvent<HTMLTextAreaElement>) => {
       dispatch(setPositivePrompt(e.target.value));
-  };
+    },
+    [dispatch]
+  );
 
   useHotkeys(
     'alt+a',
@@ -67,39 +68,67 @@ const ParamPositiveConditioning = () => {
         dispatch(clampSymmetrySteps());
         dispatch(userInvoked(activeTabName));
       }
+      if (e.key === '<') {
+        onOpen();
+      }
     },
-    [dispatch, activeTabName, isReady]
+    [isReady, dispatch, activeTabName, onOpen]
   );
 
-  const shouldShowEmbeddingPicker = useAppSelector(
-    (state: RootState) => state.ui.shouldShowEmbeddingPicker
+  const handleSelect = useCallback(
+    (v: string) => {
+      const caret = promptRef.current?.selectionStart;
+
+      if (caret === undefined) {
+        return;
+      }
+
+      let newPrompt = prompt.slice(0, caret);
+
+      if (newPrompt[newPrompt.length - 1] !== '<') {
+        newPrompt += '<';
+      }
+
+      newPrompt += `${v}>`;
+      newPrompt += prompt.slice(caret);
+
+      dispatch(setPositivePrompt(newPrompt));
+    },
+    [dispatch, prompt]
   );
 
   return (
     <Box>
-      <IAIIconButton
-        size="xs"
-        aria-label="Toggle Embedding Picker"
-        tooltip="Toggle Embedding Picker"
-        icon={<BiCode />}
-        sx={{ position: 'absolute', top: 8, right: 2, zIndex: 2 }}
-        isChecked={shouldShowEmbeddingPicker}
-        onClick={() => dispatch(toggleEmbeddingPicker())}
-      ></IAIIconButton>
       <FormControl>
+        <ParamEmbeddingPopover
+          isOpen={isOpen}
+          onClose={onClose}
+          onSelect={handleSelect}
+        >
           <IAITextarea
             id="prompt"
             name="prompt"
-          placeholder={t('parameters.positivePromptPlaceholder')}
+            ref={promptRef}
             value={prompt}
+            placeholder={t('parameters.positivePromptPlaceholder')}
             onChange={handleChangePrompt}
             onKeyDown={handleKeyDown}
             resize="vertical"
-          ref={promptRef}
             minH={32}
-          paddingRight={8}
           />
+        </ParamEmbeddingPopover>
       </FormControl>
+      {!isOpen && (
+        <Box
+          sx={{
+            position: 'absolute',
+            top: 6,
+            insetInlineEnd: 0,
+          }}
+        >
+          <AddEmbeddingButton onClick={onOpen} />
+        </Box>
+      )}
     </Box>
   );
 };

invokeai/frontend/web/src/features/ui/components/PinParametersPanelButton.tsx

@@ -1,4 +1,3 @@
-import { Tooltip } from '@chakra-ui/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import IAIIconButton, {
   IAIIconButtonProps,
@@ -25,9 +24,9 @@ const PinParametersPanelButton = (props: PinParametersPanelButtonProps) => {
   };
 
   return (
-    <Tooltip label={t('common.pinOptionsPanel')}>
     <IAIIconButton
       {...props}
+      tooltip={t('common.pinOptionsPanel')}
       aria-label={t('common.pinOptionsPanel')}
       onClick={handleClickPinOptionsPanel}
       icon={shouldPinParametersPanel ? <BsPinAngleFill /> : <BsPinAngle />}
@@ -44,7 +43,6 @@ const PinParametersPanelButton = (props: PinParametersPanelButtonProps) => {
         ...sx,
       }}
     />
-    </Tooltip>
   );
 };
 

invokeai/frontend/web/src/features/ui/components/tabs/ImageToImage/ImageToImageTabParameters.tsx

@@ -1,5 +1,4 @@
 import ParamDynamicPromptsCollapse from 'features/dynamicPrompts/components/ParamDynamicPromptsCollapse';
-import ParamEmbeddingCollapse from 'features/embedding/components/ParamEmbeddingCollapse';
 import ParamLoraCollapse from 'features/lora/components/ParamLoraCollapse';
 import ParamControlNetCollapse from 'features/parameters/components/Parameters/ControlNet/ParamControlNetCollapse';
 import ParamNegativeConditioning from 'features/parameters/components/Parameters/Core/ParamNegativeConditioning';
@@ -17,7 +16,6 @@ const ImageToImageTabParameters = () => {
     <>
       <ParamPositiveConditioning />
       <ParamNegativeConditioning />
-      <ParamEmbeddingCollapse />
       <ProcessButtons />
       <ImageToImageTabCoreParameters />
       <ParamLoraCollapse />

invokeai/frontend/web/src/features/ui/components/tabs/TextToImage/TextToImageTabParameters.tsx

@@ -1,5 +1,4 @@
 import ParamDynamicPromptsCollapse from 'features/dynamicPrompts/components/ParamDynamicPromptsCollapse';
-import ParamEmbeddingCollapse from 'features/embedding/components/ParamEmbeddingCollapse';
 import ParamLoraCollapse from 'features/lora/components/ParamLoraCollapse';
 import ParamControlNetCollapse from 'features/parameters/components/Parameters/ControlNet/ParamControlNetCollapse';
 import ParamNegativeConditioning from 'features/parameters/components/Parameters/Core/ParamNegativeConditioning';
@@ -18,7 +17,6 @@ const TextToImageTabParameters = () => {
     <>
       <ParamPositiveConditioning />
       <ParamNegativeConditioning />
-      <ParamEmbeddingCollapse />
       <ProcessButtons />
       <TextToImageTabCoreParameters />
       <ParamLoraCollapse />

invokeai/frontend/web/src/features/ui/components/tabs/UnifiedCanvas/UnifiedCanvasParameters.tsx

@@ -1,5 +1,4 @@
 import ParamDynamicPromptsCollapse from 'features/dynamicPrompts/components/ParamDynamicPromptsCollapse';
-import ParamEmbeddingCollapse from 'features/embedding/components/ParamEmbeddingCollapse';
 import ParamLoraCollapse from 'features/lora/components/ParamLoraCollapse';
 import ParamInfillAndScalingCollapse from 'features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillAndScalingCollapse';
 import ParamSeamCorrectionCollapse from 'features/parameters/components/Parameters/Canvas/SeamCorrection/ParamSeamCorrectionCollapse';
@@ -17,7 +16,6 @@ const UnifiedCanvasParameters = () => {
     <>
       <ParamPositiveConditioning />
       <ParamNegativeConditioning />
-      <ParamEmbeddingCollapse />
       <ProcessButtons />
       <UnifiedCanvasCoreParameters />
       <ParamLoraCollapse />