feat: add type-signature-based fallback for unregistered nodes

GenericNodeExtractor (previously a no-op) now inspects
RETURN_TYPES to detect MODEL loaders and CONDITIONING
encoders in nodes not registered in NODE_EXTRACTORS.

- Propagate return_types from the hook layer through the
  registry to GenericNodeExtractor.extract() and update().
- MODEL detection: scan ckpt_name/unet_name/model_path/
  model_name/gguf_name fields, validate by extension.
- CONDITIONING detection: scan text/clip_l/t5xxl/prompt
  fields, store prompt text and conditioning tensor.
- _fill_missing_metadata also checks node_cache, so
  GenericNodeExtractor-handled nodes survive cache.
This commit is contained in:
Will Miao
2026-07-25 22:14:18 +08:00
parent e6dc169a05
commit 077e70169d
4 changed files with 102 additions and 22 deletions

View File

@@ -83,7 +83,8 @@ class MetadataHook:
# Record inputs before execution # Record inputs before execution
if node_id is not None: if node_id is not None:
registry.record_node_execution(node_id, class_type, input_data_all, None) return_types = getattr(obj.__class__, 'RETURN_TYPES', None)
registry.record_node_execution(node_id, class_type, input_data_all, None, return_types=return_types)
except Exception as e: except Exception as e:
logger.error(f"Error collecting metadata (pre-execution): {str(e)}") logger.error(f"Error collecting metadata (pre-execution): {str(e)}")
@@ -114,7 +115,8 @@ class MetadataHook:
# Record outputs after execution # Record outputs after execution
if node_id is not None: if node_id is not None:
registry.update_node_execution(node_id, class_type, results) return_types = getattr(obj.__class__, 'RETURN_TYPES', None)
registry.update_node_execution(node_id, class_type, results, return_types=return_types)
except Exception as e: except Exception as e:
logger.error(f"Error collecting metadata (post-execution): {str(e)}") logger.error(f"Error collecting metadata (post-execution): {str(e)}")
@@ -166,7 +168,8 @@ class MetadataHook:
class_type = obj.__class__.__name__ class_type = obj.__class__.__name__
node_id = unique_id node_id = unique_id
if node_id is not None: if node_id is not None:
registry.record_node_execution(node_id, class_type, input_data_all, None) return_types = getattr(obj.__class__, 'RETURN_TYPES', None)
registry.record_node_execution(node_id, class_type, input_data_all, None, return_types=return_types)
except Exception as e: except Exception as e:
logger.error(f"Error collecting metadata (pre-execution): {str(e)}") logger.error(f"Error collecting metadata (pre-execution): {str(e)}")
@@ -183,7 +186,8 @@ class MetadataHook:
class_type = obj.__class__.__name__ class_type = obj.__class__.__name__
node_id = unique_id node_id = unique_id
if node_id is not None: if node_id is not None:
registry.update_node_execution(node_id, class_type, results) return_types = getattr(obj.__class__, 'RETURN_TYPES', None)
registry.update_node_execution(node_id, class_type, results, return_types=return_types)
except Exception as e: except Exception as e:
logger.error(f"Error collecting metadata (post-execution): {str(e)}") logger.error(f"Error collecting metadata (post-execution): {str(e)}")

View File

@@ -128,7 +128,7 @@ class MetadataRegistry:
cache_key = f"{node_id}:{class_type}" cache_key = f"{node_id}:{class_type}"
# Check if this node type is relevant for metadata collection # Check if this node type is relevant for metadata collection
if class_type in NODE_EXTRACTORS: if class_type in NODE_EXTRACTORS or cache_key in self.node_cache:
# Check if we have cached metadata for this node # Check if we have cached metadata for this node
if cache_key in self.node_cache: if cache_key in self.node_cache:
cached_data = self.node_cache[cache_key] cached_data = self.node_cache[cache_key]
@@ -141,7 +141,7 @@ class MetadataRegistry:
node_id node_id
] ]
def record_node_execution(self, node_id, class_type, inputs, outputs): def record_node_execution(self, node_id, class_type, inputs, outputs, return_types=None):
"""Record information about a node's execution""" """Record information about a node's execution"""
if not self.current_prompt_id: if not self.current_prompt_id:
return return
@@ -164,17 +164,18 @@ class MetadataRegistry:
# Extract node-specific metadata # Extract node-specific metadata
extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor) extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor)
extractor.extract( if type(extractor) is GenericNodeExtractor:
node_id, extractor.extract(node_id, processed_inputs, outputs,
processed_inputs, self.prompt_metadata[self.current_prompt_id],
outputs, return_types=return_types)
self.prompt_metadata[self.current_prompt_id], else:
) extractor.extract(node_id, processed_inputs, outputs,
self.prompt_metadata[self.current_prompt_id])
# Cache this node's metadata # Cache this node's metadata
self._cache_node_metadata(node_id, class_type) self._cache_node_metadata(node_id, class_type)
def update_node_execution(self, node_id, class_type, outputs): def update_node_execution(self, node_id, class_type, outputs, return_types=None):
"""Update node metadata with output information""" """Update node metadata with output information"""
if not self.current_prompt_id: if not self.current_prompt_id:
return return
@@ -185,9 +186,17 @@ class MetadataRegistry:
# Use the same extractor to update with outputs # Use the same extractor to update with outputs
extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor) extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor)
if hasattr(extractor, "update"): if hasattr(extractor, "update"):
extractor.update( if type(extractor) is GenericNodeExtractor:
node_id, processed_outputs, self.prompt_metadata[self.current_prompt_id] extractor.update(
) node_id, processed_outputs,
self.prompt_metadata[self.current_prompt_id],
return_types=return_types,
)
else:
extractor.update(
node_id, processed_outputs,
self.prompt_metadata[self.current_prompt_id],
)
# Update the cached metadata for this node # Update the cached metadata for this node
self._cache_node_metadata(node_id, class_type) self._cache_node_metadata(node_id, class_type)

View File

@@ -31,11 +31,78 @@ class NodeMetadataExtractor:
pass pass
class GenericNodeExtractor(NodeMetadataExtractor): class GenericNodeExtractor(NodeMetadataExtractor):
"""Default extractor for nodes without specific handling""" """Fallback extractor with type-signature-based detection.
When a node is not in the NODE_EXTRACTORS registry, the hook layer
passes ``return_types`` from ``obj.RETURN_TYPES``:
* ``MODEL`` output: common input fields (ckpt_name, unet_name, etc.)
are checked for a model file name and stored as checkpoint metadata.
* ``CONDITIONING`` output: common text input fields are checked for
prompt text and stored as prompt metadata.
"""
# Input field names that carry a model path in loader-style nodes.
_MODEL_NAME_FIELDS = (
"ckpt_name", "unet_name", "model_path", "model_name", "gguf_name",
)
# Extensions used by checkpoint_scanner.py — only record values that look
# like real model filenames to avoid capturing unrelated string fields.
_MODEL_EXTENSIONS = {
".ckpt", ".pt", ".pt2", ".bin", ".pth", ".safetensors", ".pkl", ".sft", ".gguf",
}
# Input field names that may carry prompt text in encoder-style nodes.
_TEXT_FIELDS = ("text", "clip_l", "t5xxl", "prompt", "positive", "negative")
@staticmethod @staticmethod
def extract(node_id, inputs, outputs, metadata): def extract(node_id, inputs, outputs, metadata, return_types=None):
pass if return_types is None:
return
# — MODEL loader detection (checkpoint / UNET / GGUF) —
if "MODEL" in return_types or any("MODEL" in str(t) for t in return_types):
for field in GenericNodeExtractor._MODEL_NAME_FIELDS:
val = inputs.get(field)
if val and isinstance(val, str) and val.strip():
name = val.strip()
if not any(name.lower().endswith(ext) for ext in GenericNodeExtractor._MODEL_EXTENSIONS):
continue
_store_checkpoint_metadata(metadata, node_id, name)
return
# — CONDITIONING encoder detection (CLIPTextEncode, Flux, custom) —
if "CONDITIONING" in return_types or any("CONDITIONING" in str(t) for t in return_types):
text = None
for field in GenericNodeExtractor._TEXT_FIELDS:
val = inputs.get(field)
if val and isinstance(val, str) and val.strip():
text = val.strip()
break
if text:
prompt_data = metadata.setdefault(PROMPTS, {})
prompt_data[node_id] = {
"text": text,
"node_id": node_id,
}
@staticmethod
def update(node_id, outputs, metadata, return_types=None):
if return_types is None:
return
if "CONDITIONING" not in return_types and not any(
"CONDITIONING" in str(t) for t in return_types
):
return
if node_id not in metadata.get(PROMPTS, {}):
return
if outputs and isinstance(outputs, list) and len(outputs) > 0:
if isinstance(outputs[0], tuple) and len(outputs[0]) > 0:
cond = outputs[0][0]
if cond is not None:
metadata[PROMPTS][node_id]["conditioning"] = cond
class CheckpointLoaderExtractor(NodeMetadataExtractor): class CheckpointLoaderExtractor(NodeMetadataExtractor):
@staticmethod @staticmethod
def extract(node_id, inputs, outputs, metadata): def extract(node_id, inputs, outputs, metadata):

View File

@@ -30,10 +30,10 @@ def test_metadata_hook_installs_and_traces_execution(monkeypatch, metadata_regis
calls = [] calls = []
def record_stub(self, node_id, class_type, inputs, outputs): def record_stub(self, node_id, class_type, inputs, outputs, return_types=None):
calls.append(("record", node_id, class_type, inputs)) calls.append(("record", node_id, class_type, inputs))
def update_stub(self, node_id, class_type, outputs): def update_stub(self, node_id, class_type, outputs, return_types=None):
calls.append(("update", node_id, class_type, outputs)) calls.append(("update", node_id, class_type, outputs))
monkeypatch.setattr(MetadataRegistry, "record_node_execution", record_stub) monkeypatch.setattr(MetadataRegistry, "record_node_execution", record_stub)