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

This commit is contained in:
psychedelicious 2024-08-07 15:32:53 +10:00
parent a42d0ce1d2
commit 3ae7250ef7
11 changed files with 206 additions and 69 deletions

View File

@ -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());
}

View File

@ -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 };
};

View File

@ -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(),
};

View File

@ -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'>) {

View File

@ -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),
};
};
}

View File

@ -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.
*/

View File

@ -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);

View File

@ -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;

View File

@ -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

View File

@ -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)) {

View File

@ -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({