779 lines
19 KiB
TypeScript
779 lines
19 KiB
TypeScript
import DialogModelDetail from 'components/DialogModelDetail.vue'
|
|
import { useLoading } from 'hooks/loading'
|
|
import { useMarkdown } from 'hooks/markdown'
|
|
import { request } from 'hooks/request'
|
|
import { defineStore } from 'hooks/store'
|
|
import { useToast } from 'hooks/toast'
|
|
import { castArray, cloneDeep } from 'lodash'
|
|
import { TreeNode } from 'primevue/treenode'
|
|
import { api, app } from 'scripts/comfyAPI'
|
|
import { BaseModel, Model, SelectEvent, WithResolved } from 'types/typings'
|
|
import { bytesToSize, formatDate, previewUrlToFile } from 'utils/common'
|
|
import { ModelGrid } from 'utils/legacy'
|
|
import { genModelKey, resolveModelTypeLoader } from 'utils/model'
|
|
import {
|
|
computed,
|
|
inject,
|
|
type InjectionKey,
|
|
MaybeRefOrGetter,
|
|
onMounted,
|
|
provide,
|
|
type Ref,
|
|
ref,
|
|
toRaw,
|
|
toValue,
|
|
unref,
|
|
} from 'vue'
|
|
import { useI18n } from 'vue-i18n'
|
|
import { configSetting } from './config'
|
|
|
|
const systemStat = ref()
|
|
|
|
type ModelFolder = Record<string, string[]>
|
|
|
|
const modelFolderProvideKey = Symbol('modelFolder') as InjectionKey<
|
|
Ref<ModelFolder>
|
|
>
|
|
|
|
export const genModelFullName = (model: BaseModel, splitter = '/') => {
|
|
return [model.subFolder, `${model.basename}${model.extension}`]
|
|
.filter(Boolean)
|
|
.join(splitter)
|
|
}
|
|
|
|
export const genModelUrl = (model: BaseModel) => {
|
|
const fullname = genModelFullName(model)
|
|
return `/model/${model.type}/${model.pathIndex}/${fullname}`
|
|
}
|
|
|
|
export const useModels = defineStore('models', (store) => {
|
|
const { toast, confirm } = useToast()
|
|
const { t } = useI18n()
|
|
const loading = useLoading()
|
|
|
|
const folders = ref<ModelFolder>({})
|
|
const initialized = ref(false)
|
|
|
|
const refreshFolders = async () => {
|
|
return request('/models').then((resData) => {
|
|
folders.value = resData
|
|
initialized.value = true
|
|
})
|
|
}
|
|
|
|
provide(modelFolderProvideKey, folders)
|
|
|
|
const models = ref<Record<string, Model[]>>({})
|
|
|
|
const refreshModels = async (folder: string) => {
|
|
loading.show(folder)
|
|
return request(`/models/${folder}`)
|
|
.then((resData) => {
|
|
models.value[folder] = resData
|
|
return resData
|
|
})
|
|
.finally(() => {
|
|
loading.hide(folder)
|
|
})
|
|
}
|
|
|
|
const refreshAllModels = async (force = false) => {
|
|
const forceRefresh = force ? refreshFolders() : Promise.resolve()
|
|
models.value = {}
|
|
const excludeScanTypes = app.ui?.settings.getSettingValue<string>(
|
|
configSetting.excludeScanTypes,
|
|
)
|
|
const customBlackList =
|
|
excludeScanTypes
|
|
?.split(',')
|
|
.map((type) => type.trim())
|
|
.filter(Boolean) ?? []
|
|
await forceRefresh.then(() =>
|
|
Promise.allSettled(
|
|
Object.keys(folders.value)
|
|
.filter((folder) => !customBlackList.includes(folder))
|
|
.map(refreshModels),
|
|
),
|
|
)
|
|
}
|
|
|
|
const updateModel = async (
|
|
model: BaseModel,
|
|
data: WithResolved<BaseModel>,
|
|
) => {
|
|
const updateData = new FormData()
|
|
let oldKey: string | null = null
|
|
let needUpdate = false
|
|
|
|
// Check current preview
|
|
if (model.preview !== data.preview) {
|
|
const preview = data.preview
|
|
if (preview) {
|
|
const previewFile = await previewUrlToFile(data.preview as string)
|
|
updateData.set('previewFile', previewFile)
|
|
} else {
|
|
updateData.set('previewFile', 'undefined')
|
|
}
|
|
needUpdate = true
|
|
}
|
|
|
|
// Check current description
|
|
if (model.description !== data.description) {
|
|
updateData.set('description', data.description)
|
|
needUpdate = true
|
|
}
|
|
|
|
// Check current name and pathIndex
|
|
if (
|
|
model.subFolder !== data.subFolder ||
|
|
model.pathIndex !== data.pathIndex
|
|
) {
|
|
oldKey = genModelKey(model)
|
|
updateData.set('type', data.type)
|
|
updateData.set('pathIndex', data.pathIndex.toString())
|
|
updateData.set('fullname', genModelFullName(data as BaseModel))
|
|
needUpdate = true
|
|
}
|
|
|
|
if (!needUpdate) {
|
|
return
|
|
}
|
|
|
|
loading.show()
|
|
|
|
await request(genModelUrl(model), {
|
|
method: 'PUT',
|
|
body: updateData,
|
|
})
|
|
.catch((err) => {
|
|
const error_message = err.message ?? err.error
|
|
toast.add({
|
|
severity: 'error',
|
|
summary: 'Error',
|
|
detail: `Failed to update model: ${error_message}`,
|
|
life: 15000,
|
|
})
|
|
throw new Error(error_message)
|
|
})
|
|
.finally(() => {
|
|
loading.hide()
|
|
})
|
|
|
|
if (oldKey) {
|
|
store.dialog.close({ key: oldKey })
|
|
}
|
|
|
|
refreshModels(data.type)
|
|
}
|
|
|
|
const deleteModel = async (model: BaseModel) => {
|
|
return new Promise((resolve) => {
|
|
confirm.require({
|
|
message: t('deleteAsk', [t('model').toLowerCase()]),
|
|
header: 'Danger',
|
|
icon: 'pi pi-info-circle',
|
|
rejectProps: {
|
|
label: t('cancel'),
|
|
severity: 'secondary',
|
|
outlined: true,
|
|
},
|
|
acceptProps: {
|
|
label: t('delete'),
|
|
severity: 'danger',
|
|
},
|
|
accept: () => {
|
|
const dialogKey = genModelKey(model)
|
|
loading.show()
|
|
request(genModelUrl(model), {
|
|
method: 'DELETE',
|
|
})
|
|
.then(() => {
|
|
toast.add({
|
|
severity: 'success',
|
|
summary: 'Success',
|
|
detail: `${model.basename} Deleted`,
|
|
life: 2000,
|
|
})
|
|
store.dialog.close({ key: dialogKey })
|
|
return refreshModels(model.type)
|
|
})
|
|
.then(() => {
|
|
resolve(void 0)
|
|
})
|
|
.catch((e) => {
|
|
toast.add({
|
|
severity: 'error',
|
|
summary: 'Error',
|
|
detail: e.message ?? 'Failed to delete model',
|
|
life: 15000,
|
|
})
|
|
})
|
|
.finally(() => {
|
|
loading.hide()
|
|
})
|
|
},
|
|
reject: () => {
|
|
resolve(void 0)
|
|
},
|
|
})
|
|
})
|
|
}
|
|
|
|
function openModelDetail(model: BaseModel) {
|
|
const filename = model.basename.replace(model.extension, '')
|
|
|
|
store.dialog.open({
|
|
key: genModelKey(model),
|
|
title: filename,
|
|
content: DialogModelDetail,
|
|
contentProps: { model: model },
|
|
})
|
|
}
|
|
|
|
function getFullPath(model: BaseModel) {
|
|
const fullname = genModelFullName(model)
|
|
const prefixPath = folders.value[model.type]?.[model.pathIndex]
|
|
return [prefixPath, fullname].filter(Boolean).join('/')
|
|
}
|
|
|
|
onMounted(() => {
|
|
api.getSystemStats().then((res) => {
|
|
systemStat.value = res
|
|
})
|
|
})
|
|
|
|
return {
|
|
initialized: initialized,
|
|
folders: folders,
|
|
data: models,
|
|
refresh: refreshAllModels,
|
|
remove: deleteModel,
|
|
update: updateModel,
|
|
openModelDetail: openModelDetail,
|
|
getFullPath: getFullPath,
|
|
}
|
|
})
|
|
|
|
declare module 'hooks/store' {
|
|
interface StoreProvider {
|
|
models: ReturnType<typeof useModels>
|
|
}
|
|
}
|
|
|
|
export const useModelFormData = (getFormData: () => BaseModel) => {
|
|
const formData = ref<BaseModel>(getFormData())
|
|
const modelData = ref<BaseModel>(getFormData())
|
|
|
|
type ResetCallback = () => void
|
|
const resetCallback = ref<ResetCallback[]>([])
|
|
|
|
const registerReset = (callback: ResetCallback) => {
|
|
resetCallback.value.push(callback)
|
|
}
|
|
|
|
const reset = () => {
|
|
formData.value = getFormData()
|
|
modelData.value = getFormData()
|
|
for (const callback of resetCallback.value) {
|
|
callback()
|
|
}
|
|
}
|
|
|
|
type SubmitCallback = (data: WithResolved<BaseModel>) => void
|
|
const submitCallback = ref<SubmitCallback[]>([])
|
|
|
|
const registerSubmit = (callback: SubmitCallback) => {
|
|
submitCallback.value.push(callback)
|
|
}
|
|
|
|
const submit = (): WithResolved<BaseModel> => {
|
|
const data: any = cloneDeep(toRaw(unref(formData)))
|
|
for (const callback of submitCallback.value) {
|
|
callback(data)
|
|
}
|
|
return data
|
|
}
|
|
|
|
const metadata = ref<Record<string, any>>({})
|
|
|
|
return {
|
|
formData,
|
|
modelData,
|
|
registerReset,
|
|
reset,
|
|
registerSubmit,
|
|
submit,
|
|
metadata,
|
|
}
|
|
}
|
|
|
|
type ModelFormInstance = ReturnType<typeof useModelFormData>
|
|
|
|
/**
|
|
* Model base info
|
|
*/
|
|
const baseInfoKey = Symbol('baseInfo') as InjectionKey<
|
|
ReturnType<typeof useModelBaseInfoEditor>
|
|
>
|
|
|
|
export const useModelBaseInfoEditor = (formInstance: ModelFormInstance) => {
|
|
const { formData: model, modelData } = formInstance
|
|
|
|
const provideModelFolders = inject(modelFolderProvideKey)
|
|
const modelFolders = computed<ModelFolder>(() => {
|
|
return provideModelFolders?.value ?? {}
|
|
})
|
|
|
|
const type = computed({
|
|
get: () => {
|
|
return model.value.type
|
|
},
|
|
set: (val) => {
|
|
model.value.type = val
|
|
},
|
|
})
|
|
|
|
const pathIndex = computed({
|
|
get: () => {
|
|
return model.value.pathIndex
|
|
},
|
|
set: (val) => {
|
|
model.value.pathIndex = val
|
|
},
|
|
})
|
|
|
|
const subFolder = computed({
|
|
get: () => {
|
|
return model.value.subFolder
|
|
},
|
|
set: (val) => {
|
|
model.value.subFolder = val
|
|
},
|
|
})
|
|
|
|
const extension = computed(() => {
|
|
return model.value.extension
|
|
})
|
|
|
|
const basename = computed({
|
|
get: () => {
|
|
return model.value.basename
|
|
},
|
|
set: (val) => {
|
|
model.value.basename = val
|
|
},
|
|
})
|
|
|
|
interface BaseInfoItem {
|
|
key: string
|
|
display: string
|
|
value: any
|
|
}
|
|
|
|
interface FieldsItem {
|
|
key: keyof Model
|
|
formatter: (val: any) => string | undefined | null
|
|
}
|
|
|
|
const baseInfo = computed(() => {
|
|
const fields: FieldsItem[] = [
|
|
{
|
|
key: 'type',
|
|
formatter: () =>
|
|
modelData.value.type in modelFolders.value
|
|
? modelData.value.type
|
|
: undefined,
|
|
},
|
|
{
|
|
key: 'pathIndex',
|
|
formatter: () => {
|
|
const modelType = model.value.type
|
|
const pathIndex = model.value.pathIndex
|
|
if (!modelType) {
|
|
return undefined
|
|
}
|
|
const folders = modelFolders.value[modelType] ?? []
|
|
return [`${folders[pathIndex]}`, model.value.subFolder]
|
|
.filter(Boolean)
|
|
.join('/')
|
|
},
|
|
},
|
|
{
|
|
key: 'basename',
|
|
formatter: (val) => `${val}${model.value.extension}`,
|
|
},
|
|
{
|
|
key: 'sizeBytes',
|
|
formatter: (val) => (val == 0 ? 'Unknown' : bytesToSize(val)),
|
|
},
|
|
{
|
|
key: 'createdAt',
|
|
formatter: (val) => val && formatDate(val),
|
|
},
|
|
{
|
|
key: 'updatedAt',
|
|
formatter: (val) => val && formatDate(val),
|
|
},
|
|
]
|
|
|
|
const information: Record<string, BaseInfoItem> = {}
|
|
for (const item of fields) {
|
|
const key = item.key
|
|
const value = model.value[key]
|
|
const display = item.formatter(value)
|
|
|
|
if (display) {
|
|
information[key] = { key, value, display }
|
|
}
|
|
}
|
|
|
|
return information
|
|
})
|
|
|
|
const result = {
|
|
type,
|
|
baseInfo,
|
|
basename,
|
|
extension,
|
|
subFolder,
|
|
pathIndex,
|
|
modelFolders,
|
|
}
|
|
|
|
provide(baseInfoKey, result)
|
|
|
|
return result
|
|
}
|
|
|
|
export const useModelBaseInfo = () => {
|
|
return inject(baseInfoKey)!
|
|
}
|
|
|
|
export const useModelFolder = (
|
|
option: {
|
|
type?: MaybeRefOrGetter<string | undefined>
|
|
} = {},
|
|
) => {
|
|
const { data: models, folders: modelFolders } = useModels()
|
|
|
|
const pathOptions = computed(() => {
|
|
const type = toValue(option.type)
|
|
|
|
if (!type) {
|
|
return []
|
|
}
|
|
|
|
const folderItems = cloneDeep(models.value[type]) ?? []
|
|
const pureFolders = folderItems.filter((item) => item.isFolder)
|
|
pureFolders.sort((a, b) => a.basename.localeCompare(b.basename))
|
|
|
|
const folders = modelFolders.value[type] ?? []
|
|
|
|
const root: TreeNode[] = []
|
|
|
|
for (const [index, folder] of folders.entries()) {
|
|
const pathIndexItem: TreeNode = {
|
|
key: folder,
|
|
label: folder,
|
|
children: [],
|
|
}
|
|
|
|
const items = pureFolders
|
|
.filter((item) => item.pathIndex === index)
|
|
.map((item) => {
|
|
const node: TreeNode = {
|
|
key: `${folder}/${genModelFullName(item)}`,
|
|
label: item.basename,
|
|
data: item,
|
|
}
|
|
return node
|
|
})
|
|
const itemMap = Object.fromEntries(items.map((item) => [item.key, item]))
|
|
|
|
for (const item of items) {
|
|
const key = item.key
|
|
const parentKey = key.split('/').slice(0, -1).join('/')
|
|
|
|
if (parentKey === folder) {
|
|
pathIndexItem.children!.push(item)
|
|
continue
|
|
}
|
|
|
|
const parentItem = itemMap[parentKey]
|
|
if (parentItem) {
|
|
parentItem.children ??= []
|
|
parentItem.children.push(item)
|
|
}
|
|
}
|
|
|
|
root.push(pathIndexItem)
|
|
}
|
|
|
|
return root
|
|
})
|
|
|
|
return {
|
|
pathOptions,
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Editable preview image.
|
|
*
|
|
* In edit mode, there are 4 methods for setting a preview picture:
|
|
* 1. default value, which is the default image of the model type
|
|
* 2. network picture
|
|
* 3. local file
|
|
* 4. no preview
|
|
*/
|
|
const previewKey = Symbol('preview') as InjectionKey<
|
|
ReturnType<typeof useModelPreviewEditor>
|
|
>
|
|
|
|
export const useModelPreviewEditor = (formInstance: ModelFormInstance) => {
|
|
const { formData: model, registerReset, registerSubmit } = formInstance
|
|
|
|
const typeOptions = ref(['default', 'network', 'local', 'none'])
|
|
const currentType = ref('default')
|
|
|
|
/**
|
|
* Default images
|
|
*/
|
|
const defaultContent = computed(() => {
|
|
return model.value.preview ? castArray(model.value.preview) : []
|
|
})
|
|
const defaultContentPage = ref(0)
|
|
|
|
/**
|
|
* Network picture url
|
|
*/
|
|
const networkContent = ref<string>()
|
|
|
|
/**
|
|
* Local file url
|
|
*/
|
|
const localContent = ref<string>()
|
|
const localContentType = ref<string>()
|
|
const updateLocalContent = async (event: SelectEvent) => {
|
|
const { files } = event
|
|
localContent.value = files[0].objectURL
|
|
localContentType.value = files[0].type
|
|
}
|
|
|
|
/**
|
|
* No preview
|
|
*/
|
|
const noPreviewContent = computed(() => {
|
|
const folder = model.value.type || 'unknown'
|
|
return `/model-manager/preview/${folder}/0/no-preview.png`
|
|
})
|
|
|
|
const preview = computed(() => {
|
|
let content: string | undefined
|
|
|
|
switch (currentType.value) {
|
|
case 'default':
|
|
content = defaultContent.value[defaultContentPage.value]
|
|
break
|
|
case 'network':
|
|
content = networkContent.value
|
|
break
|
|
case 'local':
|
|
content = localContent.value
|
|
break
|
|
default:
|
|
content = undefined
|
|
break
|
|
}
|
|
|
|
return content
|
|
})
|
|
|
|
onMounted(() => {
|
|
registerReset(() => {
|
|
currentType.value = 'default'
|
|
defaultContentPage.value = 0
|
|
networkContent.value = undefined
|
|
localContent.value = undefined
|
|
localContentType.value = undefined
|
|
})
|
|
|
|
registerSubmit((data) => {
|
|
data.preview = preview.value
|
|
})
|
|
})
|
|
|
|
const result = {
|
|
preview,
|
|
typeOptions,
|
|
currentType,
|
|
// default value
|
|
defaultContent,
|
|
defaultContentPage,
|
|
// network picture
|
|
networkContent,
|
|
// local file
|
|
localContent,
|
|
localContentType,
|
|
updateLocalContent,
|
|
// no preview
|
|
noPreviewContent,
|
|
}
|
|
|
|
provide(previewKey, result)
|
|
|
|
return result
|
|
}
|
|
|
|
export const useModelPreview = () => {
|
|
return inject(previewKey)!
|
|
}
|
|
|
|
/**
|
|
* Model description
|
|
*/
|
|
const descriptionKey = Symbol('description') as InjectionKey<
|
|
ReturnType<typeof useModelDescriptionEditor>
|
|
>
|
|
|
|
export const useModelDescriptionEditor = (formInstance: ModelFormInstance) => {
|
|
const { formData: model, metadata } = formInstance
|
|
|
|
const md = useMarkdown({ metadata: metadata.value })
|
|
|
|
const description = computed({
|
|
get: () => {
|
|
return model.value.description
|
|
},
|
|
set: (val) => {
|
|
model.value.description = val
|
|
},
|
|
})
|
|
|
|
const renderedDescription = computed(() => {
|
|
return description.value ? md.render(description.value) : undefined
|
|
})
|
|
|
|
const result = { renderedDescription, description }
|
|
|
|
provide(descriptionKey, result)
|
|
|
|
return result
|
|
}
|
|
|
|
export const useModelDescription = () => {
|
|
return inject(descriptionKey)!
|
|
}
|
|
|
|
/**
|
|
* Model metadata
|
|
*/
|
|
const metadataKey = Symbol('metadata') as InjectionKey<
|
|
ReturnType<typeof useModelMetadataEditor>
|
|
>
|
|
|
|
export const useModelMetadataEditor = (formInstance: ModelFormInstance) => {
|
|
const { formData: model } = formInstance
|
|
|
|
const metadata = computed(() => {
|
|
return model.value.metadata
|
|
})
|
|
|
|
const result = { metadata }
|
|
|
|
provide(metadataKey, result)
|
|
|
|
return result
|
|
}
|
|
|
|
export const useModelMetadata = () => {
|
|
return inject(metadataKey)!
|
|
}
|
|
|
|
export const useModelNodeAction = () => {
|
|
const { t } = useI18n()
|
|
const { toast, wrapperToastError } = useToast()
|
|
|
|
const createNode = (model: BaseModel, options: Record<string, any> = {}) => {
|
|
const nodeType = resolveModelTypeLoader(model.type)
|
|
if (!nodeType) {
|
|
throw new Error(t('unSupportedModelType', [model.type]))
|
|
}
|
|
|
|
const node = window.LiteGraph.createNode(nodeType, null, options)
|
|
const widgetIndex = node.widgets.findIndex((w) => w.type === 'combo')
|
|
if (widgetIndex > -1) {
|
|
node.widgets[widgetIndex].value = genModelFullName(model)
|
|
}
|
|
return node
|
|
}
|
|
|
|
const dragToAddModelNode = wrapperToastError(
|
|
(model: BaseModel, event: DragEvent) => {
|
|
// const target = document.elementFromPoint(event.clientX, event.clientY)
|
|
// if (
|
|
// target?.tagName.toLocaleLowerCase() === 'canvas' &&
|
|
// target.id === 'graph-canvas'
|
|
// ) {
|
|
// const pos = app.clientPosToCanvasPos([event.clientX - 20, event.clientY])
|
|
// const node = createNode({ pos })
|
|
// app.graph.add(node)
|
|
// app.canvas.selectNode(node)
|
|
// }
|
|
//
|
|
// Use the legacy method instead
|
|
const removeEmbeddingExtension = true
|
|
const strictDragToAdd = false
|
|
const splitter = systemStat.value?.system.os === 'nt' ? '\\' : '/'
|
|
|
|
ModelGrid.dragAddModel(
|
|
event,
|
|
model.type,
|
|
genModelFullName(model, splitter),
|
|
removeEmbeddingExtension,
|
|
strictDragToAdd,
|
|
)
|
|
},
|
|
)
|
|
|
|
const addModelNode = wrapperToastError((model: BaseModel) => {
|
|
const selectedNodes = app.canvas.selected_nodes
|
|
const firstSelectedNode = Object.values(selectedNodes)[0]
|
|
const offset = 25
|
|
const pos = firstSelectedNode
|
|
? [firstSelectedNode.pos[0] + offset, firstSelectedNode.pos[1] + offset]
|
|
: app.canvas.canvas_mouse
|
|
const node = createNode(model, { pos })
|
|
app.graph.add(node)
|
|
app.canvas.selectNode(node)
|
|
})
|
|
|
|
const copyModelNode = wrapperToastError((model: BaseModel) => {
|
|
const node = createNode(model)
|
|
app.canvas.copyToClipboard([node])
|
|
toast.add({
|
|
severity: 'success',
|
|
summary: 'Success',
|
|
detail: t('modelCopied'),
|
|
life: 2000,
|
|
})
|
|
})
|
|
|
|
const loadPreviewWorkflow = wrapperToastError(async (model: BaseModel) => {
|
|
const previewUrl = model.preview as string
|
|
const response = await fetch(previewUrl)
|
|
const data = await response.blob()
|
|
const type = data.type
|
|
const extension = type.split('/').pop()
|
|
const file = new File([data], `${model.basename}.${extension}`, { type })
|
|
app.handleFile(file)
|
|
})
|
|
|
|
return {
|
|
addModelNode,
|
|
dragToAddModelNode,
|
|
copyModelNode,
|
|
loadPreviewWorkflow,
|
|
}
|
|
}
|