import json from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE class MetadataProcessor: """Process and format collected metadata""" @staticmethod def find_primary_sampler(metadata): """Find the primary KSampler node (with denoise=1)""" primary_sampler = None primary_sampler_id = None # First, check for KSamplerAdvanced with add_noise="enable" for node_id, sampler_info in metadata.get(SAMPLING, {}).items(): parameters = sampler_info.get("parameters", {}) add_noise = parameters.get("add_noise") # If add_noise is "enable", this is likely the primary sampler for KSamplerAdvanced if add_noise == "enable": primary_sampler = sampler_info primary_sampler_id = node_id break # If no KSamplerAdvanced found, fall back to traditional KSampler with denoise=1 if primary_sampler is None: for node_id, sampler_info in metadata.get(SAMPLING, {}).items(): parameters = sampler_info.get("parameters", {}) denoise = parameters.get("denoise") # If denoise is 1.0, this is likely the primary sampler if denoise == 1.0 or denoise == 1: primary_sampler = sampler_info primary_sampler_id = node_id break return primary_sampler_id, primary_sampler @staticmethod def trace_node_input(prompt, node_id, input_name, target_class=None, max_depth=10): """ Trace an input connection from a node to find the source node Parameters: - prompt: The prompt object containing node connections - node_id: ID of the starting node - input_name: Name of the input to trace - target_class: Optional class name to search for (e.g., "CLIPTextEncode") - max_depth: Maximum depth to follow the node chain to prevent infinite loops Returns: - node_id of the found node, or None if not found """ if not prompt or not prompt.original_prompt or node_id not in prompt.original_prompt: return None # For depth tracking current_depth = 0 current_node_id = node_id current_input = input_name while current_depth < max_depth: if current_node_id not in prompt.original_prompt: return None node_inputs = prompt.original_prompt[current_node_id].get("inputs", {}) if current_input not in node_inputs: return None input_value = node_inputs[current_input] # Input connections are formatted as [node_id, output_index] if isinstance(input_value, list) and len(input_value) >= 2: found_node_id = input_value[0] # Connected node_id # If we're looking for a specific node class if target_class and prompt.original_prompt[found_node_id].get("class_type") == target_class: return found_node_id # If we're not looking for a specific class or haven't found it yet if not target_class: return found_node_id # Continue tracing through intermediate nodes current_node_id = found_node_id # For most conditioning nodes, the input we want to follow is named "conditioning" if "conditioning" in prompt.original_prompt[current_node_id].get("inputs", {}): current_input = "conditioning" else: # If there's no "conditioning" input, we can't trace further return found_node_id if not target_class else None else: # We've reached a node with no further connections return None current_depth += 1 # If we've reached max depth without finding target_class return None @staticmethod def find_primary_checkpoint(metadata): """Find the primary checkpoint model in the workflow""" if not metadata.get(MODELS): return None # In most workflows, there's only one checkpoint, so we can just take the first one for node_id, model_info in metadata.get(MODELS, {}).items(): if model_info.get("type") == "checkpoint": return model_info.get("name") return None @staticmethod def extract_generation_params(metadata): """Extract generation parameters from metadata using node relationships""" params = { "prompt": "", "negative_prompt": "", "seed": None, "steps": None, "cfg_scale": None, "guidance": None, # Add guidance parameter "sampler": None, "scheduler": None, "checkpoint": None, "loras": "", "size": None, "clip_skip": None } # Get the prompt object for node relationship tracing prompt = metadata.get("current_prompt") # Find the primary KSampler node primary_sampler_id, primary_sampler = MetadataProcessor.find_primary_sampler(metadata) # Directly get checkpoint from metadata instead of tracing checkpoint = MetadataProcessor.find_primary_checkpoint(metadata) if checkpoint: params["checkpoint"] = checkpoint if primary_sampler: # Extract sampling parameters sampling_params = primary_sampler.get("parameters", {}) # Handle both seed and noise_seed params["seed"] = sampling_params.get("seed") if sampling_params.get("seed") is not None else sampling_params.get("noise_seed") params["steps"] = sampling_params.get("steps") params["cfg_scale"] = sampling_params.get("cfg") params["sampler"] = sampling_params.get("sampler_name") params["scheduler"] = sampling_params.get("scheduler") # Trace connections from the primary sampler if prompt and primary_sampler_id: # Trace positive prompt - look specifically for CLIPTextEncode positive_node_id = MetadataProcessor.trace_node_input(prompt, primary_sampler_id, "positive", "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", "") # Find any FluxGuidance nodes in the positive conditioning path flux_node_id = MetadataProcessor.trace_node_input(prompt, primary_sampler_id, "positive", "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") # Trace negative prompt - look specifically for CLIPTextEncode negative_node_id = MetadataProcessor.trace_node_input(prompt, primary_sampler_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", "") # Check if the sampler itself has size information (from latent_image) if primary_sampler_id in metadata.get(SIZE, {}): width = metadata[SIZE][primary_sampler_id].get("width") height = metadata[SIZE][primary_sampler_id].get("height") if width and height: params["size"] = f"{width}x{height}" else: # Fallback to the previous trace method if needed latent_node_id = MetadataProcessor.trace_node_input(prompt, primary_sampler_id, "latent_image") if latent_node_id: # Follow chain to find EmptyLatentImage node size_found = False current_node_id = latent_node_id # Limit depth to avoid infinite loops in complex workflows max_depth = 10 for _ in range(max_depth): if current_node_id in metadata.get(SIZE, {}): width = metadata[SIZE][current_node_id].get("width") height = metadata[SIZE][current_node_id].get("height") if width and height: params["size"] = f"{width}x{height}" size_found = True break # Try to follow the chain if prompt and prompt.original_prompt and current_node_id in prompt.original_prompt: node_info = prompt.original_prompt[current_node_id] if "inputs" in node_info: # Look for a connection that might lead to size information for input_name, input_value in node_info["inputs"].items(): if isinstance(input_value, list) and len(input_value) >= 2: current_node_id = input_value[0] break else: break # No connections to follow else: break # No inputs to follow else: break # Can't follow further # Extract LoRAs using the standardized format lora_parts = [] for node_id, lora_info in metadata.get(LORAS, {}).items(): # Access the lora_list from the standardized format lora_list = lora_info.get("lora_list", []) for lora in lora_list: name = lora.get("name", "unknown") strength = lora.get("strength", 1.0) lora_parts.append(f"") params["loras"] = " ".join(lora_parts) # Set default clip_skip value params["clip_skip"] = "1" # Common default return params @staticmethod def to_dict(metadata): """Convert extracted metadata to the ComfyUI output.json format""" params = MetadataProcessor.extract_generation_params(metadata) # Convert all values to strings to match output.json format for key in params: if params[key] is not None: params[key] = str(params[key]) return params @staticmethod def to_json(metadata): """Convert metadata to JSON string""" params = MetadataProcessor.to_dict(metadata) return json.dumps(params, indent=4)