From ac4a168f1334089900682755048960acaac8fca2 Mon Sep 17 00:00:00 2001
From: Hayden <48267247+hayden-fr@users.noreply.github.com>
Date: Mon, 11 Aug 2025 09:10:20 +0800
Subject: [PATCH] feat: support download multiple actual files (#196)

---
 py/information.py                   |  7 +--
 src/components/DialogCreateTask.vue | 10 +++-
 src/components/ModelContent.vue     |  2 +-
 src/hooks/download.ts               | 75 ++++++++++++++++++++++++-----
 src/types/typings.d.ts              | 11 +++++
 5 files changed, 89 insertions(+), 16 deletions(-)

diff --git a/py/information.py b/py/information.py
index 6b6747a..c4526c2 100644
--- a/py/information.py
+++ b/py/information.py
@@ -69,8 +69,8 @@ class CivitaiModelSearcher(ModelSearcher):
         models: list[dict] = []
 
         for version in model_versions:
-            model_files: list[dict] = version.get("files", [])
-            model_files = utils.filter_with(model_files, {"type": "Model"})
+            version_files: list[dict] = version.get("files", [])
+            model_files = utils.filter_with(version_files, {"type": "Model"})
 
             shortname = version.get("name", None) if len(model_files) > 0 else None
 
@@ -108,7 +108,7 @@ class CivitaiModelSearcher(ModelSearcher):
                 description_parts.append("")
 
                 model = {
-                    "id": file.get("id"),
+                    "id": version.get("id"),
                     "shortname": shortname or basename,
                     "basename": basename,
                     "extension": extension,
@@ -122,6 +122,7 @@ class CivitaiModelSearcher(ModelSearcher):
                     "downloadPlatform": "civitai",
                     "downloadUrl": file.get("downloadUrl"),
                     "hashes": file.get("hashes"),
+                    "files": version_files if len(version_files) > 1 else None,
                 }
                 models.append(model)
 
diff --git a/src/components/DialogCreateTask.vue b/src/components/DialogCreateTask.vue
index 659fb35..2378a2b 100644
--- a/src/components/DialogCreateTask.vue
+++ b/src/components/DialogCreateTask.vue
@@ -31,12 +31,20 @@
         <KeepAlive>
           <ModelContent
             v-if="currentModel"
-            :key="currentModel.id"
+            :key="`${currentModel.id}-${currentModel.currentFileId}`"
             :model="currentModel"
             :editable="true"
             @submit="createDownTask"
           >
             <template #action>
+              <div v-if="currentModel.files" class="flex-1">
+                <ResponseSelect
+                  :model-value="currentModel.currentFileId"
+                  :items="currentModel.selectionFiles"
+                  :type="isMobile ? 'drop' : 'button'"
+                >
+                </ResponseSelect>
+              </div>
               <Button
                 icon="pi pi-download"
                 :label="$t('download')"
diff --git a/src/components/ModelContent.vue b/src/components/ModelContent.vue
index 4dcdd44..ca55723 100644
--- a/src/components/ModelContent.vue
+++ b/src/components/ModelContent.vue
@@ -17,7 +17,7 @@
         ></ModelPreview>
 
         <div class="flex flex-col gap-4 overflow-hidden">
-          <div class="flex items-center justify-end gap-4">
+          <div class="flex h-10 items-center justify-end gap-4">
             <slot name="action" :metadata="formInstance.metadata.value"></slot>
           </div>
 
diff --git a/src/hooks/download.ts b/src/hooks/download.ts
index 6afe39c..7723914 100644
--- a/src/hooks/download.ts
+++ b/src/hooks/download.ts
@@ -2,17 +2,19 @@ import { useLoading } from 'hooks/loading'
 import { request } from 'hooks/request'
 import { defineStore } from 'hooks/store'
 import { useToast } from 'hooks/toast'
+import { upperFirst } from 'lodash'
 import { api } from 'scripts/comfyAPI'
 import {
-  BaseModel,
   DownloadTask,
   DownloadTaskOptions,
   SelectOptions,
   VersionModel,
+  VersionModelFile,
 } from 'types/typings'
 import { bytesToSize } from 'utils/common'
 import { onBeforeMount, onMounted, ref, watch } from 'vue'
 import { useI18n } from 'vue-i18n'
+import yaml from 'yaml'
 
 export const useDownload = defineStore('download', (store) => {
   const { toast, confirm, wrapperToastError } = useToast()
@@ -162,12 +164,60 @@ declare module 'hooks/store' {
   }
 }
 
+type WithSelection<T> = SelectOptions & { item: T }
+
+type FileSelectionVersionModel = VersionModel & {
+  currentFileId?: number
+  selectionFiles?: WithSelection<VersionModelFile>[]
+}
+
 export const useModelSearch = () => {
   const loading = useLoading()
   const { toast } = useToast()
-  const data = ref<(SelectOptions & { item: VersionModel })[]>([])
+  const data = ref<WithSelection<FileSelectionVersionModel>[]>([])
   const current = ref<string | number>()
-  const currentModel = ref<BaseModel>()
+  const currentModel = ref<FileSelectionVersionModel>()
+
+  const genFileSelectionItem = (
+    item: VersionModel,
+  ): FileSelectionVersionModel => {
+    const fileSelectionItem: FileSelectionVersionModel = { ...item }
+    fileSelectionItem.selectionFiles = fileSelectionItem.files
+      ?.sort((file) => (file.type === 'Model' ? -1 : 1))
+      .map((file) => {
+        const parts = file.name.split('.')
+        const extension = `.${parts.pop()}`
+        const basename = parts.join('.')
+
+        const regexp = /---\n([\s\S]*?)\n---/
+        const yamlMetadataMatch = item.description.match(regexp)
+        const yamlMetadata = yaml.parse(yamlMetadataMatch?.[1] || '')
+        yamlMetadata.hashes = file.hashes
+        yamlMetadata.metadata = file.metadata
+        const yamlContent = `---\n${yaml.stringify(yamlMetadata)}---`
+        const description = item.description.replace(regexp, yamlContent)
+
+        return {
+          label: file.type === 'Model' ? upperFirst(item.type) : file.type,
+          value: file.id,
+          item: file,
+          command() {
+            if (currentModel.value) {
+              currentModel.value.basename = basename
+              currentModel.value.extension = extension
+              currentModel.value.sizeBytes = file.sizeKB * 1024
+              currentModel.value.metadata = file.metadata
+              currentModel.value.downloadUrl = file.downloadUrl
+              currentModel.value.hashes = file.hashes
+              currentModel.value.description = description
+              currentModel.value.currentFileId = file.id
+            }
+          },
+        }
+      })
+    fileSelectionItem.currentFileId = item.files?.[0]?.id
+    return fileSelectionItem
+  }
 
   const handleSearchByUrl = async (url: string) => {
     if (!url) {
@@ -177,14 +227,17 @@ export const useModelSearch = () => {
     loading.show()
     return request(`/model-info?model-page=${encodeURIComponent(url)}`, {})
       .then((resData: VersionModel[]) => {
-        data.value = resData.map((item) => ({
-          label: item.shortname,
-          value: item.id,
-          item,
-          command() {
-            current.value = item.id
-          },
-        }))
+        data.value = resData.map((item) => {
+          const resolvedItem = genFileSelectionItem(item)
+          return {
+            label: item.shortname,
+            value: item.id,
+            item: resolvedItem,
+            command() {
+              current.value = item.id
+            },
+          }
+        })
         current.value = data.value[0]?.value
         currentModel.value = data.value[0]?.item
 
diff --git a/src/types/typings.d.ts b/src/types/typings.d.ts
index 6a554cd..66ec0f3 100644
--- a/src/types/typings.d.ts
+++ b/src/types/typings.d.ts
@@ -22,11 +22,22 @@ export interface Model extends BaseModel {
   children?: Model[]
 }
 
+export interface VersionModelFile {
+  id: number
+  sizeKB: number
+  name: string
+  type: string
+  metadata: Record<string, string>
+  hashes: Record<string, string>
+  downloadUrl: string
+}
+
 export interface VersionModel extends BaseModel {
   shortname: string
   downloadPlatform: string
   downloadUrl: string
   hashes?: Record<string, string>
+  files?: VersionModelFile[]
 }
 
 export type WithResolved<T> = Omit<T, 'preview'> & {