From 40ed218c264b0ae02bcba6ef2a5244bff24cd166 Mon Sep 17 00:00:00 2001
From: Mary Hipp <maryhipp@Marys-MacBook-Air.local>
Date: Wed, 27 Sep 2023 11:59:44 -0400
Subject: [PATCH] surface usage errors for cnet and upscale, handle clearing
 cnet if error occurs

---
 .../listeners/controlNetImageProcessed.ts     | 32 +++++++++++++++++--
 .../listeners/upscaleRequested.ts             | 21 +++++++++++-
 .../controlNet/store/controlNetSlice.ts       |  4 +++
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetImageProcessed.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetImageProcessed.ts
index 814614da10..2c3bf0ff5e 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetImageProcessed.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetImageProcessed.ts
@@ -1,7 +1,11 @@
 import { logger } from 'app/logging/logger';
 import { parseify } from 'common/util/serialize';
 import { controlNetImageProcessed } from 'features/controlNet/store/actions';
-import { controlNetProcessedImageChanged } from 'features/controlNet/store/controlNetSlice';
+import {
+  clearPendingControlImages,
+  controlNetImageChanged,
+  controlNetProcessedImageChanged,
+} from 'features/controlNet/store/controlNetSlice';
 import { SAVE_IMAGE } from 'features/nodes/util/graphBuilders/constants';
 import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
@@ -11,6 +15,7 @@ import { isImageOutput } from 'services/api/guards';
 import { Graph, ImageDTO } from 'services/api/types';
 import { socketInvocationComplete } from 'services/events/actions';
 import { startAppListening } from '..';
+import { FetchBaseQueryError } from '@reduxjs/toolkit/dist/query';
 
 export const addControlNetImageProcessedListener = () => {
   startAppListening({
@@ -105,8 +110,31 @@ export const addControlNetImageProcessedListener = () => {
             })
           );
         }
-      } catch {
+      } catch (error) {
         log.error({ graph: parseify(graph) }, t('queue.graphFailedToQueue'));
+
+        // handle usage-related errors
+        if (error instanceof Object) {
+          if ('data' in error && 'status' in error) {
+            if (error.status === 403) {
+              const detail = (error.data as any)?.detail || 'Unknown Error';
+              dispatch(
+                addToast({
+                  title: t('queue.graphFailedToQueue'),
+                  status: 'error',
+                  description: detail,
+                  duration: 15000,
+                })
+              );
+              dispatch(clearPendingControlImages());
+              dispatch(
+                controlNetImageChanged({ controlNetId, controlImage: null })
+              );
+              return;
+            }
+          }
+        }
+
         dispatch(
           addToast({
             title: t('queue.graphFailedToQueue'),
diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/upscaleRequested.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/upscaleRequested.ts
index b54d8b553c..e993fa75de 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/upscaleRequested.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/upscaleRequested.ts
@@ -44,8 +44,27 @@ export const addUpscaleRequestedListener = () => {
           { enqueueResult: parseify(enqueueResult) },
           t('queue.graphQueued')
         );
-      } catch {
+      } catch (error) {
         log.error({ graph: parseify(graph) }, t('queue.graphFailedToQueue'));
+
+        // handle usage-related errors
+        if (error instanceof Object) {
+          if ('data' in error && 'status' in error) {
+            if (error.status === 403) {
+              const detail = (error.data as any)?.detail || 'Unknown Error';
+              dispatch(
+                addToast({
+                  title: t('queue.graphFailedToQueue'),
+                  status: 'error',
+                  description: detail,
+                  duration: 15000,
+                })
+              );
+              return;
+            }
+          }
+        }
+
         dispatch(
           addToast({
             title: t('queue.graphFailedToQueue'),
diff --git a/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts b/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts
index f0745eae2b..9b5dec68f3 100644
--- a/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts
+++ b/invokeai/frontend/web/src/features/controlNet/store/controlNetSlice.ts
@@ -410,6 +410,9 @@ export const controlNetSlice = createSlice({
       state.isIPAdapterEnabled = false;
       state.ipAdapterInfo = { ...initialIPAdapterState };
     },
+    clearPendingControlImages: (state) => {
+      state.pendingControlImages = [];
+    },
   },
   extraReducers: (builder) => {
     builder.addCase(controlNetImageProcessed, (state, action) => {
@@ -474,6 +477,7 @@ export const {
   ipAdapterBeginStepPctChanged,
   ipAdapterEndStepPctChanged,
   ipAdapterStateReset,
+  clearPendingControlImages,
 } = controlNetSlice.actions;
 
 export default controlNetSlice.reducer;