feat: add svd_niter parameter for SVD accuracy control

This commit is contained in:
silveroxides
2026-01-03 20:27:56 +01:00
parent 75d1f0f5c9
commit 61b2f506d0
2 changed files with 26 additions and 13 deletions
+13 -6
View File
@@ -148,6 +148,7 @@ def _svd_extract_linear_lowrank(
rank: int,
device: str,
clamp_quantile: float = 0.99,
niter: int = 2,
) -> tuple:
"""
Low-rank SVD decomposition using svd_lowrank for fixed rank extraction.
@@ -155,6 +156,9 @@ def _svd_extract_linear_lowrank(
Much faster than full SVD when rank << min(m, n) because it only computes
the top-k singular values using randomized algorithms.
Args:
niter: Power iterations for SVD accuracy (higher = more accurate but slower)
Returns:
(lora_down, lora_up, diff), "low rank" or (weight, "full")
"""
@@ -170,9 +174,8 @@ def _svd_extract_linear_lowrank(
return weight, "full"
# Use svd_lowrank for faster low-rank approximation
# niter controls accuracy (higher = more accurate but slower)
try:
U, S, V = torch.svd_lowrank(weight, q=rank, niter=2)
U, S, V = torch.svd_lowrank(weight, q=rank, niter=niter)
# U: [out_dim, rank], S: [rank], V: [in_dim, rank]
Vh = V.T # [rank, in_dim]
except Exception as e:
@@ -203,6 +206,7 @@ def _svd_extract_linear(
device: str,
max_rank: int = None,
clamp_quantile: float = 0.99,
niter: int = 2,
) -> tuple:
"""
SVD decomposition for linear layers.
@@ -221,7 +225,7 @@ def _svd_extract_linear(
target_rank = int(mode_param)
if max_rank is not None:
target_rank = min(target_rank, max_rank)
return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile)
return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
# For adaptive modes, use full SVD to analyze singular values
try:
@@ -409,6 +413,7 @@ def extract_lora_from_files(
skip_patterns_str: str = "",
mismatch_mode: str = "skip",
chunk_large_layers: bool = True,
svd_niter: int = 2,
) -> dict[str, torch.Tensor]:
"""
Extract LoRA from difference between two models.
@@ -528,7 +533,7 @@ def extract_lora_from_files(
)
else:
result, mode_str = _svd_extract_linear(
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile
weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter
)
except Exception as e:
# Try chunked extraction for large tensors
@@ -640,6 +645,8 @@ class LoRAExtractFixed(io.ComfyNode):
tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=4096,
tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations (higher = more accurate but slower)"),
*_get_common_inputs(),
],
outputs=[io.String.Output(display_name="output_path")],
@@ -647,7 +654,7 @@ class LoRAExtractFixed(io.ComfyNode):
)
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, chunk_large_layers,
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) -> io.NodeOutput:
@@ -657,7 +664,7 @@ class LoRAExtractFixed(io.ComfyNode):
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
"lora_unet_", device, save_dtype, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
+13 -7
View File
@@ -201,6 +201,7 @@ def _extract_linear(
dynamic_method: Optional[str],
dynamic_param: Optional[float],
scale: float,
niter: int = 2,
) -> Dict:
"""Extract LoRA from merged linear weight."""
out_size, in_size = weight.shape
@@ -208,7 +209,7 @@ def _extract_linear(
if dynamic_method is None:
# Fixed rank - use svd_lowrank for 10x speedup
rank = min(max_rank, min(out_size, in_size) - 1)
U, S, Vh = torch.svd_lowrank(weight, q=rank, niter=2)
U, S, Vh = torch.svd_lowrank(weight, q=rank, niter=niter)
Vh = Vh.T # svd_lowrank returns V, not Vh
new_rank = rank
new_alpha = float(scale * new_rank)
@@ -239,6 +240,7 @@ def _extract_conv(
dynamic_method: Optional[str],
dynamic_param: Optional[float],
scale: float,
niter: int = 2,
) -> Dict:
"""Extract LoRA from merged conv weight."""
out_ch, in_ch, kh, kw = weight.shape
@@ -247,7 +249,7 @@ def _extract_conv(
if dynamic_method is None:
# Fixed rank - use svd_lowrank for 10x speedup
rank = min(max_rank, min(mat.shape) - 1)
U, S, Vh = torch.svd_lowrank(mat, q=rank, niter=2)
U, S, Vh = torch.svd_lowrank(mat, q=rank, niter=niter)
Vh = Vh.T # svd_lowrank returns V, not Vh
new_rank = rank
new_alpha = float(scale * new_rank)
@@ -329,7 +331,8 @@ def resize_lora_file(
device: str,
save_dtype: torch.dtype,
output_filename: str,
verbose: bool = True
verbose: bool = True,
svd_niter: int = 2,
) -> str:
"""
Resize a LoRA file to a new rank.
@@ -343,6 +346,7 @@ def resize_lora_file(
save_dtype: Output dtype
output_filename: Output filename (without extension)
verbose: Print progress info
svd_niter: Power iterations for SVD accuracy
Returns:
Path to saved resized LoRA
@@ -393,10 +397,10 @@ def resize_lora_file(
# Merge and re-extract
if is_conv:
weight = _merge_conv(lora_down, lora_up, device)
result = _extract_conv(weight, new_rank, dynamic_method, dynamic_param, scale)
result = _extract_conv(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
else:
weight = _merge_linear(lora_down, lora_up, device)
result = _extract_linear(weight, new_rank, dynamic_method, dynamic_param, scale)
result = _extract_linear(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter)
del weight, lora_down, lora_up
@@ -463,6 +467,8 @@ class LoRAResizeFixed(io.ComfyNode):
tooltip="LoRA to resize"),
io.Int.Input("new_rank", default=64, min=1, max=1024,
tooltip="Target rank"),
io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations (higher = more accurate but slower)"),
io.String.Input("output_filename", default="resized_lora"),
io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"),
io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"),
@@ -472,11 +478,11 @@ class LoRAResizeFixed(io.ComfyNode):
)
@classmethod
def execute(cls, lora_name, new_rank, output_filename, save_dtype, device) -> io.NodeOutput:
def execute(cls, lora_name, new_rank, svd_niter, output_filename, save_dtype, device) -> io.NodeOutput:
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = resize_lora_file(lora_path, new_rank, None, None, device, dtype, output_filename)
path = resize_lora_file(lora_path, new_rank, None, None, device, dtype, output_filename, svd_niter=svd_niter)
return io.NodeOutput(path)