From b497022c3c966eaf610a47dbe9ce2bf914183f58 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Tue, 24 Feb 2026 17:39:48 +0100 Subject: [PATCH] expand on functionality and reduce restrictions when extracting LoRA --- nodes/lora_extract_svd.py | 153 ++++++++++++++++++++------------------ nodes/lora_merger.py | 110 +++++++++++++-------------- nodes/lora_resize.py | 29 +++++++- nodes/merger.py | 2 +- 4 files changed, 161 insertions(+), 133 deletions(-) diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index c53df8c..cd3e715 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -46,14 +46,14 @@ def _index_sv_ratio(S: torch.Tensor, ratio: float) -> int: 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 @@ -63,7 +63,7 @@ def _index_sv_cumulative(S: torch.Tensor, target: float, max_rank: int = None) - 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". """ @@ -75,10 +75,10 @@ def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int: 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 @@ -172,27 +172,27 @@ def _svd_extract_linear_lowrank( ) -> 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 - 1)) - + rank = max(1, min(rank, max_possible_rank)) + # Check if decomposition is worthwhile - if rank >= max_possible_rank // 2: + 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) @@ -200,7 +200,7 @@ def _svd_extract_linear_lowrank( 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: @@ -208,14 +208,14 @@ def _svd_extract_linear_lowrank( 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" @@ -239,7 +239,7 @@ def _svd_extract_linear( """ 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) @@ -257,7 +257,7 @@ def _svd_extract_linear( rank = _compute_rank(S, mode, mode_param, max_rank) # Check if decomposition is worthwhile - if rank >= min(out_dim, in_dim) // 2: + if rank >= min(out_dim, in_dim): return weight, "full" # Truncate to rank @@ -303,7 +303,7 @@ def _svd_extract_conv( # Flatten for SVD if is_1x1: - mat = weight.squeeze() # [out_ch, in_ch] + mat = weight.view(out_ch, in_ch) # [out_ch, in_ch] else: mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k] @@ -382,8 +382,11 @@ def _extract_chunked_layer( 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 = torch.stack(all_lora_down, dim=0).mean(dim=0) + combined_lora_down = all_lora_down[0] rank = combined_lora_down.shape[0] return combined_lora_up, combined_lora_down, rank @@ -476,39 +479,40 @@ def extract_lora_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)) - stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} - for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"): + def _process_layer(key): lora_name = _format_lora_key(key) if _matches_any_pattern(key, skip_patterns): - stats["skipped"] += 1 - pbar.update(1) - continue + return "skipped", None # Load tensors with pinned memory for CUDA use_pinned = device == 'cuda' - + if key not in keys_b: if mismatch_mode == "skip": - stats["skipped"] += 1 - pbar.update(1) - continue + 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 + + 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) @@ -519,10 +523,12 @@ def extract_lora_from_files( if tensor_a.shape != tensor_b.shape: if mismatch_mode == "skip": - stats["skipped"] += 1 - pbar.update(1) del tensor_a, tensor_b - continue + return "skipped", None + if mismatch_mode == "error": + raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}") + + # For zeros/fallback, use tensor_a as the difference weight_diff = tensor_a del tensor_b else: @@ -531,19 +537,16 @@ def extract_lora_from_files( # Skip small differences if min_diff > 0 and weight_diff.abs().max() < min_diff: - stats["skipped"] += 1 del weight_diff - pbar.update(1) - continue + return "skipped", None # Skip 1D tensors if weight_diff.ndim < 2: - stats["skipped"] += 1 del weight_diff - pbar.update(1) - continue + return "skipped", None is_conv = weight_diff.ndim == 4 + layer_results = {} try: if is_conv: @@ -564,32 +567,36 @@ def extract_lora_from_files( weight_diff, num_chunks, mode, linear_param, device, linear_max_rank ) if lora_up is not None: - output_sd[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() - output_sd[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() - output_sd[f"{lora_name}.alpha"] = torch.tensor(rank).to(save_torch_dtype) - stats["chunked"] += 1 + 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).to(save_torch_dtype) del weight_diff - pbar.update(1) - continue + return "chunked", layer_results print(f"[LoRA Extract] Failed: {key}: {e}") - stats["skipped"] += 1 del weight_diff - pbar.update(1) - continue + return "skipped", None # Store result if mode_str == "full": - output_sd[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous() - stats["full"] += 1 + layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous() + status = "full" else: lora_down, lora_up, _ = result - output_sd[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous() - output_sd[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() - output_sd[f"{lora_name}.alpha"] = torch.tensor(lora_down.shape[0]).to(save_torch_dtype) - stats["extracted"] += 1 + 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(lora_down.shape[0]).to(save_torch_dtype) + status = "extracted" del weight_diff + return status, layer_results + + for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"): + status, layer_sd = _process_layer(key) + stats[status] += 1 + if layer_sd: + output_sd.update(layer_sd) + if force_clear_cache: import gc gc.collect() @@ -706,7 +713,7 @@ def _format_lora_key(key: str) -> str: # Handle already correct prefixes if key.startswith("diffusion_model.") or key.startswith("transformer."): return key - + # Handle known Unet blocks (Diffusers format UNet) 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}" @@ -746,9 +753,9 @@ class LoRAExtractFixed(io.ComfyNode): description="Extract LoRA with specified fixed rank for each layer type.", inputs=[ *_get_model_inputs(), - io.Int.Input("linear_dim", default=64, min=1, max=4096, + 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=4096, + 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)"), @@ -768,7 +775,7 @@ class LoRAExtractFixed(io.ComfyNode): output_sd = extract_lora_from_files( model_a_path, model_b_path, "fixed", linear_dim, conv_dim, - "lora_unet_", device, save_dtype, linear_dim, conv_dim, + device, save_dtype, linear_dim, conv_dim, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter, lazy_load, force_clear_cache ) @@ -792,9 +799,9 @@ 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("linear_max_rank", default=128, min=1, max=3072, + 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=3072, + io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum rank for conv layers"), *_get_common_inputs(), ], @@ -812,7 +819,7 @@ class LoRAExtractRatio(io.ComfyNode): output_sd = extract_lora_from_files( model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio, - "lora_unet_", device, save_dtype, linear_max_rank, conv_max_rank, + device, save_dtype, 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 ) @@ -836,9 +843,9 @@ 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("linear_max_rank", default=128, min=1, max=3072, + 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=3072, + io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum rank for conv layers"), *_get_common_inputs(), ], @@ -856,7 +863,7 @@ class LoRAExtractQuantile(io.ComfyNode): output_sd = extract_lora_from_files( model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile, - "lora_unet_", device, save_dtype, linear_max_rank, conv_max_rank, + device, save_dtype, 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 ) @@ -878,9 +885,9 @@ class LoRAExtractKnee(io.ComfyNode): *_get_model_inputs(), io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Knee detection method"), - io.Int.Input("linear_max_rank", default=128, min=1, max=3072, + 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=3072, + io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum rank for conv layers"), *_get_common_inputs(), ], @@ -898,7 +905,7 @@ class LoRAExtractKnee(io.ComfyNode): output_sd = extract_lora_from_files( model_a_path, model_b_path, knee_method, 0, 0, - "lora_unet_", device, save_dtype, linear_max_rank, conv_max_rank, + device, save_dtype, 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 ) @@ -922,9 +929,9 @@ 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("linear_max_rank", default=128, min=1, max=3072, + 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=3072, + io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum rank for conv layers"), *_get_common_inputs(), ], @@ -942,7 +949,7 @@ class LoRAExtractFrobenius(io.ComfyNode): output_sd = extract_lora_from_files( model_a_path, model_b_path, "sv_fro", linear_target, conv_target, - "lora_unet_", device, save_dtype, linear_max_rank, conv_max_rank, + device, save_dtype, 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 ) diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index c439012..997b1be 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -13,8 +13,8 @@ from typing import List, Dict, Tuple, Optional from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation from .lora_resize import ( - detect_lora_format, - extract_lora_pairs, + detect_lora_format, + extract_lora_pairs, detect_lora_rank, _compile_patterns, _matches_any_pattern, @@ -32,7 +32,7 @@ def merge_multi_loras( ) -> str: """ Merge multiple LoRAs into a single LoRA file. - + Modes: - concatenate: Mathematically sound merge by stacking ranks. New rank = sum(input ranks). Alpha set to new rank. @@ -45,22 +45,22 @@ def merge_multi_loras( if verbose: print(f"[LoRA Multi-Merge] Preparing memory for {total_size_gb:.2f}GB operation...") print(f"[LoRA Multi-Merge] Mode: {merge_mode}") - + prepare_for_large_operation(total_size_gb * 2.5, torch.device(device)) handlers = [MemoryEfficientSafeOpen(p) for p in lora_paths] - + try: # 1. Analyze all LoRAs layer_map = {} lora_infos = [] - + for i, handler in enumerate(handlers): keys = handler.keys() fmt = detect_lora_format(keys) pairs = extract_lora_pairs(keys, fmt) rank, alpha = detect_lora_rank(handler, pairs) - + info = { "name": os.path.basename(lora_paths[i]), "format": fmt, @@ -71,10 +71,10 @@ def merge_multi_loras( "scale": alpha / rank if rank > 0 else 1.0 } lora_infos.append(info) - + if verbose: print(f"[LoRA Multi-Merge] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}, format={fmt['format']}") - + for block_name in pairs.keys(): core = _format_lora_key(block_name) if core not in layer_map: @@ -94,37 +94,37 @@ def merge_multi_loras( ups = [] is_conv = False max_rank = 0 - + # First pass: load and determine max rank for idx in indices: info = lora_infos[idx] handler = handlers[idx] orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core) block_keys = info["pairs"][orig_block] - + if "down" not in block_keys or "up" not in block_keys: continue - + t_down = handler.get_tensor(block_keys["down"]) t_up = handler.get_tensor(block_keys["up"]) - + if device == 'cuda': t_down = transfer_to_gpu_pinned(t_down, device, torch.float32) t_up = transfer_to_gpu_pinned(t_up, device, torch.float32) else: t_down = t_down.to(device=device, dtype=torch.float32) t_up = t_up.to(device=device, dtype=torch.float32) - + # Store with original info for padding/weighting is_conv = len(t_down.shape) == 4 current_rank = t_down.shape[0] max_rank = max(max_rank, current_rank) - + # Apply scale immediately for concatenate mode, or keep for later if merge_mode == "concatenate": effective_weight = info["weight"] * info["scale"] t_up = t_up * effective_weight - + downs.append((t_down, info)) ups.append((t_up, info)) @@ -148,10 +148,10 @@ def merge_multi_loras( target_down_shape[0] = max_rank target_up_shape = list(ups[0][0].shape) target_up_shape[1] = max_rank - + merged_down = torch.zeros(target_down_shape, device=device, dtype=torch.float32) merged_up = torch.zeros(target_up_shape, device=device, dtype=torch.float32) - + for (t_d, info_d), (t_u, info_u) in zip(downs, ups): r = t_d.shape[0] w = info_d["weight"] @@ -159,7 +159,7 @@ def merge_multi_loras( # Kohya just uses weights. But to be safe with different alphas, # we apply scale to the final delta. # For direct weight merge, we'll just use weights. - + if r < max_rank: # Pad down: [r, in...] -> [max_rank, in...] pad_d = [0] * (len(t_d.shape) * 2) @@ -168,15 +168,15 @@ def merge_multi_loras( # For [r, in], padding is (0,0, 0, max_rank-r) padding_d = [0, 0] * (len(t_d.shape) - 1) + [0, max_rank - r] t_d = torch.nn.functional.pad(t_d, tuple(padding_d)) - + # Pad up: [out, r...] -> [out, max_rank...] # For [out, r], padding is (0, max_rank-r, 0, 0) padding_u = [0, max_rank - r] + [0, 0] * (len(t_u.shape) - 1) t_u = torch.nn.functional.pad(t_u, tuple(padding_u)) - + merged_down += w * t_d merged_up += w * t_u - + new_rank = max_rank # Alpha is usually max of alphas or same as rank new_alpha = float(max_rank) @@ -185,12 +185,12 @@ def merge_multi_loras( output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous() output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous() output_sd[f"{core}.alpha"] = torch.tensor(new_alpha, dtype=save_dtype) - + # Cleanup del merged_down, merged_up for d, _ in downs: del d for u, _ in ups: del u - + pbar.update(1) # 3. Final Summary @@ -199,7 +199,7 @@ def merge_multi_loras( for indices in layer_map.values(): for idx in indices: matched_stats[idx] += 1 - + print(f"[LoRA Multi-Merge] --- Merge Summary ---") for i, info in enumerate(lora_infos): print(f"[LoRA Multi-Merge] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge") @@ -210,16 +210,16 @@ def merge_multi_loras( "ss_training_comment": f"Merged {len(lora_paths)} LoRAs via concatenation", "ss_network_module": "networks.lora", } - + # 4. Save output_dir = folder_paths.get_folder_paths("loras")[0] os.makedirs(output_dir, exist_ok=True) output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors") - + save_file(output_sd, output_path, metadata) if verbose: print(f"[LoRA Multi-Merge] Saved merged LoRA to {output_path}") - + return output_path finally: @@ -265,7 +265,7 @@ class LoRAMultiMerge(io.ComfyNode): # LoRA 8 io.Combo.Input("lora_8", options=["None"] + folder_paths.get_filename_list("loras"), default="None"), io.Float.Input("weight_8", default=1.0, min=-10.0, max=10.0, step=0.01), - + io.Combo.Input("merge_mode", options=["concatenate", "weighted_sum"], default="concatenate", tooltip="concatenate: safe, increases rank. weighted_sum: fixed rank (to max input rank), mathematically lossy but standard in Kohya."), io.String.Input("output_filename", default="merged_lora"), @@ -281,25 +281,25 @@ class LoRAMultiMerge(io.ComfyNode): lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4, lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8, merge_mode, output_filename, save_dtype, device) -> io.NodeOutput: - + count = int(lora_count) names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8] weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8] - + valid_paths = [] valid_weights = [] - + for i in range(count): if names[i] and names[i] != "None": path = folder_paths.get_full_path_or_raise("loras", names[i]) valid_paths.append(path) valid_weights.append(weights[i]) - + if not valid_paths: raise ValueError("No LoRAs selected for merging.") - + dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype] - + path = merge_multi_loras( lora_paths=valid_paths, lora_weights=valid_weights, @@ -308,7 +308,7 @@ class LoRAMultiMerge(io.ComfyNode): save_dtype=dtype, output_filename=output_filename ) - + return io.NodeOutput(path) @@ -359,25 +359,25 @@ class LoRAMultiMergeDARE(io.ComfyNode): lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4, lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8, drop_rate, trim_quantile, seed, output_filename, save_dtype, device) -> io.NodeOutput: - + count = int(lora_count) names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8] weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8] - + valid_paths = [] valid_weights = [] - + for i in range(count): if names[i] and names[i] != "None": path = folder_paths.get_full_path_or_raise("loras", names[i]) valid_paths.append(path) valid_weights.append(weights[i]) - + if not valid_paths: raise ValueError("No LoRAs selected for merging.") - + dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype] - + path = merge_multi_loras_dare( lora_paths=valid_paths, lora_weights=valid_weights, @@ -388,7 +388,7 @@ class LoRAMultiMergeDARE(io.ComfyNode): save_dtype=dtype, output_filename=output_filename ) - + return io.NodeOutput(path) def merge_multi_loras_dare( @@ -413,7 +413,7 @@ def merge_multi_loras_dare( prepare_for_large_operation(total_size_gb * 2.5, torch.device(device)) handlers = [MemoryEfficientSafeOpen(p) for p in lora_paths] - + rng = torch.Generator(device=device).manual_seed(seed) try: @@ -433,7 +433,7 @@ def merge_multi_loras_dare( "weight": lora_weights[i] } lora_infos.append(info) - + if verbose: print(f"[LoRA Multi-Merge DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}") @@ -454,16 +454,16 @@ def merge_multi_loras_dare( downs = [] ups = [] max_rank = 0 - + for idx in indices: info = lora_infos[idx] orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core) block_keys = info["pairs"][orig_block] if "down" not in block_keys or "up" not in block_keys: continue - + t_d = handlers[idx].get_tensor(block_keys["down"]).to(device=device, dtype=torch.float32) t_u = handlers[idx].get_tensor(block_keys["up"]).to(device=device, dtype=torch.float32) - + max_rank = max(max_rank, t_d.shape[0]) downs.append((t_d, info["weight"])) ups.append((t_u, info["weight"])) @@ -484,13 +484,13 @@ def merge_multi_loras_dare( padding[rev_dim*2 + 1] = max_rank - t.shape[dim_to_pad] t = torch.nn.functional.pad(t, tuple(padding)) padded.append(t * w) - + # DARE if drop_rate > 0: for i in range(len(padded)): mask = (torch.rand(padded[i].shape, generator=rng, device=device) > drop_rate).float() padded[i] = (padded[i] * mask) / (1 - drop_rate) - + # TIES # 1. Trim if trim_quantile > 0: @@ -500,17 +500,17 @@ def merge_multi_loras_dare( if k > 0: threshold = torch.kthvalue(flat, k).values padded[i] = torch.where(padded[i].abs() < threshold, torch.zeros_like(padded[i]), padded[i]) - + # 2. Elect & Merge stacked = torch.stack(padded) # [N, ...] signs = torch.sign(stacked) sum_signs = signs.sum(dim=0) dominant_sign = torch.sign(sum_signs) - + # Filter those matching dominant sign mask = (signs == dominant_sign) & (dominant_sign != 0) filtered = torch.where(mask, stacked, torch.zeros_like(stacked)) - + # Average matching signs count = mask.sum(dim=0) result = filtered.sum(dim=0) / torch.clamp(count, min=1.0) @@ -522,7 +522,7 @@ def merge_multi_loras_dare( output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous() output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous() output_sd[f"{core}.alpha"] = torch.tensor(float(max_rank), dtype=save_dtype) - + pbar.update(1) # Final Summary @@ -531,7 +531,7 @@ def merge_multi_loras_dare( for indices in layer_map.values(): for idx in indices: matched_stats[idx] += 1 - + print(f"[LoRA Multi-Merge DARE] --- Merge Summary ---") for i, info in enumerate(lora_infos): print(f"[LoRA Multi-Merge DARE] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge") diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index 08c7157..eef65b0 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -886,11 +886,11 @@ class LoRAMergeToModel(io.ComfyNode): node_id="LoRAMergeToModel", display_name="LoRA Merge To Model", category="ModelUtils/LoRA/Merge", - description="Merge 1-4 LoRAs into a base model and save the result. Saves to base model directory.", + description="Merge 1-8 LoRAs into a base model and save the result. Saves to base model directory.", inputs=[ io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"), tooltip="Base model the LoRAs were trained on"), - io.Combo.Input("lora_count", options=["1", "2", "3", "4"], default="2", + io.Combo.Input("lora_count", options=["1", "2", "3", "4", "5", "6", "7", "8"], default="2", tooltip="Number of LoRAs to merge"), # LoRA 1 io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"), @@ -912,6 +912,26 @@ class LoRAMergeToModel(io.ComfyNode): default="None", tooltip="Fourth LoRA"), io.Float.Input("weight_4", default=1.0, min=-10.0, max=10.0, step=0.01, tooltip="Weight strength for LoRA 4"), + # LoRA 5 + io.Combo.Input("lora_5", options=["None"] + folder_paths.get_filename_list("loras"), + default="None", tooltip="Fifth LoRA"), + io.Float.Input("weight_5", default=1.0, min=-10.0, max=10.0, step=0.01, + tooltip="Weight strength for LoRA 5"), + # LoRA 6 + io.Combo.Input("lora_6", options=["None"] + folder_paths.get_filename_list("loras"), + default="None", tooltip="Sixth LoRA"), + io.Float.Input("weight_6", default=1.0, min=-10.0, max=10.0, step=0.01, + tooltip="Weight strength for LoRA 6"), + # LoRA 7 + io.Combo.Input("lora_7", options=["None"] + folder_paths.get_filename_list("loras"), + default="None", tooltip="Seventh LoRA"), + io.Float.Input("weight_7", default=1.0, min=-10.0, max=10.0, step=0.01, + tooltip="Weight strength for LoRA 7"), + # LoRA 8 + io.Combo.Input("lora_8", options=["None"] + folder_paths.get_filename_list("loras"), + default="None", tooltip="Eighth LoRA"), + io.Float.Input("weight_8", default=1.0, min=-10.0, max=10.0, step=0.01, + tooltip="Weight strength for LoRA 8"), # Settings io.String.Input("skip_patterns", default="", multiline=True, tooltip="Regex patterns for layers to skip"), @@ -926,12 +946,13 @@ class LoRAMergeToModel(io.ComfyNode): @classmethod def execute(cls, base_model, lora_count, lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4, + lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8, skip_patterns, output_filename, save_dtype, device) -> io.NodeOutput: # Build LoRA list based on count count = int(lora_count) - lora_names = [lora_1, lora_2, lora_3, lora_4][:count] - lora_weights = [weight_1, weight_2, weight_3, weight_4][:count] + lora_names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8][:count] + lora_weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8][:count] # Filter out "None" entries valid_loras = [(name, weight) for name, weight in zip(lora_names, lora_weights) if name != "None"] diff --git a/nodes/merger.py b/nodes/merger.py index c2993de..2efa2ad 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -183,7 +183,7 @@ class MergerLogic: # Clean up tensor_a reference to allow GC del recipe_params['_tensor_a'] - + if recipe_params.get('force_clear_cache', True): import gc gc.collect()