From 4a230f720e7af3b22d575fe5b3da1fbc3ac94a06 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Mon, 2 Sep 2024 20:20:40 -0400
Subject: [PATCH] Edit group name on group creation (With Ctrl + g) (#715)

* Editor store

* Merge editors

* nit

* Edit on group creation

* nit
---
 browser_tests/interaction.spec.ts         |   3 +
 src/components/graph/GraphCanvas.vue      |   6 +-
 src/components/graph/GroupTitleEditor.vue | 104 --------------
 src/components/graph/NodeTitleEditor.vue  | 101 -------------
 src/components/graph/TitleEditor.vue      | 166 ++++++++++++++++++++++
 src/scripts/ui.ts                         |   2 +
 src/stores/graphStore.ts                  |  11 ++
 7 files changed, 184 insertions(+), 209 deletions(-)
 delete mode 100644 src/components/graph/GroupTitleEditor.vue
 delete mode 100644 src/components/graph/NodeTitleEditor.vue
 create mode 100644 src/components/graph/TitleEditor.vue
 create mode 100644 src/stores/graphStore.ts

diff --git a/browser_tests/interaction.spec.ts b/browser_tests/interaction.spec.ts
index 9948a2458..a0f1f2d23 100644
--- a/browser_tests/interaction.spec.ts
+++ b/browser_tests/interaction.spec.ts
@@ -165,6 +165,9 @@ test.describe('Node Interaction', () => {
     await comfyPage.page.keyboard.press('KeyG')
     await comfyPage.page.keyboard.up('Control')
     await comfyPage.nextFrame()
+    // Confirm group title
+    await comfyPage.page.keyboard.press('Enter')
+    await comfyPage.nextFrame()
     await expect(comfyPage.canvas).toHaveScreenshot('group-selected-nodes.png')
   })
 })
diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue
index 775ec73e4..533ba8948 100644
--- a/src/components/graph/GraphCanvas.vue
+++ b/src/components/graph/GraphCanvas.vue
@@ -5,8 +5,7 @@
         <SideToolbar />
       </template>
     </LiteGraphCanvasSplitterOverlay>
-    <NodeTitleEditor />
-    <GroupTitleEditor />
+    <TitleEditor />
     <canvas ref="canvasRef" id="graph-canvas" tabindex="1" />
   </teleport>
   <NodeSearchboxPopover v-if="nodeSearchEnabled" />
@@ -14,8 +13,7 @@
 </template>
 
 <script setup lang="ts">
-import NodeTitleEditor from '@/components/graph/NodeTitleEditor.vue'
-import GroupTitleEditor from '@/components/graph/GroupTitleEditor.vue'
+import TitleEditor from '@/components/graph/TitleEditor.vue'
 import SideToolbar from '@/components/sidebar/SideToolbar.vue'
 import LiteGraphCanvasSplitterOverlay from '@/components/LiteGraphCanvasSplitterOverlay.vue'
 import NodeSearchboxPopover from '@/components/searchbox/NodeSearchBoxPopover.vue'
diff --git a/src/components/graph/GroupTitleEditor.vue b/src/components/graph/GroupTitleEditor.vue
deleted file mode 100644
index c686ef17e..000000000
--- a/src/components/graph/GroupTitleEditor.vue
+++ /dev/null
@@ -1,104 +0,0 @@
-<template>
-  <div v-if="showInput" class="group-title-editor" :style="inputStyle">
-    <EditableText
-      :isEditing="showInput"
-      :modelValue="editedTitle"
-      @edit="onEdit"
-    />
-  </div>
-</template>
-
-<script setup lang="ts">
-import { ref, onMounted, onUnmounted, CSSProperties } from 'vue'
-import { app } from '@/scripts/app'
-import { LGraphGroup } from '@comfyorg/litegraph'
-import EditableText from '@/components/common/EditableText.vue'
-import { useSettingStore } from '@/stores/settingStore'
-import type { LiteGraphCanvasEvent } from '@comfyorg/litegraph'
-
-const settingStore = useSettingStore()
-
-const showInput = ref(false)
-const editedTitle = ref('')
-const inputStyle = ref<CSSProperties>({
-  position: 'fixed',
-  left: '0px',
-  top: '0px',
-  width: '200px',
-  height: '20px',
-  fontSize: '12px'
-})
-
-const currentGroup = ref<LGraphGroup | null>(null)
-
-const onEdit = (newValue: string) => {
-  if (currentGroup.value && newValue.trim() !== '') {
-    currentGroup.value.title = newValue.trim()
-    app.graph.setDirtyCanvas(true, true)
-  }
-  showInput.value = false
-}
-
-const canvasEventHandler = (event: LiteGraphCanvasEvent) => {
-  if (!settingStore.get('Comfy.Group.DoubleClickTitleToEdit')) {
-    return
-  }
-
-  if (event.detail.subType === 'group-double-click') {
-    const group: LGraphGroup = event.detail.group
-    const [x, y] = group.pos
-    const [w, h] = group.size
-
-    const e: MouseEvent = event.detail.originalEvent
-    // @ts-expect-error LiteGraphCanvasEvent is not typed
-    const relativeY = e.canvasY - y
-    // Only allow editing if the click is on the title bar
-    if (relativeY > group.titleHeight) {
-      return
-    }
-
-    currentGroup.value = group
-    editedTitle.value = group.title
-    showInput.value = true
-
-    const [left, top] = app.canvasPosToClientPos([x, y])
-    inputStyle.value.left = `${left}px`
-    inputStyle.value.top = `${top}px`
-
-    const width = w * app.canvas.ds.scale
-    const height = group.titleHeight * app.canvas.ds.scale
-    inputStyle.value.width = `${width}px`
-    inputStyle.value.height = `${height}px`
-
-    const fontSize = group.font_size * app.canvas.ds.scale
-    inputStyle.value.fontSize = `${fontSize}px`
-  }
-}
-
-onMounted(() => {
-  document.addEventListener('litegraph:canvas', canvasEventHandler)
-})
-
-onUnmounted(() => {
-  document.removeEventListener('litegraph:canvas', canvasEventHandler)
-})
-</script>
-
-<style scoped>
-.node-title-editor {
-  z-index: 9999;
-  padding: 0.25rem;
-}
-
-:deep(.editable-text) {
-  width: 100%;
-  height: 100%;
-}
-
-:deep(.editable-text input) {
-  width: 100%;
-  height: 100%;
-  /* Override the default font size */
-  font-size: inherit;
-}
-</style>
diff --git a/src/components/graph/NodeTitleEditor.vue b/src/components/graph/NodeTitleEditor.vue
deleted file mode 100644
index 99d466ef7..000000000
--- a/src/components/graph/NodeTitleEditor.vue
+++ /dev/null
@@ -1,101 +0,0 @@
-<template>
-  <div v-if="showInput" class="node-title-editor" :style="inputStyle">
-    <EditableText
-      :isEditing="showInput"
-      :modelValue="editedTitle"
-      @edit="onEdit"
-    />
-  </div>
-</template>
-
-<script setup lang="ts">
-import { ref, onMounted, CSSProperties } from 'vue'
-import { app } from '@/scripts/app'
-import { LGraphNode } from '@comfyorg/litegraph'
-import { ComfyExtension } from '@/types/comfy'
-import EditableText from '@/components/common/EditableText.vue'
-import { LiteGraph } from '@comfyorg/litegraph'
-import { useSettingStore } from '@/stores/settingStore'
-
-const settingStore = useSettingStore()
-
-const showInput = ref(false)
-const editedTitle = ref('')
-const inputStyle = ref<CSSProperties>({
-  position: 'fixed',
-  left: '0px',
-  top: '0px',
-  width: '200px',
-  height: '20px'
-})
-
-const currentNode = ref<LGraphNode | null>(null)
-
-const onEdit = (newValue: string) => {
-  if (currentNode.value && newValue.trim() !== '') {
-    currentNode.value.title = newValue.trim()
-    app.graph.setDirtyCanvas(true, true)
-  }
-  showInput.value = false
-}
-
-const extension: ComfyExtension = {
-  name: 'Comfy.NodeTitleEditor',
-  nodeCreated(node: LGraphNode) {
-    // Store the original callback
-    const originalCallback = node.onNodeTitleDblClick
-
-    node.onNodeTitleDblClick = function (e: MouseEvent, ...args: any[]) {
-      if (!settingStore.get('Comfy.Node.DoubleClickTitleToEdit')) {
-        return
-      }
-
-      currentNode.value = this
-      editedTitle.value = this.title
-      showInput.value = true
-
-      const isCollapsed = node.flags?.collapsed
-      const [x1, y1, nodeWidth, nodeHeight] = this.getBounding()
-      const canvasWidth =
-        // @ts-expect-error Remove after collapsed_width is exposed in LiteGraph
-        isCollapsed && node._collapsed_width ? node._collapsed_width : nodeWidth
-      const canvasHeight = LiteGraph.NODE_TITLE_HEIGHT
-
-      const [left, top] = app.canvasPosToClientPos([x1, y1])
-      inputStyle.value.left = `${left}px`
-      inputStyle.value.top = `${top}px`
-
-      const width = canvasWidth * app.canvas.ds.scale
-      const height = canvasHeight * app.canvas.ds.scale
-      inputStyle.value.width = `${width}px`
-      inputStyle.value.height = `${height}px`
-
-      // Call the original callback if it exists
-      if (typeof originalCallback === 'function') {
-        originalCallback.call(this, e, ...args)
-      }
-    }
-  }
-}
-
-onMounted(() => {
-  app.registerExtension(extension)
-})
-</script>
-
-<style scoped>
-.node-title-editor {
-  z-index: 9999;
-  padding: 0.25rem;
-}
-
-:deep(.editable-text) {
-  width: 100%;
-  height: 100%;
-}
-
-:deep(.editable-text input) {
-  width: 100%;
-  height: 100%;
-}
-</style>
diff --git a/src/components/graph/TitleEditor.vue b/src/components/graph/TitleEditor.vue
new file mode 100644
index 000000000..b565f2650
--- /dev/null
+++ b/src/components/graph/TitleEditor.vue
@@ -0,0 +1,166 @@
+<template>
+  <div
+    v-if="showInput"
+    class="group-title-editor node-title-editor"
+    :style="inputStyle"
+  >
+    <EditableText
+      :isEditing="showInput"
+      :modelValue="editedTitle"
+      @edit="onEdit"
+    />
+  </div>
+</template>
+
+<script setup lang="ts">
+import { ref, onMounted, onUnmounted, CSSProperties, watch } from 'vue'
+import { app } from '@/scripts/app'
+import { LGraphGroup, LGraphNode, LiteGraph } from '@comfyorg/litegraph'
+import EditableText from '@/components/common/EditableText.vue'
+import { ComfyExtension } from '@/types/comfy'
+import { useSettingStore } from '@/stores/settingStore'
+import type { LiteGraphCanvasEvent } from '@comfyorg/litegraph'
+import { useTitleEditorStore } from '@/stores/graphStore'
+
+const settingStore = useSettingStore()
+
+const showInput = ref(false)
+const editedTitle = ref('')
+const inputStyle = ref<CSSProperties>({
+  position: 'fixed',
+  left: '0px',
+  top: '0px',
+  width: '200px',
+  height: '20px',
+  fontSize: '12px'
+})
+
+const titleEditorStore = useTitleEditorStore()
+
+const onEdit = (newValue: string) => {
+  if (titleEditorStore.titleEditorTarget && newValue.trim() !== '') {
+    titleEditorStore.titleEditorTarget.title = newValue.trim()
+    app.graph.setDirtyCanvas(true, true)
+  }
+  showInput.value = false
+  titleEditorStore.titleEditorTarget = null
+}
+
+watch(
+  () => titleEditorStore.titleEditorTarget,
+  (target) => {
+    if (target === null) {
+      return
+    }
+    editedTitle.value = target.title
+    showInput.value = true
+
+    if (target instanceof LGraphGroup) {
+      const group = target
+      const [x, y] = group.pos
+      const [w, h] = group.size
+
+      const [left, top] = app.canvasPosToClientPos([x, y])
+      inputStyle.value.left = `${left}px`
+      inputStyle.value.top = `${top}px`
+
+      const width = w * app.canvas.ds.scale
+      const height = group.titleHeight * app.canvas.ds.scale
+      inputStyle.value.width = `${width}px`
+      inputStyle.value.height = `${height}px`
+
+      const fontSize = group.font_size * app.canvas.ds.scale
+      inputStyle.value.fontSize = `${fontSize}px`
+    } else if (target instanceof LGraphNode) {
+      const node = target
+      const isCollapsed = node.flags?.collapsed
+      const [x, y, nodeWidth, nodeHeight] = node.getBounding()
+      const canvasWidth =
+        // @ts-expect-error Remove after collapsed_width is exposed in LiteGraph
+        isCollapsed && node._collapsed_width ? node._collapsed_width : nodeWidth
+      const canvasHeight = LiteGraph.NODE_TITLE_HEIGHT
+
+      const [left, top] = app.canvasPosToClientPos([x, y])
+      inputStyle.value.left = `${left}px`
+      inputStyle.value.top = `${top}px`
+
+      const width = canvasWidth * app.canvas.ds.scale
+      const height = canvasHeight * app.canvas.ds.scale
+      inputStyle.value.width = `${width}px`
+      inputStyle.value.height = `${height}px`
+      const fontSize = 12 * app.canvas.ds.scale
+      inputStyle.value.fontSize = `${fontSize}px`
+    }
+  }
+)
+
+const canvasEventHandler = (event: LiteGraphCanvasEvent) => {
+  if (!settingStore.get('Comfy.Group.DoubleClickTitleToEdit')) {
+    return
+  }
+
+  if (event.detail.subType === 'group-double-click') {
+    const group: LGraphGroup = event.detail.group
+    const [x, y] = group.pos
+
+    const e = event.detail.originalEvent
+    // @ts-expect-error LiteGraphCanvasEvent is not typed
+    const relativeY = e.canvasY - y
+    // Only allow editing if the click is on the title bar
+    if (relativeY > group.titleHeight) {
+      return
+    }
+
+    titleEditorStore.titleEditorTarget = group
+  }
+}
+
+const extension: ComfyExtension = {
+  name: 'Comfy.NodeTitleEditor',
+  nodeCreated(node: LGraphNode) {
+    // Store the original callback
+    const originalCallback = node.onNodeTitleDblClick
+
+    node.onNodeTitleDblClick = function (e: MouseEvent, ...args: any[]) {
+      if (!settingStore.get('Comfy.Node.DoubleClickTitleToEdit')) {
+        return
+      }
+
+      titleEditorStore.titleEditorTarget = this
+
+      // Call the original callback if it exists
+      if (typeof originalCallback === 'function') {
+        originalCallback.call(this, e, ...args)
+      }
+    }
+  }
+}
+
+onMounted(() => {
+  document.addEventListener('litegraph:canvas', canvasEventHandler)
+  app.registerExtension(extension)
+})
+
+onUnmounted(() => {
+  document.removeEventListener('litegraph:canvas', canvasEventHandler)
+})
+</script>
+
+<style scoped>
+.group-title-editor.node-title-editor {
+  z-index: 9999;
+  padding: 0.25rem;
+}
+
+:deep(.editable-text) {
+  width: 100%;
+  height: 100%;
+}
+
+:deep(.editable-text input) {
+  width: 100%;
+  height: 100%;
+  /* Override the default font size */
+  font-size: inherit;
+}
+</style>
diff --git a/src/scripts/ui.ts b/src/scripts/ui.ts
index 18f5e07b7..625d61463 100644
--- a/src/scripts/ui.ts
+++ b/src/scripts/ui.ts
@@ -8,6 +8,7 @@ import { showSettingsDialog } from '@/services/dialogService'
 import { useToastStore } from '@/stores/toastStore'
 import { LGraphGroup } from '@comfyorg/litegraph'
 import { useSettingStore } from '@/stores/settingStore'
+import { useTitleEditorStore } from '@/stores/graphStore'
 
 export const ComfyDialog = _ComfyDialog
 
@@ -783,6 +784,7 @@ export class ComfyUI {
           )
           group.addNodes(Object.values(app.canvas.selected_nodes), padding)
           app.canvas.graph.add(group)
+          useTitleEditorStore().titleEditorTarget = group
         }
       })
     ]) as HTMLDivElement
diff --git a/src/stores/graphStore.ts b/src/stores/graphStore.ts
new file mode 100644
index 000000000..74b03d82a
--- /dev/null
+++ b/src/stores/graphStore.ts
@@ -0,0 +1,11 @@
+import { LGraphNode, LGraphGroup } from '@comfyorg/litegraph'
+import { defineStore } from 'pinia'
+import { ref } from 'vue'
+
+export const useTitleEditorStore = defineStore('titleEditor', () => {
+  const titleEditorTarget = ref<LGraphNode | LGraphGroup | null>(null)
+
+  return {
+    titleEditorTarget
+  }
+})