diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts index 9698f85219..e53baf5cbb 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts @@ -52,6 +52,7 @@ import { addWorkflowLoadRequestedListener } from 'app/store/middleware/listenerM import type { AppDispatch, RootState } from 'app/store/store'; import { addArchivedOrDeletedBoardListener } from './listeners/addArchivedOrDeletedBoardListener'; +import { addEnqueueRequestedUpscale } from './listeners/enqueueRequestedUpscale'; export const listenerMiddleware = createListenerMiddleware(); @@ -85,6 +86,7 @@ addGalleryOffsetChangedListener(startAppListening); addEnqueueRequestedCanvasListener(startAppListening); addEnqueueRequestedNodes(startAppListening); addEnqueueRequestedLinear(startAppListening); +addEnqueueRequestedUpscale(startAppListening) addAnyEnqueuedListener(startAppListening); addBatchEnqueuedListener(startAppListening); diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedUpscale.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedUpscale.ts new file mode 100644 index 0000000000..f94075d3ce --- /dev/null +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedUpscale.ts @@ -0,0 +1,53 @@ +import { enqueueRequested } from 'app/store/actions'; +import type { AppStartListening } from 'app/store/middleware/listenerMiddleware'; +import { isImageViewerOpenChanged } from 'features/gallery/store/gallerySlice'; +import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig'; +import { queueApi } from 'services/api/endpoints/queue'; +import { createIsAllowedToUpscaleSelector } from '../../../../../features/parameters/hooks/useIsAllowedToUpscale'; +import { toast } from '../../../../../features/toast/toast'; +import { t } from 'i18next'; +import { logger } from '../../../../logging/logger'; +import { useAppSelector } from '../../../storeHooks'; +import { buildMultidiffusionUpscsaleGraph } from '../../../../../features/nodes/util/graph/buildMultidiffusionUpscaleGraph'; + +export const addEnqueueRequestedUpscale = (startAppListening: AppStartListening) => { + startAppListening({ + predicate: (action): action is ReturnType => + enqueueRequested.match(action) && action.payload.tabName === 'upscaling', + effect: async (action, { getState, dispatch }) => { + console.log("in listener") + const log = logger('session'); + + const state = getState(); + const { shouldShowProgressInViewer } = state.ui; + const model = state.generation.model; + const { prepend } = action.payload; + + + const graph = await buildMultidiffusionUpscsaleGraph(state) + + + + // const batchConfig = prepareLinearUIBatch(state, graph, prepend); + + const req = dispatch( + queueApi.endpoints.enqueueBatch.initiate({ + batch: { + graph, + runs: 1, + }, + }, { + fixedCacheKey: 'enqueueBatch', + }) + ); + try { + await req.unwrap(); + if (shouldShowProgressInViewer) { + dispatch(isImageViewerOpenChanged(true)); + } + } finally { + req.reset(); + } + }, + }); +}; diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/buildMultidiffusionUpscaleGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/buildMultidiffusionUpscaleGraph.ts new file mode 100644 index 0000000000..c18d31fd8d --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/graph/buildMultidiffusionUpscaleGraph.ts @@ -0,0 +1,135 @@ +import { Graph, GraphType } from 'features/nodes/util/graph/generation/Graph'; +import { RootState } from '../../../../app/store/store'; +import { assert } from 'tsafe'; +import { ControlNetModelConfig, Invocation, NonNullableGraph } from '../../../../services/api/types'; +import { ESRGAN, NEGATIVE_CONDITIONING, POSITIVE_CONDITIONING } from './constants'; +import { isParamESRGANModelName } from '../../../parameters/store/postprocessingSlice'; +import { ControlNetConfig } from '../../../controlAdapters/store/types'; +import { MODEL_TYPES } from '../../types/constants'; + + +export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promise => { + const { + model, + cfgScale: cfg_scale, + scheduler, + steps, + vaePrecision, + seed, + vae, + } = state.generation; + const { positivePrompt, negativePrompt } = state.controlLayers.present; + const { upscaleModel, upscaleInitialImage, sharpness, structure, creativity } = state.upscale; + + assert(model, 'No model found in state'); + assert(upscaleModel, 'No upscale model found in state'); + assert(upscaleInitialImage, 'No initial image found in state'); + + if (!isParamESRGANModelName(upscaleModel.name)) { + throw new Error() + } + + const g = new Graph() + + const unsharp_mask_1 = g.addNode({ + id: 'unsharp_mask_1', + type: 'unsharp_mask', + image: upscaleInitialImage, + radius: 2, + strength: ((sharpness + 10) * 3.75) + 25 + }) + + const esrgan = g.addNode({ + id: ESRGAN, + type: 'esrgan', + model_name: upscaleModel.name, + tile_size: 500 + }) + + g.addEdge(unsharp_mask_1, 'image', esrgan, 'image') + + const unsharp_mask_2 = g.addNode({ + id: 'unsharp_mask_2', + type: 'unsharp_mask', + radius: 2, + strength: 50 + }) + + g.addEdge(esrgan, 'image', unsharp_mask_2, 'image',) + + const SCALE = 2 + + const resizeNode = g.addNode({ + id: 'img_resize', + type: 'img_resize', + width: upscaleInitialImage.width * SCALE, // TODO: handle floats + height: upscaleInitialImage.height * SCALE, // TODO: handle floats + resample_mode: "lanczos" + }) + + g.addEdge(unsharp_mask_2, 'image', resizeNode, "image") + + + + const sharpnessNode: Invocation<'unsharp_mask'> = { //before and after esrgan + id: 'unsharp_mask', + type: 'unsharp_mask', + image: upscaleInitialImage, + radius: 2, + strength: ((sharpness + 10) * 3.75) + 25 + }; + + const creativityNode: Invocation<'tiled_multi_diffusion_denoise_latents'> = { //before and after esrgan + id: 'tiled_multi_diffusion_denoise_latents', + type: 'tiled_multi_diffusion_denoise_latents', + tile_height: 1024, + tile_width: 1024, + tile_overlap: 128, + steps, + cfg_scale, + scheduler, + denoising_start: (((creativity * -1) + 10) * 4.99) / 100, + denoising_end: 1 + }; + + const controlnetModel = { + key: "placeholder", + hash: "placeholder", + type: "controlnet" as any, + name: "tile", + base: model.base + } + + const controlnet: Invocation<"controlnet"> = { + id: "controlnet", + type: "controlnet", + control_model: controlnetModel, + control_mode: "balanced", + resize_mode: "just_resize", + control_weight: ((((structure + 10) * 0.025) + 0.3) * 0.013) + 0.35 + } + + + const noiseNode: Invocation<'noise'> = { + id: "noise", + type: "noise", + seed, + // width: resized output width + // height: resized output height + } + + const posPrompt: Invocation<"compel"> = { + type: 'compel', + id: POSITIVE_CONDITIONING, + prompt: positivePrompt, + } + + const negPrompt: Invocation<"compel"> = { + type: 'compel', + id: NEGATIVE_CONDITIONING, + prompt: negativePrompt, + } + + return g.getGraph(); + +} \ No newline at end of file diff --git a/invokeai/frontend/web/src/features/parameters/store/upscaleSlice.ts b/invokeai/frontend/web/src/features/parameters/store/upscaleSlice.ts index a4d8fbf3af..67972e9660 100644 --- a/invokeai/frontend/web/src/features/parameters/store/upscaleSlice.ts +++ b/invokeai/frontend/web/src/features/parameters/store/upscaleSlice.ts @@ -9,12 +9,19 @@ interface UpscaleState { _version: 1; upscaleModel: ParameterSpandrelImageToImageModel | null; upscaleInitialImage: ImageDTO | null; + sharpness: number; + structure: number; + creativity: number; } const initialUpscaleState: UpscaleState = { _version: 1, upscaleModel: null, upscaleInitialImage: null, + sharpness: 0, + structure: 0, + creativity: 0 + }; export const upscaleSlice = createSlice({