[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

src/stores/coreSettings.ts

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

src/stores/modelStore.ts

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

src/stores/modelToNodeStore.ts

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