From ae7797f662582b4355e704e62a064f391d19b4c9 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 16 Apr 2024 19:41:37 +1000
Subject: [PATCH] feat(ui): re-implement with imperative konva api (wip)

---
 .../components/RegionalPromptsEditor.tsx      |   4 +-
 .../components/imperative/konvaApiDraft.tsx   | 235 ++++++++++++++++++
 .../components/imperative/mouseEventHooks.ts  | 142 +++++++++++
 .../components/imperative/test.tsx            |  22 ++
 .../regionalPrompts/util/getLayerBlobs.ts     |   2 +-
 5 files changed, 403 insertions(+), 2 deletions(-)
 create mode 100644 invokeai/frontend/web/src/features/regionalPrompts/components/imperative/konvaApiDraft.tsx
 create mode 100644 invokeai/frontend/web/src/features/regionalPrompts/components/imperative/mouseEventHooks.ts
 create mode 100644 invokeai/frontend/web/src/features/regionalPrompts/components/imperative/test.tsx

diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsEditor.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsEditor.tsx
index cf3614dbec..66fcd6aba5 100644
--- a/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsEditor.tsx
+++ b/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsEditor.tsx
@@ -4,6 +4,7 @@ import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
 import { useAppSelector } from 'app/store/storeHooks';
 import { AddLayerButton } from 'features/regionalPrompts/components/AddLayerButton';
 import { BrushSize } from 'features/regionalPrompts/components/BrushSize';
+import { StageComponent } from 'features/regionalPrompts/components/imperative/konvaApiDraft';
 import { LayerListItem } from 'features/regionalPrompts/components/LayerListItem';
 import { PromptLayerOpacity } from 'features/regionalPrompts/components/PromptLayerOpacity';
 import { RegionalPromptsStage } from 'features/regionalPrompts/components/RegionalPromptsStage';
@@ -37,7 +38,8 @@ export const RegionalPromptsEditor = memo(() => {
         ))}
       </Flex>
       <Flex>
-        <RegionalPromptsStage />
+        <StageComponent />
+        {/* <RegionalPromptsStage /> */}
       </Flex>
     </Flex>
   );
diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/imperative/konvaApiDraft.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/imperative/konvaApiDraft.tsx
new file mode 100644
index 0000000000..68e1d7edab
--- /dev/null
+++ b/invokeai/frontend/web/src/features/regionalPrompts/components/imperative/konvaApiDraft.tsx
@@ -0,0 +1,235 @@
+import { chakra } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { rgbColorToString } from 'features/canvas/util/colorToString';
+import getScaledCursorPosition from 'features/canvas/util/getScaledCursorPosition';
+import {
+  $cursorPosition,
+  layerBboxChanged,
+  layerSelected,
+  layerTranslated,
+  REGIONAL_PROMPT_LAYER_NAME,
+  REGIONAL_PROMPT_LAYER_OBJECT_GROUP_NAME,
+} from 'features/regionalPrompts/store/regionalPromptsSlice';
+import { getKonvaLayerBbox } from 'features/regionalPrompts/util/bbox';
+import Konva from 'konva';
+import type { Node, NodeConfig } from 'konva/lib/Node';
+import type { StageConfig } from 'konva/lib/Stage';
+import { atom } from 'nanostores';
+import { useLayoutEffect } from 'react';
+
+import { useMouseDown, useMouseEnter, useMouseLeave, useMouseMove, useMouseUp } from './mouseEventHooks';
+
+export const $stage = atom<Konva.Stage | null>(null);
+
+const initStage = (container: StageConfig['container']) => {
+  const stage = new Konva.Stage({
+    container,
+  });
+  $stage.set(stage);
+
+  const layer = new Konva.Layer();
+  const circle = new Konva.Circle({ id: 'cursor', radius: 5, fill: 'red' });
+  layer.add(circle);
+  stage.add(layer);
+};
+
+type Props = {
+  container: HTMLDivElement | null;
+};
+
+export const selectPromptLayerObjectGroup = (item: Node<NodeConfig>) =>
+  item.name() !== REGIONAL_PROMPT_LAYER_OBJECT_GROUP_NAME;
+
+export const LogicalStage = (props: Props) => {
+  const dispatch = useAppDispatch();
+  const width = useAppSelector((s) => s.generation.width);
+  const height = useAppSelector((s) => s.generation.height);
+  const state = useAppSelector((s) => s.regionalPrompts);
+  const stage = useStore($stage);
+  const onMouseDown = useMouseDown();
+  const onMouseUp = useMouseUp();
+  const onMouseMove = useMouseMove();
+  const onMouseEnter = useMouseEnter();
+  const onMouseLeave = useMouseLeave();
+  const cursorPosition = useStore($cursorPosition);
+
+  useLayoutEffect(() => {
+    console.log('init effect');
+    if (!props.container) {
+      return;
+    }
+    initStage(props.container);
+    return () => {
+      const stage = $stage.get();
+      if (!stage) {
+        return;
+      }
+      stage.destroy();
+    };
+  }, [props.container]);
+
+  useLayoutEffect(() => {
+    console.log('event effect');
+    if (!stage) {
+      return;
+    }
+    stage.on('mousedown', onMouseDown);
+    stage.on('mouseup', onMouseUp);
+    stage.on('mousemove', onMouseMove);
+    stage.on('mouseenter', onMouseEnter);
+    stage.on('mouseleave', onMouseLeave);
+
+    return () => {
+      stage.off('mousedown', onMouseDown);
+      stage.off('mouseup', onMouseUp);
+      stage.off('mousemove', onMouseMove);
+      stage.off('mouseenter', onMouseEnter);
+      stage.off('mouseleave', onMouseLeave);
+    };
+  }, [stage, onMouseDown, onMouseUp, onMouseMove, onMouseEnter, onMouseLeave]);
+
+  useLayoutEffect(() => {
+    console.log('stage dims effect');
+    if (!stage || !props.container) {
+      return;
+    }
+    stage.width(width);
+    stage.height(height);
+  }, [stage, width, height, props.container]);
+
+  useLayoutEffect(() => {
+    console.log('cursor effect');
+    if (!stage || !cursorPosition) {
+      return;
+    }
+    const cursor = stage.findOne('#cursor');
+    if (!cursor) {
+      return;
+    }
+    cursor.x(cursorPosition?.x);
+    cursor.y(cursorPosition?.y);
+  }, [cursorPosition, stage]);
+
+  useLayoutEffect(() => {
+    console.log('obj effect');
+    if (!stage) {
+      return;
+    }
+
+    // TODO: Handle layer getting deleted and reset
+    for (const l of state.layers) {
+      let layer = stage.findOne(`#${l.id}`) as Konva.Layer | null;
+      if (!layer) {
+        layer = new Konva.Layer({ id: l.id, name: REGIONAL_PROMPT_LAYER_NAME, draggable: true });
+        layer.on('dragmove', (e) => {
+          dispatch(layerTranslated({ layerId: l.id, x: e.target.x(), y: e.target.y() }));
+        });
+        layer.dragBoundFunc(function (pos) {
+          const cursorPos = getScaledCursorPosition(stage);
+          if (!cursorPos) {
+            return this.getAbsolutePosition();
+          }
+          // This prevents the user from dragging the object out of the stage.
+          if (cursorPos.x < 0 || cursorPos.x > stage.width() || cursorPos.y < 0 || cursorPos.y > stage.height()) {
+            return this.getAbsolutePosition();
+          }
+
+          return pos;
+        });
+        stage.add(layer);
+      }
+
+      if (state.tool === 'move') {
+        layer.listening(true);
+      } else {
+        layer.listening(l.id === state.selectedLayer);
+      }
+
+      for (const o of l.objects) {
+        if (o.kind !== 'line') {
+          return;
+        }
+        let obj = stage.findOne(`#${o.id}`) as Konva.Line | null;
+        if (!obj) {
+          obj = new Konva.Line({
+            id: o.id,
+            key: o.id,
+            strokeWidth: o.strokeWidth,
+            stroke: rgbColorToString(l.color),
+            tension: 0,
+            lineCap: 'round',
+            lineJoin: 'round',
+            shadowForStrokeEnabled: false,
+            globalCompositeOperation: o.tool === 'brush' ? 'source-over' : 'destination-out',
+            listening: false,
+            visible: l.isVisible,
+          });
+          layer.add(obj);
+        }
+        if (obj.points().length < o.points.length) {
+          obj.points(o.points);
+        }
+        if (obj.stroke() !== rgbColorToString(l.color)) {
+          obj.stroke(rgbColorToString(l.color));
+        }
+        if (obj.visible() !== l.isVisible) {
+          obj.visible(l.isVisible);
+        }
+      }
+    }
+  }, [dispatch, stage, state.tool, state.layers, state.selectedLayer]);
+
+  useLayoutEffect(() => {
+    if (!stage) {
+      return;
+    }
+
+    if (state.tool !== 'move') {
+      for (const n of stage.find('.layer-bbox')) {
+        n.visible(false);
+      }
+      return;
+    }
+
+    for (const layer of stage.find(`.${REGIONAL_PROMPT_LAYER_NAME}`) as Konva.Layer[]) {
+      const bbox = getKonvaLayerBbox(layer);
+      dispatch(layerBboxChanged({ layerId: layer.id(), bbox }));
+      let rect = layer.findOne('.layer-bbox') as Konva.Rect | null;
+      if (!rect) {
+        rect = new Konva.Rect({
+          id: `${layer.id()}-bbox`,
+          name: 'layer-bbox',
+          strokeWidth: 1,
+        });
+        layer.add(rect);
+        layer.on('mousedown', () => {
+          dispatch(layerSelected(layer.id()));
+        });
+      }
+      rect.visible(true);
+      rect.x(bbox.x);
+      rect.y(bbox.y);
+      rect.width(bbox.width);
+      rect.height(bbox.height);
+      rect.stroke(state.selectedLayer === layer.id() ? 'rgba(153, 187, 189, 1)' : 'rgba(255, 255, 255, 0.149)');
+    }
+  }, [dispatch, stage, state.tool, state.selectedLayer]);
+
+  return null;
+};
+
+const $container = atom<HTMLDivElement | null>(null);
+const containerRef = (el: HTMLDivElement | null) => {
+  $container.set(el);
+};
+
+export const StageComponent = () => {
+  const container = useStore($container);
+  return (
+    <>
+      <chakra.div ref={containerRef} tabIndex={-1} sx={{ borderWidth: 1, borderRadius: 'base' }} />
+      <LogicalStage container={container} />
+    </>
+  );
+};
diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/imperative/mouseEventHooks.ts b/invokeai/frontend/web/src/features/regionalPrompts/components/imperative/mouseEventHooks.ts
new file mode 100644
index 0000000000..35b6eb21ee
--- /dev/null
+++ b/invokeai/frontend/web/src/features/regionalPrompts/components/imperative/mouseEventHooks.ts
@@ -0,0 +1,142 @@
+import { getStore } from 'app/store/nanostores/store';
+import { useAppDispatch } from 'app/store/storeHooks';
+import getScaledCursorPosition from 'features/canvas/util/getScaledCursorPosition';
+import { $stage } from 'features/regionalPrompts/components/imperative/konvaApiDraft';
+import {
+  $cursorPosition,
+  $isMouseDown,
+  $isMouseOver,
+  lineAdded,
+  pointsAdded,
+} from 'features/regionalPrompts/store/regionalPromptsSlice';
+import type Konva from 'konva';
+import type { KonvaEventObject } from 'konva/lib/Node';
+import { useCallback } from 'react';
+
+const getTool = () => getStore().getState().regionalPrompts.tool;
+
+const getIsFocused = (stage: Konva.Stage) => {
+  return stage.container().contains(document.activeElement);
+};
+
+const syncCursorPos = (stage: Konva.Stage) => {
+  const pos = getScaledCursorPosition(stage);
+  if (!pos) {
+    return null;
+  }
+  $cursorPosition.set(pos);
+  return pos;
+};
+
+export const useMouseDown = () => {
+  const dispatch = useAppDispatch();
+  const onMouseDown = useCallback(
+    (_e: KonvaEventObject<MouseEvent | TouchEvent>) => {
+      const stage = $stage.get();
+      if (!stage) {
+        return;
+      }
+      const pos = syncCursorPos(stage);
+      if (!pos) {
+        return;
+      }
+      $isMouseDown.set(true);
+      const tool = getTool();
+      if (tool === 'brush' || tool === 'eraser') {
+        dispatch(lineAdded([pos.x, pos.y]));
+      }
+    },
+    [dispatch]
+  );
+  return onMouseDown;
+};
+
+export const useMouseUp = () => {
+  const dispatch = useAppDispatch();
+  const onMouseUp = useCallback(
+    (_e: KonvaEventObject<MouseEvent | TouchEvent>) => {
+      const stage = $stage.get();
+      if (!stage) {
+        return;
+      }
+      const tool = getTool();
+      if ((tool === 'brush' || tool === 'eraser') && $isMouseDown.get()) {
+        // Add another point to the last line.
+        $isMouseDown.set(false);
+        const pos = syncCursorPos(stage);
+        if (!pos) {
+          return;
+        }
+        dispatch(pointsAdded([pos.x, pos.y]));
+      }
+    },
+    [dispatch]
+  );
+  return onMouseUp;
+};
+
+export const useMouseMove = () => {
+  const dispatch = useAppDispatch();
+  const onMouseMove = useCallback(
+    (_e: KonvaEventObject<MouseEvent | TouchEvent>) => {
+      const stage = $stage.get();
+      if (!stage) {
+        return;
+      }
+      const pos = syncCursorPos(stage);
+      if (!pos) {
+        return;
+      }
+      const tool = getTool();
+      if (getIsFocused(stage) && $isMouseOver.get() && $isMouseDown.get() && (tool === 'brush' || tool === 'eraser')) {
+        dispatch(pointsAdded([pos.x, pos.y]));
+      }
+    },
+    [dispatch]
+  );
+  return onMouseMove;
+};
+
+export const useMouseLeave = () => {
+  const onMouseLeave = useCallback((_e: KonvaEventObject<MouseEvent | TouchEvent>) => {
+    const stage = $stage.get();
+    if (!stage) {
+      return;
+    }
+    $isMouseOver.set(false);
+    $isMouseDown.set(false);
+    $cursorPosition.set(null);
+  }, []);
+  return onMouseLeave;
+};
+
+export const useMouseEnter = () => {
+  const dispatch = useAppDispatch();
+  const onMouseEnter = useCallback(
+    (e: KonvaEventObject<MouseEvent>) => {
+      const stage = $stage.get();
+      if (!stage) {
+        return;
+      }
+      $isMouseOver.set(true);
+      const pos = syncCursorPos(stage);
+      if (!pos) {
+        return;
+      }
+      if (!getIsFocused(stage)) {
+        return;
+      }
+      if (e.evt.buttons !== 1) {
+        $isMouseDown.set(false);
+      } else {
+        $isMouseDown.set(true);
+        const tool = getTool();
+        if (tool === 'brush' || tool === 'eraser') {
+          dispatch(lineAdded([pos.x, pos.y]));
+        }
+      }
+    },
+    [dispatch]
+  );
+  return onMouseEnter;
+};
diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/imperative/test.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/imperative/test.tsx
new file mode 100644
index 0000000000..ca1db3f6c8
--- /dev/null
+++ b/invokeai/frontend/web/src/features/regionalPrompts/components/imperative/test.tsx
@@ -0,0 +1,22 @@
+import type Konva from 'konva';
+import { useCallback, useEffect, useState } from 'react';
+import { Stage } from 'react-konva';
+
+export const StageWrapper = () => {
+  const [stage, setStage] = useState<Konva.Stage | null>(null);
+  const stageRefCallback = useCallback(
+    (el: Konva.Stage | null) => {
+      setStage(el);
+    },
+    [setStage]
+  );
+  useEffect(() => {
+    if (!stage) {
+      return;
+    }
+
+    // do something with stage
+  }, [stage]);
+
+  return <Stage ref={stageRefCallback} />;
+};
diff --git a/invokeai/frontend/web/src/features/regionalPrompts/util/getLayerBlobs.ts b/invokeai/frontend/web/src/features/regionalPrompts/util/getLayerBlobs.ts
index 114f68a478..5129c3c9d0 100644
--- a/invokeai/frontend/web/src/features/regionalPrompts/util/getLayerBlobs.ts
+++ b/invokeai/frontend/web/src/features/regionalPrompts/util/getLayerBlobs.ts
@@ -19,7 +19,7 @@ export const getRegionalPromptLayerBlobs = async (
 
   // This automatically omits layers that are not rendered. Rendering is controlled by the layer's `isVisible` flag in redux.
   const regionalPromptLayers = stage.getLayers().filter((l) => {
-    console.log(l.name(), l.id())
+    console.log(l.name(), l.id());
     const isRegionalPromptLayer = l.name() === REGIONAL_PROMPT_LAYER_NAME;
     const isRequestedLayerId = layerIds ? layerIds.includes(l.id()) : true;
     return isRegionalPromptLayer && isRequestedLayerId;