From c7303adb0d821b008e7514c49b3cd7e2cf753e61 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Thu, 4 May 2023 13:27:21 +1000 Subject: [PATCH] feat(ui): fix generation mode logic --- .../web/src/common/util/arrayBuffer.ts | 19 ++- .../canvas/util/dataURLToUint8ClampedArray.ts | 26 +++ .../src/features/canvas/util/generateMask.ts | 157 +++++++++++++----- .../src/features/canvas/util/getCanvasData.ts | 69 +++----- .../features/nodes/util/buildCanvasGraph.ts | 6 +- .../web/src/services/thunks/session.ts | 2 +- 6 files changed, 178 insertions(+), 101 deletions(-) create mode 100644 invokeai/frontend/web/src/features/canvas/util/dataURLToUint8ClampedArray.ts diff --git a/invokeai/frontend/web/src/common/util/arrayBuffer.ts b/invokeai/frontend/web/src/common/util/arrayBuffer.ts index 779d9b1b17..885fc05177 100644 --- a/invokeai/frontend/web/src/common/util/arrayBuffer.ts +++ b/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) => { - const len = imageData.data.length; +export const areAnyPixelsBlack = (pixels: Uint8ClampedArray) => { + const len = pixels.length; let i = 0; for (i; i < len; ) { if ( - imageData.data[i++] === 255 && - imageData.data[i++] === 255 && - imageData.data[i++] === 255 && - imageData.data[i++] === 255 + pixels[i++] === 0 && + pixels[i++] === 0 && + pixels[i++] === 0 && + pixels[i++] === 255 ) { return true; } diff --git a/invokeai/frontend/web/src/features/canvas/util/dataURLToUint8ClampedArray.ts b/invokeai/frontend/web/src/features/canvas/util/dataURLToUint8ClampedArray.ts new file mode 100644 index 0000000000..2652d294fc --- /dev/null +++ b/invokeai/frontend/web/src/features/canvas/util/dataURLToUint8ClampedArray.ts @@ -0,0 +1,26 @@ +export const dataURLToImageData = async ( + dataURL: string, + width: number, + height: number +): Promise => + 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; + }); diff --git a/invokeai/frontend/web/src/features/canvas/util/generateMask.ts b/invokeai/frontend/web/src/features/canvas/util/generateMask.ts index db88dd1ad3..f3cc1fb237 100644 --- a/invokeai/frontend/web/src/features/canvas/util/generateMask.ts +++ b/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 => { - // 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 } => { +const generateMask = (lines: CanvasMaskLine[], boundingBox: IRect): string => { // 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; diff --git a/invokeai/frontend/web/src/features/canvas/util/getCanvasData.ts b/invokeai/frontend/web/src/features/canvas/util/getCanvasData.ts index c6192f0a09..af4ea42561 100644 --- a/invokeai/frontend/web/src/features/canvas/util/getCanvasData.ts +++ b/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( - offsetBoundingBox.x, - offsetBoundingBox.y, - offsetBoundingBox.width, - offsetBoundingBox.height + const maskDataURL = generateMask( + isMaskEnabled ? objects.filter(isCanvasMaskLine) : [], + boundingBox ); + 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); + const doesMaskHaveBlackPixels = areAnyPixelsBlack(maskImageData.data); if (state.system.enableImageDebugging) { openBase64ImageInTab([ - { base64: maskDataURL, caption: 'mask sent as init_mask' }, - { base64: baseDataURL, caption: 'image sent as init_img' }, + { base64: maskDataURL, caption: 'mask b64' }, + { base64: baseDataURL, caption: 'image b64' }, ]); } - canvasBaseLayer.scale(tempScale); - // generationParameters.init_img = imageDataURL; // generationParameters.progress_images = false; diff --git a/invokeai/frontend/web/src/features/nodes/util/buildCanvasGraph.ts b/invokeai/frontend/web/src/features/nodes/util/buildCanvasGraph.ts index 24c057c04a..8182241843 100644 --- a/invokeai/frontend/web/src/features/nodes/util/buildCanvasGraph.ts +++ b/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 => { - const c = getCanvasData(state); +export const buildCanvasGraph = async ( + state: RootState +): Promise => { + const c = await getCanvasData(state); if (!c) { moduleLog.error('Unable to create canvas graph'); diff --git a/invokeai/frontend/web/src/services/thunks/session.ts b/invokeai/frontend/web/src/services/thunks/session.ts index c92b303ff7..8e0b59b1d4 100644 --- a/invokeai/frontend/web/src/services/thunks/session.ts +++ b/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) {