From 0edbd7bcca03f2783f0195bd20f74fbde80b2a06 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Fri, 19 Jun 2026 17:13:48 +0800 Subject: [PATCH] fix(metadata): add LoraTextLoaderLM extractor so SaveImageLM records its loras (#801) --- py/metadata_collector/node_extractors.py | 50 ++++++++++++++++ .../test_metadata_collector.py | 59 +++++++++++++++++++ 2 files changed, 109 insertions(+) diff --git a/py/metadata_collector/node_extractors.py b/py/metadata_collector/node_extractors.py index 3ce433e2..352e5cf4 100644 --- a/py/metadata_collector/node_extractors.py +++ b/py/metadata_collector/node_extractors.py @@ -901,6 +901,55 @@ class LoraLoaderManagerExtractor(NodeMetadataExtractor): "node_id": node_id } +class LoraTextLoaderManagerExtractor(NodeMetadataExtractor): + """Extract LoRA metadata from LoraTextLoaderLM (LoRA Text Loader). + + The node accepts a `lora_syntax` STRING containing tags + (same format as the ComfyUI prompt), plus an optional `lora_stack`. + This extractor parses the syntax string using the same regex as the node. + """ + @staticmethod + def extract(node_id, inputs, outputs, metadata): + if not inputs: + return + + active_loras = [] + + # Process lora_stack if available (optional input) + if "lora_stack" in inputs: + lora_stack = inputs.get("lora_stack", []) + for item in lora_stack: + # lora_stack entries are (path, model_strength, clip_strength) tuples + if isinstance(item, (list, tuple)) and len(item) >= 2: + lora_path = item[0] + model_strength = item[1] + lora_name = os.path.splitext(os.path.basename(lora_path))[0] + active_loras.append({ + "name": lora_name, + "strength": round(float(model_strength), 2) + }) + + # Process lora_syntax string input + if "lora_syntax" in inputs: + lora_syntax = inputs.get("lora_syntax", "") + if lora_syntax and isinstance(lora_syntax, str): + pattern = r"]+):([^:>]+)(?::([^:>]+))?>" + matches = re.findall(pattern, lora_syntax, re.IGNORECASE) + for match in matches: + lora_name = match[0] + model_strength = float(match[1]) + active_loras.append({ + "name": lora_name, + "strength": round(model_strength, 2) + }) + + if active_loras: + metadata[LORAS][node_id] = { + "lora_list": active_loras, + "node_id": node_id + } + + class FluxGuidanceExtractor(NodeMetadataExtractor): @staticmethod def extract(node_id, inputs, outputs, metadata): @@ -1146,6 +1195,7 @@ NODE_EXTRACTORS = { "UNETLoaderLM": UNETLoaderExtractor, # LoRA Manager "LoraLoader": LoraLoaderExtractor, "LoraLoaderLM": LoraLoaderManagerExtractor, + "LoraTextLoaderLM": LoraTextLoaderManagerExtractor, "RgthreePowerLoraLoader": RgthreePowerLoraLoaderExtractor, "TensorRTLoader": TensorRTLoaderExtractor, # Conditioning diff --git a/tests/metadata_collector/test_metadata_collector.py b/tests/metadata_collector/test_metadata_collector.py index 9258f8c9..c9b08bb8 100644 --- a/tests/metadata_collector/test_metadata_collector.py +++ b/tests/metadata_collector/test_metadata_collector.py @@ -733,6 +733,65 @@ def test_lora_manager_cache_updates_when_loras_removed(metadata_registry): assert "lora_node" not in metadata[LORAS] +def test_lora_text_loader_extracts_loras_from_syntax(metadata_registry): + """LoraTextLoaderLM extractor parses tags from lora_syntax string.""" + metadata_registry.start_collection("prompt1") + + metadata_registry.record_node_execution( + "text_loader", + "LoraTextLoaderLM", + {"lora_syntax": [" "]}, + None, + ) + + metadata = metadata_registry.get_metadata("prompt1") + + assert "text_loader" in metadata[LORAS] + lora_list = metadata[LORAS]["text_loader"]["lora_list"] + assert len(lora_list) == 2 + assert lora_list[0] == {"name": "foo", "strength": 0.8} + assert lora_list[1] == {"name": "bar", "strength": 1.0} + + +def test_lora_text_loader_extracts_loras_from_lora_stack(metadata_registry): + """LoraTextLoaderLM extractor also processes the optional lora_stack input.""" + metadata_registry.start_collection("prompt1") + + metadata_registry.record_node_execution( + "stack_loader", + "LoraTextLoaderLM", + { + "lora_syntax": [""], + "lora_stack": (("/models/loras/my-lora.safetensors", 0.6, 0.5),), + }, + None, + ) + + metadata = metadata_registry.get_metadata("prompt1") + + assert "stack_loader" in metadata[LORAS] + lora_list = metadata[LORAS]["stack_loader"]["lora_list"] + assert len(lora_list) == 1 + assert lora_list[0] == {"name": "my-lora", "strength": 0.6} + + +def test_lora_text_loader_handles_empty_syntax(metadata_registry): + """LoraTextLoaderLM extractor produces no metadata when no loras are provided.""" + metadata_registry.start_collection("prompt1") + + metadata_registry.record_node_execution( + "empty_loader", + "LoraTextLoaderLM", + {"lora_syntax": [""]}, + None, + ) + + metadata = metadata_registry.get_metadata("prompt1") + + assert "empty_loader" not in metadata[LORAS] + + + def test_lora_manager_checkpoint_and_unet_loaders_extract_models(metadata_registry): metadata_registry.start_collection("prompt1")