Model downloader dialog (#569)

* API core for model downloader

* initial basic dialog for missing models

* app.ts handling for missing models

* don't explode if getModels is a 404

* actually track downloads in progress

* overall pile of improvements to the missing models view

* minor fixes

* add setting to disable missing models warning

* temporarily remove 'models' entry from default graph

to avoid missing model dialog causing issues. Also because ckpt autodownloading shouldn't be allowed

* swap the url to a title

* add model directory to display

* match settingStore commit

* check setting before scanning models list

ie avoid redundant calcs when setting is disabled anyway

src/components/dialog/content/MissingModelsWarning.vue

@@ -0,0 +1,262 @@
+<template>
+  <div class="comfy-missing-models">
+    <h4 class="warning-title">Warning: Missing Models</h4>
+    <p class="warning-description">
+      When loading the graph, the following models were not found:
+    </p>
+    <ListBox
+      :options="missingModels"
+      optionLabel="label"
+      scrollHeight="100%"
+      :class="'missing-models-list' + (props.maximized ? ' maximized' : '')"
+      :pt="{
+        list: { class: 'border-none' }
+      }"
+    >
+      <template #option="slotProps">
+        <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>
+            </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"
+              @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>
+            </div>
+            <div v-if="slotProps.option.completed" class="download-complete">
+              <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>
+            </div>
+          </div>
+        </div>
+      </template>
+    </ListBox>
+  </div>
+</template>
+
+<script setup lang="ts">
+import { ref, computed } from 'vue'
+import ListBox from 'primevue/listbox'
+import Button from 'primevue/button'
+import { api } from '@/scripts/api'
+import { DownloadModelStatus } from '@/types/apiTypes'
+
+const allowedSources = ['https://civitai.com/', 'https://huggingface.co/']
+
+interface ModelInfo {
+  name: string
+  directory: string
+  directory_invalid?: boolean
+  url: string
+  downloading?: boolean
+  completed?: boolean
+  progress?: number
+  error?: string
+}
+
+const props = defineProps<{
+  missingModels: ModelInfo[]
+  maximized: boolean
+}>()
+
+const modelDownloads = ref<Record<string, ModelInfo>>({})
+let lastModel: string | null = null
+
+const handleDownloadProgress = (detail: DownloadModelStatus) => {
+  if (detail.status === 'in_progress') {
+    const model = detail.message.split(' ', 2)[1] // TODO: better way to track which model is being downloaded?
+    lastModel = model
+    const progress = detail.progress_percentage
+    modelDownloads.value[model] = { ...modelDownloads.value[model], downloading: true, progress, completed: false }
+  } else if (detail.status === 'pending') {
+    const model = detail.message.split(' ', 4)[3]
+    lastModel = model
+    modelDownloads.value[model] = { ...modelDownloads.value[model], downloading: true, progress: 0, completed: false }
+  } else if (detail.status === 'completed') {
+    const model = detail.message.split(' ', 3)[2]
+    lastModel = model
+    modelDownloads.value[model] = { ...modelDownloads.value[model], downloading: false, progress: 100, completed: true }
+  } else if (detail.status === 'error') {
+    if (lastModel) {
+      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 download = await api.internalDownloadModel(url, directory, filename, 1)
+  handleDownloadProgress(download)
+}
+
+api.addEventListener('download_progress', (event) => {
+  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'
+        }
+      }
+      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>
+
+<style>
+:root {
+  --red-600: #dc3545;
+  --green-500: #28a745;
+}
+</style>
+
+<style scoped>
+.comfy-missing-models {
+  font-family: monospace;
+  color: var(--red-600);
+  padding: 1.5rem;
+  background-color: var(--surface-ground);
+  border-radius: var(--border-radius);
+  box-shadow: var(--card-shadow);
+}
+
+.warning-title {
+  margin-top: 0;
+  margin-bottom: 1rem;
+}
+
+.warning-description {
+  margin-bottom: 1rem;
+}
+
+.missing-models-list {
+  max-height: 300px;
+  overflow-y: auto;
+}
+
+.missing-models-list.maximized {
+  max-height: unset;
+}
+
+.missing-model-item {
+  display: flex;
+  align-items: flex-start;
+  padding: 0.5rem;
+  position: relative;
+  overflow: hidden;
+  width: 100%;
+}
+
+.missing-model-item::before {
+  content: '';
+  position: absolute;
+  top: 0;
+  left: 0;
+  height: 100%;
+  width: var(--progress);
+  background-color: var(--green-500);
+  opacity: 0.2;
+  transition: width 0.3s ease;
+}
+
+.model-info {
+  flex: 1;
+  min-width: 0;
+  z-index: 1;
+  display: flex;
+  flex-direction: column;
+  margin-right: 1rem;
+  overflow: hidden;
+}
+
+.model-details {
+  display: flex;
+  align-items: center;
+  flex-wrap: wrap;
+}
+
+.model-type {
+  font-weight: 600;
+  color: var(--text-color);
+  margin-right: 0.5rem;
+  white-space: nowrap;
+  overflow: hidden;
+  text-overflow: ellipsis;
+}
+
+.model-hint {
+  font-style: italic;
+  color: var(--text-color-secondary);
+  white-space: nowrap;
+  overflow: hidden;
+  text-overflow: ellipsis;
+}
+
+.model-error {
+  color: var(--red-600);
+  font-size: 0.8rem;
+  margin-top: 0.25rem;
+}
+
+.model-action {
+  display: flex;
+  align-items: center;
+  justify-content: flex-end;
+  z-index: 1;
+}
+
+.model-action-button {
+  min-width: 80px;
+}
+
+.download-progress, .download-complete, .download-error {
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  min-width: 80px;
+}
+
+.progress-text {
+  font-size: 0.8rem;
+  color: var(--text-color);
+}
+
+.download-complete i, .download-error i {
+  font-size: 1.2rem;
+}
+</style>

src/scripts/api.ts

@@ -1,5 +1,6 @@
 import { ComfyWorkflowJSON } from '@/types/comfyWorkflow'
 import {
+  DownloadModelStatus,
   HistoryTaskItem,
   PendingTaskItem,
   RunningTaskItem,
@@ -216,6 +217,11 @@ class ComfyApi extends EventTarget {
                 new CustomEvent('execution_cached', { detail: msg.data })
               )
               break
+            case 'download_progress':
+              this.dispatchEvent(
+                new CustomEvent('download_progress', { detail: msg.data })
+              )
+              break
             default:
               if (this.#registered.has(msg.type)) {
                 this.dispatchEvent(
@@ -319,6 +325,47 @@ class ComfyApi extends EventTarget {
     return await res.json()
   }
 
+  /**
+   * Gets a list of models in the specified folder
+   * @param {string} folder The folder to list models from, such as 'checkpoints'
+   * @returns The list of model filenames within the specified folder
+   */
+  async getModels(folder: string) {
+    const res = await this.fetchApi(`/models/${folder}`)
+    if (res.status === 404) {
+      return null
+    }
+    return await res.json()
+  }
+
+  /**
+   * Tells the server to download a model from the specified URL to the specified directory and filename
+   * @param {string} url The URL to download the model from
+   * @param {string} model_directory The main directory (eg 'checkpoints') to save the model to
+   * @param {string} model_filename The filename to save the model as
+   * @param {number} progress_interval The interval in seconds at which to report download progress (via 'download_progress' event)
+   */
+  async internalDownloadModel(
+    url: string,
+    model_directory: string,
+    model_filename: string,
+    progress_interval: number
+  ): Promise<DownloadModelStatus> {
+    const res = await this.fetchApi('/internal/models/download', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json'
+      },
+      body: JSON.stringify({
+        url,
+        model_directory,
+        model_filename,
+        progress_interval
+      })
+    })
+    return await res.json()
+  }
+
   /**
    * Loads a list of items (queue or history)
    * @param {string} type The type of items to load, queue or history

src/scripts/app.ts

@@ -45,7 +45,8 @@ import { Vector2 } from '@comfyorg/litegraph'
 import _ from 'lodash'
 import {
   showExecutionErrorDialog,
-  showLoadWorkflowWarning
+  showLoadWorkflowWarning,
+  showMissingModelsWarning
 } from '@/services/dialogService'
 import { useSettingStore } from '@/stores/settingStore'
 import { useToastStore } from '@/stores/toastStore'
@@ -134,6 +135,7 @@ export class ComfyApp {
   bodyBottom: HTMLElement
   canvasContainer: HTMLElement
   menu: ComfyAppMenu
+  modelsInFolderCache: Record<string, string[]>
 
   constructor() {
     this.vueAppReady = false
@@ -148,6 +150,7 @@ export class ComfyApp {
       parent: document.body
     })
     this.menu = new ComfyAppMenu(this)
+    this.modelsInFolderCache = {}
 
     /**
      * List of extensions that are registered with the app
@@ -2175,6 +2178,22 @@ export class ComfyApp {
     })
   }
 
+  showMissingModelsError(missingModels) {
+    if (
+      this.vueAppReady &&
+      useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')
+    ) {
+      showMissingModelsWarning({
+        missingModels,
+        maximizable: true
+      })
+    }
+    
+    this.logging.addEntry('Comfy.App', 'warn', {
+      MissingModels: missingModels
+    })
+  }
+
   async changeWorkflow(callback, workflow = null) {
     try {
       this.workflowManager.activeWorkflow?.changeTracker?.store()
@@ -2233,10 +2252,12 @@ export class ComfyApp {
     }
 
     const missingNodeTypes = []
+    const missingModels = []
     await this.#invokeExtensionsAsync(
       'beforeConfigureGraph',
       graphData,
       missingNodeTypes
+      // TODO: missingModels
     )
     for (let n of graphData.nodes) {
       // Patch T2IAdapterLoader to ControlNetLoader since they are the same node now
@@ -2251,6 +2272,19 @@ export class ComfyApp {
         n.type = sanitizeNodeName(n.type)
       }
     }
+    if (graphData.models && useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')) {
+      for (let m of graphData.models) {
+        const models_available = await this.getModelsInFolderCached(m.directory)
+        if (models_available === null) {
+          // @ts-expect-error
+          m.directory_invalid = true
+          missingModels.push(m)
+        }
+        else if (!models_available.includes(m.name)) {
+          missingModels.push(m)
+        }
+      }
+    }
 
     try {
       this.graph.configure(graphData)
@@ -2366,9 +2400,13 @@ export class ComfyApp {
       this.#invokeExtensions('loadedGraphNode', node)
     }
 
+    // TODO: Properly handle if both nodes and models are missing (sequential dialogs?)
     if (missingNodeTypes.length) {
       this.showMissingNodesError(missingNodeTypes)
     }
+    if (missingModels.length) {
+      this.showMissingModelsError(missingModels)
+    }
     await this.#invokeExtensionsAsync('afterConfigureGraph', missingNodeTypes)
     requestAnimationFrame(() => {
       this.graph.setDirtyCanvas(true, true)
@@ -2837,6 +2875,19 @@ export class ComfyApp {
     app.graph.arrange()
   }
 
+  /**
+   * Gets the list of model names in a folder, using a temporary local cache
+   */
+  async getModelsInFolderCached(folder: string): Promise<string[]> {
+    if (folder in this.modelsInFolderCache) {
+      return this.modelsInFolderCache[folder]
+    }
+    // TODO: needs a lock to avoid overlapping calls
+    const models = await api.getModels(folder)
+    this.modelsInFolderCache[folder] = models
+    return models
+  }
+
   /**
    * Registers a Comfy web extension with the app
    * @param {ComfyExtension} extension
@@ -2862,6 +2913,8 @@ export class ComfyApp {
     }
     if (this.vueAppReady) useToastStore().add(requestToastMessage)
 
+    this.modelsInFolderCache = {}
+
     const defs = await api.getNodeDefs()
 
     for (const nodeId in defs) {

src/scripts/defaultGraph.ts

@@ -133,14 +133,5 @@ export const defaultGraph: ComfyWorkflowJSON = {
   groups: [],
   config: {},
   extra: {},
-  version: 0.4,
-  models: [
-    {
-      name: 'v1-5-pruned-emaonly.ckpt',
-      url: 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt',
-      hash: 'cc6cb27103417325ff94f52b7a5d2dde45a7515b25c255d8e396c90014281516',
-      hash_type: 'SHA256',
-      directory: 'checkpoints'
-    }
-  ]
+  version: 0.4
 }

src/services/dialogService.ts

@@ -3,6 +3,7 @@
 // about importing primevue components.
 import { useDialogStore } from '@/stores/dialogStore'
 import LoadWorkflowWarning from '@/components/dialog/content/LoadWorkflowWarning.vue'
+import MissingModelsWarning from '@/components/dialog/content/MissingModelsWarning.vue'
 import SettingDialogContent from '@/components/dialog/content/SettingDialogContent.vue'
 import SettingDialogHeader from '@/components/dialog/header/SettingDialogHeader.vue'
 import type { ExecutionErrorWsMessage } from '@/types/apiTypes'
@@ -20,6 +21,17 @@ export function showLoadWorkflowWarning(props: {
   })
 }
 
+export function showMissingModelsWarning(props: {
+  missingModels: any[]
+  [key: string]: any
+}) {
+  const dialogStore = useDialogStore()
+  dialogStore.showDialog({
+    component: MissingModelsWarning,
+    props
+  })
+}
+
 export function showSettingsDialog() {
   useDialogStore().showDialog({
     headerComponent: SettingDialogHeader,

src/stores/settingStore.ts

@@ -173,6 +173,13 @@ export const useSettingStore = defineStore('setting', {
         defaultValue: true
       })
 
+      app.ui.settings.addSetting({
+        id: 'Comfy.Workflow.ShowMissingModelsWarning',
+        name: 'Show missing models warning',
+        type: 'boolean',
+        defaultValue: true
+      })
+
       app.ui.settings.addSetting({
         id: 'Comfy.Graph.ZoomSpeed',
         name: 'Canvas zoom speed',

src/types/apiTypes.ts

@@ -73,6 +73,13 @@ const zExecutionErrorWsMessage = zExecutionWsMessageBase.extend({
   current_outputs: z.any()
 })
 
+const zDownloadModelStatus = z.object({
+  status: z.string(),
+  progress_percentage: z.number(),
+  message: z.string(),
+  already_existed: z.boolean(),
+})
+
 export type StatusWsMessageStatus = z.infer<typeof zStatusWsMessageStatus>
 export type StatusWsMessage = z.infer<typeof zStatusWsMessage>
 export type ProgressWsMessage = z.infer<typeof zProgressWsMessage>
@@ -87,6 +94,8 @@ export type ExecutionInterruptedWsMessage = z.infer<
   typeof zExecutionInterruptedWsMessage
 >
 export type ExecutionErrorWsMessage = z.infer<typeof zExecutionErrorWsMessage>
+
+export type DownloadModelStatus = z.infer<typeof zDownloadModelStatus>
 // End of ws messages
 
 const zPromptInputItem = z.object({
@@ -403,6 +412,7 @@ const zSettings = z.record(z.any()).and(
       'Comfy.ConfirmClear': z.boolean(),
       'Comfy.DevMode': z.boolean(),
       'Comfy.Workflow.ShowMissingNodesWarning': z.boolean(),
+      'Comfy.Workflow.ShowMissingModelsWarning': z.boolean(),
       'Comfy.DisableFloatRounding': z.boolean(),
       'Comfy.DisableSliders': z.boolean(),
       'Comfy.DOMClippingEnabled': z.boolean(),

tests-ui/globalSetup.ts

@@ -15,7 +15,8 @@ module.exports = async function () {
 
   jest.mock('@/services/dialogService', () => {
     return {
-      showLoadWorkflowWarning: jest.fn()
+      showLoadWorkflowWarning: jest.fn(),
+      showMissingModelsWarning: jest.fn()
     }
   })
 }

tests-ui/utils/setup.ts

@@ -73,6 +73,7 @@ export function mockApi(config: APIConfig = {}) {
       userConfig.users[username + '!'] = username
       return { status: 200, json: () => username + '!' }
     }),
+    getModels: jest.fn(() => []),
     getUserConfig: jest.fn(
       () => userConfig ?? { storage: 'browser', migrated: false }
     ),