From f5039d800dfd800a85fbc9538f2817423183aaa0 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sat, 4 Feb 2023 11:56:14 -0600
Subject: [PATCH] feat(gui): add cancel button to loading card

---
 gui/src/client.ts                  |  4 +-
 gui/src/components/LoadingCard.tsx | 61 ++++++++++++++++++------------
 gui/src/components/tab/Img2Img.tsx |  2 +-
 gui/src/components/tab/Inpaint.tsx |  2 +-
 gui/src/components/tab/Txt2Img.tsx |  2 +-
 gui/src/components/tab/Upscale.tsx |  2 +-
 gui/src/state.ts                   | 50 +++++++++++++++---------
 7 files changed, 77 insertions(+), 46 deletions(-)

diff --git a/gui/src/client.ts b/gui/src/client.ts
index b22c0de5..929058e2 100644
--- a/gui/src/client.ts
+++ b/gui/src/client.ts
@@ -503,7 +503,9 @@ export function makeClient(root: string, f = fetch): ApiClient {
       const path = makeApiUrl(root, 'cancel');
       path.searchParams.append('output', params.output.key);
 
-      const res = await f(path);
+      const res = await f(path, {
+        method: 'PUT',
+      });
       return res.status === STATUS_SUCCESS;
     },
   };
diff --git a/gui/src/components/LoadingCard.tsx b/gui/src/components/LoadingCard.tsx
index c04a04c5..f05dcd0d 100644
--- a/gui/src/components/LoadingCard.tsx
+++ b/gui/src/components/LoadingCard.tsx
@@ -1,7 +1,7 @@
 import { doesExist, mustExist } from '@apextoaster/js-utils';
-import { Button, Card, CardContent, CircularProgress } from '@mui/material';
+import { Box, Button, Card, CardContent, CircularProgress, Typography } from '@mui/material';
 import * as React from 'react';
-import { useContext } from 'react';
+import { useContext, useEffect } from 'react';
 import { useMutation, useQuery } from 'react-query';
 import { useStore } from 'zustand';
 
@@ -9,6 +9,8 @@ import { ImageResponse } from '../client.js';
 import { POLL_TIME } from '../config.js';
 import { ClientContext, ConfigContext, StateContext } from '../state.js';
 
+const LOADING_PERCENT = 100;
+
 export interface LoadingCardProps {
   loading: ImageResponse;
 }
@@ -22,50 +24,61 @@ export function LoadingCard(props: LoadingCardProps) {
   const clearLoading = useStore(state, (s) => s.clearLoading);
   // eslint-disable-next-line @typescript-eslint/unbound-method
   const pushHistory = useStore(state, (s) => s.pushHistory);
+  // eslint-disable-next-line @typescript-eslint/unbound-method
+  const setReady = useStore(state, (s) => s.setReady);
 
-  async function doCancel() {
-    const cancelled = await client.cancel(props.loading);
-    if (cancelled) {
-      clearLoading();
-    }
-  }
-
-  const cancel = useMutation(doCancel);
-  const query = useQuery('ready', () => client.ready(props.loading), {
+  const cancel = useMutation(() => client.cancel(props.loading));
+  const ready = useQuery('ready', () => client.ready(props.loading), {
     // data will always be ready without this, even if the API says its not
     cacheTime: 0,
     refetchInterval: POLL_TIME,
   });
 
-  function progress() {
-    if (doesExist(query.data)) {
-      return Math.ceil(query.data.progress / props.loading.params.steps);
+  function getProgress() {
+    if (doesExist(ready.data)) {
+      return ready.data.progress;
     }
 
     return 0;
   }
 
-  function ready() {
-    return doesExist(query.data) && query.data.ready;
+  function getPercent() {
+    const pct = getProgress() / props.loading.params.steps;
+    return Math.ceil(pct * LOADING_PERCENT);
   }
 
-  React.useEffect(() => {
-    if (query.status === 'success' && query.data.ready) {
-      pushHistory(props.loading);
+  function getReady() {
+    return doesExist(ready.data) && ready.data.ready;
+  }
+
+  useEffect(() => {
+    if (cancel.status === 'success') {
+      clearLoading();
     }
-  }, [query.status, ready()]);
+  }, [cancel.status]);
+
+  useEffect(() => {
+    if (ready.status === 'success') {
+      if (ready.data.ready) {
+        pushHistory(props.loading);
+      } else {
+        setReady(props.loading, ready.data);
+      }
+    }
+  }, [ready.status, getReady(), getProgress()]);
 
   return <Card sx={{ maxWidth: params.width.default }}>
     <CardContent sx={{ height: params.height.default }}>
-      <div style={{
+      <Box sx={{
         display: 'flex',
         alignItems: 'center',
         justifyContent: 'center',
         minHeight: params.height.default,
       }}>
-        <CircularProgress value={progress()} />
-        <Button>Cancel</Button>
-      </div>
+        <CircularProgress variant='determinate' value={getPercent()} />
+        <Typography>{getProgress()} of {props.loading.params.steps}</Typography>
+        <Button onClick={() => cancel.mutate()}>Cancel</Button>
+      </Box>
     </CardContent>
   </Card>;
 }
diff --git a/gui/src/components/tab/Img2Img.tsx b/gui/src/components/tab/Img2Img.tsx
index e4cbd981..0ac4f08e 100644
--- a/gui/src/components/tab/Img2Img.tsx
+++ b/gui/src/components/tab/Img2Img.tsx
@@ -38,7 +38,7 @@ export function Img2Img() {
   // eslint-disable-next-line @typescript-eslint/unbound-method
   const setImg2Img = useStore(state, (s) => s.setImg2Img);
   // eslint-disable-next-line @typescript-eslint/unbound-method
-  const setLoading = useStore(state, (s) => s.setLoading);
+  const setLoading = useStore(state, (s) => s.pushLoading);
 
   return <Box>
     <Stack spacing={2}>
diff --git a/gui/src/components/tab/Inpaint.tsx b/gui/src/components/tab/Inpaint.tsx
index 214caeb6..b9647f1b 100644
--- a/gui/src/components/tab/Inpaint.tsx
+++ b/gui/src/components/tab/Inpaint.tsx
@@ -62,7 +62,7 @@ export function Inpaint() {
   // eslint-disable-next-line @typescript-eslint/unbound-method
   const setInpaint = useStore(state, (s) => s.setInpaint);
   // eslint-disable-next-line @typescript-eslint/unbound-method
-  const setLoading = useStore(state, (s) => s.setLoading);
+  const setLoading = useStore(state, (s) => s.pushLoading);
 
   const query = useQueryClient();
   const upload = useMutation(uploadSource, {
diff --git a/gui/src/components/tab/Txt2Img.tsx b/gui/src/components/tab/Txt2Img.tsx
index 53424bdc..1b1d4406 100644
--- a/gui/src/components/tab/Txt2Img.tsx
+++ b/gui/src/components/tab/Txt2Img.tsx
@@ -32,7 +32,7 @@ export function Txt2Img() {
   // eslint-disable-next-line @typescript-eslint/unbound-method
   const setTxt2Img = useStore(state, (s) => s.setTxt2Img);
   // eslint-disable-next-line @typescript-eslint/unbound-method
-  const setLoading = useStore(state, (s) => s.setLoading);
+  const setLoading = useStore(state, (s) => s.pushLoading);
 
   return <Box>
     <Stack spacing={2}>
diff --git a/gui/src/components/tab/Upscale.tsx b/gui/src/components/tab/Upscale.tsx
index c56d337d..9e39f63b 100644
--- a/gui/src/components/tab/Upscale.tsx
+++ b/gui/src/components/tab/Upscale.tsx
@@ -33,7 +33,7 @@ export function Upscale() {
   // eslint-disable-next-line @typescript-eslint/unbound-method
   const setSource = useStore(state, (s) => s.setUpscaleTab);
   // eslint-disable-next-line @typescript-eslint/unbound-method
-  const setLoading = useStore(state, (s) => s.setLoading);
+  const setLoading = useStore(state, (s) => s.pushLoading);
 
   return <Box>
     <Stack spacing={2}>
diff --git a/gui/src/state.ts b/gui/src/state.ts
index aa1ba1dc..8ed9aef4 100644
--- a/gui/src/state.ts
+++ b/gui/src/state.ts
@@ -49,9 +49,10 @@ interface HistorySlice {
   // TODO: hack until setLoading removes things
   clearLoading(): void;
   pushHistory(image: ImageResponse): void;
+  pushLoading(image: ImageResponse): void;
   removeHistory(image: ImageResponse): void;
   setLimit(limit: number): void;
-  setLoading(image: ImageResponse, ready?: Maybe<ReadyResponse>): void;
+  setReady(image: ImageResponse, ready: ReadyResponse): void;
 }
 
 interface ModelSlice {
@@ -135,7 +136,7 @@ export const StateContext = createContext<Maybe<StoreApi<OnnxState>>>(undefined)
 /**
  * Current state version for zustand persistence.
  */
-export const STATE_VERSION = 4;
+export const STATE_VERSION = 5;
 
 /**
  * Default parameters for the inpaint brush.
@@ -289,21 +290,20 @@ export function createStateSlices(server: ServerParams) {
         loading: [],
       }));
     },
-    setLoading(image, ready) {
-      set((prev) => {
-        const loading = [...prev.loading];
-        const idx = loading.findIndex((it) => it.image.output.key === image.output.key);
-        if (idx >= 0) {
-          loading[idx].ready = ready;
-        } else {
-          loading.push({ image, ready });
-        }
-
-        return {
-          ...prev,
-          loading,
-        };
-      });
+    pushLoading(image) {
+      set((prev) => ({
+        ...prev,
+        loading: [
+          {
+            image,
+            ready: {
+              progress: 0,
+              ready: false,
+            },
+          },
+          ...prev.loading,
+        ],
+      }));
     },
     removeHistory(image) {
       set((prev) => ({
@@ -317,6 +317,22 @@ export function createStateSlices(server: ServerParams) {
         limit,
       }));
     },
+    setReady(image, ready) {
+      set((prev) => {
+        const loading = [...prev.loading];
+        const idx = loading.findIndex((it) => it.image.output.key === image.output.key);
+        if (idx >= 0) {
+          loading[idx].ready = ready;
+        } else {
+          // TODO: error
+        }
+
+        return {
+          ...prev,
+          loading,
+        };
+      });
+    },
   });
 
   const createOutpaintSlice: Slice<OutpaintSlice> = (set) => ({