From b19cbd9111a99e702ba8c8092bb17e41026a0935 Mon Sep 17 00:00:00 2001
From: "Alex \"mcmonkey\" Goodwin"
 <4000772+mcmonkey4eva@users.noreply.github.com>
Date: Tue, 27 Aug 2024 05:56:57 -0700
Subject: [PATCH] Model Downloader Improvements (#650)

* use new download_path value in download api

from https://github.com/comfyanonymous/ComfyUI/pull/4621

* add in-UI refusal for ckpt files
---
 .../dialog/content/MissingModelsWarning.vue   | 28 ++++++++++---------
 src/types/apiTypes.ts                         |  1 +
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/src/components/dialog/content/MissingModelsWarning.vue b/src/components/dialog/content/MissingModelsWarning.vue
index 3a226ad32..c3f32a247 100644
--- a/src/components/dialog/content/MissingModelsWarning.vue
+++ b/src/components/dialog/content/MissingModelsWarning.vue
@@ -73,23 +73,18 @@ const modelDownloads = ref<Record<string, ModelInfo>>({})
 let lastModel: string | null = null
 
 const handleDownloadProgress = (detail: DownloadModelStatus) => {
+  if (detail.download_path) {
+    lastModel = detail.download_path.split('/', 2)[1]
+  }
+  if (!lastModel) return
   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 }
+    modelDownloads.value[lastModel] = { ...modelDownloads.value[lastModel], downloading: true, progress: detail.progress_percentage, 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 }
+    modelDownloads.value[lastModel] = { ...modelDownloads.value[lastModel], 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 }
+    modelDownloads.value[lastModel] = { ...modelDownloads.value[lastModel], 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 }
-    }
+    modelDownloads.value[lastModel] = { ...modelDownloads.value[lastModel], downloading: false, progress: 0, error: detail.message, completed: false }
   }
   // TODO: other statuses?
 }
@@ -115,6 +110,13 @@ const missingModels = computed(() => {
           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'
+        }
+      }
       if (model.directory_invalid) {
         return {
           label: `${model.directory} / ${model.name}`,
diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts
index 8decf09a0..3ffc0b3e0 100644
--- a/src/types/apiTypes.ts
+++ b/src/types/apiTypes.ts
@@ -77,6 +77,7 @@ const zDownloadModelStatus = z.object({
   status: z.string(),
   progress_percentage: z.number(),
   message: z.string(),
+  download_path: z.string(),
   already_existed: z.boolean()
 })