feat: Add model upload functionality (#194)

__init__.py

@@ -41,12 +41,14 @@ utils.download_web_distribution(version)
 from .py import manager
 from .py import download
 from .py import information
+from .py import upload
 
 routes = config.routes
 
 manager.ModelManager().add_routes(routes)
 download.ModelDownload().add_routes(routes)
 information.Information().add_routes(routes)
+upload.ModelUploader().add_routes(routes)
 
 
 WEB_DIRECTORY = "web"

py/upload.py

@@ -0,0 +1,79 @@
+import os
+import time
+
+import folder_paths
+
+from aiohttp import web
+
+from . import utils
+
+
+class ModelUploader:
+    def add_routes(self, routes):
+
+        @routes.get("/model-manager/supported-extensions")
+        async def fetch_model_exts(request):
+            """
+            Get model exts
+            """
+            try:
+                supported_extensions = list(folder_paths.supported_pt_extensions)
+                return web.json_response({"success": True, "data": supported_extensions})
+            except Exception as e:
+                error_msg = f"Get model supported extension failed: {str(e)}"
+                utils.print_error(error_msg)
+                return web.json_response({"success": False, "error": error_msg})
+
+        @routes.post("/model-manager/upload")
+        async def upload_model(request):
+            """
+            Upload model
+            """
+            try:
+                reader = await request.multipart()
+                await self.upload_model(reader)
+                utils.print_info(f"Upload model success")
+                return web.json_response({"success": True, "data": None})
+            except Exception as e:
+                error_msg = f"Upload model failed: {str(e)}"
+                utils.print_error(error_msg)
+                return web.json_response({"success": False, "error": error_msg})
+
+    async def upload_model(self, reader):
+        uploaded_size = 0
+        last_update_time = time.time()
+        interval = 1.0
+
+        while True:
+            part = await reader.next()
+            if part is None:
+                break
+
+            name = part.name
+            if name == "folder":
+                file_folder = await part.text()
+
+            if name == "file":
+                filename = part.filename
+                filepath = f"{file_folder}/{filename}"
+                tmp_filepath = f"{file_folder}/{filename}.tmp"
+
+                with open(tmp_filepath, "wb") as f:
+                    while True:
+                        chunk = await part.read_chunk()
+                        if not chunk:
+                            break
+                        f.write(chunk)
+                        uploaded_size += len(chunk)
+
+                        if time.time() - last_update_time >= interval:
+                            update_upload_progress = {
+                                "uploaded_size": uploaded_size,
+                            }
+                            await utils.send_json("update_upload_progress", update_upload_progress)
+
+        update_upload_progress = {
+            "uploaded_size": uploaded_size,
+        }
+        await utils.send_json("update_upload_progress", update_upload_progress)
+        os.rename(tmp_filepath, filepath)

src/App.vue

@@ -10,6 +10,7 @@ import DialogDownload from 'components/DialogDownload.vue'
 import DialogExplorer from 'components/DialogExplorer.vue'
 import DialogManager from 'components/DialogManager.vue'
 import DialogScanning from 'components/DialogScanning.vue'
+import DialogUpload from 'components/DialogUpload.vue'
 import GlobalDialogStack from 'components/GlobalDialogStack.vue'
 import GlobalLoading from 'components/GlobalLoading.vue'
 import GlobalToast from 'components/GlobalToast.vue'
@@ -64,6 +65,21 @@ onMounted(() => {
     })
   }
 
+  const openUploadDialog = () => {
+    dialog.open({
+      key: 'model-manager-upload',
+      title: t('uploadModel'),
+      content: DialogUpload,
+      headerButtons: [
+        {
+          key: 'refresh',
+          icon: 'pi pi-refresh',
+          command: refreshModelsAndConfig,
+        },
+      ],
+    })
+  }
+
   const openManagerDialog = () => {
     const { cardWidth, gutter, aspect, flat } = config
 
@@ -93,6 +109,11 @@ onMounted(() => {
           icon: 'pi pi-download',
           command: openDownloadDialog,
         },
+        {
+          key: 'upload',
+          icon: 'pi pi-upload',
+          command: openUploadDialog,
+        },
       ],
       minWidth: cardWidth * 2 + gutter + 42,
       minHeight: (cardWidth / aspect) * 0.5 + 162,

src/components/DialogUpload.vue

@@ -0,0 +1,274 @@
+<template>
+  <div class="h-full px-4">
+    <!-- <div v-show="batchScanningStep === 0" class="h-full">
+      <div class="flex h-full items-center px-8">
+        <div class="h-20 w-full opacity-60">
+          <ProgressBar mode="indeterminate" style="height: 6px"></ProgressBar>
+        </div>
+      </div>
+    </div> -->
+
+    <Stepper v-model:value="stepValue" class="flex h-full flex-col" linear>
+      <StepList>
+        <Step :value="1">{{ $t('selectModelType') }}</Step>
+        <Step :value="2">{{ $t('selectSubdirectory') }}</Step>
+        <Step :value="3">{{ $t('chooseFile') }}</Step>
+      </StepList>
+      <StepPanels class="flex-1 overflow-hidden">
+        <StepPanel :value="1" class="h-full">
+          <div class="flex h-full flex-col overflow-hidden">
+            <ResponseScroll>
+              <div class="flex flex-wrap gap-4">
+                <Button
+                  v-for="item in typeOptions"
+                  :key="item.value"
+                  :label="item.label"
+                  @click="item.command"
+                ></Button>
+              </div>
+            </ResponseScroll>
+          </div>
+        </StepPanel>
+        <StepPanel :value="2" class="h-full">
+          <div class="flex h-full flex-col overflow-hidden">
+            <ResponseScroll class="flex-1">
+              <Tree
+                class="h-full"
+                v-model:selection-keys="selectedKey"
+                :value="pathOptions"
+                selectionMode="single"
+                :pt:nodeLabel:class="'text-ellipsis overflow-hidden'"
+              ></Tree>
+            </ResponseScroll>
+
+            <div class="flex justify-between pt-6">
+              <Button
+                :label="$t('back')"
+                severity="secondary"
+                icon="pi pi-arrow-left"
+                @click="handleBackTypeSelect"
+              ></Button>
+              <Button
+                :label="$t('next')"
+                icon="pi pi-arrow-right"
+                icon-pos="right"
+                :disabled="!enabledUpload"
+                @click="handleConfirmSubdir"
+              ></Button>
+            </div>
+          </div>
+        </StepPanel>
+        <StepPanel :value="3" class="h-full">
+          <div class="flex h-full flex-col items-center justify-center">
+            <template v-if="showUploadProgress">
+              <div class="w-4/5">
+                <ProgressBar
+                  :value="uploadProgress"
+                  :pt:value:style="{ transition: 'width .1s linear' }"
+                ></ProgressBar>
+              </div>
+            </template>
+
+            <template v-else>
+              <div class="overflow-hidden break-words py-8">
+                <div class="overflow-hidden px-8">
+                  <div class="text-center">
+                    <div class="pb-2">
+                      {{ $t('selectedSpecialPath') }}
+                    </div>
+                    <div class="leading-5 opacity-60">
+                      {{ selectedModelFolder }}
+                    </div>
+                  </div>
+                </div>
+              </div>
+
+              <div class="flex items-center justify-center gap-4">
+                <Button
+                  v-for="item in uploadActions"
+                  :key="item.value"
+                  :label="item.label"
+                  :icon="item.icon"
+                  @click="item.command.call(item)"
+                ></Button>
+              </div>
+            </template>
+
+            <div class="h-1/4"></div>
+          </div>
+        </StepPanel>
+      </StepPanels>
+    </Stepper>
+  </div>
+</template>
+
+<script setup lang="ts">
+import ResponseScroll from 'components/ResponseScroll.vue'
+import { configSetting } from 'hooks/config'
+import { useModelFolder, useModels } from 'hooks/model'
+import { request } from 'hooks/request'
+import { useToast } from 'hooks/toast'
+import Button from 'primevue/button'
+import ProgressBar from 'primevue/progressbar'
+import Step from 'primevue/step'
+import StepList from 'primevue/steplist'
+import StepPanel from 'primevue/steppanel'
+import StepPanels from 'primevue/steppanels'
+import Stepper from 'primevue/stepper'
+import Tree from 'primevue/tree'
+import { api, app } from 'scripts/comfyAPI'
+import { computed, onMounted, onUnmounted, ref, toValue } from 'vue'
+import { useI18n } from 'vue-i18n'
+
+const { t } = useI18n()
+const { toast } = useToast()
+
+const stepValue = ref(1)
+
+const { folders } = useModels()
+
+const currentType = ref<string>()
+const typeOptions = computed(() => {
+  const excludeScanTypes = app.ui?.settings.getSettingValue<string>(
+    configSetting.excludeScanTypes,
+  )
+  const customBlackList =
+    excludeScanTypes
+      ?.split(',')
+      .map((type) => type.trim())
+      .filter(Boolean) ?? []
+  return Object.keys(folders.value)
+    .filter((folder) => !customBlackList.includes(folder))
+    .map((type) => {
+      return {
+        label: type,
+        value: type,
+        command: () => {
+          currentType.value = type
+          stepValue.value++
+        },
+      }
+    })
+})
+
+const { pathOptions } = useModelFolder({ type: currentType })
+
+const selectedModelFolder = ref<string>()
+const selectedKey = computed({
+  get: () => {
+    const key = selectedModelFolder.value
+    return key ? { [key]: true } : {}
+  },
+  set: (val) => {
+    const key = Object.keys(val)[0]
+    selectedModelFolder.value = key
+  },
+})
+
+const enabledUpload = computed(() => {
+  return !!selectedModelFolder.value
+})
+
+const handleBackTypeSelect = () => {
+  selectedModelFolder.value = undefined
+  currentType.value = undefined
+  stepValue.value--
+}
+
+const handleConfirmSubdir = () => {
+  stepValue.value++
+}
+
+const uploadTotalSize = ref<number>()
+const uploadSize = ref<number>()
+const uploadProgress = computed(() => {
+  const total = toValue(uploadTotalSize)
+  const size = toValue(uploadSize)
+  if (typeof total === 'number' && typeof size === 'number') {
+    return Math.floor((size / total) * 100)
+  }
+  return undefined
+})
+const showUploadProgress = computed(() => {
+  return typeof uploadProgress.value !== 'undefined'
+})
+
+const uploadActions = ref([
+  {
+    value: 'back',
+    label: t('back'),
+    icon: 'pi pi-arrow-left',
+    command: () => {
+      stepValue.value--
+    },
+  },
+  {
+    value: 'full',
+    label: t('chooseFile'),
+    command: () => {
+      const input = document.createElement('input')
+      input.type = 'file'
+      input.accept = supportedExtensions.value.join(',')
+      input.onchange = async () => {
+        const files = input.files
+        const file = files?.item(0)
+        if (!file) {
+          return
+        }
+
+        try {
+          uploadTotalSize.value = file.size
+          uploadSize.value = 0
+          const body = new FormData()
+          body.append('folder', toValue(selectedModelFolder)!)
+          body.append('file', file)
+
+          await request('/upload', {
+            method: 'POST',
+            body: body,
+          })
+        } catch (error) {
+          toast.add({
+            severity: 'error',
+            summary: 'Error',
+            detail: error.message,
+            life: 5000,
+          })
+        }
+      }
+      input.click()
+    },
+  },
+])
+
+const supportedExtensions = ref([])
+
+const fetchSupportedExtensions = async () => {
+  try {
+    const result = await request('/supported-extensions')
+    supportedExtensions.value = result ?? []
+  } catch (error) {
+    toast.add({
+      severity: 'error',
+      summary: 'Error',
+      detail: error.message,
+      life: 5000,
+    })
+  }
+}
+
+const update_process = (event: CustomEvent) => {
+  const detail = event.detail
+  uploadSize.value = detail.uploaded_size
+}
+
+onMounted(() => {
+  fetchSupportedExtensions()
+
+  api.addEventListener('update_upload_progress', update_process)
+})
+
+onUnmounted(() => {
+  api.removeEventListener('update_upload_progress', update_process)
+})
+</script>

src/i18n.ts

@@ -41,6 +41,8 @@ const messages = {
     scanMissInformation: 'Download missing information',
     scanFullInformation: 'Override full information',
     noModelsInCurrentPath: 'There are no models available in the current path',
+    uploadModel: 'Upload Model',
+    chooseFile: 'Choose File',
     sort: {
       name: 'Name',
       size: 'Largest',
@@ -116,6 +118,8 @@ const messages = {
     scanMissInformation: '下载缺失信息',
     scanFullInformation: '覆盖所有信息',
     noModelsInCurrentPath: '当前路径中没有可用的模型',
+    uploadModel: '上传模型',
+    chooseFile: '选择文件',
     sort: {
       name: '名称',
       size: '最大',

src/types/global.d.ts

@@ -12,6 +12,11 @@ declare namespace ComfyAPI {
         callback: (event: CustomEvent) => void,
         options?: AddEventListenerOptions,
       ) => void
+      removeEventListener: (
+        type: string,
+        callback: (event: CustomEvent) => void,
+        options?: AddEventListenerOptions,
+      ) => void
     }
 
     const api: ComfyApi