From c0875d066aacfa72162ddd18e47662b5ad0f76bd Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Sat, 27 Jul 2024 21:28:48 -0400
Subject: [PATCH] Node library side bar tab (#237)

* Basic tree

* Add node filter

* Fix key issue

* Add icons

* Node count

* nit

* Add node preview basics

* Node preview

* Make comfy node in node lib draggable

* Set drop target

* Add node on drop

* Drop on dynamic location

* nit

* nit

* Fix hover preview issue

* nit

* More visual diff between node and folder

* Add playwright test

* Add dep

* Get rid of screenshot test
---
 browser_tests/ComfyPage.ts                    |  46 ++++++++
 browser_tests/menu.spec.ts                    |  25 +++++
 package-lock.json                             |  28 ++++-
 package.json                                  |   1 +
 src/App.vue                                   |  33 +++++-
 src/components/primevueOverride/TreePlus.vue  | 106 ++++++++++++++++++
 src/components/sidebar/SideToolBar.vue        |   1 +
 .../sidebar/tabs/NodeLibrarySideBarTab.vue    |  99 ++++++++++++++++
 src/i18n.ts                                   |   6 +-
 src/scripts/app.ts                            |   9 ++
 src/stores/nodeDefStore.ts                    |  28 +++++
 11 files changed, 376 insertions(+), 6 deletions(-)
 create mode 100644 src/components/primevueOverride/TreePlus.vue
 create mode 100644 src/components/sidebar/tabs/NodeLibrarySideBarTab.vue

diff --git a/browser_tests/ComfyPage.ts b/browser_tests/ComfyPage.ts
index 593aef8e5..38404fab9 100644
--- a/browser_tests/ComfyPage.ts
+++ b/browser_tests/ComfyPage.ts
@@ -37,6 +37,42 @@ class ComfyNodeSearchBox {
   }
 }
 
+class NodeLibrarySideBarTab {
+  public readonly tabId: string = 'node-library'
+  constructor(public readonly page: Page) {}
+
+  get tabButton() {
+    return this.page.locator(`.${this.tabId}-tab-button`)
+  }
+
+  get selectedTabButton() {
+    return this.page.locator(
+      `.${this.tabId}-tab-button.side-bar-button-selected`
+    )
+  }
+
+  get nodeLibraryTree() {
+    return this.page.locator('.node-lib-tree')
+  }
+
+  get nodePreview() {
+    return this.page.locator('.node-lib-node-preview')
+  }
+
+  async open() {
+    if (await this.selectedTabButton.isVisible()) {
+      return
+    }
+
+    await this.tabButton.click()
+    await this.nodeLibraryTree.waitFor({ state: 'visible' })
+  }
+
+  async toggleFirstFolder() {
+    await this.page.locator('.p-tree-node-toggle-button').nth(0).click()
+  }
+}
+
 class ComfyMenu {
   public readonly sideToolBar: Locator
   public readonly themeToggleButton: Locator
@@ -46,6 +82,10 @@ class ComfyMenu {
     this.themeToggleButton = page.locator('.comfy-vue-theme-toggle')
   }
 
+  get nodeLibraryTab() {
+    return new NodeLibrarySideBarTab(this.page)
+  }
+
   async toggleTheme() {
     await this.themeToggleButton.click()
     await this.page.evaluate(() => {
@@ -96,6 +136,12 @@ export class ComfyPage {
     this.menu = new ComfyMenu(page)
   }
 
+  async getGraphNodesCount(): Promise<number> {
+    return await this.page.evaluate(() => {
+      return window['app']?.graph?._nodes?.length || 0
+    })
+  }
+
   async setup() {
     await this.goto()
     // Unify font for consistent screenshots.
diff --git a/browser_tests/menu.spec.ts b/browser_tests/menu.spec.ts
index 369cc27c4..0ab6cefa2 100644
--- a/browser_tests/menu.spec.ts
+++ b/browser_tests/menu.spec.ts
@@ -67,4 +67,29 @@ test.describe('Menu', () => {
     )
     expect(newChildrenCount).toBe(initialChildrenCount + 1)
   })
+
+  test('Sidebar node preview and drag to canvas', async ({ comfyPage }) => {
+    // Open the sidebar
+    const tab = comfyPage.menu.nodeLibraryTab
+    await tab.open()
+    await tab.toggleFirstFolder()
+
+    // Hover over a node to display the preview
+    const nodeSelector = '.p-tree-node-leaf'
+    await comfyPage.page.hover(nodeSelector)
+
+    // Verify the preview is displayed
+    const previewVisible = await comfyPage.page.isVisible(
+      '.node-lib-node-preview'
+    )
+    expect(previewVisible).toBe(true)
+
+    const count = await comfyPage.getGraphNodesCount()
+    // Drag the node onto the canvas
+    const canvasSelector = '#graph-canvas'
+    await comfyPage.page.dragAndDrop(nodeSelector, canvasSelector)
+
+    // Verify the node is added to the canvas
+    expect(await comfyPage.getGraphNodesCount()).toBe(count + 1)
+  })
 })
diff --git a/package-lock.json b/package-lock.json
index 2b59213cf..2bc4bf884 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -8,6 +8,7 @@
       "name": "comfyui-frontend",
       "version": "1.2.2",
       "dependencies": {
+        "@atlaskit/pragmatic-drag-and-drop": "^1.2.1",
         "@comfyorg/litegraph": "^0.7.29",
         "@primevue/themes": "^4.0.0-rc.2",
         "@vitejs/plugin-vue": "^5.0.5",
@@ -78,6 +79,17 @@
         "node": ">=6.0.0"
       }
     },
+    "node_modules/@atlaskit/pragmatic-drag-and-drop": {
+      "version": "1.2.1",
+      "resolved": "https://registry.npmjs.org/@atlaskit/pragmatic-drag-and-drop/-/pragmatic-drag-and-drop-1.2.1.tgz",
+      "integrity": "sha512-gW2wJblFAeg94YXITHg0YdhFM2nmFAdDmX0LKYBIm79yEbIrOiuHHukgSjII07M4U5JpJ0Ff/4BaADjN23ix+A==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "@babel/runtime": "^7.0.0",
+        "bind-event-listener": "^3.0.0",
+        "raf-schd": "^4.0.3"
+      }
+    },
     "node_modules/@babel/code-frame": {
       "version": "7.24.7",
       "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.24.7.tgz",
@@ -1755,7 +1767,6 @@
       "version": "7.24.7",
       "resolved": "https://registry.npmjs.org/@babel/runtime/-/runtime-7.24.7.tgz",
       "integrity": "sha512-UwgBRMjJP+xv857DCngvqXI3Iq6J4v0wXmwc6sapg+zyhbwmQX67LUEFrkK5tbyJ30jGuG3ZvWpBiB9LCy1kWw==",
-      "dev": true,
       "dependencies": {
         "regenerator-runtime": "^0.14.0"
       },
@@ -4055,6 +4066,12 @@
         "url": "https://github.com/sponsors/sindresorhus"
       }
     },
+    "node_modules/bind-event-listener": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmjs.org/bind-event-listener/-/bind-event-listener-3.0.0.tgz",
+      "integrity": "sha512-PJvH288AWQhKs2v9zyfYdPzlPqf5bXbGMmhmUIY9x4dAUGIWgomO771oBQNwJnMQSnUIXhKu6sgzpBRXTlvb8Q==",
+      "license": "MIT"
+    },
     "node_modules/brace-expansion": {
       "version": "1.1.11",
       "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.11.tgz",
@@ -8573,6 +8590,12 @@
         }
       ]
     },
+    "node_modules/raf-schd": {
+      "version": "4.0.3",
+      "resolved": "https://registry.npmjs.org/raf-schd/-/raf-schd-4.0.3.tgz",
+      "integrity": "sha512-tQkJl2GRWh83ui2DiPTJz9wEiMN20syf+5oKfB03yYP7ioZcJwsIK8FjrtLwH1m7C7e+Tt2yYBlrOpdT+dyeIQ==",
+      "license": "MIT"
+    },
     "node_modules/react-is": {
       "version": "18.3.1",
       "resolved": "https://registry.npmjs.org/react-is/-/react-is-18.3.1.tgz",
@@ -8641,8 +8664,7 @@
     "node_modules/regenerator-runtime": {
       "version": "0.14.1",
       "resolved": "https://registry.npmjs.org/regenerator-runtime/-/regenerator-runtime-0.14.1.tgz",
-      "integrity": "sha512-dYnhHh0nJoMfnkZs6GmmhFknAGRrLznOu5nc9ML+EJxGvrx6H7teuevqVqCuPcPK//3eDrrjQhehXVx9cnkGdw==",
-      "dev": true
+      "integrity": "sha512-dYnhHh0nJoMfnkZs6GmmhFknAGRrLznOu5nc9ML+EJxGvrx6H7teuevqVqCuPcPK//3eDrrjQhehXVx9cnkGdw=="
     },
     "node_modules/regenerator-transform": {
       "version": "0.15.2",
diff --git a/package.json b/package.json
index 5777f8b1d..f53088908 100644
--- a/package.json
+++ b/package.json
@@ -46,6 +46,7 @@
     "zip-dir": "^2.0.0"
   },
   "dependencies": {
+    "@atlaskit/pragmatic-drag-and-drop": "^1.2.1",
     "@comfyorg/litegraph": "^0.7.29",
     "@primevue/themes": "^4.0.0-rc.2",
     "@vitejs/plugin-vue": "^5.0.5",
diff --git a/src/App.vue b/src/App.vue
index 691dbdfc1..ad967f23e 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -13,7 +13,7 @@
 </template>
 
 <script setup lang="ts">
-import { computed, markRaw, onMounted, ref, watch } from 'vue'
+import { computed, markRaw, onMounted, onUnmounted, ref, watch } from 'vue'
 import NodeSearchboxPopover from '@/components/NodeSearchBoxPopover.vue'
 import SideToolBar from '@/components/sidebar/SideToolBar.vue'
 import LiteGraphCanvasSplitterOverlay from '@/components/LiteGraphCanvasSplitterOverlay.vue'
@@ -23,6 +23,9 @@ import { app } from './scripts/app'
 import { useSettingStore } from './stores/settingStore'
 import { useI18n } from 'vue-i18n'
 import { useWorkspaceStore } from './stores/workspaceStateStore'
+import NodeLibrarySideBarTab from './components/sidebar/tabs/NodeLibrarySideBarTab.vue'
+import { dropTargetForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
+import { useNodeDefStore } from './stores/nodeDefStore'
 
 const isLoading = ref(true)
 const nodeSearchEnabled = computed<boolean>(
@@ -49,6 +52,7 @@ const betaMenuEnabled = computed(
 )
 
 const { t } = useI18n()
+let dropTargetCleanup = () => {}
 const init = () => {
   useSettingStore().addSettings(app.ui.settings)
   app.vueAppReady = true
@@ -61,6 +65,29 @@ const init = () => {
     component: markRaw(QueueSideBarTab),
     type: 'vue'
   })
+  app.extensionManager.registerSidebarTab({
+    id: 'node-library',
+    icon: 'pi pi-book',
+    title: t('sideToolBar.nodeLibrary'),
+    tooltip: t('sideToolBar.nodeLibrary'),
+    component: markRaw(NodeLibrarySideBarTab),
+    type: 'vue'
+  })
+
+  dropTargetCleanup = dropTargetForElements({
+    element: document.querySelector('.graph-canvas-container'),
+    onDrop: (event) => {
+      const loc = event.location.current.input
+      // Add an offset on x to make sure after adding the node, the cursor
+      // is on the node (top left corner)
+      const pos = app.clientPosToCanvasPos([loc.clientX - 20, loc.clientY])
+      const comfyNodeName = event.source.element.getAttribute(
+        'data-comfy-node-name'
+      )
+      const nodeDef = useNodeDefStore().nodeDefsByName[comfyNodeName]
+      app.addNodeOnGraph(nodeDef, { pos })
+    }
+  })
 }
 
 onMounted(() => {
@@ -72,6 +99,10 @@ onMounted(() => {
     isLoading.value = false
   }
 })
+
+onUnmounted(() => {
+  dropTargetCleanup()
+})
 </script>
 
 <style scoped>
diff --git a/src/components/primevueOverride/TreePlus.vue b/src/components/primevueOverride/TreePlus.vue
new file mode 100644
index 000000000..de163e0ca
--- /dev/null
+++ b/src/components/primevueOverride/TreePlus.vue
@@ -0,0 +1,106 @@
+<!-- Tree with all leaf nodes draggable -->
+<script>
+import Tree from 'primevue/tree'
+import { draggable } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
+import { h, onMounted, onBeforeUnmount, computed } from 'vue'
+
+export default {
+  name: 'TreePlus',
+  extends: Tree,
+  props: {
+    dragSelector: {
+      type: String,
+      default: '.p-tree-node'
+    },
+    // Explicitly declare all v-model props
+    expandedKeys: {
+      type: Object,
+      default: () => ({})
+    },
+    selectionKeys: {
+      type: Object,
+      default: () => ({})
+    }
+  },
+  emits: ['update:expandedKeys', 'update:selectionKeys'],
+  setup(props, context) {
+    // Create computed properties for each v-model prop
+    const computedExpandedKeys = computed({
+      get: () => props.expandedKeys,
+      set: (value) => context.emit('update:expandedKeys', value)
+    })
+
+    const computedSelectionKeys = computed({
+      get: () => props.selectionKeys,
+      set: (value) => context.emit('update:selectionKeys', value)
+    })
+
+    let observer = null
+
+    const makeDraggable = (element) => {
+      if (!element._draggableCleanup) {
+        element._draggableCleanup = draggable({
+          element
+        })
+      }
+    }
+
+    const observeTreeChanges = (treeElement) => {
+      observer = new MutationObserver((mutations) => {
+        mutations.forEach((mutation) => {
+          if (mutation.type === 'childList') {
+            mutation.addedNodes.forEach((node) => {
+              if (node.nodeType === Node.ELEMENT_NODE) {
+                node.querySelectorAll(props.dragSelector).forEach(makeDraggable)
+              }
+            })
+          }
+        })
+      })
+
+      observer.observe(treeElement, { childList: true, subtree: true })
+
+      // Make existing nodes draggable
+      treeElement.querySelectorAll(props.dragSelector).forEach(makeDraggable)
+    }
+
+    onMounted(() => {
+      const treeElement = document.querySelector('.p-tree')
+      if (treeElement) {
+        observeTreeChanges(treeElement)
+      }
+    })
+
+    onBeforeUnmount(() => {
+      if (observer) {
+        observer.disconnect()
+      }
+      // Clean up draggable instances if necessary
+      const treeElement = document.querySelector('.p-tree')
+      if (treeElement) {
+        treeElement.querySelectorAll(props.dragSelector).forEach((node) => {
+          if (node._draggableCleanup) {
+            node._draggableCleanup()
+          }
+        })
+      }
+    })
+
+    return () =>
+      h(
+        Tree,
+        {
+          ...context.attrs,
+          ...props,
+          expandedKeys: computedExpandedKeys.value,
+          selectionKeys: computedSelectionKeys.value,
+          'onUpdate:expandedKeys': (value) =>
+            (computedExpandedKeys.value = value),
+          'onUpdate:selectionKeys': (value) =>
+            (computedSelectionKeys.value = value)
+        },
+        context.slots
+      )
+  }
+}
+</script>
diff --git a/src/components/sidebar/SideToolBar.vue b/src/components/sidebar/SideToolBar.vue
index 868148357..c446f31c4 100644
--- a/src/components/sidebar/SideToolBar.vue
+++ b/src/components/sidebar/SideToolBar.vue
@@ -7,6 +7,7 @@
         :icon="tab.icon"
         :tooltip="tab.tooltip"
         :selected="tab === selectedTab"
+        :class="tab.id + '-tab-button'"
         @click="onTabClick(tab)"
       />
       <div class="side-tool-bar-end">
diff --git a/src/components/sidebar/tabs/NodeLibrarySideBarTab.vue b/src/components/sidebar/tabs/NodeLibrarySideBarTab.vue
new file mode 100644
index 000000000..0697cf67b
--- /dev/null
+++ b/src/components/sidebar/tabs/NodeLibrarySideBarTab.vue
@@ -0,0 +1,99 @@
+<template>
+  <TreePlus
+    class="node-lib-tree"
+    v-model:expandedKeys="expandedKeys"
+    selectionMode="single"
+    :value="renderedRoot.children"
+    :filter="true"
+    filterMode="lenient"
+    dragSelector=".p-tree-node-leaf"
+    :pt="{
+      nodeLabel: 'node-lib-tree-node-label',
+      nodeChildren: ({ props }) => ({
+        'data-comfy-node-name': props.node?.data?.name,
+        onMouseenter: (event: MouseEvent) => {
+          hoveredComfyNodeName = props.node?.data?.name
+
+          const hoverTarget = event.target as HTMLElement
+          const targetRect = hoverTarget.getBoundingClientRect()
+          nodePreviewStyle.top = `${targetRect.top - 40}px`
+          nodePreviewStyle.left = `${targetRect.right}px`
+        },
+        onMouseleave: () => {
+          hoveredComfyNodeName = null
+        }
+      })
+    }"
+  >
+    <template #folder="{ node }">
+      <span class="folder-label">{{ node.label }}</span>
+      <Badge
+        :value="node.totalNodes"
+        severity="secondary"
+        :style="{ marginLeft: '0.5rem' }"
+      ></Badge>
+    </template>
+    <template #node="{ node }">
+      <span class="node-label">{{ node.label }}</span>
+    </template>
+  </TreePlus>
+  <div
+    v-if="hoveredComfyNode"
+    class="node-lib-node-preview"
+    :style="nodePreviewStyle"
+  >
+    <NodePreview
+      :key="hoveredComfyNode.name"
+      :nodeDef="hoveredComfyNode"
+    ></NodePreview>
+  </div>
+</template>
+
+<script setup lang="ts">
+import Badge from 'primevue/badge'
+import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
+import { computed, ref } from 'vue'
+import { TreeNode } from 'primevue/treenode'
+import TreePlus from '@/components/primevueOverride/TreePlus.vue'
+import NodePreview from '@/components/NodePreview.vue'
+
+const nodeDefStore = useNodeDefStore()
+const expandedKeys = ref({})
+const hoveredComfyNodeName = ref<string | null>(null)
+const hoveredComfyNode = computed<ComfyNodeDefImpl | null>(() => {
+  if (!hoveredComfyNodeName.value) {
+    return null
+  }
+  return nodeDefStore.nodeDefsByName[hoveredComfyNodeName.value] || null
+})
+const nodePreviewStyle = ref<Record<string, string>>({
+  position: 'absolute',
+  top: '0px',
+  left: '0px'
+})
+
+const root = computed(() => nodeDefStore.nodeTree)
+const renderedRoot = computed(() => {
+  return fillNodeInfo(root.value)
+})
+const fillNodeInfo = (node: TreeNode): TreeNode => {
+  const isLeaf = node.children === undefined || node.children.length === 0
+  const isExpanded = expandedKeys.value[node.key]
+  const icon = isLeaf
+    ? 'pi pi-circle-fill'
+    : isExpanded
+      ? 'pi pi-folder-open'
+      : 'pi pi-folder'
+  const children = node.children?.map(fillNodeInfo)
+
+  return {
+    ...node,
+    icon,
+    children,
+    type: isLeaf ? 'node' : 'folder',
+    totalNodes: isLeaf
+      ? 1
+      : children.reduce((acc, child) => acc + child.totalNodes, 0)
+  }
+}
+</script>
diff --git a/src/i18n.ts b/src/i18n.ts
index baa3414f2..d51f4996e 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -5,14 +5,16 @@ const messages = {
     sideToolBar: {
       settings: 'Settings',
       themeToggle: 'Toggle Theme',
-      queue: 'Queue'
+      queue: 'Queue',
+      nodeLibrary: 'Node Library'
     }
   },
   zh: {
     sideToolBar: {
       settings: '设置',
       themeToggle: '主题切换',
-      queue: '队列'
+      queue: '队列',
+      nodeLibrary: '节点库'
     }
   }
   // TODO: Add more languages
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index bf9bac910..2aab2d5ee 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -42,6 +42,7 @@ import {
   useNodeDefStore
 } from '@/stores/nodeDefStore'
 import { Vector2 } from '@comfyorg/litegraph'
+import _ from 'lodash'
 
 export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview'
 
@@ -2951,6 +2952,14 @@ export class ComfyApp {
     this.graph.add(node)
     return node
   }
+
+  clientPosToCanvasPos(pos: Vector2): Vector2 {
+    const rect = this.canvasContainer.getBoundingClientRect()
+    const containerOffsets = [rect.left, rect.top]
+    return _.zip(pos, this.canvas.ds.offset, containerOffsets).map(
+      ([p, o1, o2]) => p / this.canvas.ds.scale - o1 - o2
+    ) as Vector2
+  }
 }
 
 export const app = new ComfyApp()
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index 984ae1fa3..eeadf332f 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -3,6 +3,7 @@ import { ComfyNodeDef } from '@/types/apiTypes'
 import { defineStore } from 'pinia'
 import { Type, Transform, plainToClass } from 'class-transformer'
 import { ComfyWidgetConstructor } from '@/scripts/widgets'
+import { TreeNode } from 'primevue/treenode'
 
 export class BaseInputSpec<T = any> {
   name: string
@@ -232,6 +233,33 @@ export const useNodeDefStore = defineStore('nodeDef', {
     },
     nodeSearchService(state) {
       return new NodeSearchService(Object.values(state.nodeDefsByName))
+    },
+    nodeTree(state): TreeNode {
+      const root: TreeNode = {
+        key: 'root',
+        label: 'Nodes',
+        children: []
+      }
+      for (const nodeDef of Object.values(state.nodeDefsByName)) {
+        const path = nodeDef.category.split('/')
+        let current = root
+        let key = 'root'
+        for (const part of path) {
+          key += `/${part}`
+          let next = current.children.find((child) => child.label === part)
+          if (!next) {
+            next = { key, label: part, children: [] }
+            current.children.push(next)
+          }
+          current = next
+        }
+        current.children.push({
+          label: nodeDef.display_name,
+          data: nodeDef,
+          key: `${key}/${nodeDef.name}`
+        })
+      }
+      return root
     }
   },
   actions: {