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

This commit is contained in:
psychedelicious 2023-07-06 13:40:13 +10:00
parent 71dad6d404
commit 2415dc1235
14 changed files with 332 additions and 224 deletions

View File

@ -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) },

View File

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

View File

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

View File

@ -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';

View File

@ -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';

View File

@ -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(() => {

View File

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

View File

@ -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>
<IAITextarea
id="negativePrompt"
name="negativePrompt"
value={negativePrompt}
onChange={(e) => dispatch(setNegativePrompt(e.target.value))}
placeholder={t('parameters.negativePromptPlaceholder')}
fontSize="sm"
minH={16}
/>
<ParamEmbeddingPopover
isOpen={isOpen}
onClose={onClose}
onSelect={handleSelect}
>
<IAITextarea
id="negativePrompt"
name="negativePrompt"
ref={promptRef}
value={negativePrompt}
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>
);
};

View File

@ -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>) => {
dispatch(setPositivePrompt(e.target.value));
};
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>
<IAITextarea
id="prompt"
name="prompt"
placeholder={t('parameters.positivePromptPlaceholder')}
value={prompt}
onChange={handleChangePrompt}
onKeyDown={handleKeyDown}
resize="vertical"
ref={promptRef}
minH={32}
paddingRight={8}
/>
<ParamEmbeddingPopover
isOpen={isOpen}
onClose={onClose}
onSelect={handleSelect}
>
<IAITextarea
id="prompt"
name="prompt"
ref={promptRef}
value={prompt}
placeholder={t('parameters.positivePromptPlaceholder')}
onChange={handleChangePrompt}
onKeyDown={handleKeyDown}
resize="vertical"
minH={32}
/>
</ParamEmbeddingPopover>
</FormControl>
{!isOpen && (
<Box
sx={{
position: 'absolute',
top: 6,
insetInlineEnd: 0,
}}
>
<AddEmbeddingButton onClick={onOpen} />
</Box>
)}
</Box>
);
};

View File

@ -1,4 +1,3 @@
import { Tooltip } from '@chakra-ui/react';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import IAIIconButton, {
IAIIconButtonProps,
@ -25,26 +24,25 @@ const PinParametersPanelButton = (props: PinParametersPanelButtonProps) => {
};
return (
<Tooltip label={t('common.pinOptionsPanel')}>
<IAIIconButton
{...props}
aria-label={t('common.pinOptionsPanel')}
onClick={handleClickPinOptionsPanel}
icon={shouldPinParametersPanel ? <BsPinAngleFill /> : <BsPinAngle />}
variant="ghost"
size="sm"
sx={{
color: 'base.700',
_hover: {
color: 'base.550',
},
_active: {
color: 'base.500',
},
...sx,
}}
/>
</Tooltip>
<IAIIconButton
{...props}
tooltip={t('common.pinOptionsPanel')}
aria-label={t('common.pinOptionsPanel')}
onClick={handleClickPinOptionsPanel}
icon={shouldPinParametersPanel ? <BsPinAngleFill /> : <BsPinAngle />}
variant="ghost"
size="sm"
sx={{
color: 'base.700',
_hover: {
color: 'base.550',
},
_active: {
color: 'base.500',
},
...sx,
}}
/>
);
};

View File

@ -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 />

View File

@ -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 />

View File

@ -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 />