TBNilles/ComfyUI-Model-Manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Feature scan info (#53)

Browse Source 
* 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
This commit is contained in:
Hayden
2024-11-21 22:04:39 +08:00
committed by GitHub
parent 6ae7e1835f
commit 659637c6e0
20 changed files with 921 additions and 428 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
src/components/DialogCreateTask.vue
View File

@@ -70,7 +70,6 @@ import { request } from 'hooks/request'
import { useToast } from 'hooks/toast'
import Button from 'primevue/button'
import { VersionModel } from 'types/typings'
import { previewUrlToFile } from 'utils/common'
import { ref } from 'vue'
const { isMobile } = useConfig()
@@ -89,38 +88,11 @@ const searchModelsByUrl = async () => {
}
const createDownTask = async (data: VersionModel) => {
const formData = new FormData()
loading.show()
// set base info
formData.append('type', data.type)
formData.append('pathIndex', data.pathIndex.toString())
formData.append('fullname', data.fullname)
// set preview
const previewFile = await previewUrlToFile(data.preview as string).catch(
() => {
loading.hide()
toast.add({
severity: 'error',
summary: 'Error',
detail: 'Failed to download preview',
life: 15000,
})
throw new Error('Failed to download preview')
},
)
formData.append('previewFile', previewFile)
// set description
formData.append('description', data.description)
// set model download info
formData.append('downloadPlatform', data.downloadPlatform)
formData.append('downloadUrl', data.downloadUrl)
formData.append('sizeBytes', data.sizeBytes.toString())
formData.append('hashes', JSON.stringify(data.hashes))
await request('/model', {
method: 'POST',
body: formData,
body: JSON.stringify(data),
})
.then(() => {
dialog.close({ key: 'model-manager-create-task' })

184
src/hooks/config.ts
View File

@@ -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
})
},
})
},
})
},
})
})
}

263
src/hooks/download.ts
View File

@@ -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,

6
src/hooks/loading.ts
View File

@@ -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>()

15
src/hooks/markdown.ts
View File

@@ -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>

26
src/hooks/model.ts
View File

@@ -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',

1
src/scripts/comfyAPI.ts
View File

@@ -5,3 +5,4 @@ export const $el = window.comfyAPI.ui.$el
export const ComfyApp = window.comfyAPI.app.ComfyApp
export const ComfyButton = window.comfyAPI.button.ComfyButton
export const ComfyDialog = window.comfyAPI.dialog.ComfyDialog

10
src/types/global.d.ts vendored
View File

@@ -112,6 +112,7 @@ declare namespace ComfyAPI {
settings: ComfySettingsDialog
menuHamburger?: HTMLDivElement
menuContainer?: HTMLDivElement
dialog: dialog.ComfyDialog
}
type SettingInputType =
@@ -197,6 +198,15 @@ declare namespace ComfyAPI {
constructor(...buttons: (HTMLElement | ComfyButton)[]): ComfyButtonGroup
}
}
namespace dialog {
class ComfyDialog {
constructor(type = 'div', buttons: HTMLElement[] = null)
element: HTMLElement
close(): void
show(html: string | HTMLElement | HTMLElement[]): void
}
}
}
declare namespace lightGraph {

11
src/utils/common.ts
View File

@@ -26,14 +26,3 @@ export const bytesToSize = (
export const formatDate = (date: number | string | Date) => {
return dayjs(date).format('YYYY-MM-DD HH:mm:ss')
}
export const previewUrlToFile = async (url: string) => {
return fetch(url)
.then((res) => res.blob())
.then((blob) => {
const type = blob.type
const extension = type.split('/')[1]
const file = new File([blob], `preview.${extension}`, { type })
return file
})
}