added socket listeners, added more info to ui

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();

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);
+    },
+  });
+};

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ImportQueue.tsx

@@ -45,7 +45,7 @@ export const ImportQueue = () => {
     <Box mt={3} layerStyle="first" p={3} borderRadius="base" w="full" h="full">
       <Flex direction="column" gap="2">
         {data?.map((model, i) => (
-          <Flex gap="3" w="full" alignItems="center" textAlign="center">
+          <Flex key={i} gap="3" w="full" alignItems="center" textAlign="center">
             <Text w="20%" whiteSpace="nowrap" overflow="hidden" text-overflow="ellipsis">
               {model.source.repo_id}
             </Text>

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<AddMainModelResponse, AddMainModelArg>({
       query: ({ body }) => {
@@ -308,6 +308,7 @@ export const modelsApi = api.injectEndpoints({
           url: buildModelsUrl(`import`),
         };
       },
+      providesTags: ['ModelImports']
     }),
     getCheckpointConfigs: build.query<CheckpointConfigsResponse, void>({
       query: () => {

invokeai/frontend/web/src/services/api/index.ts

@@ -28,6 +28,7 @@ export const tagTypes = [
   'InvocationCacheStatus',
   'Model',
   'ModelConfig',
+  'ModelImports',
   'T2IAdapterModel',
   'MainModel',
   'VaeModel',

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');

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;

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
    */