diff --git a/nodes/adaptive_svd.py b/nodes/adaptive_svd.py new file mode 100644 index 0000000..c5064bf --- /dev/null +++ b/nodes/adaptive_svd.py @@ -0,0 +1,57 @@ +"""Bounded progressive partial SVD for adaptive extraction modes.""" + +import torch + + +ADAPTIVE_PARTIAL_MODES = {"ratio", "quantile", "sv_cumulative", "sv_fro"} + + +def _bounded_rank(singular_values, mode, parameter, output_cap, matrix_fro_sq): + capped = singular_values[:output_cap] + if capped.numel() == 0 or capped[0] < 1e-8: + return 1, True + + if mode == "ratio": + rank = int(torch.sum(capped > capped[0] / parameter).item()) if parameter > 0 else output_cap + return max(1, min(rank, output_cap)), rank < output_cap + + if mode in {"quantile", "sv_cumulative"}: + total = capped.sum() + if total < 1e-8: + return 1, True + rank = int(torch.searchsorted(torch.cumsum(capped, 0) / total, parameter).item()) + 1 + return max(1, min(rank, output_cap)), rank < output_cap + + target_sq = float(parameter) ** 2 * matrix_fro_sq + cumulative = torch.cumsum(capped.square(), 0) + reached = bool(cumulative[-1] >= target_sq) + if reached: + rank = int(torch.searchsorted(cumulative, target_sq).item()) + 1 + return max(1, min(rank, output_cap)), rank < output_cap + return output_cap, False + + +def adaptive_partial_svd(weight, mode, parameter, max_rank, probe_offset, niter): + """Return a bounded top spectrum and selected output rank.""" + max_possible = min(weight.shape) + output_cap = max(1, min(int(max_rank), max_possible)) + offset = max(1, int(probe_offset)) + initial_probe = min(max_possible, output_cap + offset) + probe_limit = min(max_possible, max(initial_probe, output_cap * 2)) + matrix_fro_sq = float(torch.sum(weight.square()).item()) if mode == "sv_fro" else 0.0 + + probe_rank = initial_probe + while True: + if probe_rank >= max_possible: + U, S, Vh = torch.linalg.svd(weight, full_matrices=False) + else: + U, S, V = torch.svd_lowrank(weight, q=probe_rank, niter=niter) + Vh = V.T + + rank, resolved_before_cap = _bounded_rank( + S, mode, parameter, output_cap, matrix_fro_sq + ) + near_probe_tail = rank >= max(1, probe_rank - offset) + if resolved_before_cap or not near_probe_tail or probe_rank >= probe_limit: + return U[:, :rank], S[:rank], Vh[:rank, :], rank + probe_rank = probe_limit diff --git a/nodes/dora_extract_wd.py b/nodes/dora_extract_wd.py index db5fcc5..15f3d69 100644 --- a/nodes/dora_extract_wd.py +++ b/nodes/dora_extract_wd.py @@ -15,10 +15,12 @@ from tqdm import tqdm from comfy_api.latest import io from .device_utils import ( estimate_model_size, prepare_for_large_operation, - cleanup_after_operation, get_device_capabilities + cleanup_after_operation ) -from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter +from .adaptive_svd import ADAPTIVE_PARTIAL_MODES, adaptive_partial_svd +from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu @@ -272,12 +274,16 @@ def _svd_extract_linear( target_rank = min(target_rank, max_rank) return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter) - if mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None: + if mode in ADAPTIVE_PARTIAL_MODES and max_rank is not None: + U, S, Vh, rank = adaptive_partial_svd( + weight, mode, mode_param, max_rank, knee_probe_offset, niter + ) + elif 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. + # Uncapped modes require the full spectrum. try: U, S, Vh = linalg.svd(weight, full_matrices=False) except Exception as e: @@ -532,91 +538,53 @@ def extract_dora_from_files( 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)) + weight_keys = sorted(k for k in keys_a if k.endswith(".weight")) + work_units = [ + (key, key, key if key in keys_b else None) + for key in weight_keys + if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) + and not (key not in keys_b and mismatch_mode == "skip") + ] + pbar = comfy.utils.ProgressBar(len(work_units)) stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} - def _process_layer(key): + def _process_layer(key, cpu_a, cpu_b): lora_name = _format_lora_key(key) + weight_diff = None + layer_device = device 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 - - # Without base, DoRA fallback just treats tensor_a as diff - dora_scale = None + weight_diff, dora_scale, layer_device = dora_difference( + cpu_a, None, device + ) 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 cpu_a.shape != cpu_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}") + raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}") - weight_diff = tensor_a - dora_scale = None - del tensor_b + weight_diff, dora_scale, layer_device = dora_difference( + cpu_a, None, device + ) else: - # 1-D weights (e.g. x_pad_token, norm .weight) have no - # weight-decomposition direction and are skipped below - # anyway; bail out before the norm, which would index - # dim=1 and raise on a 1-D tensor. - if tensor_a.ndim < 2: - del tensor_a, tensor_b + if cpu_a.ndim < 2: return "skipped", None + weight_diff, dora_scale, layer_device = dora_difference( + cpu_a, cpu_b, device + ) - # 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: W_dir = W_ft * (norm_base / norm_ft) - # Use a small epsilon to avoid division by zero - 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 + if weight_diff is None: + raise RuntimeError(f"No extraction tensor was produced for {key}") # Skip small differences if min_diff > 0 and weight_diff.abs().max() < min_diff: @@ -630,17 +598,20 @@ def extract_dora_from_files( is_conv = weight_diff.ndim == 4 layer_results = {} + svd_weight = weight_diff try: if is_conv: - result, mode_str = _svd_extract_conv( - weight_diff, mode, conv_param, device, conv_max_rank, - clamp_quantile, knee_probe_offset + result, mode_str = retry_cuda_oom_on_cpu( + lambda: _svd_extract_conv(svd_weight, mode, conv_param, layer_device, conv_max_rank, clamp_quantile, knee_probe_offset), + lambda: _svd_extract_conv(svd_weight.cpu(), mode, conv_param, "cpu", conv_max_rank, clamp_quantile, knee_probe_offset), + layer_device, ) else: - result, mode_str = _svd_extract_linear( - weight_diff, mode, linear_param, device, linear_max_rank, - clamp_quantile, svd_niter, knee_probe_offset + result, mode_str = retry_cuda_oom_on_cpu( + lambda: _svd_extract_linear(svd_weight, mode, linear_param, layer_device, linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset), + lambda: _svd_extract_linear(svd_weight.cpu(), mode, linear_param, "cpu", linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset), + layer_device, ) except Exception as e: # Try chunked extraction for large tensors @@ -649,7 +620,7 @@ 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, + weight_diff, num_chunks, mode, linear_param, layer_device, linear_max_rank, knee_probe_offset ) if lora_up is not None: @@ -686,8 +657,11 @@ def extract_dora_from_files( writer = IncrementalSafetensorsWriter(output_path) writer.__enter__() try: - for key in tqdm(weight_keys, desc="Extracting DoRA", unit="layers"): - status, layer_sd = _process_layer(key) + stream = paired_async_tensors( + handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda") + ) + for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting DoRA", unit="layers"): + status, layer_sd = _process_layer(key, cpu_a, cpu_b) stats[status] += 1 if layer_sd: writer.write_dict(layer_sd) @@ -699,6 +673,7 @@ def extract_dora_from_files( torch.cuda.empty_cache() pbar.update(1) + stats["skipped"] += len(weight_keys) - len(work_units) finally: writer.__exit__(None, None, None) @@ -780,7 +755,7 @@ def _format_lora_key(key: str) -> str: 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("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer; slower but useful under severe VRAM pressure."), 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, @@ -858,6 +833,8 @@ class DoRAExtractRatio(io.ComfyNode): tooltip="Ratio threshold for linear layers (higher = more SVs kept)"), io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for conv layers (higher = more SVs kept)"), + io.Int.Input("probe_offset", default=32, min=1, max=4096, + tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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, @@ -869,7 +846,7 @@ class DoRAExtractRatio(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_ratio, conv_ratio, 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: @@ -882,7 +859,7 @@ class DoRAExtractRatio(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=probe_offset ) return io.NodeOutput(output_path) @@ -904,6 +881,8 @@ class DoRAExtractQuantile(io.ComfyNode): tooltip="Target cumulative % for linear layers"), io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative % for conv layers"), + io.Int.Input("probe_offset", default=32, min=1, max=4096, + tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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, @@ -915,7 +894,7 @@ class DoRAExtractQuantile(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_quantile, conv_quantile, 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: @@ -928,7 +907,7 @@ class DoRAExtractQuantile(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=probe_offset ) return io.NodeOutput(output_path) @@ -998,6 +977,8 @@ class DoRAExtractFrobenius(io.ComfyNode): tooltip="Target Frobenius norm fraction for linear"), io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius norm fraction for conv"), + io.Int.Input("probe_offset", default=32, min=1, max=4096, + tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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, @@ -1009,7 +990,7 @@ class DoRAExtractFrobenius(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_target, conv_target, 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: @@ -1022,7 +1003,7 @@ class DoRAExtractFrobenius(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=probe_offset ) return io.NodeOutput(output_path) diff --git a/nodes/dora_learned_wd.py b/nodes/dora_learned_wd.py index e773c30..83d8b54 100644 --- a/nodes/dora_learned_wd.py +++ b/nodes/dora_learned_wd.py @@ -5,12 +5,8 @@ Extracts Weight-Decomposed Low-Rank Adapters using gradient descent optimization to minimize the reconstruction MSE loss, matching the true non-linear DoRA formulation. Uses analytical SVD purely as the initialization seed. """ -import fnmatch -import os -import re import math import torch -import torch.linalg as linalg import folder_paths import comfy.utils from tqdm import tqdm @@ -20,11 +16,12 @@ from .device_utils import ( cleanup_after_operation ) -from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter +from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu # Import the existing SV ranking and chunking logic so we don't have to duplicate it from .dora_extract_wd import ( - _compute_rank, _svd_extract_linear_lowrank, _svd_extract_linear, + _svd_extract_linear, _svd_extract_conv, _detect_fused_layer, _compile_patterns, _matches_any_pattern, _build_lora_output_path, _get_model_inputs, _format_lora_key, _get_common_inputs @@ -378,82 +375,54 @@ def extract_dora_learned_from_files( 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)) + weight_keys = sorted(k for k in keys_a if k.endswith(".weight")) + work_units = [ + (key, key, key if key in keys_b else None) + for key in weight_keys + if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) + and not (key not in keys_b and mismatch_mode == "skip") + ] + pbar = comfy.utils.ProgressBar(len(work_units)) stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} - def _process_layer(key): + def _process_layer(key, cpu_a, cpu_b): lora_name = _format_lora_key(key) + weight_diff = None + tensor_base = None + tensor_ft = None + layer_device = device if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns): return "skipped", None - 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") - # Cannot extract true DoRA without base. Fallback to diff. - cpu_a = handler_a.get_tensor(key) - tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32) if use_pinned else cpu_a.to(device=device, dtype=torch.float32) - del cpu_a - weight_diff = tensor_a - tensor_base = torch.zeros_like(tensor_a) - tensor_ft = tensor_a - dora_scale = None + weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = ( + dora_difference(cpu_a, None, device, keep_sources=True) + ) 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 cpu_a.shape != cpu_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}") + raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}") - weight_diff = tensor_a - tensor_base = torch.zeros_like(tensor_a) - tensor_ft = tensor_a - dora_scale = None - del tensor_b + weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = ( + dora_difference(cpu_a, None, device, keep_sources=True) + ) else: - tensor_base = tensor_b - tensor_ft = tensor_a + if cpu_a.ndim < 2: + return "skipped", None + weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = ( + dora_difference(cpu_a, cpu_b, device, keep_sources=True) + ) - 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) - - 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 w_dir + if weight_diff is None or tensor_base is None or tensor_ft is None: + raise RuntimeError(f"No extraction tensors were produced for {key}") if min_diff > 0 and weight_diff.abs().max() < min_diff: del weight_diff, tensor_base, tensor_ft @@ -465,28 +434,37 @@ def extract_dora_learned_from_files( is_conv = weight_diff.ndim == 4 layer_results = {} + svd_weight = weight_diff + optimize_base = tensor_base + optimize_ft = tensor_ft try: # 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, knee_probe_offset + result, mode_str = retry_cuda_oom_on_cpu( + lambda: _svd_extract_conv(svd_weight, mode, conv_param, layer_device, conv_max_rank, clamp_quantile, knee_probe_offset), + lambda: _svd_extract_conv(svd_weight.cpu(), mode, conv_param, "cpu", conv_max_rank, clamp_quantile, knee_probe_offset), + layer_device, ) else: - result, mode_str = _svd_extract_linear( - weight_diff, mode, linear_param, device, linear_max_rank, - clamp_quantile, svd_niter, knee_probe_offset + result, mode_str = retry_cuda_oom_on_cpu( + lambda: _svd_extract_linear(svd_weight, mode, linear_param, layer_device, linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset), + lambda: _svd_extract_linear(svd_weight.cpu(), mode, linear_param, "cpu", linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset), + layer_device, ) # 2. Run Optimization (Learned Rounding) if mode_str != "full" and dora_scale is not None: lora_down, lora_up, _ = result - lora_down, lora_up = _optimize_dora( - tensor_base, tensor_ft, lora_down, lora_up, 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 + optimize_args = ( + optimize_iters, learning_rate, optimizer_type, lr_schedule, + lr_patience, lr_factor, lr_cooldown, early_stop_loss, + early_stop_stall, early_stop_lr, + ) + lora_down, lora_up = retry_cuda_oom_on_cpu( + lambda: _optimize_dora(optimize_base, optimize_ft, lora_down, lora_up, *optimize_args, layer_name=lora_name), + lambda: _optimize_dora(optimize_base.cpu(), optimize_ft.cpu(), lora_down.cpu(), lora_up.cpu(), *optimize_args, layer_name=lora_name), + layer_device, ) result = (lora_down, lora_up, _) @@ -496,7 +474,7 @@ def extract_dora_learned_from_files( if num_chunks > 1: print(f"[Learned DoRA] Chunked: {key} ({num_chunks} chunks)") lora_up, lora_down, rank = _extract_chunked_learned_layer( - tensor_base, tensor_ft, num_chunks, mode, linear_param, device, linear_max_rank, optimize_iters, learning_rate, + tensor_base, tensor_ft, num_chunks, mode, linear_param, layer_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, @@ -534,8 +512,11 @@ def extract_dora_learned_from_files( writer = IncrementalSafetensorsWriter(output_path) writer.__enter__() try: - for key in tqdm(weight_keys, desc="Extracting Learned DoRA", unit="layers"): - status, layer_sd = _process_layer(key) + stream = paired_async_tensors( + handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda") + ) + for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting Learned DoRA", unit="layers"): + status, layer_sd = _process_layer(key, cpu_a, cpu_b) stats[status] += 1 if layer_sd: writer.write_dict(layer_sd) @@ -547,6 +528,7 @@ def extract_dora_learned_from_files( torch.cuda.empty_cache() pbar.update(1) + stats["skipped"] += len(weight_keys) - len(work_units) finally: writer.__exit__(None, None, None) @@ -642,6 +624,7 @@ class DoRALearnedExtractRatio(io.ComfyNode): *_get_model_inputs(), io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for linear layers; higher values retain more singular values."), io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."), + io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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(), @@ -652,7 +635,7 @@ class DoRALearnedExtractRatio(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_ratio, conv_ratio, 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: @@ -664,7 +647,8 @@ class DoRALearnedExtractRatio(io.ComfyNode): model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio, 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=probe_offset ) return io.NodeOutput(output_path) @@ -685,6 +669,7 @@ class DoRALearnedExtractQuantile(io.ComfyNode): *_get_model_inputs(), io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."), io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."), + io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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(), @@ -695,7 +680,7 @@ class DoRALearnedExtractQuantile(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_quantile, conv_quantile, 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: @@ -707,7 +692,8 @@ class DoRALearnedExtractQuantile(io.ComfyNode): model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile, 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=probe_offset ) return io.NodeOutput(output_path) @@ -773,6 +759,7 @@ class DoRALearnedExtractFrobenius(io.ComfyNode): *_get_model_inputs(), io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."), io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."), + io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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(), @@ -783,7 +770,7 @@ class DoRALearnedExtractFrobenius(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_target, conv_target, 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: @@ -795,7 +782,8 @@ class DoRALearnedExtractFrobenius(io.ComfyNode): model_a_path, model_b_path, "sv_fro", linear_target, conv_target, 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=probe_offset ) return io.NodeOutput(output_path) diff --git a/nodes/extraction_stream.py b/nodes/extraction_stream.py new file mode 100644 index 0000000..0d4f9e3 --- /dev/null +++ b/nodes/extraction_stream.py @@ -0,0 +1,185 @@ +"""Bounded paired UEL streaming helpers for model extraction nodes.""" + +from collections import deque + +import torch + +def is_cuda_oom(error): + return isinstance(error, torch.cuda.OutOfMemoryError) or ( + isinstance(error, RuntimeError) + and "cuda" in str(error).lower() + and "out of memory" in str(error).lower() + ) + + +def release_cuda_oom(): + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def move_source_tensors(cpu_a, cpu_b, device, dtype=torch.float32): + """Move one source pair, retrying a confirmed CUDA OOM on CPU.""" + use_cuda = str(device).startswith("cuda") + try: + if use_cuda: + tensor_a = cpu_a.to(device=device, dtype=dtype, non_blocking=True) + tensor_b = cpu_b.to(device=device, dtype=dtype, non_blocking=True) if cpu_b is not None else None + else: + tensor_a = cpu_a.to(device=device, dtype=dtype) + tensor_b = cpu_b.to(device=device, dtype=dtype) if cpu_b is not None else None + return tensor_a, tensor_b, device + except BaseException as error: + if not use_cuda or not is_cuda_oom(error): + raise + release_cuda_oom() + return ( + cpu_a.to(device="cpu", dtype=dtype), + cpu_b.to(device="cpu", dtype=dtype) if cpu_b is not None else None, + "cpu", + ) + + +def retry_cuda_oom_on_cpu(cuda_call, cpu_call, device): + """Retry one compute operation on CPU only for a confirmed CUDA OOM.""" + try: + return cuda_call() + except BaseException as error: + if not str(device).startswith("cuda") or not is_cuda_oom(error): + raise + release_cuda_oom() + return cpu_call() + + +def lora_difference(cpu_a, cpu_b, device): + """Construct A minus B with transfer and arithmetic CUDA OOM fallback.""" + tensor_a, tensor_b, actual_device = move_source_tensors( + cpu_a, cpu_b, device + ) + if tensor_b is None: + return tensor_a, actual_device + try: + return tensor_a - tensor_b, actual_device + except BaseException as error: + if not str(actual_device).startswith("cuda") or not is_cuda_oom(error): + raise + del tensor_a, tensor_b + release_cuda_oom() + return cpu_a.float() - cpu_b.float(), "cpu" + + +def dora_difference(cpu_a, cpu_b, device, *, keep_sources=False): + """Construct a directional DoRA delta with complete CUDA OOM fallback.""" + tensor_a, tensor_b, actual_device = move_source_tensors( + cpu_a, cpu_b, device + ) + + def compute(a, b): + if b is None: + base = torch.zeros_like(a) + result = (a, None, base, a) + return result if keep_sources else result[:2] + if a.ndim < 2: + result = (None, None, b, a) + return result if keep_sources else result[:2] + is_conv = a.ndim == 4 + if is_conv: + out_channels = a.shape[0] + norm_a = torch.linalg.norm(a.reshape(out_channels, -1), dim=1, keepdim=True) + norm_b = torch.linalg.norm(b.reshape(out_channels, -1), dim=1, keepdim=True) + else: + norm_a = torch.linalg.norm(a, dim=1, keepdim=True) + norm_b = torch.linalg.norm(b, dim=1, keepdim=True) + ratio = norm_b / (norm_a + 1e-8) + if is_conv: + ratio = ratio.view(-1, 1, 1, 1) + scale = norm_a.view(-1, 1, 1, 1) + else: + scale = norm_a + result = (a * ratio - b, scale, b, a) + return result if keep_sources else result[:2] + + try: + return (*compute(tensor_a, tensor_b), actual_device) + except BaseException as error: + if not str(actual_device).startswith("cuda") or not is_cuda_oom(error): + raise + del tensor_a, tensor_b + release_cuda_oom() + cpu_b_float = cpu_b.float() if cpu_b is not None else None + return (*compute(cpu_a.float(), cpu_b_float), "cpu") + + +class AsyncTensorCursor: + """Consume one ordered, bounded UEL stream and release each consumed key.""" + + def __init__(self, handler, keys, *, pin_memory: bool): + self.handler = handler + self.expected = tuple(keys) + self.position = 0 + self.pending = deque() + self.stream = None + if self.expected: + self.stream = handler.async_stream( + list(self.expected), + batch_size=1, + prefetch_batches=1, + pin_memory=pin_memory, + ) + + def take(self, expected_key): + if self.position >= len(self.expected): + raise RuntimeError(f"Unexpected tensor request: {expected_key}") + planned_key = self.expected[self.position] + if planned_key != expected_key: + raise RuntimeError( + f"UEL stream order mismatch: expected {planned_key}, requested {expected_key}" + ) + if not self.pending: + self.pending.extend(next(self.stream)) + key, tensor = self.pending.popleft() + if key != expected_key: + raise RuntimeError(f"UEL yielded {key} while {expected_key} was expected") + self.position += 1 + return tensor + + def release(self, key): + self.handler.mark_processed(key) + + def finish(self): + if self.position != len(self.expected) or self.pending: + raise RuntimeError("UEL stream ended before all planned tensors were consumed") + if self.stream is not None: + try: + next(self.stream) + except StopIteration: + return + raise RuntimeError("UEL stream yielded unplanned tensors") + + def close(self): + if self.stream is not None: + self.stream.close() + + +def paired_async_tensors(handler_a, handler_b, work_units, *, pin_memory: bool): + """Yield ordered A/B source tensors and release them after each work unit.""" + keys_a = [unit[1] for unit in work_units if unit[1] is not None] + keys_b = [unit[2] for unit in work_units if unit[2] is not None] + cursor_a = AsyncTensorCursor(handler_a, keys_a, pin_memory=pin_memory) + cursor_b = AsyncTensorCursor(handler_b, keys_b, pin_memory=pin_memory) + try: + for logical_key, key_a, key_b in work_units: + tensor_a = cursor_a.take(key_a) if key_a is not None else None + tensor_b = cursor_b.take(key_b) if key_b is not None else None + try: + yield logical_key, tensor_a, tensor_b + finally: + del tensor_a, tensor_b + if key_a is not None: + cursor_a.release(key_a) + if key_b is not None: + cursor_b.release(key_b) + cursor_a.finish() + cursor_b.finish() + finally: + cursor_a.close() + cursor_b.close() diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index 0825eaa..f22842e 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -15,11 +15,13 @@ from tqdm import tqdm from comfy_api.latest import io from .device_utils import ( estimate_model_size, prepare_for_large_operation, - cleanup_after_operation, get_device_capabilities + cleanup_after_operation ) -from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter from .quantization_guard import inspect_low_bit_input +from .adaptive_svd import ADAPTIVE_PARTIAL_MODES, adaptive_partial_svd +from .extraction_stream import lora_difference, paired_async_tensors, retry_cuda_oom_on_cpu @@ -273,12 +275,16 @@ def _svd_extract_linear( target_rank = min(target_rank, max_rank) return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter) - if mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None: + if mode in ADAPTIVE_PARTIAL_MODES and max_rank is not None: + U, S, Vh, rank = adaptive_partial_svd( + weight, mode, mode_param, max_rank, knee_probe_offset, niter + ) + elif 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. + # Uncapped modes require the full spectrum. try: U, S, Vh = linalg.svd(weight, full_matrices=False) except Exception as e: @@ -549,11 +555,20 @@ def extract_lora_from_files( and (len(handler_a.get_shape(k)) == 1 or k.endswith(direct_diff_suffixes)) )) ] - pbar = comfy.utils.ProgressBar(len(extraction_keys)) + extraction_keys.sort() + work_units = [ + (key, key, key if key in keys_b else None) + for key in extraction_keys + if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) + and not (key not in keys_b and mismatch_mode == "skip") + ] + pbar = comfy.utils.ProgressBar(len(work_units)) stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} - def _process_layer(key): + def _process_layer(key, cpu_a, cpu_b): lora_name = _format_lora_key(key) + weight_diff = None + layer_device = device if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns): return "skipped", None @@ -566,50 +581,32 @@ def extract_lora_from_files( ) # 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 + weight_diff, layer_device = lora_difference(cpu_a, None, device) if mismatch_mode == "zeros": # For zeros mode, we treat the missing base as a zeroed tensor of same shape # which is already captured by weight_diff = tensor_a pass 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 cpu_a.shape != cpu_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}") + raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}") # For zeros/fallback, use tensor_a as the difference - weight_diff = tensor_a - del tensor_b + weight_diff, layer_device = lora_difference(cpu_a, None, device) else: - weight_diff = tensor_a - tensor_b - del tensor_a, tensor_b + weight_diff, layer_device = lora_difference(cpu_a, cpu_b, device) + + if weight_diff is None: + raise RuntimeError(f"No extraction tensor was produced for {key}") # Skip small differences if min_diff > 0 and weight_diff.abs().max() < min_diff: @@ -634,17 +631,20 @@ def extract_lora_from_files( is_conv = weight_diff.ndim == 4 layer_results = {} + svd_weight = weight_diff try: if is_conv: - result, mode_str = _svd_extract_conv( - weight_diff, mode, conv_param, device, conv_max_rank, - clamp_quantile, knee_probe_offset + result, mode_str = retry_cuda_oom_on_cpu( + lambda: _svd_extract_conv(svd_weight, mode, conv_param, layer_device, conv_max_rank, clamp_quantile, knee_probe_offset), + lambda: _svd_extract_conv(svd_weight.cpu(), mode, conv_param, "cpu", conv_max_rank, clamp_quantile, knee_probe_offset), + layer_device, ) else: - result, mode_str = _svd_extract_linear( - weight_diff, mode, linear_param, device, linear_max_rank, - clamp_quantile, svd_niter, knee_probe_offset + result, mode_str = retry_cuda_oom_on_cpu( + lambda: _svd_extract_linear(svd_weight, mode, linear_param, layer_device, linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset), + lambda: _svd_extract_linear(svd_weight.cpu(), mode, linear_param, "cpu", linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset), + layer_device, ) except Exception as e: # Try chunked extraction for large tensors @@ -653,7 +653,7 @@ 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, + weight_diff, num_chunks, mode, linear_param, layer_device, linear_max_rank, knee_probe_offset ) if lora_up is not None: @@ -682,8 +682,11 @@ def extract_lora_from_files( writer = IncrementalSafetensorsWriter(output_path) writer.__enter__() try: - for key in tqdm(extraction_keys, desc="Extracting LoRA", unit="layers"): - status, layer_sd = _process_layer(key) + stream = paired_async_tensors( + handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda") + ) + for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting LoRA", unit="layers"): + status, layer_sd = _process_layer(key, cpu_a, cpu_b) stats[status] += 1 if layer_sd: writer.write_dict(layer_sd) @@ -695,6 +698,7 @@ def extract_lora_from_files( torch.cuda.empty_cache() pbar.update(1) + stats["skipped"] += len(extraction_keys) - len(work_units) finally: writer.__exit__(None, None, None) @@ -783,7 +787,7 @@ def _format_direct_diff_key(source_key: str, lora_name: str) -> str: 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("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer; slower but useful under severe VRAM pressure."), 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, @@ -865,6 +869,8 @@ class LoRAExtractRatio(io.ComfyNode): tooltip="Ratio threshold for linear layers (higher = more SVs kept)"), io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for conv layers (higher = more SVs kept)"), + io.Int.Input("probe_offset", default=32, min=1, max=4096, + tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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, @@ -876,7 +882,7 @@ class LoRAExtractRatio(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_ratio, conv_ratio, 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: @@ -890,7 +896,8 @@ class LoRAExtractRatio(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs + glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, + knee_probe_offset=probe_offset, ) return io.NodeOutput(output_path) @@ -912,6 +919,8 @@ class LoRAExtractQuantile(io.ComfyNode): tooltip="Target cumulative % for linear layers"), io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative % for conv layers"), + io.Int.Input("probe_offset", default=32, min=1, max=4096, + tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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, @@ -923,7 +932,7 @@ class LoRAExtractQuantile(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_quantile, conv_quantile, 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: @@ -937,7 +946,8 @@ class LoRAExtractQuantile(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs + glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, + knee_probe_offset=probe_offset, ) return io.NodeOutput(output_path) @@ -1008,6 +1018,8 @@ class LoRAExtractFrobenius(io.ComfyNode): tooltip="Target Frobenius norm fraction for linear"), io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius norm fraction for conv"), + io.Int.Input("probe_offset", default=32, min=1, max=4096, + tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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, @@ -1019,7 +1031,7 @@ class LoRAExtractFrobenius(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_target, conv_target, 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: @@ -1033,7 +1045,8 @@ class LoRAExtractFrobenius(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs + glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, + knee_probe_offset=probe_offset, ) return io.NodeOutput(output_path) diff --git a/nodes/text_encoder_extract.py b/nodes/text_encoder_extract.py index 8c82375..99d8def 100644 --- a/nodes/text_encoder_extract.py +++ b/nodes/text_encoder_extract.py @@ -4,11 +4,8 @@ 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 @@ -18,11 +15,11 @@ from .device_utils import ( cleanup_after_operation ) -from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter +from .extraction_stream import dora_difference, lora_difference, paired_async_tensors, retry_cuda_oom_on_cpu # 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, @@ -106,91 +103,67 @@ def extract_te_from_files( 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)) + weight_keys = sorted(k for k in keys_a if k.endswith(".weight")) + work_units = [ + (key, key, key if key in keys_b else None) + for key in weight_keys + if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns) + and not (key not in keys_b and mismatch_mode == "skip") + ] + pbar = comfy.utils.ProgressBar(len(work_units)) stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} - def _process_layer(key): + def _process_layer(key, cpu_a, cpu_b): lora_name = _format_te_lora_key(key) + weight_diff = None + layer_device = device 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) + if is_dora: + weight_diff, dora_scale, layer_device = dora_difference( + cpu_a, None, device + ) else: - weight_diff = cpu_a.to(device=device, dtype=torch.float32) - del cpu_a - - dora_scale = None + weight_diff, layer_device = lora_difference(cpu_a, None, device) + 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 cpu_a.shape != cpu_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 + raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}") 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 + weight_diff, dora_scale, layer_device = dora_difference( + cpu_a, None, device + ) else: - weight_diff = tensor_a - tensor_b + weight_diff, layer_device = lora_difference(cpu_a, None, device) dora_scale = None - del tensor_a, tensor_b + else: + if cpu_a.ndim < 2: + return "skipped", None + if is_dora: + weight_diff, dora_scale, layer_device = dora_difference( + cpu_a, cpu_b, device + ) + else: + weight_diff, layer_device = lora_difference( + cpu_a, cpu_b, device + ) + dora_scale = None + + if weight_diff is None: + raise RuntimeError(f"No extraction tensor was produced for {key}") # Skip small differences if min_diff > 0 and weight_diff.abs().max() < min_diff: @@ -204,17 +177,20 @@ def extract_te_from_files( is_conv = weight_diff.ndim == 4 layer_results = {} + svd_weight = weight_diff try: if is_conv: - result, mode_str = _svd_extract_conv( - weight_diff, mode, conv_param, device, conv_max_rank, - clamp_quantile, knee_probe_offset + result, mode_str = retry_cuda_oom_on_cpu( + lambda: _svd_extract_conv(svd_weight, mode, conv_param, layer_device, conv_max_rank, clamp_quantile, knee_probe_offset), + lambda: _svd_extract_conv(svd_weight.cpu(), mode, conv_param, "cpu", conv_max_rank, clamp_quantile, knee_probe_offset), + layer_device, ) else: - result, mode_str = _svd_extract_linear( - weight_diff, mode, linear_param, device, linear_max_rank, - clamp_quantile, svd_niter, knee_probe_offset + result, mode_str = retry_cuda_oom_on_cpu( + lambda: _svd_extract_linear(svd_weight, mode, linear_param, layer_device, linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset), + lambda: _svd_extract_linear(svd_weight.cpu(), mode, linear_param, "cpu", linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset), + layer_device, ) except Exception as e: # Try chunked extraction for large tensors @@ -223,7 +199,7 @@ 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, + weight_diff, num_chunks, mode, linear_param, layer_device, linear_max_rank, knee_probe_offset ) if lora_up is not None: @@ -258,8 +234,11 @@ def extract_te_from_files( 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) + stream = paired_async_tensors( + handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda") + ) + for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting TE Layers", unit="layers"): + status, layer_sd = _process_layer(key, cpu_a, cpu_b) stats[status] += 1 if layer_sd: writer.write_dict(layer_sd) @@ -271,6 +250,7 @@ def extract_te_from_files( torch.cuda.empty_cache() pbar.update(1) + stats["skipped"] += len(weight_keys) - len(work_units) finally: writer.__exit__(None, None, None) @@ -302,7 +282,7 @@ def _get_te_model_inputs(): 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("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer; slower but useful under severe VRAM pressure."), 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, @@ -376,6 +356,7 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode): *_get_te_model_inputs(), io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for linear layers; higher values retain more singular values."), io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."), + io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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(), @@ -385,7 +366,7 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_ratio, conv_ratio, 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: @@ -399,7 +380,8 @@ class TextEncoderLoRAExtractRatio(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=probe_offset ) return io.NodeOutput(output_path) @@ -416,6 +398,7 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode): *_get_te_model_inputs(), io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."), io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."), + io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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(), @@ -425,7 +408,7 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_quantile, conv_quantile, 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: @@ -439,7 +422,8 @@ class TextEncoderLoRAExtractQuantile(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=probe_offset ) return io.NodeOutput(output_path) @@ -498,6 +482,7 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode): *_get_te_model_inputs(), io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."), io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."), + io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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(), @@ -507,7 +492,7 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_target, conv_target, 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: @@ -521,7 +506,8 @@ class TextEncoderLoRAExtractFrobenius(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=probe_offset ) return io.NodeOutput(output_path) @@ -581,6 +567,7 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode): *_get_te_model_inputs(), io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for linear layers; higher values retain more singular values."), io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."), + io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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(), @@ -590,7 +577,7 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_ratio, conv_ratio, 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: @@ -604,7 +591,8 @@ class TextEncoderDoRAExtractRatio(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=probe_offset ) return io.NodeOutput(output_path) @@ -621,6 +609,7 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode): *_get_te_model_inputs(), io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."), io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."), + io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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(), @@ -630,7 +619,7 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_quantile, conv_quantile, 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: @@ -644,7 +633,8 @@ class TextEncoderDoRAExtractQuantile(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=probe_offset ) return io.NodeOutput(output_path) @@ -703,6 +693,7 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode): *_get_te_model_inputs(), io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."), io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."), + io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."), 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(), @@ -712,7 +703,7 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, linear_target, conv_target, 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: @@ -726,6 +717,7 @@ class TextEncoderDoRAExtractFrobenius(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=probe_offset ) return io.NodeOutput(output_path) diff --git a/tests/test_adaptive_extraction_pipeline.py b/tests/test_adaptive_extraction_pipeline.py new file mode 100644 index 0000000..533e12d --- /dev/null +++ b/tests/test_adaptive_extraction_pipeline.py @@ -0,0 +1,219 @@ +import importlib +import sys +import types +from collections import Counter +from pathlib import Path + +import pytest +import torch +from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen + + +REPO_ROOT = Path(__file__).resolve().parents[1] +PACKAGE_NAME = "modelutils_adaptive_extraction_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 adaptive(): + return _load_module("nodes.adaptive_svd") + + +def _exact_partial(weight, q, niter): + del niter + U, S, Vh = torch.linalg.svd(weight, full_matrices=False) + return U[:, :q], S[:q], Vh[:q, :].T + + +@pytest.mark.parametrize( + ("mode", "parameter", "expected"), + [("ratio", 2.0, 2), ("quantile", 0.7, 2), ("sv_fro", 0.8, 1)], +) +def test_adaptive_partial_rank_matches_reference(monkeypatch, adaptive, mode, parameter, expected): + monkeypatch.setattr(adaptive.torch, "svd_lowrank", _exact_partial) + weight = torch.diag(torch.tensor([8.0, 5.0, 2.0, 1.0, 0.5, 0.25])) + _, _, _, rank = adaptive.adaptive_partial_svd( + weight, mode, parameter, max_rank=4, probe_offset=1, niter=2 + ) + assert rank == expected + + +def test_unresolved_ratio_expands_and_caps(monkeypatch, adaptive): + probes = [] + + def tracked(weight, q, niter): + probes.append(q) + return _exact_partial(weight, q, niter) + + monkeypatch.setattr(adaptive.torch, "svd_lowrank", tracked) + weight = torch.diag(torch.linspace(10.0, 8.0, 12)) + _, _, _, rank = adaptive.adaptive_partial_svd( + weight, "ratio", 100.0, max_rank=4, probe_offset=2, niter=2 + ) + assert probes == [6, 8] + assert rank == 4 + + +def test_frobenius_uses_energy_outside_output_cap(monkeypatch, adaptive): + monkeypatch.setattr(adaptive.torch, "svd_lowrank", _exact_partial) + weight = torch.eye(10) + _, _, _, rank = adaptive.adaptive_partial_svd( + weight, "sv_fro", 0.9, max_rank=3, probe_offset=2, niter=2 + ) + assert rank == 3 + + +class _FakeStream: + def __init__(self, batches): + self._batches = iter(batches) + self.closed = False + + def __iter__(self): + return self + + def __next__(self): + return next(self._batches) + + def close(self): + self.closed = True + + +class _FakeHandler: + def __init__(self, values): + self.values = values + self.calls = [] + self.marked = [] + self.streams = [] + + def async_stream(self, keys, **kwargs): + self.calls.append((tuple(keys), kwargs)) + stream = _FakeStream([[(key, self.values[key])] for key in keys]) + self.streams.append(stream) + return stream + + def mark_processed(self, key): + self.marked.append(key) + + +def test_paired_stream_is_bounded_ordered_and_releases_once(): + streaming = _load_module("nodes.extraction_stream") + handler_a = _FakeHandler({"a1": torch.tensor(1), "a2": torch.tensor(2)}) + handler_b = _FakeHandler({"b1": torch.tensor(3)}) + units = [("one", "a1", "b1"), ("two", "a2", None)] + + observed = [] + for key, tensor_a, tensor_b in streaming.paired_async_tensors( + handler_a, handler_b, units, pin_memory=True + ): + observed.append((key, tensor_a.item(), None if tensor_b is None else tensor_b.item())) + + assert observed == [("one", 1, 3), ("two", 2, None)] + assert handler_a.calls == [(('a1', 'a2'), {"batch_size": 1, "prefetch_batches": 1, "pin_memory": True})] + assert handler_b.calls == [(('b1',), {"batch_size": 1, "prefetch_batches": 1, "pin_memory": True})] + assert Counter(handler_a.marked) == Counter({"a1": 1, "a2": 1}) + assert Counter(handler_b.marked) == Counter({"b1": 1}) + assert all(stream.closed for stream in handler_a.streams + handler_b.streams) + + +def test_cuda_oom_compute_retries_on_cpu(monkeypatch): + streaming = _load_module("nodes.extraction_stream") + monkeypatch.setattr(streaming, "release_cuda_oom", lambda: None) + calls = [] + + def fail_cuda(): + calls.append("cuda") + raise RuntimeError("CUDA out of memory") + + def run_cpu(): + calls.append("cpu") + return 7 + + assert streaming.retry_cuda_oom_on_cpu(fail_cuda, run_cpu, "cuda:0") == 7 + assert calls == ["cuda", "cpu"] + + +def test_adaptive_controls_cover_all_public_families(): + modules_and_nodes = [ + (_load_module("nodes.lora_extract_svd"), prefix) + for prefix in ("LoRAExtractRatio", "LoRAExtractQuantile", "LoRAExtractFrobenius") + ] + modules_and_nodes += [ + (_load_module("nodes.dora_extract_wd"), prefix) + for prefix in ("DoRAExtractRatio", "DoRAExtractQuantile", "DoRAExtractFrobenius") + ] + modules_and_nodes += [ + (_load_module("nodes.dora_learned_wd"), prefix) + for prefix in ("DoRALearnedExtractRatio", "DoRALearnedExtractQuantile", "DoRALearnedExtractFrobenius") + ] + modules_and_nodes += [ + (_load_module("nodes.text_encoder_extract"), prefix) + for prefix in ( + "TextEncoderLoRAExtractRatio", "TextEncoderLoRAExtractQuantile", + "TextEncoderLoRAExtractFrobenius", "TextEncoderDoRAExtractRatio", + "TextEncoderDoRAExtractQuantile", "TextEncoderDoRAExtractFrobenius", + ) + ] + + for module, node_name in modules_and_nodes: + names = [item.id for item in getattr(module, node_name).define_schema().inputs] + probe_index = names.index("probe_offset") + assert names[probe_index + 1] == "linear_max_rank" + + +def _write_uel(path, tensors): + with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer: + writer.write_dict(tensors) + + +def test_lora_extraction_uses_only_async_uel_and_releases_sources(monkeypatch, tmp_path): + extraction = _load_module("nodes.lora_extract_svd") + model_a = tmp_path / "a.safetensors" + model_b = tmp_path / "b.safetensors" + output = tmp_path / "out.safetensors" + key = "model.diffusion_model.layer.weight" + _write_uel(model_a, {key: torch.eye(3)}) + _write_uel(model_b, {key: torch.zeros(3, 3)}) + monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None) + original_async = extraction.MemoryEfficientSafeOpen.async_stream + original_mark = extraction.MemoryEfficientSafeOpen.mark_processed + async_calls = [] + marked = [] + + def track_async(self, keys, **kwargs): + async_calls.append((self.filename, tuple(keys), kwargs)) + yield from original_async(self, keys, **kwargs) + + def track_mark(self, source_key): + marked.append((self.filename, source_key)) + return original_mark(self, source_key) + + monkeypatch.setattr(extraction.MemoryEfficientSafeOpen, "async_stream", track_async) + monkeypatch.setattr(extraction.MemoryEfficientSafeOpen, "mark_processed", track_mark) + monkeypatch.setattr( + extraction.MemoryEfficientSafeOpen, + "get_tensor", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("synchronous tensor load")), + ) + + extraction.extract_lora_from_files( + str(model_a), str(model_b), "fixed", 1, 1, "cpu", "fp32", str(output), + linear_max_rank=1, conv_max_rank=1, force_clear_cache=False, + ) + + assert output.is_file() + assert len(async_calls) == 2 + assert all(options == {"batch_size": 1, "prefetch_batches": 1, "pin_memory": False} for _, _, options in async_calls) + assert Counter(marked) == Counter((filename, key) for filename, keys, _ in async_calls for key in keys) + with MemoryEfficientSafeOpen(str(output), low_memory=True) as result: + assert sorted(result.keys()) == [ + "diffusion_model.layer.lora_A.weight", + "diffusion_model.layer.lora_B.weight", + ]