InvokeAI/frontend/src/features/system/systemSlice.ts

import { createSlice } from '@reduxjs/toolkit';
import type { PayloadAction } from '@reduxjs/toolkit';
import { ExpandedIndex } from '@chakra-ui/react';
import * as InvokeAI from '../../app/invokeai';

export type LogLevel = 'info' | 'warning' | 'error';

export interface LogEntry {
  timestamp: string;
  level: LogLevel;
  message: string;
}

export interface Log {
  [index: number]: LogEntry;
}

export interface SystemState
  extends InvokeAI.SystemStatus,
    InvokeAI.SystemConfig {
  shouldDisplayInProgress: boolean;
  shouldDisplayInProgressLatents: boolean;
  log: Array<LogEntry>;
  shouldShowLogViewer: boolean;
  isGFPGANAvailable: boolean;
  isESRGANAvailable: boolean;
  isConnected: boolean;
  socketId: string;
  shouldConfirmOnDelete: boolean;
  openAccordions: ExpandedIndex;
  currentStep: number;
  totalSteps: number;
  currentIteration: number;
  totalIterations: number;
  currentStatus: string;
  currentStatusHasSteps: boolean;
  shouldDisplayGuides: boolean;
  wasErrorSeen: boolean;
  isCancelable: boolean;
}

const initialSystemState = {
  isConnected: false,
  isProcessing: false,
  log: [],
  shouldShowLogViewer: false,
  shouldDisplayInProgress: false,
  shouldDisplayInProgressLatents: false,
  shouldDisplayGuides: true,
  isGFPGANAvailable: true,
  isESRGANAvailable: true,
  socketId: '',
  shouldConfirmOnDelete: true,
  openAccordions: [0],
  currentStep: 0,
  totalSteps: 0,
  currentIteration: 0,
  totalIterations: 0,
  currentStatus: 'Disconnected',
  currentStatusHasSteps: false,
  model: '',
  model_id: '',
  model_hash: '',
  app_id: '',
  app_version: '',
  model_list: {},
  hasError: false,
  wasErrorSeen: true,
  isCancelable: true,
};

const initialState: SystemState = initialSystemState;

export const systemSlice = createSlice({
  name: 'system',
  initialState,
  reducers: {
    setShouldDisplayInProgress: (state, action: PayloadAction<boolean>) => {
      state.shouldDisplayInProgress = action.payload;
    },
    setShouldDisplayInProgressLatents: (state, action: PayloadAction<boolean>) => {
      state.shouldDisplayInProgressLatents = action.payload;
    },
    setIsProcessing: (state, action: PayloadAction<boolean>) => {
      state.isProcessing = action.payload;
    },
    setCurrentStatus: (state, action: PayloadAction<string>) => {
      state.currentStatus = action.payload;
    },
    setSystemStatus: (state, action: PayloadAction<InvokeAI.SystemStatus>) => {
      return { ...state, ...action.payload };
    },
    errorOccurred: (state) => {
      state.hasError = true;
      state.isProcessing = false;
      state.isCancelable = true;
      state.currentStep = 0;
      state.totalSteps = 0;
      state.currentIteration = 0;
      state.totalIterations = 0;
      state.currentStatusHasSteps = false;
      state.currentStatus = 'Error';
      state.wasErrorSeen = false;
    },
    errorSeen: (state) => {
      state.hasError = false;
      state.wasErrorSeen = true;
      state.currentStatus = state.isConnected ? 'Connected' : 'Disconnected';
    },
    addLogEntry: (
      state,
      action: PayloadAction<{
        timestamp: string;
        message: string;
        level?: LogLevel;
      }>
    ) => {
      const { timestamp, message, level } = action.payload;
      const logLevel = level || 'info';

      const entry: LogEntry = {
        timestamp,
        message,
        level: logLevel,
      };

      state.log.push(entry);
    },
    setShouldShowLogViewer: (state, action: PayloadAction<boolean>) => {
      state.shouldShowLogViewer = action.payload;
    },
    setIsConnected: (state, action: PayloadAction<boolean>) => {
      state.isConnected = action.payload;
      state.isProcessing = false;
      state.isCancelable = true;
      state.currentStep = 0;
      state.totalSteps = 0;
      state.currentIteration = 0;
      state.totalIterations = 0;
      state.currentStatusHasSteps = false;
      state.hasError = false;
    },
    setSocketId: (state, action: PayloadAction<string>) => {
      state.socketId = action.payload;
    },
    setShouldConfirmOnDelete: (state, action: PayloadAction<boolean>) => {
      state.shouldConfirmOnDelete = action.payload;
    },
    setOpenAccordions: (state, action: PayloadAction<ExpandedIndex>) => {
      state.openAccordions = action.payload;
    },
    setSystemConfig: (state, action: PayloadAction<InvokeAI.SystemConfig>) => {
      return {
        ...state,
        ...action.payload,
      };
    },
    setShouldDisplayGuides: (state, action: PayloadAction<boolean>) => {
      state.shouldDisplayGuides = action.payload;
    },
    processingCanceled: (state) => {
      state.isProcessing = false;
      state.isCancelable = true;
      state.currentStep = 0;
      state.totalSteps = 0;
      state.currentIteration = 0;
      state.totalIterations = 0;
      state.currentStatusHasSteps = false;
      state.currentStatus = 'Processing canceled';
    },
    setModelList: (
      state,
      action: PayloadAction<InvokeAI.ModelList | Record<string, never>>
    ) => {
      state.model_list = action.payload;
    },
    setIsCancelable: (state, action: PayloadAction<boolean>) => {
      state.isCancelable = action.payload;
    },
    modelChangeRequested: (state) => {
      state.currentStatus = 'Loading Model';
      state.isCancelable = false;
      state.isProcessing = true;
      state.currentStatusHasSteps = false;
    },
  },
});

export const {
  setShouldDisplayInProgress,
  setShouldDisplayInProgressLatents,
  setIsProcessing,
  addLogEntry,
  setShouldShowLogViewer,
  setIsConnected,
  setSocketId,
  setShouldConfirmOnDelete,
  setOpenAccordions,
  setSystemStatus,
  setCurrentStatus,
  setSystemConfig,
  setShouldDisplayGuides,
  processingCanceled,
  errorOccurred,
  errorSeen,
  setModelList,
  setIsCancelable,
  modelChangeRequested,
} = systemSlice.actions;

export default systemSlice.reducer;