From 590280dd92aaaf44b1fbfb2c3c43b6d90be4d4fb Mon Sep 17 00:00:00 2001
From: bymyself <cbyrne@comfy.org>
Date: Thu, 27 Feb 2025 17:41:36 -0700
Subject: [PATCH] node model metadata edit popover

---
 src/components/graph/SelectionToolbox.vue     |  20 ++-
 .../nodeModelsMetadata/ModelForm.vue          | 136 ++++++++++++++++++
 .../nodeModelsMetadata/NodeModelsButton.vue   |  27 ++++
 .../nodeModelsMetadata/NodeModelsPopover.vue  | 123 ++++++++++++++++
 src/locales/en/main.json                      |  17 ++-
 src/schemas/comfyWorkflowSchema.ts            |   6 +-
 src/utils/litegraphUtil.ts                    |  20 +++
 7 files changed, 340 insertions(+), 9 deletions(-)
 create mode 100644 src/components/graph/selectionToolbox/nodeModelsMetadata/ModelForm.vue
 create mode 100644 src/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsButton.vue
 create mode 100644 src/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsPopover.vue

diff --git a/src/components/graph/SelectionToolbox.vue b/src/components/graph/SelectionToolbox.vue
index 641c52cd5..2b0b84248 100644
--- a/src/components/graph/SelectionToolbox.vue
+++ b/src/components/graph/SelectionToolbox.vue
@@ -27,6 +27,9 @@
       icon="pi pi-thumbtack"
       @click="() => commandStore.execute('Comfy.Canvas.ToggleSelected.Pin')"
     />
+    <NodeModelsButton
+      v-if="modelNodeSelected && canvasStore.selectedItems.length === 1"
+    />
     <Button
       severity="danger"
       text
@@ -57,22 +60,29 @@ import Panel from 'primevue/panel'
 import { computed } from 'vue'
 
 import ColorPickerButton from '@/components/graph/selectionToolbox/ColorPickerButton.vue'
+import NodeModelsButton from '@/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsButton.vue'
 import { useRefreshableSelection } from '@/composables/useRefreshableSelection'
 import { useExtensionService } from '@/services/extensionService'
 import { ComfyCommand, useCommandStore } from '@/stores/commandStore'
 import { useCanvasStore } from '@/stores/graphStore'
-import { isLGraphGroup, isLGraphNode } from '@/utils/litegraphUtil'
+import { isLGraphGroup, isLGraphNode, isModelNode } from '@/utils/litegraphUtil'
 
 const commandStore = useCommandStore()
 const canvasStore = useCanvasStore()
 const extensionService = useExtensionService()
 const { isRefreshable, refreshSelected } = useRefreshableSelection()
-const nodeSelected = computed(() =>
-  canvasStore.selectedItems.some(isLGraphNode)
+const selectedNodes = computed(() =>
+  canvasStore.selectedItems.filter(isLGraphNode)
 )
-const groupSelected = computed(() =>
-  canvasStore.selectedItems.some(isLGraphGroup)
+const selectedModelNodes = computed(() =>
+  selectedNodes.value.filter(isModelNode)
 )
+const selectedGroups = computed(() =>
+  canvasStore.selectedItems.filter(isLGraphGroup)
+)
+const nodeSelected = computed(() => selectedNodes.value.length > 0)
+const groupSelected = computed(() => selectedGroups.value.length > 0)
+const modelNodeSelected = computed(() => selectedModelNodes.value.length > 0)
 
 const extensionToolboxCommands = computed<ComfyCommand[]>(() => {
   const commandIds = new Set<string>(
diff --git a/src/components/graph/selectionToolbox/nodeModelsMetadata/ModelForm.vue b/src/components/graph/selectionToolbox/nodeModelsMetadata/ModelForm.vue
new file mode 100644
index 000000000..9eed9d535
--- /dev/null
+++ b/src/components/graph/selectionToolbox/nodeModelsMetadata/ModelForm.vue
@@ -0,0 +1,136 @@
+<template>
+  <Form
+    :id="isFirst ? 'node-models-form' : undefined"
+    @submit="$emit('submit', $event)"
+    :resolver="zodResolver(zModelFile)"
+    :initial-values="model"
+    v-slot="$form"
+  >
+    <div class="flex flex-col gap-2">
+      <div class="p-2 surface-ground rounded-lg">
+        <Message
+          v-if="$form.name?.error || $form.directory?.error"
+          severity="error"
+          size="small"
+          variant="simple"
+        >
+          {{ $form.name?.error?.message || $form.directory?.error?.message }}
+        </Message>
+        <div class="flex flex-col gap-2">
+          <InputGroup>
+            <InputGroupAddon>
+              <i class="pi pi-file-pdf" />
+            </InputGroupAddon>
+            <FormField v-slot="$field" name="name">
+              <IftaLabel>
+                <InputText
+                  v-bind="$field"
+                  :inputId="`model-name-${index}`"
+                  class="h-full"
+                />
+                <label :for="`model-name-${index}`">
+                  {{ $t('nodeMetadata.models.fields.name') }}
+                </label>
+              </IftaLabel>
+            </FormField>
+          </InputGroup>
+          <InputGroup>
+            <InputGroupAddon>
+              <i class="pi pi-folder" />
+            </InputGroupAddon>
+            <FormField v-slot="$field" name="directory">
+              <IftaLabel>
+                <InputText
+                  v-bind="$field"
+                  :inputId="`model-directory-${index}`"
+                  class="h-full"
+                />
+                <label :for="`model-directory-${index}`">
+                  {{ $t('nodeMetadata.models.fields.directory') }}
+                </label>
+              </IftaLabel>
+            </FormField>
+          </InputGroup>
+          <InputGroup>
+            <InputGroupAddon>
+              <i class="pi pi-link" />
+            </InputGroupAddon>
+            <FormField v-slot="$field" name="url">
+              <IftaLabel>
+                <InputText v-bind="$field" :inputId="`model-url-${index}`" />
+                <label :for="`model-url-${index}`">
+                  {{ $t('nodeMetadata.models.fields.url') }}
+                </label>
+              </IftaLabel>
+            </FormField>
+          </InputGroup>
+          <Message
+            v-if="$form.url?.error"
+            severity="error"
+            size="small"
+            variant="simple"
+          >
+            {{ $form.url?.error?.message }}
+          </Message>
+        </div>
+      </div>
+      <div class="flex items-center">
+        <div class="flex-1 flex justify-center gap-2">
+          <Button
+            v-tooltip="$t('nodeMetadata.models.remove')"
+            icon="pi pi-minus"
+            severity="danger"
+            text
+            size="small"
+            @click="$emit('remove')"
+          />
+          <Button
+            v-if="isLast"
+            v-tooltip="$t('nodeMetadata.models.add')"
+            icon="pi pi-plus"
+            text
+            size="small"
+            @click="$emit('add')"
+          />
+        </div>
+        <Button
+          v-if="showSaveButton"
+          type="submit"
+          icon="pi pi-check"
+          severity="primary"
+          size="small"
+          v-tooltip="$t('nodeMetadata.models.save')"
+          form="node-models-form"
+        />
+      </div>
+    </div>
+  </Form>
+</template>
+
+<script setup lang="ts">
+import { Form, FormField, type FormSubmitEvent } from '@primevue/forms'
+// @ts-expect-error https://github.com/primefaces/primevue/issues/6722
+import { zodResolver } from '@primevue/forms/resolvers/zod'
+import { InputGroup } from 'primevue'
+import Button from 'primevue/button'
+import IftaLabel from 'primevue/iftalabel'
+import InputGroupAddon from 'primevue/inputgroupaddon'
+import InputText from 'primevue/inputtext'
+import Message from 'primevue/message'
+
+import { type ModelFile, zModelFile } from '@/schemas/comfyWorkflowSchema'
+
+defineProps<{
+  model: ModelFile
+  index: number
+  isFirst: boolean
+  isLast: boolean
+  showSaveButton: boolean
+}>()
+
+defineEmits<{
+  (e: 'submit', event: FormSubmitEvent): void
+  (e: 'remove'): void
+  (e: 'add'): void
+}>()
+</script>
diff --git a/src/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsButton.vue b/src/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsButton.vue
new file mode 100644
index 000000000..a6fe5c594
--- /dev/null
+++ b/src/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsButton.vue
@@ -0,0 +1,27 @@
+<template>
+  <div>
+    <Button
+      severity="secondary"
+      text
+      @click="overlayRef?.toggle($event)"
+      v-tooltip.top="$t('nodeMetadata.models.title')"
+    >
+      <template #icon>
+        <i class="pi pi-box" />
+      </template>
+    </Button>
+    <OverlayPanel ref="overlayRef" class="surface-card">
+      <NodeModelsPopover />
+    </OverlayPanel>
+  </div>
+</template>
+
+<script setup lang="ts">
+import Button from 'primevue/button'
+import OverlayPanel from 'primevue/overlaypanel'
+import { ref } from 'vue'
+
+import NodeModelsPopover from '@/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsPopover.vue'
+
+const overlayRef = ref()
+</script>
diff --git a/src/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsPopover.vue b/src/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsPopover.vue
new file mode 100644
index 000000000..39aa432a9
--- /dev/null
+++ b/src/components/graph/selectionToolbox/nodeModelsMetadata/NodeModelsPopover.vue
@@ -0,0 +1,123 @@
+<template>
+  <template v-if="selectedModelNode">
+    <div ref="contentRef" class="flex flex-col overflow-y-auto px-4 py-2">
+      <div class="mb-2">
+        {{ $t('nodeMetadata.models.title') }}
+      </div>
+      <div class="flex flex-col gap-2">
+        <template
+          v-for="(model, index) of nodeModels"
+          :key="`${model.name}${model.url || model.hash}`"
+        >
+          <ModelForm
+            :model="model"
+            :index="index"
+            :is-first="index === 0"
+            :is-last="isLastModel(index)"
+            :show-save-button="shouldShowSaveButton(index)"
+            @submit="submitModel(index, $event)"
+            @remove="removeModel(index)"
+            @add="nodeModels.push({ name: '', url: '', directory: '' })"
+          />
+        </template>
+        <div v-if="nodeModels.length === 0" class="flex items-center">
+          <div class="flex-1 flex justify-center">
+            <Button
+              v-tooltip="$t('nodeMetadata.models.add')"
+              icon="pi pi-plus"
+              text
+              size="small"
+              @click="nodeModels.push({ name: '', url: '', directory: '' })"
+            />
+          </div>
+          <Button
+            type="submit"
+            icon="pi pi-check"
+            severity="primary"
+            size="small"
+            v-tooltip="$t('nodeMetadata.models.save')"
+            form="node-models-form"
+          />
+        </div>
+      </div>
+    </div>
+  </template>
+</template>
+
+<script setup lang="ts">
+import type { FormSubmitEvent } from '@primevue/forms'
+import { useResizeObserver } from '@vueuse/core'
+import Button from 'primevue/button'
+import { useToast } from 'primevue/usetoast'
+import { computed, ref } from 'vue'
+import { useI18n } from 'vue-i18n'
+
+import ModelForm from '@/components/graph/selectionToolbox/nodeModelsMetadata/ModelForm.vue'
+import type { ModelFile } from '@/schemas/comfyWorkflowSchema'
+import { app } from '@/scripts/app'
+import { useCanvasStore } from '@/stores/graphStore'
+import { isModelNode } from '@/utils/litegraphUtil'
+
+const { t } = useI18n()
+const toast = useToast()
+const canvasStore = useCanvasStore()
+
+const contentRef = ref<HTMLElement>()
+const formatMaxHeight = (top: number) => `calc(100vh - ${top}px)`
+useResizeObserver(contentRef, () => {
+  if (contentRef.value) {
+    contentRef.value.style.maxHeight = formatMaxHeight(
+      contentRef.value.getBoundingClientRect().top
+    )
+  }
+})
+
+const selectedModelNode = computed(() => {
+  const nodes = canvasStore.selectedItems.filter(isModelNode)
+  if (!nodes.length) return null
+  return app.graph.getNodeById(nodes[0].id)
+})
+
+const nodeModels = ref<ModelFile[]>([
+  ...((selectedModelNode.value?.properties?.models as ModelFile[]) ?? [])
+])
+
+const isEmpty = computed(() => nodeModels.value.length === 0)
+const isLastModel = (index: number) => index === nodeModels.value.length - 1
+const shouldShowSaveButton = (index: number) =>
+  isLastModel(index) || isEmpty.value
+
+const updateNodeProperties = () => {
+  if (!selectedModelNode.value) return
+  if (!selectedModelNode.value.properties) {
+    selectedModelNode.value.properties = {}
+  }
+  selectedModelNode.value.properties.models = nodeModels.value
+}
+
+const removeModel = (index: number) => {
+  nodeModels.value.splice(index, 1)
+  updateNodeProperties()
+}
+
+const submitModel = (index: number, event: FormSubmitEvent) => {
+  if (event.valid) {
+    try {
+      nodeModels.value[index] = event.values
+      updateNodeProperties()
+      toast.add({
+        severity: 'success',
+        summary: t('nodeMetadata.models.modelUpdated'),
+        life: 3000
+      })
+    } catch (error) {
+      toast.add({
+        severity: 'error',
+        summary: t('nodeMetadata.models.modelUpdateFailed'),
+        detail: error instanceof Error ? error.message : String(error),
+        life: 3000
+      })
+    }
+  }
+}
+</script>
diff --git a/src/locales/en/main.json b/src/locales/en/main.json
index 77151e26a..484bffd2e 100644
--- a/src/locales/en/main.json
+++ b/src/locales/en/main.json
@@ -80,6 +80,21 @@
     "choose_file_to_upload": "choose file to upload",
     "capture": "capture"
   },
+  "nodeMetadata": {
+    "models" : {
+      "title": "Edit Node Models Metadata",
+      "save": "Save Models Metadata",
+      "add": "Add Model",
+      "remove": "Remove Model",
+      "modelUpdated": "Model updated",
+      "modelUpdateFailed": "Invalid model metadata. Node not updated.",
+      "fields": {
+        "name": "Name",
+        "directory": "Directory",
+        "url": "URL"
+      }
+    }
+  },
   "issueReport": {
     "submitErrorReport": "Submit Error Report (Optional)",
     "provideEmail": "Give us your email (optional)",
@@ -860,4 +875,4 @@
     "camera": "Camera",
     "light": "Light"
   }
-}
\ No newline at end of file
+}
diff --git a/src/schemas/comfyWorkflowSchema.ts b/src/schemas/comfyWorkflowSchema.ts
index 2dc440760..a2623d211 100644
--- a/src/schemas/comfyWorkflowSchema.ts
+++ b/src/schemas/comfyWorkflowSchema.ts
@@ -31,12 +31,12 @@ const zVector2 = z.union([
 ])
 
 // Definition of an AI model file used in the workflow.
-const zModelFile = z.object({
-  name: z.string(),
+export const zModelFile = z.object({
+  name: z.string().min(1, 'Name cannot be empty'),
   url: z.string().url(),
   hash: z.string().optional(),
   hash_type: z.string().optional(),
-  directory: z.string()
+  directory: z.string().min(1, 'Directory cannot be empty')
 })
 
 const zGraphState = z
diff --git a/src/utils/litegraphUtil.ts b/src/utils/litegraphUtil.ts
index 66187b4e7..aa1ce0bc1 100644
--- a/src/utils/litegraphUtil.ts
+++ b/src/utils/litegraphUtil.ts
@@ -47,6 +47,26 @@ export const isLGraphGroup = (item: unknown): item is LGraphGroup => {
   return item instanceof LGraphGroup
 }
 
+const modelOutputTypes = new Set([
+  'MODEL',
+  'CLIP',
+  'VAE',
+  'CONTROL_NET',
+  'UPSCALE_MODEL',
+  'CLIP_VISION',
+  'STYLE_MODEL',
+  'GLIGEN',
+  'HOOKS',
+  'UPSCALE_MODEL',
+  'PHOTOMAKER',
+  'SAM_MODEL'
+])
+
+export const isModelNode = (node: LGraphNode): boolean => {
+  if (!node?.outputs?.length) return false
+  return node.outputs.some((output) => modelOutputTypes.has(`${output.type}`))
+}
+
 /**
  * Get the color option of all canvas items if they are all the same.
  * @param items - The items to get the color option of.