From 45908dfbd2174eb40f95b6bd99fe9a3597b1ab8b Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Sun, 16 Jun 2024 21:12:58 +1000
Subject: [PATCH] feat(ui): r to center & fit stage on document

---
 .../controlLayers/components/StageComponent.tsx  |  6 +++++-
 .../src/features/controlLayers/konva/events.ts   | 16 ++++++++++++++++
 .../controlLayers/store/canvasV2Slice.ts         |  1 +
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx b/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx
index 62fedc296d..0cbb7fcf45 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx
@@ -9,6 +9,7 @@ import { debouncedRenderers, renderers as normalRenderers } from 'features/contr
 import {
   $bbox,
   $currentFill,
+  $document,
   $isDrawing,
   $isMouseDown,
   $lastAddedPoint,
@@ -72,6 +73,7 @@ const useStageRenderer = (stage: Konva.Stage, container: HTMLDivElement | null,
   const selectedEntityIdentifier = useAppSelector((s) => s.canvasV2.selectedEntityIdentifier);
   const maskOpacity = useAppSelector((s) => s.canvasV2.settings.maskOpacity);
   const bbox = useAppSelector((s) => s.canvasV2.bbox);
+  const document = useAppSelector((s) => s.canvasV2.document);
   const lastCursorPos = useStore($lastCursorPos);
   const lastMouseDownPos = useStore($lastMouseDownPos);
   const isMouseDown = useStore($isMouseDown);
@@ -108,7 +110,8 @@ const useStageRenderer = (stage: Konva.Stage, container: HTMLDivElement | null,
     $selectedEntity.set(selectedEntity);
     $bbox.set({ x: bbox.x, y: bbox.y, width: bbox.width, height: bbox.height });
     $currentFill.set(currentFill);
-  }, [selectedEntity, tool, bbox, currentFill]);
+    $document.set(document);
+  }, [selectedEntity, tool, bbox, currentFill, document]);
 
   const onPosChanged = useCallback(
     (arg: PosChangedArg, entityType: CanvasEntity['type']) => {
@@ -243,6 +246,7 @@ const useStageRenderer = (stage: Konva.Stage, container: HTMLDivElement | null,
       setLastMouseDownPos: $lastMouseDownPos.set,
       getSpaceKey: $spaceKey.get,
       setStageAttrs: $stageAttrs.set,
+      getDocument: $document.get,
       onBrushLineAdded,
       onEraserLineAdded,
       onPointAddedToLine,
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/events.ts b/invokeai/frontend/web/src/features/controlLayers/konva/events.ts
index 36f03b7535..de277d5675 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/events.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/events.ts
@@ -45,6 +45,7 @@ type Arg = {
   setStageAttrs: (attrs: StageAttrs) => void;
   getSelectedEntity: () => CanvasEntity | null;
   getSpaceKey: () => boolean;
+  getDocument: () => CanvasV2State['document'];
   onBrushLineAdded: (arg: BrushLineAddedArg, entityType: CanvasEntity['type']) => void;
   onEraserLineAdded: (arg: EraserLineAddedArg, entityType: CanvasEntity['type']) => void;
   onPointAddedToLine: (arg: PointAddedToLineArg, entityType: CanvasEntity['type']) => void;
@@ -144,6 +145,7 @@ export const setStageEventHandlers = ({
   setStageAttrs,
   getSelectedEntity,
   getSpaceKey,
+  getDocument,
   onBrushLineAdded,
   onEraserLineAdded,
   onPointAddedToLine,
@@ -506,6 +508,20 @@ export const setStageEventHandlers = ({
     } else if (e.key === ' ') {
       setToolBuffer(getToolState().selected);
       setTool('view');
+    } else if (e.key === 'r') {
+      // Fit & center the document on the stage
+      const width = stage.width();
+      const height = stage.height();
+      const document = getDocument();
+      const docWidthWithBuffer = document.width + 20;
+      const docHeightWithBuffer = document.height + 20;
+      const scale = Math.min(Math.min(width / docWidthWithBuffer, height / docHeightWithBuffer), 1);
+      const x = (width - docWidthWithBuffer * scale) / 2;
+      const y = (height - docHeightWithBuffer * scale) / 2;
+      stage.setAttrs({ x, y, width, height, scaleX: scale, scaleY: scale });
+      setStageAttrs({ x, y, width, height, scale });
+      scaleToolPreview(stage, getToolState());
+      renderBackgroundLayer(stage);
     }
   };
   window.addEventListener('keydown', onKeyDown);
diff --git a/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts b/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts
index 4f86bb1fb4..631cab64b3 100644
--- a/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts
@@ -318,6 +318,7 @@ export const $toolState = atom<CanvasV2State['tool']>(deepClone(initialState.too
 export const $currentFill = atom<RgbaColor>(DEFAULT_RGBA_COLOR);
 export const $selectedEntity = atom<CanvasEntity | null>(null);
 export const $bbox = atom<IRect>({ x: 0, y: 0, width: 0, height: 0 });
+export const $document = atom<CanvasV2State['document']>(deepClone(initialState.document));
 
 export const canvasV2PersistConfig: PersistConfig<CanvasV2State> = {
   name: canvasV2Slice.name,