From 3bf0dfd218d2222714ae4d0884187551073f45db Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Mon, 30 Dec 2024 16:09:44 -0500
Subject: [PATCH] [Refactor] Move addNodeOnGraph to litegraphService (#2100)

---
 src/components/graph/GraphCanvas.vue          | 13 ++++++---
 .../searchbox/NodeSearchBoxPopover.vue        |  9 ++++---
 .../sidebar/tabs/ModelLibrarySidebarTab.vue   |  6 ++---
 .../sidebar/tabs/NodeLibrarySidebarTab.vue    |  4 +--
 .../nodeLibrary/NodeBookmarkTreeExplorer.vue  |  4 +--
 src/scripts/app.ts                            | 25 +----------------
 src/services/litegraphService.ts              | 27 ++++++++++++++++++-
 7 files changed, 48 insertions(+), 40 deletions(-)

diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue
index 9db202e04..27ad9f604 100644
--- a/src/components/graph/GraphCanvas.vue
+++ b/src/components/graph/GraphCanvas.vue
@@ -69,9 +69,11 @@ import { useColorPaletteStore } from '@/stores/workspace/colorPaletteStore'
 import { useColorPaletteService } from '@/services/colorPaletteService'
 import { IS_CONTROL_WIDGET, updateControlWidgetLabel } from '@/scripts/widgets'
 import { CORE_SETTINGS } from '@/constants/coreSettings'
+import { useLitegraphService } from '@/services/litegraphService'
 
 const emit = defineEmits(['ready'])
 const canvasRef = ref<HTMLCanvasElement | null>(null)
+const litegraphService = useLitegraphService()
 const settingStore = useSettingStore()
 const nodeDefStore = useNodeDefStore()
 const workspaceStore = useWorkspaceStore()
@@ -262,7 +264,7 @@ usePragmaticDroppable(() => canvasRef.value, {
           loc.clientX - 20,
           loc.clientY
         ])
-        comfyApp.addNodeOnGraph(nodeDef, { pos })
+        litegraphService.addNodeOnGraph(nodeDef, { pos })
       } else if (node.data instanceof ComfyModelDef) {
         const model = node.data
         const pos = comfyApp.clientPosToCanvasPos([loc.clientX, loc.clientY])
@@ -283,9 +285,12 @@ usePragmaticDroppable(() => canvasRef.value, {
         if (!targetGraphNode) {
           const provider = modelToNodeStore.getNodeProvider(model.directory)
           if (provider) {
-            targetGraphNode = comfyApp.addNodeOnGraph(provider.nodeDef, {
-              pos
-            })
+            targetGraphNode = litegraphService.addNodeOnGraph(
+              provider.nodeDef,
+              {
+                pos
+              }
+            )
             targetProvider = provider
           }
         }
diff --git a/src/components/searchbox/NodeSearchBoxPopover.vue b/src/components/searchbox/NodeSearchBoxPopover.vue
index 961c2a189..f6ea23aca 100644
--- a/src/components/searchbox/NodeSearchBoxPopover.vue
+++ b/src/components/searchbox/NodeSearchBoxPopover.vue
@@ -33,7 +33,6 @@
 </template>
 
 <script setup lang="ts">
-import { app } from '@/scripts/app'
 import { computed, ref, toRaw, watchEffect } from 'vue'
 import NodeSearchBox from './NodeSearchBox.vue'
 import Dialog from 'primevue/dialog'
@@ -53,15 +52,17 @@ import type {
 import { useEventListener } from '@vueuse/core'
 import { useSearchBoxStore } from '@/stores/workspace/searchBoxStore'
 import { storeToRefs } from 'pinia'
+import { useLitegraphService } from '@/services/litegraphService'
 
 const settingStore = useSettingStore()
+const litegraphService = useLitegraphService()
 
 const { visible } = storeToRefs(useSearchBoxStore())
 const dismissable = ref(true)
 const triggerEvent = ref<LiteGraphCanvasEvent | null>(null)
 const getNewNodeLocation = (): Vector2 => {
   if (!triggerEvent.value) {
-    return app.getCanvasCenter()
+    return litegraphService.getCanvasCenter()
   }
 
   const originalEvent = (triggerEvent.value.detail as OriginalEvent)
@@ -85,7 +86,9 @@ const closeDialog = () => {
 }
 
 const addNode = (nodeDef: ComfyNodeDefImpl) => {
-  const node = app.addNodeOnGraph(nodeDef, { pos: getNewNodeLocation() })
+  const node = litegraphService.addNodeOnGraph(nodeDef, {
+    pos: getNewNodeLocation()
+  })
 
   const eventDetail = triggerEvent.value?.detail
   if (eventDetail && eventDetail.subType === 'empty-release') {
diff --git a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
index 36a95d55d..13236b034 100644
--- a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
@@ -66,9 +66,9 @@ import type {
 } from '@/types/treeExplorerTypes'
 import { computed, ref, watch, toRef, onMounted, nextTick } from 'vue'
 import type { TreeNode } from 'primevue/treenode'
-import { app } from '@/scripts/app'
 import { buildTree } from '@/utils/treeUtil'
 import { isElectron } from '@/utils/envUtil'
+import { useLitegraphService } from '@/services/litegraphService'
 
 const modelStore = useModelStore()
 const modelToNodeStore = useModelToNodeStore()
@@ -153,9 +153,7 @@ const renderedRoot = computed<TreeExplorerNode<ModelOrFolder>>(() => {
         if (node.leaf) {
           const provider = modelToNodeStore.getNodeProvider(model.directory)
           if (provider) {
-            const node = app.addNodeOnGraph(provider.nodeDef, {
-              pos: app.getCanvasCenter()
-            })
+            const node = useLitegraphService().addNodeOnGraph(provider.nodeDef)
             const widget = node.widgets.find(
               (widget) => widget.name === provider.key
             )
diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
index 1ba452937..3b7f71eb7 100644
--- a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
@@ -77,7 +77,6 @@ import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue
 import NodeBookmarkTreeExplorer from './nodeLibrary/NodeBookmarkTreeExplorer.vue'
 import TreeExplorer from '@/components/common/TreeExplorer.vue'
 import NodeTreeLeaf from '@/components/sidebar/tabs/nodeLibrary/NodeTreeLeaf.vue'
-import { app } from '@/scripts/app'
 import { sortedTree } from '@/utils/treeUtil'
 import { useTreeExpansion } from '@/hooks/treeHooks'
 import NodeSearchFilter from '@/components/searchbox/NodeSearchFilter.vue'
@@ -88,6 +87,7 @@ import type {
   TreeExplorerNode
 } from '@/types/treeExplorerTypes'
 import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
+import { useLitegraphService } from '@/services/litegraphService'
 
 const nodeDefStore = useNodeDefStore()
 const nodeBookmarkStore = useNodeBookmarkStore()
@@ -128,7 +128,7 @@ const renderedRoot = computed<TreeExplorerNode<ComfyNodeDefImpl>>(() => {
         e: MouseEvent
       ) => {
         if (node.leaf) {
-          app.addNodeOnGraph(node.data, { pos: app.getCanvasCenter() })
+          useLitegraphService().addNodeOnGraph(node.data)
         } else {
           toggleNodeOnEvent(e, node)
         }
diff --git a/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue b/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
index fb2e82b2d..f96461e0e 100644
--- a/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
+++ b/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
@@ -37,8 +37,8 @@ import type { TreeNode } from 'primevue/treenode'
 import { computed, nextTick, ref, watch } from 'vue'
 import { useI18n } from 'vue-i18n'
 import { useTreeExpansion } from '@/hooks/treeHooks'
-import { app } from '@/scripts/app'
 import { findNodeByKey } from '@/utils/treeUtil'
+import { useLitegraphService } from '@/services/litegraphService'
 
 const props = defineProps<{
   filteredNodeDefs: ComfyNodeDefImpl[]
@@ -169,7 +169,7 @@ const renderedBookmarkedRoot = computed<TreeExplorerNode<ComfyNodeDefImpl>>(
           e: MouseEvent
         ) => {
           if (node.leaf) {
-            app.addNodeOnGraph(node.data, { pos: app.getCanvasCenter() })
+            useLitegraphService().addNodeOnGraph(node.data)
           } else {
             toggleNodeOnEvent(e, node)
           }
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index b037abd4c..8849408e0 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -29,11 +29,7 @@ import {
   LGraphEventMode
 } from '@comfyorg/litegraph'
 import { ExtensionManager } from '@/types/extensionTypes'
-import {
-  ComfyNodeDefImpl,
-  SYSTEM_NODE_DEFS,
-  useNodeDefStore
-} from '@/stores/nodeDefStore'
+import { SYSTEM_NODE_DEFS, useNodeDefStore } from '@/stores/nodeDefStore'
 import { Vector2 } from '@comfyorg/litegraph'
 import _ from 'lodash'
 import { useDialogService } from '@/services/dialogService'
@@ -2049,19 +2045,6 @@ export class ComfyApp {
     this.lastExecutionError = null
   }
 
-  addNodeOnGraph(
-    nodeDef: ComfyNodeDef | ComfyNodeDefImpl,
-    options: Record<string, any> = {}
-  ): LGraphNode {
-    const node = LiteGraph.createNode(
-      nodeDef.name,
-      nodeDef.display_name,
-      options
-    )
-    this.graph.add(node)
-    return node
-  }
-
   clientPosToCanvasPos(pos: Vector2): Vector2 {
     const rect = this.canvasContainer.getBoundingClientRect()
     const containerOffsets = [rect.left, rect.top]
@@ -2078,12 +2061,6 @@ export class ComfyApp {
     ) as Vector2
   }
 
-  getCanvasCenter(): Vector2 {
-    const dpi = Math.max(window.devicePixelRatio ?? 1, 1)
-    const [x, y, w, h] = app.canvas.ds.visible_area
-    return [x + w / dpi / 2, y + h / dpi / 2]
-  }
-
   public goToNode(nodeId: NodeId) {
     const graphNode = this.graph.getNodeById(nodeId)
     if (!graphNode) return
diff --git a/src/services/litegraphService.ts b/src/services/litegraphService.ts
index dcd5ecf38..439ec03ab 100644
--- a/src/services/litegraphService.ts
+++ b/src/services/litegraphService.ts
@@ -16,6 +16,7 @@ import { ComfyApp, app, ANIM_PREVIEW_WIDGET } from '@/scripts/app'
 import { $el } from '@/scripts/ui'
 import { useToastStore } from '@/stores/toastStore'
 import { calculateImageGrid, createImageHost } from '@/scripts/ui/imagePreview'
+import { Vector2 } from '@comfyorg/litegraph'
 
 /**
  * Service that augments litegraph with ComfyUI specific functionality.
@@ -760,7 +761,31 @@ export const useLitegraphService = () => {
     }
   }
 
+  function addNodeOnGraph(
+    nodeDef: ComfyNodeDef,
+    options: Record<string, any> = {}
+  ): LGraphNode {
+    options.pos ??= getCanvasCenter()
+
+    const node = LiteGraph.createNode(
+      nodeDef.name,
+      nodeDef.display_name,
+      options
+    )
+
+    app.graph.add(node)
+    return node
+  }
+
+  function getCanvasCenter(): Vector2 {
+    const dpi = Math.max(window.devicePixelRatio ?? 1, 1)
+    const [x, y, w, h] = app.canvas.ds.visible_area
+    return [x + w / dpi / 2, y + h / dpi / 2]
+  }
+
   return {
-    registerNodeDef
+    registerNodeDef,
+    addNodeOnGraph,
+    getCanvasCenter
   }
 }