Files
Setmaster-ComfyUI-ModelPulse/py/model_types.py
T
Setmaster e9e3ade3f0 Implement Phase 1: core tracking and sidebar UI
- Add model usage tracking via PromptExecutor hook
- Support 12+ loader node types (checkpoints, LoRAs, VAEs, etc.)
- Pattern matching for custom node loaders
- Cross-platform file locking (Windows/Unix)
- JSON storage with schema migrations
- REST API endpoints for usage data
- Sidebar panel with timeframe filtering and sorting
- Stale model highlighting (>30 days unused)
2026-01-14 22:56:28 -05:00

93 lines
2.5 KiB
Python

"""
ModelPulse - Model Type Definitions
Maps ComfyUI loader node types to their model categories and input keys.
"""
# Standard ComfyUI model loaders
# Format: "NodeClassName": ("category", "input_key" or ["input_key1", "input_key2"])
MODEL_LOADERS: dict[str, tuple[str, str | list[str]]] = {
# Checkpoint loaders
"CheckpointLoaderSimple": ("checkpoint", "ckpt_name"),
"CheckpointLoader": ("checkpoint", "ckpt_name"),
# LoRA loaders
"LoraLoader": ("lora", "lora_name"),
"LoraLoaderModelOnly": ("lora", "lora_name"),
# VAE loaders
"VAELoader": ("vae", "vae_name"),
# ControlNet loaders
"ControlNetLoader": ("controlnet", "control_net_name"),
# CLIP loaders
"CLIPLoader": ("clip", "clip_name"),
"DualCLIPLoader": ("clip", ["clip_name1", "clip_name2"]),
# UNET loaders
"UNETLoader": ("unet", "unet_name"),
# Upscaler loaders
"UpscaleModelLoader": ("upscaler", "model_name"),
# Style model loaders
"StyleModelLoader": ("style_model", "style_model_name"),
# GLIGEN loaders
"GLIGENLoader": ("gligen", "gligen_name"),
}
# Patterns for detecting custom node loaders (Impact Pack, Efficiency Nodes, etc.)
# These are checked if the node class name isn't in MODEL_LOADERS
LOADER_PATTERNS: list[tuple[str, str, list[str]]] = [
# (pattern_in_class_name, category, possible_input_keys)
("Checkpoint", "checkpoint", ["ckpt_name", "checkpoint", "model_name"]),
("Lora", "lora", ["lora_name", "lora"]),
("VAE", "vae", ["vae_name", "vae"]),
("ControlNet", "controlnet", ["control_net_name", "controlnet", "control_net"]),
("CLIP", "clip", ["clip_name", "clip"]),
("UNET", "unet", ["unet_name", "unet"]),
("Upscale", "upscaler", ["model_name", "upscale_model"]),
]
# Model categories for display purposes
MODEL_CATEGORIES: dict[str, dict[str, str]] = {
"checkpoint": {
"name": "Checkpoints",
"icon": "pi-box",
},
"lora": {
"name": "LoRAs",
"icon": "pi-sparkles",
},
"vae": {
"name": "VAEs",
"icon": "pi-cog",
},
"controlnet": {
"name": "ControlNets",
"icon": "pi-sitemap",
},
"clip": {
"name": "CLIP Models",
"icon": "pi-comments",
},
"unet": {
"name": "UNETs",
"icon": "pi-server",
},
"upscaler": {
"name": "Upscalers",
"icon": "pi-expand",
},
"style_model": {
"name": "Style Models",
"icon": "pi-palette",
},
"gligen": {
"name": "GLIGEN",
"icon": "pi-th-large",
},
}