fix(ui): fix init image display buttons

- Reset and Upload buttons along top of initial image
- Also had to mess around with the control net & DnD image stuff after changing the styles
- Abstract image upload logic into hook - does not handle native HTML drag and drop upload - only the button click upload

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

@@ -12,15 +12,14 @@ import IAIIconButton from 'common/components/IAIIconButton';
 import { IAIImageLoadingFallback } from 'common/components/IAIImageFallback';
 import ImageMetadataOverlay from 'common/components/ImageMetadataOverlay';
 import { AnimatePresence } from 'framer-motion';
-import { ReactElement, SyntheticEvent, useCallback } from 'react';
+import { ReactElement, SyntheticEvent } from 'react';
 import { memo, useRef } from 'react';
-import { FaImage, FaTimes, FaUndo, FaUpload } from 'react-icons/fa';
+import { FaImage, FaUndo, FaUpload } from 'react-icons/fa';
 import { ImageDTO } from 'services/api/types';
 import { v4 as uuidv4 } from 'uuid';
 import IAIDropOverlay from './IAIDropOverlay';
-import { PostUploadAction, imageUploaded } from 'services/api/thunks/image';
+import { PostUploadAction } from 'services/api/thunks/image';
-import { useDropzone } from 'react-dropzone';
+import { useImageUploadButton } from 'common/hooks/useImageUploadButton';
-import { useAppDispatch } from 'app/store/storeHooks';
 
 type IAIDndImageProps = {
   image: ImageDTO | null | undefined;
@@ -39,6 +38,7 @@ type IAIDndImageProps = {
   minSize?: number;
   postUploadAction?: PostUploadAction;
   imageSx?: ChakraProps['sx'];
+  fitContainer?: boolean;
 };
 
 const IAIDndImage = (props: IAIDndImageProps) => {
@@ -58,8 +58,9 @@ const IAIDndImage = (props: IAIDndImageProps) => {
     minSize = 24,
     postUploadAction,
     imageSx,
+    fitContainer = false,
   } = props;
-  const dispatch = useAppDispatch();
+
   const dndId = useRef(uuidv4());
 
   const {
@@ -87,31 +88,9 @@ const IAIDndImage = (props: IAIDndImageProps) => {
     disabled: isDragDisabled || !image,
   });
 
-  const handleOnDropAccepted = useCallback(
+  const { getUploadButtonProps, getUploadInputProps } = useImageUploadButton({
-    (files: Array<File>) => {
+    postUploadAction,
-      const file = files[0];
+    isDisabled: isUploadDisabled,
-      if (!file) {
-        return;
-      }
-
-      dispatch(
-        imageUploaded({
-          file,
-          image_category: 'user',
-          is_intermediate: false,
-          postUploadAction,
-        })
-      );
-    },
-    [dispatch, postUploadAction]
-  );
-
-  const { getRootProps, getInputProps } = useDropzone({
-    accept: { 'image/png': ['.png'], 'image/jpeg': ['.jpg', '.jpeg', '.png'] },
-    onDropAccepted: handleOnDropAccepted,
-    noDrag: true,
-    multiple: false,
-    disabled: isUploadDisabled,
   });
 
   const setNodeRef = useCombinedRefs(setDroppableRef, setDraggableRef);
@@ -149,13 +128,13 @@ const IAIDndImage = (props: IAIDndImageProps) => {
           sx={{
             w: 'full',
             h: 'full',
+            position: fitContainer ? 'absolute' : 'relative',
             alignItems: 'center',
             justifyContent: 'center',
           }}
         >
           <Image
             src={image.image_url}
-            fallbackStrategy="beforeLoadOrError"
             fallback={fallback}
             onError={onError}
             objectFit="contain"
@@ -205,9 +184,9 @@ const IAIDndImage = (props: IAIDndImageProps) => {
               color: 'base.500',
               ...uploadButtonStyles,
             }}
-            {...getRootProps()}
+            {...getUploadButtonProps()}
           >
-            <input {...getInputProps()} />
+            <input {...getUploadInputProps()} />
             <Icon
               as={isUploadDisabled ? FaImage : FaUpload}
               sx={{

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

@@ -0,0 +1,64 @@
+import { useAppDispatch } from 'app/store/storeHooks';
+import { useCallback } from 'react';
+import { useDropzone } from 'react-dropzone';
+import { PostUploadAction, imageUploaded } from 'services/api/thunks/image';
+
+type UseImageUploadButtonArgs = {
+  postUploadAction?: PostUploadAction;
+  isDisabled?: boolean;
+};
+
+/**
+ * Provides image uploader functionality to any component.
+ *
+ * @example
+ * const { getUploadButtonProps, getUploadInputProps } = useImageUploadButton({
+ *   postUploadAction: {
+ *     type: 'SET_CONTROLNET_IMAGE',
+ *     controlNetId: '12345',
+ *   },
+ *   isDisabled: getIsUploadDisabled(),
+ * });
+ *
+ * // in the render function
+ * <Button {...getUploadButtonProps()} /> // will open the file dialog on click
+ * <input {...getUploadInputProps()} /> // hidden, handles native upload functionality
+ */
+export const useImageUploadButton = ({
+  postUploadAction,
+  isDisabled,
+}: UseImageUploadButtonArgs) => {
+  const dispatch = useAppDispatch();
+  const onDropAccepted = useCallback(
+    (files: File[]) => {
+      const file = files[0];
+      if (!file) {
+        return;
+      }
+
+      dispatch(
+        imageUploaded({
+          file,
+          image_category: 'user',
+          is_intermediate: false,
+          postUploadAction,
+        })
+      );
+    },
+    [dispatch, postUploadAction]
+  );
+
+  const {
+    getRootProps: getUploadButtonProps,
+    getInputProps: getUploadInputProps,
+    open: openUploader,
+  } = useDropzone({
+    accept: { 'image/png': ['.png'], 'image/jpeg': ['.jpg', '.jpeg', '.png'] },
+    onDropAccepted,
+    disabled: isDisabled,
+    noDrag: true,
+    multiple: false,
+  });
+
+  return { getUploadButtonProps, getUploadInputProps, openUploader };
+};

invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx

@@ -10,7 +10,6 @@ import { Box, ChakraProps, Flex } from '@chakra-ui/react';
 import IAIDndImage from 'common/components/IAIDndImage';
 import { createSelector } from '@reduxjs/toolkit';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
-import { AnimatePresence, motion } from 'framer-motion';
 import { IAIImageLoadingFallback } from 'common/components/IAIImageFallback';
 import IAIIconButton from 'common/components/IAIIconButton';
 import { FaUndo } from 'react-icons/fa';
@@ -105,64 +104,46 @@ const ControlNetImagePreview = (props: Props) => {
       <IAIDndImage
         image={controlImage}
         onDrop={handleDrop}
-        isDropDisabled={Boolean(
+        isDropDisabled={shouldShowProcessedImage}
-          processedControlImage && processorType !== 'none'
-        )}
         isUploadDisabled={Boolean(controlImage)}
         postUploadAction={{ type: 'SET_CONTROLNET_IMAGE', controlNetId }}
-        imageSx={imageSx}
+        imageSx={{
+          pointerEvents: shouldShowProcessedImage ? 'none' : 'auto',
+          ...imageSx,
+        }}
       />
-      <AnimatePresence>
+      {shouldShowProcessedImage && (
-        {shouldShowProcessedImage && (
+        <Box
-          <motion.div
+          sx={{
-            style={{ width: '100%' }}
+            position: 'absolute',
-            initial={{
+            top: 0,
-              opacity: 0,
+            insetInlineStart: 0,
-            }}
+            w: 'full',
-            animate={{
+            h: 'full',
-              opacity: 1,
+          }}
-              transition: { duration: 0.1 },
+        >
-            }}
+          {shouldShowProcessedImageBackdrop && (
-            exit={{
+            <Box
-              opacity: 0,
+              sx={{
-              transition: { duration: 0.1 },
+                position: 'absolute',
-            }}
+                top: 0,
-          >
+                insetInlineStart: 0,
-            <>
+                w: 'full',
-              {shouldShowProcessedImageBackdrop && (
+                h: 'full',
-                <Box
+                bg: 'base.900',
-                  sx={{
+                opacity: 0.7,
-                    position: 'absolute',
+              }}
-                    top: 0,
+            />
-                    insetInlineStart: 0,
+          )}
-                    w: 'full',
+          <IAIDndImage
-                    h: 'full',
+            image={processedControlImage}
-                    bg: 'base.900',
+            onDrop={handleDrop}
-                    opacity: 0.7,
+            payloadImage={controlImage}
-                  }}
+            isUploadDisabled={true}
-                />
+            imageSx={imageSx}
-              )}
+          />
-              <Box
+        </Box>
-                sx={{
+      )}
-                  position: 'absolute',
-                  top: 0,
-                  insetInlineStart: 0,
-                  w: 'full',
-                  h: 'full',
-                }}
-              >
-                <IAIDndImage
-                  image={processedControlImage}
-                  onDrop={handleDrop}
-                  payloadImage={controlImage}
-                  isUploadDisabled={true}
-                  imageSx={imageSx}
-                />
-              </Box>
-            </>
-          </motion.div>
-        )}
-      </AnimatePresence>
       {pendingControlImages.includes(controlNetId) && (
         <Box
           sx={{

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

@@ -1,8 +1,7 @@
-import { Box, Flex } from '@chakra-ui/react';
+import { Flex } from '@chakra-ui/react';
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppSelector } from 'app/store/storeHooks';
 import { systemSelector } from 'features/system/store/systemSelectors';
-import { isEqual } from 'lodash-es';
 
 import { gallerySelector } from '../store/gallerySelectors';
 import CurrentImageButtons from './CurrentImageButtons';
@@ -22,13 +21,8 @@ export const currentImageDisplaySelector = createSelector(
   defaultSelectorOptions
 );
 
-/**
- * Displays the current image if there is one, plus associated actions.
- */
 const CurrentImageDisplay = () => {
-  const { hasSelectedImage, hasProgressImage } = useAppSelector(
+  const { hasSelectedImage } = useAppSelector(currentImageDisplaySelector);
-    currentImageDisplaySelector
-  );
 
   return (
     <Flex
@@ -43,24 +37,8 @@ const CurrentImageDisplay = () => {
         justifyContent: 'center',
       }}
     >
-      <Flex
+      {hasSelectedImage && <CurrentImageButtons />}
-        flexDirection="column"
+      <CurrentImagePreview />
-        sx={{
-          w: 'full',
-          h: 'full',
-          alignItems: 'center',
-          justifyContent: 'center',
-          gap: 4,
-          position: 'absolute',
-        }}
-      >
-        <CurrentImagePreview />
-      </Flex>
-      {hasSelectedImage && (
-        <Box sx={{ position: 'absolute', top: 0 }}>
-          <CurrentImageButtons />
-        </Box>
-      )}
     </Flex>
   );
 };

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

@@ -51,10 +51,6 @@ const CurrentImagePreview = () => {
     shouldAntialiasProgressImage,
   } = useAppSelector(imagesSelector);
 
-  // const image = useAppSelector((state: RootState) =>
-  //   selectImagesById(state, selectedImage ?? '')
-  // );
-
   const {
     currentData: image,
     isLoading,
@@ -79,7 +75,6 @@ const CurrentImagePreview = () => {
       sx={{
         width: 'full',
         height: 'full',
-        position: 'relative',
         alignItems: 'center',
         justifyContent: 'center',
       }}
@@ -101,21 +96,13 @@ const CurrentImagePreview = () => {
           }}
         />
       ) : (
-        <Flex
+        <IAIDndImage
-          sx={{
+          image={image}
-            width: 'full',
+          onDrop={handleDrop}
-            height: 'full',
+          fallback={<IAIImageLoadingFallback sx={{ bg: 'none' }} />}
-            alignItems: 'center',
+          isUploadDisabled={true}
-            justifyContent: 'center',
+          fitContainer
-          }}
+        />
-        >
-          <IAIDndImage
-            image={image}
-            onDrop={handleDrop}
-            fallback={<IAIImageLoadingFallback sx={{ bg: 'none' }} />}
-            isUploadDisabled={true}
-          />
-        </Flex>
       )}
       {shouldShowImageDetails && image && (
         <Box

invokeai/frontend/web/src/features/parameters/components/Parameters/ImageToImage/InitialImagePreview.tsx

@@ -1,4 +1,4 @@
-import { Flex } from '@chakra-ui/react';
+import { Flex, Spacer, Text } from '@chakra-ui/react';
 import { createSelector } from '@reduxjs/toolkit';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import {
@@ -13,6 +13,10 @@ import { ImageDTO } from 'services/api/types';
 import { IAIImageLoadingFallback } from 'common/components/IAIImageFallback';
 import { useGetImageDTOQuery } from 'services/api/endpoints/images';
 import { skipToken } from '@reduxjs/toolkit/dist/query';
+import IAIIconButton from 'common/components/IAIIconButton';
+import { FaUndo, FaUpload } from 'react-icons/fa';
+import useImageUploader from 'common/hooks/useImageUploader';
+import { useImageUploadButton } from 'common/hooks/useImageUploadButton';
 
 const selector = createSelector(
   [generationSelector],
@@ -28,6 +32,7 @@ const selector = createSelector(
 const InitialImagePreview = () => {
   const { initialImage } = useAppSelector(selector);
   const dispatch = useAppDispatch();
+  const { openUploader } = useImageUploader();
 
   const {
     currentData: image,
@@ -36,6 +41,10 @@ const InitialImagePreview = () => {
     isSuccess,
   } = useGetImageDTOQuery(initialImage?.imageName ?? skipToken);
 
+  const { getUploadButtonProps, getUploadInputProps } = useImageUploadButton({
+    postUploadAction: { type: 'SET_INITIAL_IMAGE' },
+  });
+
   const handleDrop = useCallback(
     (droppedImage: ImageDTO) => {
       if (droppedImage.image_name === initialImage?.imageName) {
@@ -50,25 +59,66 @@ const InitialImagePreview = () => {
     dispatch(clearInitialImage());
   }, [dispatch]);
 
+  const handleUpload = useCallback(() => {
+    openUploader();
+  }, [openUploader]);
+
   return (
     <Flex
       sx={{
+        flexDir: 'column',
         width: 'full',
         height: 'full',
         position: 'absolute',
         alignItems: 'center',
         justifyContent: 'center',
         p: 4,
+        gap: 4,
       }}
     >
+      <Flex
+        sx={{
+          w: 'full',
+          flexWrap: 'wrap',
+          justifyContent: 'center',
+          alignItems: 'center',
+          gap: 2,
+        }}
+      >
+        <Text
+          sx={{
+            color: 'base.200',
+            fontWeight: 600,
+            fontSize: 'sm',
+            userSelect: 'none',
+          }}
+        >
+          Initial Image
+        </Text>
+        <Spacer />
+        <IAIIconButton
+          tooltip="Upload Initial Image"
+          aria-label="Upload Initial Image"
+          icon={<FaUpload />}
+          onClick={handleUpload}
+          {...getUploadButtonProps()}
+        />
+        <IAIIconButton
+          tooltip="Reset Initial Image"
+          aria-label="Reset Initial Image"
+          icon={<FaUndo />}
+          onClick={handleReset}
+          isDisabled={!initialImage}
+        />
+      </Flex>
       <IAIDndImage
         image={image}
         onDrop={handleDrop}
-        onReset={handleReset}
         fallback={<IAIImageLoadingFallback sx={{ bg: 'none' }} />}
-        postUploadAction={{ type: 'SET_INITIAL_IMAGE' }}
+        isUploadDisabled={true}
-        withResetIcon
+        fitContainer
       />
+      <input {...getUploadInputProps()} />
     </Flex>
   );
 };