From 8c9266359d16c02cd1d02abcca95401f9037e251 Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Wed, 5 Jul 2023 23:41:15 +1200
Subject: [PATCH] feat: Add Embedding Select To Linear UI

---
 .../components/IAIMantineMultiSelect.tsx      |   2 +
 .../components/ParamEmbeddingCollapse.tsx     |  13 ++
 .../components/ParamEmbeddingSelect.tsx       | 128 ++++++++++++++++++
 .../ImageToImageTabParameters.tsx             |   2 +
 .../TextToImage/TextToImageTabParameters.tsx  |   2 +
 .../UnifiedCanvas/UnifiedCanvasParameters.tsx |   2 +
 6 files changed, 149 insertions(+)
 create mode 100644 invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingCollapse.tsx
 create mode 100644 invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingSelect.tsx

diff --git a/invokeai/frontend/web/src/common/components/IAIMantineMultiSelect.tsx b/invokeai/frontend/web/src/common/components/IAIMantineMultiSelect.tsx
index 97e33f300b..7622a604d6 100644
--- a/invokeai/frontend/web/src/common/components/IAIMantineMultiSelect.tsx
+++ b/invokeai/frontend/web/src/common/components/IAIMantineMultiSelect.tsx
@@ -49,6 +49,8 @@ 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) },
diff --git a/invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingCollapse.tsx b/invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingCollapse.tsx
new file mode 100644
index 0000000000..cb3a7bc7e0
--- /dev/null
+++ b/invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingCollapse.tsx
@@ -0,0 +1,13 @@
+import { Flex } from '@chakra-ui/react';
+import IAICollapse from 'common/components/IAICollapse';
+import ParamEmbeddingSelect from './ParamEmbeddingSelect';
+
+export default function ParamEmbeddingCollapse() {
+  return (
+    <IAICollapse label="Embeddings">
+      <Flex sx={{ flexDir: 'column', gap: 2 }}>
+        <ParamEmbeddingSelect />
+      </Flex>
+    </IAICollapse>
+  );
+}
diff --git a/invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingSelect.tsx b/invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingSelect.tsx
new file mode 100644
index 0000000000..5e273b43f3
--- /dev/null
+++ b/invokeai/frontend/web/src/features/embedding/components/ParamEmbeddingSelect.tsx
@@ -0,0 +1,128 @@
+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';
diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ImageToImage/ImageToImageTabParameters.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ImageToImage/ImageToImageTabParameters.tsx
index 32b71d6187..00edbe4706 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/ImageToImage/ImageToImageTabParameters.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/ImageToImage/ImageToImageTabParameters.tsx
@@ -1,4 +1,5 @@
 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,6 +19,7 @@ const ImageToImageTabParameters = () => {
       <ParamNegativeConditioning />
       <ProcessButtons />
       <ImageToImageTabCoreParameters />
+      <ParamEmbeddingCollapse />
       <ParamLoraCollapse />
       <ParamDynamicPromptsCollapse />
       <ParamControlNetCollapse />
diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/TextToImage/TextToImageTabParameters.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/TextToImage/TextToImageTabParameters.tsx
index 6291b69a8e..1853679625 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/TextToImage/TextToImageTabParameters.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/TextToImage/TextToImageTabParameters.tsx
@@ -1,4 +1,5 @@
 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';
@@ -19,6 +20,7 @@ const TextToImageTabParameters = () => {
       <ParamNegativeConditioning />
       <ProcessButtons />
       <TextToImageTabCoreParameters />
+      <ParamEmbeddingCollapse />
       <ParamLoraCollapse />
       <ParamDynamicPromptsCollapse />
       <ParamControlNetCollapse />
diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/UnifiedCanvas/UnifiedCanvasParameters.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/UnifiedCanvas/UnifiedCanvasParameters.tsx
index 63ed4cc1cf..bb6e7f2612 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/UnifiedCanvas/UnifiedCanvasParameters.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/UnifiedCanvas/UnifiedCanvasParameters.tsx
@@ -1,4 +1,5 @@
 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';
@@ -18,6 +19,7 @@ const UnifiedCanvasParameters = () => {
       <ParamNegativeConditioning />
       <ProcessButtons />
       <UnifiedCanvasCoreParameters />
+      <ParamEmbeddingCollapse />
       <ParamLoraCollapse />
       <ParamDynamicPromptsCollapse />
       <ParamControlNetCollapse />