diff --git a/py/metadata_collector/node_extractors.py b/py/metadata_collector/node_extractors.py index ab203132..c1dbdc00 100644 --- a/py/metadata_collector/node_extractors.py +++ b/py/metadata_collector/node_extractors.py @@ -40,7 +40,7 @@ class GenericNodeExtractor(NodeMetadataExtractor): * ``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. + prompt text, and conditioning inputs are tracked through transforms. """ # Input field names that carry a model path in loader-style nodes. @@ -73,7 +73,7 @@ class GenericNodeExtractor(NodeMetadataExtractor): _store_checkpoint_metadata(metadata, node_id, name) return - # — CONDITIONING encoder detection (CLIPTextEncode, Flux, custom) — + # — CONDITIONING encoder / transform detection — if "CONDITIONING" in return_types or any("CONDITIONING" in str(t) for t in return_types): text = None for field in GenericNodeExtractor._TEXT_FIELDS: @@ -81,12 +81,18 @@ class GenericNodeExtractor(NodeMetadataExtractor): 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, - } + + input_conditionings = [ + value + for input_name, value in inputs.items() + if input_name.startswith("conditioning") and value is not None + ] + if text or input_conditionings: + prompt_metadata = _ensure_prompt_metadata(metadata, node_id) + if text: + prompt_metadata["text"] = text + if input_conditionings: + prompt_metadata["orig_conditionings"] = input_conditionings @staticmethod def update(node_id, outputs, metadata, return_types=None): @@ -98,11 +104,22 @@ class GenericNodeExtractor(NodeMetadataExtractor): 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 + output_tuple = _first_output_tuple(outputs) + if not output_tuple or len(output_tuple) < 1: + return + + output_conditioning = output_tuple[0] + if output_conditioning is None: + return + + prompt_metadata = metadata[PROMPTS][node_id] + prompt_metadata["conditioning"] = output_conditioning + _record_conditioning_source( + metadata, + node_id, + output_conditioning, + prompt_metadata.get("orig_conditionings", []), + ) class CheckpointLoaderExtractor(NodeMetadataExtractor): @staticmethod diff --git a/tests/metadata_collector/test_metadata_collector.py b/tests/metadata_collector/test_metadata_collector.py index 57631654..fe6aaf20 100644 --- a/tests/metadata_collector/test_metadata_collector.py +++ b/tests/metadata_collector/test_metadata_collector.py @@ -471,6 +471,172 @@ def test_conditioning_provenance_recovers_combined_controlnet_prompts( assert params["negative_prompt"] == "low quality" +def test_conditioning_provenance_recovers_transformed_switched_prompts( + metadata_registry, monkeypatch +): + prompt_graph = { + "encode_pos": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "expected positive", "clip": ["clip", 0]}, + }, + "encode_other_pos": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "wrong positive", "clip": ["clip", 0]}, + }, + "encode_neg": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "expected negative", "clip": ["clip", 0]}, + }, + "encode_other_neg": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "wrong negative", "clip": ["clip", 0]}, + }, + "enhancer": { + "class_type": "KreaSeedVarianceEnhancer", + "inputs": {"conditioning": ["encode_pos", 0]}, + }, + "zero_out": { + "class_type": "ConditioningZeroOut", + "inputs": {"conditioning": ["encode_neg", 0]}, + }, + "positive_switch": { + "class_type": "ComfySwitchNode", + "inputs": { + "switch": True, + "on_false": ["encode_other_pos", 0], + "on_true": ["enhancer", 0], + }, + }, + "negative_switch": { + "class_type": "ComfySwitchNode", + "inputs": { + "switch": True, + "on_false": ["encode_other_neg", 0], + "on_true": ["zero_out", 0], + }, + }, + "sampler": { + "class_type": "ClownsharKSampler_Beta", + "inputs": { + "seed": 123, + "steps": 8, + "cfg": 1.0, + "sampler_name": "linear/euler", + "scheduler": "beta57", + "denoise": 1.0, + "positive": ["positive_switch", 0], + "negative": ["negative_switch", 0], + "latent_image": { + "samples": types.SimpleNamespace(shape=(1, 4, 16, 16)) + }, + }, + }, + } + prompt = SimpleNamespace(original_prompt=prompt_graph) + + positive_conditioning = object() + other_positive_conditioning = object() + negative_conditioning = object() + other_negative_conditioning = object() + enhanced_conditioning = object() + zeroed_conditioning = object() + + monkeypatch.setattr(metadata_processor, "standalone_mode", False) + + metadata_registry.start_collection("prompt-transformed-switch") + metadata_registry.set_current_prompt(prompt) + + for node_id, text, conditioning in ( + ("encode_pos", "expected positive", positive_conditioning), + ("encode_other_pos", "wrong positive", other_positive_conditioning), + ("encode_neg", "expected negative", negative_conditioning), + ("encode_other_neg", "wrong negative", other_negative_conditioning), + ): + metadata_registry.record_node_execution( + node_id, "CLIPTextEncode", {"text": text}, None + ) + metadata_registry.update_node_execution( + node_id, "CLIPTextEncode", [(conditioning,)] + ) + + metadata_registry.record_node_execution( + "enhancer", + "KreaSeedVarianceEnhancer", + {"conditioning": positive_conditioning}, + None, + return_types=("CONDITIONING", "STRING"), + ) + metadata_registry.update_node_execution( + "enhancer", + "KreaSeedVarianceEnhancer", + [(enhanced_conditioning, "diagnostics")], + return_types=("CONDITIONING", "STRING"), + ) + metadata_registry.record_node_execution( + "zero_out", + "ConditioningZeroOut", + {"conditioning": negative_conditioning}, + None, + return_types=("CONDITIONING",), + ) + metadata_registry.update_node_execution( + "zero_out", + "ConditioningZeroOut", + [(zeroed_conditioning,)], + return_types=("CONDITIONING",), + ) + metadata_registry.record_node_execution( + "positive_switch", + "ComfySwitchNode", + { + "switch": True, + "on_false": other_positive_conditioning, + "on_true": enhanced_conditioning, + }, + None, + ) + metadata_registry.update_node_execution( + "positive_switch", "ComfySwitchNode", [(enhanced_conditioning,)] + ) + metadata_registry.record_node_execution( + "negative_switch", + "ComfySwitchNode", + { + "switch": True, + "on_false": other_negative_conditioning, + "on_true": zeroed_conditioning, + }, + None, + ) + metadata_registry.update_node_execution( + "negative_switch", "ComfySwitchNode", [(zeroed_conditioning,)] + ) + metadata_registry.record_node_execution( + "sampler", + "ClownsharKSampler_Beta", + { + "seed": 123, + "steps": 8, + "cfg": 1.0, + "sampler_name": "linear/euler", + "scheduler": "beta57", + "denoise": 1.0, + "positive": enhanced_conditioning, + "negative": zeroed_conditioning, + "latent_image": { + "samples": types.SimpleNamespace(shape=(1, 4, 16, 16)) + }, + }, + None, + ) + + metadata = metadata_registry.get_metadata("prompt-transformed-switch") + params = MetadataProcessor.extract_generation_params(metadata) + + assert params["prompt"] == "expected positive" + assert params["negative_prompt"] == "expected negative" + + def test_conditioning_provenance_recovers_kj_set_get_prompts( metadata_registry, monkeypatch ):