mirror of
https://github.com/Comfy-Org/ComfyUI_frontend.git
synced 2026-02-21 23:34:31 +00:00
92 lines
3.2 KiB
TypeScript
92 lines
3.2 KiB
TypeScript
import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
|
|
import { useNodeDefStore } from '@/stores/nodeDefStore'
|
|
import { defineStore } from 'pinia'
|
|
|
|
/** 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]?.[0]
|
|
},
|
|
|
|
/**
|
|
* Get the list of all valid node providers for the given model type name.
|
|
* @param modelType The name of the model type to get the node providers for.
|
|
* @returns The list of all valid node providers for the given model type name.
|
|
*/
|
|
getAllNodeProviders(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] ??= []
|
|
this.modelToNodeMap[modelType].push(nodeProvider)
|
|
},
|
|
|
|
/**
|
|
* Register a node provider for the given simple names.
|
|
* @param modelType The name of the model type to register the node provider for.
|
|
* @param nodeClass The node class name to register.
|
|
* @param key The key to use for the node input.
|
|
*/
|
|
quickRegister(modelType: string, nodeClass: string, key: string) {
|
|
this.registerNodeProvider(
|
|
modelType,
|
|
new ModelNodeProvider(this.nodeDefStore.nodeDefsByName[nodeClass], key)
|
|
)
|
|
},
|
|
|
|
registerDefaults() {
|
|
if (this.haveDefaultsLoaded) {
|
|
return
|
|
}
|
|
if (Object.keys(this.nodeDefStore.nodeDefsByName).length === 0) {
|
|
return
|
|
}
|
|
this.haveDefaultsLoaded = true
|
|
this.quickRegister('checkpoints', 'CheckpointLoaderSimple', 'ckpt_name')
|
|
this.quickRegister(
|
|
'checkpoints',
|
|
'ImageOnlyCheckpointLoader',
|
|
'ckpt_name'
|
|
)
|
|
this.quickRegister('loras', 'LoraLoader', 'lora_name')
|
|
this.quickRegister('loras', 'LoraLoaderModelOnly', 'lora_name')
|
|
this.quickRegister('vae', 'VAELoader', 'vae_name')
|
|
this.quickRegister('controlnet', 'ControlNetLoader', 'control_net_name')
|
|
}
|
|
}
|
|
})
|