From e380d792c7a8d463e030767b2badab57d7cd238e Mon Sep 17 00:00:00 2001
From: bymyself <cbyrne@comfy.org>
Date: Thu, 27 Feb 2025 11:25:16 -0700
Subject: [PATCH] Support models metadata in node properties (#2754)

---
 .../missing_models_from_node_properties.json  | 49 +++++++++++++++++++
 browser_tests/dialog.spec.ts                  | 13 +++++
 src/schemas/comfyWorkflowSchema.ts            |  4 +-
 src/scripts/app.ts                            | 25 ++++++++--
 4 files changed, 87 insertions(+), 4 deletions(-)
 create mode 100644 browser_tests/assets/missing_models_from_node_properties.json

diff --git a/browser_tests/assets/missing_models_from_node_properties.json b/browser_tests/assets/missing_models_from_node_properties.json
new file mode 100644
index 000000000..cf7b1e198
--- /dev/null
+++ b/browser_tests/assets/missing_models_from_node_properties.json
@@ -0,0 +1,49 @@
+{
+  "last_node_id": 1,
+  "last_link_id": 1,
+  "nodes": [
+    {
+      "id": 1,
+      "type": "CheckpointLoaderSimple",
+      "pos": [256, 256],
+      "size": [315, 98],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "MODEL",
+          "type": "MODEL",
+          "links": null
+        },
+        {
+          "name": "CLIP",
+          "type": "CLIP",
+          "links": null
+        },
+        {
+          "name": "VAE",
+          "type": "VAE",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CheckpointLoaderSimple",
+        "models": [
+          {
+            "name": "fake_model.safetensors",
+            "url": "http://localhost:8188/api/devtools/fake_model.safetensors",
+            "directory": "text_encoders"
+          }
+        ]
+      },
+      "widgets_values": ["fake_model.safetensors"]
+    }
+  ],
+  "links": [],
+  "groups": [],
+  "config": {},
+  "extra": {},
+  "version": 0.4
+}
diff --git a/browser_tests/dialog.spec.ts b/browser_tests/dialog.spec.ts
index edf656679..f6a7c1d00 100644
--- a/browser_tests/dialog.spec.ts
+++ b/browser_tests/dialog.spec.ts
@@ -78,6 +78,19 @@ test.describe('Missing models warning', () => {
     await expect(downloadButton).toBeVisible()
   })
 
+  test('Should display a warning when missing models are found in node properties', async ({
+    comfyPage
+  }) => {
+    // Load workflow that has a node with models metadata at the node level
+    await comfyPage.loadWorkflow('missing_models_from_node_properties')
+
+    const missingModelsWarning = comfyPage.page.locator('.comfy-missing-models')
+    await expect(missingModelsWarning).toBeVisible()
+
+    const downloadButton = missingModelsWarning.getByLabel('Download')
+    await expect(downloadButton).toBeVisible()
+  })
+
   test('Should not display a warning when no missing models are found', async ({
     comfyPage
   }) => {
diff --git a/src/schemas/comfyWorkflowSchema.ts b/src/schemas/comfyWorkflowSchema.ts
index c091800b4..2dc440760 100644
--- a/src/schemas/comfyWorkflowSchema.ts
+++ b/src/schemas/comfyWorkflowSchema.ts
@@ -166,7 +166,8 @@ const zProperties = z
     ['Node name for S&R']: z.string().optional(),
     cnr_id: zCnrId.optional(),
     aux_id: zAuxId.optional(),
-    ver: zVersion.optional()
+    ver: zVersion.optional(),
+    models: z.array(zModelFile).optional()
   })
   .passthrough()
 
@@ -265,6 +266,7 @@ export const zComfyWorkflow1 = z
   })
   .passthrough()
 
+export type ModelFile = z.infer<typeof zModelFile>
 export type NodeInput = z.infer<typeof zNodeInput>
 export type NodeOutput = z.infer<typeof zNodeOutput>
 export type ComfyLink = z.infer<typeof zComfyLink>
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 91863ef13..af09c159d 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -16,6 +16,7 @@ import { st } from '@/i18n'
 import type { ComfyNodeDef } from '@/schemas/apiSchema'
 import {
   type ComfyWorkflowJSON,
+  type ModelFile,
   type NodeId,
   validateComfyWorkflow
 } from '@/schemas/comfyWorkflowSchema'
@@ -1042,13 +1043,16 @@ export class ComfyApp {
     useWorkflowService().beforeLoadNewGraph()
 
     const missingNodeTypes: MissingNodeType[] = []
-    const missingModels = []
+    const missingModels: ModelFile[] = []
     await useExtensionService().invokeExtensionsAsync(
       'beforeConfigureGraph',
       graphData,
       missingNodeTypes
       // TODO: missingModels
     )
+
+    const embeddedModels: ModelFile[] = []
+
     for (let n of graphData.nodes) {
       // Patch T2IAdapterLoader to ControlNetLoader since they are the same node now
       if (n.type == 'T2IAdapterLoader') n.type = 'ControlNetLoader'
@@ -1061,13 +1065,28 @@ export class ComfyApp {
         missingNodeTypes.push(n.type)
         n.type = sanitizeNodeName(n.type)
       }
+
+      // Collect models metadata from node
+      if (n.properties?.models?.length)
+        embeddedModels.push(...n.properties.models)
     }
+
+    // Merge models from the workflow's root-level 'models' field
+    const workflowSchemaV1Models = graphData.models
+    if (workflowSchemaV1Models?.length)
+      embeddedModels.push(...workflowSchemaV1Models)
+
+    const getModelKey = (model: ModelFile) => model.url || model.hash
+    const validModels = embeddedModels.filter(getModelKey)
+    const uniqueModels = _.uniqBy(validModels, getModelKey)
+
     if (
-      graphData.models &&
+      uniqueModels.length &&
       useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')
     ) {
       const modelStore = useModelStore()
-      for (const m of graphData.models) {
+      await modelStore.loadModelFolders()
+      for (const m of uniqueModels) {
         const modelFolder = await modelStore.getLoadedModelFolder(m.directory)
         // @ts-expect-error
         if (!modelFolder) m.directory_invalid = true