From 849b97afba162de985dd3bb44dd0fe8264d8c0ef Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Thu, 10 Jul 2025 09:26:53 +0800 Subject: [PATCH] feat: Add CR_ApplyControlNetStack extractor and enhance prompt conditioning handling in metadata processing. Fixes #277 --- py/metadata_collector/metadata_processor.py | 54 ++++++++++++++------- py/metadata_collector/node_extractors.py | 35 +++++++++++++ 2 files changed, 72 insertions(+), 17 deletions(-) diff --git a/py/metadata_collector/metadata_processor.py b/py/metadata_collector/metadata_processor.py index 17d0a482..add9dbfc 100644 --- a/py/metadata_collector/metadata_processor.py +++ b/py/metadata_collector/metadata_processor.py @@ -238,25 +238,45 @@ class MetadataProcessor: pos_conditioning = metadata[PROMPTS][sampler_id].get("pos_conditioning") neg_conditioning = metadata[PROMPTS][sampler_id].get("neg_conditioning") - - # Try to match conditioning objects with those stored by CLIPTextEncodeExtractor - for prompt_node_id, prompt_data in metadata[PROMPTS].items(): - # For nodes with single conditioning output - if "conditioning" in prompt_data: - if pos_conditioning is not None and id(prompt_data["conditioning"]) == id(pos_conditioning): - result["prompt"] = prompt_data.get("text", "") + + # Helper function to recursively find prompt text for a conditioning object + def find_prompt_text_for_conditioning(conditioning_obj, is_positive=True): + if conditioning_obj is None: + return "" - if neg_conditioning is not None and id(prompt_data["conditioning"]) == id(neg_conditioning): - result["negative_prompt"] = prompt_data.get("text", "") + # Try to match conditioning objects with those stored by extractors + for prompt_node_id, prompt_data in metadata[PROMPTS].items(): + # For nodes with single conditioning output + if "conditioning" in prompt_data: + if id(prompt_data["conditioning"]) == id(conditioning_obj): + return prompt_data.get("text", "") + + # For nodes with separate pos_conditioning and neg_conditioning outputs (like TSC_EfficientLoader) + if is_positive and "positive_encoded" in prompt_data: + if id(prompt_data["positive_encoded"]) == id(conditioning_obj): + if "positive_text" in prompt_data: + return prompt_data["positive_text"] + else: + orig_conditioning = prompt_data.get("orig_pos_cond", None) + if orig_conditioning is not None: + # Recursively find the prompt text for the original conditioning + return find_prompt_text_for_conditioning(orig_conditioning, is_positive=True) + + if not is_positive and "negative_encoded" in prompt_data: + if id(prompt_data["negative_encoded"]) == id(conditioning_obj): + if "negative_text" in prompt_data: + return prompt_data["negative_text"] + else: + orig_conditioning = prompt_data.get("orig_neg_cond", None) + if orig_conditioning is not None: + # Recursively find the prompt text for the original conditioning + return find_prompt_text_for_conditioning(orig_conditioning, is_positive=False) - # For nodes with separate pos_conditioning and neg_conditioning outputs (like TSC_EfficientLoader) - if "positive_encoded" in prompt_data: - if pos_conditioning is not None and id(prompt_data["positive_encoded"]) == id(pos_conditioning): - result["prompt"] = prompt_data.get("positive_text", "") - - if "negative_encoded" in prompt_data: - if neg_conditioning is not None and id(prompt_data["negative_encoded"]) == id(neg_conditioning): - result["negative_prompt"] = prompt_data.get("negative_text", "") + return "" + + # Find prompt texts using the helper function + result["prompt"] = find_prompt_text_for_conditioning(pos_conditioning, is_positive=True) + result["negative_prompt"] = find_prompt_text_for_conditioning(neg_conditioning, is_positive=False) return result diff --git a/py/metadata_collector/node_extractors.py b/py/metadata_collector/node_extractors.py index 7bdbfc38..4e3a79d2 100644 --- a/py/metadata_collector/node_extractors.py +++ b/py/metadata_collector/node_extractors.py @@ -569,6 +569,40 @@ class CFGGuiderExtractor(NodeMetadataExtractor): metadata[SAMPLING][node_id]["parameters"]["cfg"] = cfg_value +class CR_ApplyControlNetStackExtractor(NodeMetadataExtractor): + @staticmethod + def extract(node_id, inputs, outputs, metadata): + if not inputs: + return + + # Save the original conditioning inputs + base_positive = inputs.get("base_positive") + base_negative = inputs.get("base_negative") + + if base_positive is not None or base_negative is not None: + if node_id not in metadata[PROMPTS]: + metadata[PROMPTS][node_id] = {"node_id": node_id} + + metadata[PROMPTS][node_id]["orig_pos_cond"] = base_positive + metadata[PROMPTS][node_id]["orig_neg_cond"] = base_negative + + @staticmethod + def update(node_id, outputs, metadata): + # Extract transformed conditionings from outputs + # outputs structure: [(base_positive, base_negative, show_help, )] + if outputs and isinstance(outputs, list) and len(outputs) > 0: + first_output = outputs[0] + if isinstance(first_output, tuple) and len(first_output) >= 2: + transformed_positive = first_output[0] + transformed_negative = first_output[1] + + # Save transformed conditioning objects in metadata + if node_id not in metadata[PROMPTS]: + metadata[PROMPTS][node_id] = {"node_id": node_id} + + metadata[PROMPTS][node_id]["positive_encoded"] = transformed_positive + metadata[PROMPTS][node_id]["negative_encoded"] = transformed_negative + # Registry of node-specific extractors # Keys are node class names NODE_EXTRACTORS = { @@ -595,6 +629,7 @@ NODE_EXTRACTORS = { "WAS_Text_to_Conditioning": CLIPTextEncodeExtractor, "AdvancedCLIPTextEncode": CLIPTextEncodeExtractor, # From https://github.com/BlenderNeko/ComfyUI_ADV_CLIP_emb "smZ_CLIPTextEncode": CLIPTextEncodeExtractor, # From https://github.com/shiimizu/ComfyUI_smZNodes + "CR_ApplyControlNetStack": CR_ApplyControlNetStackExtractor, # Add CR_ApplyControlNetStack # Latent "EmptyLatentImage": ImageSizeExtractor, # Flux