feat(ui): consolidate konva API

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

@@ -467,7 +467,7 @@ export const setStageEventHandlers = (manager: KonvaNodeManager): (() => void) =
         stage.position(newPos);
         setStageAttrs({ ...newPos, width: stage.width(), height: stage.height(), scale: newScale });
         manager.renderBackground();
-        manager.renderDocumentOverlay();
+        manager.renderDocumentSizeOverlay();
       }
     }
     manager.renderToolPreview();
@@ -483,7 +483,7 @@ export const setStageEventHandlers = (manager: KonvaNodeManager): (() => void) =
       scale: stage.scaleX(),
     });
     manager.renderBackground();
-    manager.renderDocumentOverlay();
+    manager.renderDocumentSizeOverlay();
     manager.renderToolPreview();
   });
 
@@ -518,10 +518,9 @@ export const setStageEventHandlers = (manager: KonvaNodeManager): (() => void) =
       setTool('view');
       setSpaceKey(true);
     } else if (e.key === 'r') {
-      manager.fitDocumentToStage();
+      manager.fitDocument();
-      manager.renderToolPreview();
       manager.renderBackground();
-      manager.renderDocumentOverlay();
+      manager.renderDocumentSizeOverlay();
     }
     manager.renderToolPreview();
   };

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

@@ -1,7 +1,6 @@
 import { getImageDataTransparency } from 'common/util/arrayBuffer';
-import { DOCUMENT_FIT_PADDING_PX } from 'features/controlLayers/konva/constants';
+import { CanvasBackground } from 'features/controlLayers/konva/renderers/background';
-import { KonvaBackground } from 'features/controlLayers/konva/renderers/background';
+import { CanvasPreview } from 'features/controlLayers/konva/renderers/preview';
-import { KonvaPreview } from 'features/controlLayers/konva/renderers/preview';
 import { konvaNodeToBlob, konvaNodeToImageData, previewBlob } from 'features/controlLayers/konva/util';
 import type {
   BrushLineAddedArg,
@@ -20,14 +19,15 @@ import type {
 import { isValidLayer } from 'features/nodes/util/graph/generation/addLayers';
 import type Konva from 'konva';
 import type { Vector2d } from 'konva/lib/types';
+import { atom } from 'nanostores';
 import { getImageDTO as defaultGetImageDTO, uploadImage as defaultUploadImage } from 'services/api/endpoints/images';
 import type { ImageCategory, ImageDTO } from 'services/api/types';
 import { assert } from 'tsafe';
 
-import { KonvaControlAdapter } from './renderers/controlAdapters';
+import { CanvasControlAdapter } from './renderers/controlAdapters';
-import { KonvaInpaintMask } from './renderers/inpaintMask';
+import { CanvasInpaintMask } from './renderers/inpaintMask';
-import { KonvaLayerAdapter } from './renderers/layers';
+import { CanvasLayer } from './renderers/layers';
-import { KonvaRegion } from './renderers/regions';
+import { CanvasRegion } from './renderers/regions';
 
 export type StateApi = {
   getToolState: () => CanvasV2State['tool'];
@@ -90,17 +90,27 @@ type Util = {
   getGenerationMode: () => GenerationMode;
 };
 
+const $nodeManager = atom<KonvaNodeManager | null>(null);
+export function getNodeManager() {
+  const nodeManager = $nodeManager.get();
+  assert(nodeManager !== null, 'Node manager not initialized');
+  return nodeManager;
+}
+export function setNodeManager(nodeManager: KonvaNodeManager) {
+  $nodeManager.set(nodeManager);
+}
+
 export class KonvaNodeManager {
   stage: Konva.Stage;
   container: HTMLDivElement;
-  controlAdapters: Map<string, KonvaControlAdapter>;
+  controlAdapters: Map<string, CanvasControlAdapter>;
-  layers: Map<string, KonvaLayerAdapter>;
+  layers: Map<string, CanvasLayer>;
-  regions: Map<string, KonvaRegion>;
+  regions: Map<string, CanvasRegion>;
-  inpaintMask: KonvaInpaintMask | null;
+  inpaintMask: CanvasInpaintMask | null;
   util: Util;
   stateApi: StateApi;
-  preview: KonvaPreview;
+  preview: CanvasPreview;
-  background: KonvaBackground;
+  background: CanvasBackground;
 
   constructor(
     stage: Konva.Stage,
@@ -122,7 +132,8 @@ export class KonvaNodeManager {
       getCompositeLayerStageClone: this._getCompositeLayerStageClone.bind(this),
       getGenerationMode: this._getGenerationMode.bind(this),
     };
-    this.preview = new KonvaPreview(
+
+    this.preview = new CanvasPreview(
       this.stage,
       this.stateApi.getBbox,
       this.stateApi.onBboxTransformed,
@@ -131,7 +142,11 @@ export class KonvaNodeManager {
       this.stateApi.getMetaKey,
       this.stateApi.getAltKey
     );
-    this.background = new KonvaBackground();
+    this.stage.add(this.preview.konvaLayer);
+
+    this.background = new CanvasBackground();
+    this.stage.add(this.background.konvaLayer);
+
     this.layers = new Map();
     this.regions = new Map();
     this.controlAdapters = new Map();
@@ -152,7 +167,7 @@ export class KonvaNodeManager {
     for (const entity of entities) {
       let adapter = this.layers.get(entity.id);
       if (!adapter) {
-        adapter = new KonvaLayerAdapter(entity, this.stateApi.onPosChanged);
+        adapter = new CanvasLayer(entity, this.stateApi.onPosChanged);
         this.layers.set(adapter.id, adapter);
         this.stage.add(adapter.konvaLayer);
       }
@@ -177,7 +192,7 @@ export class KonvaNodeManager {
     for (const entity of entities) {
       let adapter = this.regions.get(entity.id);
       if (!adapter) {
-        adapter = new KonvaRegion(entity, this.stateApi.onPosChanged);
+        adapter = new CanvasRegion(entity, this.stateApi.onPosChanged);
         this.regions.set(adapter.id, adapter);
         this.stage.add(adapter.konvaLayer);
       }
@@ -188,7 +203,7 @@ export class KonvaNodeManager {
   renderInpaintMask() {
     const inpaintMaskState = this.stateApi.getInpaintMaskState();
     if (!this.inpaintMask) {
-      this.inpaintMask = new KonvaInpaintMask(inpaintMaskState, this.stateApi.onPosChanged);
+      this.inpaintMask = new CanvasInpaintMask(inpaintMaskState, this.stateApi.onPosChanged);
       this.stage.add(this.inpaintMask.konvaLayer);
     }
     const toolState = this.stateApi.getToolState();
@@ -211,7 +226,7 @@ export class KonvaNodeManager {
     for (const entity of entities) {
       let adapter = this.controlAdapters.get(entity.id);
       if (!adapter) {
-        adapter = new KonvaControlAdapter(entity);
+        adapter = new CanvasControlAdapter(entity);
         this.controlAdapters.set(adapter.id, adapter);
         this.stage.add(adapter.konvaLayer);
       }
@@ -239,18 +254,18 @@ export class KonvaNodeManager {
     this.preview.konvaLayer.zIndex(++zIndex);
   }
 
-  renderDocumentOverlay() {
+  renderDocumentSizeOverlay() {
-    this.preview.renderDocumentOverlay(this.stage, this.stateApi.getDocument());
+    this.preview.documentSizeOverlay.render(this.stage, this.stateApi.getDocument());
   }
 
   renderBbox() {
-    this.preview.renderBbox(this.stateApi.getBbox(), this.stateApi.getToolState());
+    this.preview.bbox.render(this.stateApi.getBbox(), this.stateApi.getToolState());
   }
 
   renderToolPreview() {
-    this.preview.renderToolPreview(
+    this.preview.tool.render(
       this.stage,
-      1,
+      1, // TODO(psyche): this should be renderable entity count
       this.stateApi.getToolState(),
       this.stateApi.getCurrentFill(),
       this.stateApi.getSelectedEntity(),
@@ -261,22 +276,15 @@ export class KonvaNodeManager {
     );
   }
 
-  fitDocumentToStage(): void {
+  renderBackground() {
-    const { getDocument, setStageAttrs } = this.stateApi;
+    this.background.renderBackground(this.stage);
-    const document = getDocument();
-    // Fit & center the document on the stage
-    const width = this.stage.width();
-    const height = this.stage.height();
-    const docWidthWithBuffer = document.width + DOCUMENT_FIT_PADDING_PX * 2;
-    const docHeightWithBuffer = document.height + DOCUMENT_FIT_PADDING_PX * 2;
-    const scale = Math.min(Math.min(width / docWidthWithBuffer, height / docHeightWithBuffer), 1);
-    const x = (width - docWidthWithBuffer * scale) / 2 + DOCUMENT_FIT_PADDING_PX * scale;
-    const y = (height - docHeightWithBuffer * scale) / 2 + DOCUMENT_FIT_PADDING_PX * scale;
-    this.stage.setAttrs({ x, y, width, height, scaleX: scale, scaleY: scale });
-    setStageAttrs({ x, y, width, height, scale });
   }
 
-  fitStageToContainer(): void {
+  fitDocument() {
+    this.preview.documentSizeOverlay.fitToStage(this.stage, this.stateApi.getDocument(), this.stateApi.setStageAttrs);
+  }
+
+  fitStageToContainer() {
     this.stage.width(this.container.offsetWidth);
     this.stage.height(this.container.offsetHeight);
     this.stateApi.setStageAttrs({
@@ -287,11 +295,7 @@ export class KonvaNodeManager {
       scale: this.stage.scaleX(),
     });
     this.renderBackground();
-    this.renderDocumentOverlay();
+    this.renderDocumentSizeOverlay();
-  }
-
-  renderBackground() {
-    this.background.renderBackground(this.stage);
   }
 
   _getMaskLayerClone(): Konva.Layer {

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

@@ -1,31 +0,0 @@
-import type { KonvaNodeManager } from 'features/controlLayers/konva/nodeManager';
-
-/**
- * Gets a function to arrange the entities in the konva stage.
- * @param manager The konva node manager
- * @returns An arrange entities function
- */
-export const getArrangeEntities = (manager: KonvaNodeManager) => {
-  const { getLayersState, getControlAdaptersState, getRegionsState } = manager.stateApi;
-
-  function arrangeEntities(): void {
-    const layers = getLayersState().entities;
-    const controlAdapters = getControlAdaptersState().entities;
-    const regions = getRegionsState().entities;
-    let zIndex = 0;
-    manager.background.layer.zIndex(++zIndex);
-    for (const layer of layers) {
-      manager.get(layer.id)?.konvaLayer.zIndex(++zIndex);
-    }
-    for (const ca of controlAdapters) {
-      manager.get(ca.id)?.konvaLayer.zIndex(++zIndex);
-    }
-    for (const rg of regions) {
-      manager.get(rg.id)?.konvaLayer.zIndex(++zIndex);
-    }
-    manager.get('inpaint_mask')?.konvaLayer.zIndex(++zIndex);
-    manager.preview.konvaLayer.zIndex(++zIndex);
-  }
-
-  return arrangeEntities;
-};

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

@@ -1,6 +1,4 @@
 import { getArbitraryBaseColor } from '@invoke-ai/ui-library';
-import { BACKGROUND_LAYER_ID } from 'features/controlLayers/konva/naming';
-import type { KonvaNodeManager } from 'features/controlLayers/konva/nodeManager';
 import Konva from 'konva';
 
 const baseGridLineColor = getArbitraryBaseColor(27);
@@ -30,98 +28,7 @@ const getGridSpacing = (scale: number): number => {
   return 256;
 };
 
-/**
+export class CanvasBackground {
- * Creates the background konva layer.
- * @returns The background konva layer
- */
-export const createBackgroundLayer = (): Konva.Layer => new Konva.Layer({ id: BACKGROUND_LAYER_ID, listening: false });
-
-/**
- * Gets a render function for the background layer.
- * @param manager The konva node manager
- * @returns A function to render the background grid
- */
-export const getRenderBackground = (manager: KonvaNodeManager) => {
-  function renderBackground(): void {
-    const background = manager.background.layer;
-    background.zIndex(0);
-    const scale = manager.stage.scaleX();
-    const gridSpacing = getGridSpacing(scale);
-    const x = manager.stage.x();
-    const y = manager.stage.y();
-    const width = manager.stage.width();
-    const height = manager.stage.height();
-    const stageRect = {
-      x1: 0,
-      y1: 0,
-      x2: width,
-      y2: height,
-    };
-
-    const gridOffset = {
-      x: Math.ceil(x / scale / gridSpacing) * gridSpacing,
-      y: Math.ceil(y / scale / gridSpacing) * gridSpacing,
-    };
-
-    const gridRect = {
-      x1: -gridOffset.x,
-      y1: -gridOffset.y,
-      x2: width / scale - gridOffset.x + gridSpacing,
-      y2: height / scale - gridOffset.y + gridSpacing,
-    };
-
-    const gridFullRect = {
-      x1: Math.min(stageRect.x1, gridRect.x1),
-      y1: Math.min(stageRect.y1, gridRect.y1),
-      x2: Math.max(stageRect.x2, gridRect.x2),
-      y2: Math.max(stageRect.y2, gridRect.y2),
-    };
-
-    // find the x & y size of the grid
-    const xSize = gridFullRect.x2 - gridFullRect.x1;
-    const ySize = gridFullRect.y2 - gridFullRect.y1;
-    // compute the number of steps required on each axis.
-    const xSteps = Math.round(xSize / gridSpacing) + 1;
-    const ySteps = Math.round(ySize / gridSpacing) + 1;
-
-    const strokeWidth = 1 / scale;
-    let _x = 0;
-    let _y = 0;
-
-    background.destroyChildren();
-
-    for (let i = 0; i < xSteps; i++) {
-      _x = gridFullRect.x1 + i * gridSpacing;
-      background.add(
-        new Konva.Line({
-          x: _x,
-          y: gridFullRect.y1,
-          points: [0, 0, 0, ySize],
-          stroke: _x % 64 ? fineGridLineColor : baseGridLineColor,
-          strokeWidth,
-          listening: false,
-        })
-      );
-    }
-    for (let i = 0; i < ySteps; i++) {
-      _y = gridFullRect.y1 + i * gridSpacing;
-      background.add(
-        new Konva.Line({
-          x: gridFullRect.x1,
-          y: _y,
-          points: [0, 0, xSize, 0],
-          stroke: _y % 64 ? fineGridLineColor : baseGridLineColor,
-          strokeWidth,
-          listening: false,
-        })
-      );
-    }
-  }
-
-  return renderBackground;
-};
-
-export class KonvaBackground {
   konvaLayer: Konva.Layer;
 
   constructor() {

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

@@ -7,7 +7,7 @@ import { v4 as uuidv4 } from 'uuid';
 
 import { KonvaImage } from './objects';
 
-export class KonvaControlAdapter {
+export class CanvasControlAdapter {
   id: string;
   konvaLayer: Konva.Layer;
   konvaObjectGroup: Konva.Group;

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

@@ -9,7 +9,7 @@ import Konva from 'konva';
 import { assert } from 'tsafe';
 import { v4 as uuidv4 } from 'uuid';
 
-export class KonvaInpaintMask {
+export class CanvasInpaintMask {
   id: string;
   konvaLayer: Konva.Layer;
   konvaObjectGroup: Konva.Group;

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

@@ -7,7 +7,7 @@ import Konva from 'konva';
 import { assert } from 'tsafe';
 import { v4 as uuidv4 } from 'uuid';
 
-export class KonvaLayerAdapter {
+export class CanvasLayer {
   id: string;
   konvaLayer: Konva.Layer;
   konvaObjectGroup: Konva.Group;

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

@@ -9,7 +9,7 @@ import Konva from 'konva';
 import { assert } from 'tsafe';
 import { v4 as uuidv4 } from 'uuid';
 
-export class KonvaRegion {
+export class CanvasRegion {
   id: string;
   konvaLayer: Konva.Layer;
   konvaObjectGroup: Konva.Group;

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

@@ -5,9 +5,7 @@ import { $isDebugging } from 'app/store/nanostores/isDebugging';
 import type { RootState } from 'app/store/store';
 import { setStageEventHandlers } from 'features/controlLayers/konva/events';
 import { KonvaNodeManager, setNodeManager } from 'features/controlLayers/konva/nodeManager';
-import { KonvaBackground } from 'features/controlLayers/konva/renderers/background';
 import { updateBboxes } from 'features/controlLayers/konva/renderers/bbox';
-import { KonvaPreview } from 'features/controlLayers/konva/renderers/preview';
 import {
   $stageAttrs,
   bboxChanged,
@@ -74,7 +72,7 @@ export const initializeRenderer = (
    */
   const logIfDebugging = (message: string) => {
     if ($isDebugging.get()) {
-      _log.trace(message);
+      _log.debug(message);
     }
   };
 
@@ -338,22 +336,6 @@ export const initializeRenderer = (
   setNodeManager(manager);
   console.log(manager);
 
-  manager.background = new KonvaBackground();
-  manager.stage.add(manager.background.konvaLayer);
-  manager.preview = new KonvaPreview({
-    stage,
-    getBbox,
-    onBboxTransformed,
-    getShiftKey: $shift.get,
-    getCtrlKey: $ctrl.get,
-    getMetaKey: $meta.get,
-    getAltKey: $alt.get,
-  });
-  manager.preview.konvaLayer.add(manager.preview.bbox.group);
-  manager.preview.konvaLayer.add(manager.preview.tool.group);
-  manager.preview.konvaLayer.add(manager.preview.documentOverlay.group);
-  manager.stage.add(manager.preview.konvaLayer);
-
   const cleanupListeners = setStageEventHandlers(manager);
 
   // Calculating bounding boxes is expensive, must be debounced to not block the UI thread during a user interaction.
@@ -408,7 +390,7 @@ export const initializeRenderer = (
 
     if (isFirstRender || canvasV2.document !== prevCanvasV2.document) {
       logIfDebugging('Rendering document bounds overlay');
-      manager.renderDocumentOverlay();
+      manager.renderDocumentSizeOverlay();
     }
 
     if (isFirstRender || canvasV2.bbox !== prevCanvasV2.bbox || canvasV2.tool.selected !== prevCanvasV2.tool.selected) {
@@ -447,7 +429,7 @@ export const initializeRenderer = (
 
   // We can use a resize observer to ensure the stage always fits the container. We also need to re-render the bg and
   // document bounds overlay when the stage is resized.
-  const resizeObserver = new ResizeObserver(manager.fitStageToContainer);
+  const resizeObserver = new ResizeObserver(manager.fitStageToContainer.bind(manager));
   resizeObserver.observe(container);
   manager.fitStageToContainer();
 
@@ -455,7 +437,8 @@ export const initializeRenderer = (
 
   logIfDebugging('First render of konva stage');
   // On first render, the document should be fit to the stage.
-  manager.fitDocumentToStage();
+  manager.renderDocumentSizeOverlay();
+  manager.fitDocument();
   manager.renderToolPreview();
   renderCanvas();
 

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

@@ -1,52 +0,0 @@
-import { DOCUMENT_FIT_PADDING_PX } from 'features/controlLayers/konva/constants';
-import type { KonvaNodeManager } from 'features/controlLayers/konva/nodeManager';
-
-/**
- * Gets a function to fit the document to the stage, resetting the stage scale to 100%.
- * If the document is smaller than the stage, the stage scale is increased to fit the document.
- * @param manager The konva node manager
- * @returns A function to fit the document to the stage
- */
-export const getFitDocumentToStage = (manager: KonvaNodeManager) => {
-  function fitDocumentToStage(): void {
-    const { getDocument, setStageAttrs } = manager.stateApi;
-    const document = getDocument();
-    // Fit & center the document on the stage
-    const width = manager.stage.width();
-    const height = manager.stage.height();
-    const docWidthWithBuffer = document.width + DOCUMENT_FIT_PADDING_PX * 2;
-    const docHeightWithBuffer = document.height + DOCUMENT_FIT_PADDING_PX * 2;
-    const scale = Math.min(Math.min(width / docWidthWithBuffer, height / docHeightWithBuffer), 1);
-    const x = (width - docWidthWithBuffer * scale) / 2 + DOCUMENT_FIT_PADDING_PX * scale;
-    const y = (height - docHeightWithBuffer * scale) / 2 + DOCUMENT_FIT_PADDING_PX * scale;
-    manager.stage.setAttrs({ x, y, width, height, scaleX: scale, scaleY: scale });
-    setStageAttrs({ x, y, width, height, scale });
-  }
-
-  return fitDocumentToStage;
-};
-
-/**
- * Gets a function to fit the stage to its container element. Called during resize events.
- * @param manager The konva node manager
- * @returns A function to fit the stage to its container
- */
-export const getFitStageToContainer = (manager: KonvaNodeManager) => {
-  const { stage, container } = manager;
-  const { setStageAttrs } = manager.stateApi;
-  function fitStageToContainer(): void {
-    stage.width(container.offsetWidth);
-    stage.height(container.offsetHeight);
-    setStageAttrs({
-      x: stage.x(),
-      y: stage.y(),
-      width: stage.width(),
-      height: stage.height(),
-      scale: stage.scaleX(),
-    });
-    manager.konvaApi.renderBackground();
-    manager.renderDocumentOverlay();
-  }
-
-  return fitStageToContainer;
-};