diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index 8712ac6..758db03 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -19,6 +19,7 @@ from .device_utils import ( ) from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from .quantization_guard import inspect_low_bit_input @@ -501,16 +502,20 @@ def extract_lora_from_files( handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load) try: + low_bit_a = inspect_low_bit_input(handler_a, f"Model A ({model_a_path})", "LoRA Extract") + low_bit_b = inspect_low_bit_input(handler_b, f"Model B ({model_b_path})", "LoRA Extract") keys_a = set(handler_a.keys()) keys_b = set(handler_b.keys()) direct_diff_suffixes = (".scale", ".lin") extraction_keys = [ k for k in keys_a - if k.endswith(".weight") + if k not in low_bit_a + and k not in low_bit_b + and (k.endswith(".weight") or ( include_1d_diffs and (len(handler_a.get_shape(k)) == 1 or k.endswith(direct_diff_suffixes)) - ) + )) ] pbar = comfy.utils.ProgressBar(len(extraction_keys)) stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0} diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index 7048bb2..ceb63cd 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -12,16 +12,209 @@ 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 ( detect_lora_format, - extract_lora_pairs, detect_lora_rank, - _compile_patterns, - _matches_any_pattern, - _format_lora_key + parse_lora_layers, + select_output_dtype, ) +from .quantization_guard import inspect_low_bit_input, layer_has_low_bit, write_preserved_tensor from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +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 _strip_prefix(value: str, prefixes: List[str]) -> str: + for prefix in prefixes: + if value.startswith(prefix): + return value[len(prefix):] + return value + + +def _lora_core(block_name: str) -> str: + core = _strip_prefix(block_name, LORA_PREFIXES) + if core.endswith(".lora"): + core = core[:-5] + return core + + +def _build_layer_map(lora_infos: List[Dict], base_model_path: Optional[str]) -> Dict[str, List[Tuple[int, str]]]: + """Build a union map while optionally normalizing names against a reference model.""" + if not base_model_path: + layer_map = {} + for idx, info in enumerate(lora_infos): + for block_name in info["pairs"]: + layer_map.setdefault(block_name, []).append((idx, block_name)) + return layer_map + + with MemoryEfficientSafeOpen(base_model_path, low_memory=True) as base_handler: + base_keys = list(base_handler.keys()) + + normalized_base = {} + for base_key in base_keys: + normalized = _strip_prefix(base_key, BASE_PREFIXES) + normalized_base[normalized] = base_key + normalized_base[normalized.replace(".", "_")] = base_key + + layer_map = {} + mapped = set() + for idx, info in enumerate(lora_infos): + for block_name, block_keys in info["pairs"].items(): + core = _lora_core(block_name) + target = None + if "down" in block_keys or "up" in block_keys: + candidate = f"{core}.weight" + target = normalized_base.get(candidate) or normalized_base.get(candidate.replace(".", "_")) + elif "diff_b" in block_keys: + candidate = f"{core}.bias" + target = normalized_base.get(candidate) or normalized_base.get(candidate.replace(".", "_")) + elif "diff" in block_keys: + target = normalized_base.get(core) or normalized_base.get(core.replace(".", "_")) + if target is None: + candidate = f"{core}.weight" + target = normalized_base.get(candidate) or normalized_base.get(candidate.replace(".", "_")) + + if target is None: + continue + normalized_target = _strip_prefix(target, BASE_PREFIXES) + if "diff_b" in block_keys and normalized_target.endswith(".bias"): + output_core = normalized_target[:-5] + elif "diff" in block_keys and normalized_target.endswith(".weight"): + output_core = normalized_target[:-7] + elif "diff" in block_keys: + output_core = normalized_target + else: + output_core = normalized_target[:-7] if normalized_target.endswith(".weight") else normalized_target + output_core = f"diffusion_model.{output_core}" + layer_map.setdefault(output_core, []).append((idx, block_name)) + mapped.add((idx, block_name)) + + # Direct patches must not disappear merely because a reference model cannot resolve them. + for idx, info in enumerate(lora_infos): + for block_name, block_keys in info["pairs"].items(): + if (idx, block_name) not in mapped and ("diff" in block_keys or "diff_b" in block_keys): + layer_map.setdefault(block_name, []).append((idx, block_name)) + return layer_map + + +def _load_direct_inputs(indices, lora_infos, handlers, kind, device, include_1d_diffs): + tensors = [] + source_dtypes = [] + expected_shape = None + for idx, orig_block in indices: + block_keys = lora_infos[idx]["pairs"][orig_block] + if kind not in block_keys: + continue # Missing direct layers are implicit zero contributions. + key = block_keys[kind] + tensor = handlers[idx].get_tensor(key) + if tensor.ndim == 1 and not include_1d_diffs: + continue + if expected_shape is None: + expected_shape = tuple(tensor.shape) + elif tuple(tensor.shape) != expected_shape: + raise ValueError( + f"Direct LoRA shape mismatch for {key}: {tuple(tensor.shape)} != {expected_shape}" + ) + source_dtypes.append(handlers[idx].get_dtype(key)) + tensors.append((tensor.to(device=device, dtype=torch.float32), lora_infos[idx]["weight"])) + return tensors, source_dtypes + + +def _logical_layer_dtypes(indices, lora_infos, handlers): + source_dtypes = [] + for idx, orig_block in indices: + block_keys = lora_infos[idx]["pairs"][orig_block] + for name in ("down", "up", "alpha", "diff", "diff_b"): + if name in block_keys: + source_dtypes.append(handlers[idx].get_dtype(block_keys[name])) + return source_dtypes + + +def _merge_direct_weighted(tensors): + result = torch.zeros_like(tensors[0][0]) + for tensor, weight in tensors: + result.add_(tensor, alpha=weight) + return result + + +def _inspect_lora_inputs(handlers, lora_paths, operation): + return [ + inspect_low_bit_input( + handler, f"LoRA {i + 1} ({lora_paths[i]})", operation + ) + for i, handler in enumerate(handlers) + ] + + +def _indices_have_low_bit(indices, lora_infos): + return any( + layer_has_low_bit( + lora_infos[idx]["pairs"][block_name], + lora_infos[idx]["low_bit_keys"], + ) + for idx, block_name in indices + ) + + +def _preserve_earliest_layer(writer, indices, lora_infos, handlers, written_keys): + """Write the earliest selected adapter's original representation unchanged.""" + idx, block_name = min(indices, key=lambda item: item[0]) + block_keys = lora_infos[idx]["pairs"][block_name] + count = 0 + for name in ("down", "up", "alpha", "diff", "diff_b"): + key = block_keys.get(name) + if key is None or key in written_keys: + continue + write_preserved_tensor(writer, key, handlers[idx]) + written_keys.add(key) + count += 1 + for key in lora_infos[idx]["passthrough_keys"]: + if not key.startswith(f"{block_name}.") or key in written_keys: + continue + write_preserved_tensor(writer, key, handlers[idx]) + written_keys.add(key) + count += 1 + return count + + +def _preserve_low_bit_auxiliaries(writer, lora_infos, handlers, written_keys): + """Preserve isolated unparsed low-bit tensors from the earliest adapter.""" + count = 0 + affected_keys = { + key + for info in lora_infos + for key in info["passthrough_keys"] + if key in info["low_bit_keys"] + } + for key in sorted(affected_keys): + for idx, info in enumerate(lora_infos): + if key not in info["passthrough_keys"]: + continue + if key not in written_keys: + write_preserved_tensor(writer, key, handlers[idx]) + written_keys.add(key) + count += 1 + break + return count + + +def _ensure_guarded_layers_mapped(layer_map, lora_infos): + mapped = { + (idx, block_name) + for indices in layer_map.values() + for idx, block_name in indices + } + for idx, info in enumerate(lora_infos): + for block_name, block_keys in info["pairs"].items(): + entry = (idx, block_name) + if entry not in mapped and layer_has_low_bit(block_keys, info["low_bit_keys"]): + layer_map.setdefault(block_name, []).append(entry) + + def merge_multi_loras( lora_paths: List[str], @@ -32,6 +225,7 @@ def merge_multi_loras( output_filename: str, base_model_path: Optional[str] = None, verbose: bool = True, + include_1d_diffs: bool = False, ) -> str: """ Merge multiple LoRAs into a single LoRA file. @@ -57,12 +251,13 @@ def merge_multi_loras( # 1. Analyze all LoRAs layer_map = {} lora_infos = [] + low_bit_keys = _inspect_lora_inputs(handlers, lora_paths, "LoRA Multi-Merge") 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) + pairs, passthrough_keys = parse_lora_layers(keys) + rank, alpha = detect_lora_rank(handler, pairs, low_bit_keys[i]) info = { "name": os.path.basename(lora_paths[i]), @@ -71,68 +266,17 @@ def merge_multi_loras( "rank": rank, "alpha": alpha, "weight": lora_weights[i], - "scale": alpha / rank if rank > 0 else 1.0 + "scale": alpha / rank if rank > 0 else 1.0, + "low_bit_keys": low_bit_keys[i], + "passthrough_keys": passthrough_keys, } lora_infos.append(info) if verbose: print(f"[LoRA Multi-Merge] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}, format={fmt['format']}") - 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: - # Always use standard format with diffusion_model. prefix when reference model is provided - output_key = core_base if core_base.startswith("diffusion_model.") else "diffusion_model." + core_base - layer_map[output_key] = 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)) + layer_map = _build_layer_map(lora_infos, base_model_path) + _ensure_guarded_layers_mapped(layer_map, lora_infos) if verbose: print(f"[LoRA Multi-Merge] Total unique layers after resolving naming: {len(layer_map)}") @@ -148,6 +292,7 @@ def merge_multi_loras( pbar = comfy.utils.ProgressBar(len(layer_map)) tensor_count = 0 + written_keys = set() writer = IncrementalSafetensorsWriter(output_path, metadata=metadata) writer.__enter__() @@ -155,9 +300,15 @@ def merge_multi_loras( # 2. Merge each core layer with torch.no_grad(): for core, indices in tqdm(layer_map.items(), desc="Merging layers", unit="layers"): + if _indices_have_low_bit(indices, lora_infos): + tensor_count += _preserve_earliest_layer( + writer, indices, lora_infos, handlers, written_keys + ) + pbar.update(1) + continue + logical_source_dtypes = _logical_layer_dtypes(indices, lora_infos, handlers) downs = [] ups = [] - is_conv = False max_rank = 0 # First pass: load and determine max rank @@ -171,7 +322,6 @@ def merge_multi_loras( 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) @@ -180,7 +330,6 @@ def merge_multi_loras( 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) @@ -192,17 +341,13 @@ def merge_multi_loras( downs.append((t_down, info)) ups.append((t_up, info)) - if not downs: - pbar.update(1) - continue - - if merge_mode == "concatenate": + if downs and 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: + elif downs: # 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 @@ -245,26 +390,54 @@ def merge_multi_loras( # 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 + if downs: + layer_dtype = select_output_dtype(logical_source_dtypes, save_dtype) + writer.write_dict({ + f"{core}.lora_down.weight": merged_down.to(layer_dtype).cpu().contiguous(), + f"{core}.lora_up.weight": merged_up.to(layer_dtype).cpu().contiguous(), + f"{core}.alpha": torch.tensor(new_alpha, dtype=layer_dtype), + }) + tensor_count += 3 + written_keys.update({ + f"{core}.lora_down.weight", f"{core}.lora_up.weight", f"{core}.alpha" + }) + del merged_down, merged_up + for d, _ in downs: del d + for u, _ in ups: del u + downs.clear() + ups.clear() + + for kind in ("diff", "diff_b"): + direct_inputs, _ = _load_direct_inputs( + indices, lora_infos, handlers, kind, device, include_1d_diffs + ) + if not direct_inputs: + continue + merged_direct = _merge_direct_weighted(direct_inputs) + direct_dtype = select_output_dtype( + logical_source_dtypes, + save_dtype, + is_1d_diff=merged_direct.ndim == 1, + ) + writer.write( + f"{core}.{kind}", + merged_direct.to(direct_dtype).cpu().contiguous(), + ) + tensor_count += 1 + written_keys.add(f"{core}.{kind}") + del merged_direct + for tensor, _ in direct_inputs: + del tensor - # 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) + tensor_count += _preserve_low_bit_auxiliaries( + writer, lora_infos, handlers, written_keys + ) finally: writer.__exit__(None, None, None) @@ -334,6 +507,8 @@ class LoRAMultiMerge(io.ComfyNode): io.String.Input("output_filename", default="merged_lora"), io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"), + io.Boolean.Input("include_1d_diffs", default=False, + tooltip="Include and merge 1D direct-diff tensors as FP32."), ], outputs=[io.String.Output(display_name="output_path")], is_output_node=True, @@ -343,7 +518,7 @@ class LoRAMultiMerge(io.ComfyNode): 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: + merge_mode, output_filename, save_dtype, device, include_1d_diffs) -> io.NodeOutput: count = int(lora_count) names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8] @@ -373,7 +548,8 @@ class LoRAMultiMerge(io.ComfyNode): device=device, save_dtype=dtype, output_filename=output_filename, - base_model_path=base_path + base_model_path=base_path, + include_1d_diffs=include_1d_diffs, ) return io.NodeOutput(path) @@ -418,6 +594,8 @@ class LoRAMultiMergeDARE(io.ComfyNode): io.String.Input("output_filename", default="merged_lora_dare"), io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"), + io.Boolean.Input("include_1d_diffs", default=False, + tooltip="Include and merge 1D direct-diff tensors as FP32."), ], outputs=[io.String.Output(display_name="output_path")], is_output_node=True, @@ -427,7 +605,8 @@ class LoRAMultiMergeDARE(io.ComfyNode): 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: + drop_rate, trim_quantile, seed, output_filename, save_dtype, device, + include_1d_diffs) -> io.NodeOutput: count = int(lora_count) names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8] @@ -459,7 +638,8 @@ class LoRAMultiMergeDARE(io.ComfyNode): device=device, save_dtype=dtype, output_filename=output_filename, - base_model_path=base_path + base_model_path=base_path, + include_1d_diffs=include_1d_diffs, ) return io.NodeOutput(path) @@ -475,6 +655,7 @@ def merge_multi_loras_dare( output_filename: str, base_model_path: Optional[str] = None, verbose: bool = True, + include_1d_diffs: bool = False, ) -> str: """ Merge multiple LoRAs using DARE-Ties method applied to weights. @@ -494,78 +675,30 @@ def merge_multi_loras_dare( # 1. Analyze layer_map = {} lora_infos = [] + low_bit_keys = _inspect_lora_inputs( + handlers, lora_paths, "LoRA Multi-Merge DARE" + ) 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) + pairs, passthrough_keys = parse_lora_layers(keys) + rank, alpha = detect_lora_rank(handler, pairs, low_bit_keys[i]) info = { "name": os.path.basename(lora_paths[i]), "pairs": pairs, "rank": rank, "alpha": alpha, - "weight": lora_weights[i] + "weight": lora_weights[i], + "low_bit_keys": low_bit_keys[i], + "passthrough_keys": passthrough_keys, } lora_infos.append(info) if verbose: print(f"[LoRA Multi-Merge DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}") - 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: - # Always use standard format with diffusion_model. prefix when reference model is provided - output_key = core_base if core_base.startswith("diffusion_model.") else "diffusion_model." + core_base - layer_map[output_key] = 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)) + layer_map = _build_layer_map(lora_infos, base_model_path) + _ensure_guarded_layers_mapped(layer_map, lora_infos) if verbose: print(f"[LoRA Multi-Merge DARE] Total unique layers after resolving naming: {len(layer_map)}") @@ -577,6 +710,7 @@ def merge_multi_loras_dare( pbar = comfy.utils.ProgressBar(len(layer_map)) tensor_count = 0 + written_keys = set() writer = IncrementalSafetensorsWriter(output_path, metadata={"ss_training_comment": "Merged via DARE-Ties"}) writer.__enter__() @@ -584,6 +718,13 @@ def merge_multi_loras_dare( # 2. Merge with torch.no_grad(): for core, indices in tqdm(layer_map.items(), desc="Merging layers (DARE)", unit="layers"): + if _indices_have_low_bit(indices, lora_infos): + tensor_count += _preserve_earliest_layer( + writer, indices, lora_infos, handlers, written_keys + ) + pbar.update(1) + continue + logical_source_dtypes = _logical_layer_dtypes(indices, lora_infos, handlers) downs = [] ups = [] max_rank = 0 @@ -600,15 +741,11 @@ def merge_multi_loras_dare( 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: + if dim_to_pad is not None and 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 @@ -648,27 +785,54 @@ def merge_multi_loras_dare( 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) + if downs: + merged_down = process_ties_dare(downs, [1.0]*len(downs), 0) + merged_up = process_ties_dare(ups, [1.0]*len(ups), 1) + layer_dtype = select_output_dtype(logical_source_dtypes, save_dtype) + writer.write_dict({ + f"{core}.lora_down.weight": merged_down.to(layer_dtype).cpu().contiguous(), + f"{core}.lora_up.weight": merged_up.to(layer_dtype).cpu().contiguous(), + f"{core}.alpha": torch.tensor(float(max_rank), dtype=layer_dtype), + }) + tensor_count += 3 + written_keys.update({ + f"{core}.lora_down.weight", f"{core}.lora_up.weight", f"{core}.alpha" + }) + del merged_down, merged_up + for d, _ in downs: del d + for u, _ in ups: del u + downs.clear() + ups.clear() - 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() + for kind in ("diff", "diff_b"): + direct_inputs, _ = _load_direct_inputs( + indices, lora_infos, handlers, kind, device, include_1d_diffs + ) + if not direct_inputs: + continue + merged_direct = process_ties_dare( + direct_inputs, [1.0] * len(direct_inputs), None + ) + direct_dtype = select_output_dtype( + logical_source_dtypes, + save_dtype, + is_1d_diff=merged_direct.ndim == 1, + ) + writer.write(f"{core}.{kind}", merged_direct.to(direct_dtype).cpu().contiguous()) + tensor_count += 1 + written_keys.add(f"{core}.{kind}") + del merged_direct + for tensor, _ in direct_inputs: + del tensor import gc gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() pbar.update(1) + tensor_count += _preserve_low_bit_auxiliaries( + writer, lora_infos, handlers, written_keys + ) finally: writer.__exit__(None, None, None) @@ -732,6 +896,8 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode): io.String.Input("output_filename", default="merged_lora_dare_enhanced"), io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"), io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"), + io.Boolean.Input("include_1d_diffs", default=False, + tooltip="Include and merge 1D direct-diff tensors as FP32."), ], outputs=[io.String.Output(display_name="output_path")], is_output_node=True, @@ -741,7 +907,8 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode): 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: + mask_power, min_keep_prob, mask_smooth, trim_quantile, seed, output_filename, save_dtype, device, + include_1d_diffs) -> io.NodeOutput: count = int(lora_count) names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8] @@ -775,7 +942,8 @@ class LoRAMultiMergeDAREEnhanced(io.ComfyNode): device=device, save_dtype=dtype, output_filename=output_filename, - base_model_path=base_path + base_model_path=base_path, + include_1d_diffs=include_1d_diffs, ) return io.NodeOutput(path) @@ -794,6 +962,7 @@ def merge_multi_loras_dare_enhanced( output_filename: str, base_model_path: Optional[str] = None, verbose: bool = True, + include_1d_diffs: bool = False, ) -> str: """ Merge multiple LoRAs using Enhanced DARE-Ties method applied to weights. @@ -812,78 +981,30 @@ def merge_multi_loras_dare_enhanced( # 1. Analyze layer_map = {} lora_infos = [] + low_bit_keys = _inspect_lora_inputs( + handlers, lora_paths, "LoRA Multi-Merge Enhanced DARE" + ) 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) + pairs, passthrough_keys = parse_lora_layers(keys) + rank, alpha = detect_lora_rank(handler, pairs, low_bit_keys[i]) info = { "name": os.path.basename(lora_paths[i]), "pairs": pairs, "rank": rank, "alpha": alpha, - "weight": lora_weights[i] + "weight": lora_weights[i], + "low_bit_keys": low_bit_keys[i], + "passthrough_keys": passthrough_keys, } lora_infos.append(info) if verbose: print(f"[LoRA Multi-Merge Enhanced DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}") - 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: - # Always use standard format with diffusion_model. prefix when reference model is provided - output_key = core_base if core_base.startswith("diffusion_model.") else "diffusion_model." + core_base - layer_map[output_key] = 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)) + layer_map = _build_layer_map(lora_infos, base_model_path) + _ensure_guarded_layers_mapped(layer_map, lora_infos) if verbose: print(f"[LoRA Multi-Merge Enhanced DARE] Total unique layers after resolving naming: {len(layer_map)}") @@ -895,6 +1016,7 @@ def merge_multi_loras_dare_enhanced( pbar = comfy.utils.ProgressBar(len(layer_map)) tensor_count = 0 + written_keys = set() writer = IncrementalSafetensorsWriter(output_path, metadata={"ss_training_comment": "Merged via Enhanced DARE-Ties"}) writer.__enter__() @@ -902,6 +1024,13 @@ def merge_multi_loras_dare_enhanced( # 2. Merge with torch.no_grad(): for core, indices in tqdm(layer_map.items(), desc="Merging layers (Enhanced DARE)", unit="layers"): + if _indices_have_low_bit(indices, lora_infos): + tensor_count += _preserve_earliest_layer( + writer, indices, lora_infos, handlers, written_keys + ) + pbar.update(1) + continue + logical_source_dtypes = _logical_layer_dtypes(indices, lora_infos, handlers) downs = [] ups = [] max_rank = 0 @@ -918,10 +1047,6 @@ def merge_multi_loras_dare_enhanced( 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: @@ -948,7 +1073,7 @@ def merge_multi_loras_dare_enhanced( 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: + if dim_to_pad is not None and 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] @@ -968,27 +1093,54 @@ def merge_multi_loras_dare_enhanced( 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) + if downs: + merged_down = process_ties_dare_enhanced(downs, [1.0]*len(downs), 0) + merged_up = process_ties_dare_enhanced(ups, [1.0]*len(ups), 1) + layer_dtype = select_output_dtype(logical_source_dtypes, save_dtype) + writer.write_dict({ + f"{core}.lora_down.weight": merged_down.to(layer_dtype).cpu().contiguous(), + f"{core}.lora_up.weight": merged_up.to(layer_dtype).cpu().contiguous(), + f"{core}.alpha": torch.tensor(float(max_rank), dtype=layer_dtype), + }) + tensor_count += 3 + written_keys.update({ + f"{core}.lora_down.weight", f"{core}.lora_up.weight", f"{core}.alpha" + }) + del merged_down, merged_up + for d, _ in downs: del d + for u, _ in ups: del u + downs.clear() + ups.clear() - 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() + for kind in ("diff", "diff_b"): + direct_inputs, _ = _load_direct_inputs( + indices, lora_infos, handlers, kind, device, include_1d_diffs + ) + if not direct_inputs: + continue + merged_direct = process_ties_dare_enhanced( + direct_inputs, [1.0] * len(direct_inputs), None + ) + direct_dtype = select_output_dtype( + logical_source_dtypes, + save_dtype, + is_1d_diff=merged_direct.ndim == 1, + ) + writer.write(f"{core}.{kind}", merged_direct.to(direct_dtype).cpu().contiguous()) + tensor_count += 1 + written_keys.add(f"{core}.{kind}") + del merged_direct + for tensor, _ in direct_inputs: + del tensor import gc gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() pbar.update(1) + tensor_count += _preserve_low_bit_auxiliaries( + writer, lora_infos, handlers, written_keys + ) finally: writer.__exit__(None, None, None) diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index e77ba21..e3da014 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -15,6 +15,7 @@ from comfy_api.latest import io from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from .quantization_guard import inspect_low_bit_input, layer_has_low_bit, write_preserved_tensor from typing import Optional, Dict, Tuple, List @@ -34,6 +35,34 @@ from .lora_extract_svd import ( MIN_SV = 1e-6 +LORA_PAIR_SUFFIXES = ( + (".lora_A.default.weight", ".lora_B.default.weight", ".alpha", "peft"), + (".lora_down.weight", ".lora_up.weight", ".alpha", "comfy"), + (".lora_A.weight", ".lora_B.weight", ".alpha", "diffusers"), + (".lora.down.weight", ".lora.up.weight", ".alpha", "huggingface"), +) + + +def is_direct_diff_key(key: str) -> bool: + return key.endswith(".diff") or key.endswith(".diff_b") + + +def select_output_dtype( + source_dtypes: List[torch.dtype], + requested_dtype: torch.dtype, + force: bool = False, + is_1d_diff: bool = False, +) -> torch.dtype: + """Select an output dtype without losing precision-sensitive LoRA deltas.""" + if is_1d_diff: + return torch.float32 + if not force and torch.float32 in source_dtypes: + return torch.float32 + floating_dtypes = {torch.float16, torch.bfloat16, torch.float32, torch.float64} + if source_dtypes and not any(dtype in floating_dtypes for dtype in source_dtypes): + return source_dtypes[0] + return requested_dtype + # ============================================================================= # LoRA Format Detection - Support various LoRA sources @@ -69,59 +98,99 @@ def detect_lora_format(keys: List[str]) -> Dict: alpha_keys = [k for k in keys if '.alpha' in k or k.endswith('alpha')] format_info["has_alpha"] = len(alpha_keys) > 0 - # Detect format by key patterns - sample_keys = keys[:50] # Check more keys for full_diff detection + detected = [] + for down_suffix, up_suffix, alpha_suffix, name in LORA_PAIR_SUFFIXES: + if any(k.endswith(down_suffix) or k.endswith(up_suffix) for k in keys): + detected.append((down_suffix, up_suffix, alpha_suffix, name)) + has_full_diff = any(is_direct_diff_key(k) for k in keys) + if has_full_diff: + detected.append((None, None, None, "full_diff")) - # Full diff format: *.diff (NOT .lora_down, just straight weight diff) - if any(k.endswith('.diff') for k in sample_keys): - format_info["format"] = "full_diff" - format_info["down_suffix"] = ".diff" # Used for pairing, but no up - format_info["up_suffix"] = None # No up weight in full diff - format_info["alpha_suffix"] = None - format_info["is_full_diff"] = True - - # Kohya/A1111 format: lora_unet_down_blocks_0_*.lora_down.weight - elif any('lora_unet_' in k or 'lora_te' in k for k in sample_keys): - format_info["format"] = "kohya" - format_info["down_suffix"] = ".lora_down.weight" - format_info["up_suffix"] = ".lora_up.weight" - - # Diffusers format: transformer.*.lora_A.weight - elif any('.lora_A.weight' in k for k in sample_keys): - format_info["format"] = "diffusers" - format_info["down_suffix"] = ".lora_A.weight" - format_info["up_suffix"] = ".lora_B.weight" - format_info["alpha_suffix"] = ".alpha" - - # PEFT format: base_model.model.*.lora_A.default.weight - elif any('base_model.model.' in k and 'lora_' in k for k in sample_keys): - format_info["format"] = "peft" - # Handle nested structure - format_info["down_suffix"] = ".lora_A.default.weight" - format_info["up_suffix"] = ".lora_B.default.weight" - - # HuggingFace format: transformer.*.lora.down.weight (dots instead of underscores) - elif any('.lora.down.weight' in k for k in sample_keys): - format_info["format"] = "huggingface" - format_info["down_suffix"] = ".lora.down.weight" - format_info["up_suffix"] = ".lora.up.weight" - format_info["alpha_suffix"] = ".alpha" - - # ComfyUI native / standard - elif any('.lora_down.weight' in k for k in sample_keys): - format_info["format"] = "comfy" - format_info["down_suffix"] = ".lora_down.weight" - format_info["up_suffix"] = ".lora_up.weight" - - # Count LoRA layers - down_keys = [k for k in keys if format_info["down_suffix"] in k] - format_info["key_count"] = len(down_keys) + if len(detected) > 1: + format_info["format"] = "mixed" + elif detected: + down_suffix, up_suffix, alpha_suffix, name = detected[0] + format_info.update({ + "format": name, + "down_suffix": down_suffix, + "up_suffix": up_suffix, + "alpha_suffix": alpha_suffix, + }) + format_info["is_full_diff"] = has_full_diff and not any(item[0] for item in detected) + pairs, _ = parse_lora_layers(keys) + format_info["key_count"] = len(pairs) return format_info -def extract_lora_pairs(keys: List[str], format_info: Dict) -> Dict[str, Dict[str, str]]: +def parse_lora_layers(keys: List[str]) -> Tuple[Dict[str, Dict[str, str]], List[str]]: + """Parse low-rank, direct-diff, and pass-through tensors in one full scan.""" + layers: Dict[str, Dict[str, str]] = {} + consumed = set() + + for key in keys: + if key.endswith(".diff_b"): + block_name = key[:-7] + layers.setdefault(block_name, {})["diff_b"] = key + consumed.add(key) + continue + if key.endswith(".diff"): + block_name = key[:-5] + layers.setdefault(block_name, {})["diff"] = key + consumed.add(key) + continue + for down_suffix, up_suffix, alpha_suffix, format_name in LORA_PAIR_SUFFIXES: + if key.endswith(down_suffix): + block_name = key[:-len(down_suffix)] + layer = layers.setdefault(block_name, {}) + layer.update({ + "down": key, + "down_suffix": down_suffix, + "up_suffix": up_suffix, + "alpha_suffix": alpha_suffix, + "format": format_name, + }) + consumed.add(key) + break + if key.endswith(up_suffix): + block_name = key[:-len(up_suffix)] + layer = layers.setdefault(block_name, {}) + layer.update({ + "up": key, + "down_suffix": down_suffix, + "up_suffix": up_suffix, + "alpha_suffix": alpha_suffix, + "format": format_name, + }) + consumed.add(key) + break + + for key in keys: + if not key.endswith(".alpha"): + continue + block_name = key[:-6] + if block_name in layers and ("down" in layers[block_name] or "up" in layers[block_name]): + layers[block_name]["alpha"] = key + consumed.add(key) + + # Incomplete low-rank groups cannot be resized, but their tensors must survive. + for block_name, layer in list(layers.items()): + if ("down" in layer) == ("up" in layer): + continue + for name in ("down", "up", "alpha"): + key = layer.pop(name, None) + if key is not None: + consumed.discard(key) + for name in ("down_suffix", "up_suffix", "alpha_suffix", "format"): + layer.pop(name, None) + if not layer: + del layers[block_name] + + return layers, [key for key in keys if key not in consumed] + + +def extract_lora_pairs(keys: List[str], format_info: Optional[Dict] = None) -> Dict[str, Dict[str, str]]: """ Group LoRA keys into down/up/alpha pairs. @@ -131,50 +200,15 @@ def extract_lora_pairs(keys: List[str], format_info: Dict) -> Dict[str, Dict[str Dict[block_name, {"down": key, "up": key, "alpha": key}] For full_diff: {"diff": key, "diff_b": key} """ - down_suffix = format_info["down_suffix"] - up_suffix = format_info["up_suffix"] - alpha_suffix = format_info["alpha_suffix"] - is_full_diff = format_info.get("is_full_diff", False) - - pairs = {} - - if is_full_diff: - # Full diff format: group .diff and .diff_b - for key in keys: - if key.endswith('.diff'): - block_name = key[:-5] # Remove .diff - if block_name not in pairs: - pairs[block_name] = {} - pairs[block_name]["diff"] = key - elif key.endswith('.diff_b'): - block_name = key[:-7] # Remove .diff_b - if block_name not in pairs: - pairs[block_name] = {} - pairs[block_name]["diff_b"] = key - else: - # Standard LoRA format - for key in keys: - if down_suffix and down_suffix in key: - block_name = key.replace(down_suffix, "") - if block_name not in pairs: - pairs[block_name] = {} - pairs[block_name]["down"] = key - elif up_suffix and up_suffix in key: - block_name = key.replace(up_suffix, "") - if block_name not in pairs: - pairs[block_name] = {} - pairs[block_name]["up"] = key - elif alpha_suffix and (alpha_suffix in key or key.endswith('.alpha')): - # Handle various alpha key formats - block_name = key.replace(alpha_suffix, "").replace(".alpha", "") - if block_name not in pairs: - pairs[block_name] = {} - pairs[block_name]["alpha"] = key - + pairs, _ = parse_lora_layers(keys) return pairs -def detect_lora_rank(handler: MemoryEfficientSafeOpen, pairs: Dict) -> Tuple[int, float]: +def detect_lora_rank( + handler: MemoryEfficientSafeOpen, + pairs: Dict, + low_bit_keys: Optional[set[str]] = None, +) -> Tuple[int, float]: """ Detect the rank and alpha of an existing LoRA. @@ -196,7 +230,11 @@ def detect_lora_rank(handler: MemoryEfficientSafeOpen, pairs: Dict) -> Tuple[int network_dim = shape[0] # Get alpha if present - if network_alpha is None and "alpha" in block_keys: + if ( + network_alpha is None + and "alpha" in block_keys + and block_keys["alpha"] not in (low_bit_keys or set()) + ): alpha_tensor = handler.get_tensor(block_keys["alpha"]) network_alpha = float(alpha_tensor.item()) @@ -406,13 +444,14 @@ def resize_lora_file( handler = MemoryEfficientSafeOpen(lora_path, low_memory=lazy_load) try: - metadata = handler.metadata().copy() + low_bit_keys = inspect_low_bit_input(handler, f"LoRA ({lora_path})", "LoRA Resize") + metadata = (handler.metadata() or {}).copy() all_keys = handler.keys() # Detect format and extract pairs format_info = detect_lora_format(all_keys) - pairs = extract_lora_pairs(all_keys, format_info) - network_dim, network_alpha = detect_lora_rank(handler, pairs) + pairs, passthrough_keys = parse_lora_layers(all_keys) + network_dim, network_alpha = detect_lora_rank(handler, pairs, low_bit_keys) scale = network_alpha / network_dim if network_dim > 0 else 1.0 @@ -437,13 +476,49 @@ def resize_lora_file( output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors") fro_list = [] - pbar = comfy.utils.ProgressBar(len(pairs)) + pbar = comfy.utils.ProgressBar(len(pairs) + len(passthrough_keys)) writer = IncrementalSafetensorsWriter(output_path, metadata=metadata) writer.__enter__() try: + preserved_keys = set() with torch.no_grad(): for block_name, block_keys in tqdm(pairs.items(), desc="Resizing layers", unit="layers"): + if layer_has_low_bit(block_keys, low_bit_keys): + layer_keys = [ + block_keys[name] + for name in ("down", "up", "alpha", "diff", "diff_b") + if name in block_keys + ] + layer_keys.extend( + key for key in passthrough_keys if key.startswith(f"{block_name}.") + ) + for key in layer_keys: + if key not in preserved_keys: + write_preserved_tensor(writer, key, handler) + preserved_keys.add(key) + pbar.update(1) + continue + direct_keys = [block_keys[name] for name in ("diff", "diff_b") if name in block_keys] + layer_source_keys = [ + block_keys[name] + for name in ("down", "up", "alpha", "diff", "diff_b") + if name in block_keys + ] + layer_source_dtypes = [handler.get_dtype(key) for key in layer_source_keys] + if direct_keys: + for key in direct_keys: + tensor = handler.get_tensor(key) + target_dtype = select_output_dtype( + layer_source_dtypes, + save_dtype, + is_1d_diff=tensor.ndim == 1, + ) + writer.write(key, tensor.to(target_dtype).cpu().contiguous()) + if "down" not in block_keys or "up" not in block_keys: + pbar.update(1) + continue + if "down" not in block_keys or "up" not in block_keys: pbar.update(1) continue @@ -470,20 +545,17 @@ def resize_lora_file( 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"] - - # Preserve original key format from input LoRA verbatim - new_block_name = block_name + layer_dtype = select_output_dtype( + layer_source_dtypes, + save_dtype, + ) 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), + block_keys["down"]: result["lora_down"].to(layer_dtype), + block_keys["up"]: result["lora_up"].to(layer_dtype), } - if alpha_suffix: - block_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype) + alpha_key = block_keys.get("alpha", f"{block_name}{block_keys['alpha_suffix']}") + block_sd[alpha_key] = torch.tensor(result["new_alpha"], dtype=layer_dtype) writer.write_dict(block_sd) @@ -495,6 +567,17 @@ def resize_lora_file( torch.cuda.empty_cache() pbar.update(1) + + for key in tqdm(passthrough_keys, desc="Copying auxiliary tensors", unit="layers"): + if key in preserved_keys: + pass + elif key in low_bit_keys: + write_preserved_tensor(writer, key, handler) + else: + tensor = handler.get_tensor(key) + target_dtype = select_output_dtype([handler.get_dtype(key)], save_dtype) + writer.write(key, tensor.to(target_dtype).cpu().contiguous()) + pbar.update(1) finally: writer.__exit__(None, None, None) @@ -689,6 +772,7 @@ def merge_loras_to_model( verbose: bool = True, lazy_load: bool = True, force_clear_cache: bool = False, + include_1d_diffs: bool = False, ) -> str: """ Merge multiple LoRAs into a base model and save the result directly. @@ -728,14 +812,25 @@ def merge_loras_to_model( try: + base_low_bit_keys = inspect_low_bit_input( + base_handler, f"Base model ({base_model_path})", "LoRA Merge To Model" + ) + lora_low_bit_keys = [ + inspect_low_bit_input( + handler, f"LoRA {i + 1} ({lora_paths[i]})", "LoRA Merge To Model" + ) + for i, handler in enumerate(lora_handlers) + ] # Detect format and extract pairs for each LoRA lora_infos = [] for i, handler in enumerate(lora_handlers): keys = handler.keys() format_info = detect_lora_format(keys) - pairs = extract_lora_pairs(keys, format_info) - network_dim, network_alpha = detect_lora_rank(handler, pairs) + pairs, _ = parse_lora_layers(keys) + network_dim, network_alpha = detect_lora_rank( + handler, pairs, lora_low_bit_keys[i] + ) lora_infos.append({ "handler": handler, "format_info": format_info, @@ -743,6 +838,7 @@ def merge_loras_to_model( "network_dim": network_dim, "network_alpha": network_alpha, "weight": lora_weights[i], + "low_bit_keys": lora_low_bit_keys[i], }) if verbose: print(f"[LoRA Merge To Model] LoRA {i+1}: {format_info['format']}, {len(pairs)} layers, dim={network_dim}") @@ -778,14 +874,47 @@ def merge_loras_to_model( return result.replace(".", "_") - # Build LoRA lookup: core layer name (underscored) -> list of (info, block_keys) + base_keys = list(base_handler.keys()) + base_aliases = set() + for key in base_keys: + normalized = key + for prefix in BASE_PREFIXES: + if normalized.startswith(prefix): + normalized = normalized[len(prefix):] + break + base_aliases.add(normalized) + base_aliases.add(normalized.replace(".", "_")) + + # Build lookups for low-rank layers and exact direct patches. lora_lookup = {} + direct_lookup = {} for info in lora_infos: for block_name, block_keys in info["pairs"].items(): - core = extract_core_layer_lora(block_name) - if core not in lora_lookup: - lora_lookup[core] = [] - lora_lookup[core].append((info, block_keys)) + if "down" in block_keys or "up" in block_keys: + core = extract_core_layer_lora(block_name) + lora_lookup.setdefault(core, []).append((info, block_keys)) + + for direct_name in ("diff", "diff_b"): + if direct_name not in block_keys: + continue + direct_key = block_keys[direct_name] + core = block_name + for prefix in LORA_PREFIXES: + if core.startswith(prefix): + core = core[len(prefix):] + break + if direct_name == "diff_b": + targets = [f"{core}.bias"] + else: + # Exact non-weight state keys win; module diffs fall back to .weight. + exact_aliases = {core, core.replace(".", "_")} + targets = [core] if exact_aliases & base_aliases else [f"{core}.weight"] + contribution = (info, direct_key, block_keys) + for target in targets: + direct_lookup.setdefault(target, []).append(contribution) + underscored_target = target.replace(".", "_") + if underscored_target != target: + direct_lookup.setdefault(underscored_target, []).append(contribution) # Compile skip patterns skip_patterns = _compile_patterns(skip_patterns_str) @@ -800,7 +929,6 @@ def merge_loras_to_model( 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)) if verbose: @@ -811,6 +939,11 @@ def merge_loras_to_model( try: with torch.no_grad(): for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"): + if base_key in base_low_bit_keys: + write_preserved_tensor(writer, base_key, base_handler) + stats["copied"] += 1 + pbar.update(1) + continue # Check skip patterns if _matches_any_pattern(base_key, skip_patterns): stats["skipped"] += 1 @@ -820,13 +953,21 @@ def merge_loras_to_model( # Load base weight cpu_base = base_handler.get_tensor(base_key) - # Only process weight tensors for LoRA merging + core_with_suffix = base_key + for prefix in BASE_PREFIXES: + if core_with_suffix.startswith(prefix): + core_with_suffix = core_with_suffix[len(prefix):] + break + direct_contributions = direct_lookup.get(core_with_suffix, []) + if not direct_contributions: + direct_contributions = direct_lookup.get(core_with_suffix.replace(".", "_"), []) + + low_rank_contributions = [] if base_key.endswith(".weight"): core = extract_core_layer_base(base_key) - core_underscored = core.replace(".", "_") + low_rank_contributions = lora_lookup.get(core.replace(".", "_"), []) - # Check if any LoRA contributes to this layer - if core_underscored in lora_lookup: + if direct_contributions or low_rank_contributions: # Transfer to GPU for computation if device == 'cuda': base_weight = transfer_to_gpu_pinned(cpu_base, device, torch.float32) @@ -834,72 +975,88 @@ def merge_loras_to_model( 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) + source_dtypes = [base_handler.get_dtype(base_key)] + applied_1d_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) + for info, direct_key, block_keys in direct_contributions: + if layer_has_low_bit(block_keys, info["low_bit_keys"]): + continue + cpu_diff = info["handler"].get_tensor(direct_key) + if cpu_diff.ndim == 1 and not include_1d_diffs: del cpu_diff - effective_scale = info["weight"] + continue + if tuple(cpu_diff.shape) != tuple(base_weight.shape): + print( + f"[LoRA Merge To Model] Shape mismatch for {direct_key}: " + f"{tuple(cpu_diff.shape)} != {tuple(base_weight.shape)}; skipped" + ) + del cpu_diff + continue + source_dtypes.append(info["handler"].get_dtype(direct_key)) + applied_1d_diff = applied_1d_diff or cpu_diff.ndim == 1 + if device == 'cuda': + delta = transfer_to_gpu_pinned(cpu_diff, device, torch.float32) else: - # Standard LoRA format - if "down" not in block_keys or "up" not in block_keys: - continue + delta = cpu_diff.to(device=device, dtype=torch.float32) + del cpu_diff + base_weight = base_weight + info["weight"] * delta + del delta - 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 + # Accumulate low-rank deltas from all contributing LoRAs. + for info, block_keys in low_rank_contributions: + if "down" not in block_keys or "up" not in block_keys: + continue + if layer_has_low_bit(block_keys, info["low_bit_keys"]): + continue + source_keys = [block_keys["down"], block_keys["up"]] + if "alpha" in block_keys: + source_keys.append(block_keys["alpha"]) + source_dtypes.extend(info["handler"].get_dtype(key) for key in source_keys) - # 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 + layer_rank = lora_down.shape[0] + if "alpha" in block_keys: + alpha_tensor = info["handler"].get_tensor(block_keys["alpha"]) + layer_alpha = float(alpha_tensor.item()) + else: + layer_alpha = float(layer_rank) + effective_scale = (layer_alpha / layer_rank if layer_rank > 0 else 1.0) * info["weight"] - # Apply delta to base weight + is_conv = len(lora_down.shape) == 4 + if is_conv: + in_rank, in_size, kernel_size, _ = lora_down.shape + out_size = lora_up.shape[0] + 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 base_weight = base_weight + effective_scale * delta del delta - # Write merged weight - writer.write(base_key, base_weight.to(save_dtype).cpu().contiguous()) + target_dtype = select_output_dtype( + source_dtypes, + save_dtype, + is_1d_diff=applied_1d_diff, + ) + writer.write(base_key, base_weight.to(target_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: - # Non-weight tensor (bias, norm, etc.), write as-is - writer.write(base_key, cpu_base.to(save_dtype).contiguous()) + target_dtype = select_output_dtype( + [base_handler.get_dtype(base_key)], + save_dtype, + ) + writer.write(base_key, cpu_base.to(target_dtype).contiguous()) del cpu_base stats["copied"] += 1 @@ -989,6 +1146,8 @@ class LoRAMergeToModel(io.ComfyNode): io.Combo.Input("device", options=["cuda", "cpu"], default="cuda"), io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer (slower but saves VRAM)"), + io.Boolean.Input("include_1d_diffs", default=False, + tooltip="Apply 1D direct-diff tensors as FP32. Disabled preserves prior behavior."), ], outputs=[io.String.Output(display_name="output_path")], is_output_node=True, @@ -998,7 +1157,8 @@ class LoRAMergeToModel(io.ComfyNode): 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, lazy_load, force_clear_cache) -> io.NodeOutput: + skip_patterns, output_filename, save_dtype, device, lazy_load, force_clear_cache, + include_1d_diffs) -> io.NodeOutput: # Build LoRA list based on count count = int(lora_count) @@ -1025,5 +1185,6 @@ class LoRAMergeToModel(io.ComfyNode): skip_patterns_str=skip_patterns, lazy_load=lazy_load, force_clear_cache=force_clear_cache, + include_1d_diffs=include_1d_diffs, ) return io.NodeOutput(path) diff --git a/nodes/merger.py b/nodes/merger.py index 2357424..ddd467a 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -12,6 +12,8 @@ from .device_utils import ( ) from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from .quantization_guard import inspect_low_bit_input, write_preserved_tensor +from .lora_resize import is_direct_diff_key def load_documentation_from_file(filename): @@ -105,6 +107,18 @@ class MergerLogic: all_keys = primary_handler.keys() metadata = primary_handler.metadata() + low_bit_keys_by_name = {} + if model_type == "loras": + try: + for name, handler in handlers.items(): + low_bit_keys_by_name[name] = inspect_low_bit_input( + handler, f"LoRA ({name})", "Generic LoRA Merge" + ) + except Exception: + for handler in handlers.values(): + handler.__exit__(None, None, None) + raise + # Convert mismatch_mode string to enum mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip') mismatch_mode = MissingTensorBehavior(mismatch_mode_str) @@ -133,6 +147,7 @@ class MergerLogic: save_dtype = recipe_params.pop('save_dtype') save_torch_dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}.get(save_dtype) override_dtype = recipe_params.pop('override_dtype', False) + include_1d_diffs = recipe_params.pop('include_1d_diffs', False) def determine_target_dtype(original_dtypes, target_dtype, override): if override: @@ -193,6 +208,14 @@ class MergerLogic: with torch.no_grad(): for key in tqdm(all_keys, desc="Merging layers", unit="layers"): + preserve_low_bit = model_type == "loras" and any( + key in low_bit_keys for low_bit_keys in low_bit_keys_by_name.values() + ) + if preserve_low_bit: + _flush_batch() + write_preserved_tensor(writer, key, primary_handler) + pbar.update(1) + continue # Check discard patterns first - skip entirely if _matches_any_pattern(key, discard_patterns, glob_mode=glob_mode): discarded_keys += 1 @@ -208,6 +231,13 @@ class MergerLogic: original_dtypes.append(handler.get_dtype(key)) target_dtype = determine_target_dtype(original_dtypes, save_torch_dtype, override_dtype) + preserve_model_a_1d = model_type == "loras" and cpu_tensor.ndim == 1 and not include_1d_diffs + if preserve_model_a_1d: + target_dtype = determine_target_dtype( + [primary_handler.get_dtype(key)], save_torch_dtype, override_dtype + ) + elif model_type == "loras" and is_direct_diff_key(key) and cpu_tensor.ndim == 1: + target_dtype = torch.float32 if process_device == 'cuda': tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype) @@ -216,13 +246,14 @@ class MergerLogic: del cpu_tensor # Check exclude patterns - use Model A only, no merge - if _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode): + if preserve_model_a_1d or _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode): t = tensor_a.detach().to(target_dtype).cpu() batch_buffer[key] = t batch_bytes += t.numel() * t.element_size() batch_key_count += 1 - excluded_keys += 1 + if not preserve_model_a_1d: + excluded_keys += 1 del tensor_a if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32): _flush_batch() @@ -549,7 +580,9 @@ class LoRATwoMerger(io.ComfyNode): "When False (default), patterns are Python regex matched as substrings."), io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), - io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("override_dtype", default=False, tooltip="Force merged non-1D tensors to save_dtype. Enabled 1D direct diffs remain FP32."), + io.Boolean.Input("include_1d_diffs", default=False, + tooltip="Merge 1D tensors. When disabled, preserve Model A's 1D tensors unchanged."), ], outputs=[ io.String.Output(display_name="output_filename"), @@ -562,7 +595,8 @@ class LoRATwoMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, + include_1d_diffs: bool) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -576,6 +610,7 @@ class LoRATwoMerger(io.ComfyNode): "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, + "include_1d_diffs": include_1d_diffs, } model_names = {"model_a": model_a, "model_b": model_b} filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE) @@ -885,7 +920,9 @@ class LoRAThreeMerger(io.ComfyNode): "When False (default), patterns are Python regex matched as substrings."), io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"), io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"), - io.Boolean.Input("override_dtype", default=False, tooltip="Force the entire model to be saved as the selected save_dtype. If False (default), higher precision dtypes are preserved."), + io.Boolean.Input("override_dtype", default=False, tooltip="Force merged non-1D tensors to save_dtype. Enabled 1D direct diffs remain FP32."), + io.Boolean.Input("include_1d_diffs", default=False, + tooltip="Merge 1D tensors. When disabled, preserve Model A's 1D tensors unchanged."), ], outputs=[ io.String.Output(display_name="output_filename"), @@ -898,7 +935,8 @@ class LoRAThreeMerger(io.ComfyNode): calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float, gamma: float, delta: float, epsilon: float, zeta: float, seed: int, output_filename: str, save_dtype: str, process_device: str, exclude_patterns: str, discard_patterns: str, - glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool) -> io.NodeOutput: + glob_patterns: bool, lazy_load: bool, force_clear_cache: bool, override_dtype: bool, + include_1d_diffs: bool) -> io.NodeOutput: doc = load_documentation_from_file('merger_3_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": return io.NodeOutput("Documentation mode active. No merge performed.", doc) @@ -912,6 +950,7 @@ class LoRAThreeMerger(io.ComfyNode): "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, "glob_patterns": glob_patterns, "lazy_load": lazy_load, "force_clear_cache": force_clear_cache, + "include_1d_diffs": include_1d_diffs, } model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE) @@ -982,4 +1021,4 @@ class EmbeddingThreeMerger(io.ComfyNode): } model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE) - return io.NodeOutput(filename, doc) \ No newline at end of file + return io.NodeOutput(filename, doc) diff --git a/nodes/quantization_guard.py b/nodes/quantization_guard.py new file mode 100644 index 0000000..082168a --- /dev/null +++ b/nodes/quantization_guard.py @@ -0,0 +1,141 @@ +"""Fail-safe detection for quantized or isolated low-bit safetensors inputs.""" + +import logging +import json +import struct + +import torch + + +MAX_ISOLATED_LOW_BIT_TENSORS = 2 + + +def _low_bit_dtypes(): + names = ( + "int8", + "uint8", + "float4_e2m1fn_x2", + "float8_e4m3fn", + "float8_e4m3fnuz", + "float8_e5m2", + "float8_e5m2fnuz", + "float8_e8m0fnu", + ) + return {getattr(torch, name) for name in names if hasattr(torch, name)} + + +LOW_BIT_DTYPES = _low_bit_dtypes() +LOW_BIT_STORAGE_CODES = {"I8", "U8", "F4"} +UPSTREAM_WRITER_STORAGE_CODES = { + "F64", "F32", "F16", "BF16", "I64", "I32", "I16", "I8", "U8", + "U64", "U32", "U16", "BOOL", "C64", "F8_E5M2", "F8_E4M3", +} + + +def _header(handler): + if getattr(handler, "_header", None) is not None: + return handler._header + with open(handler.filename, "rb") as source: + header_size = struct.unpack(" str | None: + entry = _header(handler).get(key) + return entry.get("dtype") if entry else None + + +def _is_low_bit(handler, key: str) -> bool: + storage_code = storage_dtype_code(handler, key) + if storage_code in LOW_BIT_STORAGE_CODES or ( + storage_code is not None and storage_code.startswith("F8_") + ): + return True + return handler.get_dtype(key) in LOW_BIT_DTYPES + + +def inspect_low_bit_input(handler, label: str, operation: str) -> set[str]: + """Reject quantized models; return isolated low-bit keys that must not be transformed.""" + keys = list(handler.keys()) + metadata = handler.metadata() or {} + quant_keys = [key for key in keys if key.endswith(".comfy_quant")] + has_quant_metadata = "_quantization_metadata" in metadata + has_legacy_scaled_fp8 = any(key.endswith("scaled_fp8") for key in keys) + if quant_keys or has_quant_metadata or has_legacy_scaled_fp8: + raise ValueError( + f"[{operation}] {label} uses ComfyUI quantization metadata; quantized inputs are unsupported" + ) + + low_bit_keys = {key for key in keys if _is_low_bit(handler, key)} + if len(low_bit_keys) > MAX_ISOLATED_LOW_BIT_TENSORS: + raise ValueError( + f"[{operation}] {label} contains {len(low_bit_keys)} low-bit tensors; " + "casted or quantized inputs are unsupported" + ) + if low_bit_keys: + logging.warning( + "[%s] %s contains %d isolated low-bit tensor(s); " + "preserving/excluding unchanged: %s", + operation, + label, + len(low_bit_keys), + ", ".join(sorted(low_bit_keys)), + ) + return low_bit_keys + + +def layer_has_low_bit(block_keys: dict, low_bit_keys: set[str]) -> bool: + return any( + block_keys.get(name) in low_bit_keys + for name in ("down", "up", "alpha", "diff", "diff_b") + ) + + +def write_preserved_tensor(writer, key: str, handler) -> None: + """Copy one tensor's original safetensors bytes and dtype into an active writer.""" + header = _header(handler) + entry = header[key] + if entry["dtype"] in UPSTREAM_WRITER_STORAGE_CODES: + writer.write(key, handler.get_tensor(key).cpu().contiguous()) + return + + source_start, source_end = entry["data_offsets"] + byte_size = source_end - source_start + + with writer._lock: + if key in writer._manifest: + raise ValueError(f"Tensor '{key}' has already been written.") + target_start = writer._current_data_offset + target_end = target_start + byte_size + writer._manifest[key] = { + "dtype": entry["dtype"], + "shape": entry["shape"], + "data_offsets": [target_start, target_end], + } + writer._current_data_offset = target_end + + source_path = handler.filename + source_absolute = 8 + getattr(handler, "_header_size", 0) + source_start + if not getattr(handler, "_header_size", None): + with open(source_path, "rb") as source: + source_absolute = 8 + struct.unpack("