feat(ui): staging area (rendering wip)

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

@@ -2,6 +2,7 @@
 import { Flex } from '@invoke-ai/ui-library';
 import { ControlLayersToolbar } from 'features/controlLayers/components/ControlLayersToolbar';
 import { StageComponent } from 'features/controlLayers/components/StageComponent';
+import { StagingAreaToolbar } from 'features/controlLayers/components/StagingArea/StagingAreaToolbar';
 import { memo } from 'react';
 
 export const ControlLayersEditor = memo(() => {
@@ -17,6 +18,9 @@ export const ControlLayersEditor = memo(() => {
     >
       <ControlLayersToolbar />
       <StageComponent />
+      <Flex position="absolute" bottom={2} gap={2} align="center" justify="center">
+        <StagingAreaToolbar />
+      </Flex>
     </Flex>
   );
 });

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbar.tsx

@@ -0,0 +1,151 @@
+import { Button, ButtonGroup, IconButton } from '@invoke-ai/ui-library';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import {
+  stagingAreaImageAccepted,
+  stagingAreaImageDiscarded,
+  stagingAreaNextImageSelected,
+  stagingAreaPreviousImageSelected,
+  stagingAreaReset,
+} from 'features/controlLayers/store/canvasV2Slice';
+import type { CanvasV2State } from 'features/controlLayers/store/types';
+import { memo, useCallback, useMemo } from 'react';
+import { useHotkeys } from 'react-hotkeys-hook';
+import { useTranslation } from 'react-i18next';
+import { PiArrowLeftBold, PiArrowRightBold, PiCheckBold, PiTrashSimpleBold, PiXBold } from 'react-icons/pi';
+
+export const StagingAreaToolbar = memo(() => {
+  const stagingArea = useAppSelector((s) => s.canvasV2.stagingArea);
+
+  if (!stagingArea || stagingArea.images.length === 0) {
+    return null;
+  }
+
+  return <StagingAreaToolbarContent stagingArea={stagingArea} />;
+});
+
+StagingAreaToolbar.displayName = 'StagingAreaToolbar';
+
+type Props = {
+  stagingArea: NonNullable<CanvasV2State['stagingArea']>;
+};
+
+export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
+  const dispatch = useAppDispatch();
+  const images = useMemo(() => stagingArea.images, [stagingArea]);
+  const imageDTO = useMemo(() => {
+    if (stagingArea.selectedImageIndex === null) {
+      return null;
+    }
+    return images[stagingArea.selectedImageIndex] ?? null;
+  }, [images, stagingArea.selectedImageIndex]);
+
+  const { t } = useTranslation();
+
+  const onPrev = useCallback(() => {
+    dispatch(stagingAreaPreviousImageSelected());
+  }, [dispatch]);
+
+  const onNext = useCallback(() => {
+    dispatch(stagingAreaNextImageSelected());
+  }, [dispatch]);
+
+  const onAccept = useCallback(() => {
+    if (!imageDTO || !stagingArea) {
+      return;
+    }
+    dispatch(stagingAreaImageAccepted({ imageDTO }));
+  }, [dispatch, imageDTO, stagingArea]);
+
+  const onDiscardOne = useCallback(() => {
+    if (!imageDTO || !stagingArea) {
+      return;
+    }
+    if (images.length === 1) {
+      dispatch(stagingAreaReset());
+    } else {
+      dispatch(stagingAreaImageDiscarded({ imageDTO }));
+    }
+  }, [dispatch, imageDTO, images.length, stagingArea]);
+
+  const onDiscardAll = useCallback(() => {
+    if (!stagingArea) {
+      return;
+    }
+    dispatch(stagingAreaReset());
+  }, [dispatch, stagingArea]);
+
+  useHotkeys(['left'], onPrev, {
+    preventDefault: true,
+  });
+
+  useHotkeys(['right'], onNext, {
+    preventDefault: true,
+  });
+
+  useHotkeys(['enter'], onAccept, {
+    preventDefault: true,
+  });
+
+  return (
+    <>
+      <ButtonGroup borderRadius="base" shadow="dark-lg">
+        <IconButton
+          tooltip={`${t('unifiedCanvas.previous')} (Left)`}
+          aria-label={`${t('unifiedCanvas.previous')} (Left)`}
+          icon={<PiArrowLeftBold />}
+          onClick={onPrev}
+          colorScheme="invokeBlue"
+        />
+        <Button
+          colorScheme="base"
+          pointerEvents="none"
+          minW={20}
+        >{`${(stagingArea.selectedImageIndex ?? 0) + 1}/${images.length}`}</Button>
+        <IconButton
+          tooltip={`${t('unifiedCanvas.next')} (Right)`}
+          aria-label={`${t('unifiedCanvas.next')} (Right)`}
+          icon={<PiArrowRightBold />}
+          onClick={onNext}
+          colorScheme="invokeBlue"
+        />
+      </ButtonGroup>
+      <ButtonGroup borderRadius="base" shadow="dark-lg">
+        <IconButton
+          tooltip={`${t('unifiedCanvas.accept')} (Enter)`}
+          aria-label={`${t('unifiedCanvas.accept')} (Enter)`}
+          icon={<PiCheckBold />}
+          onClick={onAccept}
+          colorScheme="invokeBlue"
+        />
+        {/* <IconButton
+          tooltip={`${t('unifiedCanvas.saveToGallery')} (Shift+S)`}
+          aria-label={t('unifiedCanvas.saveToGallery')}
+          isDisabled={!imageDTO || !imageDTO.is_intermediate}
+          icon={<PiFloppyDiskBold />}
+          onClick={handleSaveToGallery}
+          colorScheme="invokeBlue"
+        /> */}
+        <IconButton
+          tooltip={`${t('unifiedCanvas.discardCurrent')}`}
+          aria-label={t('unifiedCanvas.discardCurrent')}
+          icon={<PiXBold />}
+          onClick={onDiscardOne}
+          colorScheme="invokeBlue"
+          fontSize={16}
+          isDisabled={images.length <= 1}
+        />
+        <IconButton
+          tooltip={`${t('unifiedCanvas.discardAll')} (Esc)`}
+          aria-label={t('unifiedCanvas.discardAll')}
+          icon={<PiTrashSimpleBold />}
+          onClick={onDiscardAll}
+          colorScheme="error"
+          fontSize={16}
+          isDisabled={images.length === 0}
+        />
+      </ButtonGroup>
+    </>
+  );
+});
+
+StagingAreaToolbarContent.displayName = 'StagingAreaToolbarContent';

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

@@ -53,6 +53,7 @@ export type ImageObjectRecord = {
   konvaPlaceholderGroup: Konva.Group;
   konvaPlaceholderRect: Konva.Rect;
   konvaPlaceholderText: Konva.Text;
+  imageName: string | null;
   konvaImage: Konva.Image | null; // The image is loaded asynchronously, so it may not be available immediately
   isLoading: boolean;
   isError: boolean;
@@ -69,6 +70,7 @@ type KonvaApi = {
   renderDocumentOverlay: () => void;
   renderBackground: () => void;
   renderToolPreview: () => void;
+  renderStagingArea: () => void;
   arrangeEntities: () => void;
   fitDocumentToStage: () => void;
   fitStageToContainer: () => void;
@@ -102,6 +104,10 @@ type PreviewLayer = {
     innerRect: Konva.Rect;
     outerRect: Konva.Rect;
   };
+  stagingArea: {
+    group: Konva.Group;
+    image: ImageObjectRecord | null;
+  };
 };
 
 type StateApi = {
@@ -143,6 +149,7 @@ type StateApi = {
   getControlAdaptersState: () => CanvasV2State['controlAdapters'];
   getRegionsState: () => CanvasV2State['regions'];
   getInpaintMaskState: () => CanvasV2State['inpaintMask'];
+  getStagingAreaState: () => CanvasV2State['stagingArea'];
   onInpaintMaskImageCached: (imageDTO: ImageDTO) => void;
   onRegionMaskImageCached: (id: string, imageDTO: ImageDTO) => void;
   onLayerImageCached: (imageDTO: ImageDTO) => void;

invokeai/frontend/web/src/features/controlLayers/konva/renderers/objects.ts

@@ -153,6 +153,7 @@ export const updateImageSource = async (arg: {
 
     const imageDTO = await getImageDTO(image.name);
     if (!imageDTO) {
+      objectRecord.imageName = null;
       objectRecord.isLoading = false;
       objectRecord.isError = true;
       objectRecord.konvaPlaceholderGroup.visible(true);
@@ -173,6 +174,7 @@ export const updateImageSource = async (arg: {
           image: imageEl,
         });
         objectRecord.konvaImageGroup.add(objectRecord.konvaImage);
+        objectRecord.imageName = image.name;
       }
       objectRecord.isLoading = false;
       objectRecord.isError = false;
@@ -180,6 +182,7 @@ export const updateImageSource = async (arg: {
       onLoad?.(objectRecord.konvaImage);
     };
     imageEl.onerror = () => {
+      objectRecord.imageName = null;
       objectRecord.isLoading = false;
       objectRecord.isError = true;
       objectRecord.konvaPlaceholderGroup.visible(true);
@@ -189,6 +192,7 @@ export const updateImageSource = async (arg: {
     imageEl.id = image.name;
     imageEl.src = imageDTO.image_url;
   } catch {
+    objectRecord.imageName = null;
     objectRecord.isLoading = false;
     objectRecord.isError = true;
     objectRecord.konvaPlaceholderGroup.visible(true);
@@ -218,7 +222,7 @@ export const createImageObjectGroup = (arg: {
   }
   const { id, image } = obj;
   const { width, height } = obj;
-  const konvaImageGroup = new Konva.Group({ id, name, listening: false });
+  const konvaImageGroup = new Konva.Group({ id, name, listening: false, x: obj.x, y: obj.y });
   const konvaPlaceholderGroup = new Konva.Group({ name: IMAGE_PLACEHOLDER_NAME, listening: false });
   const konvaPlaceholderRect = new Konva.Rect({
     fill: 'hsl(220 12% 45% / 1)', // 'base.500'
@@ -246,6 +250,7 @@ export const createImageObjectGroup = (arg: {
     konvaPlaceholderRect,
     konvaPlaceholderText,
     konvaImage: null,
+    imageName: null,
     isLoading: false,
     isError: false,
   });

invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts

@@ -22,6 +22,7 @@ import {
 } from 'features/controlLayers/konva/renderers/preview';
 import { getRenderRegions } from 'features/controlLayers/konva/renderers/regions';
 import { getFitDocumentToStage, getFitStageToContainer } from 'features/controlLayers/konva/renderers/stage';
+import { createStagingArea, getRenderStagingArea } from 'features/controlLayers/konva/renderers/stagingArea';
 import {
   $stageAttrs,
   bboxChanged,
@@ -259,6 +260,7 @@ export const initializeRenderer = (
   const getControlAdaptersState = () => canvasV2.controlAdapters;
   const getInpaintMaskState = () => canvasV2.inpaintMask;
   const getMaskOpacity = () => canvasV2.settings.maskOpacity;
+  const getStagingAreaState = () => canvasV2.stagingArea;
 
   // Read-write state, ephemeral interaction state
   let isDrawing = false;
@@ -307,6 +309,7 @@ export const initializeRenderer = (
     bbox: createBboxNodes(stage, getBbox, onBboxTransformed, $shift.get, $ctrl.get, $meta.get, $alt.get),
     tool: createToolPreviewNodes(),
     documentOverlay: createDocumentOverlay(),
+    stagingArea: createStagingArea(),
   };
   manager.preview.layer.add(manager.preview.bbox.group);
   manager.preview.layer.add(manager.preview.tool.group);
@@ -329,6 +332,7 @@ export const initializeRenderer = (
     getRegionsState,
     getMaskOpacity,
     getInpaintMaskState,
+    getStagingAreaState,
 
     // Read-write state
     setTool,
@@ -376,6 +380,7 @@ export const initializeRenderer = (
     renderBbox: getRenderBbox(manager),
     renderToolPreview: getRenderToolPreview(manager),
     renderDocumentOverlay: getRenderDocumentOverlay(manager),
+    renderStagingArea: getRenderStagingArea(manager),
     renderBackground: getRenderBackground(manager),
     arrangeEntities: getArrangeEntities(manager),
     fitDocumentToStage: getFitDocumentToStage(manager),

invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts

@@ -0,0 +1,41 @@
+import type { KonvaNodeManager } from 'features/controlLayers/konva/nodeManager';
+import { createImageObjectGroup, updateImageSource } from 'features/controlLayers/konva/renderers/objects';
+import { imageDTOToImageObject, imageDTOToImageWithDims } from 'features/controlLayers/store/types';
+import Konva from 'konva';
+import { assert } from 'tsafe';
+
+export const createStagingArea = (): KonvaNodeManager['preview']['stagingArea'] => {
+  const group = new Konva.Group({ id: 'staging_area_group', listening: false });
+  return { group, image: null };
+};
+
+export const getRenderStagingArea = async (manager: KonvaNodeManager) => {
+  const { getStagingAreaState } = manager.stateApi;
+  const stagingArea = getStagingAreaState();
+
+  if (!stagingArea || stagingArea.selectedImageIndex === null) {
+    if (manager.preview.stagingArea.image) {
+      manager.preview.stagingArea.image.konvaImageGroup.visible(false);
+      manager.preview.stagingArea.image = null;
+    }
+    return;
+  }
+
+  if (stagingArea.selectedImageIndex) {
+    const imageDTO = stagingArea.images[stagingArea.selectedImageIndex];
+    assert(imageDTO, 'Image must exist');
+    if (manager.preview.stagingArea.image) {
+      if (manager.preview.stagingArea.image.imageName !== imageDTO.image_name) {
+        await updateImageSource({
+          objectRecord: manager.preview.stagingArea.image,
+          image: imageDTOToImageWithDims(imageDTO),
+        });
+      }
+    } else {
+      manager.preview.stagingArea.image = await createImageObjectGroup({
+        obj: imageDTOToImageObject(imageDTO),
+        name: imageDTO.image_name,
+      });
+    }
+  }
+};

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

@@ -13,6 +13,7 @@ import { lorasReducers } from 'features/controlLayers/store/lorasReducers';
 import { paramsReducers } from 'features/controlLayers/store/paramsReducers';
 import { regionsReducers } from 'features/controlLayers/store/regionsReducers';
 import { settingsReducers } from 'features/controlLayers/store/settingsReducers';
+import { stagingAreaReducers } from 'features/controlLayers/store/stagingAreaReducers';
 import { toolReducers } from 'features/controlLayers/store/toolReducers';
 import { initialAspectRatioState } from 'features/parameters/components/ImageSize/constants';
 import type { AspectRatioState } from 'features/parameters/components/ImageSize/types';
@@ -119,6 +120,7 @@ const initialState: CanvasV2State = {
     refinerNegativeAestheticScore: 2.5,
     refinerStart: 0.8,
   },
+  stagingArea: null,
 };
 
 export const canvasV2Slice = createSlice({
@@ -136,6 +138,7 @@ export const canvasV2Slice = createSlice({
     ...toolReducers,
     ...bboxReducers,
     ...inpaintMaskReducers,
+    ...stagingAreaReducers,
     widthChanged: (state, action: PayloadAction<{ width: number; updateAspectRatio?: boolean; clamp?: boolean }>) => {
       const { width, updateAspectRatio, clamp } = action.payload;
       state.document.width = clamp ? Math.max(roundDownToMultiple(width, 8), 64) : width;
@@ -327,6 +330,15 @@ export const {
   imEraserLineAdded,
   imLinePointAdded,
   imRectAdded,
+  // Staging
+  stagingAreaInitialized,
+  stagingAreaImageAdded,
+  stagingAreaBatchIdAdded,
+  stagingAreaImageDiscarded,
+  stagingAreaImageAccepted,
+  stagingAreaReset,
+  stagingAreaNextImageSelected,
+  stagingAreaPreviousImageSelected,
 } = canvasV2Slice.actions;
 
 export const selectCanvasV2Slice = (state: RootState) => state.canvasV2;

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

@@ -231,17 +231,27 @@ export const layersReducers = {
     prepare: (payload: RectShapeAddedArg) => ({ payload: { ...payload, rectId: uuidv4() } }),
   },
   layerImageAdded: {
-    reducer: (state, action: PayloadAction<ImageObjectAddedArg & { objectId: string }>) => {
-      const { id, objectId, imageDTO } = action.payload;
+    reducer: (
+      state,
+      action: PayloadAction<ImageObjectAddedArg & { objectId: string; pos?: { x: number; y: number } }>
+    ) => {
+      const { id, objectId, imageDTO, pos } = action.payload;
       const layer = selectLayer(state, id);
       if (!layer) {
         return;
       }
-      layer.objects.push(imageDTOToImageObject(id, objectId, imageDTO));
+      const imageObject = imageDTOToImageObject(id, objectId, imageDTO);
+      if (pos) {
+        imageObject.x = pos.x;
+        imageObject.y = pos.y;
+      }
+      layer.objects.push(imageObject);
       layer.bboxNeedsUpdate = true;
       state.layers.imageCache = null;
     },
-    prepare: (payload: ImageObjectAddedArg) => ({ payload: { ...payload, objectId: uuidv4() } }),
+    prepare: (payload: ImageObjectAddedArg & { pos?: { x: number; y: number } }) => ({
+      payload: { ...payload, objectId: uuidv4() },
+    }),
   },
   layerImageCacheChanged: (state, action: PayloadAction<{ imageDTO: ImageDTO | null }>) => {
     const { imageDTO } = action.payload;

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

@@ -0,0 +1,73 @@
+import type { PayloadAction, SliceCaseReducers } from '@reduxjs/toolkit';
+import type { CanvasV2State, Rect } from 'features/controlLayers/store/types';
+import type { ImageDTO } from 'services/api/types';
+
+export const stagingAreaReducers = {
+  stagingAreaInitialized: (state, action: PayloadAction<{ bbox: Rect; batchIds: string[] }>) => {
+    const { bbox, batchIds } = action.payload;
+    state.stagingArea = {
+      bbox,
+      batchIds,
+      selectedImageIndex: null,
+      images: [],
+    };
+  },
+  stagingAreaImageAdded: (state, action: PayloadAction<{ imageDTO: ImageDTO }>) => {
+    const { imageDTO } = action.payload;
+    if (!state.stagingArea) {
+      // Should not happen
+      return;
+    }
+    state.stagingArea.images.push(imageDTO);
+    if (!state.stagingArea.selectedImageIndex) {
+      state.stagingArea.selectedImageIndex = state.stagingArea.images.length - 1;
+    }
+  },
+  stagingAreaNextImageSelected: (state) => {
+    if (!state.stagingArea) {
+      // Should not happen
+      return;
+    }
+    if (state.stagingArea.selectedImageIndex === null) {
+      if (state.stagingArea.images.length > 0) {
+        state.stagingArea.selectedImageIndex = 0;
+      }
+      return;
+    }
+    state.stagingArea.selectedImageIndex = (state.stagingArea.selectedImageIndex + 1) % state.stagingArea.images.length;
+  },
+  stagingAreaPreviousImageSelected: (state) => {
+    if (!state.stagingArea) {
+      // Should not happen
+      return;
+    }
+    if (state.stagingArea.selectedImageIndex === null) {
+      if (state.stagingArea.images.length > 0) {
+        state.stagingArea.selectedImageIndex = 0;
+      }
+      return;
+    }
+    state.stagingArea.selectedImageIndex =
+      (state.stagingArea.selectedImageIndex - 1 + state.stagingArea.images.length) % state.stagingArea.images.length;
+  },
+  stagingAreaBatchIdAdded: (state, action: PayloadAction<{ batchId: string }>) => {
+    const { batchId } = action.payload;
+    if (!state.stagingArea) {
+      // Should not happen
+      return;
+    }
+    state.stagingArea.batchIds.push(batchId);
+  },
+  stagingAreaImageDiscarded: (state, action: PayloadAction<{ imageDTO: ImageDTO }>) => {
+    const { imageDTO } = action.payload;
+    if (!state.stagingArea) {
+      // Should not happen
+      return;
+    }
+    state.stagingArea.images = state.stagingArea.images.filter((image) => image.image_name !== imageDTO.image_name);
+  },
+  stagingAreaImageAccepted: (state, _: PayloadAction<{ imageDTO: ImageDTO }>) => state,
+  stagingAreaReset: (state) => {
+    state.stagingArea = null;
+  },
+} satisfies SliceCaseReducers<CanvasV2State>;

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

@@ -785,6 +785,11 @@ export type Size = {
   height: number;
 };
 
+export type Position = {
+  x: number;
+  y: number;
+};
+
 export type LoRA = {
   id: string;
   isEnabled: boolean;
@@ -877,6 +882,12 @@ export type CanvasV2State = {
     refinerNegativeAestheticScore: number;
     refinerStart: number;
   };
+  stagingArea: {
+    bbox: Rect;
+    images: ImageDTO[];
+    selectedImageIndex: number | null;
+    batchIds: string[];
+  } | null;
 };
 
 export type StageAttrs = { x: number; y: number; width: number; height: number; scale: number };
@@ -891,7 +902,7 @@ export type EraserLineAddedArg = {
 export type BrushLineAddedArg = EraserLineAddedArg & { color: RgbaColor };
 export type PointAddedToLineArg = { id: string; point: [number, number] };
 export type RectShapeAddedArg = { id: string; rect: IRect; color: RgbaColor };
-export type ImageObjectAddedArg = { id: string; imageDTO: ImageDTO };
+export type ImageObjectAddedArg = { id: string; imageDTO: ImageDTO; pos?: Position };
 
 //#region Type guards
 export const isLine = (obj: RenderableObject): obj is BrushLine | EraserLine => {