Moving models download settings to user settings.

src/components/dialog/content/MissingModelsWarning.vue

@@ -31,30 +31,19 @@
 </template>
 
 <script setup lang="ts">
+import { computedAsync } from '@vueuse/core'
 import Checkbox from 'primevue/checkbox'
 import ListBox from 'primevue/listbox'
-import { computed, onBeforeUnmount, ref } from 'vue'
+import { onBeforeUnmount, ref } from 'vue'
 import { useI18n } from 'vue-i18n'
 
 import FileDownload from '@/components/common/FileDownload.vue'
 import NoResultsPlaceholder from '@/components/common/NoResultsPlaceholder.vue'
+import { api } from '@/scripts/api'
 import { useSettingStore } from '@/stores/settingStore'
 import { isElectron } from '@/utils/envUtil'
 
-// 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']
-// Models that fail above conditions but are still allowed
-const whiteListedUrls = new Set([
-  'https://huggingface.co/stabilityai/stable-zero123/resolve/main/stable_zero123.ckpt',
-  'https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_depth_sd14v1.pth?download=true',
-  'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
-])
+let modelsDownloadConfig: any
 
 interface ModelInfo {
   name: string
@@ -78,7 +67,15 @@ const { t } = useI18n()
 const doNotAskAgain = ref(false)
 
 const modelDownloads = ref<Record<string, ModelInfo>>({})
-const missingModels = computed(() => {
+const missingModels = computedAsync(async () => {
+  if (!modelsDownloadConfig) {
+    modelsDownloadConfig = await api.getModelsDownloadSettings()
+  }
+  // Custom models sources, extension and whitelist can be customized in user settings.
+  const allowedSources = modelsDownloadConfig.allowedSources || []
+  const allowedSuffixes = modelsDownloadConfig.allowedSuffixes || []
+  const whiteListedUrls = new Set(modelsDownloadConfig.whitelistedUrls || [])
+
   return props.missingModels.map((model) => {
     const paths = props.paths[model.directory]
     if (model.directory_invalid || !paths) {
@@ -100,14 +97,18 @@ const missingModels = computed(() => {
     }
     modelDownloads.value[model.name] = downloadInfo
     if (!whiteListedUrls.has(model.url)) {
-      if (!allowedSources.some((source) => model.url.startsWith(source))) {
+      if (
+        !allowedSources.some((source: string) => model.url.startsWith(source))
+      ) {
         return {
           label: `${model.directory} / ${model.name}`,
           url: model.url,
           error: `Download not allowed from source '${model.url}', only allowed from '${allowedSources.join("', '")}'`
         }
       }
-      if (!allowedSuffixes.some((suffix) => model.name.endsWith(suffix))) {
+      if (
+        !allowedSuffixes.some((suffix: string) => model.name.endsWith(suffix))
+      ) {
         return {
           label: `${model.directory} / ${model.name}`,
           url: model.url,

src/schemas/apiSchema.ts

@@ -411,6 +411,16 @@ const zSettings = z.record(z.any()).and(
     .optional()
 )
 
+const zModelsDownloadSettings = z.record(z.any()).and(
+  z
+    .object({
+      allowedSources: z.array(z.string()),
+      allowedSuffixes: z.array(z.string()),
+      whiteListedUrls: z.array(z.string())
+    })
+    .optional()
+)
+
 export type EmbeddingsResponse = z.infer<typeof zEmbeddingsResponse>
 export type ExtensionsResponse = z.infer<typeof zExtensionsResponse>
 export type PromptResponse = z.infer<typeof zPromptResponse>
@@ -423,3 +433,4 @@ export type UserDataFullInfo = z.infer<typeof zUserDataFullInfo>
 export type TerminalSize = z.infer<typeof zTerminalSize>
 export type LogEntry = z.infer<typeof zLogEntry>
 export type LogsRawResponse = z.infer<typeof zLogRawResponse>
+export type ModelsDownloadSettings = z.infer<typeof zModelsDownloadSettings>

src/scripts/api.ts

@@ -13,6 +13,7 @@ import type {
   HistoryTaskItem,
   LogsRawResponse,
   LogsWsMessage,
+  ModelsDownloadSettings,
   PendingTaskItem,
   ProgressWsMessage,
   PromptResponse,
@@ -898,6 +899,22 @@ export class ComfyApi extends EventTarget {
   async getCustomNodesI18n(): Promise<Record<string, any>> {
     return (await axios.get(this.apiURL('/i18n'))).data
   }
+
+  async getModelsDownloadSettings(): Promise<ModelsDownloadSettings> {
+    const config = (
+      await axios.get(this.internalURL('/models_download/config'))
+    ).data
+    if (
+      location.hostname === 'localhost' ||
+      location.hostname === '127.0.0.1'
+    ) {
+      if (!config.allowedSources) {
+        config.allowedSources = []
+      }
+      config.allowedSources.push('http://localhost:')
+    }
+    return config
+  }
 }
 
 export const api = new ComfyApi()