diff --git a/nodes/prunekeys.py b/nodes/prunekeys.py index a7b50ab..facc353 100644 --- a/nodes/prunekeys.py +++ b/nodes/prunekeys.py @@ -1,3 +1,4 @@ +from contextlib import closing import os import re import folder_paths @@ -8,12 +9,24 @@ from .utils import convert_pt_to_safetensors from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation from unifiedefficientloader import MemoryEfficientSafeOpen -from .uel_io import atomic_uel_writer +from .uel_io import atomic_uel_writer, stream_work_units from .artifact_paths import canonical_model_artifact_path +from .quantization_guard import ( + UPSTREAM_WRITER_STORAGE_CODES, + _is_low_bit, + storage_dtype_code, + write_preserved_tensor, +) -def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str, - use_regex: bool, output_filename: str) -> str: +def _prune_keys( + model_name: str, + model_type: str, + keys_to_prune_str: str, + use_regex: bool, + output_filename: str, + include_mode: bool = False, +) -> str: """Shared logic for all PruneKeys nodes.""" model_path = folder_paths.get_full_path_or_raise(model_type, model_name) @@ -27,43 +40,73 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str, else: model_path_to_load = model_path - # Prepare memory before operation - model_size_gb = estimate_model_size(model_path_to_load) - prepare_for_large_operation(model_size_gb * 1.2) - - patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()] - - if not patterns: + raw_patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()] + if not raw_patterns: raise ValueError("No keys/patterns provided to prune.") - # Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet + compiled_patterns = [] + if use_regex: + for pat in raw_patterns: + try: + compiled_patterns.append(re.compile(pat)) + except re.error as exc: + raise ValueError(f"Invalid regex pattern '{pat}': {exc}") from exc + + model_size_gb = estimate_model_size(model_path_to_load) + if model_size_gb: + prepare_for_large_operation(model_size_gb * 1.2) + output_path, output_name = canonical_model_artifact_path(model_type, output_filename) - # Stream tensors, filter on the fly, write immediately - with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler: - metadata = handler.metadata() or {} - with atomic_uel_writer(output_path, metadata) as writer: - all_keys = handler.keys() - pbar = comfy.utils.ProgressBar(len(all_keys)) - kept_keys = [ - key for key in all_keys - if not ( - any(re.search(pattern, key) for pattern in patterns) - if use_regex else any(pattern in key for pattern in patterns) - ) - ] - stream = handler.async_stream( - kept_keys, batch_size=1, prefetch_batches=1, pin_memory=False + handler = MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) + try: + metadata = (handler.metadata() or {}).copy() + all_keys = list(handler.keys()) + kept_keys = [] + for key in all_keys: + matched = ( + any(pat.search(key) is not None for pat in compiled_patterns) + if use_regex + else any(pat in key for pat in raw_patterns) ) - for batch in tqdm(stream, total=len(kept_keys), desc="Pruning keys", unit="keys"): - key, tensor = batch[0] - writer.write_batch([(key, tensor.contiguous())]) - handler.mark_processed(key) - pbar.update(1) - pbar.update(len(all_keys) - len(kept_keys)) + keep = matched if include_mode else not matched + if keep: + kept_keys.append(key) - # Cleanup after operation - cleanup_after_operation() + work_units = [] + raw_preserve_keys = set() + for key in kept_keys: + code = storage_dtype_code(handler, key) + if _is_low_bit(handler, key) or ( + code is not None and code not in UPSTREAM_WRITER_STORAGE_CODES + ): + raw_preserve_keys.add(key) + work_units.append((key, {})) + else: + work_units.append((key, {0: [key]})) + + pbar = comfy.utils.ProgressBar(len(all_keys)) + with atomic_uel_writer(output_path, metadata) as writer, closing( + stream_work_units({0: handler}, work_units, pin_memory=False) + ) as streamed: + for key, loaded in tqdm(streamed, total=len(work_units), desc="Pruning keys", unit="keys"): + if key in raw_preserve_keys: + write_preserved_tensor( + writer, key, handler, output_key=key, force_raw=True + ) + else: + tensor = loaded[(0, key)] + try: + write_preserved_tensor( + writer, key, handler, output_key=key, tensor=tensor + ) + finally: + del tensor + pbar.update(1) + pbar.update(len(all_keys) - len(work_units)) + finally: + handler.__exit__(None, None, None) + cleanup_after_operation() return output_name @@ -85,6 +128,7 @@ class ModelPruneKeys(io.ComfyNode): io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."), io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."), io.String.Input("output_filename", default="pruned_model", tooltip="Output filename without extension, written under the matching ComfyUI model directory."), + io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."), ], outputs=[ io.AnyType.Output(display_name="output_path"), @@ -92,10 +136,22 @@ class ModelPruneKeys(io.ComfyNode): ) @classmethod - def execute(cls, diffusionmodel_name: str, keys_to_prune: str, - use_regex: bool, output_filename: str) -> io.NodeOutput: - path = _prune_keys(diffusionmodel_name, "diffusion_models", - keys_to_prune, use_regex, output_filename) + def execute( + cls, + diffusionmodel_name: str, + keys_to_prune: str, + use_regex: bool, + output_filename: str, + include_mode: bool = False, + ) -> io.NodeOutput: + path = _prune_keys( + diffusionmodel_name, + "diffusion_models", + keys_to_prune, + use_regex, + output_filename, + include_mode=include_mode, + ) return io.NodeOutput(path) @@ -116,6 +172,7 @@ class TextEncoderPruneKeys(io.ComfyNode): io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."), io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."), io.String.Input("output_filename", default="pruned_textencoder", tooltip="Output filename without extension, written under the matching ComfyUI model directory."), + io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."), ], outputs=[ io.AnyType.Output(display_name="output_path"), @@ -123,10 +180,22 @@ class TextEncoderPruneKeys(io.ComfyNode): ) @classmethod - def execute(cls, textencoder_name: str, keys_to_prune: str, - use_regex: bool, output_filename: str) -> io.NodeOutput: - path = _prune_keys(textencoder_name, "text_encoders", - keys_to_prune, use_regex, output_filename) + def execute( + cls, + textencoder_name: str, + keys_to_prune: str, + use_regex: bool, + output_filename: str, + include_mode: bool = False, + ) -> io.NodeOutput: + path = _prune_keys( + textencoder_name, + "text_encoders", + keys_to_prune, + use_regex, + output_filename, + include_mode=include_mode, + ) return io.NodeOutput(path) @@ -147,6 +216,7 @@ class LoRAPruneKeys(io.ComfyNode): io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."), io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."), io.String.Input("output_filename", default="pruned_lora", tooltip="Output filename without extension, written under the matching ComfyUI model directory."), + io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."), ], outputs=[ io.AnyType.Output(display_name="output_path"), @@ -154,10 +224,22 @@ class LoRAPruneKeys(io.ComfyNode): ) @classmethod - def execute(cls, lora_name: str, keys_to_prune: str, - use_regex: bool, output_filename: str) -> io.NodeOutput: - path = _prune_keys(lora_name, "loras", - keys_to_prune, use_regex, output_filename) + def execute( + cls, + lora_name: str, + keys_to_prune: str, + use_regex: bool, + output_filename: str, + include_mode: bool = False, + ) -> io.NodeOutput: + path = _prune_keys( + lora_name, + "loras", + keys_to_prune, + use_regex, + output_filename, + include_mode=include_mode, + ) return io.NodeOutput(path) @@ -178,6 +260,7 @@ class CheckpointPruneKeys(io.ComfyNode): io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."), io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."), io.String.Input("output_filename", default="pruned_checkpoint", tooltip="Output filename without extension, written under the matching ComfyUI model directory."), + io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."), ], outputs=[ io.AnyType.Output(display_name="output_path"), @@ -185,10 +268,22 @@ class CheckpointPruneKeys(io.ComfyNode): ) @classmethod - def execute(cls, ckpt_name: str, keys_to_prune: str, - use_regex: bool, output_filename: str) -> io.NodeOutput: - path = _prune_keys(ckpt_name, "checkpoints", - keys_to_prune, use_regex, output_filename) + def execute( + cls, + ckpt_name: str, + keys_to_prune: str, + use_regex: bool, + output_filename: str, + include_mode: bool = False, + ) -> io.NodeOutput: + path = _prune_keys( + ckpt_name, + "checkpoints", + keys_to_prune, + use_regex, + output_filename, + include_mode=include_mode, + ) return io.NodeOutput(path) @@ -209,6 +304,7 @@ class EmbeddingPruneKeys(io.ComfyNode): io.String.Input("keys_to_prune", multiline=True, default="", tooltip="One tensor-key pattern per line. Blank lines are ignored."), io.Boolean.Input("use_regex", default=False, tooltip="Interpret each line as a regular expression; otherwise match literal key text."), io.String.Input("output_filename", default="pruned_embedding", tooltip="Output filename without extension, written under the matching ComfyUI model directory."), + io.Boolean.Input("include_mode", default=False, tooltip="Keep only matching keys instead of removing them. Nonmatching keys are pruned from the output."), ], outputs=[ io.AnyType.Output(display_name="output_path"), @@ -216,8 +312,20 @@ class EmbeddingPruneKeys(io.ComfyNode): ) @classmethod - def execute(cls, embedding: str, keys_to_prune: str, - use_regex: bool, output_filename: str) -> io.NodeOutput: - path = _prune_keys(embedding, "embeddings", - keys_to_prune, use_regex, output_filename) + def execute( + cls, + embedding: str, + keys_to_prune: str, + use_regex: bool, + output_filename: str, + include_mode: bool = False, + ) -> io.NodeOutput: + path = _prune_keys( + embedding, + "embeddings", + keys_to_prune, + use_regex, + output_filename, + include_mode=include_mode, + ) return io.NodeOutput(path) diff --git a/tests/test_artifact_paths.py b/tests/test_artifact_paths.py index 00bf701..2b3646e 100644 --- a/tests/test_artifact_paths.py +++ b/tests/test_artifact_paths.py @@ -11,6 +11,7 @@ REPO_ROOT = Path(__file__).resolve().parents[1] PACKAGE_NAME = "modelutils_artifact_path_tests" MIGRATED_SAVE_MODULES = ( "consensus_merger.py", + "cwb_model_graft.py", "cwb_delta_lora_merger.py", "merger.py", "lora_merger.py", diff --git a/tests/test_groups.toml b/tests/test_groups.toml index dce206b..79eba69 100644 --- a/tests/test_groups.toml +++ b/tests/test_groups.toml @@ -56,7 +56,7 @@ paths = ["nodes/lora_resize.py"] python_tests = ["tests/test_lora_resize_factor_space.py", "tests/test_lora_resize_min_rank.py"] [groups.merger_filters] -paths = ["nodes/merger.py"] +paths = ["nodes/merger.py", "nodes/prunekeys.py"] python_tests = ["tests/test_merger_filters.py"] [groups.layer_parameters] diff --git a/tests/test_merger_filters.py b/tests/test_merger_filters.py index b6f7665..3b7ac20 100644 --- a/tests/test_merger_filters.py +++ b/tests/test_merger_filters.py @@ -231,3 +231,54 @@ def test_three_diffusion_models_match_across_prefixes(monkeypatch, tmp_path, mod output = load_file(str(tmp_path / "diffusion_models" / "three_prefix_merge.safetensors")) assert list(output) == ["layer.weight"] torch.testing.assert_close(output["layer.weight"], torch.tensor([2.0, 2.0])) + + +def test_prunekeys_include_mode_and_uel_streaming(monkeypatch, tmp_path): + prunekeys = _module("nodes.prunekeys") + for cls in ( + prunekeys.ModelPruneKeys, + prunekeys.TextEncoderPruneKeys, + prunekeys.LoRAPruneKeys, + prunekeys.CheckpointPruneKeys, + prunekeys.EmbeddingPruneKeys, + ): + schema = cls.define_schema() + assert schema.inputs[-1].id == "include_mode" + assert schema.inputs[-1].default is False + + src = tmp_path / "source.safetensors" + save_file( + { + "layers.0.attn.weight": torch.tensor([1.0, 2.0]), + "layers.0.ffn.weight": torch.tensor([3.0, 4.0]), + "layers.1.attn.weight": torch.tensor([5.0, 6.0]), + }, + str(src), + metadata={"author": "test"}, + ) + monkeypatch.setattr(prunekeys.folder_paths, "get_full_path_or_raise", lambda _kind, _name: str(src)) + monkeypatch.setattr(prunekeys.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(prunekeys, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(prunekeys, "cleanup_after_operation", lambda: None) + + # Default include_mode=False removes matching keys + rel_removed = prunekeys.ModelPruneKeys.execute( + diffusionmodel_name="source.safetensors", + keys_to_prune="attn", + use_regex=False, + output_filename="removed_attn", + include_mode=False, + ).result[0] + removed_tensors = load_file(str(tmp_path / "diffusion_models" / rel_removed)) + assert set(removed_tensors) == {"layers.0.ffn.weight"} + + # include_mode=True keeps only matching keys + rel_kept = prunekeys.ModelPruneKeys.execute( + diffusionmodel_name="source.safetensors", + keys_to_prune=r"^layers\.\d+\.attn\.", + use_regex=True, + output_filename="kept_attn", + include_mode=True, + ).result[0] + kept_tensors = load_file(str(tmp_path / "diffusion_models" / rel_kept)) + assert set(kept_tensors) == {"layers.0.attn.weight", "layers.1.attn.weight"}