From bb915f0599d9ac9c32fdbc5470c6714fbd7b95ef Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sat, 16 May 2026 11:07:20 +0200 Subject: [PATCH 1/2] refactor all saves to use IncrementalSafetensorsWriter --- nodes/lora_extract_svd.py | 89 +++--- nodes/lora_merger.py | 616 ++++++++++++++++++++------------------ nodes/lora_resize.py | 326 ++++++++++---------- nodes/merger.py | 1 - nodes/prunekeys.py | 44 +-- nodes/renamekeys.py | 28 +- requirements.txt | 2 +- 7 files changed, 564 insertions(+), 542 deletions(-) diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index 2dfd499..570b02b 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -13,13 +13,12 @@ import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -from safetensors.torch import save_file 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 +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter @@ -447,6 +446,7 @@ def extract_lora_from_files( 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, @@ -458,9 +458,9 @@ def extract_lora_from_files( lazy_load: bool = True, force_clear_cache: bool = True, glob_skip_patterns: bool = False, -) -> dict[str, torch.Tensor]: +) -> None: """ - Extract LoRA from difference between two models. + Extract LoRA from difference between two models, writing incrementally to disk. Args: model_a_path: Finetuned model path @@ -470,6 +470,7 @@ def extract_lora_from_files( 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 @@ -480,8 +481,6 @@ def extract_lora_from_files( glob_skip_patterns: When True, treat skip_patterns as glob (* wildcard). When False (default), treat as Python regex. """ - output_sd = {} - save_torch_dtype = { "fp32": torch.float32, "fp16": torch.float16, @@ -611,19 +610,24 @@ def extract_lora_from_files( 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) + writer = IncrementalSafetensorsWriter(output_path) + writer.__enter__() + try: + for key in tqdm(weight_keys, desc="Extracting LoRA", 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() + if force_clear_cache: + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() - pbar.update(1) + pbar.update(1) + finally: + writer.__exit__(None, None, None) print(f"[LoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " f"{stats['full']} full, {stats['skipped']} skipped") @@ -633,22 +637,12 @@ def extract_lora_from_files( handler_b.__exit__(None, None, None) cleanup_after_operation() - return output_sd - - -def _save_lora(output_sd: dict, output_filename: str) -> str: - """Save LoRA to file.""" - if not output_sd: - raise ValueError("No LoRA weights extracted") +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) - output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors") - - save_file(output_sd, output_path) - print(f"[LoRA Extract] Saved to {output_path}") - - return output_path + return os.path.join(output_dir, f"{output_filename.strip()}.safetensors") # ============================================================================= @@ -801,15 +795,16 @@ class LoRAExtractFixed(io.ComfyNode): 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) - output_sd = extract_lora_from_files( + extract_lora_from_files( model_a_path, model_b_path, "fixed", linear_dim, conv_dim, - device, save_dtype, 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(_save_lora(output_sd, output_filename)) + return io.NodeOutput(output_path) class LoRAExtractRatio(io.ComfyNode): @@ -845,16 +840,17 @@ class LoRAExtractRatio(io.ComfyNode): 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) - output_sd = extract_lora_from_files( + extract_lora_from_files( model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio, - device, save_dtype, linear_max_rank, conv_max_rank, + 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(_save_lora(output_sd, output_filename)) + return io.NodeOutput(output_path) class LoRAExtractQuantile(io.ComfyNode): @@ -890,16 +886,17 @@ class LoRAExtractQuantile(io.ComfyNode): 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) - output_sd = extract_lora_from_files( + extract_lora_from_files( model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile, - device, save_dtype, linear_max_rank, conv_max_rank, + 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(_save_lora(output_sd, output_filename)) + return io.NodeOutput(output_path) class LoRAExtractKnee(io.ComfyNode): @@ -933,16 +930,17 @@ class LoRAExtractKnee(io.ComfyNode): 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) - output_sd = extract_lora_from_files( + extract_lora_from_files( model_a_path, model_b_path, knee_method, 0, 0, - device, save_dtype, linear_max_rank, conv_max_rank, + 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(_save_lora(output_sd, output_filename)) + return io.NodeOutput(output_path) class LoRAExtractFrobenius(io.ComfyNode): @@ -978,13 +976,14 @@ class LoRAExtractFrobenius(io.ComfyNode): 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) - output_sd = extract_lora_from_files( + extract_lora_from_files( model_a_path, model_b_path, "sv_fro", linear_target, conv_target, - device, save_dtype, linear_max_rank, conv_max_rank, + 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(_save_lora(output_sd, output_filename)) + return io.NodeOutput(output_path) diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index cde4fe1..62cc220 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -8,7 +8,6 @@ import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -from safetensors.torch import save_file from typing import List, Dict, Tuple, Optional from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation from .lora_resize import ( @@ -20,7 +19,7 @@ from .lora_resize import ( _format_lora_key ) -from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter @@ -87,120 +86,137 @@ def merge_multi_loras( if verbose: print(f"[LoRA Multi-Merge] Total unique layers after resolving naming: {len(layer_map)}") - output_sd = {} + # Build metadata and output path before loop + metadata = { + "ss_training_comment": f"Merged {len(lora_paths)} LoRAs via concatenation", + "ss_network_module": "networks.lora", + } + 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") + pbar = comfy.utils.ProgressBar(len(layer_map)) + tensor_count = 0 - # 2. Merge each core layer - with torch.no_grad(): - for core, indices in tqdm(layer_map.items(), desc="Merging layers", unit="layers"): - downs = [] - ups = [] - is_conv = False - max_rank = 0 + writer = IncrementalSafetensorsWriter(output_path, metadata=metadata) + writer.__enter__() + try: + # 2. Merge each core layer + with torch.no_grad(): + for core, indices in tqdm(layer_map.items(), desc="Merging layers", unit="layers"): + downs = [] + 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] + # 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: + 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)) + + if not downs: + pbar.update(1) 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 + # Stack ranks + merged_down = torch.cat([d[0] for d in downs], dim=0) + merged_up = torch.cat([u[0] for u in ups], dim=1) + new_rank = merged_down.shape[0] + new_alpha = float(new_rank) + else: + # weighted_sum (Kohya style) + # Result = sum( weight_i * Padded(B_i) ), sum( weight_i * Padded(A_i) ) + # Note: We apply scale to weights here to normalize different alphas + merged_down = torch.zeros_like(downs[0][0]) + # Need to handle different ranks via padding + target_down_shape = list(downs[0][0].shape) + target_down_shape[0] = max_rank + target_up_shape = list(ups[0][0].shape) + target_up_shape[1] = max_rank - downs.append((t_down, info)) - ups.append((t_up, info)) + 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"] + # We also apply sqrt(scale) to both to distribute the scale factor? + # 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) + pad_d[-1] = max_rank - r # last dim in pad is first dim in tensor (reversed) + # Wait, F.pad uses reverse order of dims. + # 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) + + # Write layer immediately + writer.write_dict({ + f"{core}.lora_down.weight": merged_down.to(save_dtype).cpu().contiguous(), + f"{core}.lora_up.weight": merged_up.to(save_dtype).cpu().contiguous(), + f"{core}.alpha": torch.tensor(new_alpha, dtype=save_dtype), + }) + tensor_count += 3 + + # Cleanup + del merged_down, merged_up + for d, _ in downs: del d + for u, _ in ups: del u + downs.clear() + ups.clear() + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() - if not downs: pbar.update(1) - continue - - if merge_mode == "concatenate": - # Stack ranks - merged_down = torch.cat([d[0] for d in downs], dim=0) - merged_up = torch.cat([u[0] for u in ups], dim=1) - new_rank = merged_down.shape[0] - new_alpha = float(new_rank) - else: - # weighted_sum (Kohya style) - # Result = sum( weight_i * Padded(B_i) ), sum( weight_i * Padded(A_i) ) - # Note: We apply scale to weights here to normalize different alphas - merged_down = torch.zeros_like(downs[0][0]) - # Need to handle different ranks via padding - target_down_shape = list(downs[0][0].shape) - 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"] - # We also apply sqrt(scale) to both to distribute the scale factor? - # 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) - pad_d[-1] = max_rank - r # last dim in pad is first dim in tensor (reversed) - # Wait, F.pad uses reverse order of dims. - # 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) - - # Store - 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 - downs.clear() - ups.clear() - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - pbar.update(1) + finally: + writer.__exit__(None, None, None) # 3. Final Summary if verbose: @@ -212,21 +228,7 @@ def merge_multi_loras( 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") - print(f"[LoRA Multi-Merge] Output state dict has {len(output_sd)} tensors") - - # 4. Metadata - metadata = { - "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] Output state dict has {tensor_count} tensors") print(f"[LoRA Multi-Merge] Saved merged LoRA to {output_path}") return output_path @@ -454,95 +456,108 @@ def merge_multi_loras_dare( if verbose: print(f"[LoRA Multi-Merge DARE] Total unique layers after resolving naming: {len(layer_map)}") - output_sd = {} + # Build output path before loop + 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") + pbar = comfy.utils.ProgressBar(len(layer_map)) + tensor_count = 0 - # 2. Merge - with torch.no_grad(): - for core, indices in tqdm(layer_map.items(), desc="Merging layers (DARE)", unit="layers"): - downs = [] - ups = [] - max_rank = 0 + writer = IncrementalSafetensorsWriter(output_path, metadata={"ss_training_comment": "Merged via DARE-Ties"}) + writer.__enter__() + try: + # 2. Merge + with torch.no_grad(): + for core, indices in tqdm(layer_map.items(), desc="Merging layers (DARE)", unit="layers"): + 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 + 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) + 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"])) + max_rank = max(max_rank, t_d.shape[0]) + downs.append((t_d, info["weight"])) + ups.append((t_u, info["weight"])) + + if not downs: + pbar.update(1) + continue + + def process_ties_dare(tensors, weights, dim_to_pad): + # Pad all to max_rank + padded = [] + for t, w in tensors: + if t.shape[dim_to_pad] < max_rank: + padding = [0] * (len(t.shape) * 2) + # dim_to_pad 0 (down) -> last pair in pad + # dim_to_pad 1 (up) -> second to last pair in pad + rev_dim = len(t.shape) - 1 - dim_to_pad + 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: + for i in range(len(padded)): + flat = padded[i].abs().flatten() + k = int(len(flat) * trim_quantile) + 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) + return result + + merged_down = process_ties_dare(downs, [1.0]*len(downs), 0) + merged_up = process_ties_dare(ups, [1.0]*len(ups), 1) + + writer.write_dict({ + f"{core}.lora_down.weight": merged_down.to(save_dtype).cpu().contiguous(), + f"{core}.lora_up.weight": merged_up.to(save_dtype).cpu().contiguous(), + f"{core}.alpha": torch.tensor(float(max_rank), dtype=save_dtype), + }) + tensor_count += 3 + + del merged_down, merged_up + for d, _ in downs: del d + for u, _ in ups: del u + downs.clear() + ups.clear() + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() - if not downs: pbar.update(1) - continue - - def process_ties_dare(tensors, weights, dim_to_pad): - # Pad all to max_rank - padded = [] - for t, w in tensors: - if t.shape[dim_to_pad] < max_rank: - padding = [0] * (len(t.shape) * 2) - # dim_to_pad 0 (down) -> last pair in pad - # dim_to_pad 1 (up) -> second to last pair in pad - rev_dim = len(t.shape) - 1 - dim_to_pad - 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: - for i in range(len(padded)): - flat = padded[i].abs().flatten() - k = int(len(flat) * trim_quantile) - 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) - return result - - merged_down = process_ties_dare(downs, [1.0]*len(downs), 0) - merged_up = process_ties_dare(ups, [1.0]*len(ups), 1) - - 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) - - del merged_down, merged_up - for d, _ in downs: del d - for u, _ in ups: del u - downs.clear() - ups.clear() - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - pbar.update(1) + finally: + writer.__exit__(None, None, None) # Final Summary if verbose: @@ -554,13 +569,8 @@ def merge_multi_loras_dare( 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") - print(f"[LoRA Multi-Merge DARE] Output state dict has {len(output_sd)} tensors") + print(f"[LoRA Multi-Merge DARE] Output state dict has {tensor_count} tensors") - # 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, {"ss_training_comment": "Merged via DARE-Ties"}) return output_path finally: @@ -707,97 +717,110 @@ def merge_multi_loras_dare_enhanced( if verbose: print(f"[LoRA Multi-Merge Enhanced DARE] Total unique layers after resolving naming: {len(layer_map)}") - output_sd = {} + # Build output path before loop + 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") + pbar = comfy.utils.ProgressBar(len(layer_map)) + tensor_count = 0 - # 2. Merge - with torch.no_grad(): - for core, indices in tqdm(layer_map.items(), desc="Merging layers (Enhanced DARE)", unit="layers"): - downs = [] - ups = [] - max_rank = 0 + writer = IncrementalSafetensorsWriter(output_path, metadata={"ss_training_comment": "Merged via Enhanced DARE-Ties"}) + writer.__enter__() + try: + # 2. Merge + with torch.no_grad(): + for core, indices in tqdm(layer_map.items(), desc="Merging layers (Enhanced DARE)", unit="layers"): + 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 + 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) + 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"])) + max_rank = max(max_rank, t_d.shape[0]) + downs.append((t_d, info["weight"])) + ups.append((t_u, info["weight"])) + + if not downs: + pbar.update(1) + continue + + def process_ties_dare_enhanced(tensors, weights, dim_to_pad): + processed = [] + for t, w in tensors: + t_val = t * w + + # Enhanced DARE + abs_t = torch.abs(t_val) + max_val = torch.max(abs_t) + if max_val > 0: + prob = torch.clamp((abs_t / max_val) ** max(mask_power, 0.001), min=min_keep_prob, max=1.0) + prob = torch.nan_to_num(prob) + + random_mask = torch.bernoulli(prob, generator=rng) + interpolated_mask = torch.lerp(random_mask, prob, mask_smooth) + + t_val = t_val * interpolated_mask + + # TIES + if trim_quantile > 0: + flat = t_val.abs().flatten() + k = max(1, int(len(flat) * trim_quantile)) + if k > 0: + threshold = torch.kthvalue(flat, k).values + t_val = torch.where(t_val.abs() < threshold, torch.zeros_like(t_val), t_val) + + # Pad + if t_val.shape[dim_to_pad] < max_rank: + padding = [0] * (len(t_val.shape) * 2) + rev_dim = len(t_val.shape) - 1 - dim_to_pad + padding[rev_dim*2 + 1] = max_rank - t_val.shape[dim_to_pad] + t_val = torch.nn.functional.pad(t_val, tuple(padding)) + + processed.append(t_val) + + stacked = torch.stack(processed) + signs = torch.sign(stacked) + sum_signs = signs.sum(dim=0) + dominant_sign = torch.sign(sum_signs) + + mask = (signs == dominant_sign) & (dominant_sign != 0) + filtered = torch.where(mask, stacked, torch.zeros_like(stacked)) + + count = mask.sum(dim=0) + result = filtered.sum(dim=0) / torch.clamp(count, min=1.0) + return result + + merged_down = process_ties_dare_enhanced(downs, [1.0]*len(downs), 0) + merged_up = process_ties_dare_enhanced(ups, [1.0]*len(ups), 1) + + writer.write_dict({ + f"{core}.lora_down.weight": merged_down.to(save_dtype).cpu().contiguous(), + f"{core}.lora_up.weight": merged_up.to(save_dtype).cpu().contiguous(), + f"{core}.alpha": torch.tensor(float(max_rank), dtype=save_dtype), + }) + tensor_count += 3 + + del merged_down, merged_up + for d, _ in downs: del d + for u, _ in ups: del u + downs.clear() + ups.clear() + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() - if not downs: pbar.update(1) - continue - - def process_ties_dare_enhanced(tensors, weights, dim_to_pad): - processed = [] - for t, w in tensors: - t_val = t * w - - # Enhanced DARE - abs_t = torch.abs(t_val) - max_val = torch.max(abs_t) - if max_val > 0: - prob = torch.clamp((abs_t / max_val) ** max(mask_power, 0.001), min=min_keep_prob, max=1.0) - prob = torch.nan_to_num(prob) - - random_mask = torch.bernoulli(prob, generator=rng) - interpolated_mask = torch.lerp(random_mask, prob, mask_smooth) - - t_val = t_val * interpolated_mask - - # TIES - if trim_quantile > 0: - flat = t_val.abs().flatten() - k = max(1, int(len(flat) * trim_quantile)) - if k > 0: - threshold = torch.kthvalue(flat, k).values - t_val = torch.where(t_val.abs() < threshold, torch.zeros_like(t_val), t_val) - - # Pad - if t_val.shape[dim_to_pad] < max_rank: - padding = [0] * (len(t_val.shape) * 2) - rev_dim = len(t_val.shape) - 1 - dim_to_pad - padding[rev_dim*2 + 1] = max_rank - t_val.shape[dim_to_pad] - t_val = torch.nn.functional.pad(t_val, tuple(padding)) - - processed.append(t_val) - - stacked = torch.stack(processed) - signs = torch.sign(stacked) - sum_signs = signs.sum(dim=0) - dominant_sign = torch.sign(sum_signs) - - mask = (signs == dominant_sign) & (dominant_sign != 0) - filtered = torch.where(mask, stacked, torch.zeros_like(stacked)) - - count = mask.sum(dim=0) - result = filtered.sum(dim=0) / torch.clamp(count, min=1.0) - return result - - merged_down = process_ties_dare_enhanced(downs, [1.0]*len(downs), 0) - merged_up = process_ties_dare_enhanced(ups, [1.0]*len(ups), 1) - - 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) - - del merged_down, merged_up - for d, _ in downs: del d - for u, _ in ups: del u - downs.clear() - ups.clear() - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - pbar.update(1) + finally: + writer.__exit__(None, None, None) # Final Summary if verbose: @@ -809,13 +832,8 @@ def merge_multi_loras_dare_enhanced( print(f"[LoRA Multi-Merge Enhanced DARE] --- Merge Summary ---") for i, info in enumerate(lora_infos): print(f"[LoRA Multi-Merge Enhanced DARE] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge") - print(f"[LoRA Multi-Merge Enhanced DARE] Output state dict has {len(output_sd)} tensors") + print(f"[LoRA Multi-Merge Enhanced DARE] Output state dict has {tensor_count} tensors") - # 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, {"ss_training_comment": "Merged via Enhanced DARE-Ties"}) return output_path finally: diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index 9e6879c..1c56dc9 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -12,11 +12,10 @@ import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -from safetensors.torch import save_file -from typing import Optional, Dict, Tuple, List from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation -from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from typing import Optional, Dict, Tuple, List @@ -419,66 +418,7 @@ def resize_lora_file( print(f"[LoRA Resize] Original dim={network_dim}, alpha={network_alpha:.1f}, scale={scale:.3f}") print(f"[LoRA Resize] Resizing with method={method_str}, max_rank={new_rank}") - output_sd = {} - fro_list = [] - pbar = comfy.utils.ProgressBar(len(pairs)) - - with torch.no_grad(): - for block_name, block_keys in tqdm(pairs.items(), desc="Resizing layers", unit="layers"): - if "down" not in block_keys or "up" not in block_keys: - pbar.update(1) - continue - - lora_down = handler.get_tensor(block_keys["down"]) - lora_up = handler.get_tensor(block_keys["up"]) - - is_conv = len(lora_down.shape) == 4 - - # Transfer to GPU with pinned memory if available - if device == 'cuda': - lora_down = transfer_to_gpu_pinned(lora_down, device, torch.float32) - lora_up = transfer_to_gpu_pinned(lora_up, device, torch.float32) - - # Merge and re-extract - if is_conv: - weight = _merge_conv(lora_down, lora_up, device) - result = _extract_conv(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter) - else: - weight = _merge_linear(lora_down, lora_up, device) - result = _extract_linear(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter) - - del weight, lora_down, lora_up - - fro_list.append(result['fro_retained']) - - # Store using same format suffixes as input - down_suffix = format_info["down_suffix"] - up_suffix = format_info["up_suffix"] - alpha_suffix = format_info["alpha_suffix"] - - # Standardize prefix using heuristics if possible - new_block_name = _format_lora_key(block_name) - - output_sd[f"{new_block_name}{down_suffix}"] = result["lora_down"].to(save_dtype) - output_sd[f"{new_block_name}{up_suffix}"] = result["lora_up"].to(save_dtype) - if alpha_suffix: - output_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype) - - del result - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - pbar.update(1) - - if verbose and fro_list: - import numpy as np - avg_fro = np.mean(fro_list) - std_fro = np.std(fro_list) - print(f"[LoRA Resize] Average Frobenius retention: {avg_fro:.1%} ± {std_fro:.3f}") - - # Update metadata + # Build metadata and output path before loop (all values known at this point) if dynamic_method: metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from dim {network_dim}" metadata["ss_network_dim"] = "Dynamic" @@ -488,14 +428,78 @@ def resize_lora_file( metadata["ss_network_dim"] = str(new_rank) metadata["ss_network_alpha"] = str(scale * new_rank) - # 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) - print(f"[LoRA Resize] Saved to {output_path}") + fro_list = [] + pbar = comfy.utils.ProgressBar(len(pairs)) + writer = IncrementalSafetensorsWriter(output_path, metadata=metadata) + writer.__enter__() + try: + with torch.no_grad(): + for block_name, block_keys in tqdm(pairs.items(), desc="Resizing layers", unit="layers"): + if "down" not in block_keys or "up" not in block_keys: + pbar.update(1) + continue + + lora_down = handler.get_tensor(block_keys["down"]) + lora_up = handler.get_tensor(block_keys["up"]) + + is_conv = len(lora_down.shape) == 4 + + # Transfer to GPU with pinned memory if available + if device == 'cuda': + lora_down = transfer_to_gpu_pinned(lora_down, device, torch.float32) + lora_up = transfer_to_gpu_pinned(lora_up, device, torch.float32) + + # Merge and re-extract + if is_conv: + weight = _merge_conv(lora_down, lora_up, device) + result = _extract_conv(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter) + else: + weight = _merge_linear(lora_down, lora_up, device) + result = _extract_linear(weight, new_rank, dynamic_method, dynamic_param, scale, svd_niter) + + del weight, lora_down, lora_up + + fro_list.append(result['fro_retained']) + + # Store using same format suffixes as input + down_suffix = format_info["down_suffix"] + up_suffix = format_info["up_suffix"] + alpha_suffix = format_info["alpha_suffix"] + + # Standardize prefix using heuristics if possible + new_block_name = _format_lora_key(block_name) + + block_sd = { + f"{new_block_name}{down_suffix}": result["lora_down"].to(save_dtype), + f"{new_block_name}{up_suffix}": result["lora_up"].to(save_dtype), + } + if alpha_suffix: + block_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype) + + writer.write_dict(block_sd) + + del result + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + pbar.update(1) + finally: + writer.__exit__(None, None, None) + + if verbose and fro_list: + import numpy as np + avg_fro = np.mean(fro_list) + std_fro = np.std(fro_list) + print(f"[LoRA Resize] Average Frobenius retention: {avg_fro:.1%} ± {std_fro:.3f}") + + print(f"[LoRA Resize] Saved to {output_path}") return output_path finally: @@ -715,7 +719,7 @@ def merge_loras_to_model( print(f"[LoRA Merge To Model] LoRA {i+1}: {format_info['format']}, {len(pairs)} layers, dim={network_dim}") # Common prefixes - BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model."] + BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."] LORA_PREFIXES = [ "lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_", "lycoris_", "diffusion_model.", "transformer.", "unet." @@ -761,7 +765,11 @@ def merge_loras_to_model( base_metadata = base_handler.metadata().copy() if base_handler.metadata() else {} base_metadata["merge_comment"] = f"Merged {len(lora_paths)} LoRAs with weights: {lora_weights}" - output_sd = {} + # Build output path before loop + base_dir = os.path.dirname(base_model_path) + os.makedirs(base_dir, exist_ok=True) + output_path = os.path.join(base_dir, f"{output_filename.strip()}.safetensors") + stats = {"merged": 0, "copied": 0, "skipped": 0} base_keys = list(base_handler.keys()) pbar = comfy.utils.ProgressBar(len(base_keys)) @@ -769,118 +777,116 @@ def merge_loras_to_model( if verbose: print(f"[LoRA Merge To Model] Processing {len(base_keys)} base model keys...") - with torch.no_grad(): - for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"): - # Check skip patterns - if _matches_any_pattern(base_key, skip_patterns): - stats["skipped"] += 1 - pbar.update(1) - continue + writer = IncrementalSafetensorsWriter(output_path, metadata=base_metadata) + writer.__enter__() + try: + with torch.no_grad(): + for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"): + # Check skip patterns + if _matches_any_pattern(base_key, skip_patterns): + stats["skipped"] += 1 + pbar.update(1) + continue - # Load base weight - cpu_base = base_handler.get_tensor(base_key) + # Load base weight + cpu_base = base_handler.get_tensor(base_key) - # Only process weight tensors for LoRA merging - if base_key.endswith(".weight"): - core = extract_core_layer_base(base_key) - core_underscored = core.replace(".", "_") + # Only process weight tensors for LoRA merging + if base_key.endswith(".weight"): + core = extract_core_layer_base(base_key) + core_underscored = core.replace(".", "_") - # Check if any LoRA contributes to this layer - if core_underscored in lora_lookup: - # Transfer to GPU for computation - if device == 'cuda': - base_weight = transfer_to_gpu_pinned(cpu_base, device, torch.float32) - else: - base_weight = cpu_base.to(device=device, dtype=torch.float32) - del cpu_base - - # Accumulate deltas from all contributing LoRAs - for info, block_keys in lora_lookup[core_underscored]: - is_full_diff = info["format_info"].get("is_full_diff", False) - - if is_full_diff: - # Full diff format - if "diff" not in block_keys: - continue - cpu_diff = info["handler"].get_tensor(block_keys["diff"]) - if device == 'cuda': - delta = transfer_to_gpu_pinned(cpu_diff, device, torch.float32) - else: - delta = cpu_diff.to(device=device, dtype=torch.float32) - del cpu_diff - effective_scale = info["weight"] + # Check if any LoRA contributes to this layer + if core_underscored in lora_lookup: + # Transfer to GPU for computation + if device == 'cuda': + base_weight = transfer_to_gpu_pinned(cpu_base, device, torch.float32) else: - # Standard LoRA format - if "down" not in block_keys or "up" not in block_keys: - continue + base_weight = cpu_base.to(device=device, dtype=torch.float32) + del cpu_base - cpu_down = info["handler"].get_tensor(block_keys["down"]) - cpu_up = info["handler"].get_tensor(block_keys["up"]) - if device == 'cuda': - lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32) - lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32) + # Accumulate deltas from all contributing LoRAs + for info, block_keys in lora_lookup[core_underscored]: + is_full_diff = info["format_info"].get("is_full_diff", False) + + if is_full_diff: + # Full diff format + if "diff" not in block_keys: + continue + cpu_diff = info["handler"].get_tensor(block_keys["diff"]) + if device == 'cuda': + delta = transfer_to_gpu_pinned(cpu_diff, device, torch.float32) + else: + delta = cpu_diff.to(device=device, dtype=torch.float32) + del cpu_diff + effective_scale = info["weight"] else: - lora_down = cpu_down.to(device=device, dtype=torch.float32) - lora_up = cpu_up.to(device=device, dtype=torch.float32) - del cpu_down, cpu_up + # Standard LoRA format + if "down" not in block_keys or "up" not in block_keys: + continue - # Get alpha - if "alpha" in block_keys: - alpha_tensor = info["handler"].get_tensor(block_keys["alpha"]) - layer_alpha = float(alpha_tensor.item()) - else: - layer_alpha = float(info["network_dim"]) - layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0 - effective_scale = layer_scale * info["weight"] + cpu_down = info["handler"].get_tensor(block_keys["down"]) + cpu_up = info["handler"].get_tensor(block_keys["up"]) + if device == 'cuda': + lora_down = transfer_to_gpu_pinned(cpu_down, device, torch.float32) + lora_up = transfer_to_gpu_pinned(cpu_up, device, torch.float32) + else: + lora_down = cpu_down.to(device=device, dtype=torch.float32) + lora_up = cpu_up.to(device=device, dtype=torch.float32) + del cpu_down, cpu_up - # Compute delta - is_conv = len(lora_down.shape) == 4 - if is_conv: - in_rank, in_size, kernel_size, k_ = lora_down.shape - out_size, out_rank, _, _ = lora_up.shape - delta = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1) - delta = delta.reshape(out_size, in_size, kernel_size, kernel_size) - else: - delta = lora_up @ lora_down - del lora_down, lora_up + # Get alpha + if "alpha" in block_keys: + alpha_tensor = info["handler"].get_tensor(block_keys["alpha"]) + layer_alpha = float(alpha_tensor.item()) + else: + layer_alpha = float(info["network_dim"]) + layer_scale = layer_alpha / info["network_dim"] if info["network_dim"] > 0 else 1.0 + effective_scale = layer_scale * info["weight"] - # Apply delta to base weight - base_weight = base_weight + effective_scale * delta - del delta + # Compute delta + is_conv = len(lora_down.shape) == 4 + if is_conv: + in_rank, in_size, kernel_size, k_ = lora_down.shape + out_size, out_rank, _, _ = lora_up.shape + delta = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1) + delta = delta.reshape(out_size, in_size, kernel_size, kernel_size) + else: + delta = lora_up @ lora_down + del lora_down, lora_up - # Store merged weight - output_sd[base_key] = base_weight.to(save_dtype).cpu().contiguous() - del base_weight - stats["merged"] += 1 + # Apply delta to base weight + base_weight = base_weight + effective_scale * delta + del delta + + # Write merged weight + writer.write(base_key, base_weight.to(save_dtype).cpu().contiguous()) + del base_weight + stats["merged"] += 1 + else: + # No LoRA contribution, write as-is + writer.write(base_key, cpu_base.to(save_dtype).contiguous()) + del cpu_base + stats["copied"] += 1 else: - # No LoRA contribution, copy as-is - output_sd[base_key] = cpu_base.to(save_dtype).contiguous() + # Non-weight tensor (bias, norm, etc.), write as-is + writer.write(base_key, cpu_base.to(save_dtype).contiguous()) del cpu_base stats["copied"] += 1 - else: - # Non-weight tensor (bias, norm, etc.), copy as-is - output_sd[base_key] = cpu_base.to(save_dtype).contiguous() - del cpu_base - stats["copied"] += 1 - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() - pbar.update(1) + pbar.update(1) + finally: + writer.__exit__(None, None, None) if verbose: print(f"[LoRA Merge To Model] Done: {stats['merged']} merged, {stats['copied']} copied, {stats['skipped']} skipped") - # Save to base model directory - base_dir = os.path.dirname(base_model_path) - os.makedirs(base_dir, exist_ok=True) - output_path = os.path.join(base_dir, f"{output_filename.strip()}.safetensors") - - save_file(output_sd, output_path, base_metadata) print(f"[LoRA Merge To Model] Saved to {output_path}") - return output_path finally: diff --git a/nodes/merger.py b/nodes/merger.py index 038ca72..bcb1d67 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -6,7 +6,6 @@ import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -from safetensors.torch import save_file from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError from .device_utils import ( estimate_model_size, prepare_for_large_operation, cleanup_after_operation diff --git a/nodes/prunekeys.py b/nodes/prunekeys.py index db398be..2201357 100644 --- a/nodes/prunekeys.py +++ b/nodes/prunekeys.py @@ -4,12 +4,11 @@ import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -from safetensors.torch import save_file from safetensors import safe_open from .utils import convert_pt_to_safetensors from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation -from unifiedefficientloader import MemoryEfficientSafeOpen +from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str, @@ -40,32 +39,33 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str, if not patterns: raise ValueError("No keys/patterns provided to prune.") - # Stream tensors and filter on the fly - pruned_tensors = {} - with MemoryEfficientSafeOpen(model_path_to_load) as handler: - all_keys = handler.keys() - pbar = comfy.utils.ProgressBar(len(all_keys)) - for key in tqdm(all_keys, desc="Pruning keys", unit="keys"): - is_match = False - if use_regex: - if any(re.search(pattern, key) for pattern in patterns): - is_match = True - else: - if any(pattern in key for pattern in patterns): - is_match = True - - if not is_match: - pruned_tensors[key] = handler.get_tensor(key) - pbar.update(1) - # Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet model_dir = folder_paths.get_folder_paths(model_type)[-1] output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors") - save_file(pruned_tensors, output_path, metadata) + # Stream tensors, filter on the fly, write immediately + writer = IncrementalSafetensorsWriter(output_path, metadata=metadata) + writer.__enter__() + try: + with MemoryEfficientSafeOpen(model_path_to_load) as handler: + all_keys = handler.keys() + pbar = comfy.utils.ProgressBar(len(all_keys)) + for key in tqdm(all_keys, desc="Pruning keys", unit="keys"): + is_match = False + if use_regex: + if any(re.search(pattern, key) for pattern in patterns): + is_match = True + else: + if any(pattern in key for pattern in patterns): + is_match = True + + if not is_match: + writer.write(key, handler.get_tensor(key).contiguous()) + pbar.update(1) + finally: + writer.__exit__(None, None, None) # Cleanup after operation - del pruned_tensors cleanup_after_operation() return output_path diff --git a/nodes/renamekeys.py b/nodes/renamekeys.py index e299a77..1bb1612 100644 --- a/nodes/renamekeys.py +++ b/nodes/renamekeys.py @@ -3,12 +3,11 @@ import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -from safetensors.torch import save_file from safetensors import safe_open from .utils import convert_pt_to_safetensors from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation -from unifiedefficientloader import MemoryEfficientSafeOpen +from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter @@ -43,24 +42,25 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str, key_map = dict(zip(old_keys, new_keys)) - # Stream tensors and rename on the fly - renamed_tensors = {} - with MemoryEfficientSafeOpen(model_path_to_load) as handler: - original_keys = handler.keys() - pbar = comfy.utils.ProgressBar(len(original_keys)) - for key in tqdm(original_keys, desc="Renaming keys", unit="keys"): - new_key = key_map.get(key, key) - renamed_tensors[new_key] = handler.get_tensor(key) - pbar.update(1) - # Use [-1] for diffusion_models to get the actual diffusion_models folder, not legacy unet model_dir = folder_paths.get_folder_paths(model_type)[-1] output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors") - save_file(renamed_tensors, output_path, metadata) + # Stream tensors, rename on the fly, write immediately + writer = IncrementalSafetensorsWriter(output_path, metadata=metadata) + writer.__enter__() + try: + with MemoryEfficientSafeOpen(model_path_to_load) as handler: + original_keys = handler.keys() + pbar = comfy.utils.ProgressBar(len(original_keys)) + for key in tqdm(original_keys, desc="Renaming keys", unit="keys"): + new_key = key_map.get(key, key) + writer.write(new_key, handler.get_tensor(key).contiguous()) + pbar.update(1) + finally: + writer.__exit__(None, None, None) # Cleanup after operation - del renamed_tensors cleanup_after_operation() return output_path diff --git a/requirements.txt b/requirements.txt index 0b044ca..c38b3f0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ requests Pillow mutagen -unifiedefficientloader==0.4.4 +unifiedefficientloader==0.4.5 av From 94d2a93b9183b91d05886d44779f07de5b9fb752 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sat, 16 May 2026 16:32:44 +0200 Subject: [PATCH 2/2] fix layer renaming issues --- nodes/lora_extract_svd.py | 40 +------ nodes/lora_merger.py | 214 +++++++++++++++++++++++++++++++++----- nodes/lora_resize.py | 4 +- 3 files changed, 192 insertions(+), 66 deletions(-) diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index 570b02b..bfb8373 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -658,42 +658,6 @@ def _get_model_inputs(): ] -def _reconstruct_dots(key: str) -> str: - """Reconstruct dot structure from underscored key using heuristics.""" - # Block indices: input_blocks_1 -> input_blocks.1 - blocks = [ - "input_blocks", "output_blocks", "middle_block", - "transformer_blocks", "single_transformer_blocks", - "double_blocks", "single_blocks", - "down_blocks", "up_blocks", "mid_block", - "attentions", "resnets", "upsamplers", "downsamplers" - ] - for b in blocks: - key = re.sub(f"{b}_(\\d+)", f"{b}.\\1", key) - - # Sequential numbers: .1_1 -> .1.1 - key = re.sub(r"(\.\d+)_(\d+)", r"\1.\2", key) - - # Separators that should be preceded by dot - separators = [ - "transformer_blocks", "single_transformer_blocks", - "attn1", "attn2", "attn", - "img_attn", "txt_attn", - "to_q", "to_k", "to_v", "to_out", - "q_proj", "k_proj", "v_proj", "out_proj", - "qkv", "proj", "ff", "net", - "norm1", "norm2", "norm3", "norm", - "time_emb_proj" - ] - for s in separators: - key = key.replace(f"_{s}", f".{s}") - - # Special case: to_out_0 -> to_out.0 - key = re.sub(r"to_out_(\d+)", r"to_out.\1", key) - - return key - - def _format_lora_key(key: str) -> str: """ Format the key for LoRA saving. @@ -722,9 +686,7 @@ def _format_lora_key(key: str) -> str: # Handle legacy lora_unet_ prefix (for resizing without base) if key.startswith("lora_unet_"): - core = key[10:] - dotted = _reconstruct_dots(core) - return f"diffusion_model.{dotted}" + return key # Handle already-prefixed keys if key.startswith("diffusion_model."): diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index 62cc220..9ea8a04 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -30,6 +30,7 @@ def merge_multi_loras( device: str, save_dtype: torch.dtype, output_filename: str, + base_model_path: Optional[str] = None, verbose: bool = True, ) -> str: """ @@ -77,11 +78,59 @@ def merge_multi_loras( 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: - layer_map[core] = [] - layer_map[core].append(i) + if base_model_path: + BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."] + LORA_PREFIXES = [ + "lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_", + "lycoris_", "diffusion_model.", "transformer.", "unet." + ] + + def extract_core_layer_base(key: str) -> str: + result = key + if result.endswith(".weight"): + result = result[:-7] + elif result.endswith(".bias"): + result = result[:-5] + for prefix in BASE_PREFIXES: + if result.startswith(prefix): + result = result[len(prefix):] + break + return result + + def extract_core_layer_lora(block_name: str) -> str: + result = block_name + for prefix in LORA_PREFIXES: + if result.startswith(prefix): + result = result[len(prefix):] + break + if result.endswith(".lora"): + result = result[:-5] + return result.replace(".", "_") + + with MemoryEfficientSafeOpen(base_model_path, low_memory=True) as base_handler: + base_keys = list(base_handler.keys()) + + lora_lookup = {} + for i, info in enumerate(lora_infos): + for block_name in info["pairs"].keys(): + core_lora = extract_core_layer_lora(block_name) + if core_lora not in lora_lookup: + lora_lookup[core_lora] = [] + lora_lookup[core_lora].append((i, block_name)) + + for base_key in base_keys: + if base_key.endswith(".weight"): + core_base = extract_core_layer_base(base_key) + core_underscored = core_base.replace(".", "_") + if core_underscored in lora_lookup: + layer_map[core_base] = lora_lookup[core_underscored] # List of (idx, block_name) + else: + for i, info in enumerate(lora_infos): + for block_name in info["pairs"].keys(): + core = block_name # Keep block_name as-is without reference + if core not in layer_map: + layer_map[core] = [] + layer_map[core].append((i, block_name)) if verbose: print(f"[LoRA Multi-Merge] Total unique layers after resolving naming: {len(layer_map)}") @@ -110,10 +159,9 @@ def merge_multi_loras( max_rank = 0 # First pass: load and determine max rank - for idx in indices: + for idx, orig_block 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: @@ -250,6 +298,8 @@ class LoRAMultiMerge(io.ComfyNode): category="ModelUtils/LoRA/Merge", description="Merge 1-8 LoRAs into a single LoRA file. This resolves different ranks and naming conventions.", inputs=[ + io.Combo.Input("base_model", options=["None"] + folder_paths.get_filename_list("diffusion_models"), default="None", + tooltip="Optional reference model to resolve key naming issues across formats. If None, input keys are preserved verbatim."), io.Combo.Input("lora_count", options=[str(i) for i in range(1, 9)], default="2", tooltip="Number of LoRAs to merge"), # LoRA 1 @@ -288,7 +338,7 @@ class LoRAMultiMerge(io.ComfyNode): ) @classmethod - def execute(cls, lora_count, + 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, merge_mode, output_filename, save_dtype, device) -> io.NodeOutput: @@ -310,6 +360,9 @@ class LoRAMultiMerge(io.ComfyNode): raise ValueError("No LoRAs selected for merging.") dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype] + base_path = None + if base_model and base_model != "None": + base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model) path = merge_multi_loras( lora_paths=valid_paths, @@ -317,7 +370,8 @@ class LoRAMultiMerge(io.ComfyNode): merge_mode=merge_mode, device=device, save_dtype=dtype, - output_filename=output_filename + output_filename=output_filename, + base_model_path=base_path ) return io.NodeOutput(path) @@ -334,6 +388,8 @@ class LoRAMultiMergeDARE(io.ComfyNode): category="ModelUtils/LoRA/Merge", description="Merge 1-8 LoRAs into a single LoRA file using DARE-Ties. Drops small values and resolves sign conflicts.", inputs=[ + io.Combo.Input("base_model", options=["None"] + folder_paths.get_filename_list("diffusion_models"), default="None", + tooltip="Optional reference model to resolve key naming issues across formats. If None, input keys are preserved verbatim."), io.Combo.Input("lora_count", options=[str(i) for i in range(1, 9)], default="2"), # LoRA 1 io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"), tooltip="First LoRA"), @@ -366,7 +422,7 @@ class LoRAMultiMergeDARE(io.ComfyNode): ) @classmethod - def execute(cls, lora_count, + 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, drop_rate, trim_quantile, seed, output_filename, save_dtype, device) -> io.NodeOutput: @@ -388,6 +444,9 @@ class LoRAMultiMergeDARE(io.ComfyNode): raise ValueError("No LoRAs selected for merging.") dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype] + base_path = None + if base_model and base_model != "None": + base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model) path = merge_multi_loras_dare( lora_paths=valid_paths, @@ -397,7 +456,8 @@ class LoRAMultiMergeDARE(io.ComfyNode): seed=seed, device=device, save_dtype=dtype, - output_filename=output_filename + output_filename=output_filename, + base_model_path=base_path ) return io.NodeOutput(path) @@ -411,6 +471,7 @@ def merge_multi_loras_dare( device: str, save_dtype: torch.dtype, output_filename: str, + base_model_path: Optional[str] = None, verbose: bool = True, ) -> str: """ @@ -448,10 +509,59 @@ def merge_multi_loras_dare( if verbose: print(f"[LoRA Multi-Merge DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}") - for block_name in pairs.keys(): - core = _format_lora_key(block_name) - if core not in layer_map: layer_map[core] = [] - layer_map[core].append(i) + if base_model_path: + BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."] + LORA_PREFIXES = [ + "lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_", + "lycoris_", "diffusion_model.", "transformer.", "unet." + ] + + def extract_core_layer_base(key: str) -> str: + result = key + if result.endswith(".weight"): + result = result[:-7] + elif result.endswith(".bias"): + result = result[:-5] + for prefix in BASE_PREFIXES: + if result.startswith(prefix): + result = result[len(prefix):] + break + return result + + def extract_core_layer_lora(block_name: str) -> str: + result = block_name + for prefix in LORA_PREFIXES: + if result.startswith(prefix): + result = result[len(prefix):] + break + if result.endswith(".lora"): + result = result[:-5] + return result.replace(".", "_") + + with MemoryEfficientSafeOpen(base_model_path, low_memory=True) as base_handler: + base_keys = list(base_handler.keys()) + + lora_lookup = {} + for i, info in enumerate(lora_infos): + for block_name in info["pairs"].keys(): + core_lora = extract_core_layer_lora(block_name) + if core_lora not in lora_lookup: + lora_lookup[core_lora] = [] + lora_lookup[core_lora].append((i, block_name)) + + for base_key in base_keys: + if base_key.endswith(".weight"): + core_base = extract_core_layer_base(base_key) + core_underscored = core_base.replace(".", "_") + if core_underscored in lora_lookup: + layer_map[core_base] = lora_lookup[core_underscored] # List of (idx, block_name) + else: + for i, info in enumerate(lora_infos): + for block_name in info["pairs"].keys(): + core = block_name + if core not in layer_map: + layer_map[core] = [] + layer_map[core].append((i, block_name)) if verbose: print(f"[LoRA Multi-Merge DARE] Total unique layers after resolving naming: {len(layer_map)}") @@ -474,9 +584,8 @@ def merge_multi_loras_dare( ups = [] max_rank = 0 - for idx in indices: + for idx, orig_block 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 @@ -589,6 +698,8 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode): category="ModelUtils/LoRA/Merge", description="Merge 1-8 LoRAs into a single LoRA file using Enhanced DARE-Ties. Uses dynamic probability masking based on value magnitudes.", inputs=[ + io.Combo.Input("base_model", options=["None"] + folder_paths.get_filename_list("diffusion_models"), default="None", + tooltip="Optional reference model to resolve key naming issues across formats. If None, input keys are preserved verbatim."), io.Combo.Input("lora_count", options=[str(i) for i in range(1, 9)], default="2"), # LoRA 1 io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"), tooltip="First LoRA"), @@ -623,7 +734,7 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode): ) @classmethod - def execute(cls, lora_count, + 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, mask_power, min_keep_prob, mask_smooth, trim_quantile, seed, output_filename, save_dtype, device) -> io.NodeOutput: @@ -645,6 +756,9 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode): raise ValueError("No LoRAs selected for merging.") dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype] + base_path = None + if base_model and base_model != "None": + base_path = folder_paths.get_full_path_or_raise("diffusion_models", base_model) path = merge_multi_loras_dare_enhanced( lora_paths=valid_paths, @@ -656,7 +770,8 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode): seed=seed, device=device, save_dtype=dtype, - output_filename=output_filename + output_filename=output_filename, + base_model_path=base_path ) return io.NodeOutput(path) @@ -673,6 +788,7 @@ def merge_multi_loras_dare_enhanced( device: str, save_dtype: torch.dtype, output_filename: str, + base_model_path: Optional[str] = None, verbose: bool = True, ) -> str: """ @@ -709,10 +825,59 @@ def merge_multi_loras_dare_enhanced( if verbose: print(f"[LoRA Multi-Merge Enhanced DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}") - for block_name in pairs.keys(): - core = _format_lora_key(block_name) - if core not in layer_map: layer_map[core] = [] - layer_map[core].append(i) + if base_model_path: + BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."] + LORA_PREFIXES = [ + "lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_", + "lycoris_", "diffusion_model.", "transformer.", "unet." + ] + + def extract_core_layer_base(key: str) -> str: + result = key + if result.endswith(".weight"): + result = result[:-7] + elif result.endswith(".bias"): + result = result[:-5] + for prefix in BASE_PREFIXES: + if result.startswith(prefix): + result = result[len(prefix):] + break + return result + + def extract_core_layer_lora(block_name: str) -> str: + result = block_name + for prefix in LORA_PREFIXES: + if result.startswith(prefix): + result = result[len(prefix):] + break + if result.endswith(".lora"): + result = result[:-5] + return result.replace(".", "_") + + with MemoryEfficientSafeOpen(base_model_path, low_memory=True) as base_handler: + base_keys = list(base_handler.keys()) + + lora_lookup = {} + for i, info in enumerate(lora_infos): + for block_name in info["pairs"].keys(): + core_lora = extract_core_layer_lora(block_name) + if core_lora not in lora_lookup: + lora_lookup[core_lora] = [] + lora_lookup[core_lora].append((i, block_name)) + + for base_key in base_keys: + if base_key.endswith(".weight"): + core_base = extract_core_layer_base(base_key) + core_underscored = core_base.replace(".", "_") + if core_underscored in lora_lookup: + layer_map[core_base] = lora_lookup[core_underscored] # List of (idx, block_name) + else: + for i, info in enumerate(lora_infos): + for block_name in info["pairs"].keys(): + core = block_name + if core not in layer_map: + layer_map[core] = [] + layer_map[core].append((i, block_name)) if verbose: print(f"[LoRA Multi-Merge Enhanced DARE] Total unique layers after resolving naming: {len(layer_map)}") @@ -735,9 +900,8 @@ def merge_multi_loras_dare_enhanced( ups = [] max_rank = 0 - for idx in indices: + for idx, orig_block 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 diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index 1c56dc9..735a284 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -471,8 +471,8 @@ def resize_lora_file( up_suffix = format_info["up_suffix"] alpha_suffix = format_info["alpha_suffix"] - # Standardize prefix using heuristics if possible - new_block_name = _format_lora_key(block_name) + # Preserve original key format from input LoRA verbatim + new_block_name = block_name block_sd = { f"{new_block_name}{down_suffix}": result["lora_down"].to(save_dtype),