feat(ui): refine ref images UI

invokeai/frontend/web/src/features/controlLayers/components/RefImage/RefImage.tsx

@@ -1,5 +1,3 @@
-import type {
-  SystemStyleObject} from '@invoke-ai/ui-library';
 import {
   Divider,
   Flex,
@@ -10,7 +8,8 @@ import {
   PopoverArrow,
   PopoverBody,
   PopoverContent,
-  Portal
+  Portal,
+  Skeleton,
 } from '@invoke-ai/ui-library';
 import { skipToken } from '@reduxjs/toolkit/query';
 import { POPPER_MODIFIERS } from 'common/components/InformationalPopover/constants';
@@ -80,24 +79,12 @@ export const RefImage = memo(() => {
 });
 RefImage.displayName = 'RefImage';
 
-const imageSx: SystemStyleObject = {
-  opacity: 0.5,
-  _hover: {
-    opacity: 1,
-  },
-  "&[data-is-open='true']": {
-    opacity: 1,
-  },
-  transitionProperty: 'opacity',
-  transitionDuration: '0.2s',
-};
-
 const Thumbnail = memo(({ disclosure }: { disclosure: UseDisclosure }) => {
   const id = useRefImageIdContext();
   const entity = useRefImageEntity(id);
   const { data: imageDTO } = useGetImageDTOQuery(entity.config.image?.image_name ?? skipToken);
 
-  if (!imageDTO || !entity.config.image) {
+  if (!entity.config.image) {
     return (
       <PopoverAnchor>
         <IconButton
@@ -113,6 +100,7 @@ const Thumbnail = memo(({ disclosure }: { disclosure: UseDisclosure }) => {
           icon={<PiImageBold />}
           colorScheme="error"
           onClick={disclosure.toggle}
+          flexShrink={0}
         />
       </PopoverAnchor>
     );
@@ -120,14 +108,21 @@ const Thumbnail = memo(({ disclosure }: { disclosure: UseDisclosure }) => {
   return (
     <PopoverAnchor>
       <Image
+        borderWidth={1}
+        borderStyle="solid"
         id={getRefImagePopoverTriggerId(id)}
         role="button"
-        src={imageDTO.thumbnail_url}
+        src={imageDTO?.thumbnail_url}
         objectFit="contain"
+        aspectRatio="1/1"
+        // width={imageDTO?.width}
+        height={imageDTO?.height}
+        fallback={<Skeleton h="full" aspectRatio="1/1" />}
         maxW="full"
         maxH="full"
         borderRadius="base"
         onClick={disclosure.toggle}
+        flexShrink={0}
         // sx={imageSx}
         // data-is-open={disclosure.isOpen}
       />

invokeai/frontend/web/src/features/controlLayers/components/RefImage/RefImageList.tsx

@@ -1,17 +1,22 @@
 /* eslint-disable i18next/no-literal-string */
 import type { FlexProps } from '@invoke-ai/ui-library';
-import { Button, Flex, IconButton, Spacer } from '@invoke-ai/ui-library';
+import { Button, Flex } from '@invoke-ai/ui-library';
+import { useAppStore } from 'app/store/nanostores/store';
 import { useAppSelector } from 'app/store/storeHooks';
+import { useImageUploadButton } from 'common/hooks/useImageUploadButton';
 import { RefImage } from 'features/controlLayers/components/RefImage/RefImage';
 import { RefImageIdContext } from 'features/controlLayers/contexts/RefImageIdContext';
-import { useAddGlobalReferenceImage } from 'features/controlLayers/hooks/addLayerHooks';
-import { selectRefImageEntityIds } from 'features/controlLayers/store/refImagesSlice';
-import { memo } from 'react';
-import { PiPlusBold } from 'react-icons/pi';
+import { getDefaultRefImageConfig } from 'features/controlLayers/hooks/addLayerHooks';
+import { refImageAdded, selectRefImageEntityIds } from 'features/controlLayers/store/refImagesSlice';
+import { imageDTOToImageWithDims } from 'features/controlLayers/store/util';
+import { addGlobalReferenceImageDndTarget } from 'features/dnd/dnd';
+import { DndDropTarget } from 'features/dnd/DndDropTarget';
+import { memo, useMemo } from 'react';
+import { PiUploadBold } from 'react-icons/pi';
+import type { ImageDTO } from 'services/api/types';
 
 export const RefImageList = memo((props: FlexProps) => {
   const ids = useAppSelector(selectRefImageEntityIds);
-  const addRefImage = useAddGlobalReferenceImage();
   return (
     <Flex gap={2} h={16} {...props}>
       {ids.map((id) => (
@@ -19,59 +24,72 @@ export const RefImageList = memo((props: FlexProps) => {
           <RefImage />
         </RefImageIdContext.Provider>
       ))}
-      <Spacer />
-      <Button
-        size="sm"
-        variant="ghost"
-        h="full"
-        borderWidth="2px !important"
-        borderStyle="dashed !important"
-        borderRadius="base"
-        leftIcon={<PiPlusBold />}
-        onClick={addRefImage}
-        isDisabled={ids.length >= 5} // Limit to 5 reference images
-      >
-        Ref Image
-      </Button>
+      {ids.length < 5 && <AddRefImageDropTargetAndButton />}
+      {ids.length >= 5 && <MaxRefImages />}
     </Flex>
   );
 });
 
 RefImageList.displayName = 'RefImageList';
 
-const AddRefImageIconButton = memo(() => {
-  const addRefImage = useAddGlobalReferenceImage();
-  return (
-    <IconButton
-      aria-label="Add reference image"
-      h="full"
-      variant="ghost"
-      aspectRatio="1/1"
-      borderWidth={2}
-      borderStyle="dashed"
-      borderRadius="base"
-      onClick={addRefImage}
-      icon={<PiPlusBold />}
-    />
-  );
-});
-AddRefImageIconButton.displayName = 'AddRefImageIconButton';
+const dndTargetData = addGlobalReferenceImageDndTarget.getData();
 
-const AddRefImageButton = memo((props) => {
-  const addRefImage = useAddGlobalReferenceImage();
+const MaxRefImages = memo(() => {
   return (
     <Button
+      position="relative"
       size="sm"
       variant="ghost"
       h="full"
-      borderWidth={2}
-      borderStyle="dashed"
+      w="full"
+      borderWidth="2px !important"
+      borderStyle="dashed !important"
       borderRadius="base"
-      leftIcon={<PiPlusBold />}
-      onClick={addRefImage}
+      isDisabled
     >
-      Ref Image
+      Max Ref Images
     </Button>
   );
 });
-AddRefImageButton.displayName = 'AddRefImageButton';
+MaxRefImages.displayName = 'MaxRefImages';
+
+const AddRefImageDropTargetAndButton = memo(() => {
+  const { dispatch, getState } = useAppStore();
+
+  const uploadOptions = useMemo(
+    () =>
+      ({
+        onUpload: (imageDTO: ImageDTO) => {
+          const config = getDefaultRefImageConfig(getState);
+          config.image = imageDTOToImageWithDims(imageDTO);
+          dispatch(refImageAdded({ overrides: { config } }));
+        },
+        allowMultiple: false,
+      }) as const,
+    [dispatch, getState]
+  );
+
+  const uploadApi = useImageUploadButton(uploadOptions);
+
+  return (
+    <>
+      <Button
+        position="relative"
+        size="sm"
+        variant="ghost"
+        h="full"
+        w="full"
+        borderWidth="2px !important"
+        borderStyle="dashed !important"
+        borderRadius="base"
+        leftIcon={<PiUploadBold />}
+        {...uploadApi.getUploadButtonProps()}
+      >
+        Reference Image
+        <input {...uploadApi.getUploadInputProps()} />
+        <DndDropTarget label="Drop" dndTarget={addGlobalReferenceImageDndTarget} dndTargetData={dndTargetData} />
+      </Button>
+    </>
+  );
+});
+AddRefImageDropTargetAndButton.displayName = 'AddRefImageDropTargetAndButton';

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

@@ -1,7 +1,10 @@
 import { logger } from 'app/logging/logger';
 import type { AppDispatch, AppGetState } from 'app/store/store';
+import { getDefaultRefImageConfig } from 'features/controlLayers/hooks/addLayerHooks';
 import { getPrefixedId } from 'features/controlLayers/konva/util';
+import { refImageAdded } from 'features/controlLayers/store/refImagesSlice';
 import type { CanvasEntityIdentifier, CanvasEntityType } from 'features/controlLayers/store/types';
+import { imageDTOToImageWithDims } from 'features/controlLayers/store/util';
 import { selectComparisonImages } from 'features/gallery/components/ImageViewer/common';
 import type { BoardId } from 'features/gallery/store/types';
 import {
@@ -152,6 +155,34 @@ export const setGlobalReferenceImageDndTarget: DndTarget<
 };
 //#endregion
 
+//#region Add Global Reference Image
+const _addGlobalReferenceImage = buildTypeAndKey('add-global-reference-image');
+export type AddGlobalReferenceImageDndTargetData = DndData<
+  typeof _addGlobalReferenceImage.type,
+  typeof _addGlobalReferenceImage.key
+>;
+export const addGlobalReferenceImageDndTarget: DndTarget<
+  AddGlobalReferenceImageDndTargetData,
+  SingleImageDndSourceData
+> = {
+  ..._addGlobalReferenceImage,
+  typeGuard: buildTypeGuard(_addGlobalReferenceImage.key),
+  getData: buildGetData(_addGlobalReferenceImage.key, _addGlobalReferenceImage.type),
+  isValid: ({ sourceData }) => {
+    if (singleImageDndSource.typeGuard(sourceData)) {
+      return true;
+    }
+    return false;
+  },
+  handler: ({ sourceData, dispatch, getState }) => {
+    const { imageDTO } = sourceData.payload;
+    const config = getDefaultRefImageConfig(getState);
+    config.image = imageDTOToImageWithDims(imageDTO);
+    dispatch(refImageAdded({ overrides: { config } }));
+  },
+};
+//#endregion
+
 //#region Set Regional Guidance Reference Image
 const _setRegionalGuidanceReferenceImage = buildTypeAndKey('set-regional-guidance-reference-image');
 export type SetRegionalGuidanceReferenceImageDndTargetData = DndData<
@@ -496,6 +527,7 @@ export const dndTargets = [
   addImageToBoardDndTarget,
   removeImageFromBoardDndTarget,
   newCanvasFromImageDndTarget,
+  addGlobalReferenceImageDndTarget,
   // Single or Multiple Image
   addImageToBoardDndTarget,
   removeImageFromBoardDndTarget,