better badges for empty/loading model library folders

src/components/common/TreeExplorer.vue

@@ -82,14 +82,18 @@ const getTreeNodeIcon = (node: TreeExplorerNode) => {
 }
 const fillNodeInfo = (node: TreeExplorerNode): RenderedTreeExplorerNode => {
   const children = node.children?.map(fillNodeInfo)
+  const totalLeaves = node.leaf
+    ? 1
+    : children.reduce((acc, child) => acc + child.totalLeaves, 0)
   return {
     ...node,
     icon: getTreeNodeIcon(node),
     children,
     type: node.leaf ? 'node' : 'folder',
-    totalLeaves: node.leaf
-      ? 1
-      : children.reduce((acc, child) => acc + child.totalLeaves, 0)
+    totalLeaves,
+    badgeText: node.getBadgeText
+      ? node.getBadgeText(node)
+      : totalLeaves.toString()
   }
 }
 const onNodeContentClick = async (

src/components/common/TreeExplorerTreeNode.vue

@@ -22,7 +22,7 @@
       </span>
       <Badge
         v-if="!props.node.leaf"
-        :value="props.node.totalLeaves"
+        :value="props.node.badgeText"
         severity="secondary"
         class="leaf-count-badge"
       />

src/components/common/__tests__/TreeExplorerTreeNode.spec.ts

@@ -26,7 +26,8 @@
     totalLeaves: 3,
     icon: 'pi pi-folder',
     type: 'folder',
-    handleRename: () => {}
+    handleRename: () => {},
+    badgeText: '3'
   } as RenderedTreeExplorerNode
 
   beforeAll(() => {
@@ -56,7 +57,7 @@
     expect(wrapper.findComponent(EditableText).props('modelValue')).toBe(
       'Test Node'
     )
     expect(wrapper.findComponent(Badge).props()['value']).toBe(3)
   })
 
   it('makes node label editable when renamingEditingNode matches', async () => {

src/components/sidebar/tabs/ModelLibrarySidebarTab.vue

@@ -137,6 +137,31 @@ const renderedRoot = computed<TreeExplorerNode<ComfyModelDef>>(() => {
           return 'pi pi-file'
         }
       },
+      getBadgeText: (node: TreeExplorerNode<ComfyModelDef>) => {
+        if (node.leaf) {
+          return ''
+        }
+        if (node.children?.length === 1) {
+          const onlyChild = node.children[0]
+          if (onlyChild.data?.is_fake_object) {
+            if (onlyChild.data.name === '(No Content)') {
+              return '0'
+            } else if (onlyChild.data.name === 'Loading') {
+              return '?'
+            }
+          }
+        }
+        const getTotalLeaves = (node: TreeExplorerNode<ComfyModelDef>) => {
+          if (node.leaf) {
+            return 1
+          }
+          return node.children.reduce(
+            (acc, child) => acc + getTotalLeaves(child),
+            0
+          )
+        }
+        return getTotalLeaves(node).toString()
+      },
       children,
       draggable: node.leaf,
       handleClick: (

src/types/treeExplorerTypes.ts

@@ -7,36 +7,39 @@ export interface TreeExplorerNode<T = any> {
   data?: T
   children?: TreeExplorerNode<T>[]
   icon?: string
+  /** Function to override what icon to use for the node */
   getIcon?: (node: TreeExplorerNode<T>) => string
-  // Function to handle renaming the node
+  /** Function to override what text to use for the leaf-count badge on a folder node */
+  getBadgeText?: (node: TreeExplorerNode<T>) => string
+  /** Function to handle renaming the node */
   handleRename?: (
     node: TreeExplorerNode<T>,
     newName: string
   ) => void | Promise<void>
-  // Function to handle deleting the node
+  /** Function to handle deleting the node */
   handleDelete?: (node: TreeExplorerNode<T>) => void | Promise<void>
-  // Function to handle adding a child node
+  /** Function to handle adding a child node */
   handleAddChild?: (
     node: TreeExplorerNode<T>,
     child: TreeExplorerNode<T>
   ) => void | Promise<void>
-  // Whether the node is draggable
+  /** Whether the node is draggable */
   draggable?: boolean
-  // Whether the node is droppable
+  /** Whether the node is droppable */
   droppable?: boolean
-  // Function to handle dropping a node
+  /** Function to handle dropping a node */
   handleDrop?: (
     node: TreeExplorerNode<T>,
     data: TreeExplorerDragAndDropData
   ) => void | Promise<void>
-  // Function to handle clicking a node
+  /** Function to handle clicking a node */
   handleClick?: (
     node: TreeExplorerNode<T>,
     event: MouseEvent
   ) => void | Promise<void>
-  // Function to handle errors
+  /** Function to handle errors */
   handleError?: (error: Error) => void
-  // Extra context menu items
+  /** Extra context menu items */
   contextMenuItems?:
     | MenuItem[]
     | ((targetNode: RenderedTreeExplorerNode) => MenuItem[])
@@ -46,8 +49,10 @@ export interface RenderedTreeExplorerNode<T = any> extends TreeExplorerNode<T> {
   children?: RenderedTreeExplorerNode<T>[]
   icon: string
   type: 'folder' | 'node'
-  // Total number of leaves in the subtree
+  /** Total number of leaves in the subtree */
   totalLeaves: number
+  /** Text to display on the leaf-count badge */
+  badgeText: string
 }
 
 export type TreeExplorerDragAndDropData<T = any> = {