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 23e23c1140..26d7949e9c 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts @@ -58,6 +58,7 @@ import { addInvocationErrorEventListener as addInvocationErrorListener } from '. import { addInvocationRetrievalErrorEventListener } from './listeners/socketio/socketInvocationRetrievalError'; import { addInvocationStartedEventListener as addInvocationStartedListener } from './listeners/socketio/socketInvocationStarted'; import { addModelLoadEventListener } from './listeners/socketio/socketModelLoad'; +import { addModelInstallEventListener } from './listeners/socketio/socketModelInstall'; import { addSocketQueueItemStatusChangedEventListener } from './listeners/socketio/socketQueueItemStatusChanged'; import { addSessionRetrievalErrorEventListener } from './listeners/socketio/socketSessionRetrievalError'; import { addSocketSubscribedEventListener as addSocketSubscribedListener } from './listeners/socketio/socketSubscribed'; @@ -135,6 +136,7 @@ addSocketDisconnectedListener(); addSocketSubscribedListener(); addSocketUnsubscribedListener(); addModelLoadEventListener(); +addModelInstallEventListener(); addSessionRetrievalErrorEventListener(); addInvocationRetrievalErrorEventListener(); addSocketQueueItemStatusChangedEventListener(); diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall.ts new file mode 100644 index 0000000000..b241df27c8 --- /dev/null +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall.ts @@ -0,0 +1,63 @@ +import { logger } from 'app/logging/logger'; +import { + socketModelInstallCompleted, + socketModelInstallDownloading, + socketModelInstallError, +} from 'services/events/actions'; +import { modelsApi } from 'services/api/endpoints/models'; +import type { components, paths } from 'services/api/schema'; + + +import { startAppListening } from '../..'; +import { createEntityAdapter } from '@reduxjs/toolkit'; + +const log = logger('socketio'); + +export const addModelInstallEventListener = () => { + startAppListening({ + actionCreator: socketModelInstallDownloading, + effect: async (action, { dispatch }) => { + const { bytes, local_path, source, timestamp, total_bytes } = action.payload.data; + let message = `Model install started: ${bytes}/${total_bytes}/${source}`; + // below doesnt work, still not sure how to update the importModels data + // dispatch( + // modelsApi.util.updateQueryData('getModelImports', undefined, (draft) => { + // importModelsAdapter.updateOne(draft, { + // id: source, + // changes: { + // bytes, + // total_bytes, + // },\q + // }); + // } + // ) + // ); + + log.debug(action.payload, message); + }, + }); + + startAppListening({ + actionCreator: socketModelInstallCompleted, + effect: (action) => { + const { key, source, timestamp } = action.payload.data; + + let message = `Model install completed: ${source}`; + + // dispatch something that marks the model installation as completed + + log.debug(action.payload, message); + }, + }); + + startAppListening({ + actionCreator: socketModelInstallError, + effect: (action) => { + const { error, error_type, source } = action.payload.data; + + // dispatch something that marks the model installation as errored + + log.debug(action.payload, error); + }, + }); +}; diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ImportQueue.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ImportQueue.tsx index 4a0407ff1a..c2dbece4b3 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ImportQueue.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ImportQueue.tsx @@ -45,7 +45,7 @@ export const ImportQueue = () => { {data?.map((model, i) => ( - + {model.source.repo_id} diff --git a/invokeai/frontend/web/src/services/api/endpoints/models.ts b/invokeai/frontend/web/src/services/api/endpoints/models.ts index 71e539e82c..2552a7e839 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/models.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/models.ts @@ -202,7 +202,7 @@ export const modelsApi = api.injectEndpoints({ body: config, }; }, - invalidatesTags: ['Model'], + invalidatesTags: ['Model', 'ModelImports'], }), addMainModels: build.mutation({ query: ({ body }) => { @@ -308,6 +308,7 @@ export const modelsApi = api.injectEndpoints({ url: buildModelsUrl(`import`), }; }, + providesTags: ['ModelImports'] }), getCheckpointConfigs: build.query({ query: () => { diff --git a/invokeai/frontend/web/src/services/api/index.ts b/invokeai/frontend/web/src/services/api/index.ts index 3b5a10492c..30cc7a9d3d 100644 --- a/invokeai/frontend/web/src/services/api/index.ts +++ b/invokeai/frontend/web/src/services/api/index.ts @@ -28,6 +28,7 @@ export const tagTypes = [ 'InvocationCacheStatus', 'Model', 'ModelConfig', + 'ModelImports', 'T2IAdapterModel', 'MainModel', 'VaeModel', diff --git a/invokeai/frontend/web/src/services/events/actions.ts b/invokeai/frontend/web/src/services/events/actions.ts index b80363315e..c306b3d2f6 100644 --- a/invokeai/frontend/web/src/services/events/actions.ts +++ b/invokeai/frontend/web/src/services/events/actions.ts @@ -11,6 +11,9 @@ import type { InvocationStartedEvent, ModelLoadCompletedEvent, ModelLoadStartedEvent, + ModelInstallDownloadingEvent, + ModelInstallCompletedEvent, + ModelInstallErrorEvent, QueueItemStatusChangedEvent, SessionRetrievalErrorEvent, } from 'services/events/types'; @@ -56,6 +59,18 @@ export const socketModelLoadCompleted = createAction<{ data: ModelLoadCompletedEvent; }>('socket/socketModelLoadCompleted'); +export const socketModelInstallDownloading = createAction<{ + data: ModelInstallDownloadingEvent; +}>('socket/socketModelInstallDownloading'); + +export const socketModelInstallCompleted = createAction<{ + data: ModelInstallCompletedEvent; +}>('socket/socketModelInstallCompleted'); + +export const socketModelInstallError = createAction<{ + data: ModelInstallErrorEvent; +}>('socket/socketModelInstallError'); + export const socketSessionRetrievalError = createAction<{ data: SessionRetrievalErrorEvent; }>('socket/socketSessionRetrievalError'); diff --git a/invokeai/frontend/web/src/services/events/types.ts b/invokeai/frontend/web/src/services/events/types.ts index 4e68131ee9..2cf73a8e69 100644 --- a/invokeai/frontend/web/src/services/events/types.ts +++ b/invokeai/frontend/web/src/services/events/types.ts @@ -45,6 +45,27 @@ export type ModelLoadCompletedEvent = { precision: string; }; +export type ModelInstallDownloadingEvent = { + bytes: string; + local_path: string; + source: string; + timestamp: number; + total_bytes: string; +}; + +export type ModelInstallCompletedEvent = { + key: number; + source: string; + timestamp: number; +}; + +export type ModelInstallErrorEvent = { + error: string; + error_type: string; + source: string; + timestamp: number; +}; + /** * A `generator_progress` socket.io event. * @@ -237,6 +258,9 @@ export type ServerToClientEvents = { graph_execution_state_complete: (payload: GraphExecutionStateCompleteEvent) => void; model_load_started: (payload: ModelLoadStartedEvent) => void; model_load_completed: (payload: ModelLoadCompletedEvent) => void; + model_install_downloading: (payload: ModelInstallDownloadingEvent) => void; + model_install_completed: (payload: ModelInstallCompletedEvent) => void; + model_install_error: (payload: ModelInstallErrorEvent) => void; session_retrieval_error: (payload: SessionRetrievalErrorEvent) => void; invocation_retrieval_error: (payload: InvocationRetrievalErrorEvent) => void; queue_item_status_changed: (payload: QueueItemStatusChangedEvent) => void; diff --git a/invokeai/frontend/web/src/services/events/util/setEventListeners.ts b/invokeai/frontend/web/src/services/events/util/setEventListeners.ts index 1f27955ca7..314f9a985f 100644 --- a/invokeai/frontend/web/src/services/events/util/setEventListeners.ts +++ b/invokeai/frontend/web/src/services/events/util/setEventListeners.ts @@ -18,6 +18,9 @@ import { socketInvocationStarted, socketModelLoadCompleted, socketModelLoadStarted, + socketModelInstallDownloading, + socketModelInstallCompleted, + socketModelInstallError, socketQueueItemStatusChanged, socketSessionRetrievalError, } from 'services/events/actions'; @@ -134,6 +137,39 @@ export const setEventListeners = (arg: SetEventListenersArg) => { ); }); + /** + * Model Install Downloading + */ + socket.on('model_install_downloading', (data) => { + dispatch( + socketModelInstallDownloading({ + data, + }) + ); + }) + + /** + * Model Install Completed + */ + socket.on('model_install_completed', (data) => { + dispatch( + socketModelInstallCompleted({ + data, + }) + ); + }) + + /** + * Model Install Error + */ + socket.on('model_install_error', (data) => { + dispatch( + socketModelInstallError({ + data, + }) + ); + }) + /** * Session retrieval error */