Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0a7578ce73 | ||
|
|
9ad992d080 | ||
|
|
3bd4977d2a | ||
|
|
10484831d0 | ||
|
|
71b7936f92 | ||
|
|
d9fc1f52f1 | ||
|
|
01570b4805 |
+16
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user