sammykumar-ComfyUI-SwissArmyKnife/dev-docs/web-js/GPU_MODEL_GROUPING_METADATA.md

Enhanced GPU Model Detection with Metadata Extraction

Status: ✅ Implemented
Date: November 29, 2025
Component: Resource Monitor - GPU Profiling
Related Files:

• nodes/utils/resource_monitor/profiler_core.py (lines 1-600)
• web/js/resource_monitor.js (tooltip display)

Problem Statement

Initial GPU model scanning detected 603 individual layer components (WanLayerNorm, WanRMSNorm, WanSelfAttention) instead of showing meaningful parent models. Users saw:

• ❌ 603 entries showing "WanLayerNorm (model, 0.0 GB)"
• ❌ No correlation to user-selected model filenames
• ❌ Technical class names instead of user-friendly names
• ❌ All layers showing 0.0 GB (incorrect sizing)

User Requirement: "The average comfyui user won't understand these levels"

Solution Overview

Enhanced three-phase approach that:

1. Reads safetensors metadata during model loading
2. Groups child layers under parent models
3. Displays user-friendly names with accurate VRAM usage

Example Output

Before: 603 × WanLayerNorm (model, 0.0 GB)

After:

• Wan2_2-T2V-A14B-HIGH (checkpoint, 2.1 GB, 145 layers)
• VACE-Depth (model, 512 MB, 45 layers)
• MyLoRA (lora, 128 MB, 32 layers)

Implementation Details

1. Safetensors Metadata Extraction

File: profiler_core.py (lines 1-20)

# Import safetensors library for metadata reading
try:
    from safetensors import safe_open
    SAFETENSORS_AVAILABLE = True
except ImportError:
    SAFETENSORS_AVAILABLE = False

Installation:

pip3 install safetensors

Purpose: Read model_type and other metadata directly from .safetensors files

2. Pending Model Loads Tracking

File: profiler_core.py (lines 136-141)

# Track pending model loads with metadata (filename, timestamp, metadata)
self._pending_model_loads: List[Dict[str, Any]] = []
self._pending_loads_max_age: float = 5.0  # Keep pending loads for 5 seconds

Purpose: Store model filenames and metadata during node execution for later correlation with loaded modules

3. Model Metadata Capture

File: profiler_core.py (_capture_model_metadata() method, lines 449-489)

Triggered: During start_node() before model execution

Process:

1. Scan node inputs for file path keys: model, checkpoint, ckpt_name, lora_name, etc.
2. Filter for .safetensors files
3. Read metadata using safe_open()
4. Store in _pending_model_loads with timestamp

Example:

# Input: {'ckpt_name': 'Wan2_2-T2V-A14B-HIGH_fp8.safetensors'}
# Reads metadata: {'model_type': 'Wan2_2-T2V-A14B-HIGH', ...}
# Stores: {
#     'filename': 'Wan2_2-T2V-A14B-HIGH_fp8.safetensors',
#     'timestamp': 1732901234.56,
#     'node_type': 'WanVideoModelLoader',
#     'metadata': {'model_type': 'Wan2_2-T2V-A14B-HIGH'}
# }

4. Enhanced GPU Scanning with Grouping

File: profiler_core.py (_scan_gpu_models() method, lines 270-444)

Five-Phase Process:

Phase 1: Collect All CUDA Modules

• Scan all Python objects using gc.get_objects()
• Filter for torch.nn.Module instances
• Build parent/child relationship map using named_children()
• Only count direct parameters (recurse=False)

for param in obj.parameters(recurse=False):  # Only direct params
    if param.device.type == 'cuda':
        size_bytes += param.nelement() * param.element_size()

Phase 2: Identify Top-Level Parents

• Find modules that are NOT children of other modules
• These represent complete models (not individual layers)

top_level_parents = set(all_cuda_modules.keys()) - set(parent_child_map.keys())

Phase 3: Calculate Total VRAM

• For each parent, traverse all descendants using depth-first search
• Sum VRAM across all child layers
• Count total layers and parameters

# Depth-first traversal through module hierarchy
stack = [parent_id]
while stack:
    current_id = stack.pop()
    total_size_bytes += current_data['size_bytes']
    layer_count += 1
    stack.extend(current_data['children'])

Phase 4: Extract Model Names

• Try module attributes: name, model_name, config._name_or_path
• Try module metadata: __metadata__['model_type']
• Match to _pending_model_loads by timestamp
• Use safetensors model_type field if available

model_name = self._extract_model_name(module, model_class)
metadata = self._match_pending_metadata(size_mb, model_class)
if metadata and 'model_type' in metadata:
    model_name = metadata['model_type']

Phase 5: Format Display Names

• Show model type, size (GB or MB), layer count
• Sort by VRAM usage (descending)
• Limit to top 10 models per GPU

display_name = f"{model_name} ({model_type}, {size_str}, {layer_count} layers)"
# Example: "Wan2_2-T2V-A14B-HIGH (checkpoint, 2.1 GB, 145 layers)"

5. Helper Methods

_read_safetensors_metadata(filepath: str)

Lines: 491-527

• Resolves relative paths to common ComfyUI model directories
• Opens safetensors file with safe_open()
• Extracts metadata dictionary
• Returns {'model_type': '...', ...} or None

Paths Checked:

possible_dirs = [
    Path('/workspace/ComfyUI/models/checkpoints'),
    Path('/workspace/ComfyUI/models/loras'),
    Path('/workspace/ComfyUI/models/vae'),
    Path('/workspace/ComfyUI/models/unet'),
    Path('models/checkpoints'),
    Path('models/loras'),
    # ... etc
]

_extract_model_name(module, default_name: str)

Lines: 529-552

• Tries common attribute patterns on PyTorch module
• Checks config._name_or_path for HuggingFace models
• Checks __metadata__['model_type'] for safetensors
• Falls back to class name if nothing found

_match_pending_metadata(size_mb: float, class_name: str)

Lines: 554-578

• Matches loaded module to pending metadata by timestamp
• Returns most recent load with metadata
• Fallback: returns most recent load even without metadata
• Used for correlating GPU modules to user-selected files

_classify_model_type(class_name: str)

Lines: 580-597

• Updated with WanVideo patterns: wanvideo, wan2
• Returns: checkpoint, lora, vae, clip, controlnet, upscaler, depth, model

Technical Challenges Solved

Challenge 1: Layer-Level Detection

Problem: Detecting 603 individual layers instead of 5-10 models

Solution: Build parent/child hierarchy and only report top-level parents

Impact: Reduced from 603 entries to ~5-10 meaningful models

Challenge 2: Zero VRAM Reporting

Problem: All layers showing 0.0 GB

Solution:

• Only count direct parameters with recurse=False
• Traverse children separately and sum VRAM
• Calculate at parent level including all descendants

Impact: Accurate VRAM reporting (e.g., 2.1 GB for Wan2 model)

Challenge 3: Filename Correlation

Problem: No way to match detected modules to user-selected files

Solution:

• Capture filenames during start_node()
• Read safetensors metadata immediately
• Store with timestamp for correlation
• Match by timing when scanning GPU

Impact: Display "Wan2_2-T2V-A14B-HIGH" instead of "WanTransformerDecoder"

Challenge 4: Block Swapping / RAM Offloading

Problem: WanVideo models use block swapping - parts may be in RAM

Solution:

• Only scan CUDA memory (param.device.type == 'cuda')
• Automatically handles offloaded models
• Only reports what's actually on GPU

Impact: Accurate representation of GPU state

Performance Optimizations

1. Scan Caching

Implementation: get_loaded_models() method (lines 600+)

# Cache scan results for 1 second
if self._model_scan_cache_timestamp:
    age = time.time() - self._model_scan_cache_timestamp
    if age < self._model_scan_cache_ttl:
        return self.loaded_models  # Return cached

Impact: Prevents expensive scans on every API request

2. Cache Invalidation

Implementation: end_node() method

# Invalidate if VRAM changed by more than 100MB
vram_change = abs(node_profile.vram_after - node_profile.vram_before)
if vram_change > 100 * 1024 * 1024:  # 100MB threshold
    self._model_scan_cache_timestamp = None

Impact: Fresh scans only when models actually load/unload

3. Garbage Collection

Implementation: _scan_gpu_models() line 288

# Force GC before scanning to get accurate state
gc.collect()

Impact: Ensures dead references are cleaned up before counting

4. Top 10 Limit

Implementation: Phase 5 of scanning

# Sort by VRAM and limit to top 10
result[gpu_id].sort(key=lambda x: x['vram_mb'], reverse=True)
result[gpu_id] = result[gpu_id][:10]

Impact: Prevents overwhelming UI with hundreds of models

5. Pending Loads Cleanup

Implementation: Lines 291-295

# Clean up old pending loads (>5 seconds)
self._pending_model_loads = [
    load for load in self._pending_model_loads
    if current_time - load.get('timestamp', 0) < 5.0
]

Impact: Prevents memory leak from accumulating metadata

Testing & Validation

Test Scenario 1: WanVideo Model Loading

Setup: Load Wan2_2-T2V-A14B-HIGH model with block swapping

Expected Output:

GPU 0: 5 models
  - Wan2_2-T2V-A14B-HIGH (checkpoint, 2.1 GB, 145 layers)
  - VACE-Depth (model, 512 MB, 45 layers)
  - MyLoRA (lora, 128 MB, 32 layers)

Validation:

• ✅ Shows grouped parent model, not 603 layers
• ✅ Accurate VRAM calculation (2.1 GB)
• ✅ Displays user-friendly name from metadata
• ✅ Shows layer count for context

Test Scenario 2: Multiple Model Types

Setup: Load checkpoint + VAE + LoRA + CLIP

Expected Output: 4-5 distinct models sorted by VRAM

Test Scenario 3: Model Unloading

Setup: Clear models and observe cache invalidation

Expected Output: Empty model list after VRAM drops

Test Scenario 4: Without Safetensors Library

Setup: Remove safetensors package

Expected Output: Graceful degradation, shows class names instead of metadata

Dependencies

Required

• torch - PyTorch module inspection
• gc - Python garbage collection

Optional

• safetensors - Metadata extraction (highly recommended)

Installation:

pip3 install safetensors

Graceful Degradation: If safetensors not available:

• Still performs grouping and VRAM calculation
• Uses class names instead of metadata
• Shows warning on startup

Future Enhancements

Enhancement 1: Multi-GPU Tooltip

Current: Single tooltip for all GPUs
Proposed: Separate tooltip per GPU with per-GPU model list

Enhancement 2: Model Age Tracking

Current: Shows only current models
Proposed: Show "loaded 2m ago" timestamp

Enhancement 3: Offloaded Model Indicator

Current: Only shows GPU-resident models
Proposed: Show RAM-offloaded models with "(offloaded)" indicator

Enhancement 4: Model File Paths

Current: Shows model name only
Proposed: Show full file path in tooltip

Enhancement 5: More Metadata Fields

Current: Only uses model_type
Proposed: Extract and display architecture, base_model, etc.

Related Documentation

• Core Implementation: GPU_MEMORY_SCANNING.md - Original scanning approach
• Widget Display: CONTROL_PANEL.md - Frontend tooltip rendering
• API Endpoint: /swissarmyknife/profiler/loaded_models - REST API for model data
• Node Tracking: profiler_core.py - Complete profiler implementation

Troubleshooting

Issue: No Models Shown

Cause: No CUDA modules in memory
Solution: Load a model and wait for node execution to complete

Issue: Shows Class Names Instead of Metadata

Cause: safetensors library not installed
Solution: pip3 install safetensors

Issue: Incorrect VRAM Numbers

Cause: Block swapping or model sharing
Solution: VRAM is per-GPU, shared layers counted once

Issue: Old Models Still Showing

Cause: Cache not invalidated
Solution: Wait 1 second or trigger VRAM change >100MB

Issue: Too Many Models Listed

Cause: Many small models loaded
Solution: Only top 10 shown, sorted by VRAM usage

Logging & Debugging

Console Output:

[SwissArmyKnife][Profiler] 📄 Captured model: Wan2_2-HIGH.safetensors (metadata: True)
[SwissArmyKnife][Profiler] 🔍 Found 603 CUDA modules, 5 top-level parents
[SwissArmyKnife][Profiler] ✅ Grouped into 5 parent models across 1 GPU(s)
[SwissArmyKnife][Profiler] GPU 0: 5 models
[SwissArmyKnife][Profiler]   - Wan2_2-T2V-A14B-HIGH (checkpoint, 2.1 GB, 145 layers)

Enable Debug Logging:

logger.setLevel(logging.DEBUG)

Commit Reference

Feature Branch: feat/Resource-monitor/CU-86b7mzbtx
Implementation Date: November 29, 2025
Key Changes:

• Added safetensors import and metadata extraction
• Enhanced _scan_gpu_models() with 5-phase approach
• Added _capture_model_metadata() for filename tracking
• Added helper methods for name extraction and matching
• Fixed VRAM calculation using parent/child traversal
• Limited display to top 10 models per GPU