nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(lora): bound knee SVD probes

Browse Source
This commit is contained in:
silveroxides
2026-08-13 16:36:13 +02:00
parent f76195f249
commit 16e493cfdf
5 changed files with 259 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/dora_extract_wd.py
+62 -23
View File
@@ -213,13 +213,34 @@ def _svd_extract_linear_lowrank(
pass
# Construct LoRA matrices
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_up = U * S.unsqueeze(0) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, None), "low rank"
return (lora_down, lora_up, diff), "low rank"
def _svd_extract_knee_lowrank(
weight: torch.Tensor,
mode: str,
max_rank: int,
probe_offset: int,
niter: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
"""Probe beyond the output rank and retry once when the knee hits the probe tail."""
max_possible_rank = min(weight.shape)
output_cap = max(1, min(max_rank, max_possible_rank))
offset = max(1, probe_offset)
probe_rank = min(max_possible_rank, output_cap + offset)
probe_limit = min(max_possible_rank, max(probe_rank, output_cap * 2))
while True:
U, S, V = torch.svd_lowrank(weight, q=probe_rank, niter=niter)
detected_rank = _compute_rank(S, mode, 0, None)
near_probe_tail = detected_rank >= max(1, probe_rank - offset)
if not near_probe_tail or probe_rank >= probe_limit:
rank = min(detected_rank, output_cap)
return U[:, :rank], S[:rank], V[:, :rank].T, rank
probe_rank = probe_limit
def _svd_extract_linear(
@@ -230,6 +251,7 @@ def _svd_extract_linear(
max_rank: int = None,
clamp_quantile: float = 0.99,
niter: int = 2,
knee_probe_offset: int = 32,
) -> tuple:
"""
SVD decomposition for linear layers.
@@ -250,14 +272,18 @@ def _svd_extract_linear(
target_rank = min(target_rank, max_rank)
return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
# For adaptive modes, use full SVD to analyze singular values
try:
U, S, Vh = linalg.svd(weight, full_matrices=False)
except Exception as e:
raise RuntimeError(f"SVD failed: {e}")
if mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None:
U, S, Vh, rank = _svd_extract_knee_lowrank(
weight, mode, max_rank, knee_probe_offset, niter
)
else:
# Other adaptive modes require the full spectrum under their current contract.
try:
U, S, Vh = linalg.svd(weight, full_matrices=False)
except Exception as e:
raise RuntimeError(f"SVD failed: {e}")
# Compute rank based on mode
rank = _compute_rank(S, mode, mode_param, max_rank)
rank = _compute_rank(S, mode, mode_param, max_rank)
# Check if decomposition is worthwhile
if rank >= min(out_dim, in_dim):
@@ -278,13 +304,10 @@ def _svd_extract_linear(
# Construct LoRA matrices
# lora_up = U @ diag(S), lora_down = Vh
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_up = U * S.unsqueeze(0) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff for optional sparse bias
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
return (lora_down, lora_up, None), "low rank"
def _svd_extract_conv(
@@ -294,6 +317,7 @@ def _svd_extract_conv(
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
knee_probe_offset: int = 32,
) -> tuple:
"""
SVD decomposition for conv2d layers.
@@ -310,7 +334,10 @@ def _svd_extract_conv(
else:
mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k]
result, mode_str = _svd_extract_linear(mat, mode, mode_param, device, max_rank, clamp_quantile)
result, mode_str = _svd_extract_linear(
mat, mode, mode_param, device, max_rank, clamp_quantile,
knee_probe_offset=knee_probe_offset,
)
if mode_str == "full":
return weight, "full"
@@ -363,6 +390,7 @@ def _extract_chunked_layer(
mode_param: float,
device: str,
max_rank: int = None,
knee_probe_offset: int = 32,
) -> tuple:
"""Extract LoRA from fused layer by chunking."""
out_dim, in_dim = weight_diff.shape
@@ -375,7 +403,10 @@ def _extract_chunked_layer(
chunk = weight_diff[i * chunk_size:(i + 1) * chunk_size, :]
try:
result, mode_str = _svd_extract_linear(chunk, mode, mode_param, device, max_rank)
result, mode_str = _svd_extract_linear(
chunk, mode, mode_param, device, max_rank,
knee_probe_offset=knee_probe_offset,
)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
@@ -458,6 +489,7 @@ def extract_dora_from_files(
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
knee_probe_offset: int = 32,
) -> None:
"""
Extract LoRA from difference between two models, writing incrementally to disk.
@@ -602,11 +634,13 @@ def extract_dora_from_files(
try:
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
weight_diff, mode, conv_param, device, conv_max_rank,
clamp_quantile, knee_probe_offset
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
weight_diff, mode, linear_param, device, linear_max_rank,
clamp_quantile, svd_niter, knee_probe_offset
)
except Exception as e:
# Try chunked extraction for large tensors
@@ -615,7 +649,8 @@ def extract_dora_from_files(
if num_chunks > 1:
print(f"[DoRA 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
weight_diff, num_chunks, mode, linear_param, device,
linear_max_rank, knee_probe_offset
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
@@ -913,6 +948,8 @@ class DoRAExtractKnee(io.ComfyNode):
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"],
default="sv_knee", tooltip="Knee detection method"),
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096,
tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
@@ -924,7 +961,8 @@ class DoRAExtractKnee(io.ComfyNode):
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
@@ -937,7 +975,8 @@ class DoRAExtractKnee(io.ComfyNode):
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns
glob_skip_patterns=glob_skip_patterns,
knee_probe_offset=knee_probe_offset,
)
return io.NodeOutput(output_path)
nodes/dora_learned_wd.py
+17 -6
View File
@@ -276,6 +276,7 @@ def _extract_chunked_learned_layer(
early_stop_stall: int = 2000,
early_stop_lr: float = 9.01e-9,
layer_name: str = "Layer",
knee_probe_offset: int = 32,
) -> tuple:
"""Extract Learned DoRA from fused layer by chunking."""
out_dim, in_dim = tensor_base.shape
@@ -296,7 +297,10 @@ def _extract_chunked_learned_layer(
weight_diff = w_dir - chunk_base
try:
result, mode_str = _svd_extract_linear(weight_diff, mode, mode_param, device, max_rank)
result, mode_str = _svd_extract_linear(
weight_diff, mode, mode_param, device, max_rank,
knee_probe_offset=knee_probe_offset,
)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
@@ -351,6 +355,7 @@ def extract_dora_learned_from_files(
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
knee_probe_offset: int = 32,
) -> None:
"""
Extract Learned DoRA from difference between two models.
@@ -465,11 +470,13 @@ def extract_dora_learned_from_files(
# 1. Run Analytical SVD (Initialization)
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
weight_diff, mode, conv_param, device, conv_max_rank,
clamp_quantile, knee_probe_offset
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
weight_diff, mode, linear_param, device, linear_max_rank,
clamp_quantile, svd_niter, knee_probe_offset
)
# 2. Run Optimization (Learned Rounding)
@@ -492,7 +499,8 @@ def extract_dora_learned_from_files(
tensor_base, tensor_ft, num_chunks, mode, linear_param, device, linear_max_rank, optimize_iters, learning_rate,
optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
early_stop_loss, early_stop_stall, early_stop_lr,
layer_name=lora_name
layer_name=lora_name,
knee_probe_offset=knee_probe_offset,
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
@@ -719,6 +727,7 @@ class DoRALearnedExtractKnee(io.ComfyNode):
inputs=[
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_learned_inputs(),
@@ -729,7 +738,8 @@ class DoRALearnedExtractKnee(io.ComfyNode):
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
linear_max_rank, conv_max_rank,
optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
@@ -741,7 +751,8 @@ class DoRALearnedExtractKnee(io.ComfyNode):
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset
)
return io.NodeOutput(output_path)
nodes/lora_extract_svd.py
+62 -23
View File
@@ -214,13 +214,34 @@ def _svd_extract_linear_lowrank(
pass
# Construct LoRA matrices
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_up = U * S.unsqueeze(0) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, None), "low rank"
return (lora_down, lora_up, diff), "low rank"
def _svd_extract_knee_lowrank(
weight: torch.Tensor,
mode: str,
max_rank: int,
probe_offset: int,
niter: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
"""Probe beyond the output rank and retry once when the knee hits the probe tail."""
max_possible_rank = min(weight.shape)
output_cap = max(1, min(max_rank, max_possible_rank))
offset = max(1, probe_offset)
probe_rank = min(max_possible_rank, output_cap + offset)
probe_limit = min(max_possible_rank, max(probe_rank, output_cap * 2))
while True:
U, S, V = torch.svd_lowrank(weight, q=probe_rank, niter=niter)
detected_rank = _compute_rank(S, mode, 0, None)
near_probe_tail = detected_rank >= max(1, probe_rank - offset)
if not near_probe_tail or probe_rank >= probe_limit:
rank = min(detected_rank, output_cap)
return U[:, :rank], S[:rank], V[:, :rank].T, rank
probe_rank = probe_limit
def _svd_extract_linear(
@@ -231,6 +252,7 @@ def _svd_extract_linear(
max_rank: int = None,
clamp_quantile: float = 0.99,
niter: int = 2,
knee_probe_offset: int = 32,
) -> tuple:
"""
SVD decomposition for linear layers.
@@ -251,14 +273,18 @@ def _svd_extract_linear(
target_rank = min(target_rank, max_rank)
return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
# For adaptive modes, use full SVD to analyze singular values
try:
U, S, Vh = linalg.svd(weight, full_matrices=False)
except Exception as e:
raise RuntimeError(f"SVD failed: {e}")
if mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None:
U, S, Vh, rank = _svd_extract_knee_lowrank(
weight, mode, max_rank, knee_probe_offset, niter
)
else:
# Other adaptive modes require the full spectrum under their current contract.
try:
U, S, Vh = linalg.svd(weight, full_matrices=False)
except Exception as e:
raise RuntimeError(f"SVD failed: {e}")
# Compute rank based on mode
rank = _compute_rank(S, mode, mode_param, max_rank)
rank = _compute_rank(S, mode, mode_param, max_rank)
# Check if decomposition is worthwhile
if rank >= min(out_dim, in_dim):
@@ -279,13 +305,10 @@ def _svd_extract_linear(
# Construct LoRA matrices
# lora_up = U @ diag(S), lora_down = Vh
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_up = U * S.unsqueeze(0) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff for optional sparse bias
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
return (lora_down, lora_up, None), "low rank"
def _svd_extract_conv(
@@ -295,6 +318,7 @@ def _svd_extract_conv(
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
knee_probe_offset: int = 32,
) -> tuple:
"""
SVD decomposition for conv2d layers.
@@ -311,7 +335,10 @@ def _svd_extract_conv(
else:
mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k]
result, mode_str = _svd_extract_linear(mat, mode, mode_param, device, max_rank, clamp_quantile)
result, mode_str = _svd_extract_linear(
mat, mode, mode_param, device, max_rank, clamp_quantile,
knee_probe_offset=knee_probe_offset,
)
if mode_str == "full":
return weight, "full"
@@ -364,6 +391,7 @@ def _extract_chunked_layer(
mode_param: float,
device: str,
max_rank: int = None,
knee_probe_offset: int = 32,
) -> tuple:
"""Extract LoRA from fused layer by chunking."""
out_dim, in_dim = weight_diff.shape
@@ -376,7 +404,10 @@ def _extract_chunked_layer(
chunk = weight_diff[i * chunk_size:(i + 1) * chunk_size, :]
try:
result, mode_str = _svd_extract_linear(chunk, mode, mode_param, device, max_rank)
result, mode_str = _svd_extract_linear(
chunk, mode, mode_param, device, max_rank,
knee_probe_offset=knee_probe_offset,
)
if mode_str == "full":
return None, None, 0
lora_down, lora_up, _ = result
@@ -460,6 +491,7 @@ def extract_lora_from_files(
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
include_1d_diffs: bool = False,
knee_probe_offset: int = 32,
) -> None:
"""
Extract LoRA from difference between two models, writing incrementally to disk.
@@ -606,11 +638,13 @@ def extract_lora_from_files(
try:
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
weight_diff, mode, conv_param, device, conv_max_rank,
clamp_quantile, knee_probe_offset
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
weight_diff, mode, linear_param, device, linear_max_rank,
clamp_quantile, svd_niter, knee_probe_offset
)
except Exception as e:
# Try chunked extraction for large tensors
@@ -619,7 +653,8 @@ def extract_lora_from_files(
if num_chunks > 1:
print(f"[LoRA 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
weight_diff, num_chunks, mode, linear_param, device,
linear_max_rank, knee_probe_offset
)
if lora_up is not None:
layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(layer_save_dtype).cpu().contiguous()
@@ -922,6 +957,8 @@ class LoRAExtractKnee(io.ComfyNode):
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"],
default="sv_knee", tooltip="Knee detection method"),
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096,
tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
@@ -933,7 +970,8 @@ class LoRAExtractKnee(io.ComfyNode):
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache,
include_1d_diffs=False) -> io.NodeOutput:
@@ -947,7 +985,8 @@ class LoRAExtractKnee(io.ComfyNode):
device, save_dtype, output_path, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs,
knee_probe_offset=knee_probe_offset,
)
return io.NodeOutput(output_path)
nodes/text_encoder_extract.py
+17 -7
View File
@@ -82,6 +82,7 @@ def extract_te_from_files(
lazy_load: bool = True,
force_clear_cache: bool = True,
glob_skip_patterns: bool = False,
knee_probe_offset: int = 32,
) -> None:
"""
Extract LoRA/DoRA from difference between two Text Encoder models, writing incrementally to disk.
@@ -207,11 +208,13 @@ def extract_te_from_files(
try:
if is_conv:
result, mode_str = _svd_extract_conv(
weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile
weight_diff, mode, conv_param, device, conv_max_rank,
clamp_quantile, knee_probe_offset
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
weight_diff, mode, linear_param, device, linear_max_rank,
clamp_quantile, svd_niter, knee_probe_offset
)
except Exception as e:
# Try chunked extraction for large tensors
@@ -220,7 +223,8 @@ def extract_te_from_files(
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
weight_diff, num_chunks, mode, linear_param, device,
linear_max_rank, knee_probe_offset
)
if lora_up is not None:
# PEFT suffixes
@@ -451,6 +455,7 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode):
inputs=[
*_get_te_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
@@ -460,7 +465,8 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode):
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
@@ -474,7 +480,8 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode):
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset
)
return io.NodeOutput(output_path)
@@ -653,6 +660,7 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode):
inputs=[
*_get_te_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."),
io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."),
io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
*_get_common_inputs(),
@@ -662,7 +670,8 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode):
)
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
def execute(cls, model_a, model_b, knee_method, knee_probe_offset,
linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
@@ -676,7 +685,8 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode):
linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
lazy_load=lazy_load, force_clear_cache=force_clear_cache,
glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset
)
return io.NodeOutput(output_path)
tests/test_knee_partial_svd.py
+101
View File
@@ -0,0 +1,101 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_knee_partial_svd_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", params=["nodes.lora_extract_svd", "nodes.dora_extract_wd"])
def extraction(request):
return _load_module(request.param)
def _fake_factors(weight, q, niter):
del niter
return (
torch.eye(weight.shape[0], q),
torch.linspace(float(q), 1.0, q),
torch.eye(weight.shape[1], q),
)
def test_knee_probe_uses_offset_and_expands_at_boundary(monkeypatch, extraction):
probes = []
detected = iter([9, 5])
def fake_svd(weight, q, niter):
probes.append(q)
return _fake_factors(weight, q, niter)
monkeypatch.setattr(extraction.torch, "svd_lowrank", fake_svd)
monkeypatch.setattr(
extraction, "_compute_rank", lambda singular_values, mode, param, cap: next(detected)
)
_, _, _, rank = extraction._svd_extract_knee_lowrank(
torch.eye(20), "sv_knee", max_rank=6, probe_offset=4, niter=2
)
assert probes == [10, 12]
assert rank == 5
def test_knee_probe_stops_when_knee_is_clear(monkeypatch, extraction):
probes = []
def fake_svd(weight, q, niter):
probes.append(q)
return _fake_factors(weight, q, niter)
monkeypatch.setattr(extraction.torch, "svd_lowrank", fake_svd)
monkeypatch.setattr(extraction, "_compute_rank", lambda *args: 3)
_, _, _, rank = extraction._svd_extract_knee_lowrank(
torch.eye(20), "sv_knee", max_rank=6, probe_offset=4, niter=2
)
assert probes == [10]
assert rank == 3
def test_lowrank_factor_construction_does_not_build_residual(monkeypatch, extraction):
monkeypatch.setattr(extraction.torch, "svd_lowrank", _fake_factors)
result, mode = extraction._svd_extract_linear_lowrank(
torch.eye(6), rank=3, device="cpu", clamp_quantile=1.0, niter=2
)
down, up, residual = result
assert mode == "low rank"
assert residual is None
assert torch.equal(up, torch.eye(6, 3) * torch.tensor([3.0, 2.0, 1.0]))
assert torch.equal(down, torch.eye(3, 6))
def test_knee_control_is_adjacent_to_method():
modules_and_nodes = [
(_load_module("nodes.lora_extract_svd"), "LoRAExtractKnee"),
(_load_module("nodes.dora_extract_wd"), "DoRAExtractKnee"),
(_load_module("nodes.dora_learned_wd"), "DoRALearnedExtractKnee"),
(_load_module("nodes.text_encoder_extract"), "TextEncoderLoRAExtractKnee"),
(_load_module("nodes.text_encoder_extract"), "TextEncoderDoRAExtractKnee"),
]
for module, node_name in modules_and_nodes:
names = [item.id for item in getattr(module, node_name).define_schema().inputs]
method_index = names.index("knee_method")
assert names[method_index + 1] == "knee_probe_offset"