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
9 changed files with 2237 additions and 12 deletions
+16
View File
@@ -21,6 +21,7 @@ from .nodes.merger import (
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
@@ -33,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,
@@ -71,6 +79,8 @@ class ModelUtilsExtension(ComfyExtension):
CheckpointThreeMerger, EmbeddingThreeMerger,
# Dedicated CWB Mergers
*CWB_MERGER_NODES,
# Two-Model Analysis
*MODEL_ANALYSIS_NODES,
# LoRA Extraction
LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile,
LoRAExtractKnee, LoRAExtractFrobenius,
@@ -80,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,
+3 -1
View File
@@ -23,7 +23,9 @@ alignment reference.
`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.
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
+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.
+6 -6
View File
@@ -72,17 +72,17 @@ CWB_PRESETS = {
"varied_merge": {
"consensus_type": "median", "power_alpha": 2.0,
"similarity_threshold": 0.0, "diversity_beta": 0.0,
"rescale_norm": True, "global_scale": 0.7,
"rescale_norm": True, "global_scale": 1.0,
},
"diverse_concept": {
"consensus_type": "median", "power_alpha": 2.0,
"similarity_threshold": 0.0, "diversity_beta": 1.0,
"rescale_norm": True, "global_scale": 0.7,
"rescale_norm": True, "global_scale": 1.0,
},
"high_diversity_concept": {
"consensus_type": "median", "power_alpha": 2.0,
"similarity_threshold": 0.0, "diversity_beta": 2.0,
"rescale_norm": True, "global_scale": 0.7,
"rescale_norm": True, "global_scale": 1.0,
},
"dsc_baseline": {
"consensus_type": "median", "power_alpha": 2.0,
@@ -105,19 +105,19 @@ CWB_PRESETS = {
"dsc_varied_merge": {
"consensus_type": "median", "power_alpha": 2.5,
"similarity_threshold": 0.0, "diversity_beta": 0.0,
"rescale_norm": True, "global_scale": 0.7,
"rescale_norm": True, "global_scale": 1.0,
"dynamic_similarity_contrast": True, "soft_comfort_bandpass": True,
},
"dsc_diverse_concept": {
"consensus_type": "median", "power_alpha": 2.0,
"similarity_threshold": 0.0, "diversity_beta": 1.5,
"rescale_norm": True, "global_scale": 0.7,
"rescale_norm": True, "global_scale": 1.0,
"dynamic_similarity_contrast": True, "soft_comfort_bandpass": True,
},
"dsc_high_diversity_concept": {
"consensus_type": "median", "power_alpha": 2.0,
"similarity_threshold": 0.0, "diversity_beta": 3.0,
"rescale_norm": True, "global_scale": 0.7,
"rescale_norm": True, "global_scale": 1.0,
"dynamic_similarity_contrast": True, "soft_comfort_bandpass": True,
},
}
File diff suppressed because it is too large Load Diff
+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.6"
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
+5
View File
@@ -69,6 +69,11 @@ def _params(output_filename, **overrides):
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)
+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)