feat(ui): staging area barely works

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear.ts

@@ -1,7 +1,7 @@
 import { enqueueRequested } from 'app/store/actions';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { getNodeManager } from 'features/controlLayers/konva/nodeManager';
-import { stagingAreaInitialized } from 'features/controlLayers/store/canvasV2Slice';
+import { stagingAreaBatchIdAdded, stagingAreaInitialized } from 'features/controlLayers/store/canvasV2Slice';
 import { isImageViewerOpenChanged } from 'features/gallery/store/gallerySlice';
 import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
 import { buildSD1Graph } from 'features/nodes/util/graph/generation/buildSD1Graph';
@@ -48,12 +48,16 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
         // TODO(psyche): update the backend schema, this is always provided
         const batchId = enqueueResult.batch.batch_id;
         assert(batchId, 'No batch ID found in enqueue result');
-        dispatch(
-          stagingAreaInitialized({
-            batchIds: [batchId],
-            bbox: getState().canvasV2.bbox,
-          })
-        );
+        if (!state.canvasV2.stagingArea) {
+          dispatch(
+            stagingAreaInitialized({
+              batchIds: [batchId],
+              bbox: state.canvasV2.bbox,
+            })
+          );
+        } else {
+          dispatch(stagingAreaBatchIdAdded({ batchId }));
+        }
       } finally {
         req.reset();
       }

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

@@ -1,6 +1,12 @@
 import { getImageDataTransparency } from 'common/util/arrayBuffer';
 import { CanvasBackground } from 'features/controlLayers/konva/renderers/background';
-import { CanvasPreview } from 'features/controlLayers/konva/renderers/preview';
+import {
+  CanvasBbox,
+  CanvasDocumentSizeOverlay,
+  CanvasPreview,
+  CanvasStagingArea,
+  CanvasTool,
+} from 'features/controlLayers/konva/renderers/preview';
 import { konvaNodeToBlob, konvaNodeToImageData, previewBlob } from 'features/controlLayers/konva/util';
 import type {
   BrushLineAddedArg,
@@ -106,7 +112,7 @@ export class KonvaNodeManager {
   controlAdapters: Map<string, CanvasControlAdapter>;
   layers: Map<string, CanvasLayer>;
   regions: Map<string, CanvasRegion>;
-  inpaintMask: CanvasInpaintMask | null;
+  inpaintMask: CanvasInpaintMask;
   util: Util;
   stateApi: StateApi;
   preview: CanvasPreview;
@@ -134,23 +140,29 @@ export class KonvaNodeManager {
     };
 
     this.preview = new CanvasPreview(
-      this.stage,
-      this.stateApi.getBbox,
-      this.stateApi.onBboxTransformed,
-      this.stateApi.getShiftKey,
-      this.stateApi.getCtrlKey,
-      this.stateApi.getMetaKey,
-      this.stateApi.getAltKey
+      new CanvasBbox(
+        this.stateApi.getBbox,
+        this.stateApi.onBboxTransformed,
+        this.stateApi.getShiftKey,
+        this.stateApi.getCtrlKey,
+        this.stateApi.getMetaKey,
+        this.stateApi.getAltKey
+      ),
+      new CanvasTool(),
+      new CanvasDocumentSizeOverlay(),
+      new CanvasStagingArea()
     );
     this.stage.add(this.preview.konvaLayer);
 
     this.background = new CanvasBackground();
     this.stage.add(this.background.konvaLayer);
 
+    this.inpaintMask = new CanvasInpaintMask(this.stateApi.getInpaintMaskState(), this.stateApi.onPosChanged);
+    this.stage.add(this.inpaintMask.konvaLayer);
+
     this.layers = new Map();
     this.regions = new Map();
     this.controlAdapters = new Map();
-    this.inpaintMask = null;
   }
 
   renderLayers() {
@@ -202,10 +214,6 @@ export class KonvaNodeManager {
 
   renderInpaintMask() {
     const inpaintMaskState = this.stateApi.getInpaintMaskState();
-    if (!this.inpaintMask) {
-      this.inpaintMask = new CanvasInpaintMask(inpaintMaskState, this.stateApi.onPosChanged);
-      this.stage.add(this.inpaintMask.konvaLayer);
-    }
     const toolState = this.stateApi.getToolState();
     const selectedEntity = this.stateApi.getSelectedEntity();
     const maskOpacity = this.stateApi.getMaskOpacity();
@@ -280,6 +288,10 @@ export class KonvaNodeManager {
     this.background.renderBackground(this.stage);
   }
 
+  renderStagingArea() {
+    this.preview.stagingArea.render(this.stateApi.getStagingAreaState());
+  }
+
   fitDocument() {
     this.preview.documentSizeOverlay.fitToStage(this.stage, this.stateApi.getDocument(), this.stateApi.setStageAttrs);
   }

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

@@ -189,8 +189,8 @@ export class KonvaImage {
           image: imageEl,
         });
         this.konvaImageGroup.add(this.konvaImage);
-        this.imageName = imageName;
       }
+      this.imageName = imageName;
       this.isLoading = false;
       this.isError = false;
       this.konvaPlaceholderGroup.visible(false);

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

@@ -101,11 +101,11 @@ export class CanvasStagingArea {
       return;
     }
 
-    if (stagingArea.selectedImageIndex) {
+    if (stagingArea.selectedImageIndex !== null) {
       const imageDTO = stagingArea.images[stagingArea.selectedImageIndex];
       assert(imageDTO, 'Image must exist');
       if (this.image) {
-        if (this.image.imageName !== imageDTO.image_name) {
+        if (!this.image.isLoading && !this.image.isError && this.image.imageName !== imageDTO.image_name) {
           await this.image.updateImageSource(imageDTO.image_name);
         }
       } else {
@@ -114,8 +114,8 @@ export class CanvasStagingArea {
           imageObject: {
             id: 'staging-area-image',
             type: 'image',
-            x: 0,
-            y: 0,
+            x: stagingArea.bbox.x,
+            y: stagingArea.bbox.y,
             width,
             height,
             filters: [],
@@ -126,6 +126,8 @@ export class CanvasStagingArea {
             },
           },
         });
+        this.group.add(this.image.konvaImageGroup);
+        await this.image.updateImageSource(imageDTO.image_name);
       }
     }
   }
@@ -375,7 +377,6 @@ export class CanvasBbox {
   NO_ANCHORS: string[] = [];
 
   constructor(
-    stage: Konva.Stage,
     getBbox: () => IRect,
     onBboxTransformed: (bbox: IRect) => void,
     getShiftKey: () => boolean,
@@ -446,6 +447,8 @@ export class CanvasBbox {
       anchorDragBoundFunc: (_oldAbsPos, newAbsPos) => {
         // This function works with absolute position - that is, a position in "physical" pixels on the screen, as opposed
         // to konva's internal coordinate system.
+        const stage = this.transformer.getStage();
+        assert(stage, 'Stage must exist');
 
         // We need to snap the anchors to the grid. If the user is holding ctrl/meta, we use the finer 8px grid.
         const gridSize = getCtrlKey() || getMetaKey() ? 8 : 64;
@@ -588,26 +591,23 @@ export class CanvasPreview {
   stagingArea: CanvasStagingArea;
 
   constructor(
-    stage: Konva.Stage,
-    getBbox: () => IRect,
-    onBboxTransformed: (bbox: IRect) => void,
-    getShiftKey: () => boolean,
-    getCtrlKey: () => boolean,
-    getMetaKey: () => boolean,
-    getAltKey: () => boolean
+    bbox: CanvasBbox,
+    tool: CanvasTool,
+    documentSizeOverlay: CanvasDocumentSizeOverlay,
+    stagingArea: CanvasStagingArea
   ) {
     this.konvaLayer = new Konva.Layer({ listening: true });
 
-    this.bbox = new CanvasBbox(stage, getBbox, onBboxTransformed, getShiftKey, getCtrlKey, getMetaKey, getAltKey);
+    this.bbox = bbox;
     this.konvaLayer.add(this.bbox.group);
 
-    this.tool = new CanvasTool();
+    this.tool = tool;
     this.konvaLayer.add(this.tool.group);
 
-    this.documentSizeOverlay = new CanvasDocumentSizeOverlay();
+    this.documentSizeOverlay = documentSizeOverlay;
     this.konvaLayer.add(this.documentSizeOverlay.group);
 
-    this.stagingArea = new CanvasStagingArea();
+    this.stagingArea = stagingArea;
     this.konvaLayer.add(this.stagingArea.group);
   }
 }

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

@@ -343,7 +343,7 @@ export const initializeRenderer = (
   // the entire state over when needed.
   const debouncedUpdateBboxes = debounce(updateBboxes, 300);
 
-  const renderCanvas = () => {
+  const renderCanvas = async () => {
     canvasV2 = store.getState().canvasV2;
 
     if (prevCanvasV2 === canvasV2 && !isFirstRender) {
@@ -408,6 +408,11 @@ export const initializeRenderer = (
       // debouncedUpdateBboxes(stage, canvasV2.layers, canvasV2.controlAdapters, canvasV2.regions, onBboxChanged);
     }
 
+    if (isFirstRender || canvasV2.stagingArea !== prevCanvasV2.stagingArea) {
+      logIfDebugging('Rendering staging area');
+      manager.renderStagingArea();
+    }
+
     if (
       isFirstRender ||
       canvasV2.layers.entities !== prevCanvasV2.layers.entities ||