diff --git a/py/metadata_collector/metadata_hook.py b/py/metadata_collector/metadata_hook.py index 8ad9a588..f1931f7c 100644 --- a/py/metadata_collector/metadata_hook.py +++ b/py/metadata_collector/metadata_hook.py @@ -83,7 +83,8 @@ class MetadataHook: # Record inputs before execution 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: logger.error(f"Error collecting metadata (pre-execution): {str(e)}") @@ -114,7 +115,8 @@ class MetadataHook: # Record outputs after execution 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: logger.error(f"Error collecting metadata (post-execution): {str(e)}") @@ -166,7 +168,8 @@ class MetadataHook: class_type = obj.__class__.__name__ node_id = unique_id 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: logger.error(f"Error collecting metadata (pre-execution): {str(e)}") @@ -183,7 +186,8 @@ class MetadataHook: class_type = obj.__class__.__name__ node_id = unique_id 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: logger.error(f"Error collecting metadata (post-execution): {str(e)}") diff --git a/py/metadata_collector/metadata_registry.py b/py/metadata_collector/metadata_registry.py index 734a18b8..ef754a06 100644 --- a/py/metadata_collector/metadata_registry.py +++ b/py/metadata_collector/metadata_registry.py @@ -128,7 +128,7 @@ class MetadataRegistry: cache_key = f"{node_id}:{class_type}" # 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 if cache_key in self.node_cache: cached_data = self.node_cache[cache_key] @@ -141,7 +141,7 @@ class MetadataRegistry: 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""" if not self.current_prompt_id: return @@ -164,17 +164,18 @@ class MetadataRegistry: # Extract node-specific metadata extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor) - extractor.extract( - node_id, - processed_inputs, - outputs, - self.prompt_metadata[self.current_prompt_id], - ) + if type(extractor) is GenericNodeExtractor: + extractor.extract(node_id, processed_inputs, outputs, + self.prompt_metadata[self.current_prompt_id], + return_types=return_types) + else: + extractor.extract(node_id, processed_inputs, outputs, + self.prompt_metadata[self.current_prompt_id]) # Cache this node's metadata 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""" if not self.current_prompt_id: return @@ -185,9 +186,17 @@ class MetadataRegistry: # Use the same extractor to update with outputs extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor) if hasattr(extractor, "update"): - extractor.update( - node_id, processed_outputs, self.prompt_metadata[self.current_prompt_id] - ) + if type(extractor) is GenericNodeExtractor: + 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 self._cache_node_metadata(node_id, class_type) diff --git a/py/metadata_collector/node_extractors.py b/py/metadata_collector/node_extractors.py index 352e5cf4..27f67e54 100644 --- a/py/metadata_collector/node_extractors.py +++ b/py/metadata_collector/node_extractors.py @@ -31,11 +31,78 @@ class NodeMetadataExtractor: pass 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 - def extract(node_id, inputs, outputs, metadata): - pass - + def extract(node_id, inputs, outputs, metadata, return_types=None): + 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): @staticmethod def extract(node_id, inputs, outputs, metadata): diff --git a/tests/metadata_collector/test_metadata_collector.py b/tests/metadata_collector/test_metadata_collector.py index c9b08bb8..e7dc35c4 100644 --- a/tests/metadata_collector/test_metadata_collector.py +++ b/tests/metadata_collector/test_metadata_collector.py @@ -30,10 +30,10 @@ def test_metadata_hook_installs_and_traces_execution(monkeypatch, metadata_regis 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)) - 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)) monkeypatch.setattr(MetadataRegistry, "record_node_execution", record_stub)