adapt embedding popover to work for trigger phrases also

This commit is contained in:
Mary Hipp 2024-02-28 11:45:34 -05:00
parent ef474a3196
commit 6eb4c1ccb6
14 changed files with 183 additions and 175 deletions

View File

@ -304,6 +304,12 @@
"method": "High Resolution Fix Method"
}
},
"prompt": {
"addPromptTrigger": "Add Prompt Trigger",
"compatibleEmbeddings": "Compatible Embeddings",
"noPromptTriggers": "No triggers available",
"noMatchingTriggers": "No matching triggers"
},
"embedding": {
"addEmbedding": "Add Embedding",
"incompatibleModel": "Incompatible base model:",

View File

@ -1,21 +0,0 @@
import type { Meta, StoryObj } from '@storybook/react';
import { EmbeddingSelect } from './EmbeddingSelect';
import type { EmbeddingSelectProps } from './types';
const meta: Meta<typeof EmbeddingSelect> = {
title: 'Feature/Prompt/EmbeddingSelect',
tags: ['autodocs'],
component: EmbeddingSelect,
};
export default meta;
type Story = StoryObj<typeof EmbeddingSelect>;
const Component = (props: EmbeddingSelectProps) => {
return <EmbeddingSelect {...props}>Invoke</EmbeddingSelect>;
};
export const Default: Story = {
render: Component,
};

View File

@ -1,67 +0,0 @@
import type { ChakraProps } from '@invoke-ai/ui-library';
import { Combobox, FormControl } from '@invoke-ai/ui-library';
import { useAppSelector } from 'app/store/storeHooks';
import { useGroupedModelCombobox } from 'common/hooks/useGroupedModelCombobox';
import type { EmbeddingSelectProps } from 'features/embedding/types';
import { t } from 'i18next';
import { memo, useCallback } from 'react';
import { useTranslation } from 'react-i18next';
import { useGetTextualInversionModelsQuery } from 'services/api/endpoints/models';
import type { TextualInversionModelConfig } from 'services/api/types';
const noOptionsMessage = () => t('embedding.noMatchingEmbedding');
export const EmbeddingSelect = memo(({ onSelect, onClose }: EmbeddingSelectProps) => {
const { t } = useTranslation();
const currentBaseModel = useAppSelector((s) => s.generation.model?.base);
const getIsDisabled = useCallback(
(embedding: TextualInversionModelConfig): boolean => {
const isCompatible = currentBaseModel === embedding.base;
const hasMainModel = Boolean(currentBaseModel);
return !hasMainModel || !isCompatible;
},
[currentBaseModel]
);
const { data, isLoading } = useGetTextualInversionModelsQuery();
const _onChange = useCallback(
(embedding: TextualInversionModelConfig | null) => {
if (!embedding) {
return;
}
onSelect(embedding.name);
},
[onSelect]
);
const { options, onChange } = useGroupedModelCombobox({
modelEntities: data,
getIsDisabled,
onChange: _onChange,
});
return (
<FormControl>
<Combobox
placeholder={isLoading ? t('common.loading') : t('embedding.addEmbedding')}
defaultMenuIsOpen
autoFocus
value={null}
options={options}
noOptionsMessage={noOptionsMessage}
onChange={onChange}
onMenuClose={onClose}
data-testid="add-embedding"
sx={selectStyles}
/>
</FormControl>
);
});
EmbeddingSelect.displayName = 'EmbeddingSelect';
const selectStyles: ChakraProps['sx'] = {
w: 'full',
};

View File

@ -10,13 +10,11 @@ export const ModelMetadata = () => {
const { data: metadata } = useGetModelMetadataQuery(selectedModelKey ?? skipToken);
return (
<>
<Flex flexDir="column" height="full" gap="3">
<Box layerStyle="second" borderRadius="base" p={3}>
<TriggerPhrases />
</Box>
<DataViewer label="metadata" data={metadata || {}} />
</Flex>
</>
<Flex flexDir="column" height="full" gap="3">
<Box layerStyle="second" borderRadius="base" p={3}>
<TriggerPhrases />
</Box>
<DataViewer label="metadata" data={metadata || {}} />
</Flex>
);
};

View File

@ -1,8 +1,8 @@
import { Box, Textarea } from '@invoke-ai/ui-library';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { AddEmbeddingButton } from 'features/embedding/AddEmbeddingButton';
import { EmbeddingPopover } from 'features/embedding/EmbeddingPopover';
import { usePrompt } from 'features/embedding/usePrompt';
import { AddPromptTriggerButton } from 'features/prompt/AddPromptTriggerButton';
import { PromptPopover } from 'features/prompt/PromptPopover';
import { usePrompt } from 'features/prompt/usePrompt';
import { PromptOverlayButtonWrapper } from 'features/parameters/components/Prompts/PromptOverlayButtonWrapper';
import { setNegativePrompt } from 'features/parameters/store/generationSlice';
import { memo, useCallback, useRef } from 'react';
@ -19,19 +19,14 @@ export const ParamNegativePrompt = memo(() => {
},
[dispatch]
);
const { onChange, isOpen, onClose, onOpen, onSelectEmbedding, onKeyDown } = usePrompt({
const { onChange, isOpen, onClose, onOpen, onSelect, onKeyDown } = usePrompt({
prompt,
textareaRef,
onChange: _onChange,
});
return (
<EmbeddingPopover
isOpen={isOpen}
onClose={onClose}
onSelect={onSelectEmbedding}
width={textareaRef.current?.clientWidth}
>
<PromptPopover isOpen={isOpen} onClose={onClose} onSelect={onSelect} width={textareaRef.current?.clientWidth}>
<Box pos="relative">
<Textarea
id="negativePrompt"
@ -45,10 +40,10 @@ export const ParamNegativePrompt = memo(() => {
variant="darkFilled"
/>
<PromptOverlayButtonWrapper>
<AddEmbeddingButton isOpen={isOpen} onOpen={onOpen} />
<AddPromptTriggerButton isOpen={isOpen} onOpen={onOpen} />
</PromptOverlayButtonWrapper>
</Box>
</EmbeddingPopover>
</PromptPopover>
);
});

View File

@ -1,9 +1,9 @@
import { Box, Textarea } from '@invoke-ai/ui-library';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { ShowDynamicPromptsPreviewButton } from 'features/dynamicPrompts/components/ShowDynamicPromptsPreviewButton';
import { AddEmbeddingButton } from 'features/embedding/AddEmbeddingButton';
import { EmbeddingPopover } from 'features/embedding/EmbeddingPopover';
import { usePrompt } from 'features/embedding/usePrompt';
import { AddPromptTriggerButton } from 'features/prompt/AddPromptTriggerButton';
import { PromptPopover } from 'features/prompt/PromptPopover';
import { usePrompt } from 'features/prompt/usePrompt';
import { PromptOverlayButtonWrapper } from 'features/parameters/components/Prompts/PromptOverlayButtonWrapper';
import { setPositivePrompt } from 'features/parameters/store/generationSlice';
import { SDXLConcatButton } from 'features/sdxl/components/SDXLPrompts/SDXLConcatButton';
@ -25,7 +25,7 @@ export const ParamPositivePrompt = memo(() => {
},
[dispatch]
);
const { onChange, isOpen, onClose, onOpen, onSelectEmbedding, onKeyDown, onFocus } = usePrompt({
const { onChange, isOpen, onClose, onOpen, onSelect, onKeyDown, onFocus } = usePrompt({
prompt,
textareaRef: textareaRef,
onChange: handleChange,
@ -42,12 +42,7 @@ export const ParamPositivePrompt = memo(() => {
useHotkeys('alt+a', focus, []);
return (
<EmbeddingPopover
isOpen={isOpen}
onClose={onClose}
onSelect={onSelectEmbedding}
width={textareaRef.current?.clientWidth}
>
<PromptPopover isOpen={isOpen} onClose={onClose} onSelect={onSelect} width={textareaRef.current?.clientWidth}>
<Box pos="relative">
<Textarea
id="prompt"
@ -56,17 +51,17 @@ export const ParamPositivePrompt = memo(() => {
value={prompt}
placeholder={t('parameters.positivePromptPlaceholder')}
onChange={onChange}
minH={28}
minH={32}
onKeyDown={onKeyDown}
variant="darkFilled"
/>
<PromptOverlayButtonWrapper>
<AddEmbeddingButton isOpen={isOpen} onOpen={onOpen} />
<AddPromptTriggerButton isOpen={isOpen} onOpen={onOpen} />
{baseModel === 'sdxl' && <SDXLConcatButton />}
<ShowDynamicPromptsPreviewButton />
</PromptOverlayButtonWrapper>
</Box>
</EmbeddingPopover>
</PromptPopover>
);
});

View File

@ -8,15 +8,15 @@ type Props = {
onOpen: () => void;
};
export const AddEmbeddingButton = memo((props: Props) => {
export const AddPromptTriggerButton = memo((props: Props) => {
const { onOpen, isOpen } = props;
const { t } = useTranslation();
return (
<Tooltip label={t('embedding.addEmbedding')}>
<Tooltip label={t('prompt.addPromptTrigger')}>
<IconButton
variant="promptOverlay"
isDisabled={isOpen}
aria-label={t('embedding.addEmbedding')}
aria-label={t('prompt.addPromptTrigger')}
icon={<PiCodeBold />}
onClick={onOpen}
/>
@ -24,4 +24,4 @@ export const AddEmbeddingButton = memo((props: Props) => {
);
});
AddEmbeddingButton.displayName = 'AddEmbeddingButton';
AddPromptTriggerButton.displayName = 'AddPromptTriggerButton';

View File

@ -1,9 +1,9 @@
import { Popover, PopoverAnchor, PopoverBody, PopoverContent } from '@invoke-ai/ui-library';
import { EmbeddingSelect } from 'features/embedding/EmbeddingSelect';
import type { EmbeddingPopoverProps } from 'features/embedding/types';
import { PromptTriggerSelect } from 'features/prompt/PromptTriggerSelect';
import type { PromptPopoverProps } from 'features/prompt/types';
import { memo } from 'react';
export const EmbeddingPopover = memo((props: EmbeddingPopoverProps) => {
export const PromptPopover = memo((props: PromptPopoverProps) => {
const { onSelect, isOpen, onClose, width, children } = props;
return (
@ -14,7 +14,7 @@ export const EmbeddingPopover = memo((props: EmbeddingPopoverProps) => {
openDelay={0}
closeDelay={0}
closeOnBlur={true}
returnFocusOnClose={true}
returnFocusOnClose={false}
isLazy
>
<PopoverAnchor>{children}</PopoverAnchor>
@ -27,11 +27,11 @@ export const EmbeddingPopover = memo((props: EmbeddingPopoverProps) => {
borderStyle="solid"
>
<PopoverBody p={0} width={`calc(${width}px - 0.25rem)`}>
<EmbeddingSelect onClose={onClose} onSelect={onSelect} />
<PromptTriggerSelect onClose={onClose} onSelect={onSelect} />
</PopoverBody>
</PopoverContent>
</Popover>
);
});
EmbeddingPopover.displayName = 'EmbeddingPopover';
PromptPopover.displayName = 'PromptPopover';

View File

@ -0,0 +1,21 @@
import type { Meta, StoryObj } from '@storybook/react';
import { PromptTriggerSelect } from './PromptTriggerSelect';
import type { PromptTriggerSelectProps } from './types';
const meta: Meta<typeof PromptTriggerSelect> = {
title: 'Feature/Prompt/PromptTriggerSelect',
tags: ['autodocs'],
component: PromptTriggerSelect,
};
export default meta;
type Story = StoryObj<typeof PromptTriggerSelect>;
const Component = (props: PromptTriggerSelectProps) => {
return <PromptTriggerSelect {...props}>Invoke</PromptTriggerSelect>;
};
export const Default: Story = {
render: Component,
};

View File

@ -0,0 +1,87 @@
import type { ChakraProps, ComboboxOnChange } from '@invoke-ai/ui-library';
import { Combobox, FormControl } from '@invoke-ai/ui-library';
import { skipToken } from '@reduxjs/toolkit/query';
import { useAppSelector } from 'app/store/storeHooks';
import type { PromptTriggerSelectProps } from 'features/prompt/types';
import { t } from 'i18next';
import { map } from 'lodash-es';
import { useMemo } from 'react';
import { memo, useCallback } from 'react';
import { useTranslation } from 'react-i18next';
import { useGetModelMetadataQuery, useGetTextualInversionModelsQuery } from 'services/api/endpoints/models';
const noOptionsMessage = () => t('prompt.noMatchingTriggers');
export const PromptTriggerSelect = memo(({ onSelect, onClose }: PromptTriggerSelectProps) => {
const { t } = useTranslation();
const currentBaseModel = useAppSelector((s) => s.generation.model?.base);
const currentModelKey = useAppSelector((s) => s.generation.model?.key);
const { data, isLoading } = useGetTextualInversionModelsQuery();
const { data: metadata } = useGetModelMetadataQuery(currentModelKey ?? skipToken);
const _onChange = useCallback<ComboboxOnChange>(
(v) => {
if (!v) {
onSelect('');
return;
}
onSelect(v.value);
},
[onSelect]
);
const embeddingOptions = useMemo(() => {
if (!data) {
return [];
}
const compatibleEmbeddingsArray = map(data.entities).filter((model) => model.base === currentBaseModel);
return [
{
label: t('prompt.compatibleEmbeddings'),
options: compatibleEmbeddingsArray.map((model) => ({ label: model.name, value: `<${model.name}>` })),
},
];
}, [data, currentBaseModel]);
const options = useMemo(() => {
if (!metadata || !metadata.trigger_phrases) {
return [...embeddingOptions];
}
const metadataOptions = [
{
label: t('modelManager.triggerPhrases'),
options: metadata.trigger_phrases.map((phrase) => ({ label: phrase, value: phrase })),
},
];
return [...metadataOptions, ...embeddingOptions];
}, [embeddingOptions, metadata]);
return (
<FormControl>
<Combobox
placeholder={isLoading ? t('common.loading') : t('prompt.addPromptTrigger')}
defaultMenuIsOpen
autoFocus
value={null}
options={options}
noOptionsMessage={noOptionsMessage}
onChange={_onChange}
onMenuClose={onClose}
data-testid="add-prompt-trigger"
sx={selectStyles}
/>
</FormControl>
);
});
PromptTriggerSelect.displayName = 'PromptTriggerSelect';
const selectStyles: ChakraProps['sx'] = {
w: 'full',
};

View File

@ -1,12 +1,13 @@
import type { PropsWithChildren } from 'react';
export type EmbeddingSelectProps = {
export type PromptTriggerSelectProps = {
onSelect: (v: string) => void;
onClose: () => void;
};
export type EmbeddingPopoverProps = PropsWithChildren &
EmbeddingSelectProps & {
export type PromptPopoverProps = PropsWithChildren &
PromptTriggerSelectProps & {
isOpen: boolean;
width?: number | string;
};

View File

@ -4,13 +4,13 @@ import type { ChangeEventHandler, KeyboardEventHandler, RefObject } from 'react'
import { useCallback } from 'react';
import { flushSync } from 'react-dom';
type UseInsertEmbeddingArg = {
type UseInsertTriggerArg = {
prompt: string;
textareaRef: RefObject<HTMLTextAreaElement>;
onChange: (v: string) => void;
};
export const usePrompt = ({ prompt, textareaRef, onChange: _onChange }: UseInsertEmbeddingArg) => {
export const usePrompt = ({ prompt, textareaRef, onChange: _onChange }: UseInsertTriggerArg) => {
const { isOpen, onClose, onOpen } = useDisclosure();
const onChange: ChangeEventHandler<HTMLTextAreaElement> = useCallback(
@ -20,13 +20,14 @@ export const usePrompt = ({ prompt, textareaRef, onChange: _onChange }: UseInser
[_onChange]
);
const insertEmbedding = useCallback(
const insertTrigger = useCallback(
(v: string) => {
console.log({ textareaRef })
if (!textareaRef.current) {
return;
}
// this is where we insert the TI trigger
// this is where we insert the trigger
const caret = textareaRef.current.selectionStart;
if (isNil(caret)) {
@ -35,46 +36,48 @@ export const usePrompt = ({ prompt, textareaRef, onChange: _onChange }: UseInser
let newPrompt = prompt.slice(0, caret);
if (newPrompt[newPrompt.length - 1] !== '<') {
newPrompt += '<';
}
newPrompt += `${v}`;
newPrompt += `${v}>`;
// we insert the cursor after the `>`
// we insert the cursor after the end of trigger
const finalCaretPos = newPrompt.length;
newPrompt += prompt.slice(caret);
console.log({ newPrompt })
// must flush dom updates else selection gets reset
flushSync(() => {
_onChange(newPrompt);
});
// set the caret position to just after the TI trigger
// set the cursor position to just after the trigger
textareaRef.current.selectionStart = finalCaretPos;
textareaRef.current.selectionEnd = finalCaretPos;
},
[textareaRef, _onChange, prompt]
);
const onFocus = useCallback(() => {
console.log("focus")
console.log(textareaRef.current)
textareaRef.current?.focus();
}, [textareaRef]);
const handleClose = useCallback(() => {
const handleClosePopover = useCallback(() => {
onClose();
onFocus();
}, [onFocus, onClose]);
const onSelectEmbedding = useCallback(
const onSelect = useCallback(
(v: string) => {
insertEmbedding(v);
handleClose();
insertTrigger(v);
handleClosePopover();
},
[handleClose, insertEmbedding]
[handleClosePopover, insertTrigger]
);
const onKeyDown: KeyboardEventHandler<HTMLTextAreaElement> = useCallback(
(e) => {
if (e.key === '<') {
@ -90,7 +93,7 @@ export const usePrompt = ({ prompt, textareaRef, onChange: _onChange }: UseInser
isOpen,
onClose,
onOpen,
onSelectEmbedding,
onSelect,
onKeyDown,
onFocus,
};

View File

@ -1,8 +1,8 @@
import { Box, Textarea } from '@invoke-ai/ui-library';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { AddEmbeddingButton } from 'features/embedding/AddEmbeddingButton';
import { EmbeddingPopover } from 'features/embedding/EmbeddingPopover';
import { usePrompt } from 'features/embedding/usePrompt';
import { AddPromptTriggerButton } from 'features/prompt/AddPromptTriggerButton';
import { PromptPopover } from 'features/prompt/PromptPopover';
import { usePrompt } from 'features/prompt/usePrompt';
import { PromptOverlayButtonWrapper } from 'features/parameters/components/Prompts/PromptOverlayButtonWrapper';
import { setNegativeStylePromptSDXL } from 'features/sdxl/store/sdxlSlice';
import { memo, useCallback, useRef } from 'react';
@ -20,7 +20,7 @@ export const ParamSDXLNegativeStylePrompt = memo(() => {
},
[dispatch]
);
const { onChange, isOpen, onClose, onOpen, onSelectEmbedding, onKeyDown, onFocus } = usePrompt({
const { onChange, isOpen, onClose, onOpen, onSelect, onKeyDown, onFocus } = usePrompt({
prompt,
textareaRef: textareaRef,
onChange: handleChange,
@ -29,12 +29,7 @@ export const ParamSDXLNegativeStylePrompt = memo(() => {
useHotkeys('alt+a', onFocus, []);
return (
<EmbeddingPopover
isOpen={isOpen}
onClose={onClose}
onSelect={onSelectEmbedding}
width={textareaRef.current?.clientWidth}
>
<PromptPopover isOpen={isOpen} onClose={onClose} onSelect={onSelect} width={textareaRef.current?.clientWidth}>
<Box pos="relative">
<Textarea
id="prompt"
@ -48,10 +43,10 @@ export const ParamSDXLNegativeStylePrompt = memo(() => {
variant="darkFilled"
/>
<PromptOverlayButtonWrapper>
<AddEmbeddingButton isOpen={isOpen} onOpen={onOpen} />
<AddPromptTriggerButton isOpen={isOpen} onOpen={onOpen} />
</PromptOverlayButtonWrapper>
</Box>
</EmbeddingPopover>
</PromptPopover>
);
});

View File

@ -1,8 +1,8 @@
import { Box, Textarea } from '@invoke-ai/ui-library';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { AddEmbeddingButton } from 'features/embedding/AddEmbeddingButton';
import { EmbeddingPopover } from 'features/embedding/EmbeddingPopover';
import { usePrompt } from 'features/embedding/usePrompt';
import { AddPromptTriggerButton } from 'features/prompt/AddPromptTriggerButton';
import { PromptPopover } from 'features/prompt/PromptPopover';
import { usePrompt } from 'features/prompt/usePrompt';
import { PromptOverlayButtonWrapper } from 'features/parameters/components/Prompts/PromptOverlayButtonWrapper';
import { setPositiveStylePromptSDXL } from 'features/sdxl/store/sdxlSlice';
import { memo, useCallback, useRef } from 'react';
@ -19,19 +19,14 @@ export const ParamSDXLPositiveStylePrompt = memo(() => {
},
[dispatch]
);
const { onChange, isOpen, onClose, onOpen, onSelectEmbedding, onKeyDown } = usePrompt({
const { onChange, isOpen, onClose, onOpen, onSelect, onKeyDown } = usePrompt({
prompt,
textareaRef: textareaRef,
onChange: handleChange,
});
return (
<EmbeddingPopover
isOpen={isOpen}
onClose={onClose}
onSelect={onSelectEmbedding}
width={textareaRef.current?.clientWidth}
>
<PromptPopover isOpen={isOpen} onClose={onClose} onSelect={onSelect} width={textareaRef.current?.clientWidth}>
<Box pos="relative">
<Textarea
id="prompt"
@ -45,10 +40,10 @@ export const ParamSDXLPositiveStylePrompt = memo(() => {
variant="darkFilled"
/>
<PromptOverlayButtonWrapper>
<AddEmbeddingButton isOpen={isOpen} onOpen={onOpen} />
<AddPromptTriggerButton isOpen={isOpen} onOpen={onOpen} />
</PromptOverlayButtonWrapper>
</Box>
</EmbeddingPopover>
</PromptPopover>
);
});