361 lines
12 KiB
Python
361 lines
12 KiB
Python
"""Static workflow widget metadata and category hints."""
|
|
|
|
import re
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from ..custom_nodes import (
|
|
get_custom_node_category_hints,
|
|
get_custom_node_widget_categories,
|
|
)
|
|
from ..type_utils import unique_ordered_strings
|
|
|
|
# Workflow widget_values do not always include file extensions. ComfyUI still
|
|
# validates these combo widgets against folder_paths by exact value at queue time.
|
|
NODE_TYPE_MODEL_WIDGET_CATEGORIES = {
|
|
"CheckpointLoaderSimple": {0: "checkpoints"},
|
|
"CheckpointLoader": {1: "checkpoints"},
|
|
"DiffusersLoader": {0: "diffusers"},
|
|
"unCLIPCheckpointLoader": {0: "checkpoints"},
|
|
"ImageOnlyCheckpointLoader": {0: "checkpoints"},
|
|
"VAELoader": {0: "vae"},
|
|
"VAELoaderKJ": {0: "vae"},
|
|
"LoraLoader": {0: "loras"},
|
|
"LoraLoaderModelOnly": {0: "loras"},
|
|
"LoraLoaderBypass": {0: "loras"},
|
|
"LoraLoaderBypassModelOnly": {0: "loras"},
|
|
"CreateHookLora": {0: "loras"},
|
|
"CreateHookLoraModelOnly": {0: "loras"},
|
|
"CreateHookModelAsLora": {0: "checkpoints"},
|
|
"CreateHookModelAsLoraModelOnly": {0: "checkpoints"},
|
|
"UNETLoader": {0: "diffusion_models"},
|
|
"LoaderGGUF": {0: "diffusion_models"},
|
|
"LoaderGGUFAdvanced": {0: "diffusion_models"},
|
|
"UnetLoaderGGUF": {0: "diffusion_models"},
|
|
"UnetLoaderGGUFAdvanced": {0: "diffusion_models"},
|
|
"LatentUpscaleModelLoader": {0: "latent_upscale_models"},
|
|
"CLIPLoader": {0: "text_encoders"},
|
|
"DualCLIPLoader": {0: "text_encoders", 1: "text_encoders"},
|
|
"CLIPLoaderGGUF": {0: "text_encoders"},
|
|
"ClipLoaderGGUF": {0: "text_encoders"},
|
|
"DualCLIPLoaderGGUF": {0: "text_encoders", 1: "text_encoders"},
|
|
"DualClipLoaderGGUF": {0: "text_encoders", 1: "text_encoders"},
|
|
"TripleCLIPLoader": {
|
|
0: "text_encoders",
|
|
1: "text_encoders",
|
|
2: "text_encoders",
|
|
},
|
|
"TripleClipLoader": {
|
|
0: "text_encoders",
|
|
1: "text_encoders",
|
|
2: "text_encoders",
|
|
},
|
|
"TripleCLIPLoaderGGUF": {
|
|
0: "text_encoders",
|
|
1: "text_encoders",
|
|
2: "text_encoders",
|
|
},
|
|
"TripleClipLoaderGGUF": {
|
|
0: "text_encoders",
|
|
1: "text_encoders",
|
|
2: "text_encoders",
|
|
},
|
|
"QuadrupleCLIPLoader": {
|
|
0: "text_encoders",
|
|
1: "text_encoders",
|
|
2: "text_encoders",
|
|
3: "text_encoders",
|
|
},
|
|
"QuadrupleClipLoader": {
|
|
0: "text_encoders",
|
|
1: "text_encoders",
|
|
2: "text_encoders",
|
|
3: "text_encoders",
|
|
},
|
|
"QuadrupleCLIPLoaderGGUF": {
|
|
0: "text_encoders",
|
|
1: "text_encoders",
|
|
2: "text_encoders",
|
|
3: "text_encoders",
|
|
},
|
|
"QuadrupleClipLoaderGGUF": {
|
|
0: "text_encoders",
|
|
1: "text_encoders",
|
|
2: "text_encoders",
|
|
3: "text_encoders",
|
|
},
|
|
"ControlNetLoader": {0: "controlnet"},
|
|
"DiffControlNetLoader": {0: "controlnet"},
|
|
"ControlNetLoaderAdvanced": {0: "controlnet"},
|
|
"ACN_ControlNetLoaderAdvanced": {0: "controlnet"},
|
|
"ACN_DiffControlNetLoaderAdvanced": {0: "controlnet"},
|
|
"CLIPVisionLoader": {0: "clip_vision"},
|
|
"StyleModelLoader": {0: "style_models"},
|
|
"GLIGENLoader": {0: "gligen"},
|
|
"UpscaleModelLoader": {0: "upscale_models"},
|
|
"SAMLoader": {0: "sams"},
|
|
"UltralyticsDetectorProvider": {0: "ultralytics"},
|
|
"AudioEncoderLoader": {0: "audio_encoders"},
|
|
"LoadBackgroundRemovalModel": {0: "background_removal"},
|
|
"LoadDA3Model": {0: "geometry_estimation"},
|
|
"FrameInterpolationModelLoader": {0: "frame_interpolation"},
|
|
"LoadMediaPipeFaceLandmarker": {0: "detection"},
|
|
"ModelPatchLoader": {0: "model_patches"},
|
|
"LoadMoGeModel": {0: "geometry_estimation"},
|
|
"PhotoMakerLoader": {0: "photomaker"},
|
|
"OpticalFlowLoader": {0: "optical_flow"},
|
|
"HypernetworkLoader": {0: "hypernetworks"},
|
|
"EmbeddingLoader": {0: "embeddings"},
|
|
"LTXVAudioVAELoader": {0: "checkpoints"},
|
|
"LowVRAMAudioVAELoader": {0: "checkpoints"},
|
|
"LTXVGemmaCLIPModelLoader": {0: "text_encoders"},
|
|
"LTXAVTextEncoderLoader": {0: "text_encoders", 1: "checkpoints"},
|
|
}
|
|
NODE_TYPE_MODEL_WIDGET_CATEGORIES.update(get_custom_node_widget_categories())
|
|
|
|
# Derived dynamically from NODE_TYPE_MODEL_WIDGET_CATEGORIES.
|
|
NODE_TYPE_TO_CATEGORY_HINTS = {
|
|
node_type: widget_map[min(widget_map.keys())]
|
|
for node_type, widget_map in NODE_TYPE_MODEL_WIDGET_CATEGORIES.items()
|
|
if widget_map
|
|
}
|
|
NODE_TYPE_TO_CATEGORY_HINTS.update(get_custom_node_category_hints())
|
|
|
|
# Model category hints by widget/input name. Workflow JSON does not always preserve
|
|
# widget names, but when it does this catches custom loaders without a node-type entry.
|
|
MODEL_WIDGET_NAME_TO_CATEGORY = {
|
|
"ckpt_name": "checkpoints",
|
|
"checkpoint": "checkpoints",
|
|
"model_name": "diffusion_models",
|
|
"unet_name": "diffusion_models",
|
|
"gguf_name": "diffusion_models",
|
|
"vae_name": "vae",
|
|
"clip_name": "text_encoders",
|
|
"clip_name1": "text_encoders",
|
|
"clip_name2": "text_encoders",
|
|
"clip_name3": "text_encoders",
|
|
"clip_name4": "text_encoders",
|
|
"clip_vision_name": "clip_vision",
|
|
"lora_name": "loras",
|
|
"existing_lora": "loras",
|
|
"control_net_name": "controlnet",
|
|
"cnet": "controlnet",
|
|
"style_model_name": "style_models",
|
|
"upscale_model_name": "upscale_models",
|
|
"gligen_name": "gligen",
|
|
"audio_encoder_name": "audio_encoders",
|
|
"bg_removal_name": "background_removal",
|
|
"photomaker_model_name": "photomaker",
|
|
"sam_model_name": "sams",
|
|
"text_encoder": "text_encoders",
|
|
"hypernetwork_name": "hypernetworks",
|
|
}
|
|
|
|
MODEL_OUTPUT_TYPE_TO_CATEGORY = {
|
|
"UPSCALE_MODEL": "upscale_models",
|
|
"LATENT_UPSCALE_MODEL": "latent_upscale_models",
|
|
"CONTROL_NET": "controlnet",
|
|
"CLIP_VISION": "clip_vision",
|
|
"STYLE_MODEL": "style_models",
|
|
"GLIGEN": "gligen",
|
|
"AUDIO_ENCODER": "audio_encoders",
|
|
"BACKGROUND_REMOVAL": "background_removal",
|
|
"DA3_MODEL": "geometry_estimation",
|
|
"MOGE_MODEL": "geometry_estimation",
|
|
"INTERP_MODEL": "frame_interpolation",
|
|
"FACE_DETECTION_MODEL": "detection",
|
|
"MODEL_PATCH": "model_patches",
|
|
"PHOTOMAKER": "photomaker",
|
|
"OPTICAL_FLOW": "optical_flow",
|
|
"SEEDVR2_DIT": "seedvr2",
|
|
"SEEDVR2_VAE": "seedvr2",
|
|
}
|
|
|
|
# Keys within dict-type widget values that contain model file references.
|
|
NESTED_MODEL_KEYS = {
|
|
"lora": "loras",
|
|
"ckpt_name": "checkpoints",
|
|
"checkpoint": "checkpoints",
|
|
"vae_name": "vae",
|
|
"clip_name": "text_encoders",
|
|
"clip_name1": "text_encoders",
|
|
"clip_name2": "text_encoders",
|
|
"clip_name3": "text_encoders",
|
|
"clip_name4": "text_encoders",
|
|
"control_net_name": "controlnet",
|
|
"cnet": "controlnet",
|
|
"model_name": "diffusion_models",
|
|
"unet_name": "diffusion_models",
|
|
"gguf_name": "diffusion_models",
|
|
"gligen_name": "gligen",
|
|
"audio_encoder_name": "audio_encoders",
|
|
"bg_removal_name": "background_removal",
|
|
"photomaker_model_name": "photomaker",
|
|
"text_encoder": "text_encoders",
|
|
}
|
|
|
|
WORKFLOW_MODEL_WIDGET_NAMES = {
|
|
"model",
|
|
"model_name",
|
|
"model_file",
|
|
"file_name",
|
|
"filename",
|
|
}
|
|
|
|
|
|
def normalize_widget_name(value: Any) -> str:
|
|
return re.sub(r"[_\s-]+", "_", str(value or "").strip().lower()).strip("_")
|
|
|
|
|
|
def _widget_item_name_candidates(item: Any) -> List[str]:
|
|
if not isinstance(item, dict):
|
|
return []
|
|
|
|
candidates = []
|
|
for key in ("name", "label", "localized_name"):
|
|
if item.get(key):
|
|
candidates.append(item.get(key))
|
|
|
|
widget = item.get("widget")
|
|
if isinstance(widget, dict):
|
|
for key in ("name", "label"):
|
|
if widget.get(key):
|
|
candidates.append(widget.get(key))
|
|
elif widget:
|
|
candidates.append(widget)
|
|
|
|
return unique_ordered_strings(candidates)
|
|
|
|
|
|
def _has_widget_input(item: Any) -> bool:
|
|
return isinstance(item, dict) and item.get("widget") is not None
|
|
|
|
|
|
def _get_widget_inputs(node: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|
return [item for item in node.get("inputs", []) if _has_widget_input(item)]
|
|
|
|
|
|
def _get_proxy_widget_entry(node: Dict[str, Any], widget_index: int) -> Any:
|
|
properties = node.get("properties", {})
|
|
proxy_widgets = properties.get("proxyWidgets", [])
|
|
if isinstance(proxy_widgets, list) and widget_index < len(proxy_widgets):
|
|
return proxy_widgets[widget_index]
|
|
return None
|
|
|
|
|
|
def _proxy_widget_name(proxy_entry: Any) -> str:
|
|
if isinstance(proxy_entry, (list, tuple)) and len(proxy_entry) >= 2:
|
|
return str(proxy_entry[1] or "").strip()
|
|
if isinstance(proxy_entry, dict):
|
|
for key in (
|
|
"name",
|
|
"widget_name",
|
|
"widgetName",
|
|
"targetWidgetName",
|
|
"inputName",
|
|
):
|
|
if proxy_entry.get(key):
|
|
return str(proxy_entry.get(key)).strip()
|
|
return ""
|
|
|
|
|
|
def _proxy_widget_node_id(proxy_entry: Any) -> str:
|
|
if isinstance(proxy_entry, (list, tuple)) and proxy_entry:
|
|
return str(proxy_entry[0] or "").strip()
|
|
if isinstance(proxy_entry, dict):
|
|
for key in ("node_id", "nodeId", "target_node_id", "targetNodeId"):
|
|
if proxy_entry.get(key) is not None:
|
|
return str(proxy_entry.get(key)).strip()
|
|
return ""
|
|
|
|
|
|
def get_widget_name_candidates(node: Dict[str, Any], widget_index: int) -> List[str]:
|
|
candidates = []
|
|
|
|
proxy_name = _proxy_widget_name(_get_proxy_widget_entry(node, widget_index))
|
|
if proxy_name:
|
|
candidates.append(proxy_name)
|
|
|
|
widgets = node.get("widgets", [])
|
|
if isinstance(widgets, list) and widget_index < len(widgets):
|
|
candidates.extend(_widget_item_name_candidates(widgets[widget_index]))
|
|
|
|
widget_inputs = _get_widget_inputs(node)
|
|
if widget_index < len(widget_inputs):
|
|
candidates.extend(_widget_item_name_candidates(widget_inputs[widget_index]))
|
|
elif not widget_inputs:
|
|
inputs = node.get("inputs", [])
|
|
if isinstance(inputs, list) and widget_index < len(inputs):
|
|
candidates.extend(_widget_item_name_candidates(inputs[widget_index]))
|
|
|
|
return unique_ordered_strings(candidates)
|
|
|
|
|
|
def get_widget_name_hint(node: Dict[str, Any], widget_index: int) -> str:
|
|
candidates = get_widget_name_candidates(node, widget_index)
|
|
return candidates[0] if candidates else ""
|
|
|
|
|
|
def is_workflow_model_widget_candidate(node: Dict[str, Any], widget_index: int) -> bool:
|
|
return any(
|
|
normalize_widget_name(candidate) in WORKFLOW_MODEL_WIDGET_NAMES
|
|
for candidate in get_widget_name_candidates(node, widget_index)
|
|
)
|
|
|
|
|
|
def _ordered_unique_categories(values: List[Any]) -> List[str]:
|
|
return unique_ordered_strings([value for value in values if value])
|
|
|
|
|
|
def get_widget_category_hint(
|
|
node: Dict[str, Any], widget_index: int
|
|
) -> Optional[str]:
|
|
"""Return a category inferred from saved widget or input names."""
|
|
for candidate in get_widget_name_candidates(node, widget_index):
|
|
category = MODEL_WIDGET_NAME_TO_CATEGORY.get(normalize_widget_name(candidate))
|
|
if category:
|
|
return category
|
|
return None
|
|
|
|
|
|
def get_node_model_widget_category_hint(
|
|
node_type: str, widget_index: int
|
|
) -> Optional[str]:
|
|
"""Return the model category for a known model selector widget."""
|
|
return NODE_TYPE_MODEL_WIDGET_CATEGORIES.get(node_type, {}).get(widget_index)
|
|
|
|
|
|
def get_node_output_category_hint(node: Dict[str, Any]) -> Optional[str]:
|
|
"""Return a category only when strongly typed model outputs are unambiguous."""
|
|
outputs = node.get("outputs", [])
|
|
if not isinstance(outputs, list):
|
|
return None
|
|
|
|
categories: List[str] = []
|
|
for output in outputs:
|
|
if not isinstance(output, dict):
|
|
continue
|
|
for key in ("type", "name", "label"):
|
|
token = str(output.get(key, "") or "").strip().upper()
|
|
category = MODEL_OUTPUT_TYPE_TO_CATEGORY.get(token)
|
|
if category:
|
|
categories.append(category)
|
|
|
|
unique_categories = _ordered_unique_categories(categories)
|
|
return unique_categories[0] if len(unique_categories) == 1 else None
|
|
|
|
|
|
__all__ = [
|
|
"MODEL_OUTPUT_TYPE_TO_CATEGORY",
|
|
"MODEL_WIDGET_NAME_TO_CATEGORY",
|
|
"NESTED_MODEL_KEYS",
|
|
"NODE_TYPE_MODEL_WIDGET_CATEGORIES",
|
|
"NODE_TYPE_TO_CATEGORY_HINTS",
|
|
"get_node_model_widget_category_hint",
|
|
"get_node_output_category_hint",
|
|
"get_widget_category_hint",
|
|
"get_widget_name_candidates",
|
|
"get_widget_name_hint",
|
|
"is_workflow_model_widget_candidate",
|
|
"normalize_widget_name",
|
|
]
|