Feature scan info (#53)

* pref: migrate fetch model info to end back

* fix(download): can't select model type

* feat: add scan model info

* feat: add trigger button in setting

* feat: add printing logs

* chore: add explanation of scan model info

src/hooks/config.ts

@@ -1,9 +1,10 @@
-import { useRequest } from 'hooks/request'
+import { request, useRequest } from 'hooks/request'
 import { defineStore } from 'hooks/store'
-import { app } from 'scripts/comfyAPI'
+import { $el, app, ComfyDialog } from 'scripts/comfyAPI'
 import { onMounted, onUnmounted, ref } from 'vue'
+import { useToast } from './toast'
 
-export const useConfig = defineStore('config', () => {
+export const useConfig = defineStore('config', (store) => {
   const mobileDeviceBreakPoint = 759
   const isMobile = ref(window.innerWidth < mobileDeviceBreakPoint)
 
@@ -36,7 +37,7 @@ export const useConfig = defineStore('config', () => {
     refresh,
   }
 
-  useAddConfigSettings()
+  useAddConfigSettings(store)
 
   return config
 })
@@ -49,7 +50,41 @@ declare module 'hooks/store' {
   }
 }
 
-function useAddConfigSettings() {
+function useAddConfigSettings(store: import('hooks/store').StoreProvider) {
+  const { toast } = useToast()
+
+  const confirm = (opts: {
+    message?: string
+    accept?: () => void
+    reject?: () => void
+  }) => {
+    const dialog = new ComfyDialog('div', [])
+
+    dialog.show(
+      $el('div', [
+        $el('p', { textContent: opts.message }),
+        $el('div.flex.gap-4', [
+          $el('button.flex-1', {
+            textContent: 'Cancel',
+            onclick: () => {
+              opts.reject?.()
+              dialog.close()
+              document.body.removeChild(dialog.element)
+            },
+          }),
+          $el('button.flex-1', {
+            textContent: 'Continue',
+            onclick: () => {
+              opts.accept?.()
+              dialog.close()
+              document.body.removeChild(dialog.element)
+            },
+          }),
+        ]),
+      ]),
+    )
+  }
+
   onMounted(() => {
     // API keys
     app.ui?.settings.addSetting({
@@ -65,5 +100,144 @@ function useAddConfigSettings() {
       type: 'text',
       defaultValue: undefined,
     })
+
+    // Migrate
+    app.ui?.settings.addSetting({
+      id: 'ModelManager.Migrate.Migrate',
+      name: 'Migrate information from cdb-boop/main',
+      defaultValue: '',
+      type: () => {
+        return $el('button.p-button.p-component.p-button-secondary', {
+          textContent: 'Migrate',
+          onclick: () => {
+            confirm({
+              message: [
+                'This operation will delete old files and override current files if it exists.',
+                // 'This may take a while and generate MANY server requests!',
+                'Continue?',
+              ].join('\n'),
+              accept: () => {
+                store.loading.loading.value = true
+                request('/migrate', {
+                  method: 'POST',
+                })
+                  .then(() => {
+                    toast.add({
+                      severity: 'success',
+                      summary: 'Complete migration',
+                      life: 2000,
+                    })
+                    store.models.refresh()
+                  })
+                  .catch((err) => {
+                    toast.add({
+                      severity: 'error',
+                      summary: 'Error',
+                      detail: err.message ?? 'Failed to migrate information',
+                      life: 15000,
+                    })
+                  })
+                  .finally(() => {
+                    store.loading.loading.value = false
+                  })
+              },
+            })
+          },
+        })
+      },
+    })
+
+    // Scan information
+    app.ui?.settings.addSetting({
+      id: 'ModelManager.ScanFiles.Full',
+      name: "Override all models' information and preview",
+      defaultValue: '',
+      type: () => {
+        return $el('button.p-button.p-component.p-button-secondary', {
+          textContent: 'Full Scan',
+          onclick: () => {
+            confirm({
+              message: [
+                'This operation will override current files.',
+                'This may take a while and generate MANY server requests!',
+                'USE AT YOUR OWN RISK! Continue?',
+              ].join('\n'),
+              accept: () => {
+                store.loading.loading.value = true
+                request('/model-info/scan', {
+                  method: 'POST',
+                  body: JSON.stringify({ scanMode: 'full' }),
+                })
+                  .then(() => {
+                    toast.add({
+                      severity: 'success',
+                      summary: 'Complete download information',
+                      life: 2000,
+                    })
+                    store.models.refresh()
+                  })
+                  .catch((err) => {
+                    toast.add({
+                      severity: 'error',
+                      summary: 'Error',
+                      detail: err.message ?? 'Failed to download information',
+                      life: 15000,
+                    })
+                  })
+                  .finally(() => {
+                    store.loading.loading.value = false
+                  })
+              },
+            })
+          },
+        })
+      },
+    })
+
+    app.ui?.settings.addSetting({
+      id: 'ModelManager.ScanFiles.Incremental',
+      name: 'Download missing information or preview',
+      defaultValue: '',
+      type: () => {
+        return $el('button.p-button.p-component.p-button-secondary', {
+          textContent: 'Diff Scan',
+          onclick: () => {
+            confirm({
+              message: [
+                'Download missing information or preview.',
+                'This may take a while and generate MANY server requests!',
+                'USE AT YOUR OWN RISK! Continue?',
+              ].join('\n'),
+              accept: () => {
+                store.loading.loading.value = true
+                request('/model-info/scan', {
+                  method: 'POST',
+                  body: JSON.stringify({ scanMode: 'diff' }),
+                })
+                  .then(() => {
+                    toast.add({
+                      severity: 'success',
+                      summary: 'Complete download information',
+                      life: 2000,
+                    })
+                    store.models.refresh()
+                  })
+                  .catch((err) => {
+                    toast.add({
+                      severity: 'error',
+                      summary: 'Error',
+                      detail: err.message ?? 'Failed to download information',
+                      life: 15000,
+                    })
+                  })
+                  .finally(() => {
+                    store.loading.loading.value = false
+                  })
+              },
+            })
+          },
+        })
+      },
+    })
   })
 }

src/hooks/download.ts

@@ -1,5 +1,4 @@
 import { useLoading } from 'hooks/loading'
-import { MarkdownTool, useMarkdown } from 'hooks/markdown'
 import { request } from 'hooks/request'
 import { defineStore } from 'hooks/store'
 import { useToast } from 'hooks/toast'
@@ -157,253 +156,8 @@ declare module 'hooks/store' {
   }
 }
 
-abstract class ModelSearch {
-  constructor(readonly md: MarkdownTool) {}
-
-  abstract search(pathname: string): Promise<VersionModel[]>
-}
-
-class Civitai extends ModelSearch {
-  async search(searchUrl: string): Promise<VersionModel[]> {
-    const { pathname, searchParams } = new URL(searchUrl)
-
-    const [, modelId] = pathname.match(/^\/models\/(\d*)/) ?? []
-    const versionId = searchParams.get('modelVersionId')
-
-    if (!modelId) {
-      return Promise.resolve([])
-    }
-
-    return fetch(`https://civitai.com/api/v1/models/${modelId}`)
-      .then((response) => response.json())
-      .then((resData) => {
-        const modelVersions: any[] = resData.modelVersions.filter(
-          (version: any) => {
-            if (versionId) {
-              return version.id == versionId
-            }
-            return true
-          },
-        )
-
-        const models: VersionModel[] = []
-
-        for (const version of modelVersions) {
-          const modelFiles: any[] = version.files.filter(
-            (file: any) => file.type === 'Model',
-          )
-
-          const shortname = modelFiles.length > 0 ? version.name : undefined
-
-          for (const file of modelFiles) {
-            const fullname = file.name
-            const extension = `.${fullname.split('.').pop()}`
-            const basename = fullname.replace(extension, '')
-
-            models.push({
-              id: file.id,
-              shortname: shortname ?? basename,
-              fullname: fullname,
-              basename: basename,
-              extension: extension,
-              preview: version.images.map((i: any) => i.url),
-              sizeBytes: file.sizeKB * 1024,
-              type: this.resolveType(resData.type),
-              pathIndex: 0,
-              description: [
-                '---',
-                ...[
-                  `website: Civitai`,
-                  `modelPage: https://civitai.com/models/${modelId}?modelVersionId=${version.id}`,
-                  `author: ${resData.creator?.username}`,
-                  version.baseModel && `baseModel: ${version.baseModel}`,
-                  file.hashes && `hashes:`,
-                  ...Object.entries(file.hashes ?? {}).map(
-                    ([key, value]) => `  ${key}: ${value}`,
-                  ),
-                  file.metadata && `metadata:`,
-                  ...Object.entries(file.metadata ?? {}).map(
-                    ([key, value]) => `  ${key}: ${value}`,
-                  ),
-                ].filter(Boolean),
-                '---',
-                '',
-                '# Trigger Words',
-                `\n${(version.trainedWords ?? ['No trigger words']).join(', ')}\n`,
-                '# About this version',
-                this.resolveDescription(
-                  version.description,
-                  '\nNo description about this version\n',
-                ),
-                `# ${resData.name}`,
-                this.resolveDescription(
-                  resData.description,
-                  'No description about this model',
-                ),
-              ].join('\n'),
-              metadata: file.metadata,
-              downloadPlatform: 'civitai',
-              downloadUrl: file.downloadUrl,
-              hashes: file.hashes,
-            })
-          }
-        }
-
-        return models
-      })
-  }
-
-  private resolveType(type: string) {
-    const mapLegacy = {
-      TextualInversion: 'embeddings',
-      LoCon: 'loras',
-      DoRA: 'loras',
-      Controlnet: 'controlnet',
-      Upscaler: 'upscale_models',
-      VAE: 'vae',
-    }
-    return mapLegacy[type] ?? `${type.toLowerCase()}s`
-  }
-
-  private resolveDescription(content: string, defaultContent: string) {
-    const mdContent = this.md.parse(content ?? '').trim()
-    return mdContent || defaultContent
-  }
-}
-
-class Huggingface extends ModelSearch {
-  async search(searchUrl: string): Promise<VersionModel[]> {
-    const { pathname } = new URL(searchUrl)
-    const [, space, name, ...restPaths] = pathname.split('/')
-
-    if (!space || !name) {
-      return Promise.resolve([])
-    }
-
-    const modelId = `${space}/${name}`
-    const restPathname = restPaths.join('/')
-
-    return fetch(`https://huggingface.co/api/models/${modelId}`)
-      .then((response) => response.json())
-      .then((resData) => {
-        const siblingFiles: string[] = resData.siblings.map(
-          (item: any) => item.rfilename,
-        )
-
-        const modelFiles: string[] = this.filterTreeFiles(
-          this.filterModelFiles(siblingFiles),
-          restPathname,
-        )
-        const images: string[] = this.filterTreeFiles(
-          this.filterImageFiles(siblingFiles),
-          restPathname,
-        ).map((filename) => {
-          return `https://huggingface.co/${modelId}/resolve/main/${filename}`
-        })
-
-        const models: VersionModel[] = []
-
-        for (const filename of modelFiles) {
-          const fullname = filename.split('/').pop()!
-          const extension = `.${fullname.split('.').pop()}`
-          const basename = fullname.replace(extension, '')
-
-          models.push({
-            id: filename,
-            shortname: filename,
-            fullname: fullname,
-            basename: basename,
-            extension: extension,
-            preview: images,
-            sizeBytes: 0,
-            type: 'unknown',
-            pathIndex: 0,
-            description: [
-              '---',
-              ...[
-                `website: HuggingFace`,
-                `modelPage: https://huggingface.co/${modelId}`,
-                `author: ${resData.author}`,
-              ].filter(Boolean),
-              '---',
-              '',
-              '# Trigger Words',
-              '\nNo trigger words\n',
-              '# About this version',
-              '\nNo description about this version\n',
-              `# ${resData.modelId}`,
-              '\nNo description about this model\n',
-            ].join('\n'),
-            metadata: {},
-            downloadPlatform: 'huggingface',
-            downloadUrl: `https://huggingface.co/${modelId}/resolve/main/${filename}?download=true`,
-          })
-        }
-
-        return models
-      })
-  }
-
-  private filterTreeFiles(files: string[], pathname: string) {
-    const [target, , ...paths] = pathname.split('/')
-
-    if (!target) return files
-
-    if (target !== 'tree' && target !== 'blob') return files
-
-    const pathPrefix = paths.join('/')
-    return files.filter((file) => {
-      return file.startsWith(pathPrefix)
-    })
-  }
-
-  private filterModelFiles(files: string[]) {
-    const extension = [
-      '.bin',
-      '.ckpt',
-      '.gguf',
-      '.onnx',
-      '.pt',
-      '.pth',
-      '.safetensors',
-    ]
-    return files.filter((file) => {
-      const ext = file.split('.').pop()
-      return ext ? extension.includes(`.${ext}`) : false
-    })
-  }
-
-  private filterImageFiles(files: string[]) {
-    const extension = [
-      '.png',
-      '.webp',
-      '.jpeg',
-      '.jpg',
-      '.jfif',
-      '.gif',
-      '.apng',
-    ]
-
-    return files.filter((file) => {
-      const ext = file.split('.').pop()
-      return ext ? extension.includes(`.${ext}`) : false
-    })
-  }
-}
-
-class UnknownWebsite extends ModelSearch {
-  async search(): Promise<VersionModel[]> {
-    return Promise.reject(
-      new Error(
-        'Unknown Website, please input a URL from huggingface.co or civitai.com.',
-      ),
-    )
-  }
-}
-
 export const useModelSearch = () => {
   const loading = useLoading()
-  const md = useMarkdown()
   const { toast } = useToast()
   const data = ref<(SelectOptions & { item: VersionModel })[]>([])
   const current = ref<string | number>()
@@ -414,22 +168,9 @@ export const useModelSearch = () => {
       return Promise.resolve([])
     }
 
-    let instance: ModelSearch = new UnknownWebsite(md)
-
-    const { hostname } = new URL(url ?? '')
-
-    if (hostname === 'civitai.com') {
-      instance = new Civitai(md)
-    }
-
-    if (hostname === 'huggingface.co') {
-      instance = new Huggingface(md)
-    }
-
     loading.show()
-    return instance
-      .search(url)
-      .then((resData) => {
+    return request(`/model-info?model-page=${encodeURIComponent(url)}`, {})
+      .then((resData: VersionModel[]) => {
         data.value = resData.map((item) => ({
           label: item.shortname,
           value: item.id,

src/hooks/loading.ts

@@ -31,6 +31,12 @@ export const useGlobalLoading = defineStore('loading', () => {
   return { loading }
 })
 
+declare module 'hooks/store' {
+  interface StoreProvider {
+    loading: ReturnType<typeof useGlobalLoading>
+  }
+}
+
 export const useLoading = () => {
   const timer = ref<NodeJS.Timeout>()
 

src/hooks/markdown.ts

@@ -1,6 +1,5 @@
 import MarkdownIt from 'markdown-it'
 import metadata_block from 'markdown-it-metadata-block'
-import TurndownService from 'turndown'
 import yaml from 'yaml'
 
 interface MarkdownOptions {
@@ -31,19 +30,7 @@ export const useMarkdown = (opts?: MarkdownOptions) => {
     return self.renderToken(tokens, idx, options)
   }
 
-  const turndown = new TurndownService({
-    headingStyle: 'atx',
-    bulletListMarker: '-',
-  })
-
-  turndown.addRule('paragraph', {
-    filter: 'p',
-    replacement: function (content) {
-      return `\n\n${content}`
-    },
-  })
-
-  return { render: md.render.bind(md), parse: turndown.turndown.bind(turndown) }
+  return { render: md.render.bind(md) }
 }
 
 export type MarkdownTool = ReturnType<typeof useMarkdown>

src/hooks/model.ts

@@ -7,7 +7,7 @@ import { useToast } from 'hooks/toast'
 import { cloneDeep } from 'lodash'
 import { app } from 'scripts/comfyAPI'
 import { BaseModel, Model, SelectEvent } from 'types/typings'
-import { bytesToSize, formatDate, previewUrlToFile } from 'utils/common'
+import { bytesToSize, formatDate } from 'utils/common'
 import { ModelGrid } from 'utils/legacy'
 import { genModelKey, resolveModelTypeLoader } from 'utils/model'
 import {
@@ -29,18 +29,17 @@ export const useModels = defineStore('models', (store) => {
   const loading = useLoading()
 
   const updateModel = async (model: BaseModel, data: BaseModel) => {
-    const formData = new FormData()
+    const updateData = new Map()
     let oldKey: string | null = null
 
     // Check current preview
     if (model.preview !== data.preview) {
-      const previewFile = await previewUrlToFile(data.preview as string)
-      formData.append('previewFile', previewFile)
+      updateData.set('previewFile', data.preview)
     }
 
     // Check current description
     if (model.description !== data.description) {
-      formData.append('description', data.description)
+      updateData.set('description', data.description)
     }
 
     // Check current name and pathIndex
@@ -49,19 +48,19 @@ export const useModels = defineStore('models', (store) => {
       model.pathIndex !== data.pathIndex
     ) {
       oldKey = genModelKey(model)
-      formData.append('type', data.type)
-      formData.append('pathIndex', data.pathIndex.toString())
-      formData.append('fullname', data.fullname)
+      updateData.set('type', data.type)
+      updateData.set('pathIndex', data.pathIndex.toString())
+      updateData.set('fullname', data.fullname)
     }
 
-    if (formData.keys().next().done) {
+    if (updateData.size === 0) {
       return
     }
 
     loading.show()
     await request(`/model/${model.type}/${model.pathIndex}/${model.fullname}`, {
       method: 'PUT',
-      body: formData,
+      body: JSON.stringify(Object.fromEntries(updateData.entries())),
     })
       .catch((err) => {
         const error_message = err.message ?? err.error
@@ -246,14 +245,17 @@ export const useModelBaseInfoEditor = (formInstance: ModelFormInstance) => {
 
   interface FieldsItem {
     key: keyof Model
-    formatter: (val: any) => string
+    formatter: (val: any) => string | undefined | null
   }
 
   const baseInfo = computed(() => {
     const fields: FieldsItem[] = [
       {
         key: 'type',
-        formatter: () => modelData.value.type,
+        formatter: () =>
+          modelData.value.type in modelFolders.value
+            ? modelData.value.type
+            : undefined,
       },
       {
         key: 'pathIndex',