From f75f774ddb2e213d9b7b6aa0f7224ce5ce7435e8 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Tue, 1 Oct 2024 05:06:43 +0900
Subject: [PATCH] Graph canvas menu (#1023)

* add graph canvas menu

* Move to corner

* Remove action bar reset zoom button

* nit

* Add setting

---------

Co-authored-by: huchenlei <chenlei.hu@mail.utoronto.ca>
---
 .../LiteGraphCanvasSplitterOverlay.vue        |  4 +-
 src/components/appMenu/AppMenu.vue            |  6 --
 src/components/graph/GraphCanvas.vue          |  7 ++
 src/components/graph/GraphCanvasMenu.vue      | 81 +++++++++++++++++++
 src/i18n.ts                                   |  7 ++
 src/scripts/domWidget.ts                      |  3 +-
 src/stores/commandStore.ts                    | 28 ++++++-
 src/stores/coreSettings.ts                    |  7 +-
 src/stores/graphStore.ts                      | 15 +++-
 9 files changed, 145 insertions(+), 13 deletions(-)
 create mode 100644 src/components/graph/GraphCanvasMenu.vue

diff --git a/src/components/LiteGraphCanvasSplitterOverlay.vue b/src/components/LiteGraphCanvasSplitterOverlay.vue
index f91663fd8..60e04b5a0 100644
--- a/src/components/LiteGraphCanvasSplitterOverlay.vue
+++ b/src/components/LiteGraphCanvasSplitterOverlay.vue
@@ -9,8 +9,8 @@
     >
       <slot name="side-bar-panel"></slot>
     </SplitterPanel>
-    <SplitterPanel class="graph-canvas-panel" :size="100">
-      <div></div>
+    <SplitterPanel class="graph-canvas-panel relative" :size="100">
+      <slot name="graph-canvas-panel"></slot>
     </SplitterPanel>
     <SplitterPanel
       class="side-bar-panel"
diff --git a/src/components/appMenu/AppMenu.vue b/src/components/appMenu/AppMenu.vue
index ac5c189c8..6aaf59dc5 100644
--- a/src/components/appMenu/AppMenu.vue
+++ b/src/components/appMenu/AppMenu.vue
@@ -61,12 +61,6 @@
               commandStore.getCommandFunction('Comfy.RefreshNodeDefinitions')()
           "
         />
-        <Button
-          v-tooltip.bottom="$t('menu.resetView')"
-          icon="pi pi-expand"
-          severity="secondary"
-          @click="() => commandStore.getCommandFunction('Comfy.ResetView')()"
-        />
       </ButtonGroup>
     </div>
   </Panel>
diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue
index 630b1da2c..731e132be 100644
--- a/src/components/graph/GraphCanvas.vue
+++ b/src/components/graph/GraphCanvas.vue
@@ -4,6 +4,9 @@
       <template #side-bar-panel>
         <SideToolbar />
       </template>
+      <template #graph-canvas-panel>
+        <GraphCanvasMenu v-if="canvasMenuEnabled" />
+      </template>
     </LiteGraphCanvasSplitterOverlay>
     <TitleEditor />
     <canvas ref="canvasRef" id="graph-canvas" tabindex="1" />
@@ -44,6 +47,7 @@ import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
 import { useCanvasStore } from '@/stores/graphStore'
 import { ComfyModelDef } from '@/stores/modelStore'
 import { useModelToNodeStore } from '@/stores/modelToNodeStore'
+import GraphCanvasMenu from '@/components/graph/GraphCanvasMenu.vue'
 
 const emit = defineEmits(['ready'])
 const canvasRef = ref<HTMLCanvasElement | null>(null)
@@ -55,6 +59,9 @@ const modelToNodeStore = useModelToNodeStore()
 const betaMenuEnabled = computed(
   () => settingStore.get('Comfy.UseNewMenu') !== 'Disabled'
 )
+const canvasMenuEnabled = computed(() =>
+  settingStore.get('Comfy.Graph.CanvasMenu')
+)
 
 watchEffect(() => {
   const canvasInfoEnabled = settingStore.get('Comfy.Graph.CanvasInfo')
diff --git a/src/components/graph/GraphCanvasMenu.vue b/src/components/graph/GraphCanvasMenu.vue
new file mode 100644
index 000000000..37e674bad
--- /dev/null
+++ b/src/components/graph/GraphCanvasMenu.vue
@@ -0,0 +1,81 @@
+<template>
+  <ButtonGroup
+    class="p-buttongroup-vertical absolute bottom-[10px] right-[10px] z-[1000] pointer-events-auto"
+  >
+    <Button
+      severity="secondary"
+      icon="pi pi-plus"
+      v-tooltip.left="t('graphCanvasMenu.zoomIn')"
+      @mousedown="repeat('Comfy.Canvas.ZoomIn')"
+      @mouseup="stopRepeat"
+    />
+    <Button
+      severity="secondary"
+      icon="pi pi-minus"
+      v-tooltip.left="t('graphCanvasMenu.zoomOut')"
+      @mousedown="repeat('Comfy.Canvas.ZoomOut')"
+      @mouseup="stopRepeat"
+    />
+    <Button
+      severity="secondary"
+      icon="pi pi-expand"
+      v-tooltip.left="t('graphCanvasMenu.resetView')"
+      @click="() => commandStore.getCommandFunction('Comfy.Canvas.ResetView')()"
+    />
+    <Button
+      severity="secondary"
+      v-tooltip.left="
+        t('graphCanvasMenu.' + (canvasStore.readOnly ? 'unlock' : 'lock'))
+      "
+      @click="
+        () => commandStore.getCommandFunction('Comfy.Canvas.ToggleLock')()
+      "
+    >
+      <template #icon>
+        <i-material-symbols:lock-outline v-if="canvasStore.readOnly" />
+        <i-material-symbols:lock-open-outline v-else />
+      </template>
+    </Button>
+  </ButtonGroup>
+</template>
+
+<script setup lang="ts">
+import ButtonGroup from 'primevue/buttongroup'
+import Button from 'primevue/button'
+import { useCommandStore } from '@/stores/commandStore'
+import { useCanvasStore } from '@/stores/graphStore'
+import { useI18n } from 'vue-i18n'
+
+const { t } = useI18n()
+const commandStore = useCommandStore()
+const canvasStore = useCanvasStore()
+
+let interval: number | null = null
+const repeat = (command: string) => {
+  if (interval) return
+  const cmd = commandStore.getCommandFunction(command)
+  cmd()
+  interval = window.setInterval(cmd, 100)
+}
+const stopRepeat = () => {
+  if (interval) {
+    clearInterval(interval)
+    interval = null
+  }
+}
+</script>
+
+<style scoped>
+.p-buttongroup-vertical {
+  display: flex;
+  flex-direction: column;
+  border-radius: var(--p-button-border-radius);
+  overflow: hidden;
+  border: 1px solid var(--p-panel-border-color);
+}
+
+.p-buttongroup-vertical .p-button {
+  margin: 0;
+  border-radius: 0;
+}
+</style>
diff --git a/src/i18n.ts b/src/i18n.ts
index 5267f8ec0..aa879d8d8 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -94,6 +94,13 @@ const messages = {
         upscale: '2 Pass Upscale',
         flux_schnell: 'Flux Schnell'
       }
+    },
+    graphCanvasMenu: {
+      zoomIn: 'Zoom In',
+      zoomOut: 'Zoom Out',
+      resetView: 'Reset View',
+      lock: 'Lock Graph',
+      unlock: 'Unlock Graph'
     }
   },
   zh: {
diff --git a/src/scripts/domWidget.ts b/src/scripts/domWidget.ts
index fbc89805e..fb1a2cb53 100644
--- a/src/scripts/domWidget.ts
+++ b/src/scripts/domWidget.ts
@@ -354,7 +354,8 @@ LGraphNode.prototype.addDOMWidget = function (
         width: `${widgetWidth - margin * 2}px`,
         height: `${(widget.computedHeight ?? 50) - margin * 2}px`,
         position: 'absolute',
-        zIndex: app.graph.nodes.indexOf(node)
+        zIndex: app.graph.nodes.indexOf(node),
+        pointerEvents: app.canvas.read_only ? 'none' : 'auto'
       })
 
       if (enableDomClipping) {
diff --git a/src/stores/commandStore.ts b/src/stores/commandStore.ts
index 6ea1dcf50..d942f2f9f 100644
--- a/src/stores/commandStore.ts
+++ b/src/stores/commandStore.ts
@@ -124,7 +124,7 @@ export const useCommandStore = defineStore('command', () => {
       }
     },
     {
-      id: 'Comfy.ResetView',
+      id: 'Comfy.Canvas.ResetView',
       icon: 'pi pi-expand',
       label: 'Reset View',
       function: () => {
@@ -180,6 +180,32 @@ export const useCommandStore = defineStore('command', () => {
       icon: 'pi pi-folder-open',
       label: 'Browse Templates',
       function: showTemplateWorkflowsDialog
+    },
+    {
+      id: 'Comfy.Canvas.ZoomIn',
+      icon: 'pi pi-plus',
+      label: 'Zoom In',
+      function: () => {
+        app.canvas.ds.changeScale(app.canvas.ds.scale + 0.1)
+        app.canvas.setDirty(true, true)
+      }
+    },
+    {
+      id: 'Comfy.Canvas.ZoomOut',
+      icon: 'pi pi-minus',
+      label: 'Zoom Out',
+      function: () => {
+        app.canvas.ds.changeScale(app.canvas.ds.scale - 0.1)
+        app.canvas.setDirty(true, true)
+      }
+    },
+    {
+      id: 'Comfy.Canvas.ToggleLock',
+      icon: 'pi pi-lock',
+      label: 'Toggle Lock',
+      function: () => {
+        app.canvas['read_only'] = !app.canvas['read_only']
+      }
     }
   ]
 
diff --git a/src/stores/coreSettings.ts b/src/stores/coreSettings.ts
index ef4e1ae25..acdc4bd13 100644
--- a/src/stores/coreSettings.ts
+++ b/src/stores/coreSettings.ts
@@ -12,7 +12,6 @@ export const CORE_SETTINGS: SettingParams[] = [
     type: 'boolean',
     defaultValue: true
   },
-
   {
     id: 'Comfy.NodeSearchBoxImpl',
     category: ['Comfy', 'Node Search Box', 'Implementation'],
@@ -381,5 +380,11 @@ export const CORE_SETTINGS: SettingParams[] = [
     type: 'hidden',
     options: ['Sidebar', 'Topbar'],
     defaultValue: 'Sidebar'
+  },
+  {
+    id: 'Comfy.Graph.CanvasMenu',
+    name: 'Show graph canvas menu',
+    type: 'boolean',
+    defaultValue: true
   }
 ]
diff --git a/src/stores/graphStore.ts b/src/stores/graphStore.ts
index 7eaa36997..6533d11a7 100644
--- a/src/stores/graphStore.ts
+++ b/src/stores/graphStore.ts
@@ -1,6 +1,6 @@
 import { LGraphNode, LGraphGroup, LGraphCanvas } from '@comfyorg/litegraph'
 import { defineStore } from 'pinia'
-import { shallowRef } from 'vue'
+import { ref, shallowRef } from 'vue'
 
 export const useTitleEditorStore = defineStore('titleEditor', () => {
   const titleEditorTarget = shallowRef<LGraphNode | LGraphGroup | null>(null)
@@ -12,8 +12,19 @@ export const useTitleEditorStore = defineStore('titleEditor', () => {
 
 export const useCanvasStore = defineStore('canvas', () => {
   const canvas = shallowRef<LGraphCanvas | null>(null)
+  const readOnly = ref(false)
+
+  document.addEventListener(
+    'litegraph:canvas',
+    (e: CustomEvent<{ subType: string; readOnly: boolean }>) => {
+      if (e.detail?.subType === 'read-only') {
+        readOnly.value = e.detail.readOnly
+      }
+    }
+  )
 
   return {
-    canvas
+    canvas,
+    readOnly
   }
 })