Setmaster
fb28135b82
- Show model file sizes (GB/MB) in the sidebar next to last used time - Use ComfyUI's folder_paths for reliable model path resolution - Support all configured model locations including extra_model_paths - Handle GGUF models that can be in unet or clip folders - Update README with new features
119 lines
3.4 KiB
Python
119 lines
3.4 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"),
|
|
|
|
# GGUF loaders (ComfyUI-GGUF by city96)
|
|
"UnetLoaderGGUF": ("gguf", "unet_name"),
|
|
"UnetLoaderGGUFAdvanced": ("gguf", "unet_name"),
|
|
"CLIPLoaderGGUF": ("gguf", "clip_name"),
|
|
"DualCLIPLoaderGGUF": ("gguf", ["clip_name1", "clip_name2"]),
|
|
"TripleCLIPLoaderGGUF": ("gguf", ["clip_name1", "clip_name2", "clip_name3"]),
|
|
"QuadrupleCLIPLoaderGGUF": ("gguf", ["clip_name1", "clip_name2", "clip_name3", "clip_name4"]),
|
|
}
|
|
|
|
# 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",
|
|
},
|
|
"gguf": {
|
|
"name": "GGUF",
|
|
"icon": "pi-database",
|
|
},
|
|
}
|
|
|
|
# Map categories to ComfyUI model folder names
|
|
CATEGORY_FOLDERS: dict[str, str] = {
|
|
"checkpoint": "checkpoints",
|
|
"lora": "loras",
|
|
"vae": "vae",
|
|
"controlnet": "controlnet",
|
|
"clip": "clip",
|
|
"unet": "unet",
|
|
"upscaler": "upscale_models",
|
|
"style_model": "style_models",
|
|
"gligen": "gligen",
|
|
"gguf": "unet", # GGUF models are stored in unet folder
|
|
}
|