feat: Add download progress to sidebar (#1490)

* feat: Add download progress to sidebar

* Removing console log

* Lint fixes

* Updating UI

* Fixing lint error

* Fixing lint error

* Fixing lint error

* PR comments

* Reverting change

---------

Co-authored-by: Oto Ciulis <oto.ciulis@gmail.com>

src/components/common/ElectronFileDownload.vue

@@ -59,6 +59,7 @@
         :disabled="props.error"
         @click="triggerCancelDownload"
         icon="pi pi-times-circle"
+        severity="danger"
         v-tooltip.top="t('electronFileDownload.cancel')"
       />
     </div>
@@ -73,6 +74,7 @@ import { ref, computed } from 'vue'
 import { formatSize } from '@/utils/formatUtil'
 import { useI18n } from 'vue-i18n'
 import { electronAPI } from '@/utils/envUtil'
+import { useElectronDownloadStore } from '@/stores/electronDownloadStore'
 
 const props = defineProps<{
   url: string
@@ -81,55 +83,37 @@ const props = defineProps<{
   error?: string
 }>()
 
-interface ModelDownload {
-  url: string
-  status: 'paused' | 'in_progress' | 'cancelled'
-  progress: number
-}
-
 const { t } = useI18n()
 const { DownloadManager } = electronAPI()
 const label = computed(() => props.label || props.url.split('/').pop())
 const hint = computed(() => props.hint || props.url)
 const download = useDownload(props.url)
-const status = ref<ModelDownload | null>(null)
 const downloadProgress = ref<number>(0)
+const status = ref<string | null>(null)
 const fileSize = computed(() =>
   download.fileSize.value ? formatSize(download.fileSize.value) : '?'
 )
+const electronDownloadStore = useElectronDownloadStore()
 const [savePath, filename] = props.label.split('/')
 
-const downloads: ModelDownload[] = await DownloadManager.getAllDownloads()
-const modelDownload = downloads.find(({ url }) => url === props.url)
+electronDownloadStore.$subscribe((mutation, { downloads }) => {
+  const download = downloads.find((download) => props.url === download.url)
 
-const updateProperties = (download: ModelDownload) => {
-  if (download.url === props.url) {
+  if (download) {
+    downloadProgress.value = Number((download.progress * 100).toFixed(1))
     status.value = download.status
-    downloadProgress.value = (download.progress * 100).toFixed(1)
   }
-}
-
-DownloadManager.onDownloadProgress((data: ModelDownload) => {
-  updateProperties(data)
 })
 
 const triggerDownload = async () => {
-  await DownloadManager.startDownload(
-    props.url,
-    filename.trim(),
-    savePath.trim()
-  )
+  await electronDownloadStore.start({
+    url: props.url,
+    savePath: savePath.trim(),
+    filename: filename.trim()
+  })
 }
 
-const triggerCancelDownload = async () => {
-  await DownloadManager.cancelDownload(props.url)
-}
-
-const triggerPauseDownload = async () => {
-  await DownloadManager.pauseDownload(props.url)
-}
-
-const triggerResumeDownload = async () => {
-  await DownloadManager.resumeDownload(props.url)
-}
+const triggerCancelDownload = () => electronDownloadStore.cancel(props.url)
+const triggerPauseDownload = () => electronDownloadStore.pause(props.url)
+const triggerResumeDownload = () => electronDownloadStore.resume(props.url)
 </script>

src/components/common/TreeExplorerTreeNode.vue

@@ -10,25 +10,27 @@
     ]"
     ref="container"
   >
-    <div class="node-content">
-      <span class="node-label">
-        <slot name="before-label" :node="props.node"></slot>
-        <EditableText
-          :modelValue="node.label"
-          :isEditing="isEditing"
-          @edit="handleRename"
+    <div class="flex flex-col w-full">
+      <div class="node-content">
+        <span class="node-label">
+          <slot name="before-label" :node="props.node"></slot>
+          <EditableText
+            :modelValue="node.label"
+            :isEditing="isEditing"
+            @edit="handleRename"
+          />
+          <slot name="after-label" :node="props.node"></slot>
+        </span>
+        <Badge
+          v-if="showNodeBadgeText"
+          :value="nodeBadgeText"
+          severity="secondary"
+          class="leaf-count-badge"
         />
-        <slot name="after-label" :node="props.node"></slot>
-      </span>
-      <Badge
-        v-if="showNodeBadgeText"
-        :value="nodeBadgeText"
-        severity="secondary"
-        class="leaf-count-badge"
-      />
-    </div>
-    <div class="node-actions">
-      <slot name="actions" :node="props.node"></slot>
+      </div>
+      <div class="node-actions">
+        <slot name="actions" :node="props.node"></slot>
+      </div>
     </div>
   </div>
 </template>

src/components/sidebar/tabs/ModelLibrarySidebarTab.vue

@@ -28,6 +28,16 @@
       />
     </template>
     <template #body>
+      <div class="mx-6 mb-4" v-if="downloads.length > 0">
+        <div class="text-lg my-4">
+          {{ t('electronFileDownload.inProgress') }}
+        </div>
+
+        <template v-for="download in downloads" :key="download.url">
+          <DownloadItem :download="download" />
+        </template>
+      </div>
+
       <TreeExplorer
         class="model-lib-tree-explorer py-0"
         :roots="renderedRoot.children"
@@ -48,6 +58,7 @@ import SearchBox from '@/components/common/SearchBox.vue'
 import TreeExplorer from '@/components/common/TreeExplorer.vue'
 import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'
 import ModelTreeLeaf from '@/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue'
+import DownloadItem from '@/components/sidebar/tabs/modelLibrary/DownloadItem.vue'
 import {
   ComfyModelDef,
   ModelFolder,
@@ -65,12 +76,19 @@ import { computed, ref, watch, toRef, onMounted, nextTick } from 'vue'
 import type { TreeNode } from 'primevue/treenode'
 import { app } from '@/scripts/app'
 import { buildTree } from '@/utils/treeUtil'
+import { useI18n } from 'vue-i18n'
+import { useElectronDownloadStore } from '@/stores/electronDownloadStore'
+import { storeToRefs } from 'pinia'
+
 const modelStore = useModelStore()
 const modelToNodeStore = useModelToNodeStore()
 const settingStore = useSettingStore()
 const searchQuery = ref<string>('')
 const expandedKeys = ref<Record<string, boolean>>({})
 const { expandNode, toggleNodeOnEvent } = useTreeExpansion(expandedKeys)
+const { t } = useI18n()
+const electronDownloadStore = useElectronDownloadStore()
+const { downloads } = storeToRefs(electronDownloadStore)
 
 const filteredModels = ref<ComfyModelDef[]>([])
 const handleSearch = async (query: string) => {
@@ -164,6 +182,7 @@ const renderedRoot = computed<TreeExplorerNode<ModelOrFolder>>(() => {
       }
     }
   }
+
   return fillNodeInfo(root.value)
 })
 

src/components/sidebar/tabs/modelLibrary/DownloadItem.vue

@@ -0,0 +1,97 @@
+<template>
+  <div class="flex flex-col">
+    <div>
+      {{ getDownloadLabel(download.savePath) }}
+    </div>
+    <div v-if="['cancelled', 'error'].includes(download.status)">
+      <Chip
+        class="h-6 text-sm font-light bg-red-700 mt-2"
+        removable
+        @remove="handleRemoveDownload"
+      >
+        {{ t('electronFileDownload.cancelled') }}
+      </Chip>
+    </div>
+    <div
+      class="mt-2 flex flex-row items-center gap-2"
+      v-if="['in_progress', 'paused', 'completed'].includes(download.status)"
+    >
+      <ProgressBar
+        class="flex-1"
+        :value="Number((download.progress * 100).toFixed(1))"
+      />
+
+      <Button
+        class="file-action-button w-[22px] h-[22px]"
+        size="small"
+        rounded
+        @click="triggerPauseDownload"
+        v-if="download.status === 'in_progress'"
+        icon="pi pi-pause"
+        v-tooltip.top="t('electronFileDownload.pause')"
+      />
+
+      <Button
+        class="file-action-button w-[22px] h-[22px]"
+        size="small"
+        rounded
+        @click="triggerResumeDownload"
+        v-if="download.status === 'paused'"
+        icon="pi pi-play"
+        v-tooltip.top="t('electronFileDownload.resume')"
+      />
+
+      <Button
+        class="file-action-button w-[22px] h-[22px] p-red"
+        size="small"
+        rounded
+        severity="danger"
+        @click="triggerCancelDownload"
+        v-if="['in_progress', 'paused'].includes(download.status)"
+        icon="pi pi-times-circle"
+        v-tooltip.top="t('electronFileDownload.cancel')"
+      />
+    </div>
+  </div>
+</template>
+
+<script setup lang="ts">
+import Button from 'primevue/button'
+import ProgressBar from 'primevue/progressbar'
+import Chip from 'primevue/chip'
+import { useI18n } from 'vue-i18n'
+const { t } = useI18n()
+import {
+  type ElectronDownload,
+  useElectronDownloadStore
+} from '@/stores/electronDownloadStore'
+const electronDownloadStore = useElectronDownloadStore()
+
+const props = defineProps<{
+  download: ElectronDownload
+}>()
+
+const getDownloadLabel = (savePath: string, filename: string) => {
+  let parts = (savePath ?? '').split('/')
+  parts = parts.length === 1 ? parts[0].split('\\') : parts
+  const name = parts.pop()
+  const dir = parts.pop()
+  return `${dir}/${name}`
+}
+
+const triggerCancelDownload = () =>
+  electronDownloadStore.cancel(props.download.url)
+const triggerPauseDownload = () =>
+  electronDownloadStore.pause(props.download.url)
+const triggerResumeDownload = () =>
+  electronDownloadStore.resume(props.download.url)
+
+const handleRemoveDownload = () => {
+  electronDownloadStore.$patch((state) => {
+    state.downloads = state.downloads.filter(
+      ({ url }) => url !== props.download.url
+    )
+    state.hasChanged = true
+  })
+}
+</script>

src/i18n.ts

@@ -119,6 +119,7 @@ const messages = {
         sortOrder: 'Sort Order'
       },
       modelLibrary: 'Model Library',
+      downloads: 'Downloads',
       queueTab: {
         showFlatList: 'Show Flat List',
         backToAllTasks: 'Back to All Tasks',
@@ -170,9 +171,12 @@ const messages = {
       toggleLinkVisibility: 'Toggle Link Visibility'
     },
     electronFileDownload: {
+      inProgress: 'In Progress',
       pause: 'Pause Download',
+      paused: 'Paused',
       resume: 'Resume Download',
-      cancel: 'Cancel Download'
+      cancel: 'Cancel Download',
+      cancelled: 'Cancelled'
     }
   },
   zh: {

src/stores/electronDownloadStore.ts

@@ -0,0 +1,69 @@
+import { ref } from 'vue'
+import { defineStore } from 'pinia'
+import { isElectron, electronAPI } from '@/utils/envUtil'
+
+export interface ElectronDownload {
+  url: string
+  status: 'paused' | 'in_progress' | 'cancelled'
+  progress: number
+  savePath: string
+  filename: string
+}
+
+/** Electron donwloads store handler */
+export const useElectronDownloadStore = defineStore('downloads', () => {
+  const downloads = ref<ElectronDownload[]>([])
+  const { DownloadManager } = electronAPI()
+
+  const findByUrl = (url: string) =>
+    downloads.value.find((download) => url === download.url)
+
+  const initialize = async () => {
+    if (isElectron()) {
+      const allDownloads: ElectronDownload[] =
+        await DownloadManager.getAllDownloads()
+
+      for (const download of allDownloads) {
+        downloads.value.push(download)
+      }
+
+      // ToDO: replace with ElectronDownload type
+      DownloadManager.onDownloadProgress((data: any) => {
+        if (!findByUrl(data.url)) {
+          downloads.value.push(data)
+        }
+
+        const download = findByUrl(data.url)
+
+        if (download) {
+          download.progress = data.progress
+          download.status = data.status
+          download.filename = data.filename
+          download.savePath = data.savePath
+        }
+      })
+    }
+  }
+
+  void initialize()
+
+  const start = ({
+    url,
+    savePath,
+    filename
+  }: Pick<ElectronDownload, 'url' | 'savePath' | 'filename'>) =>
+    DownloadManager.startDownload(url, savePath, filename)
+  const pause = (url: string) => DownloadManager.pauseDownload(url)
+  const resume = (url: string) => DownloadManager.resumeDownload(url)
+  const cancel = (url: string) => DownloadManager.cancelDownload(url)
+
+  return {
+    downloads,
+    start,
+    pause,
+    resume,
+    cancel,
+    findByUrl,
+    initialize
+  }
+})