diff --git a/py/metadata_collector/metadata_hook.py b/py/metadata_collector/metadata_hook.py index efbc22e6..8ad9a588 100644 --- a/py/metadata_collector/metadata_hook.py +++ b/py/metadata_collector/metadata_hook.py @@ -135,10 +135,13 @@ class MetadataHook: # Store the dynprompt reference for node lookups if hasattr(prompt, 'original_prompt'): registry.set_current_prompt(prompt) - + + # Store extra_data for accessing full workflow node properties + registry.set_extra_data(extra_data) + # Execute the original function return original_execute(*args, **kwargs) - + # Replace the functions execution._map_node_over_list = map_node_over_list_with_metadata execution.execute = execute_with_prompt_tracking @@ -202,6 +205,9 @@ class MetadataHook: if hasattr(prompt, 'original_prompt'): registry.set_current_prompt(prompt) + # Store extra_data for accessing full workflow node properties + registry.set_extra_data(extra_data) + # Execute the original function return await original_execute(*args, **kwargs) diff --git a/py/metadata_collector/metadata_processor.py b/py/metadata_collector/metadata_processor.py index 67984297..0d883d71 100644 --- a/py/metadata_collector/metadata_processor.py +++ b/py/metadata_collector/metadata_processor.py @@ -1,4 +1,5 @@ import json +import logging import os from .constants import IMAGES @@ -6,10 +7,62 @@ from .constants import IMAGES standalone_mode = os.environ.get("LORA_MANAGER_STANDALONE", "0") == "1" or os.environ.get("HF_HUB_DISABLE_TELEMETRY", "0") == "0" from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IS_SAMPLER +from .node_extractors import NODE_EXTRACTORS + +logger = logging.getLogger(__name__) + +# Keys that identify metadata hint marks stored in node.properties.lm_marker_role +_META_MARK_PREFIX = "meta_" +_MARK_PRIMARY_MODEL = "primary_model" +_MARK_PRIMARY_SAMPLER = "primary_sampler" +_MARK_POSITIVE_PROMPT = "positive_prompt" +_MARK_NEGATIVE_PROMPT = "negative_prompt" class MetadataProcessor: """Process and format collected metadata""" - + + @staticmethod + def _get_user_marks(metadata): + """Scan workflow nodes (from extra_data.extra_pnginfo.workflow) for user-assigned + metadata hint marks stored in node.properties.lm_marker_role. + + Returns a dict mapping mark type keys to node IDs. + Example: {'primary_model': '42', 'primary_sampler': '17'} + """ + marks: dict[str, str] = {} + + # Primary source: extra_data.extra_pnginfo.workflow.nodes (has full properties) + extra_data = metadata.get("extra_data") + if extra_data and isinstance(extra_data, dict): + extra_pnginfo = extra_data.get("extra_pnginfo", {}) + if isinstance(extra_pnginfo, dict): + workflow = extra_pnginfo.get("workflow", {}) + nodes = workflow.get("nodes", []) + for node in nodes: + node_id = str(node.get("id", "")) + role = node.get("properties", {}).get("lm_marker_role", "") + if role.startswith(_META_MARK_PREFIX): + mark_type = role[len(_META_MARK_PREFIX):] + if mark_type in marks: + logger.warning( + "Duplicate meta hint '%s': node %s (previous: %s), " + "last match wins", + mark_type, node_id, marks[mark_type], + ) + marks[mark_type] = node_id + + # Fallback: try prompt.original_prompt (API-only submissions may not have workflow) + if not marks: + prompt = metadata.get("current_prompt") + if prompt and getattr(prompt, "original_prompt", None): + for node_id, node_data in prompt.original_prompt.items(): + role = node_data.get("properties", {}).get("lm_marker_role", "") + if role.startswith(_META_MARK_PREFIX): + mark_type = role[len(_META_MARK_PREFIX):] + marks[mark_type] = node_id + + return marks + @staticmethod def find_primary_sampler(metadata, downstream_id=None): """ @@ -476,15 +529,51 @@ class MetadataProcessor: # 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, id) - - # Directly get checkpoint from metadata instead of tracing - # Pass primary_sampler_id to avoid redundant calculation - checkpoint = MetadataProcessor.find_primary_checkpoint(metadata, id, primary_sampler_id) - if checkpoint: - params["checkpoint"] = checkpoint + + # ---- User marks: override heuristic inference with user-assigned hints ---- + user_marks = MetadataProcessor._get_user_marks(metadata) + + # Find the primary KSampler node (user mark takes priority) + primary_sampler_id = None + primary_sampler = None + if _MARK_PRIMARY_SAMPLER in user_marks: + marked_id = user_marks[_MARK_PRIMARY_SAMPLER] + sampler_data = metadata.get(SAMPLING, {}).get(marked_id) + if sampler_data and sampler_data.get(IS_SAMPLER): + primary_sampler_id = marked_id + primary_sampler = sampler_data + else: + logger.warning( + "User-marked primary sampler %s has no runtime metadata, " + "falling back to heuristic", + marked_id, + ) + if primary_sampler is None: + primary_sampler_id, primary_sampler = MetadataProcessor.find_primary_sampler(metadata, id) + + # Resolve checkpoint / model (user mark takes priority) + if _MARK_PRIMARY_MODEL in user_marks: + marked_id = user_marks[_MARK_PRIMARY_MODEL] + if marked_id in metadata.get(MODELS, {}): + params["checkpoint"] = metadata[MODELS][marked_id].get("name") + else: + extra_data = metadata.get("extra_data") + extra_pnginfo = extra_data.get("extra_pnginfo", {}) if extra_data and isinstance(extra_data, dict) else {} + workflow = extra_pnginfo.get("workflow", {}) if isinstance(extra_pnginfo, dict) else {} + node_type = "unknown" + for n in workflow.get("nodes", []): + if str(n.get("id", "")) == marked_id: + node_type = n.get("type", "unknown") + break + logger.warning( + "User-marked primary model %s (type=%s, registered=%s) has no runtime metadata, " + "falling back to heuristic", + marked_id, node_type, node_type in NODE_EXTRACTORS, + ) + if params["checkpoint"] is None: + checkpoint = MetadataProcessor.find_primary_checkpoint(metadata, id, primary_sampler_id) + if checkpoint: + params["checkpoint"] = checkpoint # Check if guidance parameter exists in any sampling node for node_id, sampler_info in metadata.get(SAMPLING, {}).items(): @@ -539,7 +628,22 @@ class MetadataProcessor: # For SamplerCustom, handle any additional parameters MetadataProcessor.handle_custom_advanced_sampler(metadata, prompt, primary_sampler_id, params) - + + # ---- User marks: override prompts with explicitly tagged nodes ---- + prompts_data = metadata.get(PROMPTS, {}) + if _MARK_POSITIVE_PROMPT in user_marks: + pos_id = user_marks[_MARK_POSITIVE_PROMPT] + if pos_id in prompts_data: + prompt_text = prompts_data[pos_id].get("text") or prompts_data[pos_id].get("positive_text") + if prompt_text: + params["prompt"] = prompt_text + if _MARK_NEGATIVE_PROMPT in user_marks: + neg_id = user_marks[_MARK_NEGATIVE_PROMPT] + if neg_id in prompts_data: + prompt_text = prompts_data[neg_id].get("text") or prompts_data[neg_id].get("negative_text") + if prompt_text: + params["negative_prompt"] = prompt_text + # Size extraction is same for all sampler types # Check if the sampler itself has size information (from latent_image) if primary_sampler_id in metadata.get(SIZE, {}): diff --git a/py/metadata_collector/metadata_registry.py b/py/metadata_collector/metadata_registry.py index 2a084325..734a18b8 100644 --- a/py/metadata_collector/metadata_registry.py +++ b/py/metadata_collector/metadata_registry.py @@ -61,6 +61,7 @@ class MetadataRegistry: { "execution_order": [], "current_prompt": None, # Will store the prompt object + "extra_data": None, # Will store the API extra_data for workflow metadata "timestamp": time.time(), } ) @@ -75,6 +76,11 @@ class MetadataRegistry: # Store the prompt in the metadata for later relationship tracing self.prompt_metadata[self.current_prompt_id]["current_prompt"] = prompt + def set_extra_data(self, extra_data): + """Store the API extra_data (contains extra_pnginfo.workflow with node properties)""" + if self.current_prompt_id and self.current_prompt_id in self.prompt_metadata: + self.prompt_metadata[self.current_prompt_id]["extra_data"] = extra_data + def get_metadata(self, prompt_id=None): """Get collected metadata for a prompt""" key = prompt_id if prompt_id is not None else self.current_prompt_id diff --git a/web/comfyui/node_marker.js b/web/comfyui/node_marker.js index 2499bf37..0cdff7b5 100644 --- a/web/comfyui/node_marker.js +++ b/web/comfyui/node_marker.js @@ -7,12 +7,16 @@ import { app } from "../../scripts/app.js"; // Roles are stored in ``node.properties.lm_marker_role`` and automatically // persist with the workflow JSON. // +// Two categories: +// send_* – consumed by the standalone UI's "Send to Workflow" feature +// meta_* – consumed by the metadata processor to override heuristic inference +// // The workflow registry reads these markers and makes them available to the // standalone UI (e.g. ``sendEmbeddingToWorkflow`` also considers nodes marked // as ``send_prompt_target``). // ============================================================================= -const ROLES = { +const SEND_ROLES = { send_prompt_target: { label: "Send Prompt Target", emoji: "\uD83D\uDCDD", @@ -23,6 +27,28 @@ const ROLES = { }, }; +const META_ROLES = { + meta_primary_model: { + label: "Meta hints: Primary Model", + emoji: "\uD83D\uDCA1", + }, + meta_primary_sampler: { + label: "Meta hints: Primary Sampler", + emoji: "\uD83D\uDCA1", + }, + meta_positive_prompt: { + label: "Meta hints: Positive Prompt", + emoji: "\uD83D\uDCA1", + }, + meta_negative_prompt: { + label: "Meta hints: Negative Prompt", + emoji: "\uD83D\uDCA1", + }, +}; + +// Flat lookup for setMarker / getMarker / clearMarker +const ROLES = { ...SEND_ROLES, ...META_ROLES }; + // ---- Helpers ---------------------------------------------------------------- function getMarker(node) { @@ -54,7 +80,7 @@ function clearMarker(node) { // Restore original title: prefer stripping emoji from current title // (captures user renames after marking), fall back to saved original. const cleaned = node.title?.replace( - /^(\u2709\uFE0F?|\u2699\uFE0F?|\uD83D\uDCDD|\uD83C\uDF9B\uFE0F?|\uD83D\uDD27)\s*/, + /^(\u2709\uFE0F?|\u2699\uFE0F?|\uD83D\uDCDD|\uD83C\uDF9B\uFE0F?|\uD83D\uDD27|\uD83D\uDCA1)\s*/, '' ); if (cleaned && cleaned !== node.title) { @@ -84,16 +110,23 @@ function buildSubmenuOptions(node) { const currentRole = getMarker(node); const options = []; - for (const [key, def] of Object.entries(ROLES)) { - const isActive = currentRole === key; - options.push({ - content: `${isActive ? "\u2713 " : ""}${def.label}`, - disabled: isActive, - callback: () => setMarker(node, key), - }); - } + const buildGroup = (roles) => { + for (const [key, def] of Object.entries(roles)) { + const isActive = currentRole === key; + options.push({ + content: `${isActive ? "\u2713 " : ""}${def.label}`, + disabled: isActive, + callback: () => setMarker(node, key), + }); + } + }; + + buildGroup(SEND_ROLES); + options.push(null); // separator + buildGroup(META_ROLES); if (currentRole) { + options.push(null); // separator options.push({ content: "Clear marker", callback: () => clearMarker(node),