From 66b690e5c845607aa55c515838f691c3493f094b Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Mon, 29 Jul 2024 10:29:29 -0400
Subject: [PATCH] Manage canvas element in Vue (#255)

* Manage canvas element in Vue

* nit

* Fix unittest
---
 src/App.vue                          | 46 ++-------------------
 src/components/graph/GraphCanvas.vue | 61 ++++++++++++++++++++++++++++
 src/main.ts                          |  6 ---
 src/scripts/app.ts                   |  8 +---
 tests-ui/utils/index.ts              |  7 +++-
 5 files changed, 72 insertions(+), 56 deletions(-)
 create mode 100644 src/components/graph/GraphCanvas.vue

diff --git a/src/App.vue b/src/App.vue
index ad967f23e4..29e36de24e 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -1,22 +1,11 @@
 <template>
   <ProgressSpinner v-if="isLoading" class="spinner"></ProgressSpinner>
-  <div v-else>
-    <NodeSearchboxPopover v-if="nodeSearchEnabled" />
-    <teleport to=".graph-canvas-container">
-      <LiteGraphCanvasSplitterOverlay v-if="betaMenuEnabled">
-        <template #side-bar-panel>
-          <SideToolBar />
-        </template>
-      </LiteGraphCanvasSplitterOverlay>
-    </teleport>
-  </div>
+  <GraphCanvas v-else />
 </template>
 
 <script setup lang="ts">
-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'
+import { computed, markRaw, onMounted, ref, watch } from 'vue'
+import GraphCanvas from '@/components/graph/GraphCanvas.vue'
 import QueueSideBarTab from '@/components/sidebar/tabs/QueueSideBarTab.vue'
 import ProgressSpinner from 'primevue/progressspinner'
 import { app } from './scripts/app'
@@ -24,13 +13,8 @@ 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>(
-  () => useSettingStore().get('Comfy.NodeSearchBoxImpl') === 'default'
-)
 const theme = computed<string>(() =>
   useSettingStore().get('Comfy.ColorPalette')
 )
@@ -47,15 +31,10 @@ watch(
   },
   { immediate: true }
 )
-const betaMenuEnabled = computed(
-  () => useSettingStore().get('Comfy.UseNewMenu') !== 'Disabled'
-)
 
 const { t } = useI18n()
-let dropTargetCleanup = () => {}
 const init = () => {
   useSettingStore().addSettings(app.ui.settings)
-  app.vueAppReady = true
   app.extensionManager = useWorkspaceStore()
   app.extensionManager.registerSidebarTab({
     id: 'queue',
@@ -73,21 +52,6 @@ const init = () => {
     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(() => {
@@ -99,10 +63,6 @@ onMounted(() => {
     isLoading.value = false
   }
 })
-
-onUnmounted(() => {
-  dropTargetCleanup()
-})
 </script>
 
 <style scoped>
diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue
new file mode 100644
index 0000000000..a102da739e
--- /dev/null
+++ b/src/components/graph/GraphCanvas.vue
@@ -0,0 +1,61 @@
+<template>
+  <teleport to=".graph-canvas-container">
+    <LiteGraphCanvasSplitterOverlay v-if="betaMenuEnabled">
+      <template #side-bar-panel>
+        <SideToolBar />
+      </template>
+    </LiteGraphCanvasSplitterOverlay>
+    <canvas ref="canvasRef" id="graph-canvas" tabindex="1" />
+  </teleport>
+  <NodeSearchboxPopover v-if="nodeSearchEnabled" />
+</template>
+
+<script setup lang="ts">
+import SideToolBar from '@/components/sidebar/SideToolBar.vue'
+import LiteGraphCanvasSplitterOverlay from '@/components/LiteGraphCanvasSplitterOverlay.vue'
+import NodeSearchboxPopover from '@/components/NodeSearchBoxPopover.vue'
+import { ref, onMounted, computed, onUnmounted } from 'vue'
+import { app as comfyApp } from '@/scripts/app'
+import { useSettingStore } from '@/stores/settingStore'
+import { dropTargetForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
+import { useNodeDefStore } from '@/stores/nodeDefStore'
+
+const emit = defineEmits(['ready'])
+const canvasRef = ref<HTMLCanvasElement | null>(null)
+
+const betaMenuEnabled = computed(
+  () => useSettingStore().get('Comfy.UseNewMenu') !== 'Disabled'
+)
+const nodeSearchEnabled = computed<boolean>(
+  () => useSettingStore().get('Comfy.NodeSearchBoxImpl') === 'default'
+)
+
+let dropTargetCleanup = () => {}
+
+onMounted(async () => {
+  comfyApp.vueAppReady = true
+  await comfyApp.setup(canvasRef.value)
+  window['app'] = comfyApp
+  window['graph'] = comfyApp.graph
+
+  dropTargetCleanup = dropTargetForElements({
+    element: canvasRef.value,
+    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 = comfyApp.clientPosToCanvasPos([loc.clientX - 20, loc.clientY])
+      const comfyNodeName = event.source.element.getAttribute(
+        'data-comfy-node-name'
+      )
+      const nodeDef = useNodeDefStore().nodeDefsByName[comfyNodeName]
+      comfyApp.addNodeOnGraph(nodeDef, { pos })
+    }
+  })
+  emit('ready')
+})
+
+onUnmounted(() => {
+  dropTargetCleanup()
+})
+</script>
diff --git a/src/main.ts b/src/main.ts
index d7ba0033d9..327e9a1afc 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -8,7 +8,6 @@ import Tooltip from 'primevue/tooltip'
 import 'primeicons/primeicons.css'
 
 import App from './App.vue'
-import { app as comfyApp } from '@/scripts/app'
 import { createPinia } from 'pinia'
 import { i18n } from './i18n'
 
@@ -40,8 +39,3 @@ app
   .use(pinia)
   .use(i18n)
   .mount('#vue-app')
-
-comfyApp.setup().then(() => {
-  window['app'] = comfyApp
-  window['graph'] = comfyApp.graph
-})
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 00d30d79bd..dd96b9f72d 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -1853,17 +1853,13 @@ export class ComfyApp {
   /**
    * Set up the app on the page
    */
-  async setup() {
+  async setup(canvasEl: HTMLCanvasElement) {
+    this.canvasEl = canvasEl
     await this.#setUser()
 
     // Create and mount the LiteGraph in the DOM
     const mainCanvas = document.createElement('canvas')
     mainCanvas.style.touchAction = 'none'
-    const canvasEl = (this.canvasEl = Object.assign(mainCanvas, {
-      id: 'graph-canvas'
-    }))
-    canvasEl.tabIndex = 1
-    this.canvasContainer.prepend(canvasEl)
 
     this.resizeCanvas()
 
diff --git a/tests-ui/utils/index.ts b/tests-ui/utils/index.ts
index abbaaf69f5..ea0c086575 100644
--- a/tests-ui/utils/index.ts
+++ b/tests-ui/utils/index.ts
@@ -43,7 +43,12 @@ export async function start(config: StartConfig = {}): Promise<StartResult> {
   const { app } = await import('../../src/scripts/app')
   const { LiteGraph, LGraphCanvas } = await import('@comfyorg/litegraph')
   config.preSetup?.(app)
-  await app.setup()
+  const canvasEl = document.createElement('canvas')
+  canvasEl.style.touchAction = 'none'
+  canvasEl.id = 'graph-canvas'
+  canvasEl.tabIndex = 1
+  app.canvasContainer.prepend(canvasEl)
+  await app.setup(canvasEl)
 
   // @ts-ignore
   return { ...Ez.graph(app, LiteGraph, LGraphCanvas), app }