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'
- }
+ 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: ' + allowedSources.join(', ')
  }
-  if (model.directory_invalid) {
-  return {
-  label: `${model.directory} / ${model.name}`,
-  hint: model.url,
-  error: 'Invalid directory specified (does this require custom nodes?)'
- }
+ }
+ 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,
- 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: '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()
 })