feat(ui): use new canvas output node

This commit is contained in:
psychedelicious 2024-07-16 18:33:50 +10:00
parent 070b78501b
commit b76609ef18
12 changed files with 104 additions and 92 deletions

View File

@ -1,11 +1,10 @@
import { isAnyOf } from '@reduxjs/toolkit';
import { logger } from 'app/logging/logger';
import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
import {
layerAdded,
layerImageAdded,
sessionStagedImageAccepted,
sessionStagingCanceled,
sessionStagingAreaImageAccepted,
sessionStagingAreaReset,
} from 'features/controlLayers/store/canvasV2Slice';
import { toast } from 'features/toast/toast';
import { t } from 'i18next';
@ -14,7 +13,7 @@ import { assert } from 'tsafe';
export const addStagingListeners = (startAppListening: AppStartListening) => {
startAppListening({
matcher: isAnyOf(sessionStagingCanceled, sessionStagedImageAccepted),
actionCreator: sessionStagingAreaReset,
effect: async (_, { dispatch }) => {
const log = logger('canvas');
@ -47,10 +46,10 @@ export const addStagingListeners = (startAppListening: AppStartListening) => {
});
startAppListening({
actionCreator: sessionStagedImageAccepted,
actionCreator: sessionStagingAreaImageAccepted,
effect: async (action, api) => {
const { imageDTO } = action.payload;
const { layers, selectedEntityIdentifier, bbox } = api.getState().canvasV2;
const { index } = action.payload;
const { layers, selectedEntityIdentifier } = api.getState().canvasV2;
let layer = layers.entities.find((layer) => layer.id === selectedEntityIdentifier?.id);
if (!layer) {
@ -63,11 +62,21 @@ export const addStagingListeners = (startAppListening: AppStartListening) => {
layer = api.getState().canvasV2.layers.entities[0];
}
const stagedImage = api.getState().canvasV2.session.stagedImages[index];
assert(stagedImage, 'No staged image found to accept');
assert(layer, 'No layer found to stage image');
const { id } = layer;
api.dispatch(layerImageAdded({ id, imageDTO, pos: { x: bbox.rect.x - layer.x, y: bbox.rect.y - layer.y } }));
api.dispatch(
layerImageAdded({
id,
imageDTO: stagedImage.imageDTO,
pos: { x: stagedImage.rect.x - layer.x, y: stagedImage.rect.y - layer.y },
})
);
api.dispatch(sessionStagingAreaReset());
},
});
};

View File

@ -1,7 +1,7 @@
import { enqueueRequested } from 'app/store/actions';
import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
import { getCanvasManager } from 'features/controlLayers/konva/CanvasManager';
import { sessionStagingCanceled, sessionStartedStaging } from 'features/controlLayers/store/canvasV2Slice';
import { sessionStagingAreaReset, sessionStartedStaging } from 'features/controlLayers/store/canvasV2Slice';
import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
import { buildSD1Graph } from 'features/nodes/util/graph/generation/buildSD1Graph';
import { buildSDXLGraph } from 'features/nodes/util/graph/generation/buildSDXLGraph';
@ -49,7 +49,7 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
await req.unwrap();
} catch {
if (didStartStaging && getState().canvasV2.session.isStaging) {
dispatch(sessionStagingCanceled());
dispatch(sessionStagingAreaReset());
}
}
},

View File

@ -6,7 +6,6 @@ import { $lastProgressEvent, sessionImageStaged } from 'features/controlLayers/s
import { boardIdSelected, galleryViewChanged, imageSelected, offsetChanged } from 'features/gallery/store/gallerySlice';
import { $nodeExecutionStates, upsertExecutionState } from 'features/nodes/hooks/useExecutionState';
import { zNodeStatus } from 'features/nodes/types/invocation';
import { CANVAS_OUTPUT } from 'features/nodes/util/graph/constants';
import { boardsApi } from 'services/api/endpoints/boards';
import { imagesApi } from 'services/api/endpoints/images';
import { getCategories, getListImagesUrl } from 'services/api/util';
@ -42,7 +41,10 @@ export const addInvocationCompleteEventListener = (startAppListening: AppStartLi
}
// This complete event has an associated image output
if (data.result.type === 'image_output' && !nodeTypeDenylist.includes(data.invocation.type)) {
if (
(data.result.type === 'image_output' || data.result.type === 'canvas_v2_mask_and_crop_output') &&
!nodeTypeDenylist.includes(data.invocation.type)
) {
const { image_name } = data.result.image;
const { gallery, canvasV2 } = getState();
@ -57,9 +59,10 @@ export const addInvocationCompleteEventListener = (startAppListening: AppStartLi
imageDTORequest.unsubscribe();
// handle tab-specific logic
if (data.origin === 'canvas' && data.invocation_source_id === CANVAS_OUTPUT) {
if (data.origin === 'canvas' && data.result.type === 'canvas_v2_mask_and_crop_output') {
const { x, y, width, height } = data.result;
if (canvasV2.session.isStaging) {
dispatch(sessionImageStaged({ imageDTO }));
dispatch(sessionImageStaged({ imageDTO, rect: { x, y, width, height } }));
} else if (!canvasV2.session.isActive) {
$lastProgressEvent.set(null);
}

View File

@ -3,11 +3,11 @@ import { useStore } from '@nanostores/react';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import {
$shouldShowStagedImage,
sessionStagingCanceled,
sessionStagedImageAccepted,
sessionStagedImageDiscarded,
sessionNextStagedImageSelected,
sessionPrevStagedImageSelected,
sessionStagedImageDiscarded,
sessionStagingAreaImageAccepted,
sessionStagingAreaReset,
} from 'features/controlLayers/store/canvasV2Slice';
import { memo, useCallback, useMemo } from 'react';
import { useHotkeys } from 'react-hotkeys-hook';
@ -40,7 +40,7 @@ export const StagingAreaToolbarContent = memo(() => {
const stagingArea = useAppSelector((s) => s.canvasV2.session);
const shouldShowStagedImage = useStore($shouldShowStagedImage);
const images = useMemo(() => stagingArea.stagedImages, [stagingArea]);
const selectedImageDTO = useMemo(() => {
const selectedImage = useMemo(() => {
return images[stagingArea.selectedStagedImageIndex] ?? null;
}, [images, stagingArea.selectedStagedImageIndex]);
@ -57,25 +57,25 @@ export const StagingAreaToolbarContent = memo(() => {
}, [dispatch]);
const onAccept = useCallback(() => {
if (!selectedImageDTO) {
if (!selectedImage) {
return;
}
dispatch(sessionStagedImageAccepted({ imageDTO: selectedImageDTO }));
}, [dispatch, selectedImageDTO]);
dispatch(sessionStagingAreaImageAccepted({ index: stagingArea.selectedStagedImageIndex }));
}, [dispatch, selectedImage, stagingArea.selectedStagedImageIndex]);
const onDiscardOne = useCallback(() => {
if (!selectedImageDTO) {
if (!selectedImage) {
return;
}
if (images.length === 1) {
dispatch(sessionStagingCanceled());
dispatch(sessionStagingAreaReset());
} else {
dispatch(sessionStagedImageDiscarded({ imageDTO: selectedImageDTO }));
dispatch(sessionStagedImageDiscarded({ index: stagingArea.selectedStagedImageIndex }));
}
}, [dispatch, selectedImageDTO, images.length]);
}, [selectedImage, images.length, dispatch, stagingArea.selectedStagedImageIndex]);
const onDiscardAll = useCallback(() => {
dispatch(sessionStagingCanceled());
dispatch(sessionStagingAreaReset());
}, [dispatch]);
const onToggleShouldShowStagedImage = useCallback(() => {
@ -145,7 +145,7 @@ export const StagingAreaToolbarContent = memo(() => {
icon={<PiCheckBold />}
onClick={onAccept}
colorScheme="invokeBlue"
isDisabled={!selectedImageDTO}
isDisabled={!selectedImage}
/>
<IconButton
tooltip={shouldShowStagedImage ? t('unifiedCanvas.showResultsOn') : t('unifiedCanvas.showResultsOff')}
@ -161,7 +161,7 @@ export const StagingAreaToolbarContent = memo(() => {
icon={<PiFloppyDiskBold />}
onClick={onSaveStagingImage}
colorScheme="invokeBlue"
isDisabled={!selectedImageDTO || !selectedImageDTO.is_intermediate}
isDisabled={!selectedImage || !selectedImage.imageDTO.is_intermediate}
/>
<IconButton
tooltip={`${t('unifiedCanvas.discardCurrent')}`}
@ -170,7 +170,7 @@ export const StagingAreaToolbarContent = memo(() => {
onClick={onDiscardOne}
colorScheme="invokeBlue"
fontSize={16}
isDisabled={!selectedImageDTO}
isDisabled={!selectedImage}
/>
<IconButton
tooltip={`${t('unifiedCanvas.discardAll')} (Esc)`}

View File

@ -1,44 +1,48 @@
import { CanvasImage } from 'features/controlLayers/konva/CanvasImage';
import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
import type { StagingAreaImage } from 'features/controlLayers/store/types';
import Konva from 'konva';
import type { ImageDTO } from 'services/api/types';
export class CanvasStagingArea {
group: Konva.Group;
image: CanvasImage | null;
imageDTO: ImageDTO | null;
selectedImage: StagingAreaImage | null;
manager: CanvasManager;
constructor(manager: CanvasManager) {
this.manager = manager;
this.group = new Konva.Group({ listening: false });
this.image = null;
this.imageDTO = null;
this.selectedImage = null;
}
async render() {
const session = this.manager.stateApi.getSession();
const bboxRect = this.manager.stateApi.getBbox().rect;
const shouldShowStagedImage = this.manager.stateApi.getShouldShowStagedImage();
this.imageDTO = session.stagedImages[session.selectedStagedImageIndex] ?? null;
this.selectedImage = session.stagedImages[session.selectedStagedImageIndex] ?? null;
if (this.imageDTO) {
if (this.selectedImage) {
if (this.image) {
if (!this.image.isLoading && !this.image.isError && this.image.imageName !== this.imageDTO.image_name) {
await this.image.updateImageSource(this.imageDTO.image_name);
if (
!this.image.isLoading &&
!this.image.isError &&
this.image.imageName !== this.selectedImage.imageDTO.image_name
) {
await this.image.updateImageSource(this.selectedImage.imageDTO.image_name);
}
this.image.konvaImageGroup.x(bboxRect.x);
this.image.konvaImageGroup.y(bboxRect.y);
this.image.konvaImageGroup.x(this.selectedImage.rect.x);
this.image.konvaImageGroup.y(this.selectedImage.rect.y);
this.image.konvaImageGroup.visible(shouldShowStagedImage);
} else {
const { image_name, width, height } = this.imageDTO;
const { image_name } = this.selectedImage.imageDTO;
const { x, y, width, height } = this.selectedImage.rect;
this.image = new CanvasImage(
{
id: 'staging-area-image',
type: 'image',
x: bboxRect.x,
y: bboxRect.y,
x,
y,
width,
height,
filters: [],
@ -50,9 +54,9 @@ export class CanvasStagingArea {
},
{
onLoad: (konvaImage) => {
if (this.imageDTO) {
konvaImage.width(this.imageDTO.width);
konvaImage.height(this.imageDTO.height);
if (this.selectedImage) {
konvaImage.width(this.selectedImage.rect.width);
konvaImage.height(this.selectedImage.rect.height);
}
this.manager.stateApi.resetLastProgressEvent();
this.image?.konvaImageGroup.visible(shouldShowStagedImage);
@ -60,7 +64,7 @@ export class CanvasStagingArea {
}
);
this.group.add(this.image.konvaImageGroup);
await this.image.updateImageSource(this.imageDTO.image_name);
await this.image.updateImageSource(this.selectedImage.imageDTO.image_name);
this.image.konvaImageGroup.visible(shouldShowStagedImage);
}
} else {

View File

@ -339,8 +339,7 @@ export const {
sessionStartedStaging,
sessionImageStaged,
sessionStagedImageDiscarded,
sessionStagedImageAccepted,
sessionStagingCanceled,
sessionStagingAreaReset,
sessionNextStagedImageSelected,
sessionPrevStagedImageSelected,
// Initial image
@ -383,3 +382,6 @@ export const canvasV2PersistConfig: PersistConfig<CanvasV2State> = {
};
export const sessionRequested = createAction(`${canvasV2Slice.name}/sessionRequested`);
export const sessionStagingAreaImageAccepted = createAction<{ index: number }>(
`${canvasV2Slice.name}/sessionStagingAreaImageAccepted`
);

View File

@ -1,6 +1,5 @@
import type { PayloadAction, SliceCaseReducers } from '@reduxjs/toolkit';
import type { CanvasV2State } from 'features/controlLayers/store/types';
import type { ImageDTO } from 'services/api/types';
import type { CanvasV2State, StagingAreaImage } from 'features/controlLayers/store/types';
export const sessionReducers = {
sessionStarted: (state) => {
@ -15,9 +14,9 @@ export const sessionReducers = {
state.tool.selectedBuffer = state.tool.selected;
state.tool.selected = 'view';
},
sessionImageStaged: (state, action: PayloadAction<{ imageDTO: ImageDTO }>) => {
const { imageDTO } = action.payload;
state.session.stagedImages.push(imageDTO);
sessionImageStaged: (state, action: PayloadAction<StagingAreaImage>) => {
const { imageDTO, rect } = action.payload;
state.session.stagedImages.push({ imageDTO, rect });
state.session.selectedStagedImageIndex = state.session.stagedImages.length - 1;
},
sessionNextStagedImageSelected: (state) => {
@ -29,9 +28,9 @@ export const sessionReducers = {
(state.session.selectedStagedImageIndex - 1 + state.session.stagedImages.length) %
state.session.stagedImages.length;
},
sessionStagedImageDiscarded: (state, action: PayloadAction<{ imageDTO: ImageDTO }>) => {
const { imageDTO } = action.payload;
state.session.stagedImages = state.session.stagedImages.filter((image) => image.image_name !== imageDTO.image_name);
sessionStagedImageDiscarded: (state, action: PayloadAction<{ index: number }>) => {
const { index } = action.payload;
state.session.stagedImages = state.session.stagedImages.splice(index, 1);
state.session.selectedStagedImageIndex = Math.min(
state.session.selectedStagedImageIndex,
state.session.stagedImages.length - 1
@ -40,17 +39,7 @@ export const sessionReducers = {
state.session.isStaging = false;
}
},
sessionStagedImageAccepted: (state, _: PayloadAction<{ imageDTO: ImageDTO }>) => {
// When we finish staging, reset the tool back to the previous selection.
state.session.isStaging = false;
state.session.stagedImages = [];
state.session.selectedStagedImageIndex = 0;
if (state.tool.selectedBuffer) {
state.tool.selected = state.tool.selectedBuffer;
state.tool.selectedBuffer = null;
}
},
sessionStagingCanceled: (state) => {
sessionStagingAreaReset: (state) => {
state.session.isStaging = false;
state.session.stagedImages = [];
state.session.selectedStagedImageIndex = 0;

View File

@ -826,6 +826,11 @@ export type LoRA = {
weight: number;
};
export type StagingAreaImage = {
imageDTO: ImageDTO;
rect: Rect;
};
export type CanvasV2State = {
_version: 3;
selectedEntityIdentifier: CanvasEntityIdentifier | null;
@ -913,7 +918,7 @@ export type CanvasV2State = {
session: {
isActive: boolean;
isStaging: boolean;
stagedImages: ImageDTO[];
stagedImages: StagingAreaImage[];
selectedStagedImageIndex: number;
};
};

View File

@ -18,7 +18,7 @@ export const addInpaint = async (
compositing: CanvasV2State['compositing'],
denoising_start: number,
vaePrecision: ParameterPrecision
): Promise<Invocation<'canvas_paste_back'>> => {
): Promise<Invocation<'canvas_v2_mask_and_crop'>> => {
denoise.denoising_start = denoising_start;
const cropBbox = pick(bbox.rect, ['x', 'y', 'width', 'height']);
@ -64,10 +64,10 @@ export const addInpaint = async (
fp32: vaePrecision === 'fp32',
});
const canvasPasteBack = g.addNode({
id: 'canvas_paste_back',
type: 'canvas_paste_back',
mask_blur: compositing.maskBlur,
source_image: { image_name: initialImage.image_name },
id: 'canvas_v2_mask_and_crop',
type: 'canvas_v2_mask_and_crop',
invert: true,
crop_visible: true,
});
// Resize initial image and mask to scaled size, feed into to gradient mask
@ -88,7 +88,7 @@ export const addInpaint = async (
g.addEdge(createGradientMask, 'expanded_mask_area', resizeMaskToOriginalSize, 'image');
// Finally, paste the generated masked image back onto the original image
g.addEdge(resizeImageToOriginalSize, 'image', canvasPasteBack, 'target_image');
g.addEdge(resizeImageToOriginalSize, 'image', canvasPasteBack, 'image');
g.addEdge(resizeMaskToOriginalSize, 'image', canvasPasteBack, 'mask');
return canvasPasteBack;
@ -111,10 +111,10 @@ export const addInpaint = async (
image: { image_name: initialImage.image_name },
});
const canvasPasteBack = g.addNode({
id: 'canvas_paste_back',
type: 'canvas_paste_back',
mask_blur: compositing.maskBlur,
source_image: { image_name: initialImage.image_name },
id: 'canvas_v2_mask_and_crop',
type: 'canvas_v2_mask_and_crop',
invert: true,
crop_visible: true,
});
g.addEdge(alphaToMask, 'image', createGradientMask, 'mask');
g.addEdge(i2l, 'latents', denoise, 'latents');
@ -124,7 +124,7 @@ export const addInpaint = async (
g.addEdge(createGradientMask, 'denoise_mask', denoise, 'denoise_mask');
g.addEdge(createGradientMask, 'expanded_mask_area', canvasPasteBack, 'mask');
g.addEdge(l2i, 'image', canvasPasteBack, 'target_image');
g.addEdge(l2i, 'image', canvasPasteBack, 'image');
return canvasPasteBack;
}

View File

@ -19,7 +19,7 @@ export const addOutpaint = async (
compositing: CanvasV2State['compositing'],
denoising_start: number,
vaePrecision: ParameterPrecision
): Promise<Invocation<'canvas_paste_back'>> => {
): Promise<Invocation<'canvas_v2_mask_and_crop'>> => {
denoise.denoising_start = denoising_start;
const cropBbox = pick(bbox.rect, ['x', 'y', 'width', 'height']);
@ -99,10 +99,10 @@ export const addOutpaint = async (
...originalSize,
});
const canvasPasteBack = g.addNode({
id: 'canvas_paste_back',
type: 'canvas_paste_back',
mask_blur: compositing.maskBlur,
source_image: { image_name: initialImage.image_name },
id: 'canvas_v2_mask_and_crop',
type: 'canvas_v2_mask_and_crop',
invert: true,
crop_visible: true,
});
// Resize initial image and mask to scaled size, feed into to gradient mask
@ -112,7 +112,7 @@ export const addOutpaint = async (
g.addEdge(createGradientMask, 'expanded_mask_area', resizeOutputMaskToOriginalSize, 'image');
// Finally, paste the generated masked image back onto the original image
g.addEdge(resizeOutputImageToOriginalSize, 'image', canvasPasteBack, 'target_image');
g.addEdge(resizeOutputImageToOriginalSize, 'image', canvasPasteBack, 'image');
g.addEdge(resizeOutputMaskToOriginalSize, 'image', canvasPasteBack, 'mask');
return canvasPasteBack;
@ -145,9 +145,10 @@ export const addOutpaint = async (
image: { image_name: initialImage.image_name },
});
const canvasPasteBack = g.addNode({
id: 'canvas_paste_back',
type: 'canvas_paste_back',
mask_blur: compositing.maskBlur,
id: 'canvas_v2_mask_and_crop',
type: 'canvas_v2_mask_and_crop',
invert: true,
crop_visible: true,
});
g.addEdge(maskAlphaToMask, 'image', maskCombine, 'mask1');
g.addEdge(initialImageAlphaToMask, 'image', maskCombine, 'mask2');
@ -159,8 +160,7 @@ export const addOutpaint = async (
g.addEdge(modelLoader, 'unet', createGradientMask, 'unet');
g.addEdge(createGradientMask, 'denoise_mask', denoise, 'denoise_mask');
g.addEdge(createGradientMask, 'expanded_mask_area', canvasPasteBack, 'mask');
g.addEdge(infill, 'image', canvasPasteBack, 'source_image');
g.addEdge(l2i, 'image', canvasPasteBack, 'target_image');
g.addEdge(l2i, 'image', canvasPasteBack, 'image');
return canvasPasteBack;
}

View File

@ -122,7 +122,7 @@ export const buildSD1Graph = async (state: RootState, manager: CanvasManager): P
})
: null;
let canvasOutput: Invocation<'l2i' | 'img_nsfw' | 'img_watermark' | 'img_resize' | 'canvas_paste_back'> = l2i;
let canvasOutput: Invocation<'l2i' | 'img_nsfw' | 'img_watermark' | 'img_resize' | 'canvas_v2_mask_and_crop'> = l2i;
g.addEdge(modelLoader, 'unet', denoise, 'unet');
g.addEdge(modelLoader, 'clip', clipSkip, 'clip');

View File

@ -121,7 +121,7 @@ export const buildSDXLGraph = async (state: RootState, manager: CanvasManager):
})
: null;
let canvasOutput: Invocation<'l2i' | 'img_nsfw' | 'img_watermark' | 'img_resize' | 'canvas_paste_back'> = l2i;
let canvasOutput: Invocation<'l2i' | 'img_nsfw' | 'img_watermark' | 'img_resize' | 'canvas_v2_mask_and_crop'> = l2i;
g.addEdge(modelLoader, 'unet', denoise, 'unet');
g.addEdge(modelLoader, 'clip', posCond, 'clip');