mirror of
https://github.com/Comfy-Org/ComfyUI_frontend.git
synced 2026-02-26 01:34:07 +00:00
[Draft] Model library sidebar tab (#837)
* basic/empty model library sidebar tab in-progress * make it actually list out models * extremely primitive search impl * list out available folders (incomplete list atm) * load list dynamically * nice lil loading icon * that's not doing anything * run autoformatter * fix up some absolute vue shenanigans * swap to pi-box * is_fake_object * i think apply the tailwind thingo * trim '.safetensors' from end of display title * oop * after load, retain title if no new title is given * is_load_requested to prevent duplication * dirty initial model metadata load & preview based on node preview code * update model store tests * initial image icon for model lib * i hate this * better empty spacer * add api handler for '/models' * load model folders list instead of hardcoding * add a 'no content' placeholder for empty folders * autoformat * autoload model metadata * error handling on metadata loading * larger model icons * click a model to spawn a node for it * draggable model nodes * add a setting for whether to autoload or not * autoformat will be the death of me * cleanup promise code * make the model preview actually half-decent * revert bad unchecked change * put registration back
This commit is contained in:
committed by
Chenlei Hu
parent
bf7652227a
commit
6a158d46b8
@@ -10,8 +10,8 @@
|
||||
]"
|
||||
ref="container"
|
||||
>
|
||||
<div class="node-content truncate">
|
||||
<span class="node-label text-sm">
|
||||
<div class="node-content">
|
||||
<span class="node-label">
|
||||
<slot name="before-label" :node="props.node"></slot>
|
||||
<EditableText
|
||||
:modelValue="node.label"
|
||||
|
||||
@@ -42,6 +42,8 @@ import {
|
||||
import type { RenderedTreeExplorerNode } from '@/types/treeExplorerTypes'
|
||||
import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
|
||||
import { useCanvasStore } from '@/stores/graphStore'
|
||||
import { ComfyModelDef } from '@/stores/modelStore'
|
||||
import { useModelToNodeStore } from '@/stores/modelToNodeStore'
|
||||
|
||||
const emit = defineEmits(['ready'])
|
||||
const canvasRef = ref<HTMLCanvasElement | null>(null)
|
||||
@@ -49,7 +51,7 @@ const settingStore = useSettingStore()
|
||||
const nodeDefStore = useNodeDefStore()
|
||||
const workspaceStore = useWorkspaceStore()
|
||||
const canvasStore = useCanvasStore()
|
||||
|
||||
const modelToNodeStore = useModelToNodeStore()
|
||||
const betaMenuEnabled = computed(
|
||||
() => settingStore.get('Comfy.UseNewMenu') !== 'Disabled'
|
||||
)
|
||||
@@ -132,6 +134,22 @@ onMounted(async () => {
|
||||
loc.clientY
|
||||
])
|
||||
comfyApp.addNodeOnGraph(nodeDef, { pos })
|
||||
} else if (node.data instanceof ComfyModelDef) {
|
||||
const model = node.data
|
||||
const provider = modelToNodeStore.getNodeProvider(model.directory)
|
||||
if (provider) {
|
||||
const pos = comfyApp.clientPosToCanvasPos([
|
||||
loc.clientX - 20,
|
||||
loc.clientY
|
||||
])
|
||||
const node = comfyApp.addNodeOnGraph(provider.nodeDef, { pos })
|
||||
const widget = node.widgets.find(
|
||||
(widget) => widget.name === provider.key
|
||||
)
|
||||
if (widget) {
|
||||
widget.value = model.name
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
209
src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
Normal file
209
src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
Normal file
@@ -0,0 +1,209 @@
|
||||
<template>
|
||||
<SidebarTabTemplate :title="$t('sideToolbar.modelLibrary')">
|
||||
<template #tool-buttons> </template>
|
||||
<template #body>
|
||||
<div class="flex flex-col h-full">
|
||||
<div class="flex-shrink-0">
|
||||
<SearchBox
|
||||
class="model-lib-search-box mx-4 mt-4"
|
||||
v-model:modelValue="searchQuery"
|
||||
@search="handleSearch"
|
||||
:placeholder="$t('searchModels') + '...'"
|
||||
/>
|
||||
</div>
|
||||
<div class="flex-grow overflow-y-auto">
|
||||
<TreeExplorer
|
||||
class="model-lib-tree-explorer mt-1"
|
||||
:roots="renderedRoot.children"
|
||||
v-model:expandedKeys="expandedKeys"
|
||||
@nodeClick="handleNodeClick"
|
||||
>
|
||||
<template #node="{ node }">
|
||||
<ModelTreeLeaf :node="node" />
|
||||
</template>
|
||||
</TreeExplorer>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
</SidebarTabTemplate>
|
||||
<div id="model-library-model-preview-container" />
|
||||
</template>
|
||||
|
||||
<script setup lang="ts">
|
||||
import SearchBox from '@/components/common/SearchBox.vue'
|
||||
import { useI18n } from 'vue-i18n'
|
||||
import TreeExplorer from '@/components/common/TreeExplorer.vue'
|
||||
import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'
|
||||
import ModelTreeLeaf from '@/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue'
|
||||
import { ComfyModelDef, useModelStore } from '@/stores/modelStore'
|
||||
import { useModelToNodeStore } from '@/stores/modelToNodeStore'
|
||||
import { useSettingStore } from '@/stores/settingStore'
|
||||
import { useTreeExpansion } from '@/hooks/treeHooks'
|
||||
import type {
|
||||
RenderedTreeExplorerNode,
|
||||
TreeExplorerNode
|
||||
} from '@/types/treeExplorerTypes'
|
||||
import { computed, ref, type ComputedRef, watch, toRef } from 'vue'
|
||||
import type { TreeNode } from 'primevue/treenode'
|
||||
import { app } from '@/scripts/app'
|
||||
import { buildTree } from '@/utils/treeUtil'
|
||||
const { t } = useI18n()
|
||||
const modelStore = useModelStore()
|
||||
const modelToNodeStore = useModelToNodeStore()
|
||||
const settingStore = useSettingStore()
|
||||
const searchQuery = ref<string>('')
|
||||
const expandedKeys = ref<Record<string, boolean>>({})
|
||||
const { toggleNodeOnEvent } = useTreeExpansion(expandedKeys)
|
||||
|
||||
const root: ComputedRef<TreeNode> = computed(() => {
|
||||
let modelList: ComfyModelDef[] = []
|
||||
if (!modelStore.modelFolders.length) {
|
||||
modelStore.getModelFolders()
|
||||
}
|
||||
if (settingStore.get('Comfy.ModelLibrary.AutoLoadAll')) {
|
||||
for (let folder of modelStore.modelFolders) {
|
||||
modelStore.getModelsInFolderCached(folder)
|
||||
}
|
||||
}
|
||||
for (let folder of modelStore.modelFolders) {
|
||||
const models = modelStore.modelStoreMap[folder]
|
||||
if (models) {
|
||||
if (Object.values(models.models).length) {
|
||||
modelList.push(...Object.values(models.models))
|
||||
} else {
|
||||
const fakeModel = new ComfyModelDef('(No Content)', folder)
|
||||
fakeModel.is_fake_object = true
|
||||
modelList.push(fakeModel)
|
||||
}
|
||||
} else {
|
||||
const fakeModel = new ComfyModelDef('Loading', folder)
|
||||
fakeModel.is_fake_object = true
|
||||
modelList.push(fakeModel)
|
||||
}
|
||||
}
|
||||
if (searchQuery.value) {
|
||||
const search = searchQuery.value.toLocaleLowerCase()
|
||||
modelList = modelList.filter((model: ComfyModelDef) => {
|
||||
return model.name.toLocaleLowerCase().includes(search)
|
||||
})
|
||||
}
|
||||
const tree: TreeNode = buildTree(modelList, (model: ComfyModelDef) => {
|
||||
return [model.directory, ...model.name.replaceAll('\\', '/').split('/')]
|
||||
})
|
||||
return tree
|
||||
})
|
||||
|
||||
const renderedRoot = computed<TreeExplorerNode<ComfyModelDef>>(() => {
|
||||
const fillNodeInfo = (node: TreeNode): TreeExplorerNode<ComfyModelDef> => {
|
||||
const children = node.children?.map(fillNodeInfo)
|
||||
const model: ComfyModelDef | null =
|
||||
node.leaf && node.data ? node.data : null
|
||||
if (model?.is_fake_object) {
|
||||
if (model.name === '(No Content)') {
|
||||
return {
|
||||
key: node.key,
|
||||
label: t('noContent'),
|
||||
leaf: true,
|
||||
data: node.data,
|
||||
getIcon: (node: TreeExplorerNode<ComfyModelDef>) => {
|
||||
return 'pi pi-file'
|
||||
},
|
||||
children: []
|
||||
}
|
||||
} else {
|
||||
return {
|
||||
key: node.key,
|
||||
label: t('loading') + '...',
|
||||
leaf: true,
|
||||
data: node.data,
|
||||
getIcon: (node: TreeExplorerNode<ComfyModelDef>) => {
|
||||
return 'pi pi-spin pi-spinner'
|
||||
},
|
||||
children: []
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
key: node.key,
|
||||
label: model ? model.title : node.label,
|
||||
leaf: node.leaf,
|
||||
data: node.data,
|
||||
getIcon: (node: TreeExplorerNode<ComfyModelDef>) => {
|
||||
if (node.leaf) {
|
||||
if (node.data && node.data.image) {
|
||||
return 'pi pi-fake-spacer'
|
||||
}
|
||||
return 'pi pi-file'
|
||||
}
|
||||
},
|
||||
children,
|
||||
draggable: node.leaf,
|
||||
handleClick: (
|
||||
node: RenderedTreeExplorerNode<ComfyModelDef>,
|
||||
e: MouseEvent
|
||||
) => {
|
||||
if (node.leaf) {
|
||||
const provider = modelToNodeStore.getNodeProvider(model.directory)
|
||||
if (provider) {
|
||||
const node = app.addNodeOnGraph(provider.nodeDef, {
|
||||
pos: app.getCanvasCenter()
|
||||
})
|
||||
const widget = node.widgets.find(
|
||||
(widget) => widget.name === provider.key
|
||||
)
|
||||
if (widget) {
|
||||
widget.value = model.name
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return fillNodeInfo(root.value)
|
||||
})
|
||||
|
||||
const handleSearch = (query: string) => {
|
||||
// TODO
|
||||
}
|
||||
|
||||
const handleNodeClick = (
|
||||
node: RenderedTreeExplorerNode<ComfyModelDef>,
|
||||
e: MouseEvent
|
||||
) => {
|
||||
if (node.leaf) {
|
||||
// TODO
|
||||
} else {
|
||||
toggleNodeOnEvent(e, node)
|
||||
}
|
||||
}
|
||||
|
||||
watch(
|
||||
toRef(expandedKeys, 'value'),
|
||||
(newExpandedKeys) => {
|
||||
Object.entries(newExpandedKeys).forEach(([key, isExpanded]) => {
|
||||
if (isExpanded) {
|
||||
const folderPath = key.split('/').slice(1).join('/')
|
||||
if (folderPath && !folderPath.includes('/')) {
|
||||
// Trigger (async) load of model data for this folder
|
||||
modelStore.getModelsInFolderCached(folderPath)
|
||||
}
|
||||
}
|
||||
})
|
||||
},
|
||||
{ deep: true }
|
||||
)
|
||||
</script>
|
||||
|
||||
<style>
|
||||
.pi-fake-spacer {
|
||||
height: 1px;
|
||||
width: 16px;
|
||||
}
|
||||
</style>
|
||||
|
||||
<style scoped>
|
||||
:deep(.comfy-vue-side-bar-body) {
|
||||
background: var(--p-tree-background);
|
||||
}
|
||||
</style>
|
||||
89
src/components/sidebar/tabs/modelLibrary/ModelPreview.vue
Normal file
89
src/components/sidebar/tabs/modelLibrary/ModelPreview.vue
Normal file
@@ -0,0 +1,89 @@
|
||||
<template>
|
||||
<div class="model_preview">
|
||||
<div class="model_preview_title">
|
||||
{{ modelDef.title }}
|
||||
</div>
|
||||
<div class="model_preview_top_container">
|
||||
<div class="model_preview_architecture" v-if="modelDef.architecture_id">
|
||||
<span class="model_preview_prefix">Architecture: </span>
|
||||
{{ modelDef.architecture_id }}
|
||||
</div>
|
||||
<div class="model_preview_author" v-if="modelDef.author">
|
||||
<span class="model_preview_prefix">Author: </span>
|
||||
{{ modelDef.author }}
|
||||
</div>
|
||||
</div>
|
||||
<div class="model_preview_image" v-if="modelDef.image">
|
||||
<img :src="modelDef.image" />
|
||||
</div>
|
||||
<div class="model_preview_usage_hint" v-if="modelDef.usage_hint">
|
||||
<span class="model_preview_prefix">Usage hint: </span>
|
||||
{{ modelDef.usage_hint }}
|
||||
</div>
|
||||
<div class="model_preview_trigger_phrase" v-if="modelDef.trigger_phrase">
|
||||
<span class="model_preview_prefix">Trigger phrase: </span>
|
||||
{{ modelDef.trigger_phrase }}
|
||||
</div>
|
||||
<div class="model_preview_description" v-if="modelDef.description">
|
||||
<span class="model_preview_prefix">Description: </span>
|
||||
{{ modelDef.description }}
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<script setup lang="ts">
|
||||
import { ComfyModelDef } from '@/stores/modelStore'
|
||||
|
||||
const props = defineProps({
|
||||
modelDef: {
|
||||
type: ComfyModelDef,
|
||||
required: true
|
||||
}
|
||||
})
|
||||
|
||||
const modelDef = props.modelDef
|
||||
</script>
|
||||
<style scoped>
|
||||
.model_preview {
|
||||
background-color: var(--comfy-menu-bg);
|
||||
font-family: 'Open Sans', sans-serif;
|
||||
color: var(--descrip-text);
|
||||
border: 1px solid var(--descrip-text);
|
||||
min-width: 300px;
|
||||
max-width: 500px;
|
||||
width: fit-content;
|
||||
height: fit-content;
|
||||
z-index: 9999;
|
||||
border-radius: 12px;
|
||||
overflow: hidden;
|
||||
font-size: 12px;
|
||||
padding: 10px;
|
||||
}
|
||||
.model_preview_image {
|
||||
margin: auto;
|
||||
width: fit-content;
|
||||
}
|
||||
.model_preview_image img {
|
||||
max-width: 100%;
|
||||
max-height: 150px;
|
||||
object-fit: contain;
|
||||
}
|
||||
.model_preview_title {
|
||||
font-weight: bold;
|
||||
text-align: center;
|
||||
font-size: 14px;
|
||||
}
|
||||
.model_preview_top_container {
|
||||
text-align: center;
|
||||
}
|
||||
.model_preview_author,
|
||||
.model_preview_architecture {
|
||||
display: inline-block;
|
||||
text-align: center;
|
||||
margin: 5px;
|
||||
font-size: 10px;
|
||||
}
|
||||
.model_preview_prefix {
|
||||
font-weight: bold;
|
||||
}
|
||||
</style>
|
||||
130
src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue
Normal file
130
src/components/sidebar/tabs/modelLibrary/ModelTreeLeaf.vue
Normal file
@@ -0,0 +1,130 @@
|
||||
<template>
|
||||
<div ref="container" class="model-lib-node-container h-full w-full">
|
||||
<TreeExplorerTreeNode :node="node">
|
||||
<template #before-label>
|
||||
<span
|
||||
v-if="modelDef && modelDef.image"
|
||||
class="model-lib-model-icon-container"
|
||||
>
|
||||
<span
|
||||
class="model-lib-model-icon"
|
||||
:style="{ backgroundImage: `url(${modelDef.image})` }"
|
||||
>
|
||||
</span>
|
||||
</span>
|
||||
</template>
|
||||
</TreeExplorerTreeNode>
|
||||
|
||||
<teleport v-if="isHovered" to="#model-library-model-preview-container">
|
||||
<div class="model-lib-model-preview" :style="modelPreviewStyle">
|
||||
<ModelPreview ref="previewRef" :modelDef="modelDef"></ModelPreview>
|
||||
</div>
|
||||
</teleport>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<script setup lang="ts">
|
||||
import TreeExplorerTreeNode from '@/components/common/TreeExplorerTreeNode.vue'
|
||||
import ModelPreview from './ModelPreview.vue'
|
||||
import { ComfyModelDef } from '@/stores/modelStore'
|
||||
import { RenderedTreeExplorerNode } from '@/types/treeExplorerTypes'
|
||||
import {
|
||||
computed,
|
||||
CSSProperties,
|
||||
nextTick,
|
||||
onMounted,
|
||||
onUnmounted,
|
||||
ref
|
||||
} from 'vue'
|
||||
import { useSettingStore } from '@/stores/settingStore'
|
||||
|
||||
const props = defineProps<{
|
||||
node: RenderedTreeExplorerNode<ComfyModelDef>
|
||||
}>()
|
||||
|
||||
const modelDef = computed(() => props.node.data)
|
||||
|
||||
const previewRef = ref<InstanceType<typeof ModelPreview> | null>(null)
|
||||
const modelPreviewStyle = ref<CSSProperties>({
|
||||
position: 'absolute',
|
||||
top: '0px',
|
||||
left: '0px'
|
||||
})
|
||||
|
||||
const settingStore = useSettingStore()
|
||||
const sidebarLocation = computed<'left' | 'right'>(() =>
|
||||
settingStore.get('Comfy.Sidebar.Location')
|
||||
)
|
||||
|
||||
const handleModelHover = async () => {
|
||||
if (modelDef.value.is_fake_object) {
|
||||
return
|
||||
}
|
||||
const hoverTarget = modelContentElement.value
|
||||
const targetRect = hoverTarget.getBoundingClientRect()
|
||||
|
||||
const previewHeight = previewRef.value?.$el.offsetHeight || 0
|
||||
const availableSpaceBelow = window.innerHeight - targetRect.bottom
|
||||
|
||||
modelPreviewStyle.value.top =
|
||||
previewHeight > availableSpaceBelow
|
||||
? `${Math.max(0, targetRect.top - (previewHeight - availableSpaceBelow) - 20)}px`
|
||||
: `${targetRect.top - 40}px`
|
||||
if (sidebarLocation.value === 'left') {
|
||||
modelPreviewStyle.value.left = `${targetRect.right}px`
|
||||
} else {
|
||||
modelPreviewStyle.value.left = `${targetRect.left - 400}px`
|
||||
}
|
||||
|
||||
modelDef.value.load()
|
||||
}
|
||||
|
||||
const container = ref<HTMLElement | null>(null)
|
||||
const modelContentElement = ref<HTMLElement | null>(null)
|
||||
const isHovered = ref(false)
|
||||
const handleMouseEnter = async () => {
|
||||
if (modelDef.value.is_fake_object) {
|
||||
return
|
||||
}
|
||||
isHovered.value = true
|
||||
await nextTick()
|
||||
handleModelHover()
|
||||
}
|
||||
const handleMouseLeave = () => {
|
||||
isHovered.value = false
|
||||
}
|
||||
onMounted(() => {
|
||||
modelContentElement.value = container.value?.closest('.p-tree-node-content')
|
||||
modelContentElement.value?.addEventListener('mouseenter', handleMouseEnter)
|
||||
modelContentElement.value?.addEventListener('mouseleave', handleMouseLeave)
|
||||
if (!modelDef.value.is_fake_object) {
|
||||
modelDef.value.load()
|
||||
}
|
||||
})
|
||||
|
||||
onUnmounted(() => {
|
||||
modelContentElement.value?.removeEventListener('mouseenter', handleMouseEnter)
|
||||
modelContentElement.value?.removeEventListener('mouseleave', handleMouseLeave)
|
||||
})
|
||||
</script>
|
||||
|
||||
<style scoped>
|
||||
.model-lib-model-icon-container {
|
||||
display: inline-block;
|
||||
position: relative;
|
||||
left: 0;
|
||||
height: 1.5rem;
|
||||
vertical-align: top;
|
||||
width: 0px;
|
||||
}
|
||||
.model-lib-model-icon {
|
||||
background-size: cover;
|
||||
background-position: center;
|
||||
display: inline-block;
|
||||
position: relative;
|
||||
left: -2.5rem;
|
||||
height: 2rem;
|
||||
width: 2rem;
|
||||
vertical-align: top;
|
||||
}
|
||||
</style>
|
||||
@@ -21,6 +21,7 @@ const messages = {
|
||||
box: 'Box',
|
||||
briefcase: 'Briefcase',
|
||||
error: 'Error',
|
||||
loading: 'Loading',
|
||||
findIssues: 'Find Issues',
|
||||
copyToClipboard: 'Copy to Clipboard',
|
||||
openNewIssue: 'Open New Issue',
|
||||
@@ -39,9 +40,11 @@ const messages = {
|
||||
searchWorkflows: 'Search Workflows',
|
||||
searchSettings: 'Search Settings',
|
||||
searchNodes: 'Search Nodes',
|
||||
searchModels: 'Search Models',
|
||||
noResultsFound: 'No Results Found',
|
||||
searchFailedMessage:
|
||||
"We couldn't find any settings matching your search. Try adjusting your search terms.",
|
||||
noContent: '(No Content)',
|
||||
noTasksFound: 'No Tasks Found',
|
||||
noTasksFoundMessage: 'There are no tasks in the queue.',
|
||||
newFolder: 'New Folder',
|
||||
@@ -54,6 +57,7 @@ const messages = {
|
||||
nodeLibraryTab: {
|
||||
sortOrder: 'Sort Order'
|
||||
},
|
||||
modelLibrary: 'Model Library',
|
||||
queueTab: {
|
||||
showFlatList: 'Show Flat List',
|
||||
backToAllTasks: 'Back to All Tasks',
|
||||
|
||||
@@ -330,6 +330,18 @@ class ComfyApi extends EventTarget {
|
||||
return await res.json()
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets a list of model folder keys (eg ['checkpoints', 'loras', ...])
|
||||
* @returns The list of model folder keys
|
||||
*/
|
||||
async getModelFolders(): Promise<string[]> {
|
||||
const res = await this.fetchApi(`/models`)
|
||||
if (res.status === 404) {
|
||||
return null
|
||||
}
|
||||
return await res.json()
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets a list of models in the specified folder
|
||||
* @param {string} folder The folder to list models from, such as 'checkpoints'
|
||||
@@ -353,7 +365,22 @@ class ComfyApi extends EventTarget {
|
||||
const res = await this.fetchApi(
|
||||
`/view_metadata/${folder}?filename=${encodeURIComponent(model)}`
|
||||
)
|
||||
return await res.json()
|
||||
const rawResponse = await res.text()
|
||||
if (!rawResponse) {
|
||||
return null
|
||||
}
|
||||
try {
|
||||
return JSON.parse(rawResponse)
|
||||
} catch (error) {
|
||||
console.error(
|
||||
'Error viewing metadata',
|
||||
res.status,
|
||||
res.statusText,
|
||||
rawResponse,
|
||||
error
|
||||
)
|
||||
return null
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -254,6 +254,14 @@ export const CORE_SETTINGS: SettingParams[] = [
|
||||
step: 1
|
||||
}
|
||||
},
|
||||
{
|
||||
id: 'Comfy.ModelLibrary.AutoLoadAll',
|
||||
name: 'Automatically load all model folders',
|
||||
tooltip:
|
||||
'If true, all folders will load as soon as you open the model library (this may cause delays while it loads). If false, root level model folders will only load once you click on them.',
|
||||
type: 'boolean',
|
||||
defaultValue: false
|
||||
},
|
||||
{
|
||||
id: 'Comfy.Locale',
|
||||
name: 'Locale',
|
||||
|
||||
@@ -42,54 +42,71 @@ export class ComfyModelDef {
|
||||
image: string = ''
|
||||
/** Whether the model metadata has been loaded from the server, used for `load()` */
|
||||
has_loaded_metadata: boolean = false
|
||||
/** If true, a metadata load request has been triggered, but may or may not yet have finished loading */
|
||||
is_load_requested: boolean = false
|
||||
/** If true, this is a fake model object used as a placeholder for something (eg a loading icon) */
|
||||
is_fake_object: boolean = false
|
||||
|
||||
constructor(name: string, directory: string) {
|
||||
this.name = name
|
||||
this.title = name
|
||||
this.title = name.replaceAll('\\', '/').split('/').pop()
|
||||
if (this.title.endsWith('.safetensors')) {
|
||||
this.title = this.title.slice(0, -'.safetensors'.length)
|
||||
}
|
||||
this.directory = directory
|
||||
}
|
||||
|
||||
/** Loads the model metadata from the server, filling in this object if data is available */
|
||||
async load(): Promise<void> {
|
||||
if (this.has_loaded_metadata) {
|
||||
if (this.has_loaded_metadata || this.is_load_requested) {
|
||||
return
|
||||
}
|
||||
const metadata = await api.viewMetadata(this.directory, this.name)
|
||||
if (!metadata) {
|
||||
return
|
||||
this.is_load_requested = true
|
||||
try {
|
||||
const metadata = await api.viewMetadata(this.directory, this.name)
|
||||
if (!metadata) {
|
||||
return
|
||||
}
|
||||
this.title =
|
||||
_findInMetadata(
|
||||
metadata,
|
||||
'modelspec.title',
|
||||
'title',
|
||||
'display_name',
|
||||
'name'
|
||||
) || this.title
|
||||
this.architecture_id =
|
||||
_findInMetadata(metadata, 'modelspec.architecture', 'architecture') ||
|
||||
''
|
||||
this.author =
|
||||
_findInMetadata(metadata, 'modelspec.author', 'author') || ''
|
||||
this.description =
|
||||
_findInMetadata(metadata, 'modelspec.description', 'description') || ''
|
||||
this.resolution =
|
||||
_findInMetadata(metadata, 'modelspec.resolution', 'resolution') || ''
|
||||
this.usage_hint =
|
||||
_findInMetadata(metadata, 'modelspec.usage_hint', 'usage_hint') || ''
|
||||
this.trigger_phrase =
|
||||
_findInMetadata(
|
||||
metadata,
|
||||
'modelspec.trigger_phrase',
|
||||
'trigger_phrase'
|
||||
) || ''
|
||||
this.image =
|
||||
_findInMetadata(
|
||||
metadata,
|
||||
'modelspec.thumbnail',
|
||||
'thumbnail',
|
||||
'image',
|
||||
'icon'
|
||||
) || ''
|
||||
const tagsCommaSeparated =
|
||||
_findInMetadata(metadata, 'modelspec.tags', 'tags') || ''
|
||||
this.tags = tagsCommaSeparated.split(',').map((tag) => tag.trim())
|
||||
this.has_loaded_metadata = true
|
||||
} catch (error) {
|
||||
console.error('Error loading model metadata', this.name, this, error)
|
||||
}
|
||||
this.title =
|
||||
_findInMetadata(
|
||||
metadata,
|
||||
'modelspec.title',
|
||||
'title',
|
||||
'display_name',
|
||||
'name'
|
||||
) || this.name
|
||||
this.architecture_id =
|
||||
_findInMetadata(metadata, 'modelspec.architecture', 'architecture') || ''
|
||||
this.author = _findInMetadata(metadata, 'modelspec.author', 'author') || ''
|
||||
this.description =
|
||||
_findInMetadata(metadata, 'modelspec.description', 'description') || ''
|
||||
this.resolution =
|
||||
_findInMetadata(metadata, 'modelspec.resolution', 'resolution') || ''
|
||||
this.usage_hint =
|
||||
_findInMetadata(metadata, 'modelspec.usage_hint', 'usage_hint') || ''
|
||||
this.trigger_phrase =
|
||||
_findInMetadata(metadata, 'modelspec.trigger_phrase', 'trigger_phrase') ||
|
||||
''
|
||||
this.image =
|
||||
_findInMetadata(
|
||||
metadata,
|
||||
'modelspec.thumbnail',
|
||||
'thumbnail',
|
||||
'image',
|
||||
'icon'
|
||||
) || ''
|
||||
const tagsCommaSeparated =
|
||||
_findInMetadata(metadata, 'modelspec.tags', 'tags') || ''
|
||||
this.tags = tagsCommaSeparated.split(',').map((tag) => tag.trim())
|
||||
this.has_loaded_metadata = true
|
||||
}
|
||||
}
|
||||
|
||||
@@ -110,27 +127,42 @@ export class ModelStore {
|
||||
}
|
||||
}
|
||||
|
||||
const folderBlacklist = ['configs', 'custom_nodes']
|
||||
|
||||
/** Model store handler, wraps individual per-folder model stores */
|
||||
export const useModelStore = defineStore('modelStore', {
|
||||
state: () => ({
|
||||
modelStoreMap: {} as Record<string, ModelStore>
|
||||
modelStoreMap: {} as Record<string, ModelStore>,
|
||||
isLoading: {} as Record<string, Promise<ModelStore>>,
|
||||
modelFolders: [] as string[]
|
||||
}),
|
||||
actions: {
|
||||
async getModelsInFolderCached(folder: string): Promise<ModelStore> {
|
||||
if (folder in this.modelStoreMap) {
|
||||
return this.modelStoreMap[folder]
|
||||
}
|
||||
// TODO: needs a lock to avoid overlapping calls
|
||||
const models = await api.getModels(folder)
|
||||
if (!models) {
|
||||
return null
|
||||
if (this.isLoading[folder]) {
|
||||
return this.isLoading[folder]
|
||||
}
|
||||
const store = new ModelStore(folder, models)
|
||||
this.modelStoreMap[folder] = store
|
||||
return store
|
||||
const promise = api.getModels(folder).then((models) => {
|
||||
if (!models) {
|
||||
return null
|
||||
}
|
||||
const store = new ModelStore(folder, models)
|
||||
this.modelStoreMap[folder] = store
|
||||
this.isLoading[folder] = false
|
||||
return store
|
||||
})
|
||||
this.isLoading[folder] = promise
|
||||
return promise
|
||||
},
|
||||
clearCache() {
|
||||
this.modelStoreMap = {}
|
||||
},
|
||||
async getModelFolders() {
|
||||
this.modelFolders = (await api.getModelFolders()).filter(
|
||||
(folder) => !folderBlacklist.includes(folder)
|
||||
)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
86
src/stores/modelToNodeStore.ts
Normal file
86
src/stores/modelToNodeStore.ts
Normal file
@@ -0,0 +1,86 @@
|
||||
import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
|
||||
import { useNodeDefStore } from '@/stores/nodeDefStore'
|
||||
import { defineStore } from 'pinia'
|
||||
import { toRaw } from 'vue'
|
||||
|
||||
/** Helper class that defines how to construct a node from a model. */
|
||||
export class ModelNodeProvider {
|
||||
/** The node definition to use for this model. */
|
||||
public nodeDef: ComfyNodeDefImpl
|
||||
|
||||
/** The node input key for where to inside the model name. */
|
||||
public key: string
|
||||
|
||||
constructor(nodeDef: ComfyNodeDefImpl, key: string) {
|
||||
this.nodeDef = nodeDef
|
||||
this.key = key
|
||||
}
|
||||
}
|
||||
|
||||
/** Service for mapping model types (by folder name) to nodes. */
|
||||
export const useModelToNodeStore = defineStore('modelToNode', {
|
||||
state: () => ({
|
||||
modelToNodeMap: {} as Record<string, ModelNodeProvider>,
|
||||
nodeDefStore: useNodeDefStore(),
|
||||
haveDefaultsLoaded: false
|
||||
}),
|
||||
actions: {
|
||||
/**
|
||||
* Get the node provider for the given model type name.
|
||||
* @param modelType The name of the model type to get the node provider for.
|
||||
* @returns The node provider for the given model type name.
|
||||
*/
|
||||
getNodeProvider(modelType: string): ModelNodeProvider {
|
||||
this.registerDefaults()
|
||||
return this.modelToNodeMap[modelType]
|
||||
},
|
||||
|
||||
/**
|
||||
* Register a node provider for the given model type name.
|
||||
* @param modelType The name of the model type to register the node provider for.
|
||||
* @param nodeProvider The node provider to register.
|
||||
*/
|
||||
registerNodeProvider(modelType: string, nodeProvider: ModelNodeProvider) {
|
||||
this.registerDefaults()
|
||||
this.modelToNodeMap[modelType] = nodeProvider
|
||||
},
|
||||
|
||||
registerDefaults() {
|
||||
if (this.haveDefaultsLoaded) {
|
||||
return
|
||||
}
|
||||
if (Object.keys(this.nodeDefStore.nodeDefsByName).length === 0) {
|
||||
return
|
||||
}
|
||||
this.haveDefaultsLoaded = true
|
||||
this.registerNodeProvider(
|
||||
'checkpoints',
|
||||
new ModelNodeProvider(
|
||||
this.nodeDefStore.nodeDefsByName['CheckpointLoaderSimple'],
|
||||
'ckpt_name'
|
||||
)
|
||||
)
|
||||
this.registerNodeProvider(
|
||||
'loras',
|
||||
new ModelNodeProvider(
|
||||
this.nodeDefStore.nodeDefsByName['LoraLoader'],
|
||||
'lora_name'
|
||||
)
|
||||
)
|
||||
this.registerNodeProvider(
|
||||
'vae',
|
||||
new ModelNodeProvider(
|
||||
this.nodeDefStore.nodeDefsByName['VAELoader'],
|
||||
'vae_name'
|
||||
)
|
||||
)
|
||||
this.registerNodeProvider(
|
||||
'controlnet',
|
||||
new ModelNodeProvider(
|
||||
this.nodeDefStore.nodeDefsByName['ControlNetLoader'],
|
||||
'control_net_name'
|
||||
)
|
||||
)
|
||||
}
|
||||
}
|
||||
})
|
||||
@@ -34,6 +34,7 @@ import {
|
||||
} from '@/stores/workflowStore'
|
||||
import QueueSidebarTab from '@/components/sidebar/tabs/QueueSidebarTab.vue'
|
||||
import NodeLibrarySidebarTab from '@/components/sidebar/tabs/NodeLibrarySidebarTab.vue'
|
||||
import ModelLibrarySidebarTab from '@/components/sidebar/tabs/ModelLibrarySidebarTab.vue'
|
||||
import GlobalToast from '@/components/toast/GlobalToast.vue'
|
||||
import UnloadWindowConfirmDialog from '@/components/dialog/UnloadWindowConfirmDialog.vue'
|
||||
import BrowserTabTitle from '@/components/BrowserTabTitle.vue'
|
||||
@@ -126,6 +127,14 @@ const init = () => {
|
||||
component: markRaw(NodeLibrarySidebarTab),
|
||||
type: 'vue'
|
||||
})
|
||||
app.extensionManager.registerSidebarTab({
|
||||
id: 'model-library',
|
||||
icon: 'pi pi-box',
|
||||
title: t('sideToolbar.modelLibrary'),
|
||||
tooltip: t('sideToolbar.modelLibrary'),
|
||||
component: markRaw(ModelLibrarySidebarTab),
|
||||
type: 'vue'
|
||||
})
|
||||
app.extensionManager.registerSidebarTab({
|
||||
id: 'workflows',
|
||||
icon: 'pi pi-folder-open',
|
||||
|
||||
@@ -69,7 +69,8 @@ describe('useModelStore', () => {
|
||||
const folderStore = await store.getModelsInFolderCached('checkpoints')
|
||||
const model = folderStore.models['noinfo.safetensors']
|
||||
await model.load()
|
||||
expect(model.title).toBe('noinfo.safetensors')
|
||||
expect(model.name).toBe('noinfo.safetensors')
|
||||
expect(model.title).toBe('noinfo')
|
||||
expect(model.architecture_id).toBe('')
|
||||
expect(model.author).toBe('')
|
||||
expect(model.description).toBe('')
|
||||
|
||||
Reference in New Issue
Block a user