From 73cefa409d7da614f7930a5f02cfc7b327a86d96 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Tue, 8 Sep 2026 08:39:03 +0200 Subject: [PATCH] lora(resize) add exclusion and discard filters to resize nodes --- nodes/lora_resize.py | 73 +++++++++++++++++++++--- tests/test_lora_resize_factor_space.py | 79 ++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 8 deletions(-) diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index 1655ff2..0af7145 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -6,6 +6,8 @@ dynamic methods (sv_ratio, sv_fro, sv_cumulative). """ import os import gc +import fnmatch +import re import logging from collections import Counter, defaultdict import torch @@ -641,13 +643,32 @@ def _compute_resize( return new_rank, new_alpha, stats -def _build_resize_work_units(handler, pairs, passthrough_keys, low_bit_keys): +def _resize_filter_inputs(): + return [ + io.String.Input("exclude_patterns", default="", multiline=True, tooltip="One pattern per line. Matching layer names or source/canonical tensor keys keep the whole layer at its original rank and precision, including alpha."), + io.String.Input("discard_patterns", default="", multiline=True, tooltip="One pattern per line. Remove the whole matching layer, including factors and alpha. Discard takes precedence over exclude."), + io.Boolean.Input("glob_patterns", default=False, tooltip="Use shell-style glob patterns instead of regular expressions for exclude and discard filters."), + ] + + +def _build_resize_work_units( + handler, pairs, passthrough_keys, low_bit_keys, + exclude_patterns="", discard_patterns="", glob_patterns=False, +): """Plan raw-copy and streamed work without loading tensor payloads.""" raw_units = [] stream_units = [] claimed = set() rank_counts = Counter() preserved_companion_groups = 0 + excludes, discards = [ + [re.compile(fnmatch.translate(line.strip()) if glob_patterns else line.strip()) + for line in text.splitlines() if line.strip()] + for text in (exclude_patterns, discard_patterns) + ] + + def matches(patterns, names): + return any(pattern.search(name) for pattern in patterns for name in names) for block_name, block_keys in pairs.items(): tensor_keys = layer_tensor_keys(block_keys) @@ -664,13 +685,24 @@ def _build_resize_work_units(handler, pairs, passthrough_keys, low_bit_keys): if key.startswith(f"{block_name}.") and key not in claimed ] + group_entries = entries + related_passthrough + names = [block_name, canonical_lora_block_name(block_name)] + names.extend(name for _, key, output_key in group_entries for name in (key, output_key)) + if matches(discards, names): + claimed.update(key for _, key, _ in group_entries) + continue + excluded = matches(excludes, names) + if layer_has_low_bit(block_keys, low_bit_keys): raw_entries = entries + related_passthrough raw_units.append({"entries": raw_entries}) claimed.update(key for _, key, _ in raw_entries) continue - if layer_has_companions(block_keys): + if excluded: + stream_units.append({"kind": "preserve", "entries": group_entries}) + claimed.update(key for _, key, _ in group_entries) + elif layer_has_companions(block_keys): preserved_companion_groups += 1 stream_units.append({"kind": "preserve", "entries": entries}) elif "down" in block_keys and "up" in block_keys: @@ -690,11 +722,16 @@ def _build_resize_work_units(handler, pairs, passthrough_keys, low_bit_keys): for key in passthrough_keys: if key in claimed: continue + if matches(discards, [key]): + continue entry = (None, key, key) if key in low_bit_keys: raw_units.append({"entries": [entry]}) else: - stream_units.append({"kind": "copy", "entries": [entry]}) + stream_units.append({ + "kind": "preserve" if matches(excludes, [key]) else "copy", + "entries": [entry], + }) claimed.add(key) return raw_units, stream_units, rank_counts, preserved_companion_groups @@ -862,6 +899,9 @@ def resize_lora_file( verbose: bool = True, force_clear_cache: bool = False, min_rank: int = 1, + exclude_patterns: str = "", + discard_patterns: str = "", + glob_patterns: bool = False, ) -> str: """Resize a LoRA through one bounded UEL stream and factor-space SVD.""" if min_rank < 1: @@ -884,7 +924,8 @@ def resize_lora_file( validate_canonical_blocks(pairs, "LoRA Resize") raw_units, stream_units, rank_counts, preserved_companion_groups = ( _build_resize_work_units( - handler, pairs, passthrough_keys, low_bit_keys + handler, pairs, passthrough_keys, low_bit_keys, + exclude_patterns, discard_patterns, glob_patterns, ) ) rank_text = _rank_summary(rank_counts) @@ -1182,19 +1223,23 @@ class LoRAResizeFixed(io.ComfyNode): io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16", tooltip="Data type used to save resized LoRA tensors."), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer resize arithmetic; CUDA out-of-memory retries the affected layer on CPU."), io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"), + *_resize_filter_inputs(), ], outputs=[io.AnyType.Output(display_name="output_path")], is_output_node=True, ) @classmethod - def execute(cls, lora_name, new_rank, output_filename, save_dtype, device, force_clear_cache) -> io.NodeOutput: + def execute(cls, lora_name, new_rank, output_filename, save_dtype, device, force_clear_cache, exclude_patterns="", discard_patterns="", glob_patterns=False) -> io.NodeOutput: lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype] resize_lora_file( lora_path, new_rank, None, None, device, dtype, output_filename, force_clear_cache=force_clear_cache, + exclude_patterns=exclude_patterns, + discard_patterns=discard_patterns, + glob_patterns=glob_patterns, ) _, output_name = canonical_model_artifact_path("loras", output_filename) return io.NodeOutput(output_name) @@ -1221,19 +1266,23 @@ class LoRAResizeRatio(io.ComfyNode): io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16", tooltip="Data type used to save resized LoRA tensors."), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer resize arithmetic; CUDA out-of-memory retries the affected layer on CPU."), io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"), + *_resize_filter_inputs(), ], outputs=[io.AnyType.Output(display_name="output_path")], is_output_node=True, ) @classmethod - def execute(cls, lora_name, max_rank, ratio, output_filename, save_dtype, device, force_clear_cache) -> io.NodeOutput: + def execute(cls, lora_name, max_rank, ratio, output_filename, save_dtype, device, force_clear_cache, exclude_patterns="", discard_patterns="", glob_patterns=False) -> io.NodeOutput: lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype] resize_lora_file( lora_path, max_rank, "sv_ratio", ratio, device, dtype, output_filename, force_clear_cache=force_clear_cache, + exclude_patterns=exclude_patterns, + discard_patterns=discard_patterns, + glob_patterns=glob_patterns, ) _, output_name = canonical_model_artifact_path("loras", output_filename) return io.NodeOutput(output_name) @@ -1262,19 +1311,23 @@ class LoRAResizeFrobenius(io.ComfyNode): io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16", tooltip="Data type used to save resized LoRA tensors."), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer resize arithmetic; CUDA out-of-memory retries the affected layer on CPU."), io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"), + *_resize_filter_inputs(), ], outputs=[io.AnyType.Output(display_name="output_path")], is_output_node=True, ) @classmethod - def execute(cls, lora_name, max_rank, min_rank, target, output_filename, save_dtype, device, force_clear_cache) -> io.NodeOutput: + def execute(cls, lora_name, max_rank, min_rank, target, output_filename, save_dtype, device, force_clear_cache, exclude_patterns="", discard_patterns="", glob_patterns=False) -> io.NodeOutput: lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype] resize_lora_file( lora_path, max_rank, "sv_fro", target, device, dtype, output_filename, force_clear_cache=force_clear_cache, + exclude_patterns=exclude_patterns, + discard_patterns=discard_patterns, + glob_patterns=glob_patterns, min_rank=min_rank, ) _, output_name = canonical_model_artifact_path("loras", output_filename) @@ -1302,19 +1355,23 @@ class LoRAResizeCumulative(io.ComfyNode): io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16", tooltip="Data type used to save resized LoRA tensors."), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer resize arithmetic; CUDA out-of-memory retries the affected layer on CPU."), io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"), + *_resize_filter_inputs(), ], outputs=[io.AnyType.Output(display_name="output_path")], is_output_node=True, ) @classmethod - def execute(cls, lora_name, max_rank, target, output_filename, save_dtype, device, force_clear_cache) -> io.NodeOutput: + def execute(cls, lora_name, max_rank, target, output_filename, save_dtype, device, force_clear_cache, exclude_patterns="", discard_patterns="", glob_patterns=False) -> io.NodeOutput: lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype] resize_lora_file( lora_path, max_rank, "sv_cumulative", target, device, dtype, output_filename, force_clear_cache=force_clear_cache, + exclude_patterns=exclude_patterns, + discard_patterns=discard_patterns, + glob_patterns=glob_patterns, ) _, output_name = canonical_model_artifact_path("loras", output_filename) return io.NodeOutput(output_name) diff --git a/tests/test_lora_resize_factor_space.py b/tests/test_lora_resize_factor_space.py index 8eea70f..d2875ae 100644 --- a/tests/test_lora_resize_factor_space.py +++ b/tests/test_lora_resize_factor_space.py @@ -50,6 +50,85 @@ def _read_uel(path, async_method=None, mark_method=None): return result +@pytest.mark.parametrize("glob", [False, True]) +def test_resize_filters_preserve_and_discard_whole_layers(monkeypatch, tmp_path, resize, glob): + source = tmp_path / "filters.safetensors" + tensors = {} + for name in ("keep", "drop", "resize"): + tensors[f"diffusion_model.{name}.lora_A.weight"] = torch.eye(3, dtype=torch.float16) + tensors[f"diffusion_model.{name}.lora_B.weight"] = torch.eye(3, dtype=torch.float16) + tensors[f"diffusion_model.{name}.alpha"] = torch.tensor(1.5) + tensors["diffusion_model.drop.extra"] = torch.ones(1) + tensors["standalone"] = torch.ones(1) + _write_uel(source, tensors) + monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None) + original_async = MemoryEfficientSafeOpen.async_stream + streamed = [] + + def track(self, keys, **kwargs): + streamed.extend(keys) + yield from original_async(self, keys, **kwargs) + + monkeypatch.setattr(MemoryEfficientSafeOpen, "async_stream", track) + output = resize.resize_lora_file( + str(source), 1, None, None, "cpu", torch.bfloat16, "filtered", verbose=False, + exclude_patterns="*.keep.lora_A.weight\n*.drop.*" if glob else r"\.keep\.lora_A\.weight$" + "\n" + r"\.drop\.", + discard_patterns="*.drop.alpha\nstandalone" if glob else r"\.drop\.alpha$" + "\n^standalone$", + glob_patterns=glob, + ) + assert not any(".drop." in key or key == "standalone" for key in streamed) + actual = _read_uel(output, original_async) + assert not any(".drop." in key or key == "standalone" for key in actual) + for key in tensors: + if ".keep." in key: + assert actual[key].dtype == tensors[key].dtype + assert torch.equal(actual[key], tensors[key]) + assert actual["diffusion_model.resize.lora_A.weight"].shape[0] == 1 + + +@pytest.mark.parametrize("node_name,extra", [ + ("LoRAResizeFixed", {"new_rank": 1}), + ("LoRAResizeRatio", {"max_rank": 1, "ratio": 2.0}), + ("LoRAResizeFrobenius", {"max_rank": 1, "min_rank": 1, "target": 0.9}), + ("LoRAResizeCumulative", {"max_rank": 1, "target": 0.9}), +]) +def test_all_resize_nodes_forward_filters(monkeypatch, tmp_path, resize, node_name, extra): + node = getattr(resize, node_name) + controls = {item.id: item for item in node.define_schema().inputs} + assert controls["exclude_patterns"].default == "" + assert controls["discard_patterns"].default == "" + assert controls["glob_patterns"].default is False + calls = [] + monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(resize.folder_paths, "get_full_path_or_raise", lambda *args: "input") + monkeypatch.setattr(resize, "resize_lora_file", lambda *args, **kwargs: calls.append(kwargs)) + node.execute(lora_name="input", output_filename="out", save_dtype="fp16", device="cpu", + force_clear_cache=False, exclude_patterns="keep*", discard_patterns="drop*", + glob_patterns=True, **extra) + assert calls[0]["exclude_patterns"] == "keep*" + assert calls[0]["discard_patterns"] == "drop*" + assert calls[0]["glob_patterns"] is True + + +def test_resize_invalid_filter_fails_before_writer(monkeypatch, tmp_path, resize): + import re + + source = tmp_path / "invalid_filter.safetensors" + _write_uel(source, {"layer.diff": torch.ones(2)}) + monkeypatch.setattr(resize, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None) + + def unexpected_writer(*args, **kwargs): + pytest.fail("Writer must not open for invalid patterns") + + monkeypatch.setattr(resize, "atomic_uel_writer", unexpected_writer) + with pytest.raises(re.error): + resize.resize_lora_file(str(source), 1, None, None, "cpu", torch.float32, + "unused", exclude_patterns="[") + + def _dense_truncated(up, down, rank): dense = up.float() @ down.float() u, s, vh = torch.linalg.svd(dense, full_matrices=False)