From 8982543312dc70c8fd09d4e84b582dafb304ac45 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 30 Aug 2023 11:33:38 +1000
Subject: [PATCH] fix(ui): fix control image save button logic

---
 .../components/ControlNetImagePreview.tsx     | 26 ++++++++++++++-----
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx b/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx
index 6290d66052..cbfde49c76 100644
--- a/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx
+++ b/invokeai/frontend/web/src/features/controlNet/components/ControlNetImagePreview.tsx
@@ -18,6 +18,7 @@ import {
   useAddImageToBoardMutation,
   useChangeImageIsIntermediateMutation,
   useGetImageDTOQuery,
+  useRemoveImageFromBoardMutation,
 } from 'services/api/endpoints/images';
 import { PostUploadAction } from 'services/api/types';
 import IAIDndImageIcon from '../../../common/components/IAIDndImageIcon';
@@ -71,23 +72,36 @@ const ControlNetImagePreview = ({ isSmall, controlNet }: Props) => {
 
   const [changeIsIntermediate] = useChangeImageIsIntermediateMutation();
   const [addToBoard] = useAddImageToBoardMutation();
-
+  const [removeFromBoard] = useRemoveImageFromBoardMutation();
   const handleResetControlImage = useCallback(() => {
     dispatch(controlNetImageChanged({ controlNetId, controlImage: null }));
   }, [controlNetId, dispatch]);
 
-  const handleSaveControlImage = useCallback(() => {
+  const handleSaveControlImage = useCallback(async () => {
     if (!processedControlImage) {
       return;
     }
 
-    changeIsIntermediate({
+    await changeIsIntermediate({
       imageDTO: processedControlImage,
       is_intermediate: false,
-    });
+    }).unwrap();
 
-    addToBoard({ imageDTO: processedControlImage, board_id: autoAddBoardId });
-  }, [processedControlImage, autoAddBoardId, changeIsIntermediate, addToBoard]);
+    if (autoAddBoardId !== 'none') {
+      addToBoard({
+        imageDTO: processedControlImage,
+        board_id: autoAddBoardId,
+      });
+    } else {
+      removeFromBoard({ imageDTO: processedControlImage });
+    }
+  }, [
+    processedControlImage,
+    changeIsIntermediate,
+    autoAddBoardId,
+    addToBoard,
+    removeFromBoard,
+  ]);
 
   const handleSetControlImageToDimensions = useCallback(() => {
     if (!processedControlImage) {