From 9c42c31968764ba60904fac4c933a03dbb59712a Mon Sep 17 00:00:00 2001
From: bymyself <abolkonsky.rem@gmail.com>
Date: Fri, 24 Jan 2025 08:37:19 -0700
Subject: [PATCH] Fix missing models dialog (#2336)

---
 browser_tests/dialog.spec.ts | 56 +++++++++++++++++++++++++++++++++++-
 src/scripts/app.ts           | 22 +++++++-------
 2 files changed, 67 insertions(+), 11 deletions(-)

diff --git a/browser_tests/dialog.spec.ts b/browser_tests/dialog.spec.ts
index a90abc79d..46ad6e31d 100644
--- a/browser_tests/dialog.spec.ts
+++ b/browser_tests/dialog.spec.ts
@@ -66,9 +66,63 @@ test.describe('Missing models warning', () => {
     }, comfyPage.url)
   })
 
+  test('Should display a warning when missing models are found', async ({
+    comfyPage
+  }) => {
+    await comfyPage.loadWorkflow('missing_models')
+
+    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
+  }) => {
+    const modelFoldersRes = {
+      status: 200,
+      body: JSON.stringify([
+        {
+          name: 'clip',
+          folders: ['ComfyUI/models/clip']
+        }
+      ])
+    }
+    comfyPage.page.route(
+      '**/api/experiment/models',
+      (route) => route.fulfill(modelFoldersRes),
+      { times: 1 }
+    )
+
+    // Reload page to trigger indexing of model folders
+    await comfyPage.setup()
+
+    const clipModelsRes = {
+      status: 200,
+      body: JSON.stringify([
+        {
+          name: 'fake_model.safetensors',
+          pathIndex: 0
+        }
+      ])
+    }
+    comfyPage.page.route(
+      '**/api/experiment/models/clip',
+      (route) => route.fulfill(clipModelsRes),
+      { times: 1 }
+    )
+
+    await comfyPage.loadWorkflow('missing_models')
+
+    const missingModelsWarning = comfyPage.page.locator('.comfy-missing-models')
+    await expect(missingModelsWarning).not.toBeVisible()
+  })
+
   // Flaky test after parallelization
   // https://github.com/Comfy-Org/ComfyUI_frontend/pull/1400
-  test.skip('Should display a warning when missing models are found', async ({
+  test.skip('Should download missing model when clicking download button', async ({
     comfyPage
   }) => {
     // The fake_model.safetensors is served by
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index a49164be9..7ff50749e 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -1313,17 +1313,19 @@ export class ComfyApp {
       graphData.models &&
       useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')
     ) {
+      const modelStore = useModelStore()
       for (const m of graphData.models) {
-        const models_available = await useModelStore().getLoadedModelFolder(
-          m.directory
-        )
-        if (models_available === null) {
-          // @ts-expect-error
-          m.directory_invalid = true
-          missingModels.push(m)
-        } else if (!(m.name in models_available.models)) {
-          missingModels.push(m)
-        }
+        const modelFolder = await modelStore.getLoadedModelFolder(m.directory)
+        // @ts-expect-error
+        if (!modelFolder) m.directory_invalid = true
+
+        const modelsAvailable = modelFolder?.models
+        const modelExists =
+          modelsAvailable &&
+          Object.values(modelsAvailable).some(
+            (model) => model.file_name === m.name
+          )
+        if (!modelExists) missingModels.push(m)
       }
     }