feat(ui): txt2img, img2img, inpaint & outpaint working

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

@@ -21,7 +21,7 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
       assert(manager, 'No model found in state');
 
       let didStartStaging = false;
-      if (!state.canvasV2.session.isStaging && state.canvasV2.session.isActive) {
+      if (!state.canvasV2.session.isStaging) {
         dispatch(sessionStartedStaging());
         didStartStaging = true;
       }
@@ -49,7 +49,8 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
         );
         req.reset();
         await req.unwrap();
-      } catch {
+      } catch (error) {
+        console.log('Error in enqueueRequestedLinear', error);
         if (didStartStaging && getState().canvasV2.session.isStaging) {
           dispatch(sessionStagingAreaReset());
         }

invokeai/frontend/web/src/common/util/arrayBuffer.ts

@@ -1,16 +0,0 @@
-export const getImageDataTransparency = (imageData: ImageData) => {
-  let isFullyTransparent = true;
-  let isPartiallyTransparent = false;
-  const len = imageData.data.length;
-  for (let i = 3; i < len; i += 4) {
-    if (imageData.data[i] !== 0) {
-      isFullyTransparent = false;
-    } else {
-      isPartiallyTransparent = true;
-    }
-    if (!isFullyTransparent && isPartiallyTransparent) {
-      return { isFullyTransparent, isPartiallyTransparent };
-    }
-  }
-  return { isFullyTransparent, isPartiallyTransparent };
-};

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

@@ -28,7 +28,6 @@ export class CanvasLayerAdapter {
   renderer: CanvasObjectRenderer;
 
   isFirstRender: boolean = true;
-  bboxNeedsUpdate: boolean = true;
 
   constructor(state: CanvasLayerAdapter['state'], manager: CanvasLayerAdapter['manager']) {
     this.id = state.id;
@@ -40,6 +39,8 @@ export class CanvasLayerAdapter {
 
     this.konva = {
       layer: new Konva.Layer({
+        // We need the ID on the layer to help with building the composite initial image
+        // See `getCompositeLayerStageClone()`
         id: this.id,
         name: `${this.type}:layer`,
         listening: false,
@@ -134,7 +135,6 @@ export class CanvasLayerAdapter {
       id: this.id,
       type: this.type,
       state: deepClone(this.state),
-      bboxNeedsUpdate: this.bboxNeedsUpdate,
       transformer: this.transformer.repr(),
       renderer: this.renderer.repr(),
     };

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

@@ -13,32 +13,41 @@ import type { CanvasTransformer } from 'features/controlLayers/konva/CanvasTrans
 import {
   getCompositeLayerImage,
   getControlAdapterImage,
-  getGenerationMode,
+  getImageDataTransparency,
   getInpaintMaskImage,
   getPrefixedId,
   getRegionMaskImage,
+  konvaNodeToBlob,
+  konvaNodeToImageData,
   nanoid,
 } from 'features/controlLayers/konva/util';
 import type { Extents, ExtentsResult, GetBboxTask, WorkerLogMessage } from 'features/controlLayers/konva/worker';
 import { $lastProgressEvent, $shouldShowStagedImage } from 'features/controlLayers/store/canvasV2Slice';
-import {
-  type CanvasControlAdapterState,
-  type CanvasEntityIdentifier,
-  type CanvasInpaintMaskState,
-  type CanvasLayerState,
-  type CanvasRegionalGuidanceState,
-  type CanvasV2State,
-  type Coordinate,
-  type GenerationMode,
-  type GetLoggingContext,
-  RGBA_WHITE,
-  type RgbaColor,
+import type {
+  CanvasControlAdapterState,
+  CanvasEntityIdentifier,
+  CanvasInpaintMaskState,
+  CanvasLayerState,
+  CanvasRegionalGuidanceState,
+  CanvasV2State,
+  Coordinate,
+  GenerationMode,
+  GetLoggingContext,
+  Rect,
+  RgbaColor,
 } from 'features/controlLayers/store/types';
+import { RGBA_RED } from 'features/controlLayers/store/types';
+import { isValidLayer } from 'features/nodes/util/graph/generation/addLayers';
 import type Konva from 'konva';
 import { atom } from 'nanostores';
 import type { Logger } from 'roarr';
-import { getImageDTO as defaultGetImageDTO, uploadImage as defaultUploadImage } from 'services/api/endpoints/images';
+import {
+  getImageDTO as defaultGetImageDTO,
+  getImageDTO,
+  uploadImage as defaultUploadImage,
+} from 'services/api/endpoints/images';
 import type { ImageCategory, ImageDTO } from 'services/api/types';
+import { assert } from 'tsafe';
 
 import { CanvasBackground } from './CanvasBackground';
 import { CanvasBbox } from './CanvasBbox';
@@ -350,7 +359,8 @@ export class CanvasManager {
     if (selectedEntity) {
       // These two entity types use a compositing rect for opacity. Their fill is always white.
       if (selectedEntity.state.type === 'regional_guidance' || selectedEntity.state.type === 'inpaint_mask') {
-        currentFill = RGBA_WHITE;
+        currentFill = RGBA_RED;
+        // currentFill = RGBA_WHITE;
       }
     }
     return currentFill;
@@ -620,8 +630,96 @@ export class CanvasManager {
     return pixels / this.getStageScale();
   }
 
+  getCompositeLayerStageClone = (): Konva.Stage => {
+    const layersState = this.stateApi.getLayersState();
+    const stageClone = this.stage.clone();
+
+    stageClone.scaleX(1);
+    stageClone.scaleY(1);
+    stageClone.x(0);
+    stageClone.y(0);
+
+    const validLayers = layersState.entities.filter(isValidLayer);
+    // getLayers() returns the internal `children` array of the stage directly - calling destroy on a layer will
+    // mutate that array. We need to clone the array to avoid mutating the original.
+    for (const konvaLayer of stageClone.getLayers().slice()) {
+      if (!validLayers.find((l) => l.id === konvaLayer.id())) {
+        konvaLayer.destroy();
+      }
+    }
+
+    return stageClone;
+  };
+
+  getCompositeLayerBlob = (rect?: Rect): Promise<Blob> => {
+    return konvaNodeToBlob(this.getCompositeLayerStageClone(), rect);
+  };
+
+  getCompositeLayerImageData = (rect?: Rect): ImageData => {
+    return konvaNodeToImageData(this.getCompositeLayerStageClone(), rect);
+  };
+
+  getCompositeLayerImageDTO = async (rect?: Rect): Promise<ImageDTO> => {
+    const blob = await this.getCompositeLayerBlob(rect);
+    const imageDTO = await this.util.uploadImage(blob, 'composite-layer.png', 'general', true);
+    this.stateApi.setLayerImageCache(imageDTO);
+    return imageDTO;
+  };
+
+  getInpaintMaskBlob = (rect?: Rect): Promise<Blob> => {
+    return this.inpaintMask.renderer.getBlob({ rect });
+  };
+
+  getInpaintMaskImageData = (rect?: Rect): ImageData => {
+    return this.inpaintMask.renderer.getImageData({ rect });
+  };
+
+  getInpaintMaskImageDTO = async (rect?: Rect): Promise<ImageDTO> => {
+    const blob = await this.inpaintMask.renderer.getBlob({ rect });
+    const imageDTO = await this.util.uploadImage(blob, 'inpaint-mask.png', 'mask', true);
+    this.stateApi.setInpaintMaskImageCache(imageDTO);
+    return imageDTO;
+  };
+
+  getRegionMaskImageDTO = async (id: string, rect?: Rect): Promise<ImageDTO> => {
+    const region = this.getEntity({ id, type: 'regional_guidance' });
+    assert(region?.type === 'regional_guidance');
+    if (region.state.imageCache) {
+      const imageDTO = await getImageDTO(region.state.imageCache.name);
+      if (imageDTO) {
+        return imageDTO;
+      }
+    }
+    return region.adapter.renderer.getImageDTO({
+      rect,
+      category: 'other',
+      is_intermediate: true,
+      onUploaded: (imageDTO) => {
+        this.stateApi.setRegionMaskImageCache(region.state.id, imageDTO);
+      },
+    });
+  };
+
   getGenerationMode(): GenerationMode {
-    return getGenerationMode({ manager: this });
+    const { rect } = this.stateApi.getBbox();
+    const inpaintMaskImageData = this.getInpaintMaskImageData(rect);
+    const inpaintMaskTransparency = getImageDataTransparency(inpaintMaskImageData);
+    const compositeLayerImageData = this.getCompositeLayerImageData(rect);
+    const compositeLayerTransparency = getImageDataTransparency(compositeLayerImageData);
+    if (compositeLayerTransparency === 'FULLY_TRANSPARENT') {
+      // When the initial image is fully transparent, we are always doing txt2img
+      return 'txt2img';
+    } else if (compositeLayerTransparency === 'PARTIALLY_TRANSPARENT') {
+      // When the initial image is partially transparent, we are always outpainting
+      return 'outpaint';
+    } else if (inpaintMaskTransparency === 'FULLY_TRANSPARENT') {
+      // compositeLayerTransparency === 'OPAQUE'
+      // When the inpaint mask is fully transparent, we are doing img2img
+      return 'img2img';
+    } else {
+      // Else at least some of the inpaint mask is opaque, so we are inpainting
+      return 'inpaint';
+    }
   }
 
   getControlAdapterImage(arg: Omit<Parameters<typeof getControlAdapterImage>[0], 'manager'>) {

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

@@ -1,3 +1,4 @@
+import { deepClone } from 'common/util/deepClone';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import { CanvasObjectRenderer } from 'features/controlLayers/konva/CanvasObjectRenderer';
 import { CanvasTransformer } from 'features/controlLayers/konva/CanvasTransformer';
@@ -41,6 +42,8 @@ export class CanvasMaskAdapter {
 
     this.konva = {
       layer: new Konva.Layer({
+        // We need the ID on the layer to help with building the composite initial image
+        // See `getCompositeLayerStageClone()`
         id: this.id,
         name: `${this.type}:layer`,
         listening: false,
@@ -135,4 +138,12 @@ export class CanvasMaskAdapter {
     const isEnabled = get(arg, 'isEnabled', this.state.isEnabled);
     this.konva.layer.visible(isEnabled);
   };
+
+  repr = () => {
+    return {
+      id: this.id,
+      type: this.type,
+      state: deepClone(this.state),
+    };
+  };
 }

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

@@ -8,18 +8,20 @@ import type { CanvasLayerAdapter } from 'features/controlLayers/konva/CanvasLaye
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import type { CanvasMaskAdapter } from 'features/controlLayers/konva/CanvasMaskAdapter';
 import { CanvasRectRenderer } from 'features/controlLayers/konva/CanvasRect';
-import { getPrefixedId, konvaNodeToBlob, previewBlob } from 'features/controlLayers/konva/util';
+import { getPrefixedId, konvaNodeToBlob, konvaNodeToImageData, previewBlob } from 'features/controlLayers/konva/util';
 import {
   type CanvasBrushLineState,
   type CanvasEraserLineState,
   type CanvasImageState,
   type CanvasRectState,
   imageDTOToImageObject,
+  type Rect,
   type RgbColor,
 } from 'features/controlLayers/store/types';
 import Konva from 'konva';
 import type { Logger } from 'roarr';
 import { uploadImage } from 'services/api/endpoints/images';
+import type { ImageCategory, ImageDTO } from 'services/api/types';
 import { assert } from 'tsafe';
 
 /**
@@ -348,10 +350,8 @@ export class CanvasObjectRenderer {
   rasterize = async () => {
     this.log.debug('Rasterizing entity');
 
-    const objectGroupClone = this.konva.objectGroup.clone();
-    const interactionRectClone = this.parent.transformer.konva.proxyRect.clone();
-    const rect = interactionRectClone.getClientRect();
-    const blob = await konvaNodeToBlob(objectGroupClone, rect);
+    const rect = this.parent.transformer.getRelativeRect();
+    const blob = await this.getBlob({ rect });
     if (this.manager._isDebugging) {
       previewBlob(blob, 'Rasterized entity');
     }
@@ -365,6 +365,33 @@ export class CanvasObjectRenderer {
     });
   };
 
+  getBlob = ({ rect }: { rect?: Rect }): Promise<Blob> => {
+    return konvaNodeToBlob(this.konva.objectGroup.clone(), rect);
+  };
+
+  getImageData = ({ rect }: { rect?: Rect }): ImageData => {
+    return konvaNodeToImageData(this.konva.objectGroup.clone(), rect);
+  };
+
+  getImageDTO = async ({
+    rect,
+    category,
+    is_intermediate,
+    onUploaded,
+  }: {
+    rect?: Rect;
+    category: ImageCategory;
+    is_intermediate: boolean;
+    onUploaded?: (imageDTO: ImageDTO) => void;
+  }): Promise<ImageDTO> => {
+    const blob = await this.getBlob({ rect });
+    const imageDTO = await uploadImage(blob, `${this.id}.png`, category, is_intermediate);
+    if (onUploaded) {
+      onUploaded(imageDTO);
+    }
+    return imageDTO;
+  };
+
   /**
    * Destroys this renderer and all of its object renderers.
    */

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

@@ -685,6 +685,10 @@ export class CanvasTransformer {
     this.calculateRect();
   };
 
+  getRelativeRect = (): Rect => {
+    return this.konva.proxyRect.getClientRect({ relativeTo: this.parent.konva.layer });
+  };
+
   _enableTransform = () => {
     this.isTransformEnabled = true;
     this.konva.transformer.visible(true);

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

@@ -1,4 +1,3 @@
-import { getImageDataTransparency } from 'common/util/arrayBuffer';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import type {
   CanvasObjectState,
@@ -329,7 +328,7 @@ export const previewBlob = async (blob: Blob, label?: string) => {
 export function getInpaintMaskLayerClone(arg: { manager: CanvasManager }): Konva.Layer {
   const { manager } = arg;
   const layerClone = manager.inpaintMask.konva.layer.clone();
-  const objectGroupClone = manager.inpaintMask.konva.group.clone();
+  const objectGroupClone = manager.inpaintMask.renderer.konva.objectGroup.clone();
 
   layerClone.destroyChildren();
   layerClone.add(objectGroupClone);
@@ -347,7 +346,7 @@ export function getRegionMaskLayerClone(arg: { manager: CanvasManager; id: strin
   assert(canvasRegion, `Canvas region with id ${id} not found`);
 
   const layerClone = canvasRegion.konva.layer.clone();
-  const objectGroupClone = canvasRegion.konva.group.clone();
+  const objectGroupClone = canvasRegion.renderer.konva.objectGroup.clone();
 
   layerClone.destroyChildren();
   layerClone.add(objectGroupClone);
@@ -407,27 +406,42 @@ export function getCompositeLayerStageClone(arg: { manager: CanvasManager }): Ko
 
   const validLayers = layersState.entities.filter(isValidLayer);
   console.log(validLayers);
-  // 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) {
-      console.log('deleting', konvaLayer);
-      toDelete.push(konvaLayer);
+  // getLayers() returns the internal `children` array of the stage directly - calling destroy on a layer will
+  // mutate that array. We need to clone the array to avoid mutating the original.
+  for (const konvaLayer of stageClone.getLayers().slice()) {
+    if (!validLayers.find((l) => l.id === konvaLayer.id())) {
+      console.log('destroying', konvaLayer.id());
+      konvaLayer.destroy();
     }
   }
 
-  for (const konvaLayer of toDelete) {
-    konvaLayer.destroy();
-  }
-
   return stageClone;
 }
 
+export type Transparency = 'FULLY_TRANSPARENT' | 'PARTIALLY_TRANSPARENT' | 'OPAQUE';
+export function getImageDataTransparency(imageData: ImageData): Transparency {
+  let isFullyTransparent = true;
+  let isPartiallyTransparent = false;
+  const len = imageData.data.length;
+  for (let i = 3; i < len; i += 4) {
+    if (imageData.data[i] !== 0) {
+      isFullyTransparent = false;
+    } else {
+      isPartiallyTransparent = true;
+    }
+    if (!isFullyTransparent && isPartiallyTransparent) {
+      return 'PARTIALLY_TRANSPARENT';
+    }
+  }
+  if (isFullyTransparent) {
+    return 'FULLY_TRANSPARENT';
+  }
+  if (isPartiallyTransparent) {
+    return 'PARTIALLY_TRANSPARENT';
+  }
+  return 'OPAQUE';
+}
+
 export function getGenerationMode(arg: { manager: CanvasManager }): GenerationMode {
   const { manager } = arg;
   const { x, y, width, height } = manager.stateApi.getBbox().rect;

invokeai/frontend/web/src/features/nodes/util/graph/generation/addInpaint.ts

@@ -2,7 +2,7 @@ import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import type { CanvasV2State, Dimensions } from 'features/controlLayers/store/types';
 import type { Graph } from 'features/nodes/util/graph/generation/Graph';
 import type { ParameterPrecision } from 'features/parameters/types/parameterSchemas';
-import { isEqual, pick } from 'lodash-es';
+import { isEqual } from 'lodash-es';
 import type { Invocation } from 'services/api/types';
 
 export const addInpaint = async (
@@ -21,9 +21,8 @@ export const addInpaint = async (
 ): Promise<Invocation<'canvas_v2_mask_and_crop'>> => {
   denoise.denoising_start = denoising_start;
 
-  const cropBbox = pick(bbox.rect, ['x', 'y', 'width', 'height']);
-  const initialImage = await manager.getInitialImage({ bbox: cropBbox });
-  const maskImage = await manager.getInpaintMaskImage({ bbox: cropBbox });
+  const initialImage = await manager.getCompositeLayerImageDTO(bbox.rect);
+  const maskImage = await manager.getInpaintMaskImageDTO(bbox.rect);
 
   if (!isEqual(scaledSize, originalSize)) {
     // Scale before processing requires some resizing

invokeai/frontend/web/src/features/nodes/util/graph/generation/addOutpaint.ts

@@ -3,7 +3,7 @@ import type { CanvasV2State, Dimensions } from 'features/controlLayers/store/typ
 import type { Graph } from 'features/nodes/util/graph/generation/Graph';
 import { getInfill } from 'features/nodes/util/graph/graphBuilderUtils';
 import type { ParameterPrecision } from 'features/parameters/types/parameterSchemas';
-import { isEqual, pick } from 'lodash-es';
+import { isEqual } from 'lodash-es';
 import type { Invocation } from 'services/api/types';
 
 export const addOutpaint = async (
@@ -22,9 +22,8 @@ export const addOutpaint = async (
 ): Promise<Invocation<'canvas_v2_mask_and_crop'>> => {
   denoise.denoising_start = denoising_start;
 
-  const cropBbox = pick(bbox.rect, ['x', 'y', 'width', 'height']);
-  const initialImage = await manager.getInitialImage({ bbox: cropBbox });
-  const maskImage = await manager.getInpaintMaskImage({ bbox: cropBbox });
+  const initialImage = await manager.getCompositeLayerImageDTO(bbox.rect);
+  const maskImage = await manager.getInpaintMaskImageDTO(bbox.rect);
   const infill = getInfill(g, compositing);
 
   if (!isEqual(scaledSize, originalSize)) {

invokeai/frontend/web/src/features/nodes/util/graph/generation/addRegions.ts

@@ -1,6 +1,6 @@
 import { deepClone } from 'common/util/deepClone';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
-import type { CanvasIPAdapterState, Rect, CanvasRegionalGuidanceState } from 'features/controlLayers/store/types';
+import type { CanvasIPAdapterState, CanvasRegionalGuidanceState, Rect } from 'features/controlLayers/store/types';
 import {
   PROMPT_REGION_INVERT_TENSOR_MASK_PREFIX,
   PROMPT_REGION_MASK_TO_TENSOR_PREFIX,
@@ -44,7 +44,7 @@ export const addRegions = async (
 
   for (const region of validRegions) {
     // Upload the mask image, or get the cached image if it exists
-    const { image_name } = await manager.getRegionMaskImage({ id: region.id, bbox });
+    const { image_name } = await manager.getRegionMaskImageDTO(region.id, bbox);
 
     // The main mask-to-tensor node
     const maskToTensor = g.addNode({