feat(ui): add ref image drop targets to canvas

invoke-ai · Sep 16, 2024 · 7fe1135 · 7fe1135
1 parent 41efd81
commit 7fe1135
Show file tree

Hide file tree

Showing 4 changed files with 95 additions and 12 deletions.
diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts
@@ -1,18 +1,27 @@
 import { createAction } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { selectDefaultControlAdapter } from 'features/controlLayers/hooks/addLayerHooks';
+import { deepClone } from 'common/util/deepClone';
+import { selectDefaultControlAdapter, selectDefaultIPAdapter } from 'features/controlLayers/hooks/addLayerHooks';
+import { getPrefixedId } from 'features/controlLayers/konva/util';
 import {
   controlLayerAdded,
   entityRasterized,
   entitySelected,
   rasterLayerAdded,
+  referenceImageAdded,
   referenceImageIPAdapterImageChanged,
+  rgAdded,
   rgIPAdapterImageChanged,
 } from 'features/controlLayers/store/canvasSlice';
 import { selectCanvasSlice } from 'features/controlLayers/store/selectors';
-import type { CanvasControlLayerState, CanvasRasterLayerState } from 'features/controlLayers/store/types';
-import { imageDTOToImageObject } from 'features/controlLayers/store/util';
+import type {
+  CanvasControlLayerState,
+  CanvasRasterLayerState,
+  CanvasReferenceImageState,
+  CanvasRegionalGuidanceState,
+} from 'features/controlLayers/store/types';
+import { imageDTOToImageObject, imageDTOToImageWithDims } from 'features/controlLayers/store/util';
 import type { TypesafeDraggableData, TypesafeDroppableData } from 'features/dnd/types';
 import { isValidDrop } from 'features/dnd/util/isValidDrop';
 import { imageToCompareChanged, selectionChanged } from 'features/gallery/store/gallerySlice';
@@ -122,6 +131,36 @@ export const addImageDroppedListener = (startAppListening: AppStartListening) =>
         return;
       }
 
+      if (
+        overData.actionType === 'ADD_REGIONAL_REFERENCE_IMAGE_FROM_IMAGE' &&
+        activeData.payloadType === 'IMAGE_DTO' &&
+        activeData.payload.imageDTO
+      ) {
+        const state = getState();
+        const ipAdapter = deepClone(selectDefaultIPAdapter(state));
+        ipAdapter.image = imageDTOToImageWithDims(activeData.payload.imageDTO);
+        const overrides: Partial<CanvasRegionalGuidanceState> = {
+          referenceImages: [{ id: getPrefixedId('regional_guidance_reference_image'), ipAdapter }],
+        };
+        dispatch(rgAdded({ overrides, isSelected: true }));
+        return;
+      }
+
+      if (
+        overData.actionType === 'ADD_GLOBAL_REFERENCE_IMAGE_FROM_IMAGE' &&
+        activeData.payloadType === 'IMAGE_DTO' &&
+        activeData.payload.imageDTO
+      ) {
+        const state = getState();
+        const ipAdapter = deepClone(selectDefaultIPAdapter(state));
+        ipAdapter.image = imageDTOToImageWithDims(activeData.payload.imageDTO);
+        const overrides: Partial<CanvasReferenceImageState> = {
+          ipAdapter,
+        };
+        dispatch(referenceImageAdded({ overrides, isSelected: true }));
+        return;
+      }
+
       /**
        * Image dropped on Raster layer
        */

diff --git a/invokeai/frontend/web/src/features/controlLayers/components/CanvasDropArea.tsx b/invokeai/frontend/web/src/features/controlLayers/components/CanvasDropArea.tsx
@@ -1,6 +1,11 @@
-import { Flex } from '@invoke-ai/ui-library';
+import { Grid, GridItem } from '@invoke-ai/ui-library';
 import IAIDroppable from 'common/components/IAIDroppable';
-import type { AddControlLayerFromImageDropData, AddRasterLayerFromImageDropData } from 'features/dnd/types';
+import type {
+  AddControlLayerFromImageDropData,
+  AddGlobalReferenceImageFromImageDropData,
+  AddRasterLayerFromImageDropData,
+  AddRegionalReferenceImageFromImageDropData,
+} from 'features/dnd/types';
 import { useImageViewer } from 'features/gallery/components/ImageViewer/useImageViewer';
 import { memo } from 'react';
 
@@ -14,6 +19,16 @@ const addControlLayerFromImageDropData: AddControlLayerFromImageDropData = {
   actionType: 'ADD_CONTROL_LAYER_FROM_IMAGE',
 };
 
+const addRegionalReferenceImageFromImageDropData: AddRegionalReferenceImageFromImageDropData = {
+  id: 'add-control-layer-from-image-drop-data',
+  actionType: 'ADD_REGIONAL_REFERENCE_IMAGE_FROM_IMAGE',
+};
+
+const addGlobalReferenceImageFromImageDropData: AddGlobalReferenceImageFromImageDropData = {
+  id: 'add-control-layer-from-image-drop-data',
+  actionType: 'ADD_GLOBAL_REFERENCE_IMAGE_FROM_IMAGE',
+};
+
 export const CanvasDropArea = memo(() => {
   const imageViewer = useImageViewer();
 
@@ -23,12 +38,29 @@ export const CanvasDropArea = memo(() => {
 
   return (
     <>
-      <Flex position="absolute" top={0} right={0} bottom="50%" left={0} gap={2} pointerEvents="none">
-        <IAIDroppable dropLabel="Create Raster Layer" data={addRasterLayerFromImageDropData} />
-      </Flex>
-      <Flex position="absolute" top="50%" right={0} bottom={0} left={0} gap={2} pointerEvents="none">
-        <IAIDroppable dropLabel="Create Control Layer" data={addControlLayerFromImageDropData} />
-      </Flex>
+      <Grid
+        gridTemplateRows="1fr 1fr"
+        gridTemplateColumns="1fr 1fr"
+        position="absolute"
+        top={0}
+        right={0}
+        bottom={0}
+        left={0}
+        pointerEvents="none"
+      >
+        <GridItem position="relative">
+          <IAIDroppable dropLabel="New Raster Layer" data={addRasterLayerFromImageDropData} />
+        </GridItem>
+        <GridItem position="relative">
+          <IAIDroppable dropLabel="New Control Layer" data={addControlLayerFromImageDropData} />
+        </GridItem>
+        <GridItem position="relative">
+          <IAIDroppable dropLabel="New Regional Reference Image" data={addRegionalReferenceImageFromImageDropData} />
+        </GridItem>
+        <GridItem position="relative">
+          <IAIDroppable dropLabel="New Global Reference Image" data={addGlobalReferenceImageFromImageDropData} />
+        </GridItem>
+      </Grid>
     </>
   );
 });

diff --git a/invokeai/frontend/web/src/features/dnd/types/index.ts b/invokeai/frontend/web/src/features/dnd/types/index.ts
@@ -42,6 +42,14 @@ export type AddControlLayerFromImageDropData = BaseDropData & {
   actionType: 'ADD_CONTROL_LAYER_FROM_IMAGE';
 };
 
+export type AddRegionalReferenceImageFromImageDropData = BaseDropData & {
+  actionType: 'ADD_REGIONAL_REFERENCE_IMAGE_FROM_IMAGE';
+};
+
+export type AddGlobalReferenceImageFromImageDropData = BaseDropData & {
+  actionType: 'ADD_GLOBAL_REFERENCE_IMAGE_FROM_IMAGE';
+};
+
 export type ReplaceLayerImageDropData = BaseDropData & {
   actionType: 'REPLACE_LAYER_WITH_IMAGE';
   context: {
@@ -88,7 +96,9 @@ export type TypesafeDroppableData =
   | UpscaleInitialImageDropData
   | AddRasterLayerFromImageDropData
   | AddControlLayerFromImageDropData
-  | ReplaceLayerImageDropData;
+  | ReplaceLayerImageDropData
+  | AddRegionalReferenceImageFromImageDropData
+  | AddGlobalReferenceImageFromImageDropData;
 
 type BaseDragData = {
   id: string;

diff --git a/invokeai/frontend/web/src/features/dnd/util/isValidDrop.ts b/invokeai/frontend/web/src/features/dnd/util/isValidDrop.ts
@@ -21,6 +21,8 @@ export const isValidDrop = (overData?: TypesafeDroppableData | null, activeData?
     case 'SET_NODES_IMAGE':
     case 'SELECT_FOR_COMPARE':
     case 'REPLACE_LAYER_WITH_IMAGE':
+    case 'ADD_GLOBAL_REFERENCE_IMAGE_FROM_IMAGE':
+    case 'ADD_REGIONAL_REFERENCE_IMAGE_FROM_IMAGE':
       return payloadType === 'IMAGE_DTO';
     case 'ADD_TO_BOARD': {
       // If the board is the same, don't allow the drop