From 35a7c81fd8185b474b0f16197c7e6d80ca82fa0f Mon Sep 17 00:00:00 2001
From: "Alex \"mcmonkey\" Goodwin"
 <4000772+mcmonkey4eva@users.noreply.github.com>
Date: Tue, 24 Sep 2024 16:51:30 +0900
Subject: [PATCH] initial download-folder-selector interface (#890)

* initial download-folder-selector interface

* use primevue select

* add a folder select visibility checkbox

* slightly reduce indirection

* fix up select box updating

* revert bad upstream changes

* cleanup

* allow localhost sourced models in ui side

(for testing purposes only basically, but does no harm in deployed envs)

* add screenshot expectations to test

* Update test expectations [skip ci]

* add testing of folder select

* fix test

* don't exclude folder selector when there's only 1

since the checkbox covers that better anyway

* oo - fix checkbox

* Update test expectations [skip ci]

* testing - don't expect screenshots :(

* experimental new test code

* toHaveClass is silly

* add // comments documenting intent of allowedSources

---------

Co-authored-by: github-actions <github-actions@github.com>
---
 browser_tests/dialog.spec.ts                  |  38 +++--
 .../dialog/content/MissingModelsWarning.vue   | 133 +++++++++++++-----
 src/scripts/api.ts                            |  10 +-
 src/scripts/app.ts                            |   6 +-
 src/services/dialogService.ts                 |   1 +
 src/stores/coreSettings.ts                    |  12 --
 src/types/apiTypes.ts                         |   2 -
 7 files changed, 139 insertions(+), 63 deletions(-)

diff --git a/browser_tests/dialog.spec.ts b/browser_tests/dialog.spec.ts
index 7c56a1859..e4e54a87a 100644
--- a/browser_tests/dialog.spec.ts
+++ b/browser_tests/dialog.spec.ts
@@ -45,38 +45,54 @@ test.describe('Execution error', () => {
 
 test.describe('Missing models warning', () => {
   test.beforeEach(async ({ comfyPage }) => {
+    await comfyPage.setSetting('Comfy.Workflow.ShowMissingModelsWarning', true)
     await comfyPage.page.evaluate((url: string) => {
       return fetch(`${url}/api/devtools/cleanup_fake_model`)
     }, comfyPage.url)
-    await comfyPage.setSetting('Comfy.Workflow.ModelDownload.AllowedSources', [
-      'http://localhost:8188'
-    ])
-    await comfyPage.setSetting('Comfy.Workflow.ModelDownload.AllowedSuffixes', [
-      '.safetensors'
-    ])
   })
 
   // Regressed by https://github.com/comfyanonymous/ComfyUI/pull/4981
   test.skip('Should display a warning when missing models are found', async ({
     comfyPage
   }) => {
-    await comfyPage.setSetting('Comfy.Workflow.ShowMissingModelsWarning', true)
-
     // 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()
   })
+
+  test('Can configure download folder', async ({ comfyPage }) => {
+    await comfyPage.loadWorkflow('missing_models')
+
+    const missingModelsWarning = comfyPage.page.locator('.comfy-missing-models')
+    await expect(missingModelsWarning).toBeVisible()
+
+    const folderSelectToggle = comfyPage.page.locator(
+      '.model-path-select-checkbox'
+    )
+    const folderSelect = comfyPage.page.locator('.model-path-select')
+    await expect(folderSelectToggle).toBeVisible()
+    await expect(folderSelect).not.toBeVisible()
+
+    await folderSelectToggle.click() // show the selectors
+    await expect(folderSelect).toBeVisible()
+
+    await folderSelect.click() // open dropdown
+    await expect(folderSelect).toHaveClass(/p-select-open/)
+
+    await folderSelect.click() // close the dropdown
+    await expect(folderSelect).not.toHaveClass(/p-select-open/)
+
+    await folderSelectToggle.click() // hide the selectors
+    await expect(folderSelect).not.toBeVisible()
+  })
 })
diff --git a/src/components/dialog/content/MissingModelsWarning.vue b/src/components/dialog/content/MissingModelsWarning.vue
index 78541a860..cf9ac447f 100644
--- a/src/components/dialog/content/MissingModelsWarning.vue
+++ b/src/components/dialog/content/MissingModelsWarning.vue
@@ -4,6 +4,15 @@
     <p class="warning-description">
       When loading the graph, the following models were not found:
     </p>
+    <p class="warning-options">
+      <Checkbox
+        class="model-path-select-checkbox"
+        v-model="showFolderSelect"
+        label="Show folder selector"
+        :binary="true"
+      />
+      Show folder selector
+    </p>
     <ListBox
       :options="missingModels"
       optionLabel="label"
@@ -29,6 +38,19 @@
             </div>
           </div>
           <div class="model-action">
+            <Select
+              class="model-path-select"
+              v-if="
+                slotProps.option.action &&
+                !slotProps.option.downloading &&
+                !slotProps.option.completed &&
+                !slotProps.option.error &&
+                showFolderSelect
+              "
+              v-model="slotProps.option.folderPath"
+              :options="slotProps.option.paths"
+              @change="updateFolderPath(slotProps.option, $event)"
+            />
             <Button
               v-if="
                 slotProps.option.action &&
@@ -60,19 +82,24 @@
 
 <script setup lang="ts">
 import { ref, computed } from 'vue'
+import Checkbox from 'primevue/checkbox'
 import ListBox from 'primevue/listbox'
+import Select from 'primevue/select'
+import { SelectChangeEvent } from 'primevue/select'
 import Button from 'primevue/button'
 import { api } from '@/scripts/api'
 import { DownloadModelStatus } from '@/types/apiTypes'
-import { useSettingStore } from '@/stores/settingStore'
 
-const settingStore = useSettingStore()
-const allowedSources = settingStore.get(
-  'Comfy.Workflow.ModelDownload.AllowedSources'
-)
-const allowedSuffixes = settingStore.get(
-  'Comfy.Workflow.ModelDownload.AllowedSuffixes'
-)
+const showFolderSelect = ref(false)
+
+// TODO: Read this from server internal API rather than hardcoding here
+// as some installations may wish to use custom sources
+const allowedSources = [
+  'https://civitai.com/',
+  'https://huggingface.co/',
+  'http://localhost:' // Included for testing usage only
+]
+const allowedSuffixes = ['.safetensors', '.sft']
 
 interface ModelInfo {
   name: string
@@ -83,19 +110,26 @@ interface ModelInfo {
   completed?: boolean
   progress?: number
   error?: string
+  folder_path?: string
 }
 
 const props = defineProps<{
   missingModels: ModelInfo[]
+  paths: Record<string, string[]>
   maximized: boolean
 }>()
 
 const modelDownloads = ref<Record<string, ModelInfo>>({})
 let lastModel: string | null = null
 
+const updateFolderPath = (model: any, event: SelectChangeEvent) => {
+  const downloadInfo = modelDownloads.value[model.name]
+  downloadInfo.folder_path = event.value
+  return false
+}
 const handleDownloadProgress = (detail: DownloadModelStatus) => {
   if (detail.download_path) {
-    lastModel = detail.download_path.split('/', 2)[1]
+    lastModel = detail.download_path
   }
   if (!lastModel) return
   if (detail.status === 'in_progress') {
@@ -134,7 +168,8 @@ const handleDownloadProgress = (detail: DownloadModelStatus) => {
 const triggerDownload = async (
   url: string,
   directory: string,
-  filename: string
+  filename: string,
+  folder_path: string
 ) => {
   modelDownloads.value[filename] = {
     name: filename,
@@ -143,49 +178,75 @@ const triggerDownload = async (
     downloading: true,
     progress: 0
   }
-  const download = await api.internalDownloadModel(url, directory, filename, 1)
+  const download = await api.internalDownloadModel(
+    url,
+    directory,
+    filename,
+    1,
+    folder_path
+  )
+  lastModel = filename
   handleDownloadProgress(download)
 }
 
-api.addEventListener('download_progress', (event) => {
+api.addEventListener('download_progress', (event: CustomEvent) => {
   handleDownloadProgress(event.detail)
 })
 
 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: ' + 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) {
+    const paths = props.paths[model.directory]
+    if (model.directory_invalid || !paths) {
       return {
         label: `${model.directory} / ${model.name}`,
         hint: model.url,
         error: 'Invalid directory specified (does this require custom nodes?)'
       }
     }
+    const downloadInfo: ModelInfo = modelDownloads.value[model.name] ?? {
+      downloading: false,
+      completed: false,
+      progress: 0,
+      error: null,
+      name: model.name,
+      directory: model.directory,
+      url: model.url,
+      folder_path: paths[0]
+    }
+    modelDownloads.value[model.name] = downloadInfo
+    if (!allowedSources.some((source) => model.url.startsWith(source))) {
+      return {
+        label: `${model.directory} / ${model.name}`,
+        hint: model.url,
+        error: `Download not allowed from source '${model.url}', only allowed from '${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("', '")}'`
+      }
+    }
     return {
       label: `${model.directory} / ${model.name}`,
       hint: model.url,
-      downloading: downloadInfo?.downloading ?? false,
-      completed: downloadInfo?.completed ?? false,
-      progress: downloadInfo?.progress ?? 0,
-      error: downloadInfo?.error,
+      downloading: downloadInfo.downloading,
+      completed: downloadInfo.completed,
+      progress: downloadInfo.progress,
+      error: downloadInfo.error,
+      name: model.name,
+      paths: paths,
+      folderPath: downloadInfo.folder_path,
       action: {
         text: 'Download',
-        callback: () => triggerDownload(model.url, model.directory, model.name)
+        callback: () =>
+          triggerDownload(
+            model.url,
+            model.directory,
+            model.name,
+            downloadInfo.folder_path
+          )
       }
     }
   })
@@ -218,6 +279,10 @@ const missingModels = computed(() => {
   margin-bottom: 1rem;
 }
 
+.warning-options {
+  color: var(--fg-color);
+}
+
 .missing-models-list {
   max-height: 300px;
   overflow-y: auto;
diff --git a/src/scripts/api.ts b/src/scripts/api.ts
index 8b431c12c..30a35dcf5 100644
--- a/src/scripts/api.ts
+++ b/src/scripts/api.ts
@@ -396,7 +396,8 @@ class ComfyApi extends EventTarget {
     url: string,
     model_directory: string,
     model_filename: string,
-    progress_interval: number
+    progress_interval: number,
+    folder_path: string
   ): Promise<DownloadModelStatus> {
     const res = await this.fetchApi('/internal/models/download', {
       method: 'POST',
@@ -407,7 +408,8 @@ class ComfyApi extends EventTarget {
         url,
         model_directory,
         model_filename,
-        progress_interval
+        progress_interval,
+        folder_path
       })
     })
     return await res.json()
@@ -716,6 +718,10 @@ class ComfyApi extends EventTarget {
   async getLogs(): Promise<string> {
     return (await axios.get(this.internalURL('/logs'))).data
   }
+
+  async getFolderPaths(): Promise<Record<string, string[]>> {
+    return (await axios.get(this.internalURL('/folder_paths'))).data
+  }
 }
 
 export const api = new ComfyApi()
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index cd8a99a62..4791bc085 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -2189,10 +2189,11 @@ export class ComfyApp {
     })
   }
 
-  showMissingModelsError(missingModels) {
+  showMissingModelsError(missingModels, paths) {
     if (useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')) {
       showMissingModelsWarning({
         missingModels,
+        paths,
         maximizable: true
       })
     }
@@ -2414,7 +2415,8 @@ export class ComfyApp {
       this.showMissingNodesError(missingNodeTypes)
     }
     if (missingModels.length && showMissingModelsDialog) {
-      this.showMissingModelsError(missingModels)
+      const paths = await api.getFolderPaths()
+      this.showMissingModelsError(missingModels, paths)
     }
     await this.#invokeExtensionsAsync('afterConfigureGraph', missingNodeTypes)
     requestAnimationFrame(() => {
diff --git a/src/services/dialogService.ts b/src/services/dialogService.ts
index 21049f13a..86cc8fd4a 100644
--- a/src/services/dialogService.ts
+++ b/src/services/dialogService.ts
@@ -25,6 +25,7 @@ export function showLoadWorkflowWarning(props: {
 
 export function showMissingModelsWarning(props: {
   missingModels: any[]
+  paths: Record<string, string[]>
   [key: string]: any
 }) {
   const dialogStore = useDialogStore()
diff --git a/src/stores/coreSettings.ts b/src/stores/coreSettings.ts
index f786b60e8..48ba4c0aa 100644
--- a/src/stores/coreSettings.ts
+++ b/src/stores/coreSettings.ts
@@ -203,18 +203,6 @@ export const CORE_SETTINGS: SettingParams[] = [
     type: 'hidden',
     defaultValue: 'cover'
   },
-  {
-    id: 'Comfy.Workflow.ModelDownload.AllowedSources',
-    name: 'Allowed model download sources',
-    type: 'hidden',
-    defaultValue: ['https://huggingface.co/', 'https://civitai.com/']
-  },
-  {
-    id: 'Comfy.Workflow.ModelDownload.AllowedSuffixes',
-    name: 'Allowed model download suffixes',
-    type: 'hidden',
-    defaultValue: ['.safetensors', '.sft']
-  },
   {
     id: 'Comfy.GroupSelectedNodes.Padding',
     name: 'Group selected nodes padding',
diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts
index 0cc3eaf64..080a7bc7c 100644
--- a/src/types/apiTypes.ts
+++ b/src/types/apiTypes.ts
@@ -498,8 +498,6 @@ const zSettings = z.record(z.any()).and(
       'Comfy.Validation.Workflows': z.boolean(),
       'Comfy.Workflow.SortNodeIdOnSave': z.boolean(),
       'Comfy.Queue.ImageFit': z.enum(['contain', 'cover']),
-      'Comfy.Workflow.ModelDownload.AllowedSources': z.array(z.string()),
-      'Comfy.Workflow.ModelDownload.AllowedSuffixes': z.array(z.string()),
       'Comfy.Workflow.WorkflowTabsPosition': z.enum(['Sidebar', 'Topbar']),
       'Comfy.Node.DoubleClickTitleToEdit': z.boolean(),
       'Comfy.Window.UnloadConfirmation': z.boolean(),