From 5ee0fd3519d54c909ecde74e6e2b35481d246985 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Wed, 25 Sep 2024 11:55:50 +0900
Subject: [PATCH] Rework queue button (#968)

* Move queue button to right side

* Rework split button

* Group

* Remove unused code

* x2 buttons

* Use primevue divider

* adjust style

* Add tooltip

* Update test

* Add clearing pending tasks button to queue bar

* Fix state

* Dropdown list fix
---
 browser_tests/appMenu.spec.ts                 |   6 +-
 browser_tests/helpers/appMenu.ts              |  39 +--
 src/components/appMenu/AppMenu.vue            | 224 +++++++-----------
 src/components/appMenu/BatchCountEdit.vue     |  69 ++++++
 .../sidebar/tabs/QueueSidebarTab.vue          |  19 +-
 src/i18n.ts                                   |   2 +-
 src/scripts/ui/menu/menu.css                  |  11 -
 src/stores/commandStore.ts                    |  10 +
 src/stores/queueStore.ts                      |   2 +-
 src/stores/workspaceStateStore.ts             |   6 +-
 10 files changed, 188 insertions(+), 200 deletions(-)
 create mode 100644 src/components/appMenu/BatchCountEdit.vue

diff --git a/browser_tests/appMenu.spec.ts b/browser_tests/appMenu.spec.ts
index 61eb18512..f7b3d1495 100644
--- a/browser_tests/appMenu.spec.ts
+++ b/browser_tests/appMenu.spec.ts
@@ -12,10 +12,6 @@ test.describe('AppMenu', () => {
   })
 
   test.afterEach(async ({ comfyPage }) => {
-    const currentThemeId = await comfyPage.menu.getThemeId()
-    if (currentThemeId !== 'dark') {
-      await comfyPage.menu.toggleTheme()
-    }
     await comfyPage.setSetting('Comfy.UseNewMenu', 'Disabled')
   })
 
@@ -27,7 +23,7 @@ test.describe('AppMenu', () => {
     ws
   }) => {
     // Enable change auto-queue mode
-    let queueOpts = await comfyPage.appMenu.queueButton.toggleOptions()
+    const queueOpts = await comfyPage.appMenu.queueButton.toggleOptions()
     expect(await queueOpts.getMode()).toBe('disabled')
     await queueOpts.setMode('change')
     await comfyPage.nextFrame()
diff --git a/browser_tests/helpers/appMenu.ts b/browser_tests/helpers/appMenu.ts
index 3bf1070e2..840005dec 100644
--- a/browser_tests/helpers/appMenu.ts
+++ b/browser_tests/helpers/appMenu.ts
@@ -1,4 +1,5 @@
 import type { Page, Locator } from '@playwright/test'
+import type { AutoQueueMode } from '../../src/stores/queueStore'
 
 export class ComfyAppMenu {
   public readonly root: Locator
@@ -27,39 +28,17 @@ class ComfyQueueButton {
 }
 
 class ComfyQueueButtonOptions {
-  public readonly popup: Locator
-  public readonly modes: {
-    disabled: { input: Locator; wrapper: Locator }
-    instant: { input: Locator; wrapper: Locator }
-    change: { input: Locator; wrapper: Locator }
-  }
+  constructor(public readonly page: Page) {}
 
-  constructor(public readonly page: Page) {
-    this.popup = page.getByTestId('queue-options')
-    this.modes = (['disabled', 'instant', 'change'] as const).reduce(
-      (modes, mode) => {
-        modes[mode] = {
-          input: page.locator(`#autoqueue-${mode}`),
-          wrapper: page.getByTestId(`autoqueue-${mode}`)
-        }
-        return modes
-      },
-      {} as ComfyQueueButtonOptions['modes']
-    )
-  }
-
-  public async setMode(mode: keyof ComfyQueueButtonOptions['modes']) {
-    await this.modes[mode].input.click()
+  public async setMode(mode: AutoQueueMode) {
+    await this.page.evaluate((mode) => {
+      window['app'].extensionManager.queueSettings.mode = mode
+    }, mode)
   }
 
   public async getMode() {
-    return (
-      await Promise.all(
-        Object.entries(this.modes).map(async ([mode, opt]) => [
-          mode,
-          await opt.wrapper.getAttribute('data-p-checked')
-        ])
-      )
-    ).find(([, checked]) => checked === 'true')?.[0]
+    return await this.page.evaluate(() => {
+      return window['app'].extensionManager.queueSettings.mode
+    })
   }
 }
diff --git a/src/components/appMenu/AppMenu.vue b/src/components/appMenu/AppMenu.vue
index eb8e1899d..035c617ea 100644
--- a/src/components/appMenu/AppMenu.vue
+++ b/src/components/appMenu/AppMenu.vue
@@ -1,73 +1,47 @@
 <template>
   <Panel v-if="visible" class="app-menu">
-    <div class="app-menu-content">
-      <Popover ref="queuePopover" data-testid="queue-options">
-        <div class="queue-options">
-          <p class="batch-count">
-            <FloatLabel v-tooltip="$t('menu.batchCountTooltip')">
-              <InputNumber id="batchCount" v-model="batchCount" :min="1" />
-              <label for="batchCount">{{ $t('menu.batchCount') }}</label>
-            </FloatLabel>
-
-            <Slider
-              v-model="batchCount"
-              :min="1"
-              :max="100"
-              v-tooltip="$t('menu.batchCountTooltip')"
+    <div class="app-menu-content flex align-center">
+      <div class="queue-button-group flex">
+        <SplitButton
+          class="comfyui-queue-button"
+          :label="activeQueueModeMenuItem.label"
+          :icon="activeQueueModeMenuItem.icon"
+          severity="primary"
+          @click="queuePrompt"
+          :model="queueModeMenuItems"
+          data-testid="queue-button"
+          v-tooltip.bottom="$t('menu.queueWorkflow')"
+        >
+          <template #item="{ item }">
+            <Button
+              :label="item.label"
+              :icon="item.icon"
+              :severity="item.key === queueMode ? 'primary' : 'secondary'"
+              text
+              v-tooltip="item.tooltip"
             />
-          </p>
-
-          <Divider layout="vertical" />
-
-          <p class="auto-queue">
-            <span class="label">{{ $t('menu.autoQueue') }}</span>
-            <template v-for="mode in queueModes" :key="mode">
-              <div
-                v-tooltip="$t(`menu.${mode}Tooltip`)"
-                class="auto-queue-mode"
-              >
-                <RadioButton
-                  v-model="queueMode"
-                  :inputId="`autoqueue-${mode}`"
-                  name="dynamic"
-                  :value="mode"
-                  :data-testid="`autoqueue-${mode}`"
-                />
-                <label :for="`autoqueue-${mode}`">{{
-                  $t(`menu.${mode}`)
-                }}</label>
-              </div>
-            </template>
-          </p>
-        </div>
-      </Popover>
-      <SplitButton
-        v-tooltip.bottom="$t('menu.queueWorkflow')"
-        :label="$t('menu.generate')"
-        :icon="`pi pi-${icon}`"
-        severity="secondary"
-        @click="queuePrompt"
-        :model="[]"
-        :pt="{
-          pcDropdown: ({ instance }) => {
-            instance.onDropdownButtonClick = function (e: Event) {
-              e.preventDefault()
-              queuePopover.toggle(e)
-            }
-          }
-        }"
-        data-testid="queue-button"
-      >
-      </SplitButton>
-      <div class="separator"></div>
-      <Button
-        v-tooltip.bottom="$t('menu.interrupt')"
-        icon="pi pi-times"
-        severity="secondary"
-        :disabled="!executingPrompt"
-        @click="() => commandStore.getCommand('Comfy.Interrupt')()"
-      ></Button>
-
+          </template>
+        </SplitButton>
+        <BatchCountEdit />
+        <ButtonGroup class="execution-actions ml-2">
+          <Button
+            v-tooltip.bottom="$t('menu.interrupt')"
+            icon="pi pi-times"
+            :severity="executingPrompt ? 'danger' : 'secondary'"
+            :disabled="!executingPrompt"
+            @click="() => commandStore.getCommand('Comfy.Interrupt')()"
+          >
+          </Button>
+          <Button
+            v-tooltip.bottom="$t('sideToolbar.queueTab.clearPendingTasks')"
+            icon="pi pi-stop"
+            :severity="hasPendingTasks ? 'danger' : 'secondary'"
+            :disabled="!hasPendingTasks"
+            @click="() => commandStore.getCommand('Comfy.ClearPendingTasks')()"
+          />
+        </ButtonGroup>
+      </div>
+      <Divider layout="vertical" class="mx-2" />
       <ButtonGroup>
         <Button
           v-tooltip.bottom="$t('menu.refresh')"
@@ -89,18 +63,15 @@
 </template>
 
 <script lang="ts" setup>
-import { computed, ref } from 'vue'
+import { computed } from 'vue'
 import Panel from 'primevue/panel'
+import Divider from 'primevue/divider'
 import SplitButton from 'primevue/splitbutton'
 import Button from 'primevue/button'
-import FloatLabel from 'primevue/floatlabel'
-import InputNumber from 'primevue/inputnumber'
-import Popover from 'primevue/popover'
-import Divider from 'primevue/divider'
-import Slider from 'primevue/slider'
-import RadioButton from 'primevue/radiobutton'
 import ButtonGroup from 'primevue/buttongroup'
+import BatchCountEdit from './BatchCountEdit.vue'
 import {
+  AutoQueueMode,
   useQueuePendingTaskCountStore,
   useQueueSettingsStore
 } from '@/stores/queueStore'
@@ -108,6 +79,8 @@ import { app } from '@/scripts/app'
 import { storeToRefs } from 'pinia'
 import { useSettingStore } from '@/stores/settingStore'
 import { useCommandStore } from '@/stores/commandStore'
+import { MenuItem } from 'primevue/menuitem'
+import { useI18n } from 'vue-i18n'
 
 const settingsStore = useSettingStore()
 const commandStore = useCommandStore()
@@ -118,21 +91,46 @@ const visible = computed(
   () => settingsStore.get('Comfy.UseNewMenu') === 'Floating'
 )
 
-const queuePopover = ref(null)
-const queueModes = ['disabled', 'instant', 'change']
-
-const icon = computed(() => {
-  switch (queueMode.value) {
-    case 'instant':
-      return 'forward'
-    case 'change':
-      return 'step-forward-alt'
-    default:
-      return 'play'
+const { t } = useI18n()
+const queueModeMenuItemLookup: Record<AutoQueueMode, MenuItem> = {
+  disabled: {
+    key: 'disabled',
+    label: 'Queue',
+    icon: 'pi pi-play',
+    tooltip: t('menu.disabledTooltip'),
+    command: () => {
+      queueMode.value = 'disabled'
+    }
+  },
+  instant: {
+    key: 'instant',
+    label: 'Queue (Instant)',
+    icon: 'pi pi-forward',
+    tooltip: t('menu.instantTooltip'),
+    command: () => {
+      queueMode.value = 'instant'
+    }
+  },
+  change: {
+    key: 'change',
+    label: 'Queue (Change)',
+    icon: 'pi pi-step-forward-alt',
+    tooltip: t('menu.changeTooltip'),
+    command: () => {
+      queueMode.value = 'change'
+    }
   }
-})
+}
+
+const activeQueueModeMenuItem = computed(
+  () => queueModeMenuItemLookup[queueMode.value]
+)
+const queueModeMenuItems = computed(() =>
+  Object.values(queueModeMenuItemLookup)
+)
 
 const executingPrompt = computed(() => !!queueCountStore.count.value)
+const hasPendingTasks = computed(() => queueCountStore.count.value > 1)
 
 const queuePrompt = (e: MouseEvent) => {
   app.queuePrompt(e.shiftKey ? -1 : 0, batchCount.value)
@@ -149,12 +147,6 @@ const queuePrompt = (e: MouseEvent) => {
   z-index: 1000;
 }
 
-.app-menu-content {
-  display: flex;
-  gap: 10px;
-  align-items: center;
-}
-
 :deep(.p-panel-content) {
   padding: 10px;
 }
@@ -163,50 +155,8 @@ const queuePrompt = (e: MouseEvent) => {
   display: none;
 }
 
-.separator {
-  background-color: var(--p-content-border-color);
-  border-radius: 10px;
-  opacity: 0.75;
-  width: 5px;
-  height: 20px;
-}
-
-.queue-options {
-  display: flex;
-}
-
-.batch-count {
-  padding-top: 0.75rem;
-  display: flex;
-  flex-direction: column;
-  gap: 1em;
-}
-
-.p-slider {
-  --p-slider-border-radius: 5px;
-  margin: 5px;
-  padding: 2px;
-}
-
-.p-floatlabel label {
-  left: 2px;
-}
-
-.label {
-  font-size: 12px;
-  color: var(--p-floatlabel-focus-color);
-}
-
-.auto-queue {
-  display: flex;
-  flex-direction: column;
-  gap: 10px;
-  margin-top: 6px;
-}
-
-.auto-queue-mode {
-  display: flex;
-  align-items: center;
-  gap: 5px;
+.comfyui-queue-button :deep(.p-splitbutton-dropdown) {
+  border-top-right-radius: 0;
+  border-bottom-right-radius: 0;
 }
 </style>
diff --git a/src/components/appMenu/BatchCountEdit.vue b/src/components/appMenu/BatchCountEdit.vue
new file mode 100644
index 000000000..82860326c
--- /dev/null
+++ b/src/components/appMenu/BatchCountEdit.vue
@@ -0,0 +1,69 @@
+<template>
+  <div
+    class="batch-count"
+    :class="props.class"
+    v-tooltip.bottom="$t('menu.batchCount')"
+  >
+    <InputNumber
+      class="w-14"
+      v-model="batchCount"
+      :min="minQueueCount"
+      :max="maxQueueCount"
+      fluid
+      showButtons
+      :pt="{
+        incrementButton: {
+          class: 'w-6',
+          onmousedown: () => {
+            handleClick(true)
+          }
+        },
+        decrementButton: {
+          class: 'w-6',
+          onmousedown: () => {
+            handleClick(false)
+          }
+        }
+      }"
+    />
+  </div>
+</template>
+
+<script lang="ts" setup>
+import { useQueueSettingsStore } from '@/stores/queueStore'
+import { storeToRefs } from 'pinia'
+import InputNumber from 'primevue/inputnumber'
+
+interface Props {
+  class?: string
+}
+
+const props = withDefaults(defineProps<Props>(), {
+  class: ''
+})
+
+const queueSettingsStore = useQueueSettingsStore()
+const { batchCount } = storeToRefs(queueSettingsStore)
+const minQueueCount = 1
+const maxQueueCount = 100
+
+const handleClick = (increment: boolean) => {
+  let newCount: number
+  if (increment) {
+    const originalCount = batchCount.value - 1
+    newCount = originalCount * 2
+  } else {
+    const originalCount = batchCount.value + 1
+    newCount = Math.floor(originalCount / 2)
+  }
+
+  batchCount.value = newCount
+}
+</script>
+
+<style scoped>
+:deep(.p-inputtext) {
+  border-top-left-radius: 0;
+  border-bottom-left-radius: 0;
+}
+</style>
diff --git a/src/components/sidebar/tabs/QueueSidebarTab.vue b/src/components/sidebar/tabs/QueueSidebarTab.vue
index 590929c81..57dce8c9a 100644
--- a/src/components/sidebar/tabs/QueueSidebarTab.vue
+++ b/src/components/sidebar/tabs/QueueSidebarTab.vue
@@ -34,11 +34,10 @@
         <Button
           v-if="queueStore.hasPendingTasks"
           icon="pi pi-stop"
-          text
           severity="danger"
-          @click="clearPendingTasks"
-          class="clear-pending-button"
-          v-tooltip="$t('sideToolbar.queueTab.clearPendingTasks')"
+          text
+          @click="() => commandStore.getCommand('Comfy.ClearPendingTasks')()"
+          v-tooltip.bottom="$t('sideToolbar.queueTab.clearPendingTasks')"
         />
         <Button
           icon="pi pi-trash"
@@ -109,6 +108,7 @@ import { TaskItemImpl, useQueueStore } from '@/stores/queueStore'
 import { api } from '@/scripts/api'
 import { ComfyNode } from '@/types/comfyWorkflow'
 import { useSettingStore } from '@/stores/settingStore'
+import { useCommandStore } from '@/stores/commandStore'
 import { app } from '@/scripts/app'
 
 const IMAGE_FIT = 'Comfy.Queue.ImageFit'
@@ -116,6 +116,7 @@ const confirm = useConfirm()
 const toast = useToast()
 const queueStore = useQueueStore()
 const settingStore = useSettingStore()
+const commandStore = useCommandStore()
 const { t } = useI18n()
 
 // Expanded view: show all outputs in a flat list.
@@ -230,16 +231,6 @@ const confirmRemoveAll = (event: Event) => {
   })
 }
 
-const clearPendingTasks = async () => {
-  await queueStore.clear(['queue'])
-  toast.add({
-    severity: 'info',
-    summary: 'Confirmed',
-    detail: 'Pending tasks deleted',
-    life: 3000
-  })
-}
-
 const onStatus = async () => {
   await queueStore.update()
   updateVisibleTasks()
diff --git a/src/i18n.ts b/src/i18n.ts
index 17bc1c4e1..f1deefbc7 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -79,7 +79,7 @@ const messages = {
       change: 'On Change',
       changeTooltip: 'The workflow will be queued once a change is made',
       queueWorkflow: 'Queue workflow',
-      generate: 'Generate',
+      queue: 'Queue',
       interrupt: 'Cancel current run',
       refresh: 'Refresh node definitions',
       clipspace: 'Open Clipspace',
diff --git a/src/scripts/ui/menu/menu.css b/src/scripts/ui/menu/menu.css
index 4a203d8b5..e66d182cb 100644
--- a/src/scripts/ui/menu/menu.css
+++ b/src/scripts/ui/menu/menu.css
@@ -269,10 +269,6 @@
   display: none;
 }
 
-.lt-lg .comfyui-queue-button {
-  margin-right: 44px;
-}
-
 .lt-lg .comfyui-menu-button {
   position: absolute;
   top: 4px;
@@ -296,13 +292,6 @@
 }
 
 /** Extra small */
-.lt-sm .comfyui-queue-button {
-  margin-right: 0;
-  width: 100%;
-}
-.lt-sm .comfyui-queue-button .comfyui-button {
-  justify-content: center;
-}
 .lt-sm .comfyui-interrupt-button {
   margin-right: 45px;
 }
diff --git a/src/stores/commandStore.ts b/src/stores/commandStore.ts
index c299c406b..c2947e801 100644
--- a/src/stores/commandStore.ts
+++ b/src/stores/commandStore.ts
@@ -6,6 +6,7 @@ import { globalTracker } from '@/scripts/changeTracker'
 import { useSettingStore } from '@/stores/settingStore'
 import { useToastStore } from '@/stores/toastStore'
 import { showTemplateWorkflowsDialog } from '@/services/dialogService'
+import { useQueueStore } from './queueStore'
 
 type Command = () => void | Promise<void>
 
@@ -73,6 +74,15 @@ export const useCommandStore = defineStore('command', () => {
         life: 1000
       })
     },
+    'Comfy.ClearPendingTasks': async () => {
+      await useQueueStore().clear(['queue'])
+      useToastStore().add({
+        severity: 'info',
+        summary: 'Confirmed',
+        detail: 'Pending tasks deleted',
+        life: 3000
+      })
+    },
     'Comfy.BrowseTemplates': showTemplateWorkflowsDialog
   })
 
diff --git a/src/stores/queueStore.ts b/src/stores/queueStore.ts
index dd0931818..04c7aee55 100644
--- a/src/stores/queueStore.ts
+++ b/src/stores/queueStore.ts
@@ -9,7 +9,7 @@ import type {
   TaskOutput,
   ResultItem
 } from '@/types/apiTypes'
-import type { ComfyNode, NodeId } from '@/types/comfyWorkflow'
+import type { NodeId } from '@/types/comfyWorkflow'
 import { plainToClass } from 'class-transformer'
 import _ from 'lodash'
 import { defineStore } from 'pinia'
diff --git a/src/stores/workspaceStateStore.ts b/src/stores/workspaceStateStore.ts
index 6c98bc30e..ca8bb40fd 100644
--- a/src/stores/workspaceStateStore.ts
+++ b/src/stores/workspaceStateStore.ts
@@ -1,6 +1,7 @@
-import { SidebarTabExtension, ToastManager } from '@/types/extensionTypes'
+import type { SidebarTabExtension, ToastManager } from '@/types/extensionTypes'
 import { defineStore } from 'pinia'
 import { useToastStore } from './toastStore'
+import { useQueueSettingsStore } from './queueStore'
 
 interface WorkspaceState {
   spinner: boolean
@@ -17,6 +18,9 @@ export const useWorkspaceStore = defineStore('workspace', {
   getters: {
     toast(): ToastManager {
       return useToastStore()
+    },
+    queueSettings() {
+      return useQueueSettingsStore()
     }
   },
   actions: {