Files
ComfyUI_frontend/src/stores/modelToNodeStore.ts
Alex "mcmonkey" Goodwin 6a158d46b8 [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
2024-09-25 16:01:50 +09:00

87 lines
2.5 KiB
TypeScript

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'
)
)
}
}
})