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
This commit is contained in:
Alex "mcmonkey" Goodwin
2024-08-27 05:56:57 -07:00
committed by GitHub
parent 9cdefca481
commit b19cbd9111
2 changed files with 16 additions and 13 deletions

View File

@@ -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}`,

View File

@@ -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()
})