From 8ef3b87e59bf78db27acdda63fbd4e35121e8894 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Thu, 30 Jan 2025 11:11:22 -0800
Subject: [PATCH] Generate drag preview for node library nodes (#2387)

---
 src/components/common/TreeExplorerTreeNode.vue      | 13 ++++++++++++-
 .../sidebar/tabs/NodeLibrarySidebarTab.vue          | 10 +++++++++-
 src/hooks/canvasDropHooks.ts                        |  5 +++--
 src/types/treeExplorerTypes.ts                      |  5 +++++
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/components/common/TreeExplorerTreeNode.vue b/src/components/common/TreeExplorerTreeNode.vue
index 3235900d5..f77374211 100644
--- a/src/components/common/TreeExplorerTreeNode.vue
+++ b/src/components/common/TreeExplorerTreeNode.vue
@@ -34,6 +34,7 @@
 </template>
 
 <script setup lang="ts">
+import { setCustomNativeDragPreview } from '@atlaskit/pragmatic-drag-and-drop/element/set-custom-native-drag-preview'
 import Badge from 'primevue/badge'
 import { Ref, computed, inject, ref } from 'vue'
 
@@ -102,7 +103,17 @@ if (props.node.draggable) {
       }
     },
     onDragStart: () => emit('dragStart', props.node),
-    onDrop: () => emit('dragEnd', props.node)
+    onDrop: () => emit('dragEnd', props.node),
+    onGenerateDragPreview: props.node.renderDragPreview
+      ? ({ nativeSetDragImage }) => {
+          setCustomNativeDragPreview({
+            render: ({ container }) => {
+              return props.node.renderDragPreview(props.node, container)
+            },
+            nativeSetDragImage
+          })
+        }
+      : undefined
   })
 }
 
diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
index 7b1385604..8c0c3e1bb 100644
--- a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
@@ -66,11 +66,12 @@ import Button from 'primevue/button'
 import Divider from 'primevue/divider'
 import Popover from 'primevue/popover'
 import type { TreeNode } from 'primevue/treenode'
-import { Ref, computed, nextTick, ref } from 'vue'
+import { Ref, computed, h, nextTick, ref, render } from 'vue'
 
 import SearchBox from '@/components/common/SearchBox.vue'
 import { SearchFilter } from '@/components/common/SearchFilterChip.vue'
 import TreeExplorer from '@/components/common/TreeExplorer.vue'
+import NodePreview from '@/components/node/NodePreview.vue'
 import NodeSearchFilter from '@/components/searchbox/NodeSearchFilter.vue'
 import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'
 import NodeTreeLeaf from '@/components/sidebar/tabs/nodeLibrary/NodeTreeLeaf.vue'
@@ -125,6 +126,13 @@ const renderedRoot = computed<TreeExplorerNode<ComfyNodeDefImpl>>(() => {
       },
       children,
       draggable: node.leaf,
+      renderDragPreview: (node, container) => {
+        const vnode = h(NodePreview, { nodeDef: node.data })
+        render(vnode, container)
+        return () => {
+          render(null, container)
+        }
+      },
       handleClick: (
         node: RenderedTreeExplorerNode<ComfyNodeDefImpl>,
         e: MouseEvent
diff --git a/src/hooks/canvasDropHooks.ts b/src/hooks/canvasDropHooks.ts
index b259ab82e..d31d15c19 100644
--- a/src/hooks/canvasDropHooks.ts
+++ b/src/hooks/canvasDropHooks.ts
@@ -1,4 +1,5 @@
 import { LGraphNode } from '@comfyorg/litegraph'
+import { LiteGraph } from '@comfyorg/litegraph'
 import { Ref } from 'vue'
 
 import { usePragmaticDroppable } from '@/hooks/dndHooks'
@@ -28,8 +29,8 @@ export const useCanvasDrop = (canvasRef: Ref<HTMLCanvasElement>) => {
           // Add an offset on x to make sure after adding the node, the cursor
           // is on the node (top left corner)
           const pos = comfyApp.clientPosToCanvasPos([
-            loc.clientX - 20,
-            loc.clientY
+            loc.clientX,
+            loc.clientY + LiteGraph.NODE_TITLE_HEIGHT
           ])
           litegraphService.addNodeOnGraph(nodeDef, { pos })
         } else if (node.data instanceof ComfyModelDef) {
diff --git a/src/types/treeExplorerTypes.ts b/src/types/treeExplorerTypes.ts
index 692bfdeac..08a0e3383 100644
--- a/src/types/treeExplorerTypes.ts
+++ b/src/types/treeExplorerTypes.ts
@@ -25,6 +25,11 @@ export interface TreeExplorerNode<T = any> {
   ) => void | Promise<void>
   /** Whether the node is draggable */
   draggable?: boolean
+  /** Function to render a drag preview */
+  renderDragPreview?: (
+    node: TreeExplorerNode<T>,
+    container: HTMLElement
+  ) => void | (() => void)
   /** Whether the node is droppable */
   droppable?: boolean
   /** Function to handle dropping a node */