From 32dff2c4e3a3eadb4da63f7763143da584b23c2f Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 17 May 2024 18:41:59 +1000
Subject: [PATCH] feat(ui): copy/paste input edges when copying node

- Copy edges to selected nodes on copy
- If pasted with `ctrl/meta-shift-v`, also paste the input edges
---
 .../web/src/features/nodes/components/flow/Flow.tsx |  9 +++++++++
 .../web/src/features/nodes/hooks/useCopyPaste.ts    | 13 +++++++++++--
 .../web/src/features/nodes/store/nodesSlice.ts      |  1 +
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
index 8b33323ddd..656de737c7 100644
--- a/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
+++ b/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
@@ -231,6 +231,15 @@ export const Flow = memo(() => {
   );
   useHotkeys(['Ctrl+v', 'Meta+v'], onPasteHotkey);
 
+  const onPasteWithEdgesToNodesHotkey = useCallback(
+    (e: KeyboardEvent) => {
+      e.preventDefault();
+      pasteSelection(true);
+    },
+    [pasteSelection]
+  );
+  useHotkeys(['Ctrl+shift+v', 'Meta+shift+v'], onPasteWithEdgesToNodesHotkey);
+
   const onUndoHotkey = useCallback(() => {
     if (mayUndo) {
       dispatch(undo());
diff --git a/invokeai/frontend/web/src/features/nodes/hooks/useCopyPaste.ts b/invokeai/frontend/web/src/features/nodes/hooks/useCopyPaste.ts
index 9acd5722cf..08def1514c 100644
--- a/invokeai/frontend/web/src/features/nodes/hooks/useCopyPaste.ts
+++ b/invokeai/frontend/web/src/features/nodes/hooks/useCopyPaste.ts
@@ -4,10 +4,12 @@ import {
   $copiedEdges,
   $copiedNodes,
   $cursorPos,
+  $edgesToCopiedNodes,
   selectionPasted,
   selectNodesSlice,
 } from 'features/nodes/store/nodesSlice';
 import { findUnoccupiedPosition } from 'features/nodes/store/util/findUnoccupiedPosition';
+import { isEqual, uniqWith } from 'lodash-es';
 import { v4 as uuidv4 } from 'uuid';
 
 const copySelection = () => {
@@ -16,17 +18,24 @@ const copySelection = () => {
   const { nodes, edges } = selectNodesSlice(getState());
   const selectedNodes = nodes.filter((node) => node.selected);
   const selectedEdges = edges.filter((edge) => edge.selected);
+  const edgesToSelectedNodes = edges.filter((edge) => selectedNodes.some((node) => node.id === edge.target));
   $copiedNodes.set(selectedNodes);
   $copiedEdges.set(selectedEdges);
+  $edgesToCopiedNodes.set(edgesToSelectedNodes);
 };
 
-const pasteSelection = () => {
+const pasteSelection = (withEdgesToCopiedNodes?: boolean) => {
   const { getState, dispatch } = getStore();
   const currentNodes = selectNodesSlice(getState()).nodes;
   const cursorPos = $cursorPos.get();
 
   const copiedNodes = deepClone($copiedNodes.get());
-  const copiedEdges = deepClone($copiedEdges.get());
+  let copiedEdges = deepClone($copiedEdges.get());
+
+  if (withEdgesToCopiedNodes) {
+    const edgesToCopiedNodes = deepClone($edgesToCopiedNodes.get());
+    copiedEdges = uniqWith([...copiedEdges, ...edgesToCopiedNodes], isEqual);
+  }
 
   // Calculate an offset to reposition nodes to surround the cursor position, maintaining relative positioning
   const xCoords = copiedNodes.map((node) => node.position.x);
diff --git a/invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts b/invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts
index 644287dd29..1f61c77e83 100644
--- a/invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts
+++ b/invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts
@@ -498,6 +498,7 @@ export const $cursorPos = atom<XYPosition | null>(null);
 export const $templates = atom<Templates>({});
 export const $copiedNodes = atom<AnyNode[]>([]);
 export const $copiedEdges = atom<InvocationNodeEdge[]>([]);
+export const $edgesToCopiedNodes = atom<InvocationNodeEdge[]>([]);
 export const $pendingConnection = atom<PendingConnection | null>(null);
 export const $isUpdatingEdge = atom(false);
 export const $viewport = atom<Viewport>({ x: 0, y: 0, zoom: 1 });