From 6c1f666242deda63eb73cf17e79d88e675937f74 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Sat, 18 Mar 2023 21:27:57 +1100
Subject: [PATCH] feat(ui): nodes cancel

---
 invokeai/frontend/web/src/app/NodeAPITest.tsx | 50 +++++++++++---
 .../services/api/services/SessionsService.ts  | 26 ++++++++
 .../web/src/services/thunks/session.ts        | 65 ++++++++-----------
 3 files changed, 95 insertions(+), 46 deletions(-)

diff --git a/invokeai/frontend/web/src/app/NodeAPITest.tsx b/invokeai/frontend/web/src/app/NodeAPITest.tsx
index f90326349a..6c891a78f4 100644
--- a/invokeai/frontend/web/src/app/NodeAPITest.tsx
+++ b/invokeai/frontend/web/src/app/NodeAPITest.tsx
@@ -9,7 +9,11 @@ import {
 import { useEffect } from 'react';
 import { STATUS, ProgressImage } from 'services/apiSliceTypes';
 import { getImage } from 'services/thunks/image';
-import { createSession, invokeSession } from 'services/thunks/session';
+import {
+  cancelProcessing,
+  createSession,
+  invokeSession,
+} from 'services/thunks/session';
 import { io } from 'socket.io-client';
 import { useAppDispatch, useAppSelector } from './storeHooks';
 import { RootState } from './store';
@@ -42,11 +46,24 @@ const NodeAPITest = () => {
               id: 'a',
               type: 'txt2img',
               prompt: 'pizza',
-              steps: 10,
+              steps: 50,
+              seed: 123,
             },
             b: {
               id: 'b',
-              type: 'show_image',
+              type: 'img2img',
+              prompt: 'dog',
+              steps: 50,
+              seed: 123,
+              strength: 0.9,
+            },
+            c: {
+              id: 'c',
+              type: 'img2img',
+              prompt: 'cat',
+              steps: 50,
+              seed: 123,
+              strength: 0.9,
             },
           },
           edges: [
@@ -54,6 +71,10 @@ const NodeAPITest = () => {
               source: { node_id: 'a', field: 'image' },
               destination: { node_id: 'b', field: 'image' },
             },
+            {
+              source: { node_id: 'b', field: 'image' },
+              destination: { node_id: 'c', field: 'image' },
+            },
           ],
         },
       })
@@ -64,6 +85,10 @@ const NodeAPITest = () => {
     dispatch(invokeSession());
   };
 
+  const handleCancelProcessing = () => {
+    dispatch(cancelProcessing());
+  };
+
   useEffect(() => {
     if (!sessionId) {
       return;
@@ -97,7 +122,7 @@ const NodeAPITest = () => {
       // in the future we will want to continue building the graph and executing etc
       console.log('invocation_complete', data);
       dispatch(setProgress(null));
-      dispatch(setSessionId(null));
+      // dispatch(setSessionId(null));
       dispatch(setStatus(STATUS.idle));
 
       // think this gets a blob...
@@ -107,12 +132,12 @@ const NodeAPITest = () => {
       //     imageName: data.result.image.image_name,
       //   })
       // );
-      socket.emit('unsubscribe', { session: sessionId });
-      console.log('unsubscribe', { session: sessionId });
     });
 
     // not sure when we get this?
     socket.on('session_complete', (data) => {
+      // socket.emit('unsubscribe', { session: sessionId });
+      // console.log('unsubscribe', { session: sessionId });
       // console.log('session_complete', data);
     });
 
@@ -135,17 +160,24 @@ const NodeAPITest = () => {
       }}
     >
       <Text>Session: {sessionId ? sessionId : '...'}</Text>
+      <IAIButton
+        onClick={handleCancelProcessing}
+        // isDisabled={!sessionId}
+        colorScheme="error"
+      >
+        Cancel Processing
+      </IAIButton>
       <IAIButton
         onClick={handleCreateSession}
-        isDisabled={status === STATUS.busy || Boolean(sessionId)}
+        // isDisabled={status === STATUS.busy || Boolean(sessionId)}
         colorScheme="accent"
       >
         Create Session
       </IAIButton>
       <IAIButton
         onClick={handleInvokeSession}
-        isDisabled={status === STATUS.busy}
-        isLoading={status === STATUS.busy}
+        // isDisabled={status === STATUS.busy}
+        // isLoading={status === STATUS.busy}
         loadingText={`Invoking ${
           progress === null ? '...' : `${Math.round(progress * 100)}%`
         }`}
diff --git a/invokeai/frontend/web/src/services/api/services/SessionsService.ts b/invokeai/frontend/web/src/services/api/services/SessionsService.ts
index 6c66d4fbbd..a2a6263fe2 100644
--- a/invokeai/frontend/web/src/services/api/services/SessionsService.ts
+++ b/invokeai/frontend/web/src/services/api/services/SessionsService.ts
@@ -341,4 +341,30 @@ export class SessionsService {
     });
   }
 
+  /**
+   * Cancel Session Invoke
+   * Invokes a session
+   * @returns any Successful Response
+   * @throws ApiError
+   */
+  public static cancelSessionInvoke({
+    sessionId,
+  }: {
+    /**
+     * The id of the session to cancel
+     */
+    sessionId: string,
+  }): CancelablePromise<any> {
+    return __request(OpenAPI, {
+      method: 'DELETE',
+      url: '/api/v1/sessions/{session_id}/invoke',
+      path: {
+        'session_id': sessionId,
+      },
+      errors: {
+        422: `Validation Error`,
+      },
+    });
+  }
+
 }
diff --git a/invokeai/frontend/web/src/services/thunks/session.ts b/invokeai/frontend/web/src/services/thunks/session.ts
index 3af130e840..67cdf967af 100644
--- a/invokeai/frontend/web/src/services/thunks/session.ts
+++ b/invokeai/frontend/web/src/services/thunks/session.ts
@@ -14,26 +14,6 @@ export const createSession = createAppAsyncThunk(
   async (arg: CreateSessionArg, { getState, dispatch, ...moreThunkStuff }) => {
     const response = await SessionsService.createSession(arg);
     return response;
-  },
-  {
-    // if this returns false, the api call is skipped
-    // we can guard in many places, and maybe this isn't right for us,
-    // but just trying it here
-    condition: (arg, { getState }) => {
-      const {
-        api: { status, sessionId },
-      } = getState();
-
-      // don't create session if we are processing already
-      if (status === STATUS.busy) {
-        return false;
-      }
-
-      // don't create session if we have a sessionId
-      if (sessionId) {
-        return false;
-      }
-    },
   }
 );
 
@@ -61,22 +41,33 @@ export const invokeSession = createAppAsyncThunk(
     });
 
     return response;
-  },
-  {
-    condition: (arg, { getState }) => {
-      const {
-        api: { status, sessionId },
-      } = getState();
-
-      // don't invoke if we are processing already
-      if (status === STATUS.busy) {
-        return false;
-      }
-
-      // don't invoke if we don't have a sessionId
-      if (!sessionId) {
-        return false;
-      }
-    },
+  }
+);
+/**
+ * invokeSession
+ */
+
+export const cancelProcessing = createAppAsyncThunk(
+  'api/cancelProcessing',
+  async (_arg, { getState }) => {
+    console.log('before canceling');
+    const {
+      api: { sessionId },
+    } = getState();
+
+    // i'd really like for the typing on the condition callback below to tell this
+    // function here that sessionId will never be empty, but guess we do not get
+    // that luxury
+    if (!sessionId) {
+      return;
+    }
+
+    console.log('canceling');
+
+    const response = await SessionsService.cancelSessionInvoke({
+      sessionId,
+    });
+
+    return response;
   }
 );