Use browser download in missing model dialog (#1362)

* Remove custom backend download logic

* Add download hooks

* Download button

* Use browser download

* Update test

browser_tests/dialog.spec.ts

@@ -63,10 +63,11 @@ test.describe('Missing models warning', () => {
 
     const downloadButton = comfyPage.page.getByLabel('Download')
     await expect(downloadButton).toBeVisible()
+    const downloadPromise = comfyPage.page.waitForEvent('download')
     await downloadButton.click()
 
-    const downloadComplete = comfyPage.page.locator('.download-complete')
-    await expect(downloadComplete).toBeVisible()
+    const download = await downloadPromise
+    expect(download.suggestedFilename()).toBe('fake_model.safetensors')
   })
 })
 

src/components/common/DeviceInfo.vue

@@ -9,7 +9,7 @@
 
 <script setup lang="ts">
 import type { DeviceStats } from '@/types/apiTypes'
-import { formatMemory } from '@/utils/formatUtil'
+import { formatSize } from '@/utils/formatUtil'
 
 const props = defineProps<{
   device: DeviceStats
@@ -30,7 +30,7 @@ const formatValue = (value: any, field: string) => {
       field
     )
   ) {
-    return formatMemory(value)
+    return formatSize(value)
   }
   return value
 }

src/components/common/FileDownload.vue

@@ -0,0 +1,44 @@
+<!-- A file download button with a label and a size hint -->
+<template>
+  <div class="flex flex-row items-center gap-2">
+    <div class="file-info">
+      <div class="file-details">
+        <span class="file-type" :title="hint">{{ label }}</span>
+      </div>
+      <div v-if="props.error" class="file-error">
+        {{ props.error }}
+      </div>
+    </div>
+    <div class="file-action">
+      <Button
+        class="file-action-button"
+        :label="$t('download') + ' (' + fileSize + ')'"
+        size="small"
+        outlined
+        :disabled="props.error"
+        @click="download.triggerBrowserDownload"
+      />
+    </div>
+  </div>
+</template>
+
+<script setup lang="ts">
+import { useDownload } from '@/hooks/downloadHooks'
+import Button from 'primevue/button'
+import { computed } from 'vue'
+import { formatSize } from '@/utils/formatUtil'
+
+const props = defineProps<{
+  url: string
+  hint?: string
+  label?: string
+  error?: string
+}>()
+
+const label = computed(() => props.label || props.url.split('/').pop())
+const hint = computed(() => props.hint || props.url)
+const download = useDownload(props.url)
+const fileSize = computed(() =>
+  download.fileSize.value ? formatSize(download.fileSize.value) : '?'
+)
+</script>

src/components/common/SystemStatsPanel.vue

@@ -35,7 +35,7 @@ import TabPanel from 'primevue/tabpanel'
 import Divider from 'primevue/divider'
 import type { SystemStats } from '@/types/apiTypes'
 import DeviceInfo from '@/components/common/DeviceInfo.vue'
-import { formatMemory } from '@/utils/formatUtil'
+import { formatSize } from '@/utils/formatUtil'
 
 const props = defineProps<{
   stats: SystemStats
@@ -58,7 +58,7 @@ const systemColumns = [
 
 const formatValue = (value: any, field: string) => {
   if (['ram_total', 'ram_free'].includes(field)) {
-    return formatMemory(value)
+    return formatSize(value)
   }
   return value
 }

src/components/dialog/content/MissingModelsWarning.vue

@@ -5,67 +5,13 @@
     title="Missing Models"
     message="When loading the graph, the following models were not found"
   />
-  <ListBox
-    :options="missingModels"
-    optionLabel="label"
-    scrollHeight="100%"
-    class="comfy-missing-models"
-  >
-    <template #option="slotProps">
-      <div
-        class="missing-model-item flex flex-row items-center"
-        :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">
-          <Select
-            class="model-path-select mr-2"
-            v-if="
-              slotProps.option.action &&
-              !slotProps.option.downloading &&
-              !slotProps.option.completed &&
-              !slotProps.option.error
-            "
-            v-show="slotProps.option.paths.length > 1"
-            v-model="slotProps.option.folderPath"
-            :options="slotProps.option.paths"
-            @change="updateFolderPath(slotProps.option, $event)"
-          />
-          <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"
-            size="small"
-            outlined
-            class="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(--p-green-500)"></i>
-          </div>
-          <div v-if="slotProps.option.error" class="download-error">
-            <i class="pi pi-times" style="color: var(--p-red-600)"></i>
-          </div>
-        </div>
-      </div>
+  <ListBox :options="missingModels" class="comfy-missing-models">
+    <template #option="{ option }">
+      <FileDownload
+        :url="option.url"
+        :label="option.label"
+        :error="option.error"
+      />
     </template>
   </ListBox>
 </template>
@@ -73,12 +19,8 @@
 <script setup lang="ts">
 import { ref, computed } from 'vue'
 import ListBox from 'primevue/listbox'
-import Select from 'primevue/select'
 import NoResultsPlaceholder from '@/components/common/NoResultsPlaceholder.vue'
-import { SelectChangeEvent } from 'primevue/select'
-import Button from 'primevue/button'
-import { api } from '@/scripts/api'
-import { DownloadModelStatus } from '@/types/apiTypes'
+import FileDownload from '@/components/common/FileDownload.vue'
 
 // TODO: Read this from server internal API rather than hardcoding here
 // as some installations may wish to use custom sources
@@ -107,86 +49,13 @@ const props = defineProps<{
 }>()
 
 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
-  }
-  if (!lastModel) return
-  if (detail.status === 'in_progress') {
-    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
-    }
-  } else if (detail.status === 'completed') {
-    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
-    }
-  }
-  // TODO: other statuses?
-}
-
-const triggerDownload = async (
-  url: string,
-  directory: string,
-  filename: string,
-  folder_path: string
-) => {
-  modelDownloads.value[filename] = {
-    name: filename,
-    directory,
-    url,
-    downloading: true,
-    progress: 0
-  }
-  const download = await api.internalDownloadModel(
-    url,
-    directory,
-    filename,
-    1,
-    folder_path
-  )
-  lastModel = filename
-  handleDownloadProgress(download)
-}
-
-api.addEventListener('download_progress', (event: CustomEvent) => {
-  handleDownloadProgress(event.detail)
-})
-
 const missingModels = computed(() => {
   return props.missingModels.map((model) => {
     const paths = props.paths[model.directory]
     if (model.directory_invalid || !paths) {
       return {
         label: `${model.directory} / ${model.name}`,
-        hint: model.url,
+        url: model.url,
         error: 'Invalid directory specified (does this require custom nodes?)'
       }
     }
@@ -204,37 +73,27 @@ const missingModels = computed(() => {
     if (!allowedSources.some((source) => model.url.startsWith(source))) {
       return {
         label: `${model.directory} / ${model.name}`,
-        hint: model.url,
+        url: 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,
+        url: model.url,
         error: `Only allowed suffixes are: '${allowedSuffixes.join("', '")}'`
       }
     }
     return {
+      url: model.url,
       label: `${model.directory} / ${model.name}`,
-      hint: model.url,
       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,
-            downloadInfo.folder_path
-          )
-      }
+      folderPath: downloadInfo.folder_path
     }
   })
 })
@@ -245,85 +104,4 @@ const missingModels = computed(() => {
   max-height: 300px;
   overflow-y: auto;
 }
-
-.missing-model-item::before {
-  content: '';
-  position: absolute;
-  top: 0;
-  left: 0;
-  height: 100%;
-  width: var(--progress);
-  background-color: var(--p-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(--p-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/hooks/downloadHooks.ts

@@ -0,0 +1,44 @@
+import { onMounted, ref } from 'vue'
+
+export function useDownload(url: string, fileName?: string) {
+  const fileSize = ref<number | null>(null)
+
+  const fetchFileSize = async (): Promise<number | null> => {
+    try {
+      const response = await fetch(url, { method: 'HEAD' })
+      if (!response.ok) throw new Error('Failed to fetch file size')
+
+      const size = response.headers.get('content-length')
+      if (size) {
+        return parseInt(size)
+      } else {
+        console.error('"content-length" header not found')
+        return null
+      }
+    } catch (e) {
+      console.error('Error fetching file size:', e)
+      return null
+    }
+  }
+
+  /**
+   * Trigger browser download
+   */
+  const triggerBrowserDownload = () => {
+    const link = document.createElement('a')
+    link.href = url
+    link.download = fileName || url.split('/').pop() || 'download'
+    link.target = '_blank' // Opens in new tab if download attribute is not supported
+    link.rel = 'noopener noreferrer' // Security best practice for _blank links
+    link.click()
+  }
+
+  onMounted(async () => {
+    fileSize.value = await fetchFileSize()
+  })
+
+  return {
+    triggerBrowserDownload,
+    fileSize
+  }
+}

src/i18n.ts

@@ -2,6 +2,7 @@ import { createI18n } from 'vue-i18n'
 
 const messages = {
   en: {
+    download: 'Download',
     loadAllFolders: 'Load All Folders',
     refresh: 'Refresh',
     terminal: 'Terminal',
@@ -116,6 +117,7 @@ const messages = {
     }
   },
   zh: {
+    download: '下载',
     loadAllFolders: '加载所有文件夹',
     refresh: '刷新',
     terminal: '终端',
@@ -229,6 +231,7 @@ const messages = {
     }
   },
   ru: {
+    download: 'Скачать',
     refresh: 'Обновить',
     loadAllFolders: 'Загрузить все папки',
     terminal: 'Терминал',

src/scripts/api.ts

@@ -1,6 +1,5 @@
 import type { ComfyWorkflowJSON } from '@/types/comfyWorkflow'
 import {
-  type DownloadModelStatus,
   type HistoryTaskItem,
   type PendingTaskItem,
   type RunningTaskItem,
@@ -240,11 +239,6 @@ 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(
@@ -413,36 +407,6 @@ class ComfyApi extends EventTarget {
     }
   }
 
-  /**
-   * 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,
-    folder_path: string
-  ): 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,
-        folder_path
-      })
-    })
-    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/types/apiTypes.ts

@@ -77,14 +77,6 @@ const zExecutionErrorWsMessage = zExecutionWsMessageBase.extend({
   current_outputs: z.any()
 })
 
-const zDownloadModelStatus = z.object({
-  status: z.string(),
-  progress_percentage: z.number(),
-  message: z.string(),
-  download_path: 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>
@@ -99,8 +91,6 @@ 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({

src/utils/formatUtil.ts

@@ -60,14 +60,15 @@ export function formatNumberWithSuffix(
   return `${formattedNum}${suffixes[exp]}`
 }
 
-export function formatMemory(value?: number) {
+export function formatSize(value?: number) {
   if (value === null || value === undefined) {
     return '-'
   }
 
-  const mb = Math.round(value / (1024 * 1024))
-  if (mb >= 1024) {
-    return `${(mb / 1024).toFixed(2)} GB`
-  }
-  return `${mb} MB`
+  const bytes = value
+  if (bytes === 0) return '0 B'
+  const k = 1024
+  const sizes = ['B', 'KB', 'MB', 'GB']
+  const i = Math.floor(Math.log(bytes) / Math.log(k))
+  return `${parseFloat((bytes / Math.pow(k, i)).toFixed(2))} ${sizes[i]}`
 }