nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

lora(resize) add exclusion and discard filters to resize nodes

Browse Source
This commit is contained in:
silveroxides
2026-09-08 08:39:03 +02:00
parent 0a808b6865
commit 73cefa409d
2 changed files with 144 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/lora_resize.py
+65 -8
View File
@@ -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)
tests/test_lora_resize_factor_space.py
+79
View File
@@ -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)