From 4ab3aa9a390e8f5cc65042fdea7565b615fddd7e Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Wed, 25 Sep 2024 15:41:05 +0900
Subject: [PATCH] Backward compatibility with extension injections on legacy
 menu bar (#970)

* Compatible to legacy top menu extensions

* Rework css

* nit
---
 src/components/topbar/TopMenubar.vue   |  27 +++--
 src/components/topbar/WorkflowTabs.vue |  76 ++++++--------
 src/scripts/ui/menu/index.ts           | 131 +++----------------------
 src/scripts/ui/menu/interruptButton.ts |  32 ------
 src/scripts/ui/menu/menu.css           |   3 -
 src/scripts/ui/menu/queueButton.ts     | 103 -------------------
 src/scripts/ui/menu/queueOptions.ts    |  82 ----------------
 src/views/GraphView.vue                |   4 +-
 8 files changed, 68 insertions(+), 390 deletions(-)
 delete mode 100644 src/scripts/ui/menu/interruptButton.ts
 delete mode 100644 src/scripts/ui/menu/queueButton.ts
 delete mode 100644 src/scripts/ui/menu/queueOptions.ts

diff --git a/src/components/topbar/TopMenubar.vue b/src/components/topbar/TopMenubar.vue
index 5ec206922..b4087727b 100644
--- a/src/components/topbar/TopMenubar.vue
+++ b/src/components/topbar/TopMenubar.vue
@@ -1,16 +1,23 @@
 <template>
   <teleport to=".comfyui-body-top">
-    <div class="top-menubar comfyui-menu" v-if="betaMenuEnabled">
+    <div
+      class="top-menubar comfyui-menu flex items-center"
+      v-show="betaMenuEnabled"
+    >
       <h1 class="comfyui-logo mx-2">ComfyUI</h1>
       <Menubar
         :model="items"
         class="border-none p-0 bg-transparent"
         :pt="{
-          rootList: 'gap-0'
+          rootList: 'gap-0 flex-nowrap'
         }"
       />
       <Divider layout="vertical" class="mx-2" />
-      <WorkflowTabs v-if="workflowTabsPosition === 'Topbar'" />
+      <WorkflowTabs
+        v-if="workflowTabsPosition === 'Topbar'"
+        class="flex-grow"
+      />
+      <div class="comfyui-menu-right" ref="menuRight"></div>
     </div>
   </teleport>
 </template>
@@ -20,8 +27,9 @@ import Menubar from 'primevue/menubar'
 import Divider from 'primevue/divider'
 import WorkflowTabs from '@/components/topbar/WorkflowTabs.vue'
 import { useCoreMenuItemStore } from '@/stores/coreMenuItemStore'
-import { computed } from 'vue'
+import { computed, onMounted, ref } from 'vue'
 import { useSettingStore } from '@/stores/settingStore'
+import { app } from '@/scripts/app'
 
 const settingStore = useSettingStore()
 const workflowTabsPosition = computed(() =>
@@ -32,6 +40,14 @@ const betaMenuEnabled = computed(
 )
 const coreMenuItemsStore = useCoreMenuItemStore()
 const items = coreMenuItemsStore.menuItems
+
+const menuRight = ref<HTMLDivElement | null>(null)
+// Menu-right holds legacy topbar elements attached by custom scripts
+onMounted(() => {
+  if (menuRight.value) {
+    menuRight.value.appendChild(app.menu.element)
+  }
+})
 </script>
 
 <style scoped>
@@ -41,13 +57,10 @@ const items = coreMenuItemsStore.menuItems
   color: var(--fg-color);
   font-family: Arial, Helvetica, sans-serif;
   font-size: 0.8em;
-  display: flex;
-  align-items: center;
   box-sizing: border-box;
   z-index: 1000;
   order: 0;
   grid-column: 1/-1;
-  overflow: auto;
   max-height: 90vh;
 }
 
diff --git a/src/components/topbar/WorkflowTabs.vue b/src/components/topbar/WorkflowTabs.vue
index f92474544..50fd485df 100644
--- a/src/components/topbar/WorkflowTabs.vue
+++ b/src/components/topbar/WorkflowTabs.vue
@@ -1,32 +1,31 @@
 <template>
-  <div class="workflow-tabs">
-    <SelectButton
-      class="select-button-group bg-transparent"
-      :modelValue="selectedWorkflow"
-      @update:modelValue="onWorkflowChange"
-      :options="options"
-      optionLabel="label"
-      dataKey="value"
-    >
-      <template #option="{ option }">
-        <span
-          class="workflow-label text-sm max-w-[150px] truncate inline-block"
-          >{{ option.label }}</span
-        >
-        <div class="relative">
-          <span class="status-indicator" v-if="option.unsaved">•</span>
-          <Button
-            class="close-button p-0 w-auto"
-            icon="pi pi-times"
-            text
-            severity="secondary"
-            size="small"
-            @click.stop="onCloseWorkflow(option)"
-          />
-        </div>
-      </template>
-    </SelectButton>
-  </div>
+  <SelectButton
+    class="workflow-tabs bg-transparent flex flex-wrap"
+    :class="props.class"
+    :modelValue="selectedWorkflow"
+    @update:modelValue="onWorkflowChange"
+    :options="options"
+    optionLabel="label"
+    dataKey="value"
+  >
+    <template #option="{ option }">
+      <span
+        class="workflow-label text-sm max-w-[150px] truncate inline-block"
+        >{{ option.label }}</span
+      >
+      <div class="relative">
+        <span class="status-indicator" v-if="option.unsaved">•</span>
+        <Button
+          class="close-button p-0 w-auto"
+          icon="pi pi-times"
+          text
+          severity="secondary"
+          size="small"
+          @click.stop="onCloseWorkflow(option)"
+        />
+      </div>
+    </template>
+  </SelectButton>
 </template>
 
 <script setup lang="ts">
@@ -37,6 +36,10 @@ import SelectButton from 'primevue/selectbutton'
 import Button from 'primevue/button'
 import { computed } from 'vue'
 
+const props = defineProps<{
+  class?: string
+}>()
+
 const workflowStore = useWorkflowStore()
 interface WorkflowOption {
   label: string
@@ -79,26 +82,9 @@ const onCloseWorkflow = (option: WorkflowOption) => {
   const workflow = optionToWorkflow(option)
   app.workflowManager.closeWorkflow(workflow)
 }
-
-// Add this new function to check if a workflow is unsaved
-const isWorkflowUnsaved = (option: WorkflowOption): boolean => {
-  const workflow = optionToWorkflow(option)
-  return workflow.unsaved
-}
 </script>
 
 <style scoped>
-.select-button-group {
-  /* TODO: Make this dynamic. Take rest of space after all tool buttons */
-  max-width: 70vw;
-  overflow-x: auto;
-  overflow-y: hidden;
-
-  /* Scrollbar styling */
-  scrollbar-width: thin;
-  scrollbar-color: rgba(155, 155, 155, 0.5) transparent;
-}
-
 :deep(.p-togglebutton::before) {
   @apply hidden;
 }
diff --git a/src/scripts/ui/menu/index.ts b/src/scripts/ui/menu/index.ts
index bbca3ce47..66c068189 100644
--- a/src/scripts/ui/menu/index.ts
+++ b/src/scripts/ui/menu/index.ts
@@ -1,140 +1,37 @@
 import type { ComfyApp } from '@/scripts/app'
-import { api } from '../../api'
 import { $el } from '../../ui'
 import { downloadBlob } from '../../utils'
-import { ComfyButton } from '../components/button'
 import { ComfyButtonGroup } from '../components/buttonGroup'
-import { ComfySplitButton } from '../components/splitButton'
-import { ComfyQueueButton } from './queueButton'
-import { getInterruptButton } from './interruptButton'
 import './menu.css'
 
-const collapseOnMobile = (t) => {
-  ;(t.element ?? t).classList.add('comfyui-menu-mobile-collapse')
-  return t
-}
-const showOnMobile = (t) => {
-  ;(t.element ?? t).classList.add('lt-lg-show')
-  return t
-}
+// Import ComfyButton to make sure it's shimmed and exported by vite
+import { ComfyButton } from '../components/button'
+import { ComfySplitButton } from '../components/splitButton'
+import { ComfyPopup } from '../components/popup'
+console.debug(
+  `Keep following definitions ${ComfyButton} ${ComfySplitButton} ${ComfyPopup}`
+)
 
 export class ComfyAppMenu {
   app: ComfyApp
-  logo: HTMLElement
-  saveButton: ComfySplitButton
   actionsGroup: ComfyButtonGroup
   settingsGroup: ComfyButtonGroup
   viewGroup: ComfyButtonGroup
-  mobileMenuButton: ComfyButton
-  queueButton: ComfyQueueButton
   element: HTMLElement
 
   constructor(app: ComfyApp) {
     this.app = app
-    const getSaveButton = (t?: string) =>
-      new ComfyButton({
-        icon: 'content-save',
-        tooltip: 'Save the current workflow',
-        action: () => app.workflowManager.activeWorkflow.save(),
-        content: t
-      })
 
-    this.logo = $el('h1.comfyui-logo.nlg-hide', { title: 'ComfyUI' }, 'ComfyUI')
-    this.saveButton = new ComfySplitButton(
-      {
-        primary: getSaveButton(),
-        mode: 'hover',
-        position: 'absolute'
-      },
-      getSaveButton('Save'),
-      new ComfyButton({
-        icon: 'content-save-edit',
-        content: 'Save As',
-        tooltip: 'Save the current graph as a new workflow',
-        action: () => app.workflowManager.activeWorkflow.save(true)
-      }),
-      new ComfyButton({
-        icon: 'download',
-        content: 'Export',
-        tooltip: 'Export the current workflow as JSON',
-        action: () => this.exportWorkflow('workflow', 'workflow')
-      }),
-      new ComfyButton({
-        icon: 'api',
-        content: 'Export (API Format)',
-        tooltip:
-          'Export the current workflow as JSON for use with the ComfyUI API',
-        action: () => this.exportWorkflow('workflow_api', 'output'),
-        visibilitySetting: { id: 'Comfy.DevMode', showValue: true },
-        app
-      })
-    )
-
-    const actionButtons = [
-      new ComfyButton({
-        icon: 'refresh',
-        content: 'Refresh',
-        tooltip: 'Refresh widgets in nodes to find new models or files',
-        action: () => app.refreshComboInNodes()
-      }),
-      new ComfyButton({
-        icon: 'clipboard-edit-outline',
-        content: 'Clipspace',
-        tooltip: 'Open Clipspace window',
-        action: () => app['openClipspace']()
-      }),
-      new ComfyButton({
-        icon: 'fit-to-page-outline',
-        content: 'Reset View',
-        tooltip: 'Reset the canvas view',
-        action: () => app.resetView()
-      }),
-      new ComfyButton({
-        icon: 'cancel',
-        content: 'Clear',
-        tooltip: 'Clears current workflow',
-        action: () => {
-          if (
-            !app.ui.settings.getSettingValue('Comfy.ConfirmClear', true) ||
-            confirm('Clear workflow?')
-          ) {
-            app.clean()
-            app.graph.clear()
-            api.dispatchEvent(new CustomEvent('graphCleared'))
-          }
-        }
-      })
-    ]
-    this.actionsGroup = new ComfyButtonGroup(...actionButtons)
-
-    // Keep the settings group as there are custom scripts attaching extra
+    // Keep the group as there are custom scripts attaching extra
     // elements to it.
+    this.actionsGroup = new ComfyButtonGroup()
     this.settingsGroup = new ComfyButtonGroup()
-    const interruptButton = getInterruptButton('nlg-hide').element
-    this.viewGroup = new ComfyButtonGroup(interruptButton)
-    this.mobileMenuButton = new ComfyButton({
-      icon: 'menu',
-      action: (_, btn) => {
-        btn.icon = this.element.classList.toggle('expanded')
-          ? 'menu-open'
-          : 'menu'
-        window.dispatchEvent(new Event('resize'))
-      },
-      classList: 'comfyui-button comfyui-menu-button'
-    })
-    this.queueButton = new ComfyQueueButton(app)
+    this.viewGroup = new ComfyButtonGroup()
 
-    this.element = $el('nav.comfyui-menu.lg', { style: { display: 'none' } }, [
-      this.logo,
-      this.saveButton.element,
-      collapseOnMobile(this.actionsGroup).element,
-      $el('section.comfyui-menu-push'),
-      collapseOnMobile(this.settingsGroup).element,
-      collapseOnMobile(this.viewGroup).element,
-
-      getInterruptButton('lt-lg-show').element,
-      this.queueButton.element,
-      showOnMobile(this.mobileMenuButton).element
+    this.element = $el('div.flex.gap-2.mx-2', [
+      this.actionsGroup.element,
+      this.settingsGroup.element,
+      this.viewGroup.element
     ])
   }
 
diff --git a/src/scripts/ui/menu/interruptButton.ts b/src/scripts/ui/menu/interruptButton.ts
deleted file mode 100644
index 4d48af655..000000000
--- a/src/scripts/ui/menu/interruptButton.ts
+++ /dev/null
@@ -1,32 +0,0 @@
-import { StatusWsMessageStatus } from '@/types/apiTypes'
-import { api } from '../../api'
-import { ComfyButton } from '../components/button'
-import { useToastStore } from '@/stores/toastStore'
-
-export function getInterruptButton(visibility: string) {
-  const btn = new ComfyButton({
-    icon: 'close',
-    tooltip: 'Cancel current generation',
-    enabled: false,
-    action: async () => {
-      await api.interrupt()
-      useToastStore().add({
-        severity: 'info',
-        summary: 'Interrupted',
-        detail: 'Execution has been interrupted',
-        life: 1000
-      })
-    },
-    classList: ['comfyui-button', 'comfyui-interrupt-button', visibility]
-  })
-
-  api.addEventListener(
-    'status',
-    ({ detail }: CustomEvent<StatusWsMessageStatus>) => {
-      const sz = detail?.exec_info?.queue_remaining
-      btn.enabled = sz > 0
-    }
-  )
-
-  return btn
-}
diff --git a/src/scripts/ui/menu/menu.css b/src/scripts/ui/menu/menu.css
index e66d182cb..2c9cc1d7f 100644
--- a/src/scripts/ui/menu/menu.css
+++ b/src/scripts/ui/menu/menu.css
@@ -122,9 +122,6 @@
 }
 
 /* Menu */
-.comfyui-menu>* {
-  flex-shrink: 0;
-}
 .comfyui-menu .mdi::before {
   font-size: 18px;
 }
diff --git a/src/scripts/ui/menu/queueButton.ts b/src/scripts/ui/menu/queueButton.ts
deleted file mode 100644
index 19da770d2..000000000
--- a/src/scripts/ui/menu/queueButton.ts
+++ /dev/null
@@ -1,103 +0,0 @@
-import { ComfyButton } from '../components/button'
-import { $el } from '../../ui'
-import { api } from '../../api'
-import { ComfySplitButton } from '../components/splitButton'
-import { ComfyQueueOptions } from './queueOptions'
-import { prop } from '../../utils'
-import type { ComfyApp } from '@/scripts/app'
-import { StatusWsMessageStatus } from '@/types/apiTypes'
-
-export class ComfyQueueButton {
-  element = $el('div.comfyui-queue-button')
-  #internalQueueSize = 0
-
-  queuePrompt = async (e?: MouseEvent) => {
-    this.#internalQueueSize += this.queueOptions.batchCount
-    // Hold shift to queue front, event is undefined when auto-queue is enabled
-    await this.app.queuePrompt(
-      e?.shiftKey ? -1 : 0,
-      this.queueOptions.batchCount
-    )
-  }
-  queueOptions: ComfyQueueOptions
-  app: ComfyApp
-  autoQueueMode: string
-  graphHasChanged: boolean
-
-  constructor(app: ComfyApp) {
-    this.app = app
-
-    const queue = new ComfyButton({
-      content: $el('div', [
-        $el('span', {
-          textContent: 'Queue'
-        })
-      ]),
-      icon: 'play',
-      classList: 'comfyui-button',
-      action: this.queuePrompt
-    })
-
-    this.queueOptions = new ComfyQueueOptions(app)
-
-    const btn = new ComfySplitButton(
-      {
-        primary: queue,
-        mode: 'click',
-        position: 'absolute',
-        horizontal: 'right'
-      },
-      this.queueOptions.element
-    )
-    btn.element.classList.add('primary')
-    this.element.append(btn.element)
-
-    this.autoQueueMode = prop(this, 'autoQueueMode', '', () => {
-      switch (this.autoQueueMode) {
-        case 'instant':
-          queue.icon = 'infinity'
-          break
-        case 'change':
-          queue.icon = 'auto-mode'
-          break
-        default:
-          queue.icon = 'play'
-          break
-      }
-    })
-
-    this.queueOptions.addEventListener(
-      'autoQueueMode',
-      (e) => (this.autoQueueMode = e['detail'])
-    )
-
-    api.addEventListener('graphChanged', () => {
-      if (this.autoQueueMode === 'change') {
-        if (this.#internalQueueSize) {
-          this.graphHasChanged = true
-        } else {
-          this.graphHasChanged = false
-          this.queuePrompt()
-        }
-      }
-    })
-
-    api.addEventListener(
-      'status',
-      ({ detail }: CustomEvent<StatusWsMessageStatus>) => {
-        this.#internalQueueSize = detail?.exec_info?.queue_remaining
-        if (this.#internalQueueSize != null) {
-          if (!this.#internalQueueSize && !app.lastExecutionError) {
-            if (
-              this.autoQueueMode === 'instant' ||
-              (this.autoQueueMode === 'change' && this.graphHasChanged)
-            ) {
-              this.graphHasChanged = false
-              this.queuePrompt()
-            }
-          }
-        }
-      }
-    )
-  }
-}
diff --git a/src/scripts/ui/menu/queueOptions.ts b/src/scripts/ui/menu/queueOptions.ts
deleted file mode 100644
index 00e81f191..000000000
--- a/src/scripts/ui/menu/queueOptions.ts
+++ /dev/null
@@ -1,82 +0,0 @@
-import type { ComfyApp } from '@/scripts/app'
-import { $el } from '../../ui'
-import { prop } from '../../utils'
-
-export class ComfyQueueOptions extends EventTarget {
-  element = $el('div.comfyui-queue-options')
-  app: ComfyApp
-  batchCountInput: HTMLInputElement
-  batchCount: number
-  batchCountRange: HTMLInputElement
-  autoQueueMode: string
-  autoQueueEl: HTMLElement
-
-  constructor(app: ComfyApp) {
-    super()
-    this.app = app
-
-    this.batchCountInput = $el('input', {
-      className: 'comfyui-queue-batch-value',
-      type: 'number',
-      min: '1',
-      value: '1',
-      oninput: () => (this.batchCount = +this.batchCountInput.value)
-    })
-
-    this.batchCountRange = $el('input', {
-      type: 'range',
-      min: '1',
-      max: '100',
-      value: '1',
-      oninput: () => (this.batchCount = +this.batchCountRange.value)
-    })
-
-    this.element.append(
-      $el('div.comfyui-queue-batch', [
-        $el(
-          'label',
-          {
-            textContent: 'Batch count: '
-          },
-          this.batchCountInput
-        ),
-        this.batchCountRange
-      ])
-    )
-
-    const createOption = (text, value, checked = false) =>
-      $el(
-        'label',
-        { textContent: text },
-        $el('input', {
-          type: 'radio',
-          name: 'AutoQueueMode',
-          checked,
-          value,
-          oninput: (e) => (this.autoQueueMode = e.target['value'])
-        })
-      )
-
-    this.autoQueueEl = $el('div.comfyui-queue-mode', [
-      $el('span', 'Auto Queue:'),
-      createOption('Disabled', '', true),
-      createOption('Instant', 'instant'),
-      createOption('On Change', 'change')
-    ])
-
-    this.element.append(this.autoQueueEl)
-
-    this.batchCount = prop(this, 'batchCount', 1, () => {
-      this.batchCountInput.value = this.batchCount + ''
-      this.batchCountRange.value = this.batchCount + ''
-    })
-
-    this.autoQueueMode = prop(this, 'autoQueueMode', 'Disabled', () => {
-      this.dispatchEvent(
-        new CustomEvent('autoQueueMode', {
-          detail: this.autoQueueMode
-        })
-      )
-    })
-  }
-}
diff --git a/src/views/GraphView.vue b/src/views/GraphView.vue
index 4ce9f48d2..47b4a07ac 100644
--- a/src/views/GraphView.vue
+++ b/src/views/GraphView.vue
@@ -1,10 +1,12 @@
 <template>
+  <!-- Top menu bar needs to load before the GraphCanvas as it needs to host
+  the menu buttons added by legacy extension scripts.-->
+  <TopMenubar />
   <GraphCanvas />
   <GlobalToast />
   <UnloadWindowConfirmDialog />
   <BrowserTabTitle />
   <AppMenu />
-  <TopMenubar />
 </template>
 
 <script setup lang="ts">