refactor: enhance CFGGuider handling and add CFGGuiderExtractor for improved metadata extraction. Fixes https://github.com/willmiao/ComfyUI-Lora-Manager/issues/172

This commit is contained in:
Will Miao
2025-05-09 13:50:22 +08:00
parent fa08c9c3e4
commit 7b5b4ce082
2 changed files with 47 additions and 12 deletions

View File

@@ -187,19 +187,36 @@ class MetadataProcessor:
sampler_params = metadata[SAMPLING][sampler_node_id].get("parameters", {}) sampler_params = metadata[SAMPLING][sampler_node_id].get("parameters", {})
params["sampler"] = sampler_params.get("sampler_name") params["sampler"] = sampler_params.get("sampler_name")
# 3. Trace guider input for FluxGuidance and CLIPTextEncode # 3. Trace guider input for CFGGuider, FluxGuidance and CLIPTextEncode
guider_node_id = MetadataProcessor.trace_node_input(prompt, primary_sampler_id, "guider", max_depth=5) guider_node_id = MetadataProcessor.trace_node_input(prompt, primary_sampler_id, "guider", max_depth=5)
if guider_node_id: if guider_node_id and guider_node_id in prompt.original_prompt:
# Look for FluxGuidance along the guider path # Check if the guider node is a CFGGuider
flux_node_id = MetadataProcessor.trace_node_input(prompt, guider_node_id, "conditioning", "FluxGuidance", max_depth=5) if prompt.original_prompt[guider_node_id].get("class_type") == "CFGGuider":
if flux_node_id and flux_node_id in metadata.get(SAMPLING, {}): # Extract cfg value from the CFGGuider
flux_params = metadata[SAMPLING][flux_node_id].get("parameters", {}) if guider_node_id in metadata.get(SAMPLING, {}):
params["guidance"] = flux_params.get("guidance") cfg_params = metadata[SAMPLING][guider_node_id].get("parameters", {})
params["cfg_scale"] = cfg_params.get("cfg")
# Find CLIPTextEncode for positive prompt (through conditioning) # Find CLIPTextEncode for positive prompt
positive_node_id = MetadataProcessor.trace_node_input(prompt, guider_node_id, "conditioning", "CLIPTextEncode", max_depth=10) positive_node_id = MetadataProcessor.trace_node_input(prompt, guider_node_id, "positive", "CLIPTextEncode", max_depth=10)
if positive_node_id and positive_node_id in metadata.get(PROMPTS, {}): if positive_node_id and positive_node_id in metadata.get(PROMPTS, {}):
params["prompt"] = metadata[PROMPTS][positive_node_id].get("text", "") params["prompt"] = metadata[PROMPTS][positive_node_id].get("text", "")
# Find CLIPTextEncode for negative prompt
negative_node_id = MetadataProcessor.trace_node_input(prompt, guider_node_id, "negative", "CLIPTextEncode", max_depth=10)
if negative_node_id and negative_node_id in metadata.get(PROMPTS, {}):
params["negative_prompt"] = metadata[PROMPTS][negative_node_id].get("text", "")
else:
# Look for FluxGuidance along the guider path
flux_node_id = MetadataProcessor.trace_node_input(prompt, guider_node_id, "conditioning", "FluxGuidance", max_depth=5)
if flux_node_id and flux_node_id in metadata.get(SAMPLING, {}):
flux_params = metadata[SAMPLING][flux_node_id].get("parameters", {})
params["guidance"] = flux_params.get("guidance")
# Find CLIPTextEncode for positive prompt (through conditioning)
positive_node_id = MetadataProcessor.trace_node_input(prompt, guider_node_id, "conditioning", "CLIPTextEncode", max_depth=10)
if positive_node_id and positive_node_id in metadata.get(PROMPTS, {}):
params["prompt"] = metadata[PROMPTS][positive_node_id].get("text", "")
else: else:
# Original tracing for standard samplers # Original tracing for standard samplers

View File

@@ -362,6 +362,23 @@ class CLIPTextEncodeFluxExtractor(NodeMetadataExtractor):
metadata[SAMPLING][node_id]["parameters"]["guidance"] = guidance_value metadata[SAMPLING][node_id]["parameters"]["guidance"] = guidance_value
class CFGGuiderExtractor(NodeMetadataExtractor):
@staticmethod
def extract(node_id, inputs, outputs, metadata):
if not inputs or "cfg" not in inputs:
return
cfg_value = inputs.get("cfg")
# Store the cfg value in SAMPLING category
if SAMPLING not in metadata:
metadata[SAMPLING] = {}
if node_id not in metadata[SAMPLING]:
metadata[SAMPLING][node_id] = {"parameters": {}, "node_id": node_id}
metadata[SAMPLING][node_id]["parameters"]["cfg"] = cfg_value
# Registry of node-specific extractors # Registry of node-specific extractors
NODE_EXTRACTORS = { NODE_EXTRACTORS = {
# Sampling # Sampling
@@ -383,6 +400,7 @@ NODE_EXTRACTORS = {
"EmptyLatentImage": ImageSizeExtractor, "EmptyLatentImage": ImageSizeExtractor,
# Flux # Flux
"FluxGuidance": FluxGuidanceExtractor, # Add FluxGuidance "FluxGuidance": FluxGuidanceExtractor, # Add FluxGuidance
"CFGGuider": CFGGuiderExtractor, # Add CFGGuider
# Image # Image
"VAEDecode": VAEDecodeExtractor, # Added VAEDecode extractor "VAEDecode": VAEDecodeExtractor, # Added VAEDecode extractor
# Add other nodes as needed # Add other nodes as needed