From cb19c1c3709d54542832616866d6d1c5fe66d2a9 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 28 Jun 2024 18:43:28 +1000
Subject: [PATCH] feat(ui): staging area image visibility toggle

---
 .../StagingArea/StagingAreaToolbar.tsx        | 36 ++++++++++++++++---
 .../controlLayers/konva/nodeManager.ts        |  3 +-
 .../controlLayers/konva/renderers/renderer.ts |  6 ++++
 .../konva/renderers/stagingArea.ts            |  5 +--
 .../controlLayers/store/canvasV2Slice.ts      |  1 +
 5 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbar.tsx b/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbar.tsx
index e7faac86e8..a64a6576ed 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbar.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbar.tsx
@@ -1,6 +1,8 @@
 import { Button, ButtonGroup, IconButton } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import {
+  $shouldShowStagedImage,
   stagingAreaImageAccepted,
   stagingAreaImageDiscarded,
   stagingAreaNextImageSelected,
@@ -11,7 +13,16 @@ import type { CanvasV2State } from 'features/controlLayers/store/types';
 import { memo, useCallback, useMemo } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 import { useTranslation } from 'react-i18next';
-import { PiArrowLeftBold, PiArrowRightBold, PiCheckBold, PiTrashSimpleBold, PiXBold } from 'react-icons/pi';
+import {
+  PiArrowLeftBold,
+  PiArrowRightBold,
+  PiCheckBold,
+  PiEyeBold,
+  PiEyeSlashBold,
+  PiFloppyDiskBold,
+  PiTrashSimpleBold,
+  PiXBold,
+} from 'react-icons/pi';
 
 export const StagingAreaToolbar = memo(() => {
   const stagingArea = useAppSelector((s) => s.canvasV2.stagingArea);
@@ -31,6 +42,7 @@ type Props = {
 
 export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
   const dispatch = useAppDispatch();
+  const shouldShowStagedImage = useStore($shouldShowStagedImage);
   const images = useMemo(() => stagingArea.images, [stagingArea]);
   const imageDTO = useMemo(() => {
     if (stagingArea.selectedImageIndex === null) {
@@ -74,6 +86,12 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
     dispatch(stagingAreaReset());
   }, [dispatch, stagingArea]);
 
+  const onToggleShouldShowStagedImage = useCallback(() => {
+    $shouldShowStagedImage.set(!shouldShowStagedImage);
+  }, [shouldShowStagedImage]);
+
+  const onSaveStagingImage = useCallback(() => {}, []);
+
   useHotkeys(['left'], onPrev, {
     preventDefault: true,
   });
@@ -95,6 +113,7 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
           icon={<PiArrowLeftBold />}
           onClick={onPrev}
           colorScheme="invokeBlue"
+          isDisabled={images.length <= 1 || !shouldShowStagedImage}
         />
         <Button
           colorScheme="base"
@@ -107,6 +126,7 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
           icon={<PiArrowRightBold />}
           onClick={onNext}
           colorScheme="invokeBlue"
+          isDisabled={images.length <= 1 || !shouldShowStagedImage}
         />
       </ButtonGroup>
       <ButtonGroup borderRadius="base" shadow="dark-lg">
@@ -117,14 +137,22 @@ export const StagingAreaToolbarContent = memo(({ stagingArea }: Props) => {
           onClick={onAccept}
           colorScheme="invokeBlue"
         />
-        {/* <IconButton
+        <IconButton
+          tooltip={shouldShowStagedImage ? t('unifiedCanvas.showResultsOn') : t('unifiedCanvas.showResultsOff')}
+          aria-label={shouldShowStagedImage ? t('unifiedCanvas.showResultsOn') : t('unifiedCanvas.showResultsOff')}
+          data-alert={!shouldShowStagedImage}
+          icon={shouldShowStagedImage ? <PiEyeBold /> : <PiEyeSlashBold />}
+          onClick={onToggleShouldShowStagedImage}
+          colorScheme="invokeBlue"
+        />
+        <IconButton
           tooltip={`${t('unifiedCanvas.saveToGallery')} (Shift+S)`}
           aria-label={t('unifiedCanvas.saveToGallery')}
           isDisabled={!imageDTO || !imageDTO.is_intermediate}
           icon={<PiFloppyDiskBold />}
-          onClick={handleSaveToGallery}
+          onClick={onSaveStagingImage}
           colorScheme="invokeBlue"
-        /> */}
+        />
         <IconButton
           tooltip={`${t('unifiedCanvas.discardCurrent')}`}
           aria-label={t('unifiedCanvas.discardCurrent')}
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/nodeManager.ts b/invokeai/frontend/web/src/features/controlLayers/konva/nodeManager.ts
index e82f3144ac..6dd0a591bd 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/nodeManager.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/nodeManager.ts
@@ -52,6 +52,7 @@ export type StateApi = {
   getSelectedEntity: () => CanvasEntity | null;
   getSpaceKey: () => boolean;
   setSpaceKey: (val: boolean) => void;
+  getShouldShowStagedImage: () => boolean;
   getBbox: () => CanvasV2State['bbox'];
   getSettings: () => CanvasV2State['settings'];
   onBrushLineAdded: (arg: BrushLineAddedArg, entityType: CanvasEntity['type']) => void;
@@ -275,7 +276,7 @@ export class KonvaNodeManager {
   }
 
   renderStagingArea() {
-    this.preview.stagingArea.render(this.stateApi.getStagingAreaState());
+    this.preview.stagingArea.render(this.stateApi.getStagingAreaState(), this.stateApi.getShouldShowStagedImage());
   }
 
   fitDocument() {
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts
index 7b23fff52d..4eb5cb1c7a 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts
@@ -7,6 +7,7 @@ import { setStageEventHandlers } from 'features/controlLayers/konva/events';
 import { KonvaNodeManager, setNodeManager } from 'features/controlLayers/konva/nodeManager';
 import { updateBboxes } from 'features/controlLayers/konva/renderers/entityBbox';
 import {
+  $shouldShowStagedImage,
   $stageAttrs,
   bboxChanged,
   brushWidthChanged,
@@ -303,6 +304,7 @@ export const initializeRenderer = (
     getMaskOpacity,
     getInpaintMaskState,
     getStagingAreaState,
+    getShouldShowStagedImage: $shouldShowStagedImage.get,
 
     // Read-write state
     setTool,
@@ -443,6 +445,9 @@ export const initializeRenderer = (
 
   const unsubscribeRenderer = subscribe(renderCanvas);
 
+  // When we this flag, we need to render the staging area
+  $shouldShowStagedImage.subscribe(manager.renderStagingArea.bind(manager));
+
   logIfDebugging('First render of konva stage');
   // On first render, the document should be fit to the stage.
   manager.renderDocumentSizeOverlay();
@@ -454,6 +459,7 @@ export const initializeRenderer = (
     logIfDebugging('Cleaning up konva renderer');
     unsubscribeRenderer();
     cleanupListeners();
+    $shouldShowStagedImage.off();
     resizeObserver.disconnect();
   };
 };
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts
index 2b2deed489..93c8cdb7d7 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts
@@ -3,7 +3,6 @@ import type { CanvasV2State } from 'features/controlLayers/store/types';
 import Konva from 'konva';
 import { assert } from 'tsafe';
 
-
 export class CanvasStagingArea {
   group: Konva.Group;
   image: KonvaImage | null;
@@ -13,7 +12,7 @@ export class CanvasStagingArea {
     this.image = null;
   }
 
-  async render(stagingArea: CanvasV2State['stagingArea']) {
+  async render(stagingArea: CanvasV2State['stagingArea'], shouldShowStagedImage: boolean) {
     if (!stagingArea || stagingArea.selectedImageIndex === null) {
       if (this.image) {
         this.image.destroy();
@@ -29,6 +28,7 @@ export class CanvasStagingArea {
         if (!this.image.isLoading && !this.image.isError && this.image.imageName !== imageDTO.image_name) {
           await this.image.updateImageSource(imageDTO.image_name);
         }
+        this.image.konvaImageGroup.visible(shouldShowStagedImage);
       } else {
         const { image_name, width, height } = imageDTO;
         this.image = new KonvaImage({
@@ -49,6 +49,7 @@ export class CanvasStagingArea {
         });
         this.group.add(this.image.konvaImageGroup);
         await this.image.updateImageSource(imageDTO.image_name);
+        this.image.konvaImageGroup.visible(shouldShowStagedImage);
       }
     }
   }
diff --git a/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts b/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts
index f0ac34622f..ac4b385a25 100644
--- a/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts
@@ -357,6 +357,7 @@ export const $stageAttrs = atom<StageAttrs>({
   height: 0,
   scale: 0,
 });
+export const $shouldShowStagedImage = atom(true);
 
 export const canvasV2PersistConfig: PersistConfig<CanvasV2State> = {
   name: canvasV2Slice.name,