Missing model download dialog test (#656)

* Test prep

* Add missing model dialog test

* Basic test of download model

* Add comment

* Adjust setting in test

* Change download dir to not interfere with other tests

browser_tests/assets/missing_models.json

@@ -0,0 +1,25 @@
+{
+  "last_node_id": 0,
+  "last_link_id": 0,
+  "nodes": [],
+  "links": [],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        0,
+        0
+      ]
+    }
+  },
+  "models": [
+    {
+      "name": "fake_model.safetensors",
+      "url": "http://localhost:8188/api/devtools/fake_model.safetensors",
+      "directory": "clip"
+    }
+  ],
+  "version": 0.4
+}

browser_tests/dialog.spec.ts

@@ -25,3 +25,35 @@ test.describe('Execution error', () => {
     await expect(executionError).toBeVisible()
   })
 })
+
+test.describe('Missing models warning', () => {
+  test.beforeEach(async ({ comfyPage }) => {
+    await comfyPage.setSetting('Comfy.Workflow.ModelDownload.AllowedSources', [
+      'http://localhost:8188'
+    ])
+    await comfyPage.setSetting('Comfy.Workflow.ModelDownload.AllowedSuffixes', [
+      '.safetensors'
+    ])
+  })
+
+  test('Should display a warning when missing models are found', async ({
+    comfyPage
+  }) => {
+    // The fake_model.safetensors is served by
+    // https://github.com/Comfy-Org/ComfyUI_devtools/blob/main/__init__.py
+    await comfyPage.loadWorkflow('missing_models')
+
+    // Wait for the element with the .comfy-missing-models selector to be visible
+    const missingModelsWarning = comfyPage.page.locator('.comfy-missing-models')
+    await expect(missingModelsWarning).toBeVisible()
+
+    // Click the download button
+    const downloadButton = comfyPage.page.getByLabel('Download')
+    await expect(downloadButton).toBeVisible()
+    await downloadButton.click()
+
+    // Wait for the element with the .download-complete selector to be visible
+    const downloadComplete = comfyPage.page.locator('.download-complete')
+    await expect(downloadComplete).toBeVisible()
+  })
+})

src/components/dialog/content/MissingModelsWarning.vue

@@ -14,28 +14,42 @@
       }"
     >
       <template #option="slotProps">
-        <div class="missing-model-item" :style="{ '--progress': `${slotProps.option.progress}%` }">
+        <div
+          class="missing-model-item"
+          :style="{ '--progress': `${slotProps.option.progress}%` }"
+        >
           <div class="model-info">
             <div class="model-details">
-              <span class="model-type" :title=slotProps.option.hint>{{ slotProps.option.label }}</span>
+              <span class="model-type" :title="slotProps.option.hint">{{
+                slotProps.option.label
+              }}</span>
+            </div>
+            <div v-if="slotProps.option.error" class="model-error">
+              {{ slotProps.option.error }}
             </div>
-            <div v-if="slotProps.option.error" class="model-error">{{ slotProps.option.error }}</div>
           </div>
           <div class="model-action">
             <Button
-              v-if="slotProps.option.action && !slotProps.option.downloading && !slotProps.option.completed && !slotProps.option.error"
+              v-if="
+                slotProps.option.action &&
+                !slotProps.option.downloading &&
+                !slotProps.option.completed &&
+                !slotProps.option.error
+              "
               @click="slotProps.option.action.callback"
               :label="slotProps.option.action.text"
               class="p-button-sm p-button-outlined model-action-button"
             />
             <div v-if="slotProps.option.downloading" class="download-progress">
-              <span class="progress-text">{{ slotProps.option.progress.toFixed(2) }}%</span>
+              <span class="progress-text"
+                >{{ slotProps.option.progress.toFixed(2) }}%</span
+              >
             </div>
             <div v-if="slotProps.option.completed" class="download-complete">
-              <i class="pi pi-check" style="color: var(--green-500);"></i>
+              <i class="pi pi-check" style="color: var(--green-500)"></i>
             </div>
             <div v-if="slotProps.option.error" class="download-error">
-              <i class="pi pi-times" style="color: var(--red-600);"></i>
+              <i class="pi pi-times" style="color: var(--red-600)"></i>
             </div>
           </div>
         </div>
@@ -50,8 +64,15 @@ import ListBox from 'primevue/listbox'
 import Button from 'primevue/button'
 import { api } from '@/scripts/api'
 import { DownloadModelStatus } from '@/types/apiTypes'
+import { useSettingStore } from '@/stores/settingStore'
 
-const allowedSources = ['https://civitai.com/', 'https://huggingface.co/']
+const settingStore = useSettingStore()
+const allowedSources = settingStore.get(
+  'Comfy.Workflow.ModelDownload.AllowedSources'
+)
+const allowedSuffixes = settingStore.get(
+  'Comfy.Workflow.ModelDownload.AllowedSuffixes'
+)
 
 interface ModelInfo {
   name: string
@@ -78,19 +99,50 @@ const handleDownloadProgress = (detail: DownloadModelStatus) => {
   }
   if (!lastModel) return
   if (detail.status === 'in_progress') {
-    modelDownloads.value[lastModel] = { ...modelDownloads.value[lastModel], downloading: true, progress: detail.progress_percentage, completed: false }
+    modelDownloads.value[lastModel] = {
+      ...modelDownloads.value[lastModel],
+      downloading: true,
+      progress: detail.progress_percentage,
+      completed: false
+    }
   } else if (detail.status === 'pending') {
-    modelDownloads.value[lastModel] = { ...modelDownloads.value[lastModel], downloading: true, progress: 0, completed: false }
+    modelDownloads.value[lastModel] = {
+      ...modelDownloads.value[lastModel],
+      downloading: true,
+      progress: 0,
+      completed: false
+    }
   } else if (detail.status === 'completed') {
-    modelDownloads.value[lastModel] = { ...modelDownloads.value[lastModel], downloading: false, progress: 100, completed: true }
+    modelDownloads.value[lastModel] = {
+      ...modelDownloads.value[lastModel],
+      downloading: false,
+      progress: 100,
+      completed: true
+    }
   } else if (detail.status === 'error') {
-    modelDownloads.value[lastModel] = { ...modelDownloads.value[lastModel], downloading: false, progress: 0, error: detail.message, completed: false }
+    modelDownloads.value[lastModel] = {
+      ...modelDownloads.value[lastModel],
+      downloading: false,
+      progress: 0,
+      error: detail.message,
+      completed: false
+    }
   }
   // TODO: other statuses?
 }
 
-const triggerDownload = async (url: string, directory: string, filename: string) => {
-  modelDownloads.value[filename] = { name: filename, directory, url, downloading: true, progress: 0 }
+const triggerDownload = async (
+  url: string,
+  directory: string,
+  filename: string
+) => {
+  modelDownloads.value[filename] = {
+    name: filename,
+    directory,
+    url,
+    downloading: true,
+    progress: 0
+  }
   const download = await api.internalDownloadModel(url, directory, filename, 1)
   handleDownloadProgress(download)
 }
@@ -100,43 +152,43 @@ api.addEventListener('download_progress', (event) => {
 })
 
 const missingModels = computed(() => {
-  return props.missingModels
-    .map((model) => {
-      const downloadInfo = modelDownloads.value[model.name]
-      if (!allowedSources.some((source) => model.url.startsWith(source))) {
-        return {
-          label: `${model.directory} / ${model.name}`,
-          hint: model.url,
-          error: 'Download not allowed from this source'
-        }
-      }
-      if (!model.name.endsWith('.safetensors') && !model.name.endsWith('.sft')) {
-        return {
-          label: `${model.directory} / ${model.name}`,
-          hint: model.url,
-          error: 'Only .safetensors models are allowed'
-        }
-      }
-      if (model.directory_invalid) {
-        return {
-          label: `${model.directory} / ${model.name}`,
-          hint: model.url,
-          error: 'Invalid directory specified (does this require custom nodes?)'
-        }
-      }
+  return props.missingModels.map((model) => {
+    const downloadInfo = modelDownloads.value[model.name]
+    if (!allowedSources.some((source) => model.url.startsWith(source))) {
       return {
         label: `${model.directory} / ${model.name}`,
         hint: model.url,
-        downloading: downloadInfo?.downloading ?? false,
-        completed: downloadInfo?.completed ?? false,
-        progress: downloadInfo?.progress ?? 0,
-        error: downloadInfo?.error,
-        action: {
-          text: 'Download',
-          callback: () => triggerDownload(model.url, model.directory, model.name)
-        }
+        error:
+          'Download not allowed from this source: ' + allowedSources.join(', ')
       }
-    })
+    }
+    if (!allowedSuffixes.some((suffix) => model.name.endsWith(suffix))) {
+      return {
+        label: `${model.directory} / ${model.name}`,
+        hint: model.url,
+        error: 'Only allowed suffixes are ' + allowedSuffixes.join(', ')
+      }
+    }
+    if (model.directory_invalid) {
+      return {
+        label: `${model.directory} / ${model.name}`,
+        hint: model.url,
+        error: 'Invalid directory specified (does this require custom nodes?)'
+      }
+    }
+    return {
+      label: `${model.directory} / ${model.name}`,
+      hint: model.url,
+      downloading: downloadInfo?.downloading ?? false,
+      completed: downloadInfo?.completed ?? false,
+      progress: downloadInfo?.progress ?? 0,
+      error: downloadInfo?.error,
+      action: {
+        text: 'Download',
+        callback: () => triggerDownload(model.url, model.directory, model.name)
+      }
+    }
+  })
 })
 </script>
 
@@ -246,7 +298,9 @@ const missingModels = computed(() => {
   min-width: 80px;
 }
 
-.download-progress, .download-complete, .download-error {
+.download-progress,
+.download-complete,
+.download-error {
   display: flex;
   align-items: center;
   justify-content: center;
@@ -258,7 +312,8 @@ const missingModels = computed(() => {
   color: var(--text-color);
 }
 
-.download-complete i, .download-error i {
+.download-complete i,
+.download-error i {
   font-size: 1.2rem;
 }
 </style>

src/stores/settingStore.ts

@@ -227,6 +227,20 @@ export const useSettingStore = defineStore('setting', {
         type: 'hidden',
         defaultValue: 'cover'
       })
+
+      app.ui.settings.addSetting({
+        id: 'Comfy.Workflow.ModelDownload.AllowedSources',
+        name: 'Allowed model download sources',
+        type: 'hidden',
+        defaultValue: ['https://huggingface.co/', 'https://civitai.com/']
+      })
+
+      app.ui.settings.addSetting({
+        id: 'Comfy.Workflow.ModelDownload.AllowedSuffixes',
+        name: 'Allowed model download suffixes',
+        type: 'hidden',
+        defaultValue: ['.safetensors', '.sft']
+      })
     },
 
     set<K extends keyof Settings>(key: K, value: Settings[K]) {

src/types/apiTypes.ts

@@ -463,7 +463,9 @@ const zSettings = z.record(z.any()).and(
       'Comfy.UseNewMenu': z.any(),
       'Comfy.Validation.Workflows': z.boolean(),
       'Comfy.Workflow.SortNodeIdOnSave': z.boolean(),
-      'Comfy.Queue.ImageFit': z.enum(['contain', 'cover'])
+      'Comfy.Queue.ImageFit': z.enum(['contain', 'cover']),
+      'Comfy.Workflow.ModelDownload.AllowedSources': z.array(z.string()),
+      'Comfy.Workflow.ModelDownload.AllowedSuffixes': z.array(z.string())
     })
     .optional()
 )

src/types/comfyWorkflow.ts

@@ -25,8 +25,8 @@ export const zDataType = z.union([z.string(), z.array(z.string()), z.number()])
 const zModelFile = z.object({
   name: z.string(),
   url: z.string().url(),
-  hash: z.string(),
-  hash_type: z.string(),
+  hash: z.string().optional(),
+  hash_type: z.string().optional(),
   directory: z.string()
 })