41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
"""Parse ComfyUI /object_info into usable model lists."""
|
|
|
|
from __future__ import annotations
|
|
|
|
# Node types whose combo inputs we expose in the debug UI
|
|
_MODEL_NODES: dict[str, tuple[str, str]] = {
|
|
"checkpoints": ("CheckpointLoaderSimple", "ckpt_name"),
|
|
"unets": ("UNETLoader", "unet_name"),
|
|
"clips": ("CLIPLoader", "clip_name"),
|
|
"vaes": ("VAELoader", "vae_name"),
|
|
"loras": ("LoraLoader", "lora_name"),
|
|
}
|
|
|
|
|
|
def _combo_options(node_def: dict, input_name: str) -> list[str]:
|
|
if not isinstance(node_def, dict):
|
|
return []
|
|
required = (node_def.get("input") or {}).get("required") or {}
|
|
optional = (node_def.get("input") or {}).get("optional") or {}
|
|
spec = required.get(input_name) or optional.get(input_name)
|
|
if not spec or not isinstance(spec, (list, tuple)):
|
|
return []
|
|
first = spec[0]
|
|
if isinstance(first, list):
|
|
return [str(x) for x in first]
|
|
return []
|
|
|
|
|
|
def parse_model_lists(object_info: dict) -> dict[str, list[str]]:
|
|
out: dict[str, list[str]] = {}
|
|
for key, (node_type, input_name) in _MODEL_NODES.items():
|
|
node_def = object_info.get(node_type) or {}
|
|
options = _combo_options(node_def, input_name)
|
|
if options:
|
|
out[key] = options
|
|
return out
|
|
|
|
|
|
def list_node_types(object_info: dict) -> list[str]:
|
|
return sorted(k for k in object_info.keys() if isinstance(object_info.get(k), dict))
|