feat(ui): prompt expansion (#8140)

* initializing prompt expansion and putting response in prompt box working for all methods

* properly disable UI and show loading state on prompt box when there is a pending prompt expansion item

* misc wrapup: disable apploying prompt templates, dont block textarea resize handle

* update progress to differentiate between prompt expansion and non

* cleanup

* lint

* more cleanup

* add image to background of loading state

* add allowPromptExpansion for front-end gating

* updated readiness text for needing to accept or discard

* fix tsc

* lint

* lint

* refactor(ui): prompt expansion logic

* tidy(ui): remove unnecessary changes

* revert(ui): unused arg on useImageUploadButton

* feat(ui): simplify prompt expansion state

* set pending for dragndrop and context menu

* add readiness logic for generate tab

* missing translation

* update error handling for prompt expansion

---------

Co-authored-by: Mary Hipp <maryhipp@Marys-Air.lan>
Co-authored-by: Mary Hipp <maryhipp@Marys-MacBook-Air.local>
Co-authored-by: psychedelicious <4822129+psychedelicious@users.noreply.github.com>

Author: 4 people

invokeai/frontend/web/public/locales/en.json

@@ -225,7 +225,16 @@
     "prompt": {
         "addPromptTrigger": "Add Prompt Trigger",
         "compatibleEmbeddings": "Compatible Embeddings",
-        "noMatchingTriggers": "No matching triggers"
+        "noMatchingTriggers": "No matching triggers",
+        "generateFromImage": "Generate prompt from image",
+        "expandCurrentPrompt": "Expand Current Prompt",
+        "uploadImageForPromptGeneration": "Upload Image for Prompt Generation",
+        "expandingPrompt": "Expanding prompt...",
+        "resultTitle": "Prompt Expansion Complete",
+        "resultSubtitle": "Choose how to handle the expanded prompt:",
+        "replace": "Replace",
+        "insert": "Insert",
+        "discard": "Discard"
     },
     "queue": {
         "queue": "Queue",
@@ -342,7 +351,7 @@
         "copy": "Copy",
         "currentlyInUse": "This image is currently in use in the following features:",
         "drop": "Drop",
-        "dropOrUpload": "$t(gallery.drop) or Upload",
+        "dropOrUpload": "Drop or Upload",
         "dropToUpload": "$t(gallery.drop) to Upload",
         "deleteImage_one": "Delete Image",
         "deleteImage_other": "Delete {{count}} Images",
@@ -396,7 +405,8 @@
         "compareHelp4": "Press <Kbd>Z</Kbd> or <Kbd>Esc</Kbd> to exit.",
         "openViewer": "Open Viewer",
         "closeViewer": "Close Viewer",
-        "move": "Move"
+        "move": "Move",
+        "useForPromptGeneration": "Use for Prompt Generation"
     },
     "hotkeys": {
         "hotkeys": "Hotkeys",
@@ -938,7 +948,8 @@
         "selectModel": "Select a Model",
         "noLoRAsInstalled": "No LoRAs installed",
         "noRefinerModelsInstalled": "No SDXL Refiner models installed",
-        "defaultVAE": "Default VAE"
+        "defaultVAE": "Default VAE",
+        "noCompatibleLoRAs": "No Compatible LoRAs"
     },
     "nodes": {
         "arithmeticSequence": "Arithmetic Sequence",
@@ -1188,7 +1199,9 @@
             "canvasIsSelectingObject": "Canvas is busy (selecting object)",
             "noPrompts": "No prompts generated",
             "noNodesInGraph": "No nodes in graph",
-            "systemDisconnected": "System disconnected"
+            "systemDisconnected": "System disconnected",
+            "promptExpansionPending": "Prompt expansion in progress",
+            "promptExpansionResultPending": "Please accept or discard your prompt expansion result"
         },
         "maskBlur": "Mask Blur",
         "negativePromptPlaceholder": "Negative Prompt",
@@ -1389,7 +1402,12 @@
         "fluxKontextIncompatibleGenerationMode": "Flux Kontext supports Text to Image only. Use other models for Image to Image, Inpainting and Outpainting tasks.",
         "problemUnpublishingWorkflow": "Problem Unpublishing Workflow",
         "problemUnpublishingWorkflowDescription": "There was a problem unpublishing the workflow. Please try again.",
-        "workflowUnpublished": "Workflow Unpublished"
+        "workflowUnpublished": "Workflow Unpublished",
+        "sentToCanvas": "Sent to Canvas",
+        "sentToUpscale": "Sent to Upscale",
+        "promptGenerationStarted": "Prompt generation started",
+        "uploadAndPromptGenerationFailed": "Failed to upload image and generate prompt",
+        "promptExpansionFailed": "Prompt expansion failed"
     },
     "popovers": {
         "clipSkip": {

invokeai/frontend/web/src/app/types/invokeai.ts

@@ -78,6 +78,7 @@ export type AppConfig = {
   allowPrivateStylePresets: boolean;
   allowClientSideUpload: boolean;
   allowPublishWorkflows: boolean;
+  allowPromptExpansion: boolean;
   disabledTabs: TabName[];
   disabledFeatures: AppFeature[];
   disabledSDFeatures: SDFeature[];

invokeai/frontend/web/src/common/hooks/useImageUploadButton.tsx

@@ -21,11 +21,15 @@ type UseImageUploadButtonArgs =
       isDisabled?: boolean;
       allowMultiple: false;
       onUpload?: (imageDTO: ImageDTO) => void;
+      onUploadStarted?: (files: File) => void;
+      onError?: (error: unknown) => void;
     }
   | {
       isDisabled?: boolean;
       allowMultiple: true;
       onUpload?: (imageDTOs: ImageDTO[]) => void;
+      onUploadStarted?: (files: File[]) => void;
+      onError?: (error: unknown) => void;
     };
 
 const log = logger('gallery');
@@ -49,7 +53,13 @@ const log = logger('gallery');
  * <Button {...getUploadButtonProps()} /> // will open the file dialog on click
  * <input {...getUploadInputProps()} /> // hidden, handles native upload functionality
  */
-export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: UseImageUploadButtonArgs) => {
+export const useImageUploadButton = ({
+  onUpload,
+  isDisabled,
+  allowMultiple,
+  onUploadStarted,
+  onError,
+}: UseImageUploadButtonArgs) => {
   const autoAddBoardId = useAppSelector(selectAutoAddBoardId);
   const isClientSideUploadEnabled = useAppSelector(selectIsClientSideUploadEnabled);
   const [uploadImage, request] = useUploadImageMutation();
@@ -71,6 +81,7 @@ export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: Us
           }
           const file = files[0];
           assert(file !== undefined); // should never happen
+          onUploadStarted?.(file);
           const imageDTO = await uploadImage({
             file,
             image_category: 'user',
@@ -82,6 +93,8 @@ export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: Us
             onUpload(imageDTO);
           }
         } else {
+          onUploadStarted?.(files);
+
           let imageDTOs: ImageDTO[] = [];
           if (isClientSideUploadEnabled && files.length > 1) {
             imageDTOs = await Promise.all(files.map((file, i) => clientSideUpload(file, i)));
@@ -102,14 +115,25 @@ export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: Us
           }
         }
       } catch (error) {
+        onError?.(error);
         toast({
           id: 'UPLOAD_FAILED',
           title: t('toast.imageUploadFailed'),
           status: 'error',
         });
       }
     },
-    [allowMultiple, autoAddBoardId, onUpload, uploadImage, isClientSideUploadEnabled, clientSideUpload, t]
+    [
+      allowMultiple,
+      onUploadStarted,
+      uploadImage,
+      autoAddBoardId,
+      onUpload,
+      isClientSideUploadEnabled,
+      clientSideUpload,
+      onError,
+      t,
+    ]
   );
 
   const onDropRejected = useCallback(

invokeai/frontend/web/src/features/controlLayers/konva/CanvasStateApiModule.ts

@@ -270,7 +270,7 @@ export class CanvasStateApiModule extends CanvasModuleBase {
     outputNodeId: string;
     options?: RunGraphOptions;
   }): Promise<ImageDTO> => {
-    const dependencies = buildRunGraphDependencies(this.store, this.manager.socket);
+    const dependencies = buildRunGraphDependencies(this.store.dispatch, this.manager.socket);
 
     const { output } = await runGraph({
       dependencies,

invokeai/frontend/web/src/features/controlLayers/store/util.ts

@@ -19,6 +19,7 @@ import type {
   RgbColor,
   T2IAdapterConfig,
 } from 'features/controlLayers/store/types';
+import type { ImageField } from 'features/nodes/types/common';
 import type { ImageDTO } from 'services/api/types';
 import { assert } from 'tsafe';
 import type { PartialDeep } from 'type-fest';
@@ -59,6 +60,8 @@ export const imageDTOToImageWithDims = ({ image_name, width, height }: ImageDTO)
   height,
 });
 
+export const imageDTOToImageField = ({ image_name }: ImageDTO): ImageField => ({ image_name });
+
 const DEFAULT_RG_MASK_FILL_COLORS: RgbColor[] = [
   { r: 121, g: 157, b: 219 }, // rgb(121, 157, 219)
   { r: 131, g: 214, b: 131 }, // rgb(131, 214, 131)

invokeai/frontend/web/src/features/dnd/dnd.ts

@@ -22,6 +22,8 @@ import {
 import { fieldImageCollectionValueChanged } from 'features/nodes/store/nodesSlice';
 import { selectFieldInputInstanceSafe, selectNodesSlice } from 'features/nodes/store/selectors';
 import { type FieldIdentifier, isImageFieldCollectionInputInstance } from 'features/nodes/types/field';
+import { expandPrompt } from 'features/prompt/PromptExpansion/expand';
+import { promptExpansionApi } from 'features/prompt/PromptExpansion/state';
 import type { ImageDTO } from 'services/api/types';
 import type { JsonObject } from 'type-fest';
 
@@ -515,23 +517,48 @@ export const removeImageFromBoardDndTarget: DndTarget<
 
 //#endregion
 
+//#region Prompt Generation From Image
+const _promptGenerationFromImage = buildTypeAndKey('prompt-generation-from-image');
+export type PromptGenerationFromImageDndTargetData = DndData<
+  typeof _promptGenerationFromImage.type,
+  typeof _promptGenerationFromImage.key,
+  void
+>;
+export const promptGenerationFromImageDndTarget: DndTarget<
+  PromptGenerationFromImageDndTargetData,
+  SingleImageDndSourceData
+> = {
+  ..._promptGenerationFromImage,
+  typeGuard: buildTypeGuard(_promptGenerationFromImage.key),
+  getData: buildGetData(_promptGenerationFromImage.key, _promptGenerationFromImage.type),
+  isValid: ({ sourceData }) => {
+    if (singleImageDndSource.typeGuard(sourceData)) {
+      return true;
+    }
+    return false;
+  },
+  handler: ({ sourceData, dispatch, getState }) => {
+    const { imageDTO } = sourceData.payload;
+    promptExpansionApi.setPending(imageDTO);
+    expandPrompt({ dispatch, getState, imageDTO });
+  },
+};
+//#endregion
+
 export const dndTargets = [
-  // Single Image
   setGlobalReferenceImageDndTarget,
+  addGlobalReferenceImageDndTarget,
   setRegionalGuidanceReferenceImageDndTarget,
   setUpscaleInitialImageDndTarget,
   setNodeImageFieldImageDndTarget,
+  addImagesToNodeImageFieldCollectionDndTarget,
   setComparisonImageDndTarget,
   newCanvasEntityFromImageDndTarget,
-  replaceCanvasEntityObjectsWithImageDndTarget,
-  addImageToBoardDndTarget,
-  removeImageFromBoardDndTarget,
   newCanvasFromImageDndTarget,
-  addGlobalReferenceImageDndTarget,
-  // Single or Multiple Image
+  replaceCanvasEntityObjectsWithImageDndTarget,
   addImageToBoardDndTarget,
   removeImageFromBoardDndTarget,
-  addImagesToNodeImageFieldCollectionDndTarget,
+  promptGenerationFromImageDndTarget,
 ] as const;
 
 export type AnyDndTarget = (typeof dndTargets)[number];

invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/ImageMenuItemUseForPromptGeneration.tsx

@@ -0,0 +1,46 @@
+import { MenuItem } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
+import { useAppSelector, useAppStore } from 'app/store/storeHooks';
+import { useImageDTOContext } from 'features/gallery/contexts/ImageDTOContext';
+import { expandPrompt } from 'features/prompt/PromptExpansion/expand';
+import { promptExpansionApi } from 'features/prompt/PromptExpansion/state';
+import { selectAllowPromptExpansion } from 'features/system/store/configSlice';
+import { toast } from 'features/toast/toast';
+import { memo, useCallback } from 'react';
+import { useTranslation } from 'react-i18next';
+import { PiTextTBold } from 'react-icons/pi';
+
+export const ImageMenuItemUseForPromptGeneration = memo(() => {
+  const { t } = useTranslation();
+  const { dispatch, getState } = useAppStore();
+  const imageDTO = useImageDTOContext();
+  const { isPending } = useStore(promptExpansionApi.$state);
+  const isPromptExpansionEnabled = useAppSelector(selectAllowPromptExpansion);
+
+  const handleUseForPromptGeneration = useCallback(() => {
+    promptExpansionApi.setPending(imageDTO);
+    expandPrompt({ dispatch, getState, imageDTO });
+    toast({
+      id: 'PROMPT_GENERATION_STARTED',
+      title: t('toast.promptGenerationStarted'),
+      status: 'info',
+    });
+  }, [dispatch, getState, imageDTO, t]);
+
+  if (!isPromptExpansionEnabled) {
+    return null;
+  }
+
+  return (
+    <MenuItem
+      icon={<PiTextTBold />}
+      onClickCapture={handleUseForPromptGeneration}
+      id="use-for-prompt-generation"
+      isDisabled={isPending}
+    >
+      {t('gallery.useForPromptGeneration')}
+    </MenuItem>
+  );
+});
+
+ImageMenuItemUseForPromptGeneration.displayName = 'ImageMenuItemUseForPromptGeneration';

invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/SingleSelectionMenuItems.tsx

@@ -14,6 +14,7 @@ import { ImageMenuItemSelectForCompare } from 'features/gallery/components/Image
 import { ImageMenuItemSendToUpscale } from 'features/gallery/components/ImageContextMenu/ImageMenuItemSendToUpscale';
 import { ImageMenuItemStarUnstar } from 'features/gallery/components/ImageContextMenu/ImageMenuItemStarUnstar';
 import { ImageMenuItemUseAsRefImage } from 'features/gallery/components/ImageContextMenu/ImageMenuItemUseAsRefImage';
+import { ImageMenuItemUseForPromptGeneration } from 'features/gallery/components/ImageContextMenu/ImageMenuItemUseForPromptGeneration';
 import { ImageDTOContextProvider } from 'features/gallery/contexts/ImageDTOContext';
 import { memo } from 'react';
 import type { ImageDTO } from 'services/api/types';
@@ -38,6 +39,7 @@ const SingleSelectionMenuItems = ({ imageDTO }: SingleSelectionMenuItemsProps) =
       <ImageMenuItemMetadataRecallActions />
       <MenuDivider />
       <ImageMenuItemSendToUpscale />
+      <ImageMenuItemUseForPromptGeneration />
       <ImageMenuItemUseAsRefImage />
       <ImageMenuItemNewCanvasFromImageSubMenu />
       <ImageMenuItemNewLayerFromImageSubMenu />