From 6ae46d7c8b1f79764a33018173cd3efe55bf3e3a Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 7 Aug 2024 22:50:35 +1000
Subject: [PATCH] feat(ui): add manual scale controls

---
 invokeai/frontend/web/public/locales/en.json  |   2 +
 .../controlLayers/components/BrushWidth.tsx   |   2 +-
 .../controlLayers/components/CanvasScale.tsx  | 114 ++++++++++++++++++
 .../components/ControlLayersToolbar.tsx       |   2 +
 .../controlLayers/components/EraserWidth.tsx  |   2 +-
 .../controlLayers/konva/CanvasManager.ts      |  78 ++++++++++++
 .../features/controlLayers/konva/events.ts    |  27 +----
 7 files changed, 201 insertions(+), 26 deletions(-)
 create mode 100644 invokeai/frontend/web/src/features/controlLayers/components/CanvasScale.tsx

diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json
index 33c42d1a40..34a5d6ba9d 100644
--- a/invokeai/frontend/web/public/locales/en.json
+++ b/invokeai/frontend/web/public/locales/en.json
@@ -1653,6 +1653,8 @@
         "moveForward": "Move Forward",
         "moveBackward": "Move Backward",
         "brushSize": "Brush Size",
+        "width": "Width",
+        "zoom": "Zoom",
         "controlLayers": "Control Layers",
         "globalMaskOpacity": "Global Mask Opacity",
         "autoNegative": "Auto Negative",
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/BrushWidth.tsx b/invokeai/frontend/web/src/features/controlLayers/components/BrushWidth.tsx
index fdb7b33d84..fa72e54259 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/BrushWidth.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/BrushWidth.tsx
@@ -29,7 +29,7 @@ export const BrushWidth = memo(() => {
   );
   return (
     <FormControl w="min-content" gap={2}>
-      <FormLabel m={0}>{t('controlLayers.brushWidth')}</FormLabel>
+      <FormLabel m={0}>{t('controlLayers.width')}</FormLabel>
       <Popover isLazy>
         <PopoverTrigger>
           <CompositeNumberInput
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/CanvasScale.tsx b/invokeai/frontend/web/src/features/controlLayers/components/CanvasScale.tsx
new file mode 100644
index 0000000000..cfb7a9ff29
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/components/CanvasScale.tsx
@@ -0,0 +1,114 @@
+import {
+  CompositeSlider,
+  FormControl,
+  FormLabel,
+  IconButton,
+  NumberInput,
+  NumberInputField,
+  Popover,
+  PopoverArrow,
+  PopoverBody,
+  PopoverContent,
+  PopoverTrigger,
+} from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
+import { $canvasManager } from 'features/controlLayers/konva/CanvasManager';
+import { MAX_CANVAS_SCALE, MIN_CANVAS_SCALE } from 'features/controlLayers/konva/constants';
+import { $stageAttrs } from 'features/controlLayers/store/canvasV2Slice';
+import { clamp, round } from 'lodash-es';
+import type { KeyboardEvent } from 'react';
+import { memo, useCallback, useEffect, useState } from 'react';
+import { useTranslation } from 'react-i18next';
+import { PiArrowCounterClockwiseBold } from 'react-icons/pi';
+
+const formatPct = (v: number | string) => (isNaN(Number(v)) ? '' : `${round(Number(v), 2).toLocaleString()}%`);
+
+export const CanvasScale = memo(() => {
+  const { t } = useTranslation();
+  const canvasManager = useStore($canvasManager);
+  const stageAttrs = useStore($stageAttrs);
+  const [localScale, setLocalScale] = useState(stageAttrs.scale * 100);
+
+  const onChange = useCallback(
+    (scale: number) => {
+      if (!canvasManager) {
+        return;
+      }
+      canvasManager.setStageScale(scale / 100);
+    },
+    [canvasManager]
+  );
+
+  const onReset = useCallback(() => {
+    if (!canvasManager) {
+      return;
+    }
+
+    canvasManager.setStageScale(1);
+  }, [canvasManager]);
+
+  const onBlur = useCallback(() => {
+    if (!canvasManager) {
+      return;
+    }
+    if (isNaN(Number(localScale))) {
+      return;
+    }
+    canvasManager.setStageScale(clamp(localScale / 100, MIN_CANVAS_SCALE, MAX_CANVAS_SCALE));
+  }, [canvasManager, localScale]);
+
+  const onChangeNumberInput = useCallback((valueAsString: string, valueAsNumber: number) => {
+    setLocalScale(valueAsNumber);
+  }, []);
+
+  const onKeyDown = useCallback(
+    (e: KeyboardEvent<HTMLInputElement>) => {
+      if (e.key === 'Enter') {
+        onBlur();
+      }
+    },
+    [onBlur]
+  );
+
+  useEffect(() => {
+    setLocalScale(stageAttrs.scale * 100);
+  }, [stageAttrs.scale]);
+
+  return (
+    <FormControl w="min-content" gap={2}>
+      <FormLabel m={0}>{t('controlLayers.zoom')}</FormLabel>
+      <Popover isLazy trigger="hover" openDelay={300}>
+        <PopoverTrigger>
+          <NumberInput
+            min={MIN_CANVAS_SCALE * 100}
+            max={MAX_CANVAS_SCALE * 100}
+            value={localScale}
+            onChange={onChangeNumberInput}
+            onBlur={onBlur}
+            w="64px"
+            format={formatPct}
+            defaultValue={100}
+            onKeyDown={onKeyDown}
+          >
+            <NumberInputField textAlign="center" paddingInlineEnd={3} />
+          </NumberInput>
+        </PopoverTrigger>
+        <PopoverContent w={200} py={2} px={4}>
+          <PopoverArrow />
+          <PopoverBody>
+            <CompositeSlider
+              min={MIN_CANVAS_SCALE * 100}
+              max={MAX_CANVAS_SCALE * 100}
+              value={stageAttrs.scale * 100}
+              onChange={onChange}
+              defaultValue={100}
+            />
+          </PopoverBody>
+        </PopoverContent>
+      </Popover>
+      <IconButton aria-label="reset" onClick={onReset} icon={<PiArrowCounterClockwiseBold />} variant="link" />
+    </FormControl>
+  );
+});
+
+CanvasScale.displayName = 'CanvasScale';
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/ControlLayersToolbar.tsx b/invokeai/frontend/web/src/features/controlLayers/components/ControlLayersToolbar.tsx
index 1b8d52585e..e8d31eecf9 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/ControlLayersToolbar.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/ControlLayersToolbar.tsx
@@ -4,6 +4,7 @@ import { Flex, Switch } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { useAppSelector } from 'app/store/storeHooks';
 import { BrushWidth } from 'features/controlLayers/components/BrushWidth';
+import { CanvasScale } from 'features/controlLayers/components/CanvasScale';
 import ControlLayersSettingsPopover from 'features/controlLayers/components/ControlLayersSettingsPopover';
 import { EraserWidth } from 'features/controlLayers/components/EraserWidth';
 import { FillColorPicker } from 'features/controlLayers/components/FillColorPicker';
@@ -52,6 +53,7 @@ export const ControlLayersToolbar = memo(() => {
         {tool === 'brush' && <BrushWidth />}
         {tool === 'eraser' && <EraserWidth />}
       </Flex>
+      <CanvasScale />
       <Button onClick={bbox}>bbox</Button>
       <Switch onChange={onChangeDebugging}>debug</Switch>
       <Flex flex={1} justifyContent="center">
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/EraserWidth.tsx b/invokeai/frontend/web/src/features/controlLayers/components/EraserWidth.tsx
index 0903763f2d..deff803bb1 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/EraserWidth.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/EraserWidth.tsx
@@ -29,7 +29,7 @@ export const EraserWidth = memo(() => {
   );
   return (
     <FormControl w="min-content" gap={2}>
-      <FormLabel m={0}>{t('controlLayers.eraserWidth')}</FormLabel>
+      <FormLabel m={0}>{t('controlLayers.width')}</FormLabel>
       <Popover isLazy>
         <PopoverTrigger>
           <CompositeNumberInput
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasManager.ts b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasManager.ts
index 722c6ce470..cc779cdee6 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasManager.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasManager.ts
@@ -10,6 +10,7 @@ import { CanvasObjectRenderer } from 'features/controlLayers/konva/CanvasObjectR
 import { CanvasProgressPreview } from 'features/controlLayers/konva/CanvasProgressPreview';
 import type { CanvasRectRenderer } from 'features/controlLayers/konva/CanvasRect';
 import type { CanvasTransformer } from 'features/controlLayers/konva/CanvasTransformer';
+import { MAX_CANVAS_SCALE, MIN_CANVAS_SCALE } from 'features/controlLayers/konva/constants';
 import {
   getImageDataTransparency,
   getPrefixedId,
@@ -27,6 +28,7 @@ import type {
   CanvasRegionalGuidanceState,
   CanvasV2State,
   Coordinate,
+  Dimensions,
   GenerationMode,
   GetLoggingContext,
   Rect,
@@ -35,6 +37,7 @@ import type {
 import { RGBA_RED } from 'features/controlLayers/store/types';
 import { isValidLayer } from 'features/nodes/util/graph/generation/addLayers';
 import type Konva from 'konva';
+import { clamp } from 'lodash-es';
 import { atom } from 'nanostores';
 import type { Logger } from 'roarr';
 import {
@@ -623,15 +626,90 @@ export class CanvasManager {
     };
   };
 
+  /**
+   * Gets the center of the stage in either absolute or relative coordinates
+   * @param absolute Whether to return the center in absolute coordinates
+   */
+  getStageCenter(absolute = false): Coordinate {
+    const scale = this.getStageScale();
+    const { x, y } = this.getStagePosition();
+    const { width, height } = this.getStageSize();
+
+    const center = {
+      x: (width / 2 - x) / scale,
+      y: (height / 2 - y) / scale,
+    };
+
+    if (!absolute) {
+      return center;
+    }
+
+    return this.stage.getAbsoluteTransform().point(center);
+  }
+
+  /**
+   * Sets the scale of the stage. If center is provided, the stage will zoom in/out on that point.
+   * @param scale The new scale to set
+   * @param center The center of the stage to zoom in/out on
+   */
+  setStageScale(scale: number, center: Coordinate = this.getStageCenter(true)) {
+    const newScale = clamp(Math.round(scale * 100) / 100, MIN_CANVAS_SCALE, MAX_CANVAS_SCALE);
+
+    const { x, y } = this.getStagePosition();
+    const oldScale = this.getStageScale();
+
+    const deltaX = (center.x - x) / oldScale;
+    const deltaY = (center.y - y) / oldScale;
+
+    const newX = center.x - deltaX * newScale;
+    const newY = center.y - deltaY * newScale;
+
+    this.stage.setAttrs({
+      x: newX,
+      y: newY,
+      scaleX: newScale,
+      scaleY: newScale,
+    });
+
+    this.stateApi.$stageAttrs.set({
+      x: Math.floor(this.stage.x()),
+      y: Math.floor(this.stage.y()),
+      width: this.stage.width(),
+      height: this.stage.height(),
+      scale: this.stage.scaleX(),
+    });
+    this.background.render();
+    this.preview.tool.render();
+  }
+
+  /**
+   * Gets the scale of the stage. The stage is always scaled uniformly in x and y.
+   */
   getStageScale(): number {
     // The stage is never scaled differently in x and y
     return this.stage.scaleX();
   }
 
+  /**
+   * Gets the position of the stage.
+   */
   getStagePosition(): Coordinate {
     return this.stage.position();
   }
 
+  /**
+   * Gets the size of the stage.
+   */
+  getStageSize(): Dimensions {
+    return this.stage.size();
+  }
+
+  /**
+   * Scales a number of pixels by the current stage scale. For example, if the stage is scaled by 5, then 10 pixels
+   * would be scaled to 10px / 5 = 2 pixels.
+   * @param pixels The number of pixels to scale
+   * @returns The number of pixels scaled by the current stage scale
+   */
   getScaledPixels(pixels: number): number {
     return pixels / this.getStageScale();
   }
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/events.ts b/invokeai/frontend/web/src/features/controlLayers/konva/events.ts
index 3fd6594ddf..054b433646 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/events.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/events.ts
@@ -18,7 +18,7 @@ import type Konva from 'konva';
 import type { KonvaEventObject } from 'konva/lib/Node';
 import { clamp } from 'lodash-es';
 
-import { BRUSH_SPACING_TARGET_SCALE, CANVAS_SCALE_BY, MAX_CANVAS_SCALE, MIN_CANVAS_SCALE } from './constants';
+import { BRUSH_SPACING_TARGET_SCALE, CANVAS_SCALE_BY } from './constants';
 
 /**
  * Updates the last cursor position atom with the current cursor position, returning the new position or `null` if the
@@ -474,31 +474,10 @@ export const setStageEventHandlers = (manager: CanvasManager): (() => void) => {
       // We need the absolute cursor position - not the scaled position
       const cursorPos = stage.getPointerPosition();
       if (cursorPos) {
-        // Stage's x and y scale are always the same
-        const stageScale = stage.scaleX();
         // When wheeling on trackpad, e.evt.ctrlKey is true - in that case, let's reverse the direction
         const delta = e.evt.ctrlKey ? -e.evt.deltaY : e.evt.deltaY;
-        const mousePointTo = {
-          x: (cursorPos.x - stage.x()) / stageScale,
-          y: (cursorPos.y - stage.y()) / stageScale,
-        };
-        const newScale = clamp(stageScale * CANVAS_SCALE_BY ** delta, MIN_CANVAS_SCALE, MAX_CANVAS_SCALE);
-        const newPos = {
-          x: cursorPos.x - mousePointTo.x * newScale,
-          y: cursorPos.y - mousePointTo.y * newScale,
-        };
-
-        stage.scaleX(newScale);
-        stage.scaleY(newScale);
-        stage.position(newPos);
-        $stageAttrs.set({
-          x: newPos.x,
-          y: newPos.y,
-          width: stage.width(),
-          height: stage.height(),
-          scale: newScale,
-        });
-        manager.background.render();
+        const scale = manager.getStageScale() * CANVAS_SCALE_BY ** delta;
+        manager.setStageScale(scale, cursorPos);
       }
     }
     manager.preview.tool.render();