mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-06 22:10:14 -03:00
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:
@@ -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)}")
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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):
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user