github/InvokeAI
1
1
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-02-12 20:14:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat(ui): move socket events handling into ctx component

Browse Source
This commit is contained in:
psychedelicious
2025-06-04 23:47:29 +10:00
parent 002816653e
commit 628367b97b
3 changed files with 118 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingArea.tsx
View File

@@ -7,14 +7,12 @@ import { getQueueItemElementId } from 'features/controlLayers/components/SimpleS
import { StagingAreaContent } from 'features/controlLayers/components/SimpleSession/StagingAreaContent';
import { StagingAreaHeader } from 'features/controlLayers/components/SimpleSession/StagingAreaHeader';
import { StagingAreaNoItems } from 'features/controlLayers/components/SimpleSession/StagingAreaNoItems';
import { useProgressEvents } from 'features/controlLayers/components/SimpleSession/use-progress-events';
import { useStagingAreaKeyboardNav } from 'features/controlLayers/components/SimpleSession/use-staging-keyboard-nav';
import { memo, useEffect } from 'react';
export const StagingArea = memo(() => {
const ctx = useCanvasSessionContext();
const hasItems = useStore(ctx.$hasItems);
useProgressEvents();
useStagingAreaKeyboardNav();
useEffect(() => {

128
invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx
View File

@@ -1,3 +1,4 @@
import { useStore } from '@nanostores/react';
import { createSelector } from '@reduxjs/toolkit';
import { EMPTY_ARRAY } from 'app/store/constants';
import { useAppStore } from 'app/store/nanostores/store';
@@ -12,6 +13,7 @@ import type { PropsWithChildren } from 'react';
import { createContext, memo, useContext, useEffect, useMemo, useState } from 'react';
import { queueApi } from 'services/api/endpoints/queue';
import type { S } from 'services/api/types';
import { $socket } from 'services/events/stores';
import { assert } from 'tsafe';
export type ProgressData = {
@@ -134,13 +136,48 @@ const CanvasSessionContext = createContext<CanvasSessionContextValue | null>(nul
export const CanvasSessionContextProvider = memo(
({ session, children }: PropsWithChildren<{ session: SimpleSessionIdentifier | AdvancedSessionIdentifier }>) => {
/**
* For best performance and interop with the Canvas, which is outside react but needs to interact with the react
* app, all canvas session state is packaged as nanostores atoms. The trickiest part is syncing the queue items
* with a nanostores atom.
*/
/**
* App store
*/
const store = useAppStore();
const [$items] = useState(() => atom<S['SessionQueueItem'][]>([]));
const [$hasItems] = useState(() => computed([$items], (items) => items.length > 0));
const [$autoSwitch] = useState(() => atom(true));
const [$selectedItemId] = useState(() => atom<number | null>(null));
const [$progressData] = useState(() => atom<Record<string, ProgressData>>({}));
const [$selectedItem] = useState(() =>
const socket = useStore($socket);
/**
* Manually-synced atom containing the queue items for the current session.
*/
const $items = useState(() => atom<S['SessionQueueItem'][]>([]))[0];
/**
* Whether auto-switch is enabled.
*/
const $autoSwitch = useState(() => atom(true))[0];
/**
* An ephemeral store of progress events and images for all items in the current session.
*/
const $progressData = useState(() => atom<Record<string, ProgressData>>({}))[0];
/**
* The currently selected queue item's ID, or null if one is not selected.
*/
const $selectedItemId = useState(() => atom<number | null>(null))[0];
/**
* Whether there are any items. Computed from the queue items array.
*/
const $hasItems = useState(() => computed([$items], (items) => items.length > 0))[0];
/**
* The currently selected queue item, or null if one is not selected.
*/
const $selectedItem = useState(() =>
computed([$items, $selectedItemId], (items, selectedItemId) => {
if (items.length === 0) {
return null;
@@ -150,8 +187,12 @@ export const CanvasSessionContextProvider = memo(
}
return items.find(({ item_id }) => item_id === selectedItemId) ?? null;
})
);
const [$selectedItemIndex] = useState(() =>
)[0];
/**
* The currently selected queue item's index in the list of items, or null if one is not selected.
*/
const $selectedItemIndex = useState(() =>
computed([$items, $selectedItemId], (items, selectedItemId) => {
if (items.length === 0) {
return null;
@@ -161,20 +202,59 @@ export const CanvasSessionContextProvider = memo(
}
return items.findIndex(({ item_id }) => item_id === selectedItemId) ?? null;
})
);
)[0];
/**
* A redux selector to select all queue items from the RTK Query cache. It's important that this returns stable
* references if possible to reduce re-renders. All derivations of the queue items (e.g. filtering out canceled
* items) should be done in a nanostores computed.
*/
const selectQueueItems = useMemo(
() =>
createSelector(
queueApi.endpoints.listAllQueueItems.select({ destination: session.id }),
({ data }) => data?.filter((item) => item.status !== 'canceled') ?? EMPTY_ARRAY
({ data }) => data ?? EMPTY_ARRAY
),
[session.id]
);
// Set up socket listeners
useEffect(() => {
if (!socket) {
return;
}
const onProgress = (data: S['InvocationProgressEvent']) => {
if (data.destination !== session.id) {
return;
}
setProgress($progressData, data);
};
const onQueueItemStatusChanged = (data: S['QueueItemStatusChangedEvent']) => {
if (data.destination !== session.id) {
return;
}
if (data.status === 'completed' && $autoSwitch.get()) {
$selectedItemId.set(data.item_id);
}
};
socket.on('invocation_progress', onProgress);
socket.on('queue_item_status_changed', onQueueItemStatusChanged);
return () => {
socket.off('invocation_progress', onProgress);
socket.off('queue_item_status_changed', onQueueItemStatusChanged);
};
}, [$autoSwitch, $progressData, $selectedItemId, session.id, socket]);
// Set up state subscriptions and effects
useEffect(() => {
// Seed the $items atom with the initial query cache state
$items.set(selectQueueItems(store.getState()));
// Manually keep the $items atom in sync as the query cache is updated
const unsubReduxSyncToItemsAtom = store.subscribe(() => {
const prevItems = $items.get();
const items = selectQueueItems(store.getState());
@@ -183,29 +263,57 @@ export const CanvasSessionContextProvider = memo(
}
});
// Handle cases that could result in a nonexistent queue item being selected.
const unsubEnsureSelectedItemIdExists = effect([$items, $selectedItemId], (items, selectedItemId) => {
// If there are no items, cannot have a selected item.
if (items.length === 0) {
$selectedItemId.set(null);
return;
}
// If there is no selected item but there are items, select the first one.
if (selectedItemId === null && items.length > 0) {
$selectedItemId.set(items[0]?.item_id ?? null);
return;
}
// If an item is selected and it is not in the list of items, un-set it. This effect will run again and we'll
// the above case, selecting the first item if there are any.
if (selectedItemId !== null && items.findIndex(({ item_id }) => item_id === selectedItemId) === -1) {
$selectedItemId.set(null);
return;
}
});
// Clean up the progress data when a queue item is discarded.
const unsubCleanUpProgressData = effect([$items, $progressData], (items, progressData) => {
const toDelete: string[] = [];
for (const datum of Object.values(progressData)) {
if (items.findIndex(({ session_id }) => session_id === datum.sessionId) === -1) {
toDelete.push(datum.sessionId);
}
}
if (toDelete.length === 0) {
return;
}
const newProgressData = { ...progressData };
for (const sessionId of toDelete) {
delete newProgressData[sessionId];
}
// This will re-trigger the effect - maybe this could just be a listener on $items? Brain hurt
$progressData.set(newProgressData);
});
// Create an RTK Query subscription. Without this, the query cache selector will never return anything bc RTK
// doesn't know we care about it.
const { unsubscribe: unsubQueueItemsQuery } = store.dispatch(
queueApi.endpoints.listAllQueueItems.initiate({ destination: session.id })
);
// Clean up all subscriptions and top-level (i.e. non-computed/derived state)
return () => {
unsubQueueItemsQuery();
unsubReduxSyncToItemsAtom();
unsubEnsureSelectedItemIdExists();
unsubCleanUpProgressData();
$items.set([]);
$progressData.set({});
$selectedItemId.set(null);

48
invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/use-progress-events.ts
View File

@@ -1,48 +0,0 @@
import { useStore } from '@nanostores/react';
import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
import { useEffect } from 'react';
import type { S } from 'services/api/types';
import { $socket, setProgress } from 'services/events/stores';
export const useProgressEvents = () => {
const ctx = useCanvasSessionContext();
const socket = useStore($socket);
useEffect(() => {
if (!socket) {
return;
}
const onQueueItemStatusChanged = (data: S['QueueItemStatusChangedEvent']) => {
if (data.destination !== ctx.session.id) {
return;
}
if (data.status === 'completed' && ctx.$autoSwitch.get()) {
ctx.$selectedItemId.set(data.item_id);
}
};
socket.on('queue_item_status_changed', onQueueItemStatusChanged);
return () => {
socket.off('queue_item_status_changed', onQueueItemStatusChanged);
};
}, [ctx.$autoSwitch, ctx.$progressData, ctx.$selectedItemId, ctx.session.id, socket]);
useEffect(() => {
if (!socket) {
return;
}
const onProgress = (data: S['InvocationProgressEvent']) => {
if (data.destination !== ctx.session.id) {
return;
}
// TODO: clear progress when done w/ it memory leak
setProgress(ctx.$progressData, data);
};
socket.on('invocation_progress', onProgress);
return () => {
socket.off('invocation_progress', onProgress);
};
}, [ctx.$progressData, ctx.session.id, socket]);
};