From 777a6d9ce357c0c0f5ed60affbbdf43ae1b6fc08 Mon Sep 17 00:00:00 2001
From: oto-ciulis-tt <62436273+oto-ciulis-tt@users.noreply.github.com>
Date: Fri, 8 Nov 2024 12:59:35 -0800
Subject: [PATCH] feat:use electron api to download models (#1473)

* enh: Use electron API to download models

* Adding tooltips

* PR comments

---------

Co-authored-by: Oto Ciulis <oto.ciulis@gmail.com>
---
 .../common/ElectronFileDownload.vue           | 135 ++++++++++++++++++
 .../dialog/content/MissingModelsWarning.vue   |   9 ++
 src/i18n.ts                                   |   5 +
 3 files changed, 149 insertions(+)
 create mode 100644 src/components/common/ElectronFileDownload.vue

diff --git a/src/components/common/ElectronFileDownload.vue b/src/components/common/ElectronFileDownload.vue
new file mode 100644
index 000000000..8b2430286
--- /dev/null
+++ b/src/components/common/ElectronFileDownload.vue
@@ -0,0 +1,135 @@
+<!-- A Electron-backed download button with a label, size hint and progress bar -->
+<template>
+  <div class="flex flex-col">
+    <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="triggerDownload"
+          v-if="status === null || status === 'error'"
+          icon="pi pi-download"
+        />
+      </div>
+    </div>
+    <div
+      class="flex flex-row items-center gap-2"
+      v-if="status === 'in_progress' || status === 'paused'"
+    >
+      <ProgressBar class="flex-1" :value="downloadProgress" />
+
+      <Button
+        class="file-action-button"
+        size="small"
+        outlined
+        :disabled="props.error"
+        @click="triggerPauseDownload"
+        v-if="status === 'in_progress'"
+        icon="pi pi-pause-circle"
+        v-tooltip.top="t('electronFileDownload.pause')"
+      />
+
+      <Button
+        class="file-action-button"
+        size="small"
+        outlined
+        :disabled="props.error"
+        @click="triggerResumeDownload"
+        v-if="status === 'paused'"
+        icon="pi pi-play-circle"
+        v-tooltip.top="t('electronFileDownload.resume')"
+      />
+
+      <Button
+        class="file-action-button"
+        size="small"
+        outlined
+        :disabled="props.error"
+        @click="triggerCancelDownload"
+        icon="pi pi-times-circle"
+        v-tooltip.top="t('electronFileDownload.cancel')"
+      />
+    </div>
+  </div>
+</template>
+
+<script setup lang="ts">
+import { useDownload } from '@/hooks/downloadHooks'
+import Button from 'primevue/button'
+import ProgressBar from 'primevue/progressbar'
+import { ref, computed } from 'vue'
+import { formatSize } from '@/utils/formatUtil'
+import { useI18n } from 'vue-i18n'
+import { electronAPI } from '@/utils/envUtil'
+
+const props = defineProps<{
+  url: string
+  hint?: string
+  label?: string
+  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 fileSize = computed(() =>
+  download.fileSize.value ? formatSize(download.fileSize.value) : '?'
+)
+const [savePath, filename] = props.label.split('/')
+
+const downloads: ModelDownload[] = await DownloadManager.getAllDownloads()
+const modelDownload = downloads.find(({ url }) => url === props.url)
+
+const updateProperties = (download: ModelDownload) => {
+  if (download.url === props.url) {
+    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()
+  )
+}
+
+const triggerCancelDownload = async () => {
+  await DownloadManager.cancelDownload(props.url)
+}
+
+const triggerPauseDownload = async () => {
+  await DownloadManager.pauseDownload(props.url)
+}
+
+const triggerResumeDownload = async () => {
+  await DownloadManager.resumeDownload(props.url)
+}
+</script>
diff --git a/src/components/dialog/content/MissingModelsWarning.vue b/src/components/dialog/content/MissingModelsWarning.vue
index a8b1b7f9c..3131e2221 100644
--- a/src/components/dialog/content/MissingModelsWarning.vue
+++ b/src/components/dialog/content/MissingModelsWarning.vue
@@ -7,7 +7,15 @@
   />
   <ListBox :options="missingModels" class="comfy-missing-models">
     <template #option="{ option }">
+      <Suspense v-if="isElectron()">
+        <ElectronFileDownload
+          :url="option.url"
+          :label="option.label"
+          :error="option.error"
+        />
+      </Suspense>
       <FileDownload
+        v-else
         :url="option.url"
         :label="option.label"
         :error="option.error"
@@ -21,6 +29,7 @@ import { ref, computed } from 'vue'
 import ListBox from 'primevue/listbox'
 import NoResultsPlaceholder from '@/components/common/NoResultsPlaceholder.vue'
 import FileDownload from '@/components/common/FileDownload.vue'
+import { isElectron } from '@/utils/envUtil'
 
 // TODO: Read this from server internal API rather than hardcoding here
 // as some installations may wish to use custom sources
diff --git a/src/i18n.ts b/src/i18n.ts
index 3f3848d89..55fe90840 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -119,6 +119,11 @@ const messages = {
       selectMode: 'Select Mode',
       panMode: 'Pan Mode',
       toggleLinkVisibility: 'Toggle Link Visibility'
+    },
+    electronFileDownload: {
+      pause: 'Pause Download',
+      resume: 'Resume Download',
+      cancel: 'Cancel Download'
     }
   },
   zh: {