diff --git a/README.md b/README.md index 14f30d3..c2dd6ea 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,18 @@ A collection of ComfyUI custom nodes for inspecting, modifying, merging, and cre - **LoRA Extraction** – Extract LoRA adapters from model pairs using various SVD rank selection methods (Fixed, Ratio, Quantile, Knee-detection, Frobenius-norm) - **Diffusion Model Dtype Conversion** – Stream models to fp32, fp16, or bf16 while preserving excluded tensor dtypes +## Layer filters + +Nodes with `exclude_patterns` or `skip_patterns` have an appended `include_mode` toggle +(off by default). Turn it on to process only layers matching that same field, +using the node's existing regex or glob syntax. An empty include filter selects +nothing. Nonmatches follow the node's usual exclusion behavior: preserve the +source/anchor layer for pattern-exclusion mergers, resize, or conversion; omit +it from analysis or extraction. LoRA Merge to Model retains its existing skip +semantics: nonmatching base tensors are omitted from the saved model, while +guarded low-bit base tensors are always preserved. Its filter uses regex only. +`discard_patterns`, where available, still takes precedence. + ## Example Workflows

diff --git a/docs/model_analysis.md b/docs/model_analysis.md index d428696..116dbc2 100644 --- a/docs/model_analysis.md +++ b/docs/model_analysis.md @@ -4,6 +4,12 @@ These nodes compare two safetensors models without merging them or writing an output file. Separate nodes are provided for checkpoints, diffusion models, text encoders, LoRAs, and embeddings. +`exclude_patterns` accepts one regex per line, or glob patterns when +`glob_patterns` is enabled. By default, matches are excluded from metrics and +topology counts. Enable `include_mode` to analyze only matches instead; an empty +include filter selects no tensors. The same field and matching syntax are used +in both modes. + ## Outputs `comparison_report` contains conventional numerical differences and topology diff --git a/nodes/consensus_merger.py b/nodes/consensus_merger.py index 1538d6d..a3485a1 100644 --- a/nodes/consensus_merger.py +++ b/nodes/consensus_merger.py @@ -883,6 +883,7 @@ def _cwb_merge_metadata( "force_clear_cache", "override_dtype", "include_1d_diffs", + "include_mode", "counterfactual_weight_sweep", "embedding_coalesce", "target_vector_count", @@ -897,9 +898,10 @@ def _cwb_merge_metadata( metadata["merge_options"][key] = params[key] exclude_patterns = params["exclude_patterns"].split() discard_patterns = params["discard_patterns"].split() - if exclude_patterns or discard_patterns: + if exclude_patterns or discard_patterns or params.get("include_mode", False): metadata["filters"] = { "exclude_patterns": exclude_patterns, + "include_mode": params.get("include_mode", False), "discard_patterns": discard_patterns, "glob_patterns": params["glob_patterns"], } @@ -1794,6 +1796,7 @@ class ConsensusMergerLogic: requested_dtype = _requested_dtype(params["save_dtype"]) mismatch_mode = params["mismatch_mode"] glob_mode = params["glob_patterns"] + include_mode = params.get("include_mode", False) exclude = _compile_patterns(params["exclude_patterns"], glob_mode=glob_mode) discard = _compile_patterns(params["discard_patterns"], glob_mode=glob_mode) pbar = comfy.utils.ProgressBar(len(keys)) @@ -1834,7 +1837,8 @@ class ConsensusMergerLogic: secondary_only = not primary_owned preserve_index = 0 if primary_owned else source_indices[0] guarded = any(key in low_bit_sets[i] for i in source_indices) - excluded = _matches_any_pattern(key, exclude, glob_mode=glob_mode) + matched = _matches_any_pattern(key, exclude, glob_mode=glob_mode) + excluded = not matched if include_mode else matched if guarded or excluded: write_preserved_tensor( writer, @@ -2058,6 +2062,7 @@ class ConsensusMergerLogic: mismatch_mode = params["mismatch_mode"] include_1d = params.get("include_1d_diffs", False) glob_mode = params["glob_patterns"] + include_mode = params.get("include_mode", False) exclude = _compile_patterns(params["exclude_patterns"], glob_mode=glob_mode) discard = _compile_patterns(params["discard_patterns"], glob_mode=glob_mode) written = set() @@ -2266,10 +2271,11 @@ class ConsensusMergerLogic: keys is not None and has_non_alpha_low_bit(keys, low_bit_sets[index]) for index, _, keys in matches ) - excluded = any( + matched = any( _matches_any_pattern(key, exclude, glob_mode=glob_mode) for key in layer_keys ) + excluded = not matched if include_mode else matched companion_bearing = any( keys is not None and layer_has_companions(keys) for _, _, keys in matches @@ -2647,6 +2653,11 @@ def _common_inputs( default=False, tooltip="Force generated tensors to save_dtype; guarded tensors and enabled 1D direct diffs are exempt.", ), + io.Boolean.Input( + "include_mode", + default=False, + tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are merged; nonmatching layers are preserved from the anchor.", + ), ]) return inputs @@ -2705,11 +2716,13 @@ class _CWBMergerNode(io.ComfyNode): cls.DEFAULT_PRESET, ) if cls.LORA_MODE: + include_mode = inputs.pop() inputs.append(io.Boolean.Input( "include_1d_diffs", default=False, tooltip="CWB-merge 1D direct diffs as FP32. Disabled preserves Model A.", )) + inputs.append(include_mode) return io.Schema( node_id=cls.NODE_ID, display_name=cls.DISPLAY_NAME, @@ -2838,6 +2851,7 @@ class CWBLoRAMultiMerger(io.ComfyNode): io.Boolean.Input("override_dtype", default=False, tooltip="Force generated floating factors to the requested save dtype."), io.Boolean.Input("include_1d_diffs", default=False, tooltip="CWB-merge 1D direct differences as FP32 instead of preserving the anchor."), io.Boolean.Input("counterfactual_weight_sweep", default=False, tooltip="Evaluate alpha, beta, similarity-threshold, DSC, and comfort-bandpass weight combinations from each already-computed consensus similarity vector and append them to the report without additional model loads or saved outputs."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are merged; nonmatching layers are preserved from the anchor."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), diff --git a/nodes/cwb_delta_lora_merger.py b/nodes/cwb_delta_lora_merger.py index f1cf247..11f7614 100644 --- a/nodes/cwb_delta_lora_merger.py +++ b/nodes/cwb_delta_lora_merger.py @@ -202,6 +202,7 @@ def _build_layer_units( exclude_patterns, discard_patterns, glob_patterns, + include_mode=False, ): units = [] stream_keys = [[] for _ in handlers] @@ -213,7 +214,8 @@ def _build_layer_units( if _matches_pattern(output_key, discard_patterns, glob_patterns): rejected += 1 continue - preserve = _matches_pattern(output_key, exclude_patterns, glob_patterns) + matched = _matches_pattern(output_key, exclude_patterns, glob_patterns) + preserve = not matched if include_mode else matched candidates = [] shapes = {} for index, block in sorted(mapped): @@ -414,6 +416,7 @@ class DeltaCWBLoRAMergerLogic: discard_patterns, glob_patterns, force_clear_cache, + include_mode=False, ): prepare_for_large_operation( sum(estimate_model_size(path) for path in lora_paths) * 2.5, @@ -439,6 +442,7 @@ class DeltaCWBLoRAMergerLogic: compiled_excludes, compiled_discards, glob_patterns, + include_mode, ) report.rejected_unsupported_layers = rejected cursors = [ @@ -547,6 +551,7 @@ class DeltaCWBLoRAMergerLogic: kwargs["discard_patterns"], kwargs["glob_patterns"], kwargs["force_clear_cache"], + kwargs.get("include_mode", False), ) return io.NodeOutput(filename, report) @@ -638,6 +643,11 @@ def _fixed_inputs(count, default_filename): default=True, tooltip="Clear Python and CUDA caches before each layer to reduce retained memory.", ), + io.Boolean.Input( + "include_mode", + default=False, + tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are merged; nonmatching layers are preserved from LoRA 1.", + ), ]) return inputs diff --git a/nodes/dora_extract_wd.py b/nodes/dora_extract_wd.py index d687e57..ac45a14 100644 --- a/nodes/dora_extract_wd.py +++ b/nodes/dora_extract_wd.py @@ -497,6 +497,7 @@ def extract_dora_from_files( force_clear_cache: bool = True, glob_skip_patterns: bool = False, knee_probe_offset: int = 32, + include_mode: bool = False, ) -> None: """ Extract LoRA from difference between two models, writing incrementally to disk. @@ -543,7 +544,7 @@ def extract_dora_from_files( work_units = [ (key, key, key if key in keys_b else None) for key in weight_keys - if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) + if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)) and not (key not in keys_b and mismatch_mode == "skip") ] pbar = comfy.utils.ProgressBar(len(work_units)) @@ -554,7 +555,7 @@ def extract_dora_from_files( weight_diff = None layer_device = device - if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns): + if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)): return "skipped", None # Load tensors with pinned memory for CUDA @@ -773,6 +774,7 @@ def _get_common_inputs(): io.Boolean.Input("glob_skip_patterns", default=False, tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). " "When False (default), patterns are Python regex matched as substrings."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Skip Patterns as a whitelist instead. Only matching layers are extracted; an empty whitelist extracts nothing."), ] @@ -803,7 +805,7 @@ class DoRAExtractFixed(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -813,7 +815,7 @@ class DoRAExtractFixed(io.ComfyNode): model_a_path, model_b_path, "fixed", linear_dim, conv_dim, device, save_dtype, output_path, linear_dim, conv_dim, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter, - lazy_load, force_clear_cache, glob_skip_patterns + lazy_load, force_clear_cache, glob_skip_patterns, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -850,7 +852,7 @@ class DoRAExtractRatio(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -861,7 +863,7 @@ class DoRAExtractRatio(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset + glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -898,7 +900,7 @@ class DoRAExtractQuantile(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -909,7 +911,7 @@ class DoRAExtractQuantile(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset + glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -945,7 +947,7 @@ class DoRAExtractKnee(io.ComfyNode): def execute(cls, model_a, model_b, knee_method, knee_probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -956,7 +958,7 @@ class DoRAExtractKnee(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, + glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=knee_probe_offset, ) @@ -994,7 +996,7 @@ class DoRAExtractFrobenius(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -1005,7 +1007,7 @@ class DoRAExtractFrobenius(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset + glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset, include_mode=include_mode ) return io.NodeOutput(output_name) diff --git a/nodes/dora_learned_wd.py b/nodes/dora_learned_wd.py index c66c0a9..8368c36 100644 --- a/nodes/dora_learned_wd.py +++ b/nodes/dora_learned_wd.py @@ -354,6 +354,7 @@ def extract_dora_learned_from_files( force_clear_cache: bool = True, glob_skip_patterns: bool = False, knee_probe_offset: int = 32, + include_mode: bool = False, ) -> None: """ Extract Learned DoRA from difference between two models. @@ -380,7 +381,7 @@ def extract_dora_learned_from_files( work_units = [ (key, key, key if key in keys_b else None) for key in weight_keys - if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) + if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)) and not (key not in keys_b and mismatch_mode == "skip") ] pbar = comfy.utils.ProgressBar(len(work_units)) @@ -393,7 +394,7 @@ def extract_dora_learned_from_files( tensor_ft = None layer_device = device - if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns): + if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)): return "skipped", None if key not in keys_b: @@ -597,7 +598,7 @@ class DoRALearnedExtractFixed(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_dim, conv_dim, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, svd_niter, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -607,7 +608,7 @@ class DoRALearnedExtractFixed(io.ComfyNode): model_a_path, model_b_path, "fixed", linear_dim, conv_dim, device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_dim, conv_dim, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter, - lazy_load, force_clear_cache, glob_skip_patterns + lazy_load, force_clear_cache, glob_skip_patterns, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -641,7 +642,7 @@ class DoRALearnedExtractRatio(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, - output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -651,7 +652,7 @@ class DoRALearnedExtractRatio(io.ComfyNode): model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio, device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=probe_offset ) @@ -686,7 +687,7 @@ class DoRALearnedExtractQuantile(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, - output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -696,7 +697,7 @@ class DoRALearnedExtractQuantile(io.ComfyNode): model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile, device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=probe_offset ) @@ -731,7 +732,7 @@ class DoRALearnedExtractKnee(io.ComfyNode): def execute(cls, model_a, model_b, knee_method, knee_probe_offset, linear_max_rank, conv_max_rank, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, - output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -742,7 +743,7 @@ class DoRALearnedExtractKnee(io.ComfyNode): device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset + glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -776,7 +777,7 @@ class DoRALearnedExtractFrobenius(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, - output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -786,7 +787,7 @@ class DoRALearnedExtractFrobenius(io.ComfyNode): model_a_path, model_b_path, "sv_fro", linear_target, conv_target, device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=probe_offset ) diff --git a/nodes/dtype_conversion.py b/nodes/dtype_conversion.py index c20a92d..6dca6c4 100644 --- a/nodes/dtype_conversion.py +++ b/nodes/dtype_conversion.py @@ -55,6 +55,7 @@ def convert_diffusion_model_dtype( exclude_patterns: str, output_filename: str, reference_model: str | None = None, + include_mode: bool = False, ) -> tuple[str, str]: try: destination_dtype = TARGET_DTYPES[target_dtype] @@ -104,7 +105,7 @@ def convert_diffusion_model_dtype( key, source = batch[0] original_dtypes[str(source.dtype).removeprefix("torch.")] += 1 output = source - if _is_excluded(key, exclusions): + if _is_excluded(key, exclusions) != include_mode: counts["excluded"] += 1 elif not source.is_floating_point(): counts["non_floating"] += 1 @@ -175,7 +176,8 @@ class DiffusionModelDtypeConversion(io.ComfyNode): description=( "Streams a diffusion-model safetensors file into fp32, fp16, " "or bf16. Floating tensors matching an exclusion regex retain " - "their original dtype; non-floating tensors are always preserved." + "their original dtype, or enable Include Mode to convert only matches; " + "non-floating tensors are always preserved." ), inputs=[ io.Combo.Input( @@ -185,7 +187,7 @@ class DiffusionModelDtypeConversion(io.ComfyNode): ), io.Combo.Input( "target_dtype", options=["fp32", "fp16", "bf16"], - tooltip="Target dtype for floating tensors not matched by an exclusion pattern.", + tooltip="Target dtype for floating tensors selected by the layer filter.", ), io.Combo.Input( "reference_model", @@ -199,13 +201,18 @@ class DiffusionModelDtypeConversion(io.ComfyNode): multiline=True, tooltip=( "Optional Python regex patterns, one per line. Matching " - "tensor keys retain their original dtype." + "tensor keys retain their original dtype, or are the only keys " + "converted when Include Mode is enabled." ), ), io.String.Input( "output_filename", default="converted_model", tooltip="Output filename without extension, written under ComfyUI's diffusion-model directory.", ), + io.Boolean.Input( + "include_mode", default=False, + tooltip="Use Exclude Patterns as an include-only filter. Only matching tensors are converted; an empty filter converts nothing. Nonmatches retain their original dtype.", + ), ], outputs=[ io.AnyType.Output(display_name="output_path"), @@ -222,6 +229,7 @@ class DiffusionModelDtypeConversion(io.ComfyNode): reference_model: str, exclude_patterns: str, output_filename: str, + include_mode: bool = False, ) -> io.NodeOutput: output_path, report = convert_diffusion_model_dtype( model_name, @@ -229,5 +237,6 @@ class DiffusionModelDtypeConversion(io.ComfyNode): exclude_patterns, output_filename, reference_model, + include_mode, ) return io.NodeOutput(output_path, report) diff --git a/nodes/lodestone_merger.py b/nodes/lodestone_merger.py index 0612ebc..7aefe76 100644 --- a/nodes/lodestone_merger.py +++ b/nodes/lodestone_merger.py @@ -92,7 +92,7 @@ def _is_cuda_oom(error): ) -def _build_units(layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns): +def _build_units(layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns, include_mode=False): units = [] stream_keys = [[] for _ in handlers] for core, mapped in layer_map.items(): @@ -101,7 +101,8 @@ def _build_units(layer_map, infos, handlers, mismatch_mode, exclude_patterns, di output_key = canonical_lora_key(core, "diff") if _matches(output_key, discard_patterns, glob_patterns): continue - preserve = _matches(output_key, exclude_patterns, glob_patterns) + matched = _matches(output_key, exclude_patterns, glob_patterns) + preserve = not matched if include_mode else matched sources = [] for index, block in sorted(mapped): roles = infos[index]["pairs"][block] @@ -178,7 +179,7 @@ def _process(unit, loaded, method, save_dtype, requested_device): def merge_lodestone_loras(lora_paths, method, device, save_dtype, output_filename, mismatch_mode="skip", exclude_patterns="", discard_patterns="", - glob_patterns=False, verbose=True): + glob_patterns=False, verbose=True, include_mode=False): operation = "Lodestone LoRA Merge" prepare_for_large_operation(sum(estimate_model_size(path) for path in lora_paths) * 2.5, torch.device(device)) handlers = [MemoryEfficientSafeOpen(path, low_memory=True) for path in lora_paths] @@ -192,7 +193,7 @@ def merge_lodestone_loras(lora_paths, method, device, save_dtype, output_filenam layer_map = build_lora_layer_map(infos, None) validate_canonical_blocks(layer_map, operation) units, stream_keys = _build_units( - layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns + layer_map, infos, handlers, mismatch_mode, exclude_patterns, discard_patterns, glob_patterns, include_mode ) cursors = [ AsyncTensorCursor(handler, keys, pin_memory=torch.device(device).type == "cuda") @@ -244,6 +245,7 @@ def _controls(default_filename): io.String.Input("exclude_patterns", default="", multiline=True, tooltip="Optional layer-name patterns to copy from LoRA 1 instead of merging. Enter one pattern per line; useful for protecting specific blocks."), io.String.Input("discard_patterns", default="", multiline=True, tooltip="Optional layer-name patterns to leave out of the saved LoRA completely. Enter one pattern per line; discarded layers cannot affect the model."), io.Boolean.Input("glob_patterns", default=False, tooltip="Choose filter syntax. Off uses regular expressions; on uses shell-style globs where * matches any text and ? matches one character."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching layers are merged; nonmatching layers are preserved from LoRA 1."), ] @@ -252,7 +254,7 @@ def _execute(kwargs, count): dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[kwargs["save_dtype"]] merge_lodestone_loras(paths, kwargs["calc_mode"], kwargs["process_device"], dtype, kwargs["output_filename"], kwargs["mismatch_mode"], - kwargs["exclude_patterns"], kwargs["discard_patterns"], kwargs["glob_patterns"]) + kwargs["exclude_patterns"], kwargs["discard_patterns"], kwargs["glob_patterns"], include_mode=kwargs.get("include_mode", False)) return canonical_model_artifact_path("loras", kwargs["output_filename"])[1] diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index 242cbeb..7f47a8a 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -499,6 +499,7 @@ def extract_lora_from_files( glob_skip_patterns: bool = False, include_1d_diffs: bool = False, knee_probe_offset: int = 32, + include_mode: bool = False, ) -> None: """ Extract LoRA from difference between two models, writing incrementally to disk. @@ -560,7 +561,7 @@ def extract_lora_from_files( work_units = [ (key, key, key if key in keys_b else None) for key in extraction_keys - if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) + if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)) and not (key not in keys_b and mismatch_mode == "skip") ] pbar = comfy.utils.ProgressBar(len(work_units)) @@ -571,7 +572,7 @@ def extract_lora_from_files( weight_diff = None layer_device = device - if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns): + if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)): return "skipped", None source_dtypes = [handler_a.get_dtype(key)] @@ -808,6 +809,7 @@ def _get_common_inputs(): "When False (default), patterns are Python regex matched as substrings."), io.Boolean.Input("include_1d_diffs", default=False, tooltip="Include 1D tensors as FP32 and keys ending in .scale or .lin as full .diff patches"), + io.Boolean.Input("include_mode", default=False, tooltip="Use Skip Patterns as a whitelist instead. Only matching layers are extracted; an empty whitelist extracts nothing."), ] @@ -839,7 +841,7 @@ class LoRAExtractFixed(io.ComfyNode): def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, - include_1d_diffs=False) -> io.NodeOutput: + include_1d_diffs=False, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -849,7 +851,7 @@ class LoRAExtractFixed(io.ComfyNode): model_a_path, model_b_path, "fixed", linear_dim, conv_dim, device, save_dtype, output_path, linear_dim, conv_dim, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter, - lazy_load, force_clear_cache, glob_skip_patterns, include_1d_diffs + lazy_load, force_clear_cache, glob_skip_patterns, include_1d_diffs, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -887,7 +889,7 @@ class LoRAExtractRatio(io.ComfyNode): def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, - include_1d_diffs=False) -> io.NodeOutput: + include_1d_diffs=False, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -898,7 +900,7 @@ class LoRAExtractRatio(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, + glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, include_mode=include_mode, knee_probe_offset=probe_offset, ) @@ -937,7 +939,7 @@ class LoRAExtractQuantile(io.ComfyNode): def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, - include_1d_diffs=False) -> io.NodeOutput: + include_1d_diffs=False, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -948,7 +950,7 @@ class LoRAExtractQuantile(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, + glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, include_mode=include_mode, knee_probe_offset=probe_offset, ) @@ -986,7 +988,7 @@ class LoRAExtractKnee(io.ComfyNode): linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, - include_1d_diffs=False) -> io.NodeOutput: + include_1d_diffs=False, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -997,7 +999,7 @@ class LoRAExtractKnee(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, + glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, include_mode=include_mode, knee_probe_offset=knee_probe_offset, ) @@ -1036,7 +1038,7 @@ class LoRAExtractFrobenius(io.ComfyNode): def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, - include_1d_diffs=False) -> io.NodeOutput: + include_1d_diffs=False, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) @@ -1047,7 +1049,7 @@ class LoRAExtractFrobenius(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, + glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, include_mode=include_mode, knee_probe_offset=probe_offset, ) diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index 0af7145..2870a7e 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -648,12 +648,14 @@ def _resize_filter_inputs(): 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."), + io.Boolean.Input("include_mode", default=False, tooltip="Treat exclude patterns as an include list. Matching layers are resized; nonmatching layers are preserved unchanged. Discard patterns still remove matching layers."), ] def _build_resize_work_units( handler, pairs, passthrough_keys, low_bit_keys, exclude_patterns="", discard_patterns="", glob_patterns=False, + include_mode=False, ): """Plan raw-copy and streamed work without loading tensor payloads.""" raw_units = [] @@ -691,7 +693,8 @@ def _build_resize_work_units( if matches(discards, names): claimed.update(key for _, key, _ in group_entries) continue - excluded = matches(excludes, names) + matched = matches(excludes, names) + preserve = matched != include_mode if layer_has_low_bit(block_keys, low_bit_keys): raw_entries = entries + related_passthrough @@ -699,7 +702,7 @@ def _build_resize_work_units( claimed.update(key for _, key, _ in raw_entries) continue - if excluded: + if preserve: stream_units.append({"kind": "preserve", "entries": group_entries}) claimed.update(key for _, key, _ in group_entries) elif layer_has_companions(block_keys): @@ -725,11 +728,12 @@ def _build_resize_work_units( if matches(discards, [key]): continue entry = (None, key, key) + matched = matches(excludes, [key]) if key in low_bit_keys: raw_units.append({"entries": [entry]}) else: stream_units.append({ - "kind": "preserve" if matches(excludes, [key]) else "copy", + "kind": "preserve" if matched != include_mode else "copy", "entries": [entry], }) claimed.add(key) @@ -902,6 +906,7 @@ def resize_lora_file( exclude_patterns: str = "", discard_patterns: str = "", glob_patterns: bool = False, + include_mode: bool = False, ) -> str: """Resize a LoRA through one bounded UEL stream and factor-space SVD.""" if min_rank < 1: @@ -925,7 +930,7 @@ def resize_lora_file( raw_units, stream_units, rank_counts, preserved_companion_groups = ( _build_resize_work_units( handler, pairs, passthrough_keys, low_bit_keys, - exclude_patterns, discard_patterns, glob_patterns, + exclude_patterns, discard_patterns, glob_patterns, include_mode, ) ) rank_text = _rank_summary(rank_counts) @@ -948,6 +953,9 @@ def resize_lora_file( ) metadata["ss_network_dim"] = "Dynamic" metadata["ss_network_alpha"] = "Dynamic" + metadata["modelutils_resize_filter_mode"] = ( + "include" if include_mode else "exclude" + ) output_path, _ = canonical_model_artifact_path("loras", output_filename) @@ -1230,7 +1238,7 @@ class LoRAResizeFixed(io.ComfyNode): ) @classmethod - def execute(cls, lora_name, new_rank, output_filename, save_dtype, device, force_clear_cache, exclude_patterns="", discard_patterns="", glob_patterns=False) -> io.NodeOutput: + def execute(cls, lora_name, new_rank, output_filename, save_dtype, device, force_clear_cache, exclude_patterns="", discard_patterns="", glob_patterns=False, include_mode=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] @@ -1240,6 +1248,7 @@ class LoRAResizeFixed(io.ComfyNode): exclude_patterns=exclude_patterns, discard_patterns=discard_patterns, glob_patterns=glob_patterns, + include_mode=include_mode, ) _, output_name = canonical_model_artifact_path("loras", output_filename) return io.NodeOutput(output_name) @@ -1273,7 +1282,7 @@ class LoRAResizeRatio(io.ComfyNode): ) @classmethod - 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: + def execute(cls, lora_name, max_rank, ratio, output_filename, save_dtype, device, force_clear_cache, exclude_patterns="", discard_patterns="", glob_patterns=False, include_mode=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] @@ -1283,6 +1292,7 @@ class LoRAResizeRatio(io.ComfyNode): exclude_patterns=exclude_patterns, discard_patterns=discard_patterns, glob_patterns=glob_patterns, + include_mode=include_mode, ) _, output_name = canonical_model_artifact_path("loras", output_filename) return io.NodeOutput(output_name) @@ -1318,7 +1328,7 @@ class LoRAResizeFrobenius(io.ComfyNode): ) @classmethod - 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: + 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, include_mode=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] @@ -1328,6 +1338,7 @@ class LoRAResizeFrobenius(io.ComfyNode): exclude_patterns=exclude_patterns, discard_patterns=discard_patterns, glob_patterns=glob_patterns, + include_mode=include_mode, min_rank=min_rank, ) _, output_name = canonical_model_artifact_path("loras", output_filename) @@ -1362,7 +1373,7 @@ class LoRAResizeCumulative(io.ComfyNode): ) @classmethod - 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: + def execute(cls, lora_name, max_rank, target, output_filename, save_dtype, device, force_clear_cache, exclude_patterns="", discard_patterns="", glob_patterns=False, include_mode=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] @@ -1372,6 +1383,7 @@ class LoRAResizeCumulative(io.ComfyNode): exclude_patterns=exclude_patterns, discard_patterns=discard_patterns, glob_patterns=glob_patterns, + include_mode=include_mode, ) _, output_name = canonical_model_artifact_path("loras", output_filename) return io.NodeOutput(output_name) @@ -1396,6 +1408,7 @@ def merge_loras_to_model( force_clear_cache: bool = False, include_1d_diffs: bool = False, return_report: bool = False, + include_mode: bool = False, ) -> str | tuple[str, str]: """ Merge multiple LoRAs into a base model and save the result directly. @@ -1416,6 +1429,7 @@ def merge_loras_to_model( lazy_load: Low memory mode: load tensors from disk on demand force_clear_cache: Clear CUDA cache after each layer return_report: Return the path and merge report instead of only the path + include_mode: Use skip_patterns_str to select matching base keys instead Returns: Path to saved merged model, optionally paired with the merge report @@ -1664,7 +1678,7 @@ def merge_loras_to_model( if planned_base_key.endswith(".weight"): core = extract_core_layer_base(planned_base_key) planned_low_rank = lora_lookup.get(core.replace(".", "_"), []) - skipped = any(pattern.search(planned_base_key) for pattern in skip_patterns) + skipped = any(pattern.search(planned_base_key) for pattern in skip_patterns) != include_mode entries = {} if planned_base_key not in base_low_bit_keys and not skipped: entries[-1] = [planned_base_key] @@ -1745,20 +1759,21 @@ def merge_loras_to_model( (pattern.pattern for pattern in skip_patterns if pattern.search(base_key)), None, ) - if skip_pattern is not None: + if (skip_pattern is not None) != include_mode: + skip_reason = "no include pattern matched" if include_mode else f"pattern={skip_pattern}" stats["skipped"] += 1 base_outcomes["OMITTED BY SKIP PATTERN"].append( - f"{base_key}: pattern={skip_pattern}" + f"{base_key}: {skip_reason}" ) for info, block_name, _, direct_key, _ in direct_contributions: record_group_event( info, block_name, "skipped", base_key, - f"target omitted by pattern {skip_pattern}; source={direct_key}", + f"target omitted: {skip_reason}; source={direct_key}", ) for info, block_name, block_keys in low_rank_contributions: record_group_event( info, block_name, "skipped", base_key, - f"target omitted by pattern {skip_pattern}; sources=" + f"target omitted: {skip_reason}; sources=" + ", ".join(sorted(layer_tensor_keys(block_keys).values())), ) pbar.update(1) @@ -2220,7 +2235,7 @@ class LoRAMergeToModel(io.ComfyNode): tooltip="Weight strength for LoRA 8"), # Settings io.String.Input("skip_patterns", default="", multiline=True, - tooltip="Regex patterns for layers to skip"), + tooltip="Regex patterns for base tensor keys to omit, or to include when Include Mode is enabled. Guarded low-bit base tensors are always preserved."), io.String.Input("output_filename", default="merged_model", tooltip="Output filename without extension, written under ComfyUI's diffusion-model directory."), io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16", tooltip="Data type used to save the model after applying the LoRA."), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda", tooltip="Device used for per-layer LoRA application; CUDA out-of-memory retries the affected layer on CPU."), @@ -2228,6 +2243,8 @@ class LoRAMergeToModel(io.ComfyNode): io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"), io.Boolean.Input("include_1d_diffs", default=False, tooltip="Apply 1D direct-diff tensors as FP32. Disabled preserves prior behavior."), + io.Boolean.Input("include_mode", default=False, + tooltip="Use Skip Patterns as an include-only filter. Only matching base tensors are processed and saved; an empty filter selects nothing. Guarded low-bit base tensors remain preserved."), ], outputs=[ io.AnyType.Output(display_name="output_path"), @@ -2241,7 +2258,7 @@ class LoRAMergeToModel(io.ComfyNode): lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4, lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8, skip_patterns, output_filename, save_dtype, device, lazy_load, force_clear_cache, - include_1d_diffs) -> io.NodeOutput: + include_1d_diffs, include_mode=False) -> io.NodeOutput: # Build LoRA list based on count count = int(lora_count) @@ -2270,6 +2287,7 @@ class LoRAMergeToModel(io.ComfyNode): force_clear_cache=force_clear_cache, include_1d_diffs=include_1d_diffs, return_report=True, + include_mode=include_mode, ) _, output_name = canonical_model_artifact_path("diffusion_models", output_filename) return io.NodeOutput(output_name, report) diff --git a/nodes/merger.py b/nodes/merger.py index 6db5040..fa23f0e 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -128,6 +128,7 @@ class MergerLogic: # Compile filter patterns glob_mode = recipe_params.get('glob_patterns', False) + include_mode = recipe_params.get('include_mode', False) exclude_patterns = _compile_patterns(recipe_params.get('exclude_patterns', ''), glob_mode=glob_mode) discard_patterns = _compile_patterns(recipe_params.get('discard_patterns', ''), glob_mode=glob_mode) @@ -216,7 +217,11 @@ class MergerLogic: model_type == "loras" and len(primary_handler.get_shape(key)) == 1 and not include_1d_diffs - ) or _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode) + ) or ( + not _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode) + if include_mode else + _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode) + ) requested_names = [primary_model_name] if not preserve_a: requested_names.extend( @@ -297,7 +302,11 @@ class MergerLogic: del cpu_tensor # Check exclude patterns - use Model A only, no merge - if preserve_model_a_1d or _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode): + if preserve_model_a_1d or ( + not _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode) + if include_mode else + _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode) + ): t = tensor_a.detach().to(target_dtype).cpu() writer.write_batch([(key, t)]) @@ -418,6 +427,7 @@ class CheckpointTwoMerger(io.ComfyNode): io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -430,7 +440,7 @@ class CheckpointTwoMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -443,6 +453,7 @@ class CheckpointTwoMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, } model_names = {"model_a": model_a, "model_b": model_b} @@ -484,6 +495,7 @@ class ModelTwoMerger(io.ComfyNode): io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -496,7 +508,7 @@ class ModelTwoMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -509,6 +521,7 @@ class ModelTwoMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, } model_names = {"model_a": model_a, "model_b": model_b} @@ -550,6 +563,7 @@ class TextEncoderTwoMerger(io.ComfyNode): io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -562,7 +576,7 @@ class TextEncoderTwoMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -575,6 +589,7 @@ class TextEncoderTwoMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, } model_names = {"model_a": model_a, "model_b": model_b} @@ -618,6 +633,7 @@ class LoRATwoMerger(io.ComfyNode): io.Boolean.Input("override_dtype", default=False, tooltip="Force merged non-1D tensors to save_dtype. Enabled 1D direct diffs remain FP32."), io.Boolean.Input("include_1d_diffs", default=False, tooltip="Merge 1D tensors. When disabled, preserve Model A's 1D tensors unchanged."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -631,7 +647,7 @@ class LoRATwoMerger(io.ComfyNode): gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, - include_1d_diffs: bool) -> io.NodeOutput: + include_1d_diffs: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -644,6 +660,7 @@ class LoRATwoMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "include_1d_diffs": include_1d_diffs, } @@ -686,6 +703,7 @@ class EmbeddingTwoMerger(io.ComfyNode): io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -698,7 +716,7 @@ class EmbeddingTwoMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -711,6 +729,7 @@ class EmbeddingTwoMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, } model_names = {"model_a": model_a, "model_b": model_b} @@ -755,6 +774,7 @@ class CheckpointThreeMerger(io.ComfyNode): io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -767,7 +787,7 @@ class CheckpointThreeMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_3_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -780,6 +800,7 @@ class CheckpointThreeMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, } model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} @@ -822,6 +843,7 @@ class ModelThreeMerger(io.ComfyNode): io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -834,7 +856,7 @@ class ModelThreeMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_3_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -847,6 +869,7 @@ class ModelThreeMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, } model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} @@ -889,6 +912,7 @@ class TextEncoderThreeMerger(io.ComfyNode): io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -901,7 +925,7 @@ class TextEncoderThreeMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_3_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -914,6 +938,7 @@ class TextEncoderThreeMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, } model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} @@ -958,6 +983,7 @@ class LoRAThreeMerger(io.ComfyNode): io.Boolean.Input("override_dtype", default=False, tooltip="Force merged non-1D tensors to save_dtype. Enabled 1D direct diffs remain FP32."), io.Boolean.Input("include_1d_diffs", default=False, tooltip="Merge 1D tensors. When disabled, preserve Model A's 1D tensors unchanged."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -971,7 +997,7 @@ class LoRAThreeMerger(io.ComfyNode): gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, - include_1d_diffs: bool) -> io.NodeOutput: + include_1d_diffs: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_3_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -984,6 +1010,7 @@ class LoRAThreeMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, "include_1d_diffs": include_1d_diffs, } @@ -1027,6 +1054,7 @@ class EmbeddingThreeMerger(io.ComfyNode): io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Exclude Patterns as a whitelist instead. Only matching tensors are merged; nonmatching tensors are preserved from Model A."), ], outputs=[ io.AnyType.Output(display_name="output_filename"), @@ -1039,7 +1067,7 @@ class EmbeddingThreeMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, include_mode: bool = False) -> io.NodeOutput: doc = load_documentation_from_file('merger_3_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -1052,6 +1080,7 @@ class EmbeddingThreeMerger(io.ComfyNode): "device": process_device, "dtype": torch.float32, "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, + "include_mode": include_mode, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, } model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} diff --git a/nodes/model_analysis.py b/nodes/model_analysis.py index d0052d8..02a7ef6 100644 --- a/nodes/model_analysis.py +++ b/nodes/model_analysis.py @@ -781,13 +781,14 @@ class ModelAnalysisLogic: params.get("exclude_patterns", ""), glob_mode=glob_mode, ) + include_mode = bool(params.get("include_mode", False)) excluded_a = { key for key in keys_a - if _matches_any_pattern(key, excluded_patterns, glob_mode=glob_mode) + if _matches_any_pattern(key, excluded_patterns, glob_mode=glob_mode) != include_mode } excluded_b = { key for key in keys_b - if _matches_any_pattern(key, excluded_patterns, glob_mode=glob_mode) + if _matches_any_pattern(key, excluded_patterns, glob_mode=glob_mode) != include_mode } keys_a.difference_update(excluded_a) keys_b.difference_update(excluded_b) @@ -1092,12 +1093,17 @@ def _common_inputs(model_type: str, *, alignment_control: bool): "exclude_patterns", default="", multiline=True, - tooltip="One pattern per line. Matching tensors are excluded from all comparison metrics and topology counts. Uses regex unless Glob Patterns is enabled.", + tooltip="One pattern per line. Matching tensors are excluded from all comparison metrics and topology counts, or are the only tensors analyzed when Include Mode is enabled. Uses regex unless Glob Patterns is enabled.", ), io.Boolean.Input( "glob_patterns", default=False, - tooltip="Interpret exclusion entries as shell-style glob patterns instead of regular expressions.", + tooltip="Interpret layer-filter entries as shell-style glob patterns instead of regular expressions.", + ), + io.Boolean.Input( + "include_mode", + default=False, + tooltip="Use Exclude Patterns as an include-only filter. Analyze only matching tensors; an empty filter selects nothing.", ), ]) return inputs diff --git a/nodes/text_encoder_extract.py b/nodes/text_encoder_extract.py index 6f3b41c..d2c1d12 100644 --- a/nodes/text_encoder_extract.py +++ b/nodes/text_encoder_extract.py @@ -81,6 +81,7 @@ def extract_te_from_files( force_clear_cache: bool = True, glob_skip_patterns: bool = False, knee_probe_offset: int = 32, + include_mode: bool = False, ) -> None: """ Extract LoRA/DoRA from difference between two Text Encoder models, writing incrementally to disk. @@ -108,7 +109,7 @@ def extract_te_from_files( work_units = [ (key, key, key if key in keys_b else None) for key in weight_keys - if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) + if (not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)) and not (key not in keys_b and mismatch_mode == "skip") ] pbar = comfy.utils.ProgressBar(len(work_units)) @@ -119,7 +120,7 @@ def extract_te_from_files( weight_diff = None layer_device = device - if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns): + if (_matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) if not include_mode else not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)): return "skipped", None # Load tensors with pinned memory for CUDA @@ -300,6 +301,7 @@ def _get_common_inputs(): io.Boolean.Input("glob_skip_patterns", default=False, tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). " "When False (default), patterns are Python regex matched as substrings."), + io.Boolean.Input("include_mode", default=False, tooltip="Use Skip Patterns as a whitelist instead. Only matching layers are extracted; an empty whitelist extracts nothing."), ] @@ -329,7 +331,7 @@ class TextEncoderLoRAExtractFixed(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -341,7 +343,7 @@ class TextEncoderLoRAExtractFixed(io.ComfyNode): linear_max_rank=linear_dim, conv_max_rank=conv_dim, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, svd_niter=svd_niter, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -370,7 +372,7 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -382,7 +384,7 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode): linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=probe_offset ) return io.NodeOutput(output_name) @@ -412,7 +414,7 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -424,7 +426,7 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode): linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=probe_offset ) return io.NodeOutput(output_name) @@ -454,7 +456,7 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode): def execute(cls, model_a, model_b, knee_method, knee_probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -467,7 +469,7 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode): clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset + glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -496,7 +498,7 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -508,7 +510,7 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode): linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=probe_offset ) return io.NodeOutput(output_name) @@ -540,7 +542,7 @@ class TextEncoderDoRAExtractFixed(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -552,7 +554,7 @@ class TextEncoderDoRAExtractFixed(io.ComfyNode): linear_max_rank=linear_dim, conv_max_rank=conv_dim, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, svd_niter=svd_niter, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -581,7 +583,7 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -593,7 +595,7 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode): linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=probe_offset ) return io.NodeOutput(output_name) @@ -623,7 +625,7 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -635,7 +637,7 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode): linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=probe_offset ) return io.NodeOutput(output_name) @@ -665,7 +667,7 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode): def execute(cls, model_a, model_b, knee_method, knee_probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -678,7 +680,7 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode): clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset + glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset, include_mode=include_mode ) return io.NodeOutput(output_name) @@ -707,7 +709,7 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode): @classmethod def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, - save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_mode=False) -> io.NodeOutput: model_a_path = folder_paths.get_full_path_or_raise("text_encoders", model_a) model_b_path = folder_paths.get_full_path_or_raise("text_encoders", model_b) @@ -719,7 +721,7 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode): linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, + lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns, include_mode=include_mode, knee_probe_offset=probe_offset ) return io.NodeOutput(output_name) diff --git a/tests/test_adaptive_extraction_pipeline.py b/tests/test_adaptive_extraction_pipeline.py index 533e12d..076c0af 100644 --- a/tests/test_adaptive_extraction_pipeline.py +++ b/tests/test_adaptive_extraction_pipeline.py @@ -217,3 +217,38 @@ def test_lora_extraction_uses_only_async_uel_and_releases_sources(monkeypatch, t "diffusion_model.layer.lora_A.weight", "diffusion_model.layer.lora_B.weight", ] + + +@pytest.mark.parametrize("module_name, function_name, extra", [ + ("nodes.lora_extract_svd", "extract_lora_from_files", {"include_1d_diffs": False}), + ("nodes.dora_extract_wd", "extract_dora_from_files", {}), + ("nodes.dora_learned_wd", "extract_dora_learned_from_files", {"optimize_iters": 0}), + ("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": False}), + ("nodes.text_encoder_extract", "extract_te_from_files", {"is_dora": True}), +]) +def test_extraction_include_mode_selects_only_matching_layers(monkeypatch, tmp_path, module_name, function_name, extra): + module = _load_module(module_name) + first, second = tmp_path / "first.safetensors", tmp_path / "second.safetensors" + match, other = "model.encoder.match.weight", "model.encoder.other.weight" + _write_uel(first, {match: torch.eye(2), other: torch.eye(2)}) + _write_uel(second, {match: torch.zeros(2, 2), other: torch.zeros(2, 2)}) + monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(module, "cleanup_after_operation", lambda: None) + common = dict(linear_max_rank=1, conv_max_rank=1, force_clear_cache=False, + skip_patterns_str="match", include_mode=True, **extra) + selected = tmp_path / "selected.safetensors" + getattr(module, function_name)(str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(selected), **common) + with MemoryEfficientSafeOpen(str(selected), low_memory=True) as result: + keys = list(result.keys()) + assert keys + assert all("match" in key for key in keys) + empty = tmp_path / "empty.safetensors" + common["skip_patterns_str"] = "" + getattr(module, function_name)(str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(empty), **common) + with MemoryEfficientSafeOpen(str(empty), low_memory=True) as result: + assert list(result.keys()) == [] + globbed = tmp_path / "globbed.safetensors" + common.update(skip_patterns_str="*.match.weight", glob_skip_patterns=True) + getattr(module, function_name)(str(first), str(second), "fixed", 1, 1, "cpu", "fp32", str(globbed), **common) + with MemoryEfficientSafeOpen(str(globbed), low_memory=True) as result: + assert list(result.keys()) == keys diff --git a/tests/test_consensus_merger.py b/tests/test_consensus_merger.py index 31fcd20..3a82d7f 100644 --- a/tests/test_consensus_merger.py +++ b/tests/test_consensus_merger.py @@ -651,9 +651,10 @@ def test_all_cwb_schemas_and_compact_control_contract(cwb): if node is cwb.CWBLoRAMultiMerger: assert ids[:4] == ["execution_mode", "lora_count", "lora_1", "lora_2"] assert schema.inputs[1].default == "2" - assert ids[-2:] == [ + assert ids[-3:] == [ "include_1d_diffs", "counterfactual_weight_sweep", + "include_mode", ] continue if node in (cwb.CWBEmbeddingSelfCoalesce, cwb.CWBEmbeddingMultiMerger): @@ -667,10 +668,11 @@ def test_all_cwb_schemas_and_compact_control_contract(cwb): if node.INPUT_COUNT == 3: assert ids[3] == "model_c" if node.LORA_MODE: - assert schema.inputs[-1].id == "include_1d_diffs" - assert schema.inputs[-1].default is False + assert schema.inputs[-2].id == "include_1d_diffs" else: assert "include_1d_diffs" not in ids + assert schema.inputs[-1].id == "include_mode" + assert schema.inputs[-1].default is False def test_use_case_preset_names_match_complete_settings(cwb): @@ -1261,6 +1263,37 @@ def test_secondary_shape_conflict_skip_preserves_earliest_source(monkeypatch, tm torch.testing.assert_close(tensor, torch.tensor([2.0, 3.0])) +@pytest.mark.parametrize("lora_mode", [False, True]) +@pytest.mark.parametrize("patterns, glob_patterns, selected", [ + (r"\.selected", False, True), + ("*.selected*", True, True), + ("", False, False), + ("no_match", False, False), +]) +def test_include_filter_preserves_nonmatches_and_discards_first( + monkeypatch, tmp_path, cwb, lora_mode, patterns, glob_patterns, selected, +): + paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")} + for name, value in (("a", 2.0), ("b", 6.0)): + save_file({ + f"diffusion_model.{layer}.diff": torch.full((2, 2), value) + for layer in ("selected", "other", "selected_discard") + }, paths[name]) + _patch_io(monkeypatch, cwb, tmp_path, paths) + category = "loras" if lora_mode else "diffusion_models" + result = cwb.ConsensusMergerLogic.execute( + ["a", "b"], category, + _params("include_filter", exclude_patterns=patterns, include_mode=True, + glob_patterns=glob_patterns, discard_patterns="*discard*" if glob_patterns else "discard", + power_alpha=0.0), + lora_mode=lora_mode, + ) + tensors = load_file(str(_result_path(tmp_path, category, result))) + assert set(tensors) == {"diffusion_model.selected.diff", "diffusion_model.other.diff"} + torch.testing.assert_close(tensors["diffusion_model.selected.diff"], torch.full((2, 2), 4.0 if selected else 2.0)) + torch.testing.assert_close(tensors["diffusion_model.other.diff"], torch.full((2, 2), 2.0)) + + def test_secondary_only_filters_and_low_bit_preservation(monkeypatch, tmp_path, cwb): a = tmp_path / "filter_a.safetensors" b = tmp_path / "filter_b.safetensors" diff --git a/tests/test_cwb_delta_lora_merger.py b/tests/test_cwb_delta_lora_merger.py index 7c09d98..e05541b 100644 --- a/tests/test_cwb_delta_lora_merger.py +++ b/tests/test_cwb_delta_lora_merger.py @@ -186,6 +186,36 @@ def test_streamed_merge_normalizes_alpha_once_and_supports_mixed_rank( assert metadata["output_representation"] == "full_direct_difference" +@pytest.mark.parametrize("patterns, glob_patterns, selected", [ + (r"\.selected", False, True), + ("*.selected*", True, True), + ("", False, False), + ("no_match", False, False), +]) +def test_include_filter_preserves_nonmatches_and_discards_first( + monkeypatch, tmp_path, delta_cwb, patterns, glob_patterns, selected, +): + paths = [tmp_path / f"{name}.safetensors" for name in ("a", "b")] + for path, value in zip(paths, (2.0, 6.0)): + _write(path, { + f"diffusion_model.{layer}.diff": torch.full((2, 2), value) + for layer in ("selected", "other", "selected_discard") + }) + monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None) + filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute( + [str(path) for path in paths], [path.name for path in paths], _settings(delta_cwb), + "balanced_mean", False, "skip", "include_filter", torch.float32, + "full_difference", 384, "cpu", patterns, + "*discard*" if glob_patterns else "discard", glob_patterns, False, include_mode=True, + ) + tensors = _read(tmp_path / "loras" / filename) + assert set(tensors) == {"diffusion_model.selected.diff", "diffusion_model.other.diff"} + torch.testing.assert_close(tensors["diffusion_model.selected.diff"], torch.full((2, 2), 4.0 if selected else 2.0)) + torch.testing.assert_close(tensors["diffusion_model.other.diff"], torch.full((2, 2), 2.0)) + + def test_existing_diff_and_1d_output_remain_fp32(monkeypatch, tmp_path, delta_cwb): first = tmp_path / "diff_first.safetensors" second = tmp_path / "diff_second.safetensors" @@ -350,7 +380,7 @@ def test_exact_node_schema_contract(delta_cwb): "lora_1", "lora_2", "cwb_preset", "cwb_config", "mismatch_mode", "output_filename", "save_dtype", "output_representation", "output_rank", "process_device", "exclude_patterns", - "discard_patterns", "glob_patterns", "force_clear_cache", + "discard_patterns", "glob_patterns", "force_clear_cache", "include_mode", ] multi_ids = [value.id for value in schemas[2].inputs] assert multi_ids[:9] == ["lora_count", *[f"lora_{index}" for index in range(1, 9)]] @@ -360,8 +390,8 @@ def test_exact_node_schema_contract(delta_cwb): assert schemas[0].inputs[7].default == "full_difference" assert schemas[0].inputs[8].default == 384 assert schemas[0].inputs[9].default == "cuda" - assert schemas[0].inputs[-2].default is False - assert schemas[0].inputs[-1].default is True + assert schemas[0].inputs[-2].default is True + assert schemas[0].inputs[-1].default is False assert "agreement" in schemas[0].inputs[2].tooltip assert "Only a failed CUDA layer retries on CPU" in schemas[0].inputs[9].tooltip diff --git a/tests/test_dtype_conversion.py b/tests/test_dtype_conversion.py index deb5325..32f806f 100644 --- a/tests/test_dtype_conversion.py +++ b/tests/test_dtype_conversion.py @@ -68,7 +68,9 @@ def test_schema_is_diffusion_model_scoped(conversion): "reference_model", "exclude_patterns", "output_filename", + "include_mode", ] + assert schema.inputs[-1].default is False assert schema.inputs[1].options == ["fp32", "fp16", "bf16"] assert [item.display_name for item in schema.outputs] == [ "output_path", @@ -152,6 +154,44 @@ def test_converts_floating_and_preserves_excluded_and_nonfloating( assert "Non-floating tensors preserved: 1" in report +@pytest.mark.parametrize("patterns, selected", [ + (r"^blocks\.0\.weight$", {"blocks.0.weight"}), + (r"^blocks\.0\.weight$" + "\n" + r"^blocks\.0\.norm\.weight$", {"blocks.0.weight", "blocks.0.norm.weight"}), + ("", set()), + ("no_such_layer", set()), +]) +@pytest.mark.parametrize("use_reference", [False, True]) +def test_include_filter_only_converts_matches( + monkeypatch, tmp_path, conversion, patterns, selected, use_reference, +): + source = tmp_path / "source.safetensors" + reference = tmp_path / "reference.safetensors" + tensors = { + "blocks.0.weight": torch.arange(4, dtype=torch.float32), + "blocks.0.norm.weight": torch.ones(3, dtype=torch.float32), + "position_ids": torch.tensor([1, 2], dtype=torch.int64), + } + _write_uel(source, tensors) + _write_uel(reference, {"blocks.0.weight": tensors["blocks.0.weight"].to(torch.bfloat16)}) + _patch_runtime(monkeypatch, tmp_path, conversion, source) + monkeypatch.setattr( + conversion.folder_paths, "get_full_path_or_raise", + lambda _, name: str(reference if name == "reference" else source), + ) + + output_name, _ = conversion.DiffusionModelDtypeConversion.execute( + source.name, "fp16", "reference" if use_reference else "None", + patterns, "include_converted", include_mode=True, + ) + result, _ = _read_uel(tmp_path / "diffusion_models" / output_name) + for key, tensor in tensors.items(): + expected_dtype = tensor.dtype + if key in selected: + expected_dtype = torch.bfloat16 if use_reference and key == "blocks.0.weight" else torch.float16 + assert result[key].dtype == expected_dtype + torch.testing.assert_close(result[key].to(tensor.dtype), tensor) + + def test_uses_bounded_async_stream_and_releases_every_tensor( monkeypatch, tmp_path, conversion ): diff --git a/tests/test_groups.toml b/tests/test_groups.toml index 7d28a0e..0e8ec0d 100644 --- a/tests/test_groups.toml +++ b/tests/test_groups.toml @@ -48,6 +48,10 @@ python_tests = [ 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"] +python_tests = ["tests/test_merger_filters.py"] + [groups.dtype_conversion] paths = ["nodes/dtype_conversion.py"] python_tests = ["tests/test_dtype_conversion.py"] diff --git a/tests/test_lodestone_merger.py b/tests/test_lodestone_merger.py index bd4b47e..8da6f1c 100644 --- a/tests/test_lodestone_merger.py +++ b/tests/test_lodestone_merger.py @@ -129,6 +129,35 @@ def test_streamed_merge_writes_only_alpha_normalized_full_delta(monkeypatch, tmp torch.testing.assert_close(tensors["diffusion_model.layer.diff"], torch.diag(torch.tensor([1.0, 4.0]))) +@pytest.mark.parametrize("patterns, glob_patterns, selected", [ + (r"\.selected", False, True), + ("*.selected*", True, True), + ("", False, False), + ("no_match", False, False), +]) +def test_include_filter_preserves_nonmatches_and_discards_first( + monkeypatch, tmp_path, lodestone, patterns, glob_patterns, selected, +): + paths = [tmp_path / f"{name}.safetensors" for name in ("a", "b")] + for path, value in zip(paths, (2.0, 6.0)): + _write(path, { + f"diffusion_model.{layer}.diff": torch.full((2, 2), value) + for layer in ("selected", "other", "selected_discard") + }) + monkeypatch.setattr(lodestone.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(lodestone, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(lodestone, "cleanup_after_operation", lambda: None) + output = lodestone.merge_lodestone_loras( + [str(path) for path in paths], "sum", "cpu", torch.float32, "include_filter", + exclude_patterns=patterns, glob_patterns=glob_patterns, include_mode=True, + discard_patterns="*discard*" if glob_patterns else "discard", verbose=False, + ) + tensors = _read(output) + assert set(tensors) == {"diffusion_model.selected.diff", "diffusion_model.other.diff"} + torch.testing.assert_close(tensors["diffusion_model.selected.diff"], torch.full((2, 2), 8.0 if selected else 2.0)) + torch.testing.assert_close(tensors["diffusion_model.other.diff"], torch.full((2, 2), 2.0)) + + def test_dedicated_nodes_do_not_modify_factor_merger_schema(lodestone): standard = _load_module("nodes.lora_merger") assert [node.define_schema().node_id for node in lodestone.LODESTONE_MERGER_NODES] == [ diff --git a/tests/test_lora_diff_dtype_contract.py b/tests/test_lora_diff_dtype_contract.py index e4ec4ea..eaf9a0c 100644 --- a/tests/test_lora_diff_dtype_contract.py +++ b/tests/test_lora_diff_dtype_contract.py @@ -417,18 +417,68 @@ def test_merge_to_model_report_lists_skipped_tensor_and_pattern( assert "[SKIPPED BY PATTERN] (1)" in report +@pytest.mark.parametrize("patterns, selected", [ + (r"selected\.weight$", True), + ("", False), + ("no_match", False), +]) +def test_merge_to_model_include_filter_limits_stream_and_output( + monkeypatch, tmp_path, modules, patterns, selected, +): + resize, _ = modules + _patch_output(monkeypatch, resize, tmp_path) + base = tmp_path / "include_base.safetensors" + adapter = tmp_path / "include_adapter.safetensors" + save_file({ + f"model.diffusion_model.{name}.weight": torch.zeros((2, 2)) + for name in ("selected", "other") + }, str(base)) + save_file({ + f"diffusion_model.{name}.{factor}": torch.ones(shape) + for name in ("selected", "other") + for factor, shape in (("lora_A.weight", (1, 2)), ("lora_B.weight", (2, 1))) + }, str(adapter)) + planned_keys = [] + real_stream = resize.stream_work_units + + def capture_units(handlers, units, **kwargs): + units = list(units) + planned_keys.extend(key for _, entries in units for keys in entries.values() for key in keys) + return real_stream(handlers, units, **kwargs) + + monkeypatch.setattr(resize, "stream_work_units", capture_units) + path, report = resize.merge_loras_to_model( + [str(adapter)], [1.0], str(base), "cpu", torch.float32, + "include_result", skip_patterns_str=patterns, verbose=False, + return_report=True, include_mode=True, + ) + tensors = load_file(path) + assert set(tensors) == ({"model.diffusion_model.selected.weight"} if selected else set()) + if selected: + torch.testing.assert_close(tensors["model.diffusion_model.selected.weight"], torch.ones((2, 2))) + assert len(planned_keys) == 3 + assert all("selected" in key for key in planned_keys) + else: + assert planned_keys == [] + assert "no include pattern matched" in report + + +@pytest.mark.parametrize("include_mode", [False, True]) def test_merge_to_model_node_returns_path_blank_line_report( - monkeypatch, modules + monkeypatch, modules, include_mode, ): resize, _ = modules monkeypatch.setattr( resize.folder_paths, "get_full_path_or_raise", lambda category, name: f"{category}/{name}", ) - monkeypatch.setattr( - resize, "merge_loras_to_model", - lambda **kwargs: ("saved/model.safetensors", "SUCCESS\nDETAILS"), - ) + received = {} + + def merge(**kwargs): + received.update(kwargs) + return "saved/model.safetensors", "SUCCESS\nDETAILS" + + monkeypatch.setattr(resize, "merge_loras_to_model", merge) result = resize.LoRAMergeToModel.execute( "base.safetensors", "1", @@ -436,9 +486,11 @@ def test_merge_to_model_node_returns_path_blank_line_report( "None", 1.0, "None", 1.0, "None", 1.0, "None", 1.0, "None", 1.0, "None", 1.0, "None", 1.0, "", "merged", "fp32", "cpu", True, False, False, + include_mode=include_mode, ) assert result.result == ("merged.safetensors", "SUCCESS\nDETAILS") + assert received["include_mode"] is include_mode def test_mixed_parser_scans_every_key_and_retains_unpaired(modules): @@ -486,7 +538,10 @@ def test_include_1d_switches_are_appended_and_default_off(modules, generic_modul for node_class in node_classes: final_input = node_class.define_schema().inputs[-1] - assert final_input.id == "include_1d_diffs" + expected = "include_mode" if node_class in { + generic.LoRATwoMerger, generic.LoRAThreeMerger, resize.LoRAMergeToModel + } else "include_1d_diffs" + assert final_input.id == expected assert final_input.default is False outputs = resize.LoRAMergeToModel.define_schema().outputs diff --git a/tests/test_lora_resize_factor_space.py b/tests/test_lora_resize_factor_space.py index d2875ae..3bc4b3b 100644 --- a/tests/test_lora_resize_factor_space.py +++ b/tests/test_lora_resize_factor_space.py @@ -88,6 +88,55 @@ def test_resize_filters_preserve_and_discard_whole_layers(monkeypatch, tmp_path, assert actual["diffusion_model.resize.lora_A.weight"].shape[0] == 1 +@pytest.mark.parametrize("glob", [False, True]) +def test_resize_include_mode_only_resizes_matching_layers(monkeypatch, tmp_path, resize, glob): + source = tmp_path / "include_filters.safetensors" + tensors = {} + for name in ("include", "discard", "skip"): + 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["standalone.include"] = torch.ones(1) + tensors["standalone.skip"] = 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) + + output = resize.resize_lora_file( + str(source), 1, None, None, "cpu", torch.bfloat16, "included", verbose=False, + exclude_patterns="*include*\n*discard*\nstandalone.include" if glob else r"\.(include|discard)\." + "\n^standalone\\.include$", + discard_patterns="*discard*" if glob else r"\.discard\.", + glob_patterns=glob, + include_mode=True, + ) + + actual = _read_uel(output) + assert actual["diffusion_model.include.lora_A.weight"].shape[0] == 1 + assert torch.equal(actual["standalone.include"], tensors["standalone.include"]) + assert not any(".discard." in key for key in actual) + for key in tensors: + if ".skip." in key or key == "standalone.skip": + assert torch.equal(actual[key], tensors[key]) + + +def test_resize_include_mode_empty_patterns_preserves_all_tensors(monkeypatch, tmp_path, resize): + source = tmp_path / "empty_include.safetensors" + tensors = {"standalone": torch.ones(2)} + _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) + + output = resize.resize_lora_file( + str(source), 1, None, None, "cpu", torch.float32, "empty", verbose=False, + include_mode=True, + ) + + actual = _read_uel(output) + assert torch.equal(actual["standalone"], tensors["standalone"]) + + @pytest.mark.parametrize("node_name,extra", [ ("LoRAResizeFixed", {"new_rank": 1}), ("LoRAResizeRatio", {"max_rank": 1, "ratio": 2.0}), @@ -100,16 +149,19 @@ def test_all_resize_nodes_forward_filters(monkeypatch, tmp_path, resize, node_na assert controls["exclude_patterns"].default == "" assert controls["discard_patterns"].default == "" assert controls["glob_patterns"].default is False + assert controls["include_mode"].default is False + assert node.define_schema().inputs[-1].id == "include_mode" 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) + glob_patterns=True, include_mode=True, **extra) assert calls[0]["exclude_patterns"] == "keep*" assert calls[0]["discard_patterns"] == "drop*" assert calls[0]["glob_patterns"] is True + assert calls[0]["include_mode"] is True def test_resize_invalid_filter_fails_before_writer(monkeypatch, tmp_path, resize): diff --git a/tests/test_merger_filters.py b/tests/test_merger_filters.py new file mode 100644 index 0000000..6bd8c4e --- /dev/null +++ b/tests/test_merger_filters.py @@ -0,0 +1,70 @@ +import importlib +import sys +import types +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file, save_file + + +REPO_ROOT = Path(__file__).resolve().parents[1] +PACKAGE_NAME = "modelutils_merger_filter_tests" + + +def _module(name): + if PACKAGE_NAME not in sys.modules: + package = types.ModuleType(PACKAGE_NAME) + package.__path__ = [str(REPO_ROOT)] + sys.modules[PACKAGE_NAME] = package + return importlib.import_module(f"{PACKAGE_NAME}.{name}") + + +@pytest.fixture(scope="module") +def modules(): + return _module("nodes.merger"), _module("nodes.merger_ops") + + +def _params(name, *, patterns="", include=False, discard="", glob=False): + return { + "mismatch_mode": "skip", "alignment_mode": "pad/crop", "alpha": 0.5, + "beta": 0.0, "gamma": 0.5, "delta": 2.0, "epsilon": 0.01, + "zeta": 0.0, "seed": 0, "output_filename": name, "save_dtype": "fp32", + "override_dtype": True, "device": "cpu", "dtype": torch.float32, + "exclude_patterns": patterns, "include_mode": include, + "discard_patterns": discard, "glob_patterns": glob, "lazy_load": True, + "force_clear_cache": False, + } + + +def _run(monkeypatch, tmp_path, merger, operations, params): + paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")} + save_file({"match.weight": torch.tensor([2.0]), "other.weight": torch.tensor([4.0])}, paths["a"]) + save_file({"match.weight": torch.tensor([6.0]), "other.weight": torch.tensor([8.0])}, paths["b"]) + monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name)) + monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)]) + monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None) + params.update({"model_a": "a", "model_b": "b"}) + merger.MergerLogic.execute_merge({"model_a": "a", "model_b": "b"}, "Weight-Sum", operations.TWO_MODEL_MODES, params, "checkpoints") + return load_file(str(tmp_path / "checkpoints" / f"{params['output_filename']}.safetensors")) + + +def test_include_filter_controls_preflight_and_streaming(monkeypatch, tmp_path, modules): + merger, operations = modules + selected = _run(monkeypatch, tmp_path, merger, operations, _params("selected", patterns="match", include=True)) + assert selected["match.weight"].item() == 4.0 + assert selected["other.weight"].item() == 4.0 + + empty = _run(monkeypatch, tmp_path, merger, operations, _params("empty", include=True)) + assert empty["match.weight"].item() == 2.0 + assert empty["other.weight"].item() == 4.0 + + discarded = _run(monkeypatch, tmp_path, merger, operations, _params("discarded", patterns="match", include=True, discard="match")) + assert "match.weight" not in discarded + assert discarded["other.weight"].item() == 4.0 + + globbed = _run(monkeypatch, tmp_path, merger, operations, _params("globbed", patterns="*.weight", include=True, glob=True)) + assert globbed["match.weight"].item() == 4.0 + assert globbed["other.weight"].item() == 6.0 diff --git a/tests/test_model_analysis.py b/tests/test_model_analysis.py index 145f86b..a9d5326 100644 --- a/tests/test_model_analysis.py +++ b/tests/test_model_analysis.py @@ -53,6 +53,8 @@ def test_five_node_schema_contract(analysis): assert schema.inputs[0].options == ["ANALYZE", "DOCUMENTATION ONLY"] assert "lazy_load" not in [item.id for item in schema.inputs] input_ids = [item.id for item in schema.inputs] + assert input_ids[-3:] == ["exclude_patterns", "glob_patterns", "include_mode"] + assert schema.inputs[-1].default is False if node in (analysis.LoRAModelAnalysis, analysis.EmbeddingModelAnalysis): assert input_ids[3] == "cwb_similarity_alignment" else: @@ -227,7 +229,14 @@ def test_lora_cross_format_keys_compare_as_one_logical_pair(monkeypatch, tmp_pat assert all(len(call[1]) == 2 for call in async_calls) -def test_excluded_lora_blocks_do_not_enter_metrics(monkeypatch, tmp_path, analysis): +@pytest.mark.parametrize("patterns, glob_patterns, include_mode", [ + (r"^diffusion_model\.blocks\.4\.", False, False), + (r"^diffusion_model\.blocks\.3\.", False, True), + ("diffusion_model.blocks.3.*", True, True), +]) +def test_excluded_lora_blocks_do_not_enter_metrics( + monkeypatch, tmp_path, analysis, patterns, glob_patterns, include_mode, +): path_a = tmp_path / "lora_a.safetensors" path_b = tmp_path / "lora_b.safetensors" tensors_a = {} @@ -247,7 +256,7 @@ def test_excluded_lora_blocks_do_not_enter_metrics(monkeypatch, tmp_path, analys comparison, cwb_report = analysis.ModelAnalysisLogic.execute( "a", "b", "loras", - _params(exclude_patterns=r"^diffusion_model\.blocks\.4\."), + _params(exclude_patterns=patterns, glob_patterns=glob_patterns, include_mode=include_mode), lora_mode=True, ) @@ -262,6 +271,40 @@ def test_excluded_lora_blocks_do_not_enter_metrics(monkeypatch, tmp_path, analys assert "blocks_4_" not in cwb_report +@pytest.mark.parametrize("patterns, glob_patterns, expected_keys", [ + (r"^blocks\.0\.", False, ["blocks.0.weight"]), + ("blocks.0.*", True, ["blocks.0.weight"]), + ("", False, []), + ("no_such_layer", True, []), +]) +def test_include_filter_limits_analysis_inventory( + monkeypatch, tmp_path, analysis, patterns, glob_patterns, expected_keys, +): + paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")} + for name, value in (("a", 1.0), ("b", 2.0)): + save_file({ + "blocks.0.weight": torch.full((2, 2), value), + "blocks.1.weight": torch.full((2, 2), value), + f"only_{name}.weight": torch.full((2, 2), value), + }, paths[name]) + monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name]) + monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None) + monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None) + streamed = [] + original_stream = analysis.MemoryEfficientSafeOpen.async_stream + + def stream(handler, keys, *args, **kwargs): + streamed.extend(keys) + return original_stream(handler, keys, *args, **kwargs) + + monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", stream) + analysis.ModelAnalysisLogic.execute( + "a", "b", "diffusion_models", + _params(exclude_patterns=patterns, glob_patterns=glob_patterns, include_mode=True), + ) + assert sorted(streamed) == sorted(expected_keys * 2) + + def test_embedding_similarity_alignment_is_explicit_opt_in(analysis): a = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) b = torch.tensor([[0.0, 1.0], [1.0, 0.0]]) diff --git a/tests/test_schema_tooltips.py b/tests/test_schema_tooltips.py index e708820..1d44991 100644 --- a/tests/test_schema_tooltips.py +++ b/tests/test_schema_tooltips.py @@ -1,6 +1,13 @@ import ast +import importlib +import inspect +import sys +import types from pathlib import Path +import pytest +from comfy_api.latest import io + NODE_DIR = Path(__file__).resolve().parents[1] / "nodes" @@ -40,3 +47,102 @@ def test_every_schema_input_has_a_direct_static_tooltip(): assert not missing, "Inputs without tooltips:\n" + "\n".join(missing) assert not non_static, "Tooltips must be direct non-empty string literals:\n" + "\n".join(non_static) + + +FILTER_MODULES = [ + "merger", "consensus_merger", "cwb_delta_lora_merger", "lodestone_merger", + "lora_resize", "dtype_conversion", "model_analysis", + "lora_extract_svd", "dora_extract_wd", "dora_learned_wd", "text_encoder_extract", +] + + +def test_filter_schema_audit_covers_all_direct_filter_definitions(): + discovered = set() + for path in NODE_DIR.glob("*.py"): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for call in ast.walk(tree): + if ( + isinstance(call, ast.Call) + and isinstance(call.func, ast.Attribute) + and call.func.attr == "Input" + and call.args + and isinstance(call.args[0], ast.Constant) + and call.args[0].value in {"exclude_patterns", "skip_patterns"} + ): + discovered.add(path.stem) + assert discovered <= set(FILTER_MODULES), f"Unaudited filter modules: {discovered - set(FILTER_MODULES)}" + + +def _load_filter_module(module_name): + package_name = "modelutils_filter_schema_tests" + if package_name not in sys.modules: + package = types.ModuleType(package_name) + package.__path__ = [str(NODE_DIR.parent)] + sys.modules[package_name] = package + return importlib.import_module(f"{package_name}.nodes.{module_name}") + + +@pytest.mark.parametrize("module_name", FILTER_MODULES) +def test_layer_filter_schemas_append_disabled_include_toggle(module_name): + module = _load_filter_module(module_name) + checked = [] + for name, node in vars(module).items(): + if name.startswith("_") or not isinstance(node, type): + continue + if node.__module__ != module.__name__ or not issubclass(node, io.ComfyNode): + continue + schema = node.define_schema() + ids = [item.id for item in schema.inputs] + if not {"exclude_patterns", "skip_patterns"}.intersection(ids): + continue + assert ids.count("include_mode") == 1, schema.node_id + assert ids[-1] == "include_mode", schema.node_id + assert isinstance(schema.inputs[-1], io.Boolean.Input), schema.node_id + assert schema.inputs[-1].default is False, schema.node_id + checked.append(schema.node_id) + assert checked, f"No filter nodes checked in {module_name}" + + +@pytest.mark.parametrize("module_name, helper_name, expected_count", [ + ("lora_extract_svd", "extract_lora_from_files", 5), + ("dora_extract_wd", "extract_dora_from_files", 5), + ("dora_learned_wd", "extract_dora_learned_from_files", 5), + ("text_encoder_extract", "extract_te_from_files", 10), +]) +@pytest.mark.parametrize("include_mode", [False, True]) +def test_all_extraction_nodes_forward_include_mode_without_shifting_parameters( + monkeypatch, module_name, helper_name, expected_count, include_mode, +): + module = _load_filter_module(module_name) + signature = inspect.signature(getattr(module, helper_name)) + assert list(signature.parameters)[-1] == "include_mode" + captured = [] + + def extract(*args, **kwargs): + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + captured.append(bound.arguments) + + monkeypatch.setattr(module, helper_name, extract) + monkeypatch.setattr(module.folder_paths, "get_full_path_or_raise", lambda category, name: name) + monkeypatch.setattr(module, "_build_lora_output_path", lambda name: (name, name)) + for name, node in vars(module).items(): + if name.startswith("_") or not isinstance(node, type): + continue + if node.__module__ != module.__name__ or not issubclass(node, io.ComfyNode): + continue + schema = node.define_schema() + values = {} + for widget in schema.inputs: + value = getattr(widget, "default", None) + if value is None: + options = getattr(widget, "options", []) + value = options[0] if options else widget.id + values[widget.id] = value + values["include_mode"] = include_mode + values["skip_patterns"] = "selected_layer" + node.execute(**values) + assert captured[-1]["include_mode"] is include_mode, schema.node_id + assert captured[-1]["skip_patterns_str"] == "selected_layer", schema.node_id + assert captured[-1]["knee_probe_offset"] == values.get("knee_probe_offset", 32), schema.node_id + assert len(captured) == expected_count