tidy(ui): organise files

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

@@ -1,6 +1,5 @@
 import { Flex } from '@invoke-ai/ui-library';
 import { logger } from 'app/logging/logger';
-import { $isDebugging } from 'app/store/nanostores/isDebugging';
 import { useAppStore } from 'app/store/storeHooks';
 import { HeadsUpDisplay } from 'features/controlLayers/components/HeadsUpDisplay';
 import { CanvasManager, setCanvasManager } from 'features/controlLayers/konva/CanvasManager';
@@ -9,7 +8,7 @@ import { memo, useCallback, useEffect, useLayoutEffect, useState } from 'react';
 import { useDevicePixelRatio } from 'use-device-pixel-ratio';
 import { v4 as uuidv4 } from 'uuid';
 
-const log = logger('konva');
+const log = logger('canvas');
 
 // This will log warnings when layers > 5 - maybe use `import.meta.env.MODE === 'development'` instead?
 Konva.showWarnings = false;
@@ -19,23 +18,14 @@ const useStageRenderer = (stage: Konva.Stage, container: HTMLDivElement | null,
   const dpr = useDevicePixelRatio({ round: false });
 
   useLayoutEffect(() => {
-    /**
+    log.debug('Initializing renderer');
-     * Logs a message to the console if debugging is enabled.
-     */
-    const logIfDebugging = (message: string) => {
-      if ($isDebugging.get()) {
-        log.debug(message);
-      }
-    };
-
-    logIfDebugging('Initializing renderer');
     if (!container) {
       // Nothing to clean up
-      logIfDebugging('No stage container, skipping initialization');
+      log.debug('No stage container, skipping initialization');
       return () => {};
     }
 
-    const manager = new CanvasManager(stage, container, store, logIfDebugging);
+    const manager = new CanvasManager(stage, container, store);
     setCanvasManager(manager);
     const cleanup = manager.initialize();
     return cleanup;

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

@@ -1,9 +1,14 @@
 import type { Store } from '@reduxjs/toolkit';
+import { logger } from 'app/logging/logger';
 import type { RootState } from 'app/store/store';
-import { getImageDataTransparency } from 'common/util/arrayBuffer';
+import {
+  getGenerationMode,
+  getImageSourceImage,
+  getInpaintMaskImage,
+  getRegionMaskImage,
+} from 'features/controlLayers/konva/util';
 import { $lastProgressEvent, $shouldShowStagedImage } from 'features/controlLayers/store/canvasV2Slice';
-import type { CanvasV2State, GenerationMode, Rect } from 'features/controlLayers/store/types';
+import type { CanvasV2State } from 'features/controlLayers/store/types';
-import { isValidLayer } from 'features/nodes/util/graph/generation/addLayers';
 import type Konva from 'konva';
 import { atom } from 'nanostores';
 import { getImageDTO as defaultGetImageDTO, uploadImage as defaultUploadImage } from 'services/api/endpoints/images';
@@ -19,10 +24,11 @@ import { CanvasLayer } from './CanvasLayer';
 import { CanvasPreview } from './CanvasPreview';
 import { CanvasRegion } from './CanvasRegion';
 import { CanvasStagingArea } from './CanvasStagingArea';
+import { CanvasStateApi } from './CanvasStateApi';
 import { CanvasTool } from './CanvasTool';
 import { setStageEventHandlers } from './events';
-import { StateApi } from './StateApi';
+
-import { konvaNodeToBlob, konvaNodeToImageData, previewBlob } from './util';
+const log = logger('canvas');
 
 type Util = {
   getImageDTO: (imageName: string) => Promise<ImageDTO | null>;
@@ -52,27 +58,24 @@ export class CanvasManager {
   regions: Map<string, CanvasRegion>;
   inpaintMask: CanvasInpaintMask;
   util: Util;
-  stateApi: StateApi;
+  stateApi: CanvasStateApi;
   preview: CanvasPreview;
   background: CanvasBackground;
   private store: Store<RootState>;
   private isFirstRender: boolean;
   private prevState: CanvasV2State;
-  private log: (message: string) => void;
 
   constructor(
     stage: Konva.Stage,
     container: HTMLDivElement,
     store: Store<RootState>,
-    log: (message: string) => void,
     getImageDTO: Util['getImageDTO'] = defaultGetImageDTO,
     uploadImage: Util['uploadImage'] = defaultUploadImage
   ) {
-    this.log = log;
     this.stage = stage;
     this.container = container;
     this.store = store;
-    this.stateApi = new StateApi(this.store, this.log);
+    this.stateApi = new CanvasStateApi(this.store);
     this.prevState = this.stateApi.getState();
     this.isFirstRender = true;
 
@@ -207,7 +210,7 @@ export class CanvasManager {
     const state = this.stateApi.getState();
 
     if (this.prevState === state && !this.isFirstRender) {
-      this.log('No changes detected, skipping render');
+      log.debug('No changes detected, skipping render');
       return;
     }
 
@@ -217,7 +220,7 @@ export class CanvasManager {
       state.tool.selected !== this.prevState.tool.selected ||
       state.selectedEntityIdentifier?.id !== this.prevState.selectedEntityIdentifier?.id
     ) {
-      this.log('Rendering layers');
+      log.debug('Rendering layers');
       this.renderLayers();
     }
 
@@ -228,7 +231,7 @@ export class CanvasManager {
       state.tool.selected !== this.prevState.tool.selected ||
       state.selectedEntityIdentifier?.id !== this.prevState.selectedEntityIdentifier?.id
     ) {
-      this.log('Rendering regions');
+      log.debug('Rendering regions');
       this.renderRegions();
     }
 
@@ -239,7 +242,7 @@ export class CanvasManager {
       state.tool.selected !== this.prevState.tool.selected ||
       state.selectedEntityIdentifier?.id !== this.prevState.selectedEntityIdentifier?.id
     ) {
-      this.log('Rendering inpaint mask');
+      log.debug('Rendering inpaint mask');
       this.renderInpaintMask();
     }
 
@@ -248,12 +251,12 @@ export class CanvasManager {
       state.controlAdapters.entities !== this.prevState.controlAdapters.entities ||
       state.selectedEntityIdentifier?.id !== this.prevState.selectedEntityIdentifier?.id
     ) {
-      this.log('Rendering control adapters');
+      log.debug('Rendering control adapters');
       this.renderControlAdapters();
     }
 
     if (this.isFirstRender || state.document !== this.prevState.document) {
-      this.log('Rendering document bounds overlay');
+      log.debug('Rendering document bounds overlay');
       this.preview.documentSizeOverlay.render();
     }
 
@@ -262,7 +265,7 @@ export class CanvasManager {
       state.bbox !== this.prevState.bbox ||
       state.tool.selected !== this.prevState.tool.selected
     ) {
-      this.log('Rendering generation bbox');
+      log.debug('Rendering generation bbox');
       this.preview.bbox.render();
     }
 
@@ -272,12 +275,12 @@ export class CanvasManager {
       state.controlAdapters !== this.prevState.controlAdapters ||
       state.regions !== this.prevState.regions
     ) {
-      // this.log('Updating entity bboxes');
+      // log.debug('Updating entity bboxes');
       // debouncedUpdateBboxes(stage, canvasV2.layers, canvasV2.controlAdapters, canvasV2.regions, onBboxChanged);
     }
 
     if (this.isFirstRender || state.stagingArea !== this.prevState.stagingArea) {
-      this.log('Rendering staging area');
+      log.debug('Rendering staging area');
       this.preview.stagingArea.render();
     }
 
@@ -289,7 +292,7 @@ export class CanvasManager {
       state.inpaintMask !== this.prevState.inpaintMask ||
       state.selectedEntityIdentifier?.id !== this.prevState.selectedEntityIdentifier?.id
     ) {
-      this.log('Arranging entities');
+      log.debug('Arranging entities');
       this.arrangeEntities();
     }
 
@@ -301,7 +304,7 @@ export class CanvasManager {
   };
 
   initialize = () => {
-    this.log('Initializing renderer');
+    log.debug('Initializing renderer');
     this.stage.container(this.container);
 
     const cleanupListeners = setStageEventHandlers(this);
@@ -316,18 +319,18 @@ export class CanvasManager {
 
     // When we this flag, we need to render the staging area
     $shouldShowStagedImage.subscribe((shouldShowStagedImage, prevShouldShowStagedImage) => {
-      this.log('Rendering staging area');
+      log.debug('Rendering staging area');
       if (shouldShowStagedImage !== prevShouldShowStagedImage) {
         this.preview.stagingArea.render();
       }
     });
 
     $lastProgressEvent.subscribe(() => {
-      this.log('Rendering staging area');
+      log.debug('Rendering staging area');
       this.preview.stagingArea.render();
     });
 
-    this.log('First render of konva stage');
+    log.debug('First render of konva stage');
     // On first render, the document should be fit to the stage.
     this.preview.documentSizeOverlay.render();
     this.preview.documentSizeOverlay.fitToStage();
@@ -335,7 +338,7 @@ export class CanvasManager {
     this.render();
 
     return () => {
-      this.log('Cleaning up konva renderer');
+      log.debug('Cleaning up konva renderer');
       unsubscribeRenderer();
       cleanupListeners();
       $shouldShowStagedImage.off();
@@ -343,164 +346,19 @@ export class CanvasManager {
     };
   };
 
-  getInpaintMaskLayerClone(): Konva.Layer {
+  getGenerationMode() {
-    const layerClone = this.inpaintMask.layer.clone();
+    return getGenerationMode({ manager: this });
-    const objectGroupClone = this.inpaintMask.group.clone();
-
-    layerClone.destroyChildren();
-    layerClone.add(objectGroupClone);
-
-    objectGroupClone.opacity(1);
-    objectGroupClone.cache();
-
-    return layerClone;
   }
 
-  getRegionMaskLayerClone(arg: { id: string }): Konva.Layer {
+  getRegionMaskImage(arg: Omit<Parameters<typeof getRegionMaskImage>[0], 'manager'>) {
-    const { id } = arg;
+    return getRegionMaskImage({ ...arg, manager: this });
-
-    const canvasRegion = this.regions.get(id);
-    assert(canvasRegion, `Canvas region with id ${id} not found`);
-
-    const layerClone = canvasRegion.layer.clone();
-    const objectGroupClone = canvasRegion.group.clone();
-
-    layerClone.destroyChildren();
-    layerClone.add(objectGroupClone);
-
-    objectGroupClone.opacity(1);
-    objectGroupClone.cache();
-
-    return layerClone;
   }
 
-  getCompositeLayerStageClone(): Konva.Stage {
+  getInpaintMaskImage(arg: Omit<Parameters<typeof getInpaintMaskImage>[0], 'manager'>) {
-    const layersState = this.stateApi.getLayersState();
+    return getInpaintMaskImage({ ...arg, manager: this });
-
-    const stageClone = this.stage.clone();
-
-    stageClone.scaleX(1);
-    stageClone.scaleY(1);
-    stageClone.x(0);
-    stageClone.y(0);
-
-    const validLayers = layersState.entities.filter(isValidLayer);
-
-    // Konva bug (?) - when iterating over the array returned from `stage.getLayers()`, if you destroy a layer, the array
-    // is mutated in-place and the next iteration will skip the next layer. To avoid this, we first collect the layers
-    // to delete in a separate array and then destroy them.
-    // TODO(psyche): Maybe report this?
-    const toDelete: Konva.Layer[] = [];
-
-    for (const konvaLayer of stageClone.getLayers()) {
-      const layer = validLayers.find((l) => l.id === konvaLayer.id());
-      if (!layer) {
-        toDelete.push(konvaLayer);
-      }
   }
 
-    for (const konvaLayer of toDelete) {
+  getImageSourceImage(arg: Omit<Parameters<typeof getImageSourceImage>[0], 'manager'>) {
-      konvaLayer.destroy();
+    return getImageSourceImage({ ...arg, manager: this });
-    }
-
-    return stageClone;
-  }
-
-  getGenerationMode(): GenerationMode {
-    const { x, y, width, height } = this.stateApi.getBbox();
-    const inpaintMaskLayer = this.getInpaintMaskLayerClone();
-    const inpaintMaskImageData = konvaNodeToImageData(inpaintMaskLayer, { x, y, width, height });
-    const inpaintMaskTransparency = getImageDataTransparency(inpaintMaskImageData);
-    const compositeLayer = this.getCompositeLayerStageClone();
-    const compositeLayerImageData = konvaNodeToImageData(compositeLayer, { x, y, width, height });
-    const compositeLayerTransparency = getImageDataTransparency(compositeLayerImageData);
-    if (compositeLayerTransparency.isPartiallyTransparent) {
-      if (compositeLayerTransparency.isFullyTransparent) {
-        return 'txt2img';
-      }
-      return 'outpaint';
-    } else {
-      if (!inpaintMaskTransparency.isFullyTransparent) {
-        return 'inpaint';
-      }
-      return 'img2img';
-    }
-  }
-
-  async getRegionMaskImage(arg: { id: string; bbox?: Rect; preview?: boolean }): Promise<ImageDTO> {
-    const { id, bbox, preview = false } = arg;
-    const region = this.stateApi.getRegionsState().entities.find((entity) => entity.id === id);
-    assert(region, `Region entity state with id ${id} not found`);
-
-    // if (region.imageCache) {
-    //   const imageDTO = await this.util.getImageDTO(region.imageCache.name);
-    //   if (imageDTO) {
-    //     return imageDTO;
-    //   }
-    // }
-
-    const layerClone = this.getRegionMaskLayerClone({ id });
-    const blob = await konvaNodeToBlob(layerClone, bbox);
-
-    if (preview) {
-      previewBlob(blob, `region ${region.id} mask`);
-    }
-
-    layerClone.destroy();
-
-    const imageDTO = await this.util.uploadImage(blob, `${region.id}_mask.png`, 'mask', true);
-    this.stateApi.onRegionMaskImageCached(region.id, imageDTO);
-    return imageDTO;
-  }
-
-  async getInpaintMaskImage(arg: { bbox?: Rect; preview?: boolean }): Promise<ImageDTO> {
-    const { bbox, preview = false } = arg;
-    // const inpaintMask = this.stateApi.getInpaintMaskState();
-
-    // if (inpaintMask.imageCache) {
-    //   const imageDTO = await this.util.getImageDTO(inpaintMask.imageCache.name);
-    //   if (imageDTO) {
-    //     return imageDTO;
-    //   }
-    // }
-
-    const layerClone = this.getInpaintMaskLayerClone();
-    const blob = await konvaNodeToBlob(layerClone, bbox);
-
-    if (preview) {
-      previewBlob(blob, 'inpaint mask');
-    }
-
-    layerClone.destroy();
-
-    const imageDTO = await this.util.uploadImage(blob, 'inpaint_mask.png', 'mask', true);
-    this.stateApi.onInpaintMaskImageCached(imageDTO);
-    return imageDTO;
-  }
-
-  async getImageSourceImage(arg: { bbox?: Rect; preview?: boolean }): Promise<ImageDTO> {
-    const { bbox, preview = false } = arg;
-    // const { imageCache } = this.stateApi.getLayersState();
-
-    // if (imageCache) {
-    //   const imageDTO = await this.util.getImageDTO(imageCache.name);
-    //   if (imageDTO) {
-    //     return imageDTO;
-    //   }
-    // }
-
-    const stageClone = this.getCompositeLayerStageClone();
-
-    const blob = await konvaNodeToBlob(stageClone, bbox);
-
-    if (preview) {
-      previewBlob(blob, 'image source');
-    }
-
-    stageClone.destroy();
-
-    const imageDTO = await this.util.uploadImage(blob, 'base_layer.png', 'general', true);
-    this.stateApi.onLayerImageCached(imageDTO);
-    return imageDTO;
   }
 }

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

@@ -1,20 +1,71 @@
-import { $alt, $ctrl, $meta, $shift } from "@invoke-ai/ui-library";
+import { $alt, $ctrl, $meta, $shift } from '@invoke-ai/ui-library';
-import type { Store } from "@reduxjs/toolkit";
+import type { Store } from '@reduxjs/toolkit';
-import type { RootState } from "app/store/store";
+import { logger } from 'app/logging/logger';
-import { $isDrawing, $isMouseDown, $lastAddedPoint, $lastCursorPos, $lastMouseDownPos, $lastProgressEvent, $shouldShowStagedImage, $spaceKey, $stageAttrs, bboxChanged, brushWidthChanged, caBboxChanged, caTranslated, eraserWidthChanged, imBboxChanged, imBrushLineAdded, imEraserLineAdded, imImageCacheChanged, imLinePointAdded, imRectAdded, imScaled, imTranslated, layerBboxChanged, layerBrushLineAdded, layerEraserLineAdded, layerImageCacheChanged, layerLinePointAdded, layerRectAdded, layerScaled, layerTranslated, rgBboxChanged, rgBrushLineAdded, rgEraserLineAdded, rgImageCacheChanged, rgLinePointAdded, rgRectAdded, rgScaled, rgTranslated, toolBufferChanged, toolChanged } from "features/controlLayers/store/canvasV2Slice";
+import type { RootState } from 'app/store/store';
-import type { BboxChangedArg, BrushLineAddedArg, CanvasEntity, EraserLineAddedArg, PointAddedToLineArg, PosChangedArg, RectShapeAddedArg, ScaleChangedArg, Tool } from "features/controlLayers/store/types";
+import {
-import type { IRect } from "konva/lib/types";
+  $isDrawing,
-import type { RgbaColor } from "react-colorful";
+  $isMouseDown,
-import type { ImageDTO } from "services/api/types";
+  $lastAddedPoint,
+  $lastCursorPos,
+  $lastMouseDownPos,
+  $lastProgressEvent,
+  $shouldShowStagedImage,
+  $spaceKey,
+  $stageAttrs,
+  bboxChanged,
+  brushWidthChanged,
+  caBboxChanged,
+  caTranslated,
+  eraserWidthChanged,
+  imBboxChanged,
+  imBrushLineAdded,
+  imEraserLineAdded,
+  imImageCacheChanged,
+  imLinePointAdded,
+  imRectAdded,
+  imScaled,
+  imTranslated,
+  layerBboxChanged,
+  layerBrushLineAdded,
+  layerEraserLineAdded,
+  layerImageCacheChanged,
+  layerLinePointAdded,
+  layerRectAdded,
+  layerScaled,
+  layerTranslated,
+  rgBboxChanged,
+  rgBrushLineAdded,
+  rgEraserLineAdded,
+  rgImageCacheChanged,
+  rgLinePointAdded,
+  rgRectAdded,
+  rgScaled,
+  rgTranslated,
+  toolBufferChanged,
+  toolChanged,
+} from 'features/controlLayers/store/canvasV2Slice';
+import type {
+  BboxChangedArg,
+  BrushLineAddedArg,
+  CanvasEntity,
+  EraserLineAddedArg,
+  PointAddedToLineArg,
+  PosChangedArg,
+  RectShapeAddedArg,
+  ScaleChangedArg,
+  Tool,
+} from 'features/controlLayers/store/types';
+import type { IRect } from 'konva/lib/types';
+import type { RgbaColor } from 'react-colorful';
+import type { ImageDTO } from 'services/api/types';
 
+const log = logger('canvas');
 
-export class StateApi {
+export class CanvasStateApi {
   private store: Store<RootState>;
-  private log: (message: string) => void;
 
-  constructor(store: Store<RootState>, log: (message: string) => void) {
+  constructor(store: Store<RootState>) {
     this.store = store;
-    this.log = log;
   }
 
   // Reminder - use arrow functions to avoid binding issues
@@ -23,7 +74,7 @@ export class StateApi {
   };
 
   onPosChanged = (arg: PosChangedArg, entityType: CanvasEntity['type']) => {
-    this.log('onPosChanged');
+    log.debug('onPosChanged');
     if (entityType === 'layer') {
       this.store.dispatch(layerTranslated(arg));
     } else if (entityType === 'control_adapter') {
@@ -35,7 +86,7 @@ export class StateApi {
     }
   };
   onScaleChanged = (arg: ScaleChangedArg, entityType: CanvasEntity['type']) => {
-    this.log('onScaleChanged');
+    log.debug('onScaleChanged');
     if (entityType === 'layer') {
       this.store.dispatch(layerScaled(arg));
     } else if (entityType === 'inpaint_mask') {
@@ -45,7 +96,7 @@ export class StateApi {
     }
   };
   onBboxChanged = (arg: BboxChangedArg, entityType: CanvasEntity['type']) => {
-    this.log('Entity bbox changed');
+    log.debug('Entity bbox changed');
     if (entityType === 'layer') {
       this.store.dispatch(layerBboxChanged(arg));
     } else if (entityType === 'control_adapter') {
@@ -57,7 +108,7 @@ export class StateApi {
     }
   };
   onBrushLineAdded = (arg: BrushLineAddedArg, entityType: CanvasEntity['type']) => {
-    this.log('Brush line added');
+    log.debug('Brush line added');
     if (entityType === 'layer') {
       this.store.dispatch(layerBrushLineAdded(arg));
     } else if (entityType === 'regional_guidance') {
@@ -67,7 +118,7 @@ export class StateApi {
     }
   };
   onEraserLineAdded = (arg: EraserLineAddedArg, entityType: CanvasEntity['type']) => {
-    this.log('Eraser line added');
+    log.debug('Eraser line added');
     if (entityType === 'layer') {
       this.store.dispatch(layerEraserLineAdded(arg));
     } else if (entityType === 'regional_guidance') {
@@ -77,7 +128,7 @@ export class StateApi {
     }
   };
   onPointAddedToLine = (arg: PointAddedToLineArg, entityType: CanvasEntity['type']) => {
-    this.log('Point added to line');
+    log.debug('Point added to line');
     if (entityType === 'layer') {
       this.store.dispatch(layerLinePointAdded(arg));
     } else if (entityType === 'regional_guidance') {
@@ -87,7 +138,7 @@ export class StateApi {
     }
   };
   onRectShapeAdded = (arg: RectShapeAddedArg, entityType: CanvasEntity['type']) => {
-    this.log('Rect shape added');
+    log.debug('Rect shape added');
     if (entityType === 'layer') {
       this.store.dispatch(layerRectAdded(arg));
     } else if (entityType === 'regional_guidance') {
@@ -97,35 +148,35 @@ export class StateApi {
     }
   };
   onBboxTransformed = (bbox: IRect) => {
-    this.log('Generation bbox transformed');
+    log.debug('Generation bbox transformed');
     this.store.dispatch(bboxChanged(bbox));
   };
   onBrushWidthChanged = (width: number) => {
-    this.log('Brush width changed');
+    log.debug('Brush width changed');
     this.store.dispatch(brushWidthChanged(width));
   };
   onEraserWidthChanged = (width: number) => {
-    this.log('Eraser width changed');
+    log.debug('Eraser width changed');
     this.store.dispatch(eraserWidthChanged(width));
   };
   onRegionMaskImageCached = (id: string, imageDTO: ImageDTO) => {
-    this.log('Region mask image cached');
+    log.debug('Region mask image cached');
     this.store.dispatch(rgImageCacheChanged({ id, imageDTO }));
   };
   onInpaintMaskImageCached = (imageDTO: ImageDTO) => {
-    this.log('Inpaint mask image cached');
+    log.debug('Inpaint mask image cached');
     this.store.dispatch(imImageCacheChanged({ imageDTO }));
   };
   onLayerImageCached = (imageDTO: ImageDTO) => {
-    this.log('Layer image cached');
+    log.debug('Layer image cached');
     this.store.dispatch(layerImageCacheChanged({ imageDTO }));
   };
   setTool = (tool: Tool) => {
-    this.log('Tool selection changed');
+    log.debug('Tool selection changed');
     this.store.dispatch(toolChanged(tool));
   };
   setToolBuffer = (toolBuffer: Tool | null) => {
-    this.log('Tool buffer changed');
+    log.debug('Tool buffer changed');
     this.store.dispatch(toolBufferChanged(toolBuffer));
   };
 

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

@@ -1,3 +1,5 @@
+import { getImageDataTransparency } from 'common/util/arrayBuffer';
+import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import {
   CA_LAYER_NAME,
   INPAINT_MASK_LAYER_ID,
@@ -11,10 +13,12 @@ import {
   RG_LAYER_NAME,
   RG_LAYER_RECT_SHAPE_NAME,
 } from 'features/controlLayers/konva/naming';
-import type { Rect, RgbaColor } from 'features/controlLayers/store/types';
+import type { GenerationMode, Rect, RgbaColor } from 'features/controlLayers/store/types';
+import { isValidLayer } from 'features/nodes/util/graph/generation/addLayers';
 import Konva from 'konva';
 import type { KonvaEventObject } from 'konva/lib/Node';
 import type { Vector2d } from 'konva/lib/types';
+import type { ImageDTO } from 'services/api/types';
 import { assert } from 'tsafe';
 
 /**
@@ -282,3 +286,181 @@ export const previewBlob = async (blob: Blob, label?: string) => {
   }
   w.document.write(`<img src="${url}" style="border: 1px solid red;" />`);
 };
+
+export function getInpaintMaskLayerClone(arg: { manager: CanvasManager }): Konva.Layer {
+  const { manager } = arg;
+  const layerClone = manager.inpaintMask.layer.clone();
+  const objectGroupClone = manager.inpaintMask.group.clone();
+
+  layerClone.destroyChildren();
+  layerClone.add(objectGroupClone);
+
+  objectGroupClone.opacity(1);
+  objectGroupClone.cache();
+
+  return layerClone;
+}
+
+export function getRegionMaskLayerClone(arg: { manager: CanvasManager; id: string }): Konva.Layer {
+  const { id, manager } = arg;
+
+  const canvasRegion = manager.regions.get(id);
+  assert(canvasRegion, `Canvas region with id ${id} not found`);
+
+  const layerClone = canvasRegion.layer.clone();
+  const objectGroupClone = canvasRegion.group.clone();
+
+  layerClone.destroyChildren();
+  layerClone.add(objectGroupClone);
+
+  objectGroupClone.opacity(1);
+  objectGroupClone.cache();
+
+  return layerClone;
+}
+
+export function getCompositeLayerStageClone(arg: { manager: CanvasManager }): Konva.Stage {
+  const { manager } = arg;
+
+  const layersState = manager.stateApi.getLayersState();
+
+  const stageClone = manager.stage.clone();
+
+  stageClone.scaleX(1);
+  stageClone.scaleY(1);
+  stageClone.x(0);
+  stageClone.y(0);
+
+  const validLayers = layersState.entities.filter(isValidLayer);
+
+  // Konva bug (?) - when iterating over the array returned from `stage.getLayers()`, if you destroy a layer, the array
+  // is mutated in-place and the next iteration will skip the next layer. To avoid this, we first collect the layers
+  // to delete in a separate array and then destroy them.
+  // TODO(psyche): Maybe report this?
+  const toDelete: Konva.Layer[] = [];
+
+  for (const konvaLayer of stageClone.getLayers()) {
+    const layer = validLayers.find((l) => l.id === konvaLayer.id());
+    if (!layer) {
+      toDelete.push(konvaLayer);
+    }
+  }
+
+  for (const konvaLayer of toDelete) {
+    konvaLayer.destroy();
+  }
+
+  return stageClone;
+}
+
+export function getGenerationMode(arg: { manager: CanvasManager }): GenerationMode {
+  const { manager } = arg;
+  const { x, y, width, height } = manager.stateApi.getBbox();
+  const inpaintMaskLayer = getInpaintMaskLayerClone(arg);
+  const inpaintMaskImageData = konvaNodeToImageData(inpaintMaskLayer, { x, y, width, height });
+  const inpaintMaskTransparency = getImageDataTransparency(inpaintMaskImageData);
+  const compositeLayer = getCompositeLayerStageClone(arg);
+  const compositeLayerImageData = konvaNodeToImageData(compositeLayer, { x, y, width, height });
+  const compositeLayerTransparency = getImageDataTransparency(compositeLayerImageData);
+  if (compositeLayerTransparency.isPartiallyTransparent) {
+    if (compositeLayerTransparency.isFullyTransparent) {
+      return 'txt2img';
+    }
+    return 'outpaint';
+  } else {
+    if (!inpaintMaskTransparency.isFullyTransparent) {
+      return 'inpaint';
+    }
+    return 'img2img';
+  }
+}
+
+export async function getRegionMaskImage(arg: {
+  manager: CanvasManager;
+  id: string;
+  bbox?: Rect;
+  preview?: boolean;
+}): Promise<ImageDTO> {
+  const { manager, id, bbox, preview = false } = arg;
+  const region = manager.stateApi.getRegionsState().entities.find((entity) => entity.id === id);
+  assert(region, `Region entity state with id ${id} not found`);
+
+  // if (region.imageCache) {
+  //   const imageDTO = await this.util.getImageDTO(region.imageCache.name);
+  //   if (imageDTO) {
+  //     return imageDTO;
+  //   }
+  // }
+
+  const layerClone = getRegionMaskLayerClone({ id, manager });
+  const blob = await konvaNodeToBlob(layerClone, bbox);
+
+  if (preview) {
+    previewBlob(blob, `region ${region.id} mask`);
+  }
+
+  layerClone.destroy();
+
+  const imageDTO = await manager.util.uploadImage(blob, `${region.id}_mask.png`, 'mask', true);
+  manager.stateApi.onRegionMaskImageCached(region.id, imageDTO);
+  return imageDTO;
+}
+
+export async function getInpaintMaskImage(arg: {
+  manager: CanvasManager;
+  bbox?: Rect;
+  preview?: boolean;
+}): Promise<ImageDTO> {
+  const { manager, bbox, preview = false } = arg;
+  // const inpaintMask = this.stateApi.getInpaintMaskState();
+
+  // if (inpaintMask.imageCache) {
+  //   const imageDTO = await this.util.getImageDTO(inpaintMask.imageCache.name);
+  //   if (imageDTO) {
+  //     return imageDTO;
+  //   }
+  // }
+
+  const layerClone = getInpaintMaskLayerClone({ manager });
+  const blob = await konvaNodeToBlob(layerClone, bbox);
+
+  if (preview) {
+    previewBlob(blob, 'inpaint mask');
+  }
+
+  layerClone.destroy();
+
+  const imageDTO = await manager.util.uploadImage(blob, 'inpaint_mask.png', 'mask', true);
+  manager.stateApi.onInpaintMaskImageCached(imageDTO);
+  return imageDTO;
+}
+
+export async function getImageSourceImage(arg: {
+  manager: CanvasManager;
+  bbox?: Rect;
+  preview?: boolean;
+}): Promise<ImageDTO> {
+  const { manager, bbox, preview = false } = arg;
+  // const { imageCache } = this.stateApi.getLayersState();
+
+  // if (imageCache) {
+  //   const imageDTO = await this.util.getImageDTO(imageCache.name);
+  //   if (imageDTO) {
+  //     return imageDTO;
+  //   }
+  // }
+
+  const stageClone = getCompositeLayerStageClone({ manager });
+
+  const blob = await konvaNodeToBlob(stageClone, bbox);
+
+  if (preview) {
+    previewBlob(blob, 'image source');
+  }
+
+  stageClone.destroy();
+
+  const imageDTO = await manager.util.uploadImage(blob, 'base_layer.png', 'general', true);
+  manager.stateApi.onLayerImageCached(imageDTO);
+  return imageDTO;
+}