feat(ui): fix generation mode logic

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

@@ -1,10 +1,11 @@
-export const getImageDataTransparency = (imageData: ImageData) => {
+export const getImageDataTransparency = (pixels: Uint8ClampedArray) => {
+  console.log(pixels);
   let isFullyTransparent = true;
   let isPartiallyTransparent = false;
-  const len = imageData.data.length;
+  const len = pixels.length;
   let i = 3;
   for (i; i < len; i += 4) {
-    if (imageData.data[i] === 255) {
+    if (pixels[i] === 255) {
       isFullyTransparent = false;
     } else {
       isPartiallyTransparent = true;
@@ -16,15 +17,15 @@ export const getImageDataTransparency = (imageData: ImageData) => {
   return { isFullyTransparent, isPartiallyTransparent };
 };
 
-export const areAnyPixelsBlack = (imageData: ImageData) => {
+export const areAnyPixelsBlack = (pixels: Uint8ClampedArray) => {
-  const len = imageData.data.length;
+  const len = pixels.length;
   let i = 0;
   for (i; i < len; ) {
     if (
-      imageData.data[i++] === 255 &&
+      pixels[i++] === 0 &&
-      imageData.data[i++] === 255 &&
+      pixels[i++] === 0 &&
-      imageData.data[i++] === 255 &&
+      pixels[i++] === 0 &&
-      imageData.data[i++] === 255
+      pixels[i++] === 255
     ) {
       return true;
     }

invokeai/frontend/web/src/features/canvas/util/dataURLToUint8ClampedArray.ts

@@ -0,0 +1,26 @@
+export const dataURLToImageData = async (
+  dataURL: string,
+  width: number,
+  height: number
+): Promise<ImageData> =>
+  new Promise((resolve, reject) => {
+    const canvas = document.createElement('canvas');
+    canvas.width = width;
+    canvas.height = height;
+    const ctx = canvas.getContext('2d');
+    const image = new Image();
+
+    if (!ctx) {
+      canvas.remove();
+      reject('Unable to get context');
+      return;
+    }
+
+    image.onload = function () {
+      ctx.drawImage(image, 0, 0);
+      canvas.remove();
+      resolve(ctx.getImageData(0, 0, width, height));
+    };
+
+    image.src = dataURL;
+  });

invokeai/frontend/web/src/features/canvas/util/generateMask.ts

@@ -1,6 +1,108 @@
+// import { CanvasMaskLine } from 'features/canvas/store/canvasTypes';
+// import Konva from 'konva';
+// import { Stage } from 'konva/lib/Stage';
+// import { IRect } from 'konva/lib/types';
+
+// /**
+//  * Generating a mask image from InpaintingCanvas.tsx is not as simple
+//  * as calling toDataURL() on the canvas, because the mask may be represented
+//  * by colored lines or transparency, or the user may have inverted the mask
+//  * display.
+//  *
+//  * So we need to regenerate the mask image by creating an offscreen canvas,
+//  * drawing the mask and compositing everything correctly to output a valid
+//  * mask image.
+//  */
+// export const getStageDataURL = (stage: Stage, boundingBox: IRect): string => {
+//   // create an offscreen canvas and add the mask to it
+//   // const { stage, offscreenContainer } = buildMaskStage(lines, boundingBox);
+
+//   const dataURL = stage.toDataURL({ ...boundingBox });
+
+//   // const imageData = stage
+//   //   .toCanvas()
+//   //   .getContext('2d')
+//   //   ?.getImageData(
+//   //     boundingBox.x,
+//   //     boundingBox.y,
+//   //     boundingBox.width,
+//   //     boundingBox.height
+//   //   );
+
+//   // offscreenContainer.remove();
+
+//   // return { dataURL, imageData };
+
+//   return dataURL;
+// };
+
+// export const getStageImageData = (
+//   stage: Stage,
+//   boundingBox: IRect
+// ): ImageData | undefined => {
+//   const imageData = stage
+//     .toCanvas()
+//     .getContext('2d')
+//     ?.getImageData(
+//       boundingBox.x,
+//       boundingBox.y,
+//       boundingBox.width,
+//       boundingBox.height
+//     );
+
+//   return imageData;
+// };
+
+// export const buildMaskStage = (
+//   lines: CanvasMaskLine[],
+//   boundingBox: IRect
+// ): { stage: Stage; offscreenContainer: HTMLDivElement } => {
+//   // create an offscreen canvas and add the mask to it
+//   const { width, height } = boundingBox;
+
+//   const offscreenContainer = document.createElement('div');
+
+//   const stage = new Konva.Stage({
+//     container: offscreenContainer,
+//     width: width,
+//     height: height,
+//   });
+
+//   const baseLayer = new Konva.Layer();
+//   const maskLayer = new Konva.Layer();
+
+//   // composite the image onto the mask layer
+//   baseLayer.add(
+//     new Konva.Rect({
+//       ...boundingBox,
+//       fill: 'white',
+//     })
+//   );
+
+//   lines.forEach((line) =>
+//     maskLayer.add(
+//       new Konva.Line({
+//         points: line.points,
+//         stroke: 'black',
+//         strokeWidth: line.strokeWidth * 2,
+//         tension: 0,
+//         lineCap: 'round',
+//         lineJoin: 'round',
+//         shadowForStrokeEnabled: false,
+//         globalCompositeOperation:
+//           line.tool === 'brush' ? 'source-over' : 'destination-out',
+//       })
+//     )
+//   );
+
+//   stage.add(baseLayer);
+//   stage.add(maskLayer);
+
+//   return { stage, offscreenContainer };
+// };
+
 import { CanvasMaskLine } from 'features/canvas/store/canvasTypes';
 import Konva from 'konva';
-import { Stage } from 'konva/lib/Stage';
 import { IRect } from 'konva/lib/types';
 
 /**
@@ -13,50 +115,7 @@ import { IRect } from 'konva/lib/types';
  * drawing the mask and compositing everything correctly to output a valid
  * mask image.
  */
-export const getStageDataURL = (stage: Stage, boundingBox: IRect): string => {
+const generateMask = (lines: CanvasMaskLine[], boundingBox: IRect): string => {
-  // create an offscreen canvas and add the mask to it
-  // const { stage, offscreenContainer } = buildMaskStage(lines, boundingBox);
-
-  const dataURL = stage.toDataURL({ ...boundingBox });
-
-  // const imageData = stage
-  //   .toCanvas()
-  //   .getContext('2d')
-  //   ?.getImageData(
-  //     boundingBox.x,
-  //     boundingBox.y,
-  //     boundingBox.width,
-  //     boundingBox.height
-  //   );
-
-  // offscreenContainer.remove();
-
-  // return { dataURL, imageData };
-
-  return dataURL;
-};
-
-export const getStageImageData = (
-  stage: Stage,
-  boundingBox: IRect
-): ImageData | undefined => {
-  const imageData = stage
-    .toCanvas()
-    .getContext('2d')
-    ?.getImageData(
-      boundingBox.x,
-      boundingBox.y,
-      boundingBox.width,
-      boundingBox.height
-    );
-
-  return imageData;
-};
-
-export const buildMaskStage = (
-  lines: CanvasMaskLine[],
-  boundingBox: IRect
-): { stage: Stage; offscreenContainer: HTMLDivElement } => {
   // create an offscreen canvas and add the mask to it
   const { width, height } = boundingBox;
 
@@ -98,5 +157,11 @@ export const buildMaskStage = (
   stage.add(baseLayer);
   stage.add(maskLayer);
 
-  return { stage, offscreenContainer };
+  const dataURL = stage.toDataURL({ ...boundingBox });
+
+  offscreenContainer.remove();
+
+  return dataURL;
 };
+
+export default generateMask;

invokeai/frontend/web/src/features/canvas/util/getCanvasData.ts

@@ -1,22 +1,18 @@
 import { RootState } from 'app/store/store';
 import { getCanvasBaseLayer, getCanvasStage } from './konvaInstanceProvider';
 import { isCanvasMaskLine } from '../store/canvasTypes';
-import {
-  buildMaskStage,
-  getStageDataURL,
-  getStageImageData,
-} from './generateMask';
 import { log } from 'app/logging/useLogger';
 import {
   areAnyPixelsBlack,
   getImageDataTransparency,
 } from 'common/util/arrayBuffer';
 import openBase64ImageInTab from 'common/util/openBase64ImageInTab';
-import { masks } from 'dateformat';
+import generateMask from './generateMask';
+import { dataURLToImageData } from './dataURLToUint8ClampedArray';
 
 const moduleLog = log.child({ namespace: 'getCanvasDataURLs' });
 
-export const getCanvasData = (state: RootState) => {
+export const getCanvasData = async (state: RootState) => {
   const canvasBaseLayer = getCanvasBaseLayer();
   const canvasStage = getCanvasStage();
 
@@ -65,57 +61,44 @@ export const getCanvasData = (state: RootState) => {
     height: boundingBox.height,
   };
 
-  const { stage: maskStage, offscreenContainer } = buildMaskStage(
-    isMaskEnabled ? objects.filter(isCanvasMaskLine) : [],
-    offsetBoundingBox
-  );
-
-  const maskDataURL = maskStage.toDataURL(offsetBoundingBox);
-
-  const maskImageData = maskStage
-    .toCanvas()
-    .getContext('2d')
-    ?.getImageData(
-      offsetBoundingBox.x,
-      offsetBoundingBox.y,
-      offsetBoundingBox.width,
-      offsetBoundingBox.height
-    );
-
-  offscreenContainer.remove();
-
-  if (!maskImageData) {
-    return;
-  }
-
   const baseDataURL = canvasBaseLayer.toDataURL(offsetBoundingBox);
 
-  const ctx = canvasBaseLayer.getContext();
+  canvasBaseLayer.scale(tempScale);
 
-  const baseImageData = ctx.getImageData(
+  const maskDataURL = generateMask(
-    offsetBoundingBox.x,
+    isMaskEnabled ? objects.filter(isCanvasMaskLine) : [],
-    offsetBoundingBox.y,
+    boundingBox
-    offsetBoundingBox.width,
-    offsetBoundingBox.height
   );
 
+  const baseImageData = await dataURLToImageData(
+    baseDataURL,
+    boundingBox.width,
+    boundingBox.height
+  );
+
+  const maskImageData = await dataURLToImageData(
+    maskDataURL,
+    boundingBox.width,
+    boundingBox.height
+  );
+
+  console.log('baseImageData', baseImageData);
+  console.log('maskImageData', maskImageData);
+
   const {
     isPartiallyTransparent: baseIsPartiallyTransparent,
     isFullyTransparent: baseIsFullyTransparent,
-  } = getImageDataTransparency(baseImageData);
+  } = getImageDataTransparency(baseImageData.data);
 
-  // const doesMaskHaveBlackPixels = false;
+  const doesMaskHaveBlackPixels = areAnyPixelsBlack(maskImageData.data);
-  const doesMaskHaveBlackPixels = areAnyPixelsBlack(maskImageData);
 
   if (state.system.enableImageDebugging) {
     openBase64ImageInTab([
-      { base64: maskDataURL, caption: 'mask sent as init_mask' },
+      { base64: maskDataURL, caption: 'mask b64' },
-      { base64: baseDataURL, caption: 'image sent as init_img' },
+      { base64: baseDataURL, caption: 'image b64' },
     ]);
   }
 
-  canvasBaseLayer.scale(tempScale);
-
   // generationParameters.init_img = imageDataURL;
   // generationParameters.progress_images = false;
 

invokeai/frontend/web/src/features/nodes/util/buildCanvasGraph.ts

@@ -16,8 +16,10 @@ const moduleLog = log.child({ namespace: 'buildCanvasGraph' });
 /**
  * Builds the Canvas workflow graph.
  */
-export const buildCanvasGraph = (state: RootState): Graph | undefined => {
+export const buildCanvasGraph = async (
-  const c = getCanvasData(state);
+  state: RootState
+): Promise<Graph | undefined> => {
+  const c = await getCanvasData(state);
 
   if (!c) {
     moduleLog.error('Unable to create canvas graph');

invokeai/frontend/web/src/services/thunks/session.ts

@@ -47,7 +47,7 @@ export const canvasGraphBuilt = createAppAsyncThunk(
   'api/canvasGraphBuilt',
   async (_, { dispatch, getState, rejectWithValue }) => {
     try {
-      const graph = buildCanvasGraph(getState());
+      const graph = await buildCanvasGraph(getState());
       dispatch(sessionCreated({ graph }));
       return graph;
     } catch (err: any) {