feat(ui): rework image uploads (wip)

Author: psychedelicious

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

@@ -1,3 +1,4 @@
+import { Flex, Icon, type SystemStyleObject } from '@invoke-ai/ui-library';
 import { logger } from 'app/logging/logger';
 import { useAppSelector } from 'app/store/storeHooks';
 import { selectAutoAddBoardId } from 'features/gallery/store/gallerySelectors';
@@ -7,14 +8,22 @@ import { useCallback } from 'react';
 import type { FileRejection } from 'react-dropzone';
 import { useDropzone } from 'react-dropzone';
 import { useTranslation } from 'react-i18next';
-import { useUploadImageMutation } from 'services/api/endpoints/images';
-import type { PostUploadAction } from 'services/api/types';
+import { PiUploadSimpleBold } from 'react-icons/pi';
+import { uploadImages, useUploadImageMutation } from 'services/api/endpoints/images';
+import type { ImageDTO } from 'services/api/types';
+import { assert } from 'tsafe';
 
-type UseImageUploadButtonArgs = {
-  postUploadAction?: PostUploadAction;
-  isDisabled?: boolean;
-  allowMultiple?: boolean;
-};
+type UseImageUploadButtonArgs =
+  | {
+      isDisabled?: boolean;
+      allowMultiple: false;
+      onUpload?: (imageDTO: ImageDTO) => void;
+    }
+  | {
+      isDisabled?: boolean;
+      allowMultiple: true;
+      onUpload?: (imageDTOs: ImageDTO[]) => void;
+    };
 
 const log = logger('gallery');
 
@@ -37,30 +46,46 @@ const log = logger('gallery');
  * <Button {...getUploadButtonProps()} /> // will open the file dialog on click
  * <input {...getUploadInputProps()} /> // hidden, handles native upload functionality
  */
-export const useImageUploadButton = ({
-  postUploadAction,
-  isDisabled,
-  allowMultiple = false,
-}: UseImageUploadButtonArgs) => {
+export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: UseImageUploadButtonArgs) => {
   const autoAddBoardId = useAppSelector(selectAutoAddBoardId);
   const [uploadImage] = useUploadImageMutation();
   const maxImageUploadCount = useAppSelector(selectMaxImageUploadCount);
   const { t } = useTranslation();
 
   const onDropAccepted = useCallback(
-    (files: File[]) => {
-      for (const [i, file] of files.entries()) {
-        uploadImage({
+    async (files: File[]) => {
+      if (!allowMultiple) {
+        if (files.length > 1) {
+          log.warn('Multiple files dropped but only one allowed');
+          return;
+        }
+        const file = files[0];
+        assert(file !== undefined); // should never happen
+        const imageDTO = await uploadImage({
           file,
           image_category: 'user',
           is_intermediate: false,
-          postUploadAction: postUploadAction ?? { type: 'TOAST' },
           board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
-          isFirstUploadOfBatch: i === 0,
-        });
+        }).unwrap();
+        if (onUpload) {
+          onUpload(imageDTO);
+        }
+      } else {
+        //
+        const imageDTOs = await uploadImages(
+          files.map((file) => ({
+            file,
+            image_category: 'user',
+            is_intermediate: false,
+            board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
+          }))
+        );
+        if (onUpload) {
+          onUpload(imageDTOs);
+        }
       }
     },
-    [autoAddBoardId, postUploadAction, uploadImage]
+    [allowMultiple, autoAddBoardId, onUpload, uploadImage]
   );
 
   const onDropRejected = useCallback(
@@ -105,3 +130,34 @@ export const useImageUploadButton = ({
 
   return { getUploadButtonProps, getUploadInputProps, openUploader };
 };
+
+const sx = {
+  w: 'full',
+  h: 'full',
+  alignItems: 'center',
+  justifyContent: 'center',
+  borderRadius: 'base',
+  transitionProperty: 'common',
+  transitionDuration: '0.1s',
+  '&[data-disabled=false]': {
+    color: 'base.500',
+  },
+  '&[data-disabled=true]': {
+    cursor: 'pointer',
+    bg: 'base.700',
+    _hover: {
+      bg: 'base.650',
+      color: 'base.300',
+    },
+  },
+} satisfies SystemStyleObject;
+
+export const UploadImageButton = (props: UseImageUploadButtonArgs) => {
+  const uploadApi = useImageUploadButton(props);
+  return (
+    <Flex sx={sx} {...uploadApi.getUploadButtonProps()}>
+      <Icon as={PiUploadSimpleBold} boxSize={16} />
+      <input {...uploadApi.getUploadInputProps()} />
+    </Flex>
+  );
+};

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

@@ -1,6 +1,7 @@
 import { Flex, IconButton } from '@invoke-ai/ui-library';
 import { createMemoizedAppSelector } from 'app/store/createMemoizedSelector';
-import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { useAppStore } from 'app/store/nanostores/store';
+import { useAppSelector } from 'app/store/storeHooks';
 import { useImageUploadButton } from 'common/hooks/useImageUploadButton';
 import { BeginEndStepPct } from 'features/controlLayers/components/common/BeginEndStepPct';
 import { Weight } from 'features/controlLayers/components/common/Weight';
@@ -21,10 +22,11 @@ import { getFilterForModel } from 'features/controlLayers/store/filters';
 import { selectIsFLUX } from 'features/controlLayers/store/paramsSlice';
 import { selectCanvasSlice, selectEntityOrThrow } from 'features/controlLayers/store/selectors';
 import type { CanvasEntityIdentifier, ControlModeV2 } from 'features/controlLayers/store/types';
+import { replaceCanvasEntityObjectsWithImage } from 'features/imageActions/actions';
 import { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiBoundingBoxBold, PiShootingStarFill, PiUploadBold } from 'react-icons/pi';
-import type { ControlNetModelConfig, PostUploadAction, T2IAdapterModelConfig } from 'services/api/types';
+import type { ControlNetModelConfig, ImageDTO, T2IAdapterModelConfig } from 'services/api/types';
 
 const useControlLayerControlAdapter = (entityIdentifier: CanvasEntityIdentifier<'control_layer'>) => {
   const selectControlAdapter = useMemo(
@@ -41,7 +43,7 @@ const useControlLayerControlAdapter = (entityIdentifier: CanvasEntityIdentifier<
 
 export const ControlLayerControlAdapter = memo(() => {
   const { t } = useTranslation();
-  const dispatch = useAppDispatch();
+  const { dispatch, getState } = useAppStore();
   const entityIdentifier = useEntityIdentifierContext('control_layer');
   const controlAdapter = useControlLayerControlAdapter(entityIdentifier);
   const filter = useEntityFilter(entityIdentifier);
@@ -113,11 +115,17 @@ export const ControlLayerControlAdapter = memo(() => {
 
   const pullBboxIntoLayer = usePullBboxIntoLayer(entityIdentifier);
   const isBusy = useCanvasIsBusy();
-  const postUploadAction = useMemo<PostUploadAction>(
-    () => ({ type: 'REPLACE_LAYER_WITH_IMAGE', entityIdentifier }),
-    [entityIdentifier]
+  const uploadOptions = useMemo(
+    () =>
+      ({
+        onUpload: (imageDTO: ImageDTO) => {
+          replaceCanvasEntityObjectsWithImage({ entityIdentifier, imageDTO, dispatch, getState });
+        },
+        allowMultiple: false,
+      }) as const,
+    [dispatch, entityIdentifier, getState]
   );
-  const uploadApi = useImageUploadButton({ postUploadAction });
+  const uploadApi = useImageUploadButton(uploadOptions);
 
   return (
     <Flex flexDir="column" gap={3} position="relative" w="full">

invokeai/frontend/web/src/features/controlLayers/components/ControlLayer/ControlLayerControlAdapter.tsx

@@ -2,11 +2,9 @@ import type { SystemStyleObject } from '@invoke-ai/ui-library';
 import { Flex } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { skipToken } from '@reduxjs/toolkit/query';
+import { UploadImageButton } from 'common/hooks/useImageUploadButton';
 import type { ImageWithDims } from 'features/controlLayers/store/types';
-import type {
-  setGlobalReferenceImageDndTarget,
-  setRegionalGuidanceReferenceImageDndTarget,
-} from 'features/dnd/dnd';
+import type { setGlobalReferenceImageDndTarget, setRegionalGuidanceReferenceImageDndTarget } from 'features/dnd/dnd';
 import { DndDropTarget } from 'features/dnd/DndDropTarget';
 import { DndImage } from 'features/dnd/DndImage';
 import { DndImageIcon } from 'features/dnd/DndImageIcon';
@@ -58,8 +56,16 @@ export const IPAdapterImagePreview = memo(
       }
     }, [handleResetControlImage, isError, isConnected]);
 
+    const onUpload = useCallback(
+      (imageDTO: ImageDTO) => {
+        onChangeImage(imageDTO);
+      },
+      [onChangeImage]
+    );
+
     return (
       <Flex sx={sx} data-error={!imageDTO && !image?.image_name}>
+        {!imageDTO && <UploadImageButton allowMultiple={false} onUpload={onUpload} />}
         {imageDTO && (
           <>
             <DndImage imageDTO={imageDTO} />

invokeai/frontend/web/src/features/controlLayers/components/IPAdapter/IPAdapterImagePreview.tsx

@@ -157,12 +157,11 @@ export const DndDropTarget = memo(<T extends AnyDndTarget>(props: Props<T>) => {
               continue;
             }
             const imageDTO = await uploadImage({
-              blob: file,
-              fileName: file.name,
+              type: 'file',
+              file: file,
               image_category: 'user',
               is_intermediate: false,
             });
-            // Dnd.Util.handleDrop(Dnd.Source.singleImage.getData({ imageDTO }), targetData);
           }
         },
       }),

invokeai/frontend/web/src/features/dnd/DndDropTarget.tsx

@@ -3,13 +3,13 @@ import { useAppSelector } from 'app/store/storeHooks';
 import { useImageUploadButton } from 'common/hooks/useImageUploadButton';
 import { selectMaxImageUploadCount } from 'features/system/store/configSlice';
 import { t } from 'i18next';
+import { useMemo } from 'react';
 import { PiUploadBold } from 'react-icons/pi';
 
-const options = { postUploadAction: { type: 'TOAST' }, allowMultiple: true } as const;
-
 export const GalleryUploadButton = () => {
-  const uploadApi = useImageUploadButton(options);
   const maxImageUploadCount = useAppSelector(selectMaxImageUploadCount);
+  const uploadOptions = useMemo(() => ({ allowMultiple: maxImageUploadCount !== 1 }), [maxImageUploadCount]);
+  const uploadApi = useImageUploadButton(uploadOptions);
   return (
     <>
       <IconButton

invokeai/frontend/web/src/features/gallery/components/GalleryUploadButton.tsx

@@ -11,9 +11,9 @@ import type {
   ImageDTO,
   ListImagesArgs,
   ListImagesResponse,
-  PostUploadAction,
 } from 'services/api/types';
 import { getCategories, getListImagesUrl } from 'services/api/util';
+import type { JsonObject } from 'type-fest';
 
 import type { ApiTagDescription } from '..';
 import { api, buildV1Url, LIST_TAG } from '..';
@@ -267,7 +267,6 @@ export const imagesApi = api.injectEndpoints({
         file: File;
         image_category: ImageCategory;
         is_intermediate: boolean;
-        postUploadAction?: PostUploadAction;
         session_id?: string;
         board_id?: string;
         crop_visible?: boolean;
@@ -623,30 +622,60 @@ export const getImageMetadata = (
   return req.unwrap();
 };
 
-export type UploadOptions = {
-  blob: Blob;
-  fileName: string;
+export type UploadImageArg = {
+  file: File;
   image_category: ImageCategory;
   is_intermediate: boolean;
+  session_id?: string;
+  board_id?: string;
   crop_visible?: boolean;
-  board_id?: BoardId;
-  metadata?: SerializableObject;
+  metadata?: JsonObject;
 };
-export const uploadImage = (arg: UploadOptions): Promise<ImageDTO> => {
-  const { blob, fileName, image_category, is_intermediate, crop_visible = false, board_id, metadata } = arg;
+
+export const uploadImage = (arg: UploadImageArg): Promise<ImageDTO> => {
+  const { file, image_category, is_intermediate, crop_visible = false, board_id, metadata, session_id } = arg;
 
   const { dispatch } = getStore();
-  const file = new File([blob], fileName, { type: 'image/png' });
+
   const req = dispatch(
-    imagesApi.endpoints.uploadImage.initiate({
-      file,
-      image_category,
-      is_intermediate,
-      crop_visible,
-      board_id,
-      metadata,
-    })
+    imagesApi.endpoints.uploadImage.initiate(
+      {
+        file,
+        image_category,
+        is_intermediate,
+        crop_visible,
+        board_id,
+        metadata,
+        session_id,
+      },
+      { track: false }
+    )
   );
-  req.reset();
   return req.unwrap();
 };
+
+export const uploadImages = async (args: UploadImageArg[]): Promise<ImageDTO[]> => {
+  const { dispatch } = getStore();
+  const results = await Promise.allSettled(
+    args.map((arg, i) => {
+      const { file, image_category, is_intermediate, crop_visible = false, board_id, metadata, session_id } = arg;
+      const req = dispatch(
+        imagesApi.endpoints.uploadImage.initiate(
+          {
+            file,
+            image_category,
+            is_intermediate,
+            crop_visible,
+            board_id,
+            metadata,
+            session_id,
+            isFirstUploadOfBatch: i === 0,
+          },
+          { track: false }
+        )
+      );
+      return req.unwrap();
+    })
+  );
+  return results.filter((r): r is PromiseFulfilledResult<ImageDTO> => r.status === 'fulfilled').map((r) => r.value);
+};