Compare commits

..
Author SHA1 Message Date
silveroxides 0a7578ce73 feat(analysis): add model similarity nodes
Separate CWB-derived diagnostics from merge weighting while streaming each logical layer once.
2026-08-04 15:52:42 +02:00
silveroxides 9ad992d080 feat(extract): register text encoder nodes 2026-08-03 18:57:57 +02:00
silveroxides 3bd4977d2a fix(cwb): keep preset scale neutral
Runtime conditioning scale reduction must not be carried into persisted model or LoRA weights. Keep named merger presets at 1.0 and reserve scaling for explicit custom configuration.
2026-08-03 18:48:14 +02:00
silveroxides 10484831d0 fix(package): use requirements manifest
Remove duplicate dependency metadata from pyproject.toml and rely on requirements.txt as the custom-node installation source.

Release as 0.5.2.
2026-08-02 20:43:32 +02:00
silveroxides 71b7936f92 fix(package): source registry requirements
Use requirements.txt as the single dependency source so Registry releases cannot drift from runtime installation constraints.

Release as 0.5.1.
2026-08-02 20:40:42 +02:00
Silver d9fc1f52f1 Bump version from 0.4.6 to 0.5.0 2026-08-02 20:33:22 +02:00
Silver 01570b4805 Merge pull request #19 from silveroxides/feature/cwb-secondary-layer-union
Feature/cwb secondary layer union
2026-08-02 20:32:58 +02:00
silveroxides 89f4def83e feat(cwb): preserve secondary-only layers
- build generic and LoRA worklists from the union of inputs
- copy single-source layers exactly and merge available providers
- keep logical LoRA groups canonical, atomic, and collision-safe
- retain mismatch, dtype, quantization, streaming, and OOM contracts
2026-08-02 20:26:41 +02:00
silveroxides aca6147f35 chore(package): exclude tests 2026-08-02 18:12:01 +02:00
silveroxides f390c12776 fix(merge): enforce LoRA tensor contracts
- preserve direct diffs, precision, and guarded low-bit tensors
- stream paired CWB layers with atomic CPU OOM fallback
- save generated merge artifacts under models_dir
2026-08-02 18:11:19 +02:00
silveroxides bed9c6e8ff feat(merge): add dedicated CWB mergers 2026-08-01 13:14:14 +02:00
silveroxides 49b66ea64e fix(lora): enforce adapter tensor contracts 2026-08-01 12:15:39 +02:00
silveroxides f828ca7dfc chore: bump version to 0.4.6 2026-07-31 05:55:30 +02:00
silveroxides 47be3fa1bd fix(lora): preserve FP32 source layers 2026-07-31 05:55:21 +02:00
silveroxides e972b53cf0 chore: bump version to 0.4.5 2026-07-31 05:49:02 +02:00
silveroxides 4cb065b9d2 feat(lora): add opt-in full 1D diffs 2026-07-31 05:48:53 +02:00
Silver febab8ffcd Merge pull request #16 from Booyaka101/fix/dora-extract-1d-tensor-crash
Fix DoRA extraction crash on 1-D weights (x_pad_token, norm weights)
2026-07-11 11:26:54 +02:00
Silver 06365311e2 Merge pull request #18 from silveroxides/fix/loading-refactor
Refactor loading strategies in mergers and LoRA resize nodes.
2026-07-11 11:23:07 +02:00
Booyaka101 22e7701411 fix: skip 1-D weights before the norm in DoRA extraction
The matched-shape branch computed torch.linalg.norm(tensor_a, dim=1)
for every non-4-D tensor, including 1-D weights (e.g. x_pad_token,
norm .weight). A 1-D tensor has no dim=1, so this raised
'Dimension out of range ... but got 1' and aborted the whole
extraction. The code already skips 1-D tensors a few lines later, so
bail out before the norm instead.

Fixes #15
2026-06-06 18:21:50 +08:00
18 changed files with 7668 additions and 494 deletions
+1
View File
@@ -0,0 +1 @@
tests/
+19
View File
@@ -20,6 +20,8 @@ from .nodes.merger import (
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger
)
from .nodes.consensus_merger import CWB_MERGER_NODES
from .nodes.model_analysis import MODEL_ANALYSIS_NODES
from .nodes.lora_extract_svd import (
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius
@@ -32,6 +34,13 @@ from .nodes.dora_learned_wd import (
DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile,
DoRALearnedExtractKnee, DoRALearnedExtractFrobenius
)
from .nodes.text_encoder_extract import (
TextEncoderLoRAExtractFixed, TextEncoderLoRAExtractRatio,
TextEncoderLoRAExtractQuantile, TextEncoderLoRAExtractKnee,
TextEncoderLoRAExtractFrobenius, TextEncoderDoRAExtractFixed,
TextEncoderDoRAExtractRatio, TextEncoderDoRAExtractQuantile,
TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius,
)
from .nodes.lora_resize import (
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
@@ -68,6 +77,10 @@ class ModelUtilsExtension(ComfyExtension):
# Three-Model Mergers
ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger,
CheckpointThreeMerger, EmbeddingThreeMerger,
# Dedicated CWB Mergers
*CWB_MERGER_NODES,
# Two-Model Analysis
*MODEL_ANALYSIS_NODES,
# LoRA Extraction
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius,
@@ -77,6 +90,12 @@ class ModelUtilsExtension(ComfyExtension):
# Learned DoRA Extraction
DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile,
DoRALearnedExtractKnee, DoRALearnedExtractFrobenius,
# Text Encoder LoRA/DoRA Extraction
TextEncoderLoRAExtractFixed, TextEncoderLoRAExtractRatio,
TextEncoderLoRAExtractQuantile, TextEncoderLoRAExtractKnee,
TextEncoderLoRAExtractFrobenius, TextEncoderDoRAExtractFixed,
TextEncoderDoRAExtractRatio, TextEncoderDoRAExtractQuantile,
TextEncoderDoRAExtractKnee, TextEncoderDoRAExtractFrobenius,
# LoRA Resize
LoRAResizeFixed, LoRAResizeRatio,
LoRAResizeFrobenius, LoRAResizeCumulative,
+111
View File
@@ -0,0 +1,111 @@
# Consensus-Weighted Blending Mergers
The dedicated CWB nodes merge two or three safetensors files using
Consensus-Weighted Blending. They calculate an element-wise mean or median
consensus, compare each aligned source vector to that consensus, and normalize
similarity-derived weights before producing the output vector.
## Node families
Two-input and three-input variants are provided for checkpoints, standalone
diffusion models, text encoders, LoRAs, and embeddings.
Model A supplies output metadata and anchors shared-layer shape, naming, and
mismatch behavior. The output key set is the union of all inputs. A valid
tensor found only in a secondary input is copied unchanged. When Model A lacks
a tensor that is present in multiple secondary inputs, those available sources
are CWB-merged without inserting a zero contribution for Model A. Embedding
mergers additionally use the longest compatible first dimension as their
alignment reference.
## Presets and custom controls
`baseline` is the default. The standard clarity, smooth, varied, and diversity
presets are joined by `power_blend` and the DSC presets demonstrated by the CWB
conditioning implementation. Select `custom` to make every manual CWB control
authoritative. Every named merger preset keeps `global_scale` at `1.0`; scaling
persisted model or LoRA weights is only performed when explicitly selected in
the custom controls.
Similarity alignment greedily matches source rows to the reference by cosine
similarity. Index alignment keeps absolute row positions. Position weight can
bias similarity matches toward nearby normalized positions. Preserve Common
Prefix copies the longest numerically identical leading row span directly from
the first input.
Inputs have equal prior weight. The similarity calculation supplies the
effective contribution for each aligned vector. If every candidate is rejected
or receives zero weight, CWB falls back to equal weighting.
## Missing tensors and filtering
- `skip` preserves Model A when an anchored input is missing or incompatible.
In embedding-union mode it omits that source for the affected key.
- `zeros` supplies a compatible zero tensor for the missing source.
- `error` stops the merge.
- Exclude patterns preserve the anchored source tensor.
- Discard patterns remove the matching output tensor or logical LoRA layer.
The missing-input modes above apply to Model-A-owned layers. Model A being
absent from a secondary-only layer is not itself a mismatch. A shape conflict
between multiple available secondary providers does use the selected mismatch
mode, with the earliest available input as its preservation anchor.
Patterns use regular expressions by default. Enable Glob Patterns to use glob
syntax instead.
## LoRA behavior
LoRA inputs use the same parser as the resize and multi-merge nodes. It accepts
all seven low-rank A/B spellings recognized by ComfyUI's `LoRAAdapter`, plus
direct `.diff`, `.diff_b`, `.w_norm`, `.b_norm`, and `.set_weight` layers.
Recognized diffusion-model outputs use the canonical
`diffusion_model.<layer>.lora_A.weight` / `.lora_B.weight` convention.
Secondary formats are matched by normalized logical layer name. Differing
ranks are zero-padded like the DARE/TIES mergers: A/down tensors on dimension 0
and B/up tensors on dimension 1. An existing anchor-source alpha key is updated
to the resulting maximum rank; no alpha key is invented when the anchor has
none.
Similarity alignment operates on paired latent-rank components. A rows and B
columns share one mapping; fixed input/output feature axes are never reordered.
Input alpha scaling is absorbed into the B factor before blending, and global
scale is applied once to the resulting pair rather than once per factor.
Incomplete pairs and unrecognized tensors are copied from the earliest
available source unchanged.
Companion-bearing groups (`lora_mid`, reshape, DoRA, or set-weight) are kept
atomically from their earliest available source rather than partially
transformed. Secondary-only logical groups are included automatically: a group
from one source is copied unchanged, while compatible groups from multiple
secondary sources are merged using only those sources. Recognized roles are
written with the same canonical output convention as shared layers.
`include_1d_diffs` is disabled by default. Disabled shared 1D direct layers
preserve Model A, while secondary-only layers preserve their earliest source.
Enabled 1D direct layers participate as complete tensors and are always saved
in FP32.
## Precision and quantized inputs
CWB computes in FP32. Unless Override Dtype is enabled, any participating FP32
source keeps the generated logical layer in FP32. Guarded low-bit tensors and
enabled 1D direct diffs are never downcast by the override.
ComfyUI/comfy-kitchen quantized inputs are rejected before an output writer is
opened. Files with three or more bare INT8, UINT8, FP8, or FP4 tensors are also
rejected. One or two isolated low-bit tensors produce one warning per input and
are preserved without CWB arithmetic; an affected LoRA causes its complete
logical layer to be preserved.
## Streaming and failure safety
Model parameters and direct patches use fixed-coordinate blending; similarity
row matching is limited to embeddings and LoRA rank components. Each logical
layer is loaded, transferred to the processing device, completed, copied back
to CPU, queued for writing, and released before the next layer. Outputs are
written to a temporary sibling and atomically published only after every layer
and asynchronous write succeeds. Failed merges do not replace or finalize the
requested output. A CUDA OOM retries only the affected layer on CPU after
releasing the failed GPU working set; failure of that retry aborts atomically.
Generated files are published under the matching category inside
`folder_paths.models_dir`, never ComfyUI's general output directory.
+75
View File
@@ -0,0 +1,75 @@
# Two-Model Similarity Analysis
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.
## Outputs
`comparison_report` contains conventional numerical differences and topology
findings. `cwb_report` separately contains diagnostics derived from the
similarity, alignment, and consensus stages of Consensus-Weighted Blending.
`documentation` contains this reference.
The reports begin with a global summary, followed by inferred blockwise values
and one entry for every comparable tensor. Block names are approximated from
the tensor-key hierarchy through its first numeric path component; they are not
an architecture-specific interpretation.
## Standard metrics
- MAE, MSE, and RMSE describe absolute error at the original weight scale.
- Maximum absolute difference identifies the largest individual change.
- Relative L2 is `||A-B|| / ((||A||+||B||)/2)` and is scale-aware.
- Cosine similarity compares direction; Pearson correlation compares centered
variation.
- Exact equality, sign agreement, norm ratio, and finite-value coverage expose
differences hidden by averages.
The global section reports both parameter-weighted metrics and equal-layer
averages. No arbitrary combined similarity score is assigned. Undefined
zero-norm or constant-tensor statistics are shown as `N/A`.
Only shared, same-shape floating tensors enter numerical aggregates. Missing
keys and shape mismatches are listed separately and are never padded, cropped,
or replaced with zeros. Non-floating tensors receive an exact-equality check.
Quantized models carrying ComfyUI quantization metadata are rejected because
encoded storage values do not represent directly comparable model weights.
## CWB diagnostics
The CWB report provides pairwise row-cosine statistics and each input's affinity
to both mean and median consensus. It stops before CWB's merge-weighting stage:
no thresholds, contribution weights, diversity weighting, vector
reconstitution, norm rescaling, or output writer are used.
Checkpoint, diffusion-model, and text-encoder tensors always retain their fixed
coordinate alignment. Embedding and LoRA nodes expose one optional
`cwb_similarity_alignment` control. When enabled, it performs
threshold-independent greedy one-to-one matching derived from CWB and reports
coverage, matched-score statistics, and improvement over index alignment. This
operation is quadratic and may be slow for large embeddings.
LoRA down/up tensors are treated as one logical layer. Rank-component matching
uses the product of down and up cosine similarity, including CWB's paired-sign
equivalence for comparison. No source tensor is modified.
## Memory and failure handling
Inputs always use Unified Efficient Loader's low-memory streaming mode. Matching
A/B tensors are loaded concurrently through bounded `async_stream()` pipelines
with `batch_size=1` and `prefetch_batches=1`, but no subsequent logical layer is
requested early. A complete LoRA down/up pair is loaded once and used for both
reports. CUDA inputs use pinned transfer while retaining the current CPU source
tensors for fallback.
Calculations use FP32. When CUDA processing runs out of memory, CUDA memory is
cleared and the current logical layer is retried once on CPU. After success or
failure, all tensor references for that layer are released and every loaded key
is marked processed exactly once before advancing. Both reports list every CPU
fallback. A failed CPU retry or any non-OOM error stops analysis with the layer
name in the error.
`top_weight_differences` bounds the scalar-coordinate list; every comparable
tensor is still present in the layerwise report. `DOCUMENTATION ONLY` loads no
model data.
File diff suppressed because it is too large Load Diff
+8
View File
@@ -551,6 +551,14 @@ def extract_dora_from_files(
dora_scale = None
del tensor_b
else:
# 1-D weights (e.g. x_pad_token, norm .weight) have no
# weight-decomposition direction and are skipped below
# anyway; bail out before the norm, which would index
# dim=1 and raise on a 1-D tensor.
if tensor_a.ndim < 2:
del tensor_a, tensor_b
return "skipped", None
# Weight Decomposition calculation
is_conv_dim = tensor_a.ndim == 4
+70 -21
View File
@@ -19,6 +19,7 @@ from .device_utils import (
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
from .quantization_guard import inspect_low_bit_input
@@ -458,6 +459,7 @@ def extract_lora_from_files(
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
include_1d_diffs: bool = False,
) -> None:
"""
Extract LoRA from difference between two models, writing incrementally to disk.
@@ -469,7 +471,7 @@ def extract_lora_from_files(
linear_param: Mode parameter for linear layers
conv_param: Mode parameter for conv layers
device: Computation device
save_dtype: Output dtype
save_dtype: Output dtype for layers whose sources are not FP32
output_path: Full path to output safetensors file
linear_max_rank: Max rank for linear layers
conv_max_rank: Max rank for conv layers
@@ -480,6 +482,8 @@ def extract_lora_from_files(
chunk_large_layers: Enable chunked extraction
glob_skip_patterns: When True, treat skip_patterns as glob (* wildcard).
When False (default), treat as Python regex.
include_1d_diffs: Save 1D tensors in FP32 and keys ending in .scale or
.lin as full .diff patches instead of skipping them.
"""
save_torch_dtype = {
"fp32": torch.float32,
@@ -498,10 +502,22 @@ def extract_lora_from_files(
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
low_bit_a = inspect_low_bit_input(handler_a, f"Model A ({model_a_path})", "LoRA Extract")
low_bit_b = inspect_low_bit_input(handler_b, f"Model B ({model_b_path})", "LoRA Extract")
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = [k for k in keys_a if k.endswith(".weight")]
pbar = comfy.utils.ProgressBar(len(weight_keys))
direct_diff_suffixes = (".scale", ".lin")
extraction_keys = [
k for k in keys_a
if k not in low_bit_a
and k not in low_bit_b
and (k.endswith(".weight")
or (
include_1d_diffs
and (len(handler_a.get_shape(k)) == 1 or k.endswith(direct_diff_suffixes))
))
]
pbar = comfy.utils.ProgressBar(len(extraction_keys))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key):
@@ -510,6 +526,13 @@ def extract_lora_from_files(
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
return "skipped", None
source_dtypes = [handler_a.get_dtype(key)]
if key in keys_b:
source_dtypes.append(handler_b.get_dtype(key))
layer_save_dtype = (
torch.float32 if torch.float32 in source_dtypes else save_torch_dtype
)
# Load tensors with pinned memory for CUDA
use_pinned = device == 'cuda'
@@ -561,7 +584,18 @@ def extract_lora_from_files(
del weight_diff
return "skipped", None
# Skip 1D tensors
# SVD requires matrix-shaped tensors. Optionally preserve 1D weights
# and scale/lin parameters as exact runtime-applicable diff patches.
is_direct_diff = weight_diff.ndim == 1 or key.endswith(direct_diff_suffixes)
if include_1d_diffs and is_direct_diff:
direct_diff_dtype = torch.float32 if weight_diff.ndim == 1 else layer_save_dtype
layer_results = {
_format_direct_diff_key(key, lora_name):
weight_diff.to(direct_diff_dtype).cpu().contiguous()
}
del weight_diff
return "full", layer_results
if weight_diff.ndim < 2:
del weight_diff
return "skipped", None
@@ -588,8 +622,8 @@ def extract_lora_from_files(
weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(layer_save_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(layer_save_dtype).cpu().contiguous()
del weight_diff
return "chunked", layer_results
@@ -599,12 +633,12 @@ def extract_lora_from_files(
# Store result
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.diff"] = weight_diff.to(layer_save_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, _ = result
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(layer_save_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(layer_save_dtype).cpu().contiguous()
status = "extracted"
del weight_diff
@@ -613,7 +647,7 @@ def extract_lora_from_files(
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
for key in tqdm(extraction_keys, desc="Extracting LoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
@@ -704,6 +738,13 @@ def _format_lora_key(key: str) -> str:
return f"diffusion_model.{key}"
def _format_direct_diff_key(source_key: str, lora_name: str) -> str:
"""Map a full delta to ComfyUI's canonical direct-patch key."""
if source_key.endswith(".bias"):
return f"{lora_name[:-5]}.diff_b"
return f"{lora_name}.diff"
def _get_common_inputs():
return [
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
@@ -716,13 +757,16 @@ def _get_common_inputs():
tooltip="Skip layers with max difference below this"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.String.Input("output_filename", default="extracted_lora"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16",
tooltip="Output dtype; layers stored as FP32 in either source remain FP32"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"),
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_1d_diffs", default=False,
tooltip="Include 1D tensors as FP32 and keys ending in .scale or .lin as full .diff patches"),
]
@@ -753,7 +797,8 @@ class LoRAExtractFixed(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_1d_diffs=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)
@@ -763,7 +808,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
lazy_load, force_clear_cache, glob_skip_patterns, include_1d_diffs
)
return io.NodeOutput(output_path)
@@ -798,7 +843,8 @@ class LoRAExtractRatio(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, 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_1d_diffs=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)
@@ -809,7 +855,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
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
)
return io.NodeOutput(output_path)
@@ -844,7 +890,8 @@ class LoRAExtractQuantile(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, 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_1d_diffs=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)
@@ -855,7 +902,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
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
)
return io.NodeOutput(output_path)
@@ -888,7 +935,8 @@ class LoRAExtractKnee(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, knee_method, 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_1d_diffs=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)
@@ -899,7 +947,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
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
)
return io.NodeOutput(output_path)
@@ -934,7 +982,8 @@ class LoRAExtractFrobenius(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, 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_1d_diffs=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)
@@ -945,7 +994,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
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
)
return io.NodeOutput(output_path)
+484 -247
View File
File diff suppressed because it is too large Load Diff
+899 -210
View File
File diff suppressed because it is too large Load Diff
+48 -11
View File
@@ -12,6 +12,8 @@ from .device_utils import (
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
from .quantization_guard import inspect_low_bit_input, write_preserved_tensor
from .lora_resize import is_direct_diff_key
def load_documentation_from_file(filename):
@@ -105,6 +107,18 @@ class MergerLogic:
all_keys = primary_handler.keys()
metadata = primary_handler.metadata()
low_bit_keys_by_name = {}
if model_type == "loras":
try:
for name, handler in handlers.items():
low_bit_keys_by_name[name] = inspect_low_bit_input(
handler, f"LoRA ({name})", "Generic LoRA Merge"
)
except Exception:
for handler in handlers.values():
handler.__exit__(None, None, None)
raise
# Convert mismatch_mode string to enum
mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip')
mismatch_mode = MissingTensorBehavior(mismatch_mode_str)
@@ -133,6 +147,7 @@ class MergerLogic:
save_dtype = recipe_params.pop('save_dtype')
save_torch_dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}.get(save_dtype)
override_dtype = recipe_params.pop('override_dtype', False)
include_1d_diffs = recipe_params.pop('include_1d_diffs', False)
def determine_target_dtype(original_dtypes, target_dtype, override):
if override:
@@ -161,8 +176,7 @@ class MergerLogic:
error_keys = []
output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
output_dir = folder_paths.get_folder_paths(output_folder)[-1]
output_dir = os.path.join(folder_paths.models_dir, output_folder)
os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
@@ -193,6 +207,14 @@ class MergerLogic:
with torch.no_grad():
for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
preserve_low_bit = model_type == "loras" and any(
key in low_bit_keys for low_bit_keys in low_bit_keys_by_name.values()
)
if preserve_low_bit:
_flush_batch()
write_preserved_tensor(writer, key, primary_handler)
pbar.update(1)
continue
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns, glob_mode=glob_mode):
discarded_keys += 1
@@ -208,6 +230,13 @@ class MergerLogic:
original_dtypes.append(handler.get_dtype(key))
target_dtype = determine_target_dtype(original_dtypes, save_torch_dtype, override_dtype)
preserve_model_a_1d = model_type == "loras" and cpu_tensor.ndim == 1 and not include_1d_diffs
if preserve_model_a_1d:
target_dtype = determine_target_dtype(
[primary_handler.get_dtype(key)], save_torch_dtype, override_dtype
)
elif model_type == "loras" and is_direct_diff_key(key) and cpu_tensor.ndim == 1:
target_dtype = torch.float32
if process_device == 'cuda':
tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype)
@@ -216,13 +245,14 @@ class MergerLogic:
del cpu_tensor
# Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode):
if preserve_model_a_1d or _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode):
t = tensor_a.detach().to(target_dtype).cpu()
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
batch_key_count += 1
excluded_keys += 1
if not preserve_model_a_1d:
excluded_keys += 1
del tensor_a
if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
_flush_batch()
@@ -303,8 +333,7 @@ class MergerLogic:
handler.__exit__(None, None, None)
output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type
# Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet
output_dir = folder_paths.get_folder_paths(output_folder)[-1]
output_dir = os.path.join(folder_paths.models_dir, output_folder)
os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
@@ -549,7 +578,9 @@ class LoRATwoMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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("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."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -562,7 +593,8 @@ class LoRATwoMerger(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_1d_diffs: bool) -> 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)
@@ -576,6 +608,7 @@ class LoRATwoMerger(io.ComfyNode):
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
"include_1d_diffs": include_1d_diffs,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -885,7 +918,9 @@ class LoRAThreeMerger(io.ComfyNode):
"When False (default), patterns are Python regex matched as substrings."),
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("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."),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -898,7 +933,8 @@ class LoRAThreeMerger(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_1d_diffs: bool) -> 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)
@@ -912,6 +948,7 @@ class LoRAThreeMerger(io.ComfyNode):
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"glob_patterns": glob_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
"include_1d_diffs": include_1d_diffs,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -982,4 +1019,4 @@ class EmbeddingThreeMerger(io.ComfyNode):
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
return io.NodeOutput(filename, doc)
return io.NodeOutput(filename, doc)
File diff suppressed because it is too large Load Diff
+145
View File
@@ -0,0 +1,145 @@
"""Fail-safe detection for quantized or isolated low-bit safetensors inputs."""
import logging
import json
import struct
import torch
MAX_ISOLATED_LOW_BIT_TENSORS = 2
def _low_bit_dtypes():
names = (
"int8",
"uint8",
"float4_e2m1fn_x2",
"float8_e4m3fn",
"float8_e4m3fnuz",
"float8_e5m2",
"float8_e5m2fnuz",
"float8_e8m0fnu",
)
return {getattr(torch, name) for name in names if hasattr(torch, name)}
LOW_BIT_DTYPES = _low_bit_dtypes()
LOW_BIT_STORAGE_CODES = {"I8", "U8", "F4"}
UPSTREAM_WRITER_STORAGE_CODES = {
"F64", "F32", "F16", "BF16", "I64", "I32", "I16", "I8", "U8",
"U64", "U32", "U16", "BOOL", "C64", "F8_E5M2", "F8_E4M3",
}
def _header(handler):
if getattr(handler, "_header", None) is not None:
return handler._header
with open(handler.filename, "rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
return json.loads(source.read(header_size).decode("utf-8"))
def storage_dtype_code(handler, key: str) -> str | None:
entry = _header(handler).get(key)
return entry.get("dtype") if entry else None
def _is_low_bit(handler, key: str) -> bool:
storage_code = storage_dtype_code(handler, key)
if storage_code in LOW_BIT_STORAGE_CODES or (
storage_code is not None and storage_code.startswith("F8_")
):
return True
return handler.get_dtype(key) in LOW_BIT_DTYPES
def inspect_low_bit_input(handler, label: str, operation: str) -> set[str]:
"""Reject quantized models; return isolated low-bit keys that must not be transformed."""
keys = list(handler.keys())
metadata = handler.metadata() or {}
quant_keys = [key for key in keys if key.endswith(".comfy_quant")]
has_quant_metadata = "_quantization_metadata" in metadata
has_legacy_scaled_fp8 = any(key.endswith("scaled_fp8") for key in keys)
if quant_keys or has_quant_metadata or has_legacy_scaled_fp8:
raise ValueError(
f"[{operation}] {label} uses ComfyUI quantization metadata; quantized inputs are unsupported"
)
low_bit_keys = {key for key in keys if _is_low_bit(handler, key)}
if len(low_bit_keys) > MAX_ISOLATED_LOW_BIT_TENSORS:
raise ValueError(
f"[{operation}] {label} contains {len(low_bit_keys)} low-bit tensors; "
"casted or quantized inputs are unsupported"
)
if low_bit_keys:
logging.warning(
"[%s] %s contains %d isolated low-bit tensor(s); "
"preserving/excluding unchanged: %s",
operation,
label,
len(low_bit_keys),
", ".join(sorted(low_bit_keys)),
)
return low_bit_keys
def layer_has_low_bit(block_keys: dict, low_bit_keys: set[str]) -> bool:
return any(
block_keys.get(name) in low_bit_keys
for name in (
"down", "up", "alpha", "mid", "reshape", "dora_scale",
"diff", "diff_b", "w_norm", "b_norm", "set_weight",
)
)
def write_preserved_tensor(writer, key: str, handler, output_key: str | None = None) -> None:
"""Copy one tensor's original safetensors bytes and dtype into an active writer."""
destination = output_key or key
header = _header(handler)
entry = header[key]
if entry["dtype"] in UPSTREAM_WRITER_STORAGE_CODES:
writer.write(destination, handler.get_tensor(key).cpu().contiguous())
return
source_start, source_end = entry["data_offsets"]
byte_size = source_end - source_start
with writer._lock:
if destination in writer._manifest:
raise ValueError(f"Tensor '{destination}' has already been written.")
target_start = writer._current_data_offset
target_end = target_start + byte_size
writer._manifest[destination] = {
"dtype": entry["dtype"],
"shape": entry["shape"],
"data_offsets": [target_start, target_end],
}
writer._current_data_offset = target_end
source_path = handler.filename
source_absolute = 8 + getattr(handler, "_header_size", 0) + source_start
if not getattr(handler, "_header_size", None):
with open(source_path, "rb") as source:
source_absolute = 8 + struct.unpack("<Q", source.read(8))[0] + source_start
target_absolute = 8 + writer.max_header_bytes + target_start
def copy_raw():
try:
with open(source_path, "rb") as source:
source.seek(source_absolute)
remaining = byte_size
with writer._lock:
writer._file.seek(target_absolute)
while remaining:
chunk = source.read(min(8 * 1024 * 1024, remaining))
if not chunk:
raise EOFError(f"Unexpected EOF while preserving tensor '{key}'")
writer._file.write(chunk)
remaining -= len(chunk)
finally:
writer._semaphore.release()
writer._semaphore.acquire()
writer._futures.append(writer._executor.submit(copy_raw))
+721
View File
@@ -0,0 +1,721 @@
"""
Text Encoder LoRA and DoRA Extraction.
Specifically targets extracting text encoder deltas as LoRA/DoRA with custom mapping rules
for layer name prefixes and PEFT weight formatting.
"""
import fnmatch
import os
import re
import torch
import torch.linalg as linalg
import folder_paths
import comfy.utils
from tqdm import tqdm
from comfy_api.latest import io
from .device_utils import (
estimate_model_size, prepare_for_large_operation,
cleanup_after_operation
)
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
# Import SVD rank utilities from standard lora extract file to avoid duplication
from .lora_extract_svd import (
_compute_rank,
_svd_extract_linear,
_svd_extract_conv,
_extract_chunked_layer,
_detect_fused_layer,
_compile_patterns,
_matches_any_pattern
)
def _format_te_lora_key(key: str) -> str:
"""
Format key specifically for Text Encoder LoRA/DoRA saving.
Corrects layer name prefixing based on specific target mappings.
"""
if key.endswith(".weight"):
key = key[:-7]
rules = [
("model.language_model.layers", "text_encoders.transformer.model.layers"),
("model.visual.blocks", "text_encoders.transformer.visual.blocks"),
("model.visual.deepstack_merger_list", "text_encoders.transformer.visual.deepstack_merger_list"),
("model.visual.merger", "text_encoders.transformer.visual.merger"),
("model.visual.pos_embed", "text_encoders.transformer.visual.pos_embed"),
("model.layers", "text_encoders.transformer.model.layers"),
("vision_model.embeddings.position_embedding", "text_encoders.transformer.vision_model.embeddings.position_embedding"),
("vision_model.encoder.layers", "text_encoders.transformer.vision_model.encoder.layers"),
]
for prefix, target in rules:
if key.startswith(prefix):
return target + key[len(prefix):]
# Prepend text_encoders. as default generic fallback
if key.startswith("text_encoders."):
return key
return f"text_encoders.{key}"
def extract_te_from_files(
model_a_path: str,
model_b_path: str,
mode: str,
linear_param: float,
conv_param: float,
device: str,
save_dtype: str,
output_path: str,
is_dora: bool = False,
linear_max_rank: int = None,
conv_max_rank: int = None,
clamp_quantile: float = 0.99,
min_diff: float = 0.0,
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
chunk_large_layers: bool = True,
svd_niter: int = 2,
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
) -> None:
"""
Extract LoRA/DoRA from difference between two Text Encoder models, writing incrementally to disk.
"""
save_torch_dtype = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}.get(save_dtype, torch.float16)
skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns)
# Prepare memory before heavy operation
total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path)
print(f"[TE Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = [k for k in keys_a if k.endswith(".weight")]
pbar = comfy.utils.ProgressBar(len(weight_keys))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
def _process_layer(key):
lora_name = _format_te_lora_key(key)
if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
return "skipped", None
# Load tensors with pinned memory for CUDA
use_pinned = device == 'cuda'
if key not in keys_b:
if mismatch_mode == "skip":
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
cpu_a = handler_a.get_tensor(key)
if use_pinned:
weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
else:
weight_diff = cpu_a.to(device=device, dtype=torch.float32)
del cpu_a
dora_scale = None
else:
cpu_a = handler_a.get_tensor(key)
cpu_b = handler_b.get_tensor(key)
if use_pinned:
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
else:
tensor_a = cpu_a.to(device=device, dtype=torch.float32)
tensor_b = cpu_b.to(device=device, dtype=torch.float32)
del cpu_a, cpu_b
if tensor_a.shape != tensor_b.shape:
if mismatch_mode == "skip":
del tensor_a, tensor_b
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
weight_diff = tensor_a
dora_scale = None
del tensor_b
else:
# Skip 1D weights
if tensor_a.ndim < 2:
del tensor_a, tensor_b
return "skipped", None
if is_dora:
# Weight Decomposition calculation
is_conv_dim = tensor_a.ndim == 4
if is_conv_dim:
out_ch = tensor_a.shape[0]
flat_a = tensor_a.reshape(out_ch, -1)
flat_b = tensor_b.reshape(out_ch, -1)
norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
else:
norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
# Compute directional diff
eps = 1e-8
norm_ratio = norm_b / (norm_a + eps)
if is_conv_dim:
norm_ratio = norm_ratio.view(-1, 1, 1, 1)
dora_scale = norm_a.view(-1, 1, 1, 1)
else:
dora_scale = norm_a
w_dir = tensor_a * norm_ratio
weight_diff = w_dir - tensor_b
del tensor_a, tensor_b, w_dir
else:
weight_diff = tensor_a - tensor_b
dora_scale = None
del tensor_a, tensor_b
# Skip small differences
if min_diff > 0 and weight_diff.abs().max() < min_diff:
del weight_diff
return "skipped", None
# Skip 1D tensors
if weight_diff.ndim < 2:
del weight_diff
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
try:
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
)
except Exception as e:
# Try chunked extraction for large tensors
if chunk_large_layers and not is_conv:
num_chunks = _detect_fused_layer(key, weight_diff.shape)
if num_chunks > 1:
print(f"[TE Extract] Chunked: {key} ({num_chunks} chunks)")
lora_up, lora_down, rank = _extract_chunked_layer(
weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
)
if lora_up is not None:
# PEFT suffixes
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
if is_dora and dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
del weight_diff
return "chunked", layer_results
print(f"[TE Extract] Failed: {key}: {e}")
del weight_diff
return "skipped", None
# Store result
if mode_str == "full":
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, _ = result
# Write to PEFT weight suffixes
layer_results[f"{lora_name}.lora_A.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
if is_dora and dora_scale is not None:
layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous()
status = "extracted"
del weight_diff
return status, layer_results
writer = IncrementalSafetensorsWriter(output_path)
writer.__enter__()
try:
for key in tqdm(weight_keys, desc="Extracting TE Layers", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
writer.write_dict(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
finally:
writer.__exit__(None, None, None)
print(f"[TE Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
f"{stats['full']} full, {stats['skipped']} skipped")
finally:
handler_a.__exit__(None, None, None)
handler_b.__exit__(None, None, None)
cleanup_after_operation()
def _build_lora_output_path(output_filename: str) -> str:
"""Build output path for LoRA file."""
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
return os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
def _get_te_model_inputs():
return [
io.Combo.Input("model_a", options=folder_paths.get_filename_list("text_encoders"),
tooltip="Finetuned Text Encoder model (A - B = LoRA)"),
io.Combo.Input("model_b", options=folder_paths.get_filename_list("text_encoders"),
tooltip="Base Text Encoder model (A - B = LoRA)"),
]
def _get_common_inputs():
return [
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("chunk_large_layers", default=False,
tooltip="Split large fused layers (QKV, MLP) into chunks"),
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
tooltip="Clamp outlier singular values"),
io.Float.Input("min_diff", default=0.0, min=0.0, max=1.0, step=0.001,
tooltip="Skip layers with max difference below this"),
io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"),
io.String.Input("output_filename", default="extracted_te_lora"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"),
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."),
]
# =============================================================================
# LoRA Text Encoder Extract Nodes
# =============================================================================
class TextEncoderLoRAExtractFixed(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractFixed",
display_name="TE LoRA Extract (Fixed Rank)",
category="ModelUtils/LoRA Extract (TE)",
description="Extract Text Encoder LoRA with specified fixed rank for each layer type.",
inputs=[
*_get_te_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD power iterations"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, is_dora=False,
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
)
return io.NodeOutput(output_path)
class TextEncoderLoRAExtractRatio(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractRatio",
display_name="TE LoRA Extract (Ratio)",
category="ModelUtils/LoRA Extract (TE)",
description="Keep singular values > max(S) / ratio.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, is_dora=False,
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
)
return io.NodeOutput(output_path)
class TextEncoderLoRAExtractQuantile(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractQuantile",
display_name="TE LoRA Extract (Quantile)",
category="ModelUtils/LoRA Extract (TE)",
description="Keep enough singular values to reach target percentage.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, is_dora=False,
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
)
return io.NodeOutput(output_path)
class TextEncoderLoRAExtractKnee(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractKnee",
display_name="TE LoRA Extract (Knee)",
category="ModelUtils/LoRA Extract (TE)",
description="Automatically find optimal rank using knee detection on singular value curve.",
inputs=[
*_get_te_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, is_dora=False,
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
)
return io.NodeOutput(output_path)
class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderLoRAExtractFrobenius",
display_name="TE LoRA Extract (Frobenius)",
category="ModelUtils/LoRA Extract (TE)",
description="Preserve target fraction of Frobenius norm.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, is_dora=False,
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
)
return io.NodeOutput(output_path)
# =============================================================================
# DoRA Text Encoder Extract Nodes
# =============================================================================
class TextEncoderDoRAExtractFixed(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractFixed",
display_name="TE DoRA Extract (Fixed Rank)",
category="ModelUtils/DoRA Extract (TE)",
description="Extract Text Encoder DoRA with specified fixed rank for each layer type.",
inputs=[
*_get_te_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD power iterations"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, output_path, is_dora=True,
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
)
return io.NodeOutput(output_path)
class TextEncoderDoRAExtractRatio(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractRatio",
display_name="TE DoRA Extract (Ratio)",
category="ModelUtils/DoRA Extract (TE)",
description="Keep singular values > max(S) / ratio.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, output_path, is_dora=True,
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
)
return io.NodeOutput(output_path)
class TextEncoderDoRAExtractQuantile(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractQuantile",
display_name="TE DoRA Extract (Quantile)",
category="ModelUtils/DoRA Extract (TE)",
description="Keep enough singular values to reach target percentage.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, output_path, is_dora=True,
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
)
return io.NodeOutput(output_path)
class TextEncoderDoRAExtractKnee(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractKnee",
display_name="TE DoRA Extract (Knee)",
category="ModelUtils/DoRA Extract (TE)",
description="Automatically find optimal rank using knee detection on singular value curve.",
inputs=[
*_get_te_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee"),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, knee_method, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, is_dora=True,
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
)
return io.NodeOutput(output_path)
class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="TextEncoderDoRAExtractFrobenius",
display_name="TE DoRA Extract (Frobenius)",
category="ModelUtils/DoRA Extract (TE)",
description="Preserve target fraction of Frobenius norm.",
inputs=[
*_get_te_model_inputs(),
io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
is_output_node=True,
)
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, 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:
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)
output_path = _build_lora_output_path(output_filename)
extract_te_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, output_path, is_dora=True,
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
)
return io.NodeOutput(output_path)
+1 -5
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.4.4"
version = "0.5.2"
license = { file = "LICENSE" }
classifiers = [
@@ -15,10 +15,6 @@ classifiers = [
"Programming Language :: Python :: 3.12",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
"unifiedefficientloader==0.4.4"
]
[project.urls]
Repository = "https://github.com/silveroxides/ComfyUI-ModelUtils"
# Used by Comfy Registry https://comfyregistry.org
+845
View File
@@ -0,0 +1,845 @@
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_cwb_tests"
def _load_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 cwb():
return _load_module("nodes.consensus_merger")
def _patch_io(monkeypatch, module, tmp_path, paths):
monkeypatch.setattr(module.folder_paths, "get_full_path", lambda _, name: paths.get(name))
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
def _result_path(tmp_path, category, result):
return tmp_path / category / result
def _params(output_filename, **overrides):
params = {
"cwb_preset": "custom",
"consensus_type": "mean",
"alignment_method": "index",
"alignment_threshold": 0.4,
"similarity_threshold": 0.0,
"power_alpha": 2.0,
"diversity_beta": 0.0,
"rescale_norm": False,
"global_scale": 1.0,
"dynamic_similarity_contrast": False,
"soft_comfort_bandpass": False,
"position_weight": 0.0,
"preserve_common_prefix": False,
"mismatch_mode": "skip",
"output_filename": output_filename,
"save_dtype": "fp16",
"process_device": "cpu",
"exclude_patterns": "",
"discard_patterns": "",
"glob_patterns": False,
"lazy_load": True,
"force_clear_cache": False,
"override_dtype": False,
"include_1d_diffs": False,
}
params.update(overrides)
return params
def test_preset_resolution_matches_extended_contract(cwb):
assert all(
preset["global_scale"] == pytest.approx(1.0)
for preset in cwb.CWB_PRESETS.values()
)
power = cwb.resolve_cwb_settings(cwb_preset="power_blend")
assert power.consensus_type == "median"
assert power.alignment_threshold == pytest.approx(0.9)
assert power.power_alpha == pytest.approx(8.0)
assert power.dynamic_similarity_contrast is True
assert power.soft_comfort_bandpass is False
baseline = cwb.resolve_cwb_settings(
cwb_preset="baseline",
alignment_method="index",
dynamic_similarity_contrast=True,
soft_comfort_bandpass=True,
)
assert baseline.alignment_method == "similarity"
assert baseline.dynamic_similarity_contrast is False
assert baseline.soft_comfort_bandpass is False
varied = cwb.resolve_cwb_settings(
cwb_preset="varied_merge",
global_scale=1.25,
position_weight=0.2,
preserve_common_prefix=True,
)
assert varied.global_scale == pytest.approx(1.25)
assert varied.position_weight == pytest.approx(0.2)
assert varied.preserve_common_prefix is True
def test_cwb_math_index_prefix_and_scale(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
global_scale=1.0,
)
first = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
second = torch.tensor([[3.0, 0.0], [0.0, 3.0]])
torch.testing.assert_close(
cwb.merge_cwb_tensors([first, second], settings),
torch.tensor([[2.0, 0.0], [0.0, 2.0]]),
)
prefix_settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
global_scale=2.0,
preserve_common_prefix=True,
)
prefixed = cwb.merge_cwb_tensors(
[torch.tensor([[9.0, 9.0], [1.0, 0.0]]),
torch.tensor([[9.0, 9.0], [3.0, 0.0]])],
prefix_settings,
)
torch.testing.assert_close(prefixed[0], torch.tensor([9.0, 9.0]))
torch.testing.assert_close(prefixed[1], torch.tensor([4.0, 0.0]))
def test_similarity_alignment_reorders_matching_rows(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.5,
)
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reversed_source = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
torch.testing.assert_close(
cwb.merge_cwb_tensors([reference, reversed_source], settings),
torch.tensor([[1.5, 0.0], [0.0, 1.5]]),
)
scores = torch.ones((2, 2))
biased = cwb._position_biased_scores(scores, 1.0)
assert biased[0, 0] > biased[0, 1]
def test_fixed_coordinate_merge_never_reorders_model_rows(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.5,
)
reference = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reversed_source = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
expected = torch.stack([
cwb.merge_consensus_group(
torch.stack([reference[row], reversed_source[row]]), settings
)
for row in range(reference.shape[0])
])
torch.testing.assert_close(
cwb.merge_cwb_tensors(
[reference, reversed_source],
settings,
allow_similarity_alignment=False,
),
expected,
)
def test_lora_similarity_alignment_pairs_a_rows_with_b_columns(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="similarity",
alignment_threshold=0.5,
)
reference_down = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
reference_up = torch.tensor([[2.0, 0.0], [0.0, 3.0]])
source_down = torch.tensor([[0.0, 1.0], [-1.0, 0.0]])
source_up = torch.tensor([[0.0, -2.0], [3.0, 0.0]])
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
[reference_down, source_down],
[reference_up, source_up],
settings,
reference_index=0,
)
torch.testing.assert_close(merged_down, reference_down)
torch.testing.assert_close(merged_up, reference_up)
def test_lora_krea_shape_scores_only_rank_components(monkeypatch, cwb):
rank = 256
output_features = 36864
down = torch.ones((rank, 2))
up = torch.ones((output_features, rank))
matrix_shapes = []
def bounded_mm(left, right):
matrix_shapes.append((left.shape, right.shape))
return torch.eye(left.shape[0], right.shape[1], dtype=left.dtype)
monkeypatch.setattr(cwb.torch, "mm", bounded_mm)
monkeypatch.setattr(
cwb,
"merge_consensus_group",
lambda stacked, settings, **kwargs: stacked[0].clone(),
)
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
alignment_method="similarity",
alignment_threshold=0.5,
)
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
[down, down], [up, up], settings, reference_index=0
)
assert merged_down.shape == down.shape
assert merged_up.shape == up.shape
assert matrix_shapes
assert all(left[0] == rank and right[1] == rank for left, right in matrix_shapes)
def test_lora_alpha_and_global_scale_apply_once_to_pair(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
global_scale=3.0,
)
down = torch.ones((1, 1))
up = torch.ones((1, 1))
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
[(0, down, up, 2.0), (1, down, up, 2.0)],
settings,
"cpu",
torch.float32,
)
assert rank == 1
torch.testing.assert_close(merged_down, torch.ones((1, 1)))
torch.testing.assert_close(merged_up, torch.full((1, 1), 6.0))
def test_lora_pair_alignment_supports_convolution_factors(cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
)
down = torch.arange(6, dtype=torch.float32).reshape(2, 3, 1, 1)
up = torch.arange(8, dtype=torch.float32).reshape(4, 2, 1, 1)
merged_down, merged_up = cwb.merge_cwb_lora_pairs(
[down, down], [up, up], settings, reference_index=0
)
torch.testing.assert_close(merged_down, down)
torch.testing.assert_close(merged_up, up)
def test_lora_cuda_oom_retries_current_pair_on_cpu(monkeypatch, caplog, cwb):
settings = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
alignment_method="index",
)
original_to_compute = cwb._to_compute
attempted_devices = []
def fail_cuda(tensor, device):
attempted_devices.append(device)
if str(device).startswith("cuda"):
raise torch.OutOfMemoryError("injected allocation failure")
return original_to_compute(tensor, device)
monkeypatch.setattr(cwb, "_to_compute", fail_cuda)
monkeypatch.setattr(cwb, "_release_failed_cuda_operation", lambda: None)
down = torch.ones((1, 2))
up = torch.ones((2, 1))
with caplog.at_level("WARNING"):
merged_down, merged_up, rank = cwb._merge_lora_pair_to_cpu(
[(0, down, up, 1.0), (1, down, up, 1.0)],
settings,
"cuda",
torch.float32,
operation_label="diffusion_model.foo",
)
assert rank == 1
torch.testing.assert_close(merged_down, down)
torch.testing.assert_close(merged_up, up)
assert attempted_devices[0] == "cuda"
assert "cpu" in attempted_devices
assert "retrying this layer on CPU" in caplog.text
assert "diffusion_model.foo" in caplog.text
def test_norm_rescale_and_dsc_bandpass_are_finite(cwb):
rescale = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="mean",
rescale_norm=True,
)
merged = cwb.merge_consensus_group(
torch.tensor([[2.0, 0.0], [0.0, 2.0]]), rescale
)
assert torch.linalg.vector_norm(merged).item() == pytest.approx(2.0)
dsc = cwb.resolve_cwb_settings(
cwb_preset="custom",
consensus_type="median",
diversity_beta=1.5,
dynamic_similarity_contrast=True,
soft_comfort_bandpass=True,
)
result = cwb.merge_consensus_group(
torch.tensor([[1.0, 0.0], [0.8, 0.2], [0.0, 1.0]]), dsc
)
assert torch.isfinite(result).all()
def test_all_ten_schemas_and_lora_switch_contract(cwb):
assert len(cwb.CWB_MERGER_NODES) == 10
assert len({node.NODE_ID for node in cwb.CWB_MERGER_NODES}) == 10
for node in cwb.CWB_MERGER_NODES:
schema = node.define_schema()
ids = [value.id for value in schema.inputs]
assert ids[:3] == ["execution_mode", "model_a", "model_b"]
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
else:
assert "include_1d_diffs" not in ids
def test_streaming_dtype_nonfloat_and_secondary_preservation(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "a.safetensors"
b = tmp_path / "b.safetensors"
save_file({
"shared": torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32),
"only_a": torch.tensor([5.0, 6.0], dtype=torch.float16),
"metadata_tensor": torch.tensor([1, 2], dtype=torch.int64),
}, str(a), metadata={"source": "A"})
save_file({
"shared": torch.tensor([[3.0, 0.0], [0.0, 3.0]], dtype=torch.float16),
"extra": torch.ones((2, 2)),
"metadata_tensor": torch.tensor([9, 9], dtype=torch.int64),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("anchored")
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
assert set(tensors) == {"shared", "only_a", "extra", "metadata_tensor"}
torch.testing.assert_close(tensors["shared"], torch.tensor([[2.0, 0.0], [0.0, 2.0]]))
assert tensors["shared"].dtype == torch.float32
torch.testing.assert_close(tensors["only_a"], torch.tensor([5.0, 6.0], dtype=torch.float16))
torch.testing.assert_close(tensors["metadata_tensor"], torch.tensor([1, 2]))
torch.testing.assert_close(tensors["extra"], torch.ones((2, 2)))
overridden = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params("anchored_override", override_dtype=True),
)
overridden_tensors = load_file(
str(_result_path(tmp_path, "diffusion_models", overridden))
)
assert overridden_tensors["shared"].dtype == torch.float16
assert overridden_tensors["metadata_tensor"].dtype == torch.int64
def test_failed_merge_preserves_existing_output_and_removes_temporary_file(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "atomic_a.safetensors"
b = tmp_path / "atomic_b.safetensors"
output = tmp_path / "diffusion_models" / "atomic_output.safetensors"
output.parent.mkdir(parents=True)
save_file({"layer": torch.ones((2, 2))}, str(a))
save_file({"layer": torch.full((2, 2), 2.0)}, str(b))
save_file({"existing": torch.tensor([7.0])}, str(output))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
def fail_merge(*args, **kwargs):
raise RuntimeError("injected CWB failure")
monkeypatch.setattr(cwb, "merge_cwb_tensors", fail_merge)
with pytest.raises(RuntimeError, match="injected CWB failure"):
cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("atomic_output")
)
existing = load_file(str(output))
assert set(existing) == {"existing"}
torch.testing.assert_close(existing["existing"], torch.tensor([7.0]))
assert not list(output.parent.glob(".atomic_output.safetensors.*.tmp"))
def test_embedding_union_uses_longest_first_dimension(monkeypatch, tmp_path, cwb):
a = tmp_path / "embed_a.safetensors"
b = tmp_path / "embed_b.safetensors"
save_file({"emb": torch.tensor([[1.0, 0.0], [0.0, 1.0]])}, str(a))
save_file({
"emb": torch.tensor([[3.0, 0.0], [0.0, 3.0], [4.0, 4.0]]),
"secondary_only": torch.tensor([[7.0, 0.0]]),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"embeddings",
_params("embedding_union"),
embedding_union=True,
)
tensors = load_file(str(_result_path(tmp_path, "embeddings", result)))
assert tensors["emb"].shape == (3, 2)
torch.testing.assert_close(tensors["emb"][:2], torch.tensor([[2.0, 0.0], [0.0, 2.0]]))
torch.testing.assert_close(tensors["emb"][2], torch.tensor([4.0, 4.0]))
torch.testing.assert_close(tensors["secondary_only"], torch.tensor([[7.0, 0.0]]))
def test_three_input_streaming_merge(monkeypatch, tmp_path, cwb):
paths = {}
for name, value in (("a", 1.0), ("b", 3.0), ("c", 5.0)):
path = tmp_path / f"three_{name}.safetensors"
save_file({"layer": torch.tensor([[value, 0.0]])}, str(path))
paths[name] = str(path)
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "diffusion_models", _params("three_inputs")
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["layer"]
torch.testing.assert_close(tensor, torch.tensor([[3.0, 0.0]]))
def test_two_input_preserves_secondary_only_tensor_exactly(monkeypatch, tmp_path, cwb):
a = tmp_path / "secondary_a.safetensors"
b = tmp_path / "secondary_b.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"shared": torch.tensor([1.0]),
"secondary": torch.tensor([2.0, 3.0], dtype=torch.float32),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params(
"secondary_exact",
global_scale=2.0,
override_dtype=True,
save_dtype="fp16",
),
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["secondary"]
torch.testing.assert_close(tensor, torch.tensor([2.0, 3.0]))
assert tensor.dtype == torch.float32
def test_three_input_merges_available_secondary_sources(monkeypatch, tmp_path, cwb):
a = tmp_path / "union_a.safetensors"
b = tmp_path / "union_b.safetensors"
c = tmp_path / "union_c.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"shared": torch.tensor([1.0]),
"secondary_shared": torch.tensor([2.0]),
"b_only": torch.tensor([7.0]),
}, str(b))
save_file({
"shared": torch.tensor([1.0]),
"secondary_shared": torch.tensor([4.0]),
"c_only": torch.tensor([9.0]),
}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "diffusion_models", _params("secondary_union")
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
torch.testing.assert_close(tensors["secondary_shared"], torch.tensor([3.0]))
torch.testing.assert_close(tensors["b_only"], torch.tensor([7.0]))
torch.testing.assert_close(tensors["c_only"], torch.tensor([9.0]))
def test_secondary_shape_conflict_skip_preserves_earliest_source(monkeypatch, tmp_path, cwb):
a = tmp_path / "shape_a.safetensors"
b = tmp_path / "shape_b.safetensors"
c = tmp_path / "shape_c.safetensors"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({"secondary": torch.tensor([2.0, 3.0])}, str(b))
save_file({"secondary": torch.tensor([[8.0, 9.0]])}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"],
"diffusion_models",
_params("secondary_shape_skip", mismatch_mode="skip"),
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["secondary"]
torch.testing.assert_close(tensor, torch.tensor([2.0, 3.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"
save_file({"shared": torch.tensor([1.0])}, str(a))
save_file({
"secondary.keep": torch.tensor([4, 5], dtype=torch.uint8),
"secondary.drop": torch.tensor([6.0]),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"diffusion_models",
_params("secondary_filters", discard_patterns=r"secondary\.drop"),
)
tensors = load_file(str(_result_path(tmp_path, "diffusion_models", result)))
assert "secondary.drop" not in tensors
torch.testing.assert_close(
tensors["secondary.keep"], torch.tensor([4, 5], dtype=torch.uint8)
)
def test_lora_cross_format_companion_group_preserves_model_a_canonically(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_a.safetensors"
b = tmp_path / "lora_b.safetensors"
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]], dtype=torch.float32),
"diffusion_model.foo.lora_B.weight": torch.tensor([[1.0], [2.0]], dtype=torch.float32),
"diffusion_model.foo.dora_scale": torch.tensor([1.0, 1.0]),
}, str(a))
save_file({
"lora_unet_foo.lora_down.weight": torch.tensor([[3.0, 4.0], [5.0, 6.0]]),
"lora_unet_foo.lora_up.weight": torch.tensor([[3.0, 5.0], [4.0, 6.0]]),
"lora_unet_secondary.lora_down.weight": torch.ones((1, 2)),
"lora_unet_secondary.lora_up.weight": torch.ones((2, 1)),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("lora_cross_format"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
assert set(tensors) == {
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.dora_scale",
"lora_unet_secondary.lora_A.weight",
"lora_unet_secondary.lora_B.weight",
}
assert tensors["diffusion_model.foo.lora_A.weight"].shape == (1, 2)
assert tensors["diffusion_model.foo.lora_B.weight"].shape == (2, 1)
assert tensors["diffusion_model.foo.lora_A.weight"].dtype == torch.float32
torch.testing.assert_close(
tensors["diffusion_model.foo.dora_scale"], torch.tensor([1.0, 1.0])
)
torch.testing.assert_close(
tensors["lora_unet_secondary.lora_A.weight"], torch.ones((1, 2))
)
torch.testing.assert_close(
tensors["lora_unet_secondary.lora_B.weight"], torch.ones((2, 1))
)
def test_lora_secondary_only_pair_is_preserved_without_scaling(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_secondary_a.safetensors"
b = tmp_path / "lora_secondary_b.safetensors"
save_file({
"diffusion_model.base.diff": torch.tensor([[1.0]]),
}, str(a))
save_file({
"transformer.extra.lora_down.weight": torch.tensor([[2.0, 3.0]]),
"transformer.extra.lora_up.weight": torch.tensor([[4.0], [5.0]]),
"transformer.extra.alpha": torch.tensor(1.0),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params(
"lora_secondary_exact",
global_scale=2.0,
override_dtype=True,
save_dtype="fp16",
),
lora_mode=True,
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"],
torch.tensor([[2.0, 3.0]]),
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"],
torch.tensor([[4.0], [5.0]]),
)
assert tensors["diffusion_model.extra.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.extra.alpha"].item() == pytest.approx(1.0)
def test_lora_three_input_merges_b_and_c_group_without_a(monkeypatch, tmp_path, cwb):
a = tmp_path / "lora_union_a.safetensors"
b = tmp_path / "lora_union_b.safetensors"
c = tmp_path / "lora_union_c.safetensors"
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
save_file({
"diffusion_model.extra.lora_A.weight": torch.full((1, 2), 1.0),
"diffusion_model.extra.lora_B.weight": torch.full((2, 1), 1.0),
}, str(b))
save_file({
"transformer.extra.lora_down.weight": torch.full((1, 2), 3.0),
"transformer.extra.lora_up.weight": torch.full((2, 1), 3.0),
}, str(c))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b), "c": str(c)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b", "c"], "loras", _params("lora_secondary_union"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"], torch.full((1, 2), 2.0)
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"], torch.full((2, 1), 2.0)
)
def test_lora_secondary_direct_and_companion_groups_are_atomic(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "lora_atomic_a.safetensors"
b = tmp_path / "lora_atomic_b.safetensors"
save_file({"diffusion_model.base.diff": torch.tensor([[1.0]])}, str(a))
save_file({
"diffusion_model.norm.diff": torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
"diffusion_model.extra.lora_A.weight": torch.tensor([[4.0, 5.0]]),
"diffusion_model.extra.lora_B.weight": torch.tensor([[6.0], [7.0]]),
"diffusion_model.extra.dora_scale": torch.tensor([8.0, 9.0]),
}, str(b))
_patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)})
result = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params("lora_secondary_atomic", global_scale=3.0),
lora_mode=True,
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.norm.diff"],
torch.tensor([2.0, 3.0], dtype=torch.bfloat16),
)
assert tensors["diffusion_model.norm.diff"].dtype == torch.bfloat16
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_A.weight"], torch.tensor([[4.0, 5.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.extra.lora_B.weight"], torch.tensor([[6.0], [7.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.extra.dora_scale"], torch.tensor([8.0, 9.0])
)
def test_lora_rejects_duplicate_normalized_cores_in_secondary_input(cwb):
with pytest.raises(ValueError, match="same logical layer 'foo'.*input 2"):
cwb._secondary_lora_map(
{
"diffusion_model.foo": {},
"transformer.foo": {},
},
input_index=1,
)
def test_lora_peft_prefix_matching_and_existing_alpha(monkeypatch, tmp_path, cwb):
a = tmp_path / "peft_a.safetensors"
b = tmp_path / "peft_b.safetensors"
save_file({
"base_model.model.diffusion_model.foo.lora_A.weight": torch.ones((1, 2)),
"base_model.model.diffusion_model.foo.lora_B.weight": torch.ones((2, 1)),
"base_model.model.diffusion_model.foo.alpha": torch.tensor(1.0),
}, str(a))
save_file({
"diffusion_model.foo.lora_down.weight": torch.ones((2, 2)),
"diffusion_model.foo.lora_up.weight": torch.ones((2, 2)),
"diffusion_model.foo.alpha": torch.tensor(2.0),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("peft_prefix"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
assert set(tensors) == {
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.alpha",
}
assert tensors["diffusion_model.foo.alpha"].item() == pytest.approx(2.0)
def test_lora_mochi_inputs_emit_preferred_canonical_output(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "mochi_a.safetensors"
b = tmp_path / "mochi_b.safetensors"
for path, value in ((a, 1.0), (b, 3.0)):
save_file({
"diffusion_model.foo.lora_A": torch.full((1, 2), value),
"diffusion_model.foo.lora_B": torch.full((2, 1), value),
"diffusion_model.foo.alpha": torch.tensor(1.0),
}, str(path))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("mochi_canonical"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
assert set(tensors) == {
"diffusion_model.foo.lora_A.weight",
"diffusion_model.foo.lora_B.weight",
"diffusion_model.foo.alpha",
}
def test_lora_1d_direct_switch_is_default_fallback_and_fp32(
monkeypatch, tmp_path, cwb
):
a = tmp_path / "direct_a.safetensors"
b = tmp_path / "direct_b.safetensors"
save_file({"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16)}, str(a))
save_file({"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16)}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
default = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("direct_default"), lora_mode=True
)
default_tensor = load_file(str(_result_path(tmp_path, "loras", default)))[
"diffusion_model.norm.diff"
]
torch.testing.assert_close(default_tensor, torch.tensor([1.0, 2.0], dtype=torch.bfloat16))
assert default_tensor.dtype == torch.bfloat16
enabled = cwb.ConsensusMergerLogic.execute(
["a", "b"],
"loras",
_params("direct_enabled", include_1d_diffs=True, override_dtype=True),
lora_mode=True,
)
enabled_tensor = load_file(str(_result_path(tmp_path, "loras", enabled)))[
"diffusion_model.norm.diff"
]
torch.testing.assert_close(enabled_tensor, torch.tensor([2.0, 3.0]))
assert enabled_tensor.dtype == torch.float32
def test_quantized_marker_errors_before_output(monkeypatch, tmp_path, cwb):
a = tmp_path / "quant_a.safetensors"
b = tmp_path / "quant_b.safetensors"
save_file({
"layer.weight": torch.ones((2, 2)),
"layer.comfy_quant": torch.tensor([1], dtype=torch.uint8),
}, str(a))
save_file({"layer.weight": torch.ones((2, 2))}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("must_not_exist")
)
assert not (tmp_path / "diffusion_models" / "must_not_exist.safetensors").exists()
def test_isolated_low_bit_preserves_anchored_tensor(monkeypatch, tmp_path, cwb):
a = tmp_path / "guard_a.safetensors"
b = tmp_path / "guard_b.safetensors"
save_file({"layer": torch.tensor([1.0, 2.0])}, str(a))
save_file({"layer": torch.tensor([8, 9], dtype=torch.uint8)}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "diffusion_models", _params("guarded_generic")
)
tensor = load_file(str(_result_path(tmp_path, "diffusion_models", result)))["layer"]
torch.testing.assert_close(tensor, torch.tensor([1.0, 2.0]))
assert tensor.dtype == torch.float32
def test_isolated_low_bit_preserves_complete_lora_layer(monkeypatch, tmp_path, cwb):
a = tmp_path / "guard_lora_a.safetensors"
b = tmp_path / "guard_lora_b.safetensors"
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.foo.lora_B.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.foo.alpha": torch.tensor(1.0),
"diffusion_model.foo.dora_scale": torch.tensor([5.0, 6.0]),
}, str(a))
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[9.0, 9.0]]),
"diffusion_model.foo.lora_B.weight": torch.tensor([[9.0], [9.0]]),
"diffusion_model.foo.alpha": torch.tensor(1, dtype=torch.uint8),
}, str(b))
paths = {"a": str(a), "b": str(b)}
_patch_io(monkeypatch, cwb, tmp_path, paths)
result = cwb.ConsensusMergerLogic.execute(
["a", "b"], "loras", _params("guarded_lora"), lora_mode=True
)
tensors = load_file(str(_result_path(tmp_path, "loras", result)))
torch.testing.assert_close(
tensors["diffusion_model.foo.lora_A.weight"], torch.tensor([[1.0, 2.0]])
)
torch.testing.assert_close(
tensors["diffusion_model.foo.lora_B.weight"], torch.tensor([[3.0], [4.0]])
)
assert tensors["diffusion_model.foo.alpha"].item() == pytest.approx(1.0)
torch.testing.assert_close(
tensors["diffusion_model.foo.dora_scale"], torch.tensor([5.0, 6.0])
)
+899
View File
@@ -0,0 +1,899 @@
import importlib
import logging
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_contract_tests"
def _load_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 _load_module("nodes.lora_resize"), _load_module("nodes.lora_merger")
@pytest.fixture(scope="module")
def generic_modules():
return _load_module("nodes.merger"), _load_module("nodes.merger_ops")
def _patch_output(monkeypatch, module, output_dir):
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(output_dir)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(output_dir))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
COMFY_LORA_PAIR_SUFFIXES = [
(".lora_down.weight", ".lora_up.weight"),
("_lora.down.weight", "_lora.up.weight"),
(".lora_A.weight", ".lora_B.weight"),
(".lora.down.weight", ".lora.up.weight"),
(".lora_A", ".lora_B"),
(".lora_linear_layer.down.weight", ".lora_linear_layer.up.weight"),
(".lora_A.default.weight", ".lora_B.default.weight"),
]
@pytest.mark.parametrize("down_suffix,up_suffix", COMFY_LORA_PAIR_SUFFIXES)
def test_all_comfy_lora_pair_spellings_parse_to_canonical_keys(
modules, down_suffix, up_suffix
):
resize, _ = modules
block = "diffusion_model.layer"
layers, passthrough = resize.parse_lora_layers([
f"{block}{down_suffix}",
f"{block}{up_suffix}",
f"{block}.alpha",
])
assert not passthrough
assert layers[block]["down"] == f"{block}{down_suffix}"
assert layers[block]["up"] == f"{block}{up_suffix}"
assert resize.canonical_lora_key(block, "down") == f"{block}.lora_A.weight"
assert resize.canonical_lora_key(block, "up") == f"{block}.lora_B.weight"
@pytest.mark.parametrize("down_suffix,up_suffix", COMFY_LORA_PAIR_SUFFIXES)
def test_merge_to_model_applies_every_comfy_lora_pair_spelling(
monkeypatch, tmp_path, modules, down_suffix, up_suffix
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
variant = COMFY_LORA_PAIR_SUFFIXES.index((down_suffix, up_suffix))
base = tmp_path / f"base_{variant}.safetensors"
adapter = tmp_path / f"adapter_{variant}.safetensors"
block = "diffusion_model.layer"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2), dtype=torch.bfloat16),
}, str(base))
save_file({
f"{block}{down_suffix}": torch.tensor([[1.0, 2.0]]),
f"{block}{up_suffix}": torch.tensor([[3.0], [4.0]]),
f"{block}.alpha": torch.tensor(0.5),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16,
f"merged_variant_{variant}", verbose=False,
)
result = load_file(output)["model.diffusion_model.layer.weight"]
torch.testing.assert_close(
result,
torch.tensor([[1.5, 3.0], [2.0, 4.0]]),
)
assert result.dtype == torch.float32
def test_merge_to_model_applies_comfy_direct_norm_and_set_forms(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "direct_base.safetensors"
adapter = tmp_path / "direct_adapter.safetensors"
save_file({
"model.diffusion_model.norm.weight": torch.zeros(2),
"model.diffusion_model.norm.bias": torch.zeros(2),
"model.diffusion_model.replace.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.norm.w_norm": torch.tensor([1.0, 2.0]),
"diffusion_model.norm.b_norm": torch.tensor([3.0, 4.0]),
"diffusion_model.replace.set_weight": torch.full((2, 2), 7.0),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16,
"direct_forms", verbose=False, include_1d_diffs=True,
)
result = load_file(output)
torch.testing.assert_close(
result["model.diffusion_model.norm.weight"], torch.tensor([0.5, 1.0])
)
torch.testing.assert_close(
result["model.diffusion_model.norm.bias"], torch.tensor([1.5, 2.0])
)
torch.testing.assert_close(
result["model.diffusion_model.replace.weight"], torch.full((2, 2), 7.0)
)
def test_merge_to_model_uses_comfy_reshape_companion(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "reshape_base.safetensors"
adapter = tmp_path / "reshape_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((1, 2)),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.layer.lora_B.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.layer.reshape_weight": torch.tensor([2, 2]),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16,
"reshape_companion", verbose=False,
)
result = load_file(output)["model.diffusion_model.layer.weight"]
assert result.shape == (2, 2)
torch.testing.assert_close(result, torch.tensor([[3.0, 6.0], [4.0, 8.0]]))
def test_merge_to_model_reports_success_and_separates_base_only_biases(
monkeypatch, tmp_path, modules, capsys
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "summary_base.safetensors"
adapter = tmp_path / "summary_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
"model.diffusion_model.layer.bias": torch.zeros(2),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"summary", verbose=True, return_report=True,
)
output = capsys.readouterr().out
assert Path(path).is_file()
assert report.startswith("SUCCESS: 1/1 adapter groups fully applied")
assert "[PATCHED]" not in report
assert "[APPLIED]" not in report
assert "L1: summary_adapter.safetensors, strength=1.0" in report
assert "[BASE-ONLY BIASES RETAINED] (1)" in report
assert "model.diffusion_model.layer.bias" in report
assert "SUCCESS: 1/1 groups applied" in output
assert "1 tensors patched, 1 base tensors not patched, 0 adapter issues" in output
assert "model.diffusion_model.layer.bias" not in output
def test_success_report_stays_compact_and_console_is_summary_only(
monkeypatch, tmp_path, modules, capsys
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "many_success_base.safetensors"
adapter = tmp_path / "many_success_adapter.safetensors"
save_file({
f"model.diffusion_model.layer_{index}.weight": torch.zeros((2, 2))
for index in range(40)
}, str(base))
save_file({
f"diffusion_model.layer_{index}.diff": torch.ones((2, 2))
for index in range(40)
}, str(adapter))
_, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"many_success", verbose=True, return_report=True,
)
console = capsys.readouterr().out
assert report.startswith("SUCCESS: 40/40")
assert "layer_0.weight" not in report
assert "ADAPTER GROUP ISSUES\n- None" in report
assert len(report) < 1000
assert max(map(len, report.splitlines())) < 160
assert "layer_0.weight" not in console
assert len(console.splitlines()) <= 9
def test_merge_to_model_report_identifies_shape_mismatch_and_unresolved_group(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "diagnostic_base.safetensors"
adapter = tmp_path / "diagnostic_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.layer.diff": torch.ones(3),
"diffusion_model.orphan.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.orphan.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
_, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"diagnostic", verbose=False, include_1d_diffs=True, return_report=True,
)
assert report.startswith("COMPLETED WITH ERRORS: 0/2")
assert "[MAPPED BUT UNCHANGED] (1)" in report
assert "[SHAPE MISMATCH] (1)" in report
assert "source_shape=(3,), target_shape=(2, 2)" in report
assert "[UNRESOLVED] (1)" in report
assert "diffusion_model.orphan" in report
assert "no matching base tensor" in report
def test_merge_to_model_detects_comfy_adapter_logged_failure_atomically(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "logged_failure_base.safetensors"
adapter = tmp_path / "logged_failure_adapter.safetensors"
base_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
save_file({"model.diffusion_model.layer.weight": base_weight}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
class LoggedFailureAdapter:
def calculate_weight(self, weight, key, *args, **kwargs):
weight.add_(99.0)
logging.error("ERROR lora %s synthetic internal failure", key)
return weight
class FakeLoRAAdapter:
@staticmethod
def load(*args, **kwargs):
return LoggedFailureAdapter()
monkeypatch.setattr(resize, "LoRAAdapter", FakeLoRAAdapter)
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"logged_failure", verbose=False, return_report=True,
)
torch.testing.assert_close(
load_file(path)["model.diffusion_model.layer.weight"], base_weight
)
assert report.startswith("COMPLETED WITH ERRORS: 0/1")
assert "[CALCULATION FAILURE] (1)" in report
assert "synthetic internal failure" in report
assert "[MAPPED BUT UNCHANGED] (1)" in report
def test_merge_to_model_report_classifies_disabled_1d_as_exclusion(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "disabled_1d_base.safetensors"
adapter = tmp_path / "disabled_1d_adapter.safetensors"
save_file({"model.diffusion_model.norm.weight": torch.zeros(2)}, str(base))
save_file({"diffusion_model.norm.diff": torch.ones(2)}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"disabled_1d", verbose=False, return_report=True,
)
torch.testing.assert_close(
load_file(path)["model.diffusion_model.norm.weight"], torch.zeros(2)
)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 0/1")
assert "[1D DISABLED] (1)" in report
assert "source=diffusion_model.norm.diff" in report
assert "include_1d_diffs is disabled" in report
def test_merge_to_model_report_marks_partially_applied_group(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "partial_base.safetensors"
adapter = tmp_path / "partial_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
"model.diffusion_model.layer.bias": torch.zeros(2),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
"diffusion_model.layer.diff_b": torch.ones(3),
}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"partial", verbose=False, include_1d_diffs=True, return_report=True,
)
result = load_file(path)
torch.testing.assert_close(
result["model.diffusion_model.layer.weight"], torch.ones((2, 2))
)
torch.testing.assert_close(
result["model.diffusion_model.layer.bias"], torch.zeros(2)
)
assert report.startswith("COMPLETED WITH ERRORS: 0/1")
assert "[PARTIALLY APPLIED] (1)" in report
assert "source_shape=(3,), target_shape=(2,)" in report
def test_merge_to_model_detects_raised_adapter_failure_atomically(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "raised_failure_base.safetensors"
adapter = tmp_path / "raised_failure_adapter.safetensors"
base_weight = torch.eye(2)
save_file({"model.diffusion_model.layer.weight": base_weight}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
class RaisedFailureAdapter:
def calculate_weight(self, weight, key, *args, **kwargs):
weight.zero_()
raise RuntimeError("synthetic raised failure")
class FakeLoRAAdapter:
@staticmethod
def load(*args, **kwargs):
return RaisedFailureAdapter()
monkeypatch.setattr(resize, "LoRAAdapter", FakeLoRAAdapter)
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"raised_failure", verbose=False, return_report=True,
)
torch.testing.assert_close(
load_file(path)["model.diffusion_model.layer.weight"], base_weight
)
assert "[CALCULATION FAILURE] (1)" in report
assert "RuntimeError: synthetic raised failure" in report
def test_merge_to_model_report_lists_skipped_tensor_and_pattern(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "skip_base.safetensors"
adapter = tmp_path / "skip_adapter.safetensors"
save_file({
"model.diffusion_model.layer.weight": torch.zeros((2, 2)),
"model.diffusion_model.layer.bias": torch.zeros(2),
}, str(base))
save_file({
"diffusion_model.layer.lora_A.weight": torch.ones((1, 2)),
"diffusion_model.layer.lora_B.weight": torch.ones((2, 1)),
}, str(adapter))
path, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"skip", skip_patterns_str=r"layer\.weight$", verbose=False,
return_report=True,
)
assert "model.diffusion_model.layer.weight" not in load_file(path)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 0/1")
assert "[OMITTED BY SKIP PATTERN] (1)" in report
assert r"pattern=layer\.weight$" in report
assert "[SKIPPED BY PATTERN] (1)" in report
def test_merge_to_model_node_returns_path_blank_line_report(
monkeypatch, modules
):
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"),
)
result = resize.LoRAMergeToModel.execute(
"base.safetensors", "1",
"one.safetensors", 1.0,
"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,
)
assert result.result == ("saved/model.safetensors\n\nSUCCESS\nDETAILS",)
def test_mixed_parser_scans_every_key_and_retains_unpaired(modules):
resize, _ = modules
filler = [f"metadata.tensor_{index}" for index in range(60)]
keys = filler + [
"diffusion_model.linear.lora_down.weight",
"diffusion_model.linear.lora_up.weight",
"diffusion_model.linear.alpha",
"diffusion_model.norm.diff",
"diffusion_model.block.diff_b",
"diffusion_model.orphan.lora_down.weight",
"diffusion_model.linear.dora_scale",
]
info = resize.detect_lora_format(keys)
layers, passthrough = resize.parse_lora_layers(keys)
assert info["format"] == "mixed"
assert layers["diffusion_model.linear"]["down"].endswith("lora_down.weight")
assert layers["diffusion_model.norm"]["diff"].endswith(".diff")
assert layers["diffusion_model.block"]["diff_b"].endswith(".diff_b")
assert "diffusion_model.orphan.lora_down.weight" in passthrough
assert layers["diffusion_model.linear"]["dora_scale"].endswith(".dora_scale")
def test_dtype_contract(modules):
resize, _ = modules
assert resize.select_output_dtype([torch.bfloat16], torch.float16, is_1d_diff=True) == torch.float32
assert resize.select_output_dtype([torch.float32, torch.bfloat16], torch.float16) == torch.float32
assert resize.select_output_dtype([torch.float32], torch.bfloat16, force=True) == torch.bfloat16
def test_include_1d_switches_are_appended_and_default_off(modules, generic_modules):
resize, multi = modules
generic, _ = generic_modules
node_classes = [
generic.LoRATwoMerger,
generic.LoRAThreeMerger,
multi.LoRAMultiMerge,
multi.LoRAMultiMergeDARE,
multi.LoRAMultiMergeDAREEnhanced,
resize.LoRAMergeToModel,
]
for node_class in node_classes:
final_input = node_class.define_schema().inputs[-1]
assert final_input.id == "include_1d_diffs"
assert final_input.default is False
assert len(resize.LoRAMergeToModel.define_schema().outputs) == 1
def test_resize_preserves_mixed_direct_auxiliary_and_fp32(monkeypatch, tmp_path, modules):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "mixed.safetensors"
save_file({
"diffusion_model.linear.lora_down.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4),
"diffusion_model.linear.lora_up.weight": torch.arange(8, dtype=torch.float32).reshape(4, 2),
"diffusion_model.linear.alpha": torch.tensor(2.0, dtype=torch.float32),
"diffusion_model.norm.diff": torch.tensor([0.125, -0.25], dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(0.5, dtype=torch.bfloat16),
"diffusion_model.lin.diff": torch.full((2, 2), 1e-8, dtype=torch.float32),
"diffusion_model.linear.dora_scale": torch.ones(4, dtype=torch.float32),
"diffusion_model.orphan.lora_down.weight": torch.ones((1, 2), dtype=torch.bfloat16),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "resized", verbose=False
)
tensors = load_file(output)
assert tensors["diffusion_model.linear.lora_A.weight"].shape[0] == 2
assert tensors["diffusion_model.linear.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.linear.lora_B.weight"].dtype == torch.float32
assert tensors["diffusion_model.linear.alpha"].dtype == torch.float32
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.scale.diff"].dtype == torch.float16
assert tensors["diffusion_model.lin.diff"].dtype == torch.float32
assert tensors["diffusion_model.linear.dora_scale"].dtype == torch.float32
assert tensors["diffusion_model.orphan.lora_down.weight"].dtype == torch.float16
pure_source = tmp_path / "pure_diff.safetensors"
save_file({
"diffusion_model.pure.diff": torch.ones((2, 2), dtype=torch.bfloat16),
"diffusion_model.pure_norm.diff": torch.ones(2, dtype=torch.bfloat16),
}, str(pure_source))
pure_output = resize.resize_lora_file(
str(pure_source), 1, None, None, "cpu", torch.float16, "pure_resized", verbose=False
)
pure_tensors = load_file(pure_output)
assert set(pure_tensors) == {"diffusion_model.pure.diff", "diffusion_model.pure_norm.diff"}
assert pure_tensors["diffusion_model.pure.diff"].dtype == torch.float16
assert pure_tensors["diffusion_model.pure_norm.diff"].dtype == torch.float32
def test_resize_normalizes_mochi_pair_to_preferred_output(
monkeypatch, tmp_path, modules
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "mochi.safetensors"
save_file({
"diffusion_model.layer.lora_A": torch.ones((2, 2)),
"diffusion_model.layer.lora_B": torch.ones((2, 2)),
"diffusion_model.layer.alpha": torch.tensor(2.0),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float32,
"mochi_normalized", verbose=False,
)
tensors = load_file(output)
assert set(tensors) == {
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
"diffusion_model.layer.alpha",
}
def test_standard_multi_merge_normalizes_mochi_pair_to_preferred_output(
monkeypatch, tmp_path, modules
):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "multi_mochi_a.safetensors"
second = tmp_path / "multi_mochi_b.safetensors"
for path, value in ((first, 1.0), (second, 3.0)):
save_file({
"diffusion_model.layer.lora_A": torch.full((1, 2), value),
"diffusion_model.layer.lora_B": torch.full((2, 1), value),
"diffusion_model.layer.alpha": torch.tensor(1.0),
}, str(path))
output = merger.merge_multi_loras(
[str(first), str(second)], [1.0, 1.0], "weighted_sum", "cpu",
torch.float32, "multi_mochi_normalized", verbose=False,
)
tensors = load_file(output)
assert set(tensors) == {
"diffusion_model.layer.lora_A.weight",
"diffusion_model.layer.lora_B.weight",
"diffusion_model.layer.alpha",
}
def test_standard_multi_merge_emits_weighted_direct_union(monkeypatch, tmp_path, modules):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "first.safetensors"
second = tmp_path / "second.safetensors"
base = tmp_path / "multi_base.safetensors"
save_file({
"model.diffusion_model.norm.weight": torch.zeros(2),
"model.diffusion_model.matrix.weight": torch.zeros((2, 2)),
"model.diffusion_model.only_first.weight": torch.zeros(()),
"model.diffusion_model.scale": torch.zeros(()),
"model.diffusion_model.lowrank.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.ones((2, 2), dtype=torch.float32),
"diffusion_model.only_first.diff": torch.tensor(4.0, dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(1.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.float32),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.float32),
"diffusion_model.lowrank.alpha": torch.tensor(1.0, dtype=torch.float32),
}, str(first))
save_file({
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.full((2, 2), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.full((1, 2), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_up.weight": torch.full((2, 1), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.alpha": torch.tensor(1.0, dtype=torch.bfloat16),
}, str(second))
output = merger.merge_multi_loras(
[str(first), str(second)], [2.0, -1.0], "concatenate", "cpu",
torch.float16, "merged", base_model_path=str(base), verbose=False,
include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["diffusion_model.norm.diff"], torch.tensor([-1.0, 0.0]))
torch.testing.assert_close(tensors["diffusion_model.matrix.diff"], torch.zeros((2, 2)))
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.matrix.diff"].dtype == torch.float32
assert tensors["diffusion_model.only_first.diff"].item() == pytest.approx(8.0)
assert tensors["diffusion_model.only_first.diff"].dtype == torch.float16
assert tensors["diffusion_model.scale.diff"].item() == pytest.approx(2.0)
assert tensors["diffusion_model.lowrank.lora_A.weight"].shape[0] == 2
assert tensors["diffusion_model.lowrank.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_B.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.alpha"].dtype == torch.float32
default_output = merger.merge_multi_loras(
[str(first), str(second)], [2.0, -1.0], "concatenate", "cpu",
torch.float16, "merged_default", base_model_path=str(base), verbose=False,
)
default_tensors = load_file(default_output)
assert "diffusion_model.norm.diff" not in default_tensors
assert "diffusion_model.matrix.diff" in default_tensors
assert "diffusion_model.scale.diff" in default_tensors
assert "diffusion_model.lowrank.lora_A.weight" in default_tensors
def test_multi_merge_rejects_direct_shape_mismatch(monkeypatch, tmp_path, modules):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "shape_first.safetensors"
second = tmp_path / "shape_second.safetensors"
save_file({"diffusion_model.bad.diff": torch.ones(2)}, str(first))
save_file({"diffusion_model.bad.diff": torch.ones(3)}, str(second))
with pytest.raises(ValueError, match="Direct LoRA shape mismatch"):
merger.merge_multi_loras(
[str(first), str(second)], [1.0, 1.0], "weighted_sum", "cpu",
torch.float16, "bad_shapes", verbose=False, include_1d_diffs=True,
)
@pytest.mark.parametrize("enhanced", [False, True])
def test_dare_variants_emit_direct_layers(monkeypatch, tmp_path, modules, enhanced):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / f"dare_first_{enhanced}.safetensors"
second = tmp_path / f"dare_second_{enhanced}.safetensors"
save_file({
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.float32),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.float32),
}, str(first))
save_file({
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.float32),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.bfloat16),
}, str(second))
if enhanced:
output = merger.merge_multi_loras_dare_enhanced(
[str(first), str(second)], [1.0, 1.0], 1.0, 1.0, 0.0, 0.0, 7,
"cpu", torch.float16, "enhanced", verbose=False, include_1d_diffs=True,
)
else:
output = merger.merge_multi_loras_dare(
[str(first), str(second)], [1.0, 1.0], 0.0, 0.0, 7,
"cpu", torch.float16, "dare", verbose=False, include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["diffusion_model.norm.diff"], torch.tensor([2.0, 3.0]))
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_A.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_B.weight"].dtype == torch.float32
if enhanced:
default_output = merger.merge_multi_loras_dare_enhanced(
[str(first), str(second)], [1.0, 1.0], 1.0, 1.0, 0.0, 0.0, 7,
"cpu", torch.float16, "enhanced_default", verbose=False,
)
else:
default_output = merger.merge_multi_loras_dare(
[str(first), str(second)], [1.0, 1.0], 0.0, 0.0, 7,
"cpu", torch.float16, "dare_default", verbose=False,
)
default_tensors = load_file(default_output)
assert "diffusion_model.norm.diff" not in default_tensors
assert "diffusion_model.lowrank.lora_A.weight" in default_tensors
def test_merge_to_model_applies_generic_direct_keys_and_preserves_dtype(
monkeypatch, tmp_path, modules, caplog
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "base.safetensors"
adapter = tmp_path / "adapter.safetensors"
save_file({
"model.diffusion_model.block.weight": torch.zeros((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.norm.weight": torch.zeros(2, dtype=torch.bfloat16),
"model.diffusion_model.block.bias": torch.zeros(2, dtype=torch.bfloat16),
"model.diffusion_model.scale": torch.tensor(1.0, dtype=torch.bfloat16),
"model.diffusion_model.lin": torch.zeros((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.bad.weight": torch.ones((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.untouched": torch.tensor(1e-8, dtype=torch.float32),
}, str(base))
save_file({
"diffusion_model.block.diff": torch.full((2, 2), 1e-8, dtype=torch.float32),
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.block.diff_b": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(0.5, dtype=torch.bfloat16),
"diffusion_model.lin.diff": torch.ones((2, 2), dtype=torch.float32),
"diffusion_model.bad.diff": torch.ones(3, dtype=torch.float32),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16, "merged_model",
verbose=False, include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["model.diffusion_model.norm.weight"], torch.tensor([0.5, 1.0]))
torch.testing.assert_close(tensors["model.diffusion_model.block.bias"], torch.tensor([1.5, 2.0]))
assert tensors["model.diffusion_model.scale"].item() == pytest.approx(1.25)
torch.testing.assert_close(tensors["model.diffusion_model.lin"], torch.full((2, 2), 0.5))
torch.testing.assert_close(tensors["model.diffusion_model.bad.weight"].float(), torch.ones((2, 2)))
assert "Shape mismatch" in caplog.text
assert tensors["model.diffusion_model.block.weight"].dtype == torch.float32
assert tensors["model.diffusion_model.norm.weight"].dtype == torch.float32
assert tensors["model.diffusion_model.block.bias"].dtype == torch.float32
assert tensors["model.diffusion_model.scale"].dtype == torch.float16
assert tensors["model.diffusion_model.lin"].dtype == torch.float32
assert tensors["model.diffusion_model.untouched"].dtype == torch.float32
default_output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16, "merged_model_default",
verbose=False,
)
default_tensors = load_file(default_output)
torch.testing.assert_close(
default_tensors["model.diffusion_model.norm.weight"].float(), torch.zeros(2)
)
torch.testing.assert_close(
default_tensors["model.diffusion_model.block.bias"].float(), torch.zeros(2)
)
assert default_tensors["model.diffusion_model.scale"].item() == pytest.approx(1.25)
torch.testing.assert_close(
default_tensors["model.diffusion_model.lin"], torch.full((2, 2), 0.5)
)
def _generic_params(output_filename, override_dtype, include_1d_diffs=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": output_filename,
"save_dtype": "bf16",
"override_dtype": override_dtype,
"device": "cpu",
"dtype": torch.float32,
"exclude_patterns": "",
"discard_patterns": "",
"glob_patterns": False,
"lazy_load": True,
"force_clear_cache": False,
"include_1d_diffs": include_1d_diffs,
}
def test_generic_two_and_three_mergers_enforce_direct_dtype_contract(
monkeypatch, tmp_path, generic_modules
):
generic, operations = generic_modules
paths = {}
for name, tensors in {
"a": {
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.ones((2, 2), dtype=torch.float32),
},
"b": {
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.full((2, 2), 3.0, dtype=torch.bfloat16),
"diffusion_model.secondary_only.diff": torch.ones(1, dtype=torch.float32),
},
"c": {
"diffusion_model.norm.diff": torch.tensor([0.5, 1.0], dtype=torch.float32),
"diffusion_model.matrix.diff": torch.full((2, 2), 0.5, dtype=torch.bfloat16),
},
}.items():
path = tmp_path / f"{name}.safetensors"
save_file(tensors, str(path))
paths[name] = str(path)
monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths.get(name))
monkeypatch.setattr(generic.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(generic.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None)
two_params = _generic_params("generic_two", override_dtype=True)
two_params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"},
"Weight-Sum",
operations.TWO_MODEL_MODES,
two_params,
"loras",
)
two = load_file(str(tmp_path / "loras" / "generic_two.safetensors"))
torch.testing.assert_close(
two["diffusion_model.norm.diff"].float(), torch.tensor([1.0, 2.0])
)
assert two["diffusion_model.norm.diff"].dtype == torch.bfloat16
assert two["diffusion_model.matrix.diff"].dtype == torch.bfloat16
assert "diffusion_model.secondary_only.diff" not in two
two_enabled_params = _generic_params(
"generic_two_enabled", override_dtype=True, include_1d_diffs=True
)
two_enabled_params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"},
"Weight-Sum",
operations.TWO_MODEL_MODES,
two_enabled_params,
"loras",
)
two_enabled = load_file(
str(tmp_path / "loras" / "generic_two_enabled.safetensors")
)
torch.testing.assert_close(
two_enabled["diffusion_model.norm.diff"], torch.tensor([2.0, 3.0])
)
assert two_enabled["diffusion_model.norm.diff"].dtype == torch.float32
three_params = _generic_params("generic_three", override_dtype=False)
three_params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"},
"Add-Difference",
operations.THREE_MODEL_MODES,
three_params,
"loras",
)
three = load_file(str(tmp_path / "loras" / "generic_three.safetensors"))
torch.testing.assert_close(
three["diffusion_model.norm.diff"].float(), torch.tensor([1.0, 2.0])
)
assert three["diffusion_model.norm.diff"].dtype == torch.bfloat16
assert three["diffusion_model.matrix.diff"].dtype == torch.float32
three_enabled_params = _generic_params(
"generic_three_enabled", override_dtype=False, include_1d_diffs=True
)
three_enabled_params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"},
"Add-Difference",
operations.THREE_MODEL_MODES,
three_enabled_params,
"loras",
)
three_enabled = load_file(
str(tmp_path / "loras" / "generic_three_enabled.safetensors")
)
torch.testing.assert_close(
three_enabled["diffusion_model.norm.diff"], torch.tensor([2.25, 3.5])
)
assert three_enabled["diffusion_model.norm.diff"].dtype == torch.float32
+437
View File
@@ -0,0 +1,437 @@
import importlib
import json
import logging
import struct
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_quant_guard_tests"
def _load_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 {
"guard": _load_module("nodes.quantization_guard"),
"extract": _load_module("nodes.lora_extract_svd"),
"resize": _load_module("nodes.lora_resize"),
"multi": _load_module("nodes.lora_merger"),
"generic": _load_module("nodes.merger"),
"operations": _load_module("nodes.merger_ops"),
}
def _patch_output(monkeypatch, module, output_dir):
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(output_dir)])
monkeypatch.setattr(module.folder_paths, "models_dir", str(output_dir))
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
def _inspect_file(path, modules, label="input"):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
return modules["guard"].inspect_low_bit_input(handler, label, "Guard Test")
finally:
handler.__exit__(None, None, None)
def _round_trip_preserved(path, output, modules):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
with modules["resize"].IncrementalSafetensorsWriter(str(output)) as writer:
modules["guard"].write_preserved_tensor(writer, "low", handler)
finally:
handler.__exit__(None, None, None)
return load_file(str(output))["low"]
def _tensor_payload(path, key="low"):
with open(path, "rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
header = json.loads(source.read(header_size).decode("utf-8"))
start, end = header[key]["data_offsets"]
source.seek(8 + header_size + start)
return source.read(end - start)
@pytest.mark.parametrize("dtype", [torch.int8, torch.uint8])
def test_int8_uint8_are_detected_and_warn_once(tmp_path, modules, caplog, dtype):
path = tmp_path / f"{dtype}.safetensors"
save_file({"low": torch.arange(4, dtype=dtype), "normal": torch.ones(1)}, str(path))
with caplog.at_level(logging.WARNING):
detected = _inspect_file(path, modules, "threshold input")
assert detected == {"low"}
records = [record for record in caplog.records if "threshold input" in record.message]
assert len(records) == 1
assert "1 isolated low-bit tensor" in records[0].message
assert "low" in records[0].message
def test_exact_two_tensor_threshold(tmp_path, modules):
for count in range(4):
path = tmp_path / f"threshold_{count}.safetensors"
tensors = {f"low_{i}": torch.ones(2, dtype=torch.uint8) for i in range(count)}
tensors["normal"] = torch.ones(1)
save_file(tensors, str(path))
if count <= 2:
assert len(_inspect_file(path, modules)) == count
else:
with pytest.raises(ValueError, match="3 low-bit tensors.*casted or quantized"):
_inspect_file(path, modules)
@pytest.mark.parametrize(
("tensors", "metadata"),
[
({"layer.comfy_quant": torch.tensor([123], dtype=torch.uint8)}, None),
({"normal": torch.ones(1)}, {"_quantization_metadata": "{}"}),
({"scaled_fp8": torch.ones(2)}, None),
],
)
def test_recognized_comfy_quantization_hard_errors(tmp_path, modules, tensors, metadata):
path = tmp_path / f"quant_{len(list(tmp_path.iterdir()))}.safetensors"
save_file(tensors, str(path), metadata=metadata)
with pytest.raises(ValueError, match="ComfyUI quantization metadata.*unsupported"):
_inspect_file(path, modules, "recognized input")
def test_supported_fp8_storage_round_trip(tmp_path, modules):
candidates = [
dtype
for name in (
"float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz",
"float8_e8m0fnu",
)
if (dtype := getattr(torch, name, None)) is not None
]
supported = []
for index, dtype in enumerate(candidates):
path = tmp_path / f"low_float_{index}.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
continue
supported.append(dtype)
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / f"preserved_{index}.safetensors", modules)
assert preserved.dtype == dtype
torch.testing.assert_close(preserved.float(), loaded.float())
if candidates and not supported:
pytest.skip("Installed safetensors cannot write any installed torch FP8 dtype")
def test_fp4_storage_when_supported(tmp_path, modules):
dtype = getattr(torch, "float4_e2m1fn_x2", None)
if dtype is None:
pytest.skip("Installed PyTorch has no torch.float4_e2m1fn_x2")
path = tmp_path / "fp4.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
pytest.skip("Installed safetensors cannot write torch.float4_e2m1fn_x2")
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / "preserved_fp4.safetensors", modules)
assert preserved.dtype == dtype
assert _tensor_payload(path) == _tensor_payload(tmp_path / "preserved_fp4.safetensors")
def test_realistic_packed_uint8_comfy_sidecar_is_rejected(tmp_path, modules):
path = tmp_path / "packed_comfy.safetensors"
save_file({
"layer.weight": torch.arange(8, dtype=torch.uint8).reshape(2, 4),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
"layer.weight_scale": torch.ones(2),
}, str(path))
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
_inspect_file(path, modules)
def test_resize_preserves_complete_low_bit_layer_and_resizes_normal(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "resize_quant.safetensors"
guarded_down = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"diffusion_model.guarded.lora_down.weight": guarded_down,
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 2)),
"diffusion_model.guarded.alpha": torch.tensor(2.0),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.lora_down.weight": torch.ones((2, 2)),
"diffusion_model.normal.lora_up.weight": torch.ones((2, 2)),
"isolated_aux": torch.tensor([7], dtype=torch.int8),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "resize_guarded", verbose=False
)
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_A.weight"], guarded_down)
assert result["diffusion_model.guarded.lora_A.weight"].dtype == torch.uint8
assert result["diffusion_model.guarded.lora_B.weight"].dtype == torch.float32
assert result["diffusion_model.guarded.dora_scale"].dtype == torch.float32
assert result["diffusion_model.normal.lora_A.weight"].shape[0] == 1
assert result["isolated_aux"].dtype == torch.int8
def test_extract_excludes_low_bit_candidate_and_extracts_float_layer(
monkeypatch, tmp_path, modules
):
extract = modules["extract"]
monkeypatch.setattr(extract, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(extract, "cleanup_after_operation", lambda: None)
model_a = tmp_path / "extract_a.safetensors"
model_b = tmp_path / "extract_b.safetensors"
output = tmp_path / "extracted.safetensors"
save_file({
"model.diffusion_model.normal.weight": torch.eye(2),
"model.diffusion_model.quant.weight": torch.ones((2, 2), dtype=torch.uint8),
}, str(model_a))
save_file({
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
"model.diffusion_model.quant.weight": torch.zeros((2, 2), dtype=torch.uint8),
}, str(model_b))
extract.extract_lora_from_files(
str(model_a), str(model_b), "fixed", 1, 1, "cpu", "fp32", str(output),
force_clear_cache=False,
)
result = load_file(str(output))
assert any("normal" in key for key in result)
assert not any("quant" in key for key in result)
def test_extract_bias_diff_round_trips_through_merge_to_model(
monkeypatch, tmp_path, modules
):
extract = modules["extract"]
resize = modules["resize"]
monkeypatch.setattr(extract, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(extract, "cleanup_after_operation", lambda: None)
_patch_output(monkeypatch, resize, tmp_path)
finetuned = tmp_path / "bias_finetuned.safetensors"
base = tmp_path / "bias_base.safetensors"
adapter = tmp_path / "bias_adapter.safetensors"
save_file({
"model.diffusion_model.layer.bias": torch.tensor([0.25, -0.5]),
}, str(finetuned))
save_file({
"model.diffusion_model.layer.bias": torch.tensor([0.0, 0.5]),
}, str(base))
extract.extract_lora_from_files(
str(finetuned), str(base), "fixed", 1, 1, "cpu", "fp16", str(adapter),
force_clear_cache=False, include_1d_diffs=True,
)
extracted = load_file(str(adapter))
assert set(extracted) == {"diffusion_model.layer.diff_b"}
assert extracted["diffusion_model.layer.diff_b"].dtype == torch.float32
merged_path = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float32,
"bias_round_trip", verbose=False, include_1d_diffs=True,
)
merged = load_file(merged_path)
torch.testing.assert_close(
merged["model.diffusion_model.layer.bias"],
load_file(str(finetuned))["model.diffusion_model.layer.bias"],
)
def _run_multi_variant(module, variant, paths, output_name):
common = (paths, [1.0, 1.0])
if variant == "standard":
return module.merge_multi_loras(
*common, "weighted_sum", "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
if variant == "dare":
return module.merge_multi_loras_dare(
*common, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
return module.merge_multi_loras_dare_enhanced(
*common, 1.0, 1.0, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
@pytest.mark.parametrize("variant", ["standard", "dare", "enhanced"])
def test_multi_variants_preserve_earliest_complete_affected_layer(
monkeypatch, tmp_path, modules, variant
):
multi = modules["multi"]
_patch_output(monkeypatch, multi, tmp_path)
first = tmp_path / f"{variant}_first.safetensors"
second = tmp_path / f"{variant}_second.safetensors"
first_down = torch.tensor([[1.0, 2.0]])
save_file({
"diffusion_model.guarded.lora_down.weight": first_down,
"diffusion_model.guarded.lora_up.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(first))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 2]], dtype=torch.uint8),
"diffusion_model.guarded.lora_up.weight": torch.tensor([[9.0], [9.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([9.0, 9.0]),
"diffusion_model.normal.diff": torch.full((2, 2), 3.0),
}, str(second))
output = _run_multi_variant(multi, variant, [str(first), str(second)], f"guard_{variant}")
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_A.weight"], first_down)
torch.testing.assert_close(
result["diffusion_model.guarded.lora_B.weight"], torch.tensor([[3.0], [4.0]])
)
torch.testing.assert_close(
result["diffusion_model.guarded.dora_scale"], torch.tensor([0.125, 0.25])
)
assert "diffusion_model.normal.diff" in result
def test_hard_error_happens_before_resize_output_creation(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "three_low.safetensors"
save_file({f"low_{i}": torch.ones(1, dtype=torch.uint8) for i in range(3)}, str(source))
output = tmp_path / "must_not_exist.safetensors"
with pytest.raises(ValueError, match="3 low-bit tensors"):
resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
)
assert not output.exists()
def test_recognized_quantization_creates_no_resize_output(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "recognized_quant.safetensors"
save_file({
"layer.weight": torch.ones((2, 2)),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
}, str(source))
output = tmp_path / "recognized_must_not_exist.safetensors"
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
)
assert not output.exists()
def _generic_params(output_filename):
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": output_filename, "save_dtype": "bf16",
"override_dtype": True, "device": "cpu", "dtype": torch.float32,
"exclude_patterns": "", "discard_patterns": "", "glob_patterns": False,
"lazy_load": True, "force_clear_cache": False, "include_1d_diffs": True,
}
def test_generic_merge_preserves_model_a_when_either_source_key_is_low_bit(
monkeypatch, tmp_path, modules
):
generic = modules["generic"]
operations = modules["operations"]
model_a = tmp_path / "generic_guard_a.safetensors"
model_b = tmp_path / "generic_guard_b.safetensors"
a_tensors = {
"a_low": torch.tensor([[1, 2]], dtype=torch.uint8),
"b_low": torch.tensor([[0.125, 0.25]], dtype=torch.float32),
"normal": torch.ones((1, 2)),
}
save_file(a_tensors, str(model_a))
save_file({
"a_low": torch.tensor([[9.0, 9.0]]),
"b_low": torch.tensor([[3, 4]], dtype=torch.int8),
"normal": torch.full((1, 2), 3.0),
}, str(model_b))
paths = {"a": str(model_a), "b": str(model_b)}
monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(generic.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(generic.folder_paths, "models_dir", str(tmp_path))
monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None)
params = _generic_params("generic_guarded")
params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"}, "Weight-Sum",
operations.TWO_MODEL_MODES, params, "loras",
)
result = load_file(
str(tmp_path / "loras" / "generic_guarded.safetensors")
)
torch.testing.assert_close(result["a_low"], a_tensors["a_low"])
torch.testing.assert_close(result["b_low"], a_tensors["b_low"])
assert result["a_low"].dtype == torch.uint8
assert result["b_low"].dtype == torch.float32
torch.testing.assert_close(result["normal"].float(), torch.full((1, 2), 2.0))
def test_merge_to_model_preserves_low_bit_base_and_skips_low_bit_adapter_layer(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "merge_guard_base.safetensors"
adapter = tmp_path / "merge_guard_adapter.safetensors"
base_quant = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"model.diffusion_model.base_quant.weight": base_quant,
"model.diffusion_model.guarded.weight": torch.zeros((2, 2)),
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 1]], dtype=torch.int8),
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 1)),
"diffusion_model.guarded.diff": torch.full((2, 2), 5.0),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(adapter))
output, report = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16, "merge_guarded",
verbose=False, include_1d_diffs=True, return_report=True,
)
result = load_file(output)
torch.testing.assert_close(result["model.diffusion_model.base_quant.weight"], base_quant)
assert result["model.diffusion_model.base_quant.weight"].dtype == torch.uint8
torch.testing.assert_close(
result["model.diffusion_model.guarded.weight"].float(), torch.zeros((2, 2))
)
torch.testing.assert_close(
result["model.diffusion_model.normal.weight"].float(), torch.ones((2, 2))
)
assert report.startswith("COMPLETED WITH EXCLUSIONS: 1/2")
assert "[GUARDED] (1)" in report
assert "diffusion_model.guarded.lora_down.weight" in report
assert "torch.int8" in report
assert "model.diffusion_model.base_quant.weight" in report
+282
View File
@@ -0,0 +1,282 @@
import importlib
import math
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_analysis_tests"
def _load_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 analysis():
return _load_module("nodes.model_analysis")
def _params(**overrides):
params = {
"cwb_similarity_alignment": False,
"top_weight_differences": 3,
"process_device": "cpu",
"force_clear_cache": False,
}
params.update(overrides)
return params
def test_five_node_schema_contract(analysis):
assert len(analysis.MODEL_ANALYSIS_NODES) == 5
assert len({node.NODE_ID for node in analysis.MODEL_ANALYSIS_NODES}) == 5
for node in analysis.MODEL_ANALYSIS_NODES:
schema = node.define_schema()
assert schema.category == "ModelUtils/Analysis"
assert [item.id for item in schema.inputs[:3]] == [
"execution_mode", "model_a", "model_b",
]
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]
if node in (analysis.LoRAModelAnalysis, analysis.EmbeddingModelAnalysis):
assert input_ids[3] == "cwb_similarity_alignment"
else:
assert "cwb_similarity_alignment" not in input_ids
assert not {
"cwb_preset", "consensus_type", "alignment_threshold",
"similarity_threshold", "power_alpha", "diversity_beta",
"dynamic_similarity_contrast", "soft_comfort_bandpass",
"position_weight", "preserve_common_prefix",
} & set(input_ids)
assert [item.display_name for item in schema.outputs] == [
"comparison_report", "cwb_report", "documentation",
]
def test_documentation_mode_has_three_separate_outputs(analysis):
result = analysis.DiffusionModelAnalysis.execute(execution_mode="DOCUMENTATION ONLY")
assert len(result.result) == 3
assert "No comparison performed" in result[0]
assert "No CWB analysis performed" in result[1]
assert result[2].startswith("# Two-Model Similarity Analysis")
def test_raw_metrics_have_independent_expected_values(analysis):
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[1.0, 0.0], [5.0, 4.0]])
raw, top = analysis._compute_raw(a, b, 2)
metrics = analysis._metrics(raw)
assert metrics["mae"] == pytest.approx(1.0)
assert metrics["mse"] == pytest.approx(2.0)
assert metrics["rmse"] == pytest.approx(math.sqrt(2.0))
assert metrics["max"] == pytest.approx(2.0)
assert metrics["exact"] == pytest.approx(0.5)
assert len(top) == 2
assert {entry[1] for entry in top} == {1, 2}
def test_cwb_derived_metrics_have_independent_expected_values(analysis):
a = torch.tensor([[1.0, 0.0], [1.0, 0.0]])
b = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
values = analysis._compute_cwb(a, b, False).values()
assert values["pair_mean"] == pytest.approx(0.5)
assert values["pair_min"] == pytest.approx(0.0)
assert values["pair_max"] == pytest.approx(1.0)
assert values["mean_aff_a"] == pytest.approx((1.0 + 2 ** -0.5) / 2)
assert values["mean_aff_b"] == pytest.approx((1.0 + 2 ** -0.5) / 2)
assert values["median_aff_a"] == pytest.approx(0.5)
assert values["median_aff_b"] == pytest.approx(0.5)
assert values["alignment"] is None
def test_zero_constant_and_nonfinite_metrics_are_explicit(analysis):
raw, _ = analysis._compute_raw(torch.zeros(3), torch.zeros(3), 0)
metrics = analysis._metrics(raw)
assert metrics["cosine"] == pytest.approx(1.0)
assert metrics["pearson"] is None
assert metrics["norm_ratio"] == pytest.approx(1.0)
raw, _ = analysis._compute_raw(
torch.tensor([1.0, float("nan"), float("inf")]),
torch.tensor([2.0, 4.0, 5.0]),
0,
)
metrics = analysis._metrics(raw)
assert raw.finite == 1
assert raw.nonfinite_a == 2
assert metrics["coverage"] == pytest.approx(1 / 3)
def test_end_to_end_reports_and_topology_classification(monkeypatch, tmp_path, analysis):
path_a = tmp_path / "a.safetensors"
path_b = tmp_path / "b.safetensors"
save_file({
"model.blocks.0.weight": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
"only_a": torch.ones(1),
"shape": torch.ones(2),
"integer": torch.tensor([1, 2], dtype=torch.int64),
}, str(path_a), metadata={"source": "a"})
save_file({
"model.blocks.0.weight": torch.tensor([[1.0, 0.0], [5.0, 4.0]]),
"only_b": torch.ones(1),
"shape": torch.ones(3),
"integer": torch.tensor([1, 3], dtype=torch.int64),
}, str(path_b), metadata={"source": "b"})
paths = {"a": str(path_a), "b": str(path_b)}
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)
processed = []
async_calls = []
original_mark_processed = analysis.MemoryEfficientSafeOpen.mark_processed
original_async_stream = analysis.MemoryEfficientSafeOpen.async_stream
def track_processed(loader, key):
processed.append((loader.filename, key))
return original_mark_processed(loader, key)
def track_async_stream(loader, keys, **kwargs):
async_calls.append((loader.filename, tuple(keys), kwargs))
yield from original_async_stream(loader, keys, **kwargs)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", track_processed)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_async_stream)
comparison, cwb_report = analysis.ModelAnalysisLogic.execute(
"a", "b", "diffusion_models", _params()
)
assert "Comparable floating tensors: 1" in comparison
assert "mae: 1" in comparison
assert "mse: 2" in comparison
assert "[A ONLY] (1)" in comparison
assert "[B ONLY] (1)" in comparison
assert "[SHAPE MISMATCH] (1)" in comparison
assert "[NON-FLOATING] (1)" in comparison
assert "model.blocks.0 |" in comparison
assert "CWB-DERIVED SIMILARITY SUMMARY" in cwb_report
assert "merge weighting" in cwb_report
assert "PRESET-IMPLIED CONTRIBUTIONS" not in cwb_report
assert "Average contribution" not in cwb_report
assert len(processed) == 4
assert {key for _, key in processed} == {"model.blocks.0.weight", "integer"}
assert len(async_calls) == 4
assert all(len(call[1]) == 1 for call in async_calls)
assert all(call[2]["batch_size"] == 1 for call in async_calls)
assert all(call[2]["prefetch_batches"] == 1 for call in async_calls)
assert not list(tmp_path.glob("**/*analysis*"))
def test_lora_cross_format_keys_compare_as_one_logical_pair(monkeypatch, tmp_path, analysis):
path_a = tmp_path / "lora_a.safetensors"
path_b = tmp_path / "lora_b.safetensors"
save_file({
"diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
"diffusion_model.foo.lora_B.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
}, str(path_a))
save_file({
"base_model.model.diffusion_model.foo.lora_down.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
"base_model.model.diffusion_model.foo.lora_up.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
}, str(path_b))
paths = {"a": str(path_a), "b": str(path_b)}
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)
processed = []
async_calls = []
original_mark_processed = analysis.MemoryEfficientSafeOpen.mark_processed
original_async_stream = analysis.MemoryEfficientSafeOpen.async_stream
def track_processed(loader, key):
processed.append((loader.filename, key))
return original_mark_processed(loader, key)
def track_async_stream(loader, keys, **kwargs):
async_calls.append((loader.filename, tuple(keys), kwargs))
yield from original_async_stream(loader, keys, **kwargs)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "mark_processed", track_processed)
monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_async_stream)
comparison, cwb_report = analysis.ModelAnalysisLogic.execute(
"a", "b", "loras", _params(), lora_mode=True,
)
assert "Comparable floating tensors: 2" in comparison
assert "[A ONLY] (0)" in comparison
assert "[B ONLY] (0)" in comparison
assert "foo.down |" in comparison
assert "foo.up |" in comparison
assert "foo.[down+up] |" in cwb_report
assert len(processed) == 4
assert len({entry for entry in processed}) == 4
assert len(async_calls) == 2
assert all(len(call[1]) == 2 for call in async_calls)
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]])
aligned = analysis._compute_cwb(a, b, True).values()
indexed = analysis._compute_cwb(a, b, False).values()
assert aligned["alignment"] == pytest.approx(1.0)
assert aligned["alignment_gain"] == pytest.approx(1.0)
assert indexed["pair_mean"] == pytest.approx(0.0)
assert indexed["alignment"] is None
def test_cuda_oom_retries_once_on_cpu(monkeypatch, analysis):
calls = []
def fake_compute(a, b, device, allow_alignment, top_count):
calls.append(device)
if device == "cuda":
raise torch.cuda.OutOfMemoryError("synthetic CUDA out of memory")
return analysis.RawStats(finite=1, total=1), analysis.CWBStats(), []
monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute)
monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None)
result = analysis._analyze_with_fallback(
"layer", torch.ones(1), torch.ones(1), "cuda", False, 1,
)
assert calls == ["cuda", "cpu"]
assert result[-1] is True
def test_non_oom_error_is_not_retried(monkeypatch, analysis):
calls = []
def fake_compute(*args):
calls.append(args[2])
raise RuntimeError("unrelated failure")
monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute)
with pytest.raises(RuntimeError, match="unrelated failure"):
analysis._analyze_with_fallback(
"layer", torch.ones(1), torch.ones(1), "cuda", False, 1,
)
assert calls == ["cuda"]
def test_failed_cpu_retry_names_layer_and_chains_oom(monkeypatch, analysis):
def fake_compute(a, b, device, allow_alignment, top_count):
if device == "cuda":
raise torch.cuda.OutOfMemoryError("synthetic CUDA out of memory")
raise ValueError("CPU calculation failed")
monkeypatch.setattr(analysis, "_compute_pair_on_device", fake_compute)
monkeypatch.setattr(analysis, "_release_cuda_oom", lambda: None)
with pytest.raises(RuntimeError, match="CPU fallback failed while analyzing 'named.layer'") as caught:
analysis._analyze_with_fallback(
"named.layer", torch.ones(1), torch.ones(1), "cuda", False, 1,
)
assert isinstance(caught.value.__cause__, torch.cuda.OutOfMemoryError)