From 1512fa7390580aef86fb717b14c16af1b4e9e6d1 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Thu, 28 May 2026 03:27:47 +0200 Subject: [PATCH 1/2] Add DoRA extraction nodes --- __init__.py | 7 + nodes/dora_extract_wd.py | 981 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 988 insertions(+) create mode 100644 nodes/dora_extract_wd.py diff --git a/__init__.py b/__init__.py index c004ead..fb1f91e 100644 --- a/__init__.py +++ b/__init__.py @@ -24,6 +24,10 @@ from .nodes.lora_extract_svd import ( LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, LoRAExtractKnee, LoRAExtractFrobenius ) +from .nodes.dora_extract_wd import ( + DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile, + DoRAExtractKnee, DoRAExtractFrobenius +) from .nodes.lora_resize import ( LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFrobenius, LoRAResizeCumulative, @@ -61,6 +65,9 @@ class ModelUtilsExtension(ComfyExtension): # LoRA Extraction LoRAExtractFixed, LoRAExtractRatio, LoRAExtractQuantile, LoRAExtractKnee, LoRAExtractFrobenius, + # DoRA Extraction + DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile, + DoRAExtractKnee, DoRAExtractFrobenius, # LoRA Resize LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFrobenius, LoRAResizeCumulative, diff --git a/nodes/dora_extract_wd.py b/nodes/dora_extract_wd.py new file mode 100644 index 0000000..0597bf9 --- /dev/null +++ b/nodes/dora_extract_wd.py @@ -0,0 +1,981 @@ +""" +DoRA Extraction via SVD and Weight Decomposition. + +Native implementation of Weight-Decomposed Low-Rank Adapter extraction +with multiple rank selection modes. Extracted models use the DoRA structure. +""" +import fnmatch +import os +import re +import torch +import torch.linalg as linalg +import folder_paths +import comfy.utils +from tqdm import tqdm +from comfy_api.latest import io +from .device_utils import ( + estimate_model_size, prepare_for_large_operation, + cleanup_after_operation, get_device_capabilities +) + +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter + + + + +# ============================================================================= +# Rank Index Functions - Determine LoRA rank from singular values +# ============================================================================= + +def _index_sv_fixed(S: torch.Tensor, dim: int) -> int: + """Fixed rank mode - use specified dimension.""" + return max(1, min(dim, len(S) - 1)) + + +def _index_sv_threshold(S: torch.Tensor, threshold: float) -> int: + """Threshold mode - keep singular values above absolute threshold.""" + rank = int(torch.sum(S > threshold).item()) + return max(1, min(rank, len(S) - 1)) + + +def _index_sv_ratio(S: torch.Tensor, ratio: float) -> int: + """Ratio mode - keep singular values > max(S) / ratio (kohya convention).""" + if ratio <= 0: + return len(S) - 1 + min_sv = S[0] / ratio + rank = int(torch.sum(S > min_sv).item()) + return max(1, min(rank, len(S) - 1)) + + +def _index_sv_cumulative(S: torch.Tensor, target: float, max_rank: int = None) -> int: + """Cumulative mode - keep enough SVs to reach target % of total. + + Calculates relative to max_rank if provided, otherwise relative to full. + """ + if max_rank is not None and max_rank < len(S): + total = torch.sum(S[:max_rank]) + else: + total = torch.sum(S) + + if total < 1e-8: + return 1 + cumsum = torch.cumsum(S, dim=0) / total + rank = int(torch.searchsorted(cumsum, target).item()) + 1 + return max(1, min(rank, len(S) - 1)) + + +def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int: + """Frobenius norm mode - preserve target fraction of Frobenius norm. + + Calculates relative to max_rank if provided, otherwise relative to full. + This means "retain target% of what's achievable within max_rank". + """ + if max_rank is not None and max_rank < len(S): + # Calculate relative to what's achievable within max_rank + S_capped = S[:max_rank] + S_sq = S_capped.pow(2) + total_sq = torch.sum(S_sq) + else: + S_sq = S.pow(2) + total_sq = torch.sum(S_sq) + + if total_sq < 1e-8: + return 1 + + # Cumsum of all S (not capped) to find where we reach target + cumsum = torch.cumsum(S.pow(2), dim=0) / total_sq + rank = int(torch.searchsorted(cumsum, target ** 2).item()) + 1 + return max(1, min(rank, len(S) - 1)) + + +def _index_sv_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int: + """Knee detection - find elbow point on singular value curve.""" + n = len(S) + if n < 3: + return 1 + s_max, s_min = S[0].item(), S[-1].item() + if s_max - s_min < min_sv: + return 1 + # Normalize to [0, 1] + s_norm = (S - s_min) / (s_max - s_min) + x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype) + # Distance from diagonal + distances = (x_norm + s_norm - 1).abs() + rank = torch.argmax(distances).item() + 1 + return max(1, min(rank, n - 1)) + + +def _index_sv_cumulative_knee(S: torch.Tensor, min_sv: float = 1e-8) -> int: + """Knee detection on cumulative singular value curve.""" + n = len(S) + if n < 3: + return 1 + total = torch.sum(S) + if total < min_sv: + return 1 + cumsum = torch.cumsum(S, dim=0) / total + y_min, y_max = cumsum[0].item(), cumsum[-1].item() + if y_max - y_min < min_sv: + return 1 + y_norm = (cumsum - y_min) / (y_max - y_min) + x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype) + distances = (y_norm - x_norm).abs() + rank = torch.argmax(distances).item() + 1 + return max(1, min(rank, n - 1)) + + +def _index_sv_rel_decrease(S: torch.Tensor, tau: float = 0.1) -> int: + """Relative decrease mode - stop when consecutive SV ratio drops below tau.""" + if len(S) < 2: + return 1 + ratios = S[1:] / (S[:-1] + 1e-8) + for k in range(len(ratios)): + if ratios[k] < tau: + return k + 1 + return len(S) + + +def _compute_rank(S: torch.Tensor, mode: str, mode_param: float, + max_rank: int = None) -> int: + """Compute rank based on mode and parameters.""" + if mode == "fixed": + rank = _index_sv_fixed(S, int(mode_param)) + elif mode == "threshold": + rank = _index_sv_threshold(S, mode_param) + elif mode == "ratio": + rank = _index_sv_ratio(S, mode_param) + elif mode == "quantile" or mode == "sv_cumulative": + rank = _index_sv_cumulative(S, mode_param, max_rank) + elif mode == "sv_fro": + rank = _index_sv_fro(S, mode_param, max_rank) + elif mode == "sv_knee": + rank = _index_sv_knee(S) + elif mode == "sv_cumulative_knee": + rank = _index_sv_cumulative_knee(S) + elif mode == "sv_rel_decrease": + rank = _index_sv_rel_decrease(S, mode_param) + else: + rank = _index_sv_fixed(S, int(mode_param) if mode_param else 64) + + if max_rank is not None: + rank = min(rank, max_rank) + return rank + + +# ============================================================================= +# SVD Extraction Functions +# ============================================================================= + +def _svd_extract_linear_lowrank( + weight: torch.Tensor, + rank: int, + device: str, + clamp_quantile: float = 0.99, + niter: int = 2, +) -> tuple: + """ + Low-rank SVD decomposition using svd_lowrank for fixed rank extraction. + + 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") + """ + weight = weight.to(device=device, dtype=torch.float32) + out_dim, in_dim = weight.shape + + # Clamp rank to valid range + max_possible_rank = min(out_dim, in_dim) + rank = max(1, min(rank, max_possible_rank)) + + # Check if decomposition is worthwhile + if rank >= max_possible_rank: + return weight, "full" + + # Use svd_lowrank for faster low-rank approximation + try: + 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: + raise RuntimeError(f"svd_lowrank failed: {e}") + + # Clamp outliers + if clamp_quantile < 1.0 and len(S) > 0: + try: + max_val = torch.quantile(S, clamp_quantile) + S = S.clamp(max=max_val) + except RuntimeError: + pass + + # Construct LoRA matrices + lora_up = U @ torch.diag(S) # [out_dim, rank] + lora_down = Vh # [rank, in_dim] + + # Compute reconstruction diff + diff = weight - (lora_up @ lora_down) + + return (lora_down, lora_up, diff), "low rank" + + +def _svd_extract_linear( + weight: torch.Tensor, + mode: str, + mode_param: float, + device: str, + max_rank: int = None, + clamp_quantile: float = 0.99, + niter: int = 2, +) -> tuple: + """ + SVD decomposition for linear layers. + + For fixed rank mode, uses svd_lowrank which is much faster. + For adaptive modes, uses full SVD to analyze all singular values. + + Returns: + (lora_down, lora_up, rank), "low rank" or (weight, "full") + """ + weight = weight.to(device=device, dtype=torch.float32) + out_dim, in_dim = weight.shape + + # For fixed rank, use the faster svd_lowrank + if mode == "fixed" and mode_param > 0: + 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, 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}") + + # Compute rank based on mode + rank = _compute_rank(S, mode, mode_param, max_rank) + + # Check if decomposition is worthwhile + if rank >= min(out_dim, in_dim): + return weight, "full" + + # Truncate to rank + U = U[:, :rank] + S = S[:rank] + Vh = Vh[:rank, :] + + # Clamp outliers + if clamp_quantile < 1.0 and len(S) > 0: + try: + max_val = torch.quantile(S, clamp_quantile) + S = S.clamp(max=max_val) + except RuntimeError: + pass # Skip clamping if quantile fails + + # Construct LoRA matrices + # lora_up = U @ diag(S), lora_down = Vh + lora_up = U @ torch.diag(S) # [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" + + +def _svd_extract_conv( + weight: torch.Tensor, + mode: str, + mode_param: float, + device: str, + max_rank: int = None, + clamp_quantile: float = 0.99, +) -> tuple: + """ + SVD decomposition for conv2d layers. + + Returns: + (lora_down, lora_up, rank), "low rank" or (weight, "full") + """ + out_ch, in_ch, kh, kw = weight.shape + is_1x1 = (kh == 1 and kw == 1) + + # Flatten for SVD + if is_1x1: + mat = weight.view(out_ch, in_ch) # [out_ch, in_ch] + 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) + + if mode_str == "full": + return weight, "full" + + lora_down, lora_up, diff = result + rank = lora_down.shape[0] + + # Reshape for conv + if is_1x1: + lora_down = lora_down.unsqueeze(-1).unsqueeze(-1) # [rank, in_ch, 1, 1] + lora_up = lora_up.unsqueeze(-1).unsqueeze(-1) # [out_ch, rank, 1, 1] + else: + lora_down = lora_down.reshape(rank, in_ch, kh, kw) + lora_up = lora_up.reshape(out_ch, rank, 1, 1) + + return (lora_down, lora_up, diff), "low rank" + + +# ============================================================================= +# Chunked Extraction for Large Layers +# ============================================================================= + +def _detect_fused_layer(key: str, shape: tuple) -> int: + """Detect fused QKV/MLP layers and return number of chunks.""" + if len(shape) != 2: + return 1 + + out_dim, in_dim = shape + key_lower = key.lower() + + # QKV layers (3 chunks) + if "qkv" in key_lower and out_dim % 3 == 0 and out_dim // 3 >= in_dim // 4: + return 3 + + # KV layers (2 chunks) + if "kv" in key_lower and "qkv" not in key_lower and out_dim % 2 == 0: + return 2 + + # Fused MLP (2 chunks) + if ("linear1" in key_lower or "gate_up" in key_lower) and out_dim % 2 == 0 and out_dim >= in_dim * 2: + return 2 + + return 1 + + +def _extract_chunked_layer( + weight_diff: torch.Tensor, + num_chunks: int, + mode: str, + mode_param: float, + device: str, + max_rank: int = None, +) -> tuple: + """Extract LoRA from fused layer by chunking.""" + out_dim, in_dim = weight_diff.shape + chunk_size = out_dim // num_chunks + + all_lora_up = [] + all_lora_down = [] + + for i in range(num_chunks): + chunk = weight_diff[i * chunk_size:(i + 1) * chunk_size, :] + + try: + result, mode_str = _svd_extract_linear(chunk, mode, mode_param, device, max_rank) + if mode_str == "full": + return None, None, 0 + lora_down, lora_up, _ = result + all_lora_down.append(lora_down) + all_lora_up.append(lora_up) + except Exception: + return None, None, 0 + + # Combine chunks + # Note: Using the first chunk's down matrix as the basis for all chunks + # is a better approximation than a simple mean if the chunks share the same input space (like QKV). + # Ideally, we would use a more sophisticated joint SVD, but this is a reasonable fallback. + combined_lora_up = torch.cat(all_lora_up, dim=0) + combined_lora_down = all_lora_down[0] + rank = combined_lora_down.shape[0] + + return combined_lora_up, combined_lora_down, rank + + +# ============================================================================= +# Pattern Matching +# ============================================================================= + +def _compile_patterns(pattern_string: str, glob_mode: bool = False) -> list: + """Compile patterns from whitespace-separated string. + + In regex mode (default) returns compiled re.Pattern objects. + In glob mode returns plain strings; fnmatch handles matching. + """ + if not pattern_string or not pattern_string.strip(): + return [] + patterns = [] + for p in pattern_string.split(): + p = p.strip() + if not p: + continue + if glob_mode: + patterns.append(p) + else: + try: + patterns.append(re.compile(p)) + except re.error: + pass + return patterns + + +def _matches_any_pattern(key: str, patterns: list, glob_mode: bool = False) -> bool: + """Check if key matches any pattern. + + Glob mode: each pattern is matched as a substring glob using fnmatch + (dots are literal, * matches any sequence of characters). + Regex mode: each pattern is a compiled re.Pattern, matched via search(). + """ + if glob_mode: + return any(fnmatch.fnmatch(key, f"*{p}*") for p in patterns) + return any(p.search(key) for p in patterns) + + +# ============================================================================= +# Main Extraction Function +# ============================================================================= + +def extract_dora_from_files( + model_a_path: str, + model_b_path: str, + mode: str, + linear_param: float, + conv_param: float, + device: str, + save_dtype: str, + output_path: str, + linear_max_rank: int = None, + conv_max_rank: int = None, + clamp_quantile: float = 0.99, + min_diff: float = 0.0, + skip_patterns_str: str = "", + mismatch_mode: str = "skip", + chunk_large_layers: bool = True, + svd_niter: int = 2, + lazy_load: bool = True, + force_clear_cache: bool = True, + glob_skip_patterns: bool = False, +) -> None: + """ + Extract LoRA from difference between two models, writing incrementally to disk. + + Args: + model_a_path: Finetuned model path + model_b_path: Base model path + mode: Rank selection mode + linear_param: Mode parameter for linear layers + conv_param: Mode parameter for conv layers + device: Computation device + save_dtype: Output dtype + output_path: Full path to output safetensors file + linear_max_rank: Max rank for linear layers + conv_max_rank: Max rank for conv layers + clamp_quantile: Quantile for weight clamping + min_diff: Minimum difference threshold + skip_patterns_str: Patterns for layers to skip + mismatch_mode: How to handle mismatches + chunk_large_layers: Enable chunked extraction + glob_skip_patterns: When True, treat skip_patterns as glob (* wildcard). + When False (default), treat as Python regex. + """ + save_torch_dtype = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, + }.get(save_dtype, torch.float16) + + skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns) + + # Prepare memory before heavy operation + total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path) + print(f"[DoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...") + prepare_for_large_operation(total_size_gb * 1.5, torch.device(device)) # 1.5x for SVD headroom + + handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load) + handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load) + + try: + keys_a = set(handler_a.keys()) + keys_b = set(handler_b.keys()) + weight_keys = [k for k in keys_a if k.endswith(".weight")] + pbar = comfy.utils.ProgressBar(len(weight_keys)) + stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} + + def _process_layer(key): + lora_name = _format_lora_key(key) + + if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns): + return "skipped", None + + # Load tensors with pinned memory for CUDA + use_pinned = device == 'cuda' + + if key not in keys_b: + if mismatch_mode == "skip": + return "skipped", None + if mismatch_mode == "error": + raise ValueError(f"Key {key} not found in model B") + + cpu_a = handler_a.get_tensor(key) + if use_pinned: + weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32) + else: + weight_diff = cpu_a.to(device=device, dtype=torch.float32) + del cpu_a + + # Without base, DoRA fallback just treats tensor_a as diff + dora_scale = None + else: + cpu_a = handler_a.get_tensor(key) + cpu_b = handler_b.get_tensor(key) + + if use_pinned: + tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32) + tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32) + else: + tensor_a = cpu_a.to(device=device, dtype=torch.float32) + tensor_b = cpu_b.to(device=device, dtype=torch.float32) + del cpu_a, cpu_b + + if tensor_a.shape != tensor_b.shape: + if mismatch_mode == "skip": + del tensor_a, tensor_b + return "skipped", None + if mismatch_mode == "error": + raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}") + + weight_diff = tensor_a + dora_scale = None + del tensor_b + else: + # 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 + + # Skip small differences + if min_diff > 0 and weight_diff.abs().max() < min_diff: + del weight_diff + return "skipped", None + + # Skip 1D tensors + if weight_diff.ndim < 2: + del weight_diff + return "skipped", None + + is_conv = weight_diff.ndim == 4 + layer_results = {} + + try: + if is_conv: + result, mode_str = _svd_extract_conv( + weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile + ) + else: + result, mode_str = _svd_extract_linear( + weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter + ) + except Exception as e: + # Try chunked extraction for large tensors + if chunk_large_layers and not is_conv: + num_chunks = _detect_fused_layer(key, weight_diff.shape) + if num_chunks > 1: + print(f"[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 + ) + if lora_up is not None: + layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype) + if dora_scale is not None: + layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous() + del weight_diff + return "chunked", layer_results + + print(f"[DoRA Extract] Failed: {key}: {e}") + del weight_diff + return "skipped", None + + # Store result + if mode_str == "full": + layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous() + status = "full" + else: + lora_down, lora_up, rank_or_diff = result + # rank is the size of lora_down's first dimension + rank = lora_down.shape[0] + layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype) + if dora_scale is not None: + layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous() + status = "extracted" + + del weight_diff + return status, layer_results + + writer = IncrementalSafetensorsWriter(output_path) + writer.__enter__() + try: + for key in tqdm(weight_keys, desc="Extracting DoRA", unit="layers"): + status, layer_sd = _process_layer(key) + stats[status] += 1 + if layer_sd: + writer.write_dict(layer_sd) + + if force_clear_cache: + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + pbar.update(1) + finally: + writer.__exit__(None, None, None) + + print(f"[DoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " + f"{stats['full']} full, {stats['skipped']} skipped") + + finally: + handler_a.__exit__(None, None, None) + handler_b.__exit__(None, None, None) + cleanup_after_operation() + + +def _build_lora_output_path(output_filename: str) -> str: + """Build output path for LoRA file.""" + output_dir = folder_paths.get_folder_paths("loras")[0] + os.makedirs(output_dir, exist_ok=True) + return os.path.join(output_dir, f"{output_filename.strip()}.safetensors") + + +# ============================================================================= +# Node Definitions +# ============================================================================= + +def _get_model_inputs(): + return [ + io.Combo.Input("model_a", options=folder_paths.get_filename_list("diffusion_models"), + tooltip="Finetuned model (A - B = LoRA)"), + io.Combo.Input("model_b", options=folder_paths.get_filename_list("diffusion_models"), + tooltip="Base model (A - B = LoRA)"), + ] + + +def _format_lora_key(key: str) -> str: + """ + Format the key for LoRA saving. + Always standardizes to 'diffusion_model.' prefix for maximum compatibility. + + ComfyUI maps 'diffusion_model.' generically for all model types, making + it universally compatible. The 'transformer.' prefix only works for a subset of + model types via model-specific mapping code and is intentionally not used here. + """ + if key.endswith(".weight"): + key = key[:-7] + + # Handle ComfyUI Checkpoint wrapper (ldm/sgm format: model.diffusion_model.*) + if key.startswith("model.diffusion_model."): + inner_key = key[22:] # strip len("model.diffusion_model.") + return f"diffusion_model.{inner_key}" + + # Handle net.* wrapper + if key.startswith("net."): + inner_key = key[4:] # strip len("net.") + return f"diffusion_model.{inner_key}" + + # Handle direct Diffusers keys (transformer_blocks.* / single_transformer_blocks.* without prefix) + if key.startswith("transformer_blocks") or key.startswith("single_transformer_blocks"): + return f"diffusion_model.{key}" + + # Handle legacy lora_unet_ prefix (for resizing without base) + if key.startswith("lora_unet_"): + return key + + # Handle already-prefixed keys + if key.startswith("diffusion_model."): + return key + + # Absorb now-invalid transformer. prefix into diffusion_model. for compatibility + if key.startswith("transformer."): + return f"diffusion_model.{key[12:]}" + + # Handle known Diffusers UNet blocks + if any(key.startswith(p) for p in ["down_blocks", "up_blocks", "mid_block", "conv_in", "conv_out", "time_embedding", "class_embedding"]): + return f"diffusion_model.{key}" + + # Default fallback + return f"diffusion_model.{key}" + + +def _get_common_inputs(): + return [ + io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), + io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), + io.Boolean.Input("chunk_large_layers", default=False, + tooltip="Split large fused layers (QKV, MLP) into chunks"), + io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01, + tooltip="Clamp outlier singular values"), + io.Float.Input("min_diff", default=0.0, min=0.0, max=1.0, step=0.001, + tooltip="Skip layers with max difference below this"), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.String.Input("output_filename", default="extracted_lora"), + io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"), + io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"), + io.String.Input("skip_patterns", default="", multiline=True, + tooltip="Patterns for layers to skip (regex or glob depending on glob_skip_patterns)"), + io.Boolean.Input("glob_skip_patterns", default=False, + tooltip="When True, skip_patterns use glob syntax (* = any sequence, ? = any char, dots are literal). " + "When False (default), patterns are Python regex matched as substrings."), + ] + + +class DoRAExtractFixed(io.ComfyNode): + """Extract DoRA with fixed rank.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRAExtractFixed", + display_name="DoRA Extract (Fixed Rank)", + category="ModelUtils/DoRA", + description="Extract DoRA with specified fixed rank for each layer type.", + inputs=[ + *_get_model_inputs(), + io.Int.Input("linear_dim", default=64, min=1, max=16384, + tooltip="Rank for linear/attention layers"), + io.Int.Input("conv_dim", default=32, min=1, max=16384, + tooltip="Rank for conv layers"), + io.Int.Input("svd_niter", default=2, min=0, max=10, + tooltip="SVD power iterations (higher = more accurate but slower)"), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers, + clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_from_files( + model_a_path, model_b_path, "fixed", linear_dim, conv_dim, + device, save_dtype, output_path, linear_dim, conv_dim, + clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter, + lazy_load, force_clear_cache, glob_skip_patterns + ) + + return io.NodeOutput(output_path) + + +class DoRAExtractRatio(io.ComfyNode): + """Extract DoRA with singular value ratio threshold.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRAExtractRatio", + display_name="DoRA Extract (Ratio)", + category="ModelUtils/DoRA", + description="Keep singular values > max(S) / ratio.", + inputs=[ + *_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 = 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("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, + tooltip="Maximum rank for conv layers"), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_from_files( + model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio, + 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 + ) + + return io.NodeOutput(output_path) + + +class DoRAExtractQuantile(io.ComfyNode): + """Extract DoRA with cumulative singular value quantile.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRAExtractQuantile", + display_name="DoRA Extract (Quantile)", + category="ModelUtils/DoRA", + description="Keep enough singular values to reach target percentage.", + inputs=[ + *_get_model_inputs(), + io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, + 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("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, + tooltip="Maximum rank for conv layers"), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_from_files( + model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile, + 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 + ) + + return io.NodeOutput(output_path) + + +class DoRAExtractKnee(io.ComfyNode): + """Extract DoRA with automatic knee detection.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRAExtractKnee", + display_name="DoRA Extract (Knee Detection)", + category="ModelUtils/DoRA", + description="Automatically find optimal rank using knee detection on singular value curve.", + inputs=[ + *_get_model_inputs(), + io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], + default="sv_knee", tooltip="Knee detection method"), + 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, + tooltip="Maximum rank for conv layers"), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_from_files( + model_a_path, model_b_path, knee_method, 0, 0, + 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 + ) + + return io.NodeOutput(output_path) + + +class DoRAExtractFrobenius(io.ComfyNode): + """Extract DoRA preserving target fraction of Frobenius norm.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRAExtractFrobenius", + display_name="DoRA Extract (Frobenius)", + category="ModelUtils/DoRA", + description="Preserve target fraction of Frobenius norm.", + inputs=[ + *_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 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("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, + tooltip="Maximum rank for conv layers"), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, + chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_from_files( + model_a_path, model_b_path, "sv_fro", linear_target, conv_target, + 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 + ) + + return io.NodeOutput(output_path) From 3a505e704fa20d0e96e8c756ebdf1f2072e07508 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Fri, 29 May 2026 14:36:26 +0200 Subject: [PATCH 2/2] Add work in progress experimental learned extraction nodes --- __init__.py | 7 + nodes/dora_learned_wd.py | 790 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 797 insertions(+) create mode 100644 nodes/dora_learned_wd.py diff --git a/__init__.py b/__init__.py index fb1f91e..22ab52e 100644 --- a/__init__.py +++ b/__init__.py @@ -28,6 +28,10 @@ from .nodes.dora_extract_wd import ( DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile, DoRAExtractKnee, DoRAExtractFrobenius ) +from .nodes.dora_learned_wd import ( + DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile, + DoRALearnedExtractKnee, DoRALearnedExtractFrobenius +) from .nodes.lora_resize import ( LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFrobenius, LoRAResizeCumulative, @@ -68,6 +72,9 @@ class ModelUtilsExtension(ComfyExtension): # DoRA Extraction DoRAExtractFixed, DoRAExtractRatio, DoRAExtractQuantile, DoRAExtractKnee, DoRAExtractFrobenius, + # Learned DoRA Extraction + DoRALearnedExtractFixed, DoRALearnedExtractRatio, DoRALearnedExtractQuantile, + DoRALearnedExtractKnee, DoRALearnedExtractFrobenius, # LoRA Resize LoRAResizeFixed, LoRAResizeRatio, LoRAResizeFrobenius, LoRAResizeCumulative, diff --git a/nodes/dora_learned_wd.py b/nodes/dora_learned_wd.py new file mode 100644 index 0000000..be38f07 --- /dev/null +++ b/nodes/dora_learned_wd.py @@ -0,0 +1,790 @@ +""" +Learned DoRA Extraction via Optimization. + +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 +from comfy_api.latest import io +from .device_utils import ( + estimate_model_size, prepare_for_large_operation, + cleanup_after_operation +) + +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter + +# Import 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_conv, _detect_fused_layer, _compile_patterns, + _matches_any_pattern, _build_lora_output_path, _get_model_inputs, + _format_lora_key, _get_common_inputs +) + + +def _adaptive_lr_update_cosine(curr_lr: float, improved: bool, worse_loss_counter: int, iteration: int, tensor_shape: tuple, min_lr: float, early_stop_stall: int, lr_factor: float = 0.95, lr_cooldown: int = 1) -> tuple: + M, N = tensor_shape + shape_ratio = abs(M - N) / max(M, N) if max(M, N) > 0 else 0 + + t = min(worse_loss_counter / max(early_stop_stall, 1), 1.0) + u_factor = (1 + math.cos(2 * math.pi * t)) / 2 + + if improved: + base_boost = 1.25 + distance = base_boost - 1.0 + scaled_distance = distance * (1.0 + 0.5 * shape_ratio) + boost_mult = 1.0 + scaled_distance * (1.0 - u_factor) + return min(curr_lr * boost_mult, 100.0), True + else: + cooldown = max(lr_cooldown, 1) + if iteration % cooldown != 0: + return curr_lr, False + + min_decay = lr_factor - 0.03 * shape_ratio + max_decay = 0.995 + decay_mult = min_decay + (max_decay - min_decay) * u_factor + return max(curr_lr * decay_mult, min_lr), True + + +def _compute_shape_aware_plateau_params(M: int, N: int, lr_patience: int, lr_factor: float, lr_cooldown: int, lr_shape_influence: float = 1.0) -> tuple: + if lr_shape_influence > 0: + aspect_ratio = max(M, N) / min(M, N) if min(M, N) > 0 else 1.0 + ar_factor = math.sqrt(aspect_ratio) + blend = lr_shape_influence + + effective_patience = lr_patience + raw_factor = lr_factor + aggressive_factor = raw_factor**ar_factor + effective_factor = raw_factor + (aggressive_factor - raw_factor) * blend + effective_cooldown = lr_cooldown + else: + effective_patience = lr_patience + effective_factor = lr_factor + effective_cooldown = lr_cooldown + + return effective_patience, effective_factor, effective_cooldown + + +def _optimize_dora( + tensor_base: torch.Tensor, + tensor_ft: torch.Tensor, + init_down: torch.Tensor, + init_up: torch.Tensor, + optimize_iters: int, + learning_rate: float, + optimizer_type: str = "prodigy", + lr_schedule: str = "plateau", + lr_patience: int = 2, + lr_factor: float = 0.9, + lr_cooldown: int = 2, + early_stop_loss: float = 1e-6, + early_stop_stall: int = 2000, + early_stop_lr: float = 9.01e-9, + layer_name: str = "Layer", +) -> tuple: + """ + Run gradient descent to optimize lora_down and lora_up against the actual DoRA formula. + + Args: + tensor_base: Base model weight + tensor_ft: Finetuned model weight + init_down: SVD-initialized lora_down matrix + init_up: SVD-initialized lora_up matrix + optimize_iters: Number of iterations for AdamW + learning_rate: Learning rate for AdamW + optimizer_type: Optimizer choice (adamw, radam, prodigy) + lr_schedule: LR Schedule choice (exponential, plateau, adaptive) + early_stop_loss: Target loss to stop optimizing + early_stop_stall: Number of non-improving iterations before stopping + + Returns: + (optimized_down, optimized_up) + """ + if optimize_iters <= 0: + return init_down, init_up + + is_conv = tensor_base.ndim == 4 + if is_conv: + out_ch, in_ch, kh, kw = tensor_base.shape + flat_base = tensor_base.reshape(out_ch, -1) + flat_ft = tensor_ft.reshape(out_ch, -1) + + # Flatten init matrices for matmul: (out, rank) @ (rank, in*k*k) + rank = init_down.shape[0] + flat_up = init_up.reshape(out_ch, rank) + flat_down = init_down.reshape(rank, -1) + else: + flat_base = tensor_base + flat_ft = tensor_ft + flat_up = init_up + flat_down = init_down + + # Magnitude of the finetuned model (fixed DoRA scale) + m = torch.linalg.norm(flat_ft, dim=1, keepdim=True).detach() + eps = 1e-8 + + # Set up learnable parameters + with torch.inference_mode(False), torch.enable_grad(): + # Clone inference tensors to normal tensors for autograd + flat_base_opt = flat_base.clone() + flat_ft_opt = flat_ft.clone() + m_opt = m.clone() + + # Free original tensors to save memory + del flat_base, flat_ft, m + + B = torch.nn.Parameter(flat_up.clone()) + A = torch.nn.Parameter(flat_down.clone()) + + opt_type = optimizer_type.lower() + if opt_type == "radam": + optimizer = torch.optim.RAdam([B, A], lr=learning_rate) + elif opt_type == "prodigy": + try: + from prodigyplus.prodigy_plus_schedulefree import ProdigyPlusScheduleFree + optimizer = ProdigyPlusScheduleFree([B, A], lr=learning_rate, use_schedulefree=False) + except ImportError: + print("ProdigyPlusScheduleFree optimizer not found. Falling back to AdamW. To install: pip install prodigyplus") + optimizer = torch.optim.AdamW([B, A], lr=learning_rate) + else: + optimizer = torch.optim.AdamW([B, A], lr=learning_rate) + + # Learning rate scheduling and early stopping state + best_loss = float('inf') + worse_loss_counter = 0 + plateau_counter = 0 + cooldown_counter = 0 + current_lr = learning_rate + schedule_type = lr_schedule.lower() + + M_dim, N_dim = flat_base_opt.shape[0], flat_base_opt.shape[1] + eff_patience, eff_factor, eff_cooldown = _compute_shape_aware_plateau_params(M_dim, N_dim, lr_patience, lr_factor, lr_cooldown) + + # Optimization loop + opt_pbar = tqdm(range(optimize_iters), desc=f" └─ Optimizing {layer_name}", leave=False, dynamic_ncols=True, mininterval=0.5) + for i in opt_pbar: + optimizer.zero_grad() + + # 1. Reconstruct directional weight: W_dir = W_base + B @ A + w_approx = flat_base_opt + (B @ A) + + # 2. Normalize directional weight + approx_norm = torch.linalg.norm(w_approx, dim=1, keepdim=True) + eps + + # 3. Scale by target magnitude + w_reconstructed = m_opt * (w_approx / approx_norm) + + # 4. Compute MSE loss against actual finetuned weight + loss = torch.nn.functional.mse_loss(w_reconstructed, flat_ft_opt) + loss_val = loss.item() + + if i == 0: + initial_loss = loss_val + best_loss = loss_val + + loss.backward() + optimizer.step() + + # --- Early Stopping & Schedulers --- + improved = False + if loss_val < best_loss * 0.9999: # 0.01% improvement threshold + best_loss = loss_val + worse_loss_counter = 0 + plateau_counter = 0 + improved = True + else: + worse_loss_counter += 1 + plateau_counter += 1 + + if loss_val <= early_stop_loss: + break + + if worse_loss_counter > early_stop_stall: + break + + if current_lr <= early_stop_lr: + break + + # LR Schedulers + if schedule_type == "adaptive": + current_lr, updated = _adaptive_lr_update_cosine( + current_lr, improved, worse_loss_counter, i, (M_dim, N_dim), early_stop_lr, early_stop_stall, lr_factor, lr_cooldown + ) + if updated: + for param_group in optimizer.param_groups: + param_group['lr'] = current_lr + elif schedule_type == "plateau": + if cooldown_counter > 0: + cooldown_counter -= 1 + elif plateau_counter >= max(eff_patience, 1): # Decay if patience threshold reached + current_lr = max(current_lr * eff_factor, early_stop_lr) + for param_group in optimizer.param_groups: + param_group['lr'] = current_lr + plateau_counter = 0 # reset counter after decay + cooldown_counter = eff_cooldown + elif schedule_type == "exponential": + current_lr = max(current_lr * 0.99, early_stop_lr) + for param_group in optimizer.param_groups: + param_group['lr'] = current_lr + + if i % 10 == 0 or i == optimize_iters - 1: + if schedule_type == "plateau": + opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "plat": f"{plateau_counter}/{eff_patience}"}, refresh=False) + else: + opt_pbar.set_postfix({"loss": f"{loss_val:.4e}", "lr": f"{current_lr:.2e}", "stall": f"{worse_loss_counter}"}, refresh=False) + + final_loss = loss.item() + opt_pbar.close() + print(f" └─ {layer_name} | Loss: {initial_loss:.6f} -> {final_loss:.6f} (stopped at iter {i+1})") + + final_up = B.detach() + final_down = A.detach() + + # Reshape back to original dimensions + if is_conv: + final_up = final_up.reshape(out_ch, rank, 1, 1) + final_down = final_down.reshape(rank, in_ch, kh, kw) + + return final_down, final_up + + +def _extract_chunked_learned_layer( + tensor_base: torch.Tensor, + tensor_ft: torch.Tensor, + num_chunks: int, + mode: str, + mode_param: float, + device: str, + max_rank: int = None, + optimize_iters: int = 500, + learning_rate: float = 1.0, + optimizer_type: str = "prodigy", + lr_schedule: str = "plateau", + lr_patience: int = 2, + lr_factor: float = 0.9, + lr_cooldown: int = 2, + early_stop_loss: float = 1e-6, + early_stop_stall: int = 2000, + early_stop_lr: float = 9.01e-9, + layer_name: str = "Layer", +) -> tuple: + """Extract Learned DoRA from fused layer by chunking.""" + out_dim, in_dim = tensor_base.shape + chunk_size = out_dim // num_chunks + + all_lora_up = [] + all_lora_down = [] + + for i in range(num_chunks): + chunk_base = tensor_base[i * chunk_size:(i + 1) * chunk_size, :] + chunk_ft = tensor_ft[i * chunk_size:(i + 1) * chunk_size, :] + + norm_base = torch.linalg.norm(chunk_base, dim=1, keepdim=True) + norm_ft = torch.linalg.norm(chunk_ft, dim=1, keepdim=True) + eps = 1e-8 + + w_dir = chunk_ft * (norm_base / (norm_ft + eps)) + weight_diff = w_dir - chunk_base + + try: + result, mode_str = _svd_extract_linear(weight_diff, mode, mode_param, device, max_rank) + if mode_str == "full": + return None, None, 0 + lora_down, lora_up, _ = result + + # Apply learned optimization + lora_down, lora_up = _optimize_dora( + chunk_base, chunk_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=f"{layer_name} (Chunk {i+1}/{num_chunks})" + ) + + all_lora_down.append(lora_down) + all_lora_up.append(lora_up) + except Exception: + return None, None, 0 + + combined_lora_up = torch.cat(all_lora_up, dim=0) + combined_lora_down = all_lora_down[0] + rank = combined_lora_down.shape[0] + + return combined_lora_up, combined_lora_down, rank + + +def extract_dora_learned_from_files( + model_a_path: str, + model_b_path: str, + mode: str, + linear_param: float, + conv_param: float, + device: str, + save_dtype: str, + output_path: str, + optimize_iters: int = 500, + learning_rate: float = 1.0, + optimizer_type: str = "prodigy", + lr_schedule: str = "plateau", + lr_patience: int = 2, + lr_factor: float = 0.9, + lr_cooldown: int = 2, + early_stop_loss: float = 1e-6, + early_stop_stall: int = 2000, + early_stop_lr: float = 9.01e-9, + linear_max_rank: int = None, + conv_max_rank: int = None, + clamp_quantile: float = 0.99, + min_diff: float = 0.0, + skip_patterns_str: str = "", + mismatch_mode: str = "skip", + chunk_large_layers: bool = True, + svd_niter: int = 2, + lazy_load: bool = True, + force_clear_cache: bool = True, + glob_skip_patterns: bool = False, +) -> None: + """ + Extract Learned DoRA from difference between two models. + """ + save_torch_dtype = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, + }.get(save_dtype, torch.float16) + + skip_patterns = _compile_patterns(skip_patterns_str, glob_mode=glob_skip_patterns) + + total_size_gb = estimate_model_size(model_a_path) + estimate_model_size(model_b_path) + print(f"[Learned DoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...") + prepare_for_large_operation(total_size_gb * 1.5, torch.device(device)) + + handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load) + handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load) + + try: + keys_a = set(handler_a.keys()) + keys_b = set(handler_b.keys()) + weight_keys = [k for k in keys_a if k.endswith(".weight")] + pbar = comfy.utils.ProgressBar(len(weight_keys)) + stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} + + def _process_layer(key): + lora_name = _format_lora_key(key) + + 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 + else: + cpu_a = handler_a.get_tensor(key) + cpu_b = handler_b.get_tensor(key) + + if use_pinned: + tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32) + tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32) + else: + tensor_a = cpu_a.to(device=device, dtype=torch.float32) + tensor_b = cpu_b.to(device=device, dtype=torch.float32) + del cpu_a, cpu_b + + if tensor_a.shape != tensor_b.shape: + if mismatch_mode == "skip": + del tensor_a, tensor_b + return "skipped", None + if mismatch_mode == "error": + raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}") + + weight_diff = tensor_a + tensor_base = torch.zeros_like(tensor_a) + tensor_ft = tensor_a + dora_scale = None + del tensor_b + else: + tensor_base = tensor_b + tensor_ft = tensor_a + + 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 min_diff > 0 and weight_diff.abs().max() < min_diff: + del weight_diff, tensor_base, tensor_ft + return "skipped", None + + if weight_diff.ndim < 2: + del weight_diff, tensor_base, tensor_ft + return "skipped", None + + is_conv = weight_diff.ndim == 4 + layer_results = {} + + 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 + ) + else: + result, mode_str = _svd_extract_linear( + weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter + ) + + # 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 + ) + result = (lora_down, lora_up, _) + + except Exception as e: + if chunk_large_layers and not is_conv: + num_chunks = _detect_fused_layer(key, weight_diff.shape) + 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, + optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown, + early_stop_loss, early_stop_stall, early_stop_lr, + layer_name=lora_name + ) + if lora_up is not None: + layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype) + if dora_scale is not None: + layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous() + del weight_diff, tensor_base, tensor_ft + return "chunked", layer_results + + print(f"[Learned DoRA] Failed: {key}: {e}") + del weight_diff, tensor_base, tensor_ft + return "skipped", None + + if mode_str == "full": + layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous() + status = "full" + else: + lora_down, lora_up, _ = result + rank = lora_down.shape[0] + layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() + layer_results[f"{lora_name}.alpha"] = torch.tensor(rank, dtype=save_torch_dtype) + if dora_scale is not None: + layer_results[f"{lora_name}.dora_scale"] = dora_scale.to(save_torch_dtype).cpu().contiguous() + status = "extracted" + + del weight_diff, tensor_base, tensor_ft + return status, layer_results + + 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) + stats[status] += 1 + if layer_sd: + writer.write_dict(layer_sd) + + if force_clear_cache: + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + pbar.update(1) + finally: + writer.__exit__(None, None, None) + + print(f"[Learned DoRA] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " + f"{stats['full']} full, {stats['skipped']} skipped") + + finally: + handler_a.__exit__(None, None, None) + handler_b.__exit__(None, None, None) + cleanup_after_operation() + + +# ============================================================================= +# Node Definitions +# ============================================================================= + +def _get_learned_inputs(): + return [ + io.Int.Input("optimize_iters", default=500, min=0, max=10000, + tooltip="Number of gradient descent iterations (0 = skip optimization)"), + io.Float.Input("learning_rate", default=1.0, min=0.0001, max=100.0, step=0.001, + tooltip="Base learning rate. AdamW prefers ~0.01, Prodigy prefers ~1.0"), + io.Combo.Input("optimizer", options=["prodigy", "adamw", "radam"], default="prodigy", + tooltip="Optimization algorithm"), + io.Combo.Input("lr_schedule", options=["plateau", "adaptive", "exponential", "constant"], default="plateau", + tooltip="Learning rate scheduling strategy"), + io.Int.Input("lr_patience", default=2, min=0, max=1000, tooltip="Steps to wait before decaying LR (Plateau)"), + io.Float.Input("lr_factor", default=0.9, min=0.01, max=1.0, step=0.01, tooltip="Factor to decay LR by (Plateau)"), + io.Int.Input("lr_cooldown", default=2, min=0, max=1000, tooltip="Steps to wait after decay before checking again (Plateau)"), + io.Float.Input("early_stop_loss", default=1e-6, min=0.0, max=1.0, step=1e-9, + tooltip="Stop early if MSE loss drops below this value"), + io.Int.Input("early_stop_stall", default=2000, min=0, max=10000, + tooltip="Stop early if loss doesn't improve for this many iterations"), + io.Float.Input("early_stop_lr", default=9.01e-9, min=0.0, max=1.0, step=1e-10, + tooltip="Stop early if learning rate drops below this value"), + ] + + +class DoRALearnedExtractFixed(io.ComfyNode): + """Extract Learned DoRA with fixed rank.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRALearnedExtractFixed", + display_name="DoRA Learned Extract (Fixed Rank) (WIP)", + category="ModelUtils/DoRA", + description="Extract DoRA via MSE Loss minimization with fixed rank.", + is_experimental=True, + inputs=[ + *_get_model_inputs(), + io.Int.Input("linear_dim", default=64, min=1, max=16384, tooltip="Rank for linear/attention layers"), + io.Int.Input("conv_dim", default=32, min=1, max=16384, tooltip="Rank for conv layers"), + *_get_learned_inputs(), + io.Int.Input("svd_niter", default=2, min=0, max=10, tooltip="SVD iterations for initialization"), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_dim, conv_dim, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, svd_niter, chunk_large_layers, + clamp_quantile, min_diff, mismatch_mode, output_filename, + save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_learned_from_files( + model_a_path, model_b_path, "fixed", linear_dim, conv_dim, + 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_dim, conv_dim, + clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter, + lazy_load, force_clear_cache, glob_skip_patterns + ) + + return io.NodeOutput(output_path) + + +class DoRALearnedExtractRatio(io.ComfyNode): + """Extract Learned DoRA with singular value ratio threshold.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRALearnedExtractRatio", + display_name="DoRA Learned Extract (Ratio) (WIP)", + category="ModelUtils/DoRA", + description="Learned extraction keeping SVs > max(S) / ratio.", + is_experimental=True, + inputs=[ + *_get_model_inputs(), + io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1), + io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_learned_inputs(), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank, + 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: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_learned_from_files( + 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 + ) + + return io.NodeOutput(output_path) + + +class DoRALearnedExtractQuantile(io.ComfyNode): + """Extract Learned DoRA with cumulative singular value quantile.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRALearnedExtractQuantile", + display_name="DoRA Learned Extract (Quantile) (WIP)", + category="ModelUtils/DoRA", + description="Learned extraction keeping enough SVs to reach target percentage.", + is_experimental=True, + inputs=[ + *_get_model_inputs(), + io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01), + io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_learned_inputs(), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank, + 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: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_learned_from_files( + 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 + ) + + return io.NodeOutput(output_path) + + +class DoRALearnedExtractKnee(io.ComfyNode): + """Extract Learned DoRA with automatic knee detection.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRALearnedExtractKnee", + display_name="DoRA Learned Extract (Knee Detection) (WIP)", + category="ModelUtils/DoRA", + description="Learned extraction automatically finding optimal rank using knee detection.", + is_experimental=True, + inputs=[ + *_get_model_inputs(), + io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee"), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_learned_inputs(), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, + 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: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_learned_from_files( + 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 + ) + + return io.NodeOutput(output_path) + + +class DoRALearnedExtractFrobenius(io.ComfyNode): + """Extract Learned DoRA preserving target fraction of Frobenius norm.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DoRALearnedExtractFrobenius", + display_name="DoRA Learned Extract (Frobenius) (WIP)", + category="ModelUtils/DoRA", + description="Learned extraction preserving target fraction of Frobenius norm.", + is_experimental=True, + inputs=[ + *_get_model_inputs(), + io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01), + io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01), + io.Int.Input("linear_max_rank", default=128, min=1, max=16384), + io.Int.Input("conv_max_rank", default=128, min=1, max=16384), + *_get_learned_inputs(), + *_get_common_inputs(), + ], + outputs=[io.String.Output(display_name="output_path")], + is_output_node=True, + ) + + @classmethod + def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank, + 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: + + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + output_path = _build_lora_output_path(output_filename) + + extract_dora_learned_from_files( + 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 + ) + + return io.NodeOutput(output_path)