From 6c6d86a30b86a17f4a3f63ae435dfb437a3cf87e Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Wed, 19 Feb 2025 15:25:46 -0500
Subject: [PATCH] Selection toolbox color picker button (#2637)

Co-authored-by: github-actions <github-actions@github.com>
---
 browser_tests/selectionToolbox.spec.ts        |  93 ++++++++++++
 package-lock.json                             |   8 +-
 package.json                                  |   2 +-
 .../common/ColorCustomizationSelector.vue     |   9 +-
 src/components/graph/SelectionToolbox.vue     |   2 +
 .../selectionToolbox/ColorPickerButton.vue    | 140 ++++++++++++++++++
 src/locales/en/main.json                      |   6 +
 src/locales/fr/main.json                      |   6 +
 src/locales/ja/main.json                      |   6 +
 src/locales/ko/main.json                      |   1 +
 src/locales/ru/main.json                      |   6 +
 src/locales/zh/main.json                      |   6 +
 src/utils/litegraphUtil.ts                    |  22 ++-
 13 files changed, 299 insertions(+), 8 deletions(-)
 create mode 100644 src/components/graph/selectionToolbox/ColorPickerButton.vue

diff --git a/browser_tests/selectionToolbox.spec.ts b/browser_tests/selectionToolbox.spec.ts
index dfa53262a..18e64a626 100644
--- a/browser_tests/selectionToolbox.spec.ts
+++ b/browser_tests/selectionToolbox.spec.ts
@@ -91,4 +91,97 @@ test.describe('Selection Toolbox', () => {
       )
     ).not.toBeVisible()
   })
+
+  test.describe('Color Picker', () => {
+    test('displays color picker button and allows color selection', async ({
+      comfyPage
+    }) => {
+      // Select a node
+      await comfyPage.selectNodes(['KSampler'])
+
+      // Color picker button should be visible
+      const colorPickerButton = comfyPage.page.locator(
+        '.selection-toolbox .pi-circle-fill'
+      )
+      await expect(colorPickerButton).toBeVisible()
+
+      // Click color picker button
+      await colorPickerButton.click()
+
+      // Color picker dropdown should be visible
+      const colorPickerDropdown = comfyPage.page.locator(
+        '.color-picker-container'
+      )
+      await expect(colorPickerDropdown).toBeVisible()
+
+      // Select a color (e.g., blue)
+      const blueColorOption = colorPickerDropdown.locator(
+        'i[data-testid="blue"]'
+      )
+      await blueColorOption.click()
+
+      // Dropdown should close after selection
+      await expect(colorPickerDropdown).not.toBeVisible()
+
+      // Node should have the selected color class/style
+      // Note: Exact verification method depends on how color is applied to nodes
+      const selectedNode = (await comfyPage.getNodeRefsByTitle('KSampler'))[0]
+      expect(selectedNode.getProperty('color')).not.toBeNull()
+    })
+
+    test('color picker shows current color of selected nodes', async ({
+      comfyPage
+    }) => {
+      // Select multiple nodes
+      await comfyPage.selectNodes(['KSampler', 'CLIP Text Encode (Prompt)'])
+
+      const colorPickerButton = comfyPage.page.locator(
+        '.selection-toolbox .pi-circle-fill'
+      )
+
+      // Initially should show default color
+      await expect(colorPickerButton).not.toHaveAttribute('color')
+
+      // Click color picker and select a color
+      await colorPickerButton.click()
+      const redColorOption = comfyPage.page.locator(
+        '.color-picker-container i[data-testid="red"]'
+      )
+      await redColorOption.click()
+
+      // Button should now show the selected color
+      await expect(colorPickerButton).toHaveCSS(
+        'color',
+        'rgb(85, 51, 51)' // Red color, adjust if different
+      )
+    })
+
+    test('color picker shows mixed state for differently colored selections', async ({
+      comfyPage
+    }) => {
+      // Select first node and color it
+      await comfyPage.selectNodes(['KSampler'])
+      await comfyPage.page.locator('.selection-toolbox .pi-circle-fill').click()
+      await comfyPage.page
+        .locator('.color-picker-container i[data-testid="blue"]')
+        .click()
+      await comfyPage.selectNodes(['KSampler'])
+
+      // Select second node and color it differently
+      await comfyPage.selectNodes(['CLIP Text Encode (Prompt)'])
+      await comfyPage.page.locator('.selection-toolbox .pi-circle-fill').click()
+      await comfyPage.page
+        .locator('.color-picker-container i[data-testid="red"]')
+        .click()
+
+      // Select both nodes
+      await comfyPage.selectNodes(['KSampler', 'CLIP Text Encode (Prompt)'])
+
+      // Color picker should show null/mixed state
+      const colorPickerButton = comfyPage.page.locator(
+        '.selection-toolbox .pi-circle-fill'
+      )
+      await expect(colorPickerButton).not.toHaveAttribute('color')
+    })
+  })
 })
diff --git a/package-lock.json b/package-lock.json
index 483368193..691488d80 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -11,7 +11,7 @@
       "dependencies": {
         "@atlaskit/pragmatic-drag-and-drop": "^1.3.1",
         "@comfyorg/comfyui-electron-types": "^0.4.20",
-        "@comfyorg/litegraph": "^0.8.87",
+        "@comfyorg/litegraph": "^0.8.89",
         "@primevue/forms": "^4.2.5",
         "@primevue/themes": "^4.2.5",
         "@sentry/vue": "^8.48.0",
@@ -1944,9 +1944,9 @@
       "license": "GPL-3.0-only"
     },
     "node_modules/@comfyorg/litegraph": {
-      "version": "0.8.87",
-      "resolved": "https://registry.npmjs.org/@comfyorg/litegraph/-/litegraph-0.8.87.tgz",
-      "integrity": "sha512-hEBe8Cc8C3PkWLfUxxhuO7zitYYCq3dO9mX8DfoK6On8EBE+1UijugVKfTWHuB/Yii4rN8yck/CI9yOYvCuD7Q==",
+      "version": "0.8.89",
+      "resolved": "https://registry.npmjs.org/@comfyorg/litegraph/-/litegraph-0.8.89.tgz",
+      "integrity": "sha512-/s5UUfZc3OOLmQQpAdRgPUkgK7vEqoClovzGIDmO3N++xkgbCmr1MGo8FQvC0+oqg56t5Ve1F+yYS54lwGZm0A==",
       "license": "MIT"
     },
     "node_modules/@cspotcode/source-map-support": {
diff --git a/package.json b/package.json
index a7c458dce..b4ce9d735 100644
--- a/package.json
+++ b/package.json
@@ -84,7 +84,7 @@
   "dependencies": {
     "@atlaskit/pragmatic-drag-and-drop": "^1.3.1",
     "@comfyorg/comfyui-electron-types": "^0.4.20",
-    "@comfyorg/litegraph": "^0.8.87",
+    "@comfyorg/litegraph": "^0.8.89",
     "@primevue/forms": "^4.2.5",
     "@primevue/themes": "^4.2.5",
     "@sentry/vue": "^8.48.0",
diff --git a/src/components/common/ColorCustomizationSelector.vue b/src/components/common/ColorCustomizationSelector.vue
index 50ca9d402..413549513 100644
--- a/src/components/common/ColorCustomizationSelector.vue
+++ b/src/components/common/ColorCustomizationSelector.vue
@@ -34,15 +34,20 @@ import ColorPicker from 'primevue/colorpicker'
 import SelectButton from 'primevue/selectbutton'
 import { computed, onMounted, ref, watch } from 'vue'
 
-const { modelValue, colorOptions } = defineProps<{
+const {
+  modelValue,
+  colorOptions,
+  allowCustom = true
+} = defineProps<{
   modelValue: string | null
   colorOptions: { name: Exclude<string, '_custom'>; value: string }[]
+  allowCustom?: boolean
 }>()
 
 const customColorOption = { name: '_custom', value: '' }
 const colorOptionsWithCustom = computed(() => [
   ...colorOptions,
-  customColorOption
+  ...(allowCustom ? [customColorOption] : [])
 ])
 
 const emit = defineEmits<{
diff --git a/src/components/graph/SelectionToolbox.vue b/src/components/graph/SelectionToolbox.vue
index da0a0205f..d52025ee8 100644
--- a/src/components/graph/SelectionToolbox.vue
+++ b/src/components/graph/SelectionToolbox.vue
@@ -6,6 +6,7 @@
       content: 'p-0 flex flex-row'
     }"
   >
+    <ColorPickerButton />
     <Button
       v-if="nodeSelected"
       severity="secondary"
@@ -46,6 +47,7 @@ import Button from 'primevue/button'
 import Panel from 'primevue/panel'
 import { computed } from 'vue'
 
+import ColorPickerButton from '@/components/graph/selectionToolbox/ColorPickerButton.vue'
 import { useRefreshableSelection } from '@/composables/useRefreshableSelection'
 import { useCommandStore } from '@/stores/commandStore'
 import { useCanvasStore } from '@/stores/graphStore'
diff --git a/src/components/graph/selectionToolbox/ColorPickerButton.vue b/src/components/graph/selectionToolbox/ColorPickerButton.vue
new file mode 100644
index 000000000..f77132978
--- /dev/null
+++ b/src/components/graph/selectionToolbox/ColorPickerButton.vue
@@ -0,0 +1,140 @@
+<template>
+  <div class="relative">
+    <Button
+      severity="secondary"
+      text
+      @click="() => (showColorPicker = !showColorPicker)"
+    >
+      <template #icon>
+        <div class="flex items-center gap-1">
+          <i class="pi pi-circle-fill" :style="{ color: currentColor }" />
+          <i class="pi pi-chevron-down" :style="{ fontSize: '0.5rem' }" />
+        </div>
+      </template>
+    </Button>
+    <div
+      v-if="showColorPicker"
+      class="color-picker-container absolute -top-10 left-1/2"
+    >
+      <SelectButton
+        :modelValue="selectedColorOption"
+        @update:modelValue="applyColor"
+        :options="colorOptions"
+        optionLabel="name"
+        dataKey="value"
+      >
+        <template #option="{ option }">
+          <i
+            class="pi pi-circle-fill"
+            :style="{
+              color: isLightTheme ? option.value.light : option.value.dark
+            }"
+            v-tooltip.top="option.localizedName"
+            :data-testid="option.name"
+          />
+        </template>
+      </SelectButton>
+    </div>
+  </div>
+</template>
+
+<script setup lang="ts">
+import type { ColorOption as CanvasColorOption } from '@comfyorg/litegraph'
+import { LGraphCanvas, LiteGraph, isColorable } from '@comfyorg/litegraph'
+import Button from 'primevue/button'
+import SelectButton from 'primevue/selectbutton'
+import { computed, ref, watch } from 'vue'
+import { useI18n } from 'vue-i18n'
+
+import { useCanvasStore } from '@/stores/graphStore'
+import { useColorPaletteStore } from '@/stores/workspace/colorPaletteStore'
+import { adjustColor } from '@/utils/colorUtil'
+import { getItemsColorOption } from '@/utils/litegraphUtil'
+
+const { t } = useI18n()
+const canvasStore = useCanvasStore()
+const colorPaletteStore = useColorPaletteStore()
+const isLightTheme = computed(
+  () => colorPaletteStore.completedActivePalette.light_theme
+)
+const toLightThemeColor = (color: string) =>
+  adjustColor(color, { lightness: 0.5 })
+
+const showColorPicker = ref(false)
+
+type ColorOption = {
+  name: string
+  localizedName: string
+  value: {
+    dark: string
+    light: string
+  }
+}
+
+const NO_COLOR_OPTION: ColorOption = {
+  name: 'noColor',
+  localizedName: t('color.noColor'),
+  value: {
+    dark: LiteGraph.NODE_DEFAULT_BGCOLOR,
+    light: toLightThemeColor(LiteGraph.NODE_DEFAULT_BGCOLOR)
+  }
+}
+const colorOptions: ColorOption[] = [
+  NO_COLOR_OPTION,
+  ...Object.entries(LGraphCanvas.node_colors).map(([name, color]) => ({
+    name,
+    localizedName: t(`color.${name}`),
+    value: {
+      dark: color.bgcolor,
+      light: toLightThemeColor(color.bgcolor)
+    }
+  }))
+]
+
+const selectedColorOption = ref<ColorOption | null>(null)
+const applyColor = (colorOption: ColorOption | null) => {
+  const colorName = colorOption?.name ?? NO_COLOR_OPTION.name
+  const canvasColorOption =
+    colorName === NO_COLOR_OPTION.name
+      ? null
+      : LGraphCanvas.node_colors[colorName]
+
+  for (const item of canvasStore.selectedItems) {
+    if (isColorable(item)) {
+      item.setColorOption(canvasColorOption)
+    }
+  }
+
+  canvasStore.canvas?.setDirty(true, true)
+  currentColorOption.value = canvasColorOption
+  showColorPicker.value = false
+}
+
+const currentColorOption = ref<CanvasColorOption | null>(null)
+const currentColor = computed(() =>
+  currentColorOption.value
+    ? isLightTheme.value
+      ? toLightThemeColor(currentColorOption.value?.bgcolor)
+      : currentColorOption.value?.bgcolor
+    : null
+)
+
+watch(
+  () => canvasStore.selectedItems,
+  (newSelectedItems) => {
+    showColorPicker.value = false
+    selectedColorOption.value = null
+    currentColorOption.value = getItemsColorOption(newSelectedItems)
+  }
+)
+</script>
+
+<style scoped>
+.color-picker-container {
+  transform: translateX(-50%);
+}
+
+:deep(.p-togglebutton) {
+  @apply py-2 px-1;
+}
+</style>
diff --git a/src/locales/en/main.json b/src/locales/en/main.json
index 2a819e3f5..8d12c6ae6 100644
--- a/src/locales/en/main.json
+++ b/src/locales/en/main.json
@@ -93,12 +93,18 @@
     }
   },
   "color": {
+    "noColor": "No Color",
     "default": "Default",
     "blue": "Blue",
     "green": "Green",
     "red": "Red",
     "pink": "Pink",
     "yellow": "Yellow",
+    "brown": "Brown",
+    "pale_blue": "Pale Blue",
+    "cyan": "Cyan",
+    "purple": "Purple",
+    "black": "Black",
     "custom": "Custom"
   },
   "contextMenu": {
diff --git a/src/locales/fr/main.json b/src/locales/fr/main.json
index f75fabac2..19f9511f9 100644
--- a/src/locales/fr/main.json
+++ b/src/locales/fr/main.json
@@ -5,11 +5,17 @@
     "successMessage": "Copié dans le presse-papiers"
   },
   "color": {
+    "black": "Noir",
     "blue": "Bleu",
+    "brown": "Marron",
     "custom": "Personnalisé",
+    "cyan": "Cyan",
     "default": "Par défaut",
     "green": "Vert",
+    "noColor": "Pas de couleur",
+    "pale_blue": "Bleu pâle",
     "pink": "Rose",
+    "purple": "Violet",
     "red": "Rouge",
     "yellow": "Jaune"
   },
diff --git a/src/locales/ja/main.json b/src/locales/ja/main.json
index 8b4288e4d..ac229b2f7 100644
--- a/src/locales/ja/main.json
+++ b/src/locales/ja/main.json
@@ -5,11 +5,17 @@
     "successMessage": "クリップボードにコピーしました"
   },
   "color": {
+    "black": "黒",
     "blue": "青",
+    "brown": "ブラウン",
     "custom": "カスタム",
+    "cyan": "シアン",
     "default": "デフォルト",
     "green": "緑",
+    "noColor": "色なし",
+    "pale_blue": "淡い青",
     "pink": "ピンク",
+    "purple": "パープル",
     "red": "赤",
     "yellow": "黄色"
   },
diff --git a/src/locales/ko/main.json b/src/locales/ko/main.json
index d1ff55f02..3d6b0cfde 100644
--- a/src/locales/ko/main.json
+++ b/src/locales/ko/main.json
@@ -1,4 +1,5 @@
 {
+  "0": "{",
   "clipboard": {
     "errorMessage": "클립보드에 복사하지 못했습니다",
     "errorNotSupported": "브라우저가 클립보드 API를 지원하지 않습니다.",
diff --git a/src/locales/ru/main.json b/src/locales/ru/main.json
index 2362e85ee..e71eb0bb0 100644
--- a/src/locales/ru/main.json
+++ b/src/locales/ru/main.json
@@ -5,11 +5,17 @@
     "successMessage": "Скопировано в буфер обмена"
   },
   "color": {
+    "black": "Черный",
     "blue": "Синий",
+    "brown": "Коричневый",
     "custom": "Пользовательский",
+    "cyan": "Голубой",
     "default": "По умолчанию",
     "green": "Зелёный",
+    "noColor": "Без цвета",
+    "pale_blue": "Бледно-синий",
     "pink": "Розовый",
+    "purple": "Фиолетовый",
     "red": "Красный",
     "yellow": "Жёлтый"
   },
diff --git a/src/locales/zh/main.json b/src/locales/zh/main.json
index 344197353..179b10f40 100644
--- a/src/locales/zh/main.json
+++ b/src/locales/zh/main.json
@@ -5,11 +5,17 @@
     "successMessage": "已复制到剪贴板"
   },
   "color": {
+    "black": "黑色",
     "blue": "蓝色",
+    "brown": "棕色",
     "custom": "自定义",
+    "cyan": "青色",
     "default": "默认",
     "green": "绿色",
+    "noColor": "无色",
+    "pale_blue": "淡蓝色",
     "pink": "粉色",
+    "purple": "紫色",
     "red": "红色",
     "yellow": "黄色"
   },
diff --git a/src/utils/litegraphUtil.ts b/src/utils/litegraphUtil.ts
index aa099c10b..96b0f18c6 100644
--- a/src/utils/litegraphUtil.ts
+++ b/src/utils/litegraphUtil.ts
@@ -1,5 +1,7 @@
-import type { IWidget, LGraphNode } from '@comfyorg/litegraph'
+import type { ColorOption, IWidget } from '@comfyorg/litegraph'
+import { LGraphNode, isColorable } from '@comfyorg/litegraph'
 import type { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets'
+import _ from 'lodash'
 
 export function isImageNode(node: LGraphNode) {
   return (
@@ -22,3 +24,21 @@ export const isLGraphNode = (item: unknown): item is LGraphNode => {
   const name = item?.constructor?.name
   return name === 'ComfyNode' || name === 'LGraphNode'
 }
+
+/**
+ * Get the color option of all canvas items if they are all the same.
+ * @param items - The items to get the color option of.
+ * @returns The color option of the item.
+ */
+export const getItemsColorOption = (items: unknown[]): ColorOption | null => {
+  const validItems = _.filter(items, isColorable)
+  if (_.isEmpty(validItems)) return null
+
+  const colorOptions = _.map(validItems, (item) => item.getColorOption())
+
+  return _.every(colorOptions, (option) =>
+    _.isEqual(option, _.head(colorOptions))
+  )
+    ? _.head(colorOptions)!
+    : null
+}