feat(ui): rough out canvas mode

This commit is contained in:
psychedelicious 2024-08-22 19:13:27 +10:00
parent 86e8ce9139
commit 0596d25e07
15 changed files with 227 additions and 137 deletions

View File

@ -1645,6 +1645,11 @@
"storeNotInitialized": "Store is not initialized"
},
"controlLayers": {
"generateMode": "Generate",
"generateModeDesc": "Create individual images. Generated images are added directly to the gallery.",
"composeMode": "Compose",
"composeModeDesc": "Compose your work iterative. Generated images are added back to the canvas.",
"autoSave": "Auto-save to Gallery",
"resetCanvas": "Reset Canvas",
"resetAll": "Reset All",
"clearCaches": "Clear Caches",

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) {
if (!state.canvasV2.session.isStaging && state.canvasV2.session.mode === 'compose') {
dispatch(sessionStartedStaging());
didStartStaging = true;
}
@ -50,7 +50,6 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
req.reset();
await req.unwrap();
} catch (error) {
console.log('Error in enqueueRequestedLinear', error);
if (didStartStaging && getState().canvasV2.session.isStaging) {
dispatch(sessionStagingAreaReset());
}

View File

@ -0,0 +1,26 @@
import { Button, ButtonGroup } from '@invoke-ai/ui-library';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { sessionModeChanged } from 'features/controlLayers/store/canvasV2Slice';
import { memo, useCallback } from 'react';
import { useTranslation } from 'react-i18next';
export const CanvasModeSwitcher = memo(() => {
const { t } = useTranslation();
const dispatch = useAppDispatch();
const mode = useAppSelector((s) => s.canvasV2.session.mode);
const onClickGenerate = useCallback(() => dispatch(sessionModeChanged({ mode: 'generate' })), [dispatch]);
const onClickCompose = useCallback(() => dispatch(sessionModeChanged({ mode: 'compose' })), [dispatch]);
return (
<ButtonGroup variant="outline">
<Button onClick={onClickGenerate} colorScheme={mode === 'generate' ? 'invokeBlue' : 'base'}>
{t('controlLayers.generateMode')}
</Button>
<Button onClick={onClickCompose} colorScheme={mode === 'compose' ? 'invokeBlue' : 'base'}>
{t('controlLayers.composeMode')}
</Button>
</ButtonGroup>
);
});
CanvasModeSwitcher.displayName = 'CanvasModeSwitcher';

View File

@ -23,7 +23,7 @@ export const CanvasResetViewButton = memo(() => {
if (!canvasManager) {
return;
}
canvasManager.stage.resetView();
canvasManager.stage.fitLayersToStage();
}, [canvasManager]);
const onReset = useCallback(() => {

View File

@ -1,6 +1,7 @@
/* eslint-disable i18next/no-literal-string */
import { Flex, Spacer } from '@invoke-ai/ui-library';
import { useAppSelector } from 'app/store/storeHooks';
import { CanvasModeSwitcher } from 'features/controlLayers/components/CanvasModeSwitcher';
import { CanvasResetViewButton } from 'features/controlLayers/components/CanvasResetViewButton';
import { CanvasScale } from 'features/controlLayers/components/CanvasScale';
import { CanvasSettingsPopover } from 'features/controlLayers/components/Settings/CanvasSettingsPopover';
@ -23,11 +24,12 @@ export const ControlLayersToolbar = memo(() => {
<ToolChooser />
{tool === 'brush' && <ToolBrushWidth />}
{tool === 'eraser' && <ToolEraserWidth />}
<ToolFillColorPicker />
<Spacer />
<CanvasScale />
<CanvasResetViewButton />
<Spacer />
<ToolFillColorPicker />
<CanvasModeSwitcher />
<UndoRedoButtonGroup />
<CanvasSettingsPopover />
<ViewerToggleMenu />

View File

@ -70,13 +70,6 @@ export class CanvasProgressImageModule {
return;
}
const { isStaging } = this.manager.stateApi.getSession();
if (!isStaging) {
release();
return;
}
this.isLoading = true;
const { x, y, width, height } = this.manager.stateApi.getBbox().rect;

View File

@ -33,6 +33,7 @@ export class CanvasStageModule {
const resizeObserver = new ResizeObserver(this.fitStageToContainer);
resizeObserver.observe(this.container);
this.fitStageToContainer();
this.fitLayersToStage();
return () => {
this.log.debug('Destroying stage');
@ -91,10 +92,20 @@ export class CanvasStageModule {
}
};
resetView() {
this.log.trace('Resetting view');
const { width, height } = this.getSize();
fitBboxToStage = () => {
this.log.trace('Fitting bbox to stage');
const bbox = this.manager.stateApi.getBbox();
this.fitRect(bbox.rect);
};
fitLayersToStage() {
this.log.trace('Fitting layers to stage');
const rect = this.getVisibleRect();
this.fitRect(rect);
}
fitRect = (rect: Rect) => {
const { width, height } = this.getSize();
const padding = 20; // Padding in absolute pixels
@ -118,7 +129,7 @@ export class CanvasStageModule {
y,
scale,
});
}
};
/**
* Gets the center of the stage in either absolute or relative coordinates

View File

@ -132,6 +132,7 @@ const initialState: CanvasV2State = {
refinerStart: 0.8,
},
session: {
mode: 'generate',
isStaging: false,
stagedImages: [],
selectedStagedImageIndex: 0,
@ -474,6 +475,7 @@ export const {
clipToBboxChanged,
canvasReset,
settingsDynamicGridToggled,
settingsAutoSaveToggled,
// All entities
entitySelected,
entityNameChanged,
@ -601,6 +603,7 @@ export const {
sessionStagingAreaReset,
sessionNextStagedImageSelected,
sessionPrevStagedImageSelected,
sessionModeChanged,
} = canvasV2Slice.actions;
export const selectCanvasV2Slice = (state: RootState) => state.canvasV2;

View File

@ -1,5 +1,5 @@
import type { PayloadAction, SliceCaseReducers } from '@reduxjs/toolkit';
import type { CanvasV2State, StagingAreaImage } from 'features/controlLayers/store/types';
import type { CanvasV2State, SessionMode, StagingAreaImage } from 'features/controlLayers/store/types';
export const sessionReducers = {
sessionStartedStaging: (state) => {
@ -45,4 +45,8 @@ export const sessionReducers = {
state.tool.selectedBuffer = null;
}
},
sessionModeChanged: (state, action: PayloadAction<{ mode: SessionMode }>) => {
const { mode } = action.payload;
state.session.mode = mode;
},
} satisfies SliceCaseReducers<CanvasV2State>;

View File

@ -8,4 +8,7 @@ export const settingsReducers = {
settingsDynamicGridToggled: (state) => {
state.settings.dynamicGrid = !state.settings.dynamicGrid;
},
settingsAutoSaveToggled: (state) => {
state.settings.autoSave = !state.settings.autoSave;
},
} satisfies SliceCaseReducers<CanvasV2State>;

View File

@ -834,6 +834,8 @@ export type StagingAreaImage = {
offsetY: number;
};
export type SessionMode = 'generate' | 'compose';
export type CanvasV2State = {
_version: 3;
selectedEntityIdentifier: CanvasEntityIdentifier | null;
@ -929,6 +931,7 @@ export type CanvasV2State = {
refinerStart: number;
};
session: {
mode: SessionMode;
isStaging: boolean;
stagedImages: StagingAreaImage[];
selectedStagedImageIndex: number;

View File

@ -41,7 +41,7 @@ export const buildSD1Graph = async (state: RootState, manager: CanvasManager): P
const generationMode = manager.compositor.getGenerationMode();
log.debug({ generationMode }, 'Building SD1/SD2 graph');
const { bbox, params } = state.canvasV2;
const { bbox, params, session, settings } = state.canvasV2;
const {
model,
@ -249,10 +249,11 @@ export const buildSD1Graph = async (state: RootState, manager: CanvasManager): P
canvasOutput = addWatermarker(g, canvasOutput);
}
// This is the terminal node and must always save to gallery.
const shouldSaveToGallery = session.mode === 'generate' || settings.autoSave;
g.updateNode(canvasOutput, {
id: CANVAS_OUTPUT,
is_intermediate: false,
is_intermediate: !shouldSaveToGallery,
use_cache: false,
board: getBoardField(state),
});

View File

@ -40,7 +40,7 @@ export const buildSDXLGraph = async (state: RootState, manager: CanvasManager):
const generationMode = manager.compositor.getGenerationMode();
log.debug({ generationMode }, 'Building SDXL graph');
const { bbox, params } = state.canvasV2;
const { bbox, params, session, settings } = state.canvasV2;
const {
model,
@ -246,10 +246,11 @@ export const buildSDXLGraph = async (state: RootState, manager: CanvasManager):
canvasOutput = addWatermarker(g, canvasOutput);
}
// This is the terminal node and must always save to gallery.
const shouldSaveToGallery = session.mode === 'generate' || settings.autoSave;
g.updateNode(canvasOutput, {
id: CANVAS_OUTPUT,
is_intermediate: false,
is_intermediate: !shouldSaveToGallery,
use_cache: false,
board: getBoardField(state),
});

View File

@ -0,0 +1,145 @@
import { logger } from 'app/logging/logger';
import type { AppDispatch, RootState } from 'app/store/store';
import type { SerializableObject } from 'common/types';
import { deepClone } from 'common/util/deepClone';
import { sessionImageStaged } from 'features/controlLayers/store/canvasV2Slice';
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 { boardsApi } from 'services/api/endpoints/boards';
import { getImageDTO, imagesApi } from 'services/api/endpoints/images';
import type { ImageDTO } from 'services/api/types';
import { getCategories, getListImagesUrl } from 'services/api/util';
import type { InvocationCompleteEvent, InvocationDenoiseProgressEvent } from 'services/events/types';
const log = logger('events');
export const buildOnInvocationComplete = (
getState: () => RootState,
dispatch: AppDispatch,
nodeTypeDenylist: string[],
setLastProgressEvent: (event: InvocationDenoiseProgressEvent | null) => void,
setLastCanvasProgressEvent: (event: InvocationDenoiseProgressEvent | null) => void
) => {
const addImageToGallery = (imageDTO: ImageDTO) => {
if (imageDTO.is_intermediate) {
return;
}
// update the total images for the board
dispatch(
boardsApi.util.updateQueryData('getBoardImagesTotal', imageDTO.board_id ?? 'none', (draft) => {
// eslint-disable-next-line @typescript-eslint/no-unused-vars
draft.total += 1;
})
);
dispatch(
imagesApi.util.invalidateTags([
{ type: 'Board', id: imageDTO.board_id ?? 'none' },
{
type: 'ImageList',
id: getListImagesUrl({
board_id: imageDTO.board_id ?? 'none',
categories: getCategories(imageDTO),
}),
},
])
);
const { shouldAutoSwitch, galleryView, selectedBoardId } = getState().gallery;
// If auto-switch is enabled, select the new image
if (shouldAutoSwitch) {
// if auto-add is enabled, switch the gallery view and board if needed as the image comes in
if (galleryView !== 'images') {
dispatch(galleryViewChanged('images'));
}
if (imageDTO.board_id && imageDTO.board_id !== selectedBoardId) {
dispatch(
boardIdSelected({
boardId: imageDTO.board_id,
selectedImageName: imageDTO.image_name,
})
);
}
dispatch(offsetChanged({ offset: 0 }));
if (!imageDTO.board_id && selectedBoardId !== 'none') {
dispatch(
boardIdSelected({
boardId: 'none',
selectedImageName: imageDTO.image_name,
})
);
}
dispatch(imageSelected(imageDTO));
}
};
return async (data: InvocationCompleteEvent) => {
log.debug(
{ data } as SerializableObject,
`Invocation complete (${data.invocation.type}, ${data.invocation_source_id})`
);
const { result, invocation_source_id } = data;
// Update the node execution states - the image output is handled below
if (data.origin === 'workflows') {
const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]);
if (nes) {
nes.status = zNodeStatus.enum.COMPLETED;
if (nes.progress !== null) {
nes.progress = 1;
}
nes.outputs.push(result);
upsertExecutionState(nes.nodeId, nes);
}
}
// This complete event has an associated image output
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 { session } = getState().canvasV2;
const imageDTO = await getImageDTO(image_name);
if (!imageDTO) {
log.error({ data } as SerializableObject, 'Failed to fetch image DTO after generation');
return;
}
if (data.origin === 'canvas') {
if (data.invocation_source_id !== 'canvas_output') {
// Not a canvas output image - ignore
return;
}
if (session.mode === 'compose' && session.isStaging) {
if (data.result.type === 'canvas_v2_mask_and_crop_output') {
const { offset_x, offset_y } = data.result;
if (session.isStaging) {
dispatch(sessionImageStaged({ stagingAreaImage: { imageDTO, offsetX: offset_x, offsetY: offset_y } }));
}
} else if (data.result.type === 'image_output') {
if (session.isStaging) {
dispatch(sessionImageStaged({ stagingAreaImage: { imageDTO, offsetX: 0, offsetY: 0 } }));
}
}
addImageToGallery(imageDTO);
} else {
addImageToGallery(imageDTO);
setLastCanvasProgressEvent(null);
}
}
}
setLastProgressEvent(null);
};
};

View File

@ -7,8 +7,6 @@ import { $queueId } from 'app/store/nanostores/queueId';
import type { AppDispatch, RootState } from 'app/store/store';
import type { SerializableObject } from 'common/types';
import { deepClone } from 'common/util/deepClone';
import { sessionImageStaged } from 'features/controlLayers/store/canvasV2Slice';
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 ErrorToastDescription, { getTitleFromErrorType } from 'features/toast/ErrorToastDescription';
@ -17,11 +15,9 @@ import { t } from 'i18next';
import { forEach } from 'lodash-es';
import { atom, computed } from 'nanostores';
import { api, LIST_TAG } from 'services/api';
import { boardsApi } from 'services/api/endpoints/boards';
import { imagesApi } from 'services/api/endpoints/images';
import { modelsApi } from 'services/api/endpoints/models';
import { queueApi, queueItemsAdapter } from 'services/api/endpoints/queue';
import { getCategories, getListImagesUrl } from 'services/api/util';
import { buildOnInvocationComplete } from 'services/events/onInvocationComplete';
import type { ClientToServerEvents, InvocationDenoiseProgressEvent, ServerToClientEvents } from 'services/events/types';
import type { Socket } from 'socket.io-client';
@ -147,116 +143,14 @@ export const setEventListeners = ({ socket, dispatch, getState, setIsConnected }
}
});
socket.on('invocation_complete', async (data) => {
log.debug(
{ data } as SerializableObject,
`Invocation complete (${data.invocation.type}, ${data.invocation_source_id})`
);
const { result, invocation_source_id } = data;
if (data.origin === 'workflows') {
const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]);
if (nes) {
nes.status = zNodeStatus.enum.COMPLETED;
if (nes.progress !== null) {
nes.progress = 1;
}
nes.outputs.push(result);
upsertExecutionState(nes.nodeId, nes);
}
}
// This complete event has an associated image output
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();
// This populates the `getImageDTO` cache
const imageDTORequest = dispatch(
imagesApi.endpoints.getImageDTO.initiate(image_name, {
forceRefetch: true,
})
);
const imageDTO = await imageDTORequest.unwrap();
imageDTORequest.unsubscribe();
// handle tab-specific logic
if (data.origin === 'canvas' && data.invocation_source_id === 'canvas_output') {
if (data.result.type === 'canvas_v2_mask_and_crop_output') {
const { offset_x, offset_y } = data.result;
if (canvasV2.session.isStaging) {
dispatch(sessionImageStaged({ stagingAreaImage: { imageDTO, offsetX: offset_x, offsetY: offset_y } }));
}
} else if (data.result.type === 'image_output') {
if (canvasV2.session.isStaging) {
dispatch(sessionImageStaged({ stagingAreaImage: { imageDTO, offsetX: 0, offsetY: 0 } }));
}
}
}
if (!imageDTO.is_intermediate) {
// update the total images for the board
dispatch(
boardsApi.util.updateQueryData('getBoardImagesTotal', imageDTO.board_id ?? 'none', (draft) => {
// eslint-disable-next-line @typescript-eslint/no-unused-vars
draft.total += 1;
})
);
dispatch(
imagesApi.util.invalidateTags([
{ type: 'Board', id: imageDTO.board_id ?? 'none' },
{
type: 'ImageList',
id: getListImagesUrl({
board_id: imageDTO.board_id ?? 'none',
categories: getCategories(imageDTO),
}),
},
])
);
const { shouldAutoSwitch } = gallery;
// If auto-switch is enabled, select the new image
if (shouldAutoSwitch) {
// if auto-add is enabled, switch the gallery view and board if needed as the image comes in
if (gallery.galleryView !== 'images') {
dispatch(galleryViewChanged('images'));
}
if (imageDTO.board_id && imageDTO.board_id !== gallery.selectedBoardId) {
dispatch(
boardIdSelected({
boardId: imageDTO.board_id,
selectedImageName: imageDTO.image_name,
})
);
}
dispatch(offsetChanged({ offset: 0 }));
if (!imageDTO.board_id && gallery.selectedBoardId !== 'none') {
dispatch(
boardIdSelected({
boardId: 'none',
selectedImageName: imageDTO.image_name,
})
);
}
dispatch(imageSelected(imageDTO));
}
}
}
$lastProgressEvent.set(null);
});
const onInvocationComplete = buildOnInvocationComplete(
getState,
dispatch,
nodeTypeDenylist,
$lastProgressEvent.set,
$lastCanvasProgressEvent.set
);
socket.on('invocation_complete', onInvocationComplete);
socket.on('model_load_started', (data) => {
const { config, submodel_type } = data;