nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(lora): enforce adapter tensor contracts

Browse Source
This commit is contained in:
silveroxides
2026-08-01 12:15:39 +02:00
parent f828ca7dfc
commit 49b66ea64e
7 changed files with 1747 additions and 423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/lora_extract_svd.py
+7 -2
View File
@@ -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}
nodes/lora_merger.py
+396 -244
View File
@@ -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)
nodes/lora_resize.py
+331 -170
View File
@@ -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)
nodes/merger.py
+46 -7
View File
@@ -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)
return io.NodeOutput(filename, doc)
nodes/quantization_guard.py
+141
View File
@@ -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("<Q", source.read(8))[0]
return json.loads(source.read(header_size).decode("utf-8"))
def storage_dtype_code(handler, key: str) -> 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("<Q", source.read(8))[0] + source_start
target_absolute = 8 + writer.max_header_bytes + target_start
def copy_raw():
try:
with open(source_path, "rb") as source:
source.seek(source_absolute)
remaining = byte_size
with writer._lock:
writer._file.seek(target_absolute)
while remaining:
chunk = source.read(min(8 * 1024 * 1024, remaining))
if not chunk:
raise EOFError(f"Unexpected EOF while preserving tensor '{key}'")
writer._file.write(chunk)
remaining -= len(chunk)
finally:
writer._semaphore.release()
writer._semaphore.acquire()
writer._futures.append(writer._executor.submit(copy_raw))
tests/test_lora_diff_dtype_contract.py
+436
View File
@@ -0,0 +1,436 @@
import importlib
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_contract_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def modules():
return _load_module("nodes.lora_resize"), _load_module("nodes.lora_merger")
@pytest.fixture(scope="module")
def generic_modules():
return _load_module("nodes.merger"), _load_module("nodes.merger_ops")
def _patch_output(monkeypatch, module, output_dir):
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(output_dir)])
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
def test_mixed_parser_scans_every_key_and_retains_unpaired(modules):
resize, _ = modules
filler = [f"metadata.tensor_{index}" for index in range(60)]
keys = filler + [
"diffusion_model.linear.lora_down.weight",
"diffusion_model.linear.lora_up.weight",
"diffusion_model.linear.alpha",
"diffusion_model.norm.diff",
"diffusion_model.block.diff_b",
"diffusion_model.orphan.lora_down.weight",
"diffusion_model.linear.dora_scale",
]
info = resize.detect_lora_format(keys)
layers, passthrough = resize.parse_lora_layers(keys)
assert info["format"] == "mixed"
assert layers["diffusion_model.linear"]["down"].endswith("lora_down.weight")
assert layers["diffusion_model.norm"]["diff"].endswith(".diff")
assert layers["diffusion_model.block"]["diff_b"].endswith(".diff_b")
assert "diffusion_model.orphan.lora_down.weight" in passthrough
assert "diffusion_model.linear.dora_scale" in passthrough
def test_dtype_contract(modules):
resize, _ = modules
assert resize.select_output_dtype([torch.bfloat16], torch.float16, is_1d_diff=True) == torch.float32
assert resize.select_output_dtype([torch.float32, torch.bfloat16], torch.float16) == torch.float32
assert resize.select_output_dtype([torch.float32], torch.bfloat16, force=True) == torch.bfloat16
def test_include_1d_switches_are_appended_and_default_off(modules, generic_modules):
resize, multi = modules
generic, _ = generic_modules
node_classes = [
generic.LoRATwoMerger,
generic.LoRAThreeMerger,
multi.LoRAMultiMerge,
multi.LoRAMultiMergeDARE,
multi.LoRAMultiMergeDAREEnhanced,
resize.LoRAMergeToModel,
]
for node_class in node_classes:
final_input = node_class.define_schema().inputs[-1]
assert final_input.id == "include_1d_diffs"
assert final_input.default is False
def test_resize_preserves_mixed_direct_auxiliary_and_fp32(monkeypatch, tmp_path, modules):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "mixed.safetensors"
save_file({
"diffusion_model.linear.lora_down.weight": torch.arange(8, dtype=torch.float32).reshape(2, 4),
"diffusion_model.linear.lora_up.weight": torch.arange(8, dtype=torch.float32).reshape(4, 2),
"diffusion_model.linear.alpha": torch.tensor(2.0, dtype=torch.float32),
"diffusion_model.norm.diff": torch.tensor([0.125, -0.25], dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(0.5, dtype=torch.bfloat16),
"diffusion_model.lin.diff": torch.full((2, 2), 1e-8, dtype=torch.float32),
"diffusion_model.linear.dora_scale": torch.ones(4, dtype=torch.float32),
"diffusion_model.orphan.lora_down.weight": torch.ones((1, 2), dtype=torch.bfloat16),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "resized", verbose=False
)
tensors = load_file(output)
assert tensors["diffusion_model.linear.lora_down.weight"].shape[0] == 1
assert tensors["diffusion_model.linear.lora_down.weight"].dtype == torch.float32
assert tensors["diffusion_model.linear.lora_up.weight"].dtype == torch.float32
assert tensors["diffusion_model.linear.alpha"].dtype == torch.float32
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.scale.diff"].dtype == torch.float16
assert tensors["diffusion_model.lin.diff"].dtype == torch.float32
assert tensors["diffusion_model.linear.dora_scale"].dtype == torch.float32
assert tensors["diffusion_model.orphan.lora_down.weight"].dtype == torch.float16
pure_source = tmp_path / "pure_diff.safetensors"
save_file({
"diffusion_model.pure.diff": torch.ones((2, 2), dtype=torch.bfloat16),
"diffusion_model.pure_norm.diff": torch.ones(2, dtype=torch.bfloat16),
}, str(pure_source))
pure_output = resize.resize_lora_file(
str(pure_source), 1, None, None, "cpu", torch.float16, "pure_resized", verbose=False
)
pure_tensors = load_file(pure_output)
assert set(pure_tensors) == {"diffusion_model.pure.diff", "diffusion_model.pure_norm.diff"}
assert pure_tensors["diffusion_model.pure.diff"].dtype == torch.float16
assert pure_tensors["diffusion_model.pure_norm.diff"].dtype == torch.float32
def test_standard_multi_merge_emits_weighted_direct_union(monkeypatch, tmp_path, modules):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "first.safetensors"
second = tmp_path / "second.safetensors"
base = tmp_path / "multi_base.safetensors"
save_file({
"model.diffusion_model.norm.weight": torch.zeros(2),
"model.diffusion_model.matrix.weight": torch.zeros((2, 2)),
"model.diffusion_model.only_first.weight": torch.zeros(()),
"model.diffusion_model.scale": torch.zeros(()),
"model.diffusion_model.lowrank.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.ones((2, 2), dtype=torch.float32),
"diffusion_model.only_first.diff": torch.tensor(4.0, dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(1.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.float32),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.float32),
"diffusion_model.lowrank.alpha": torch.tensor(1.0, dtype=torch.float32),
}, str(first))
save_file({
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.full((2, 2), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.full((1, 2), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_up.weight": torch.full((2, 1), 2.0, dtype=torch.bfloat16),
"diffusion_model.lowrank.alpha": torch.tensor(1.0, dtype=torch.bfloat16),
}, str(second))
output = merger.merge_multi_loras(
[str(first), str(second)], [2.0, -1.0], "concatenate", "cpu",
torch.float16, "merged", base_model_path=str(base), verbose=False,
include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["diffusion_model.norm.diff"], torch.tensor([-1.0, 0.0]))
torch.testing.assert_close(tensors["diffusion_model.matrix.diff"], torch.zeros((2, 2)))
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.matrix.diff"].dtype == torch.float32
assert tensors["diffusion_model.only_first.diff"].item() == pytest.approx(8.0)
assert tensors["diffusion_model.only_first.diff"].dtype == torch.float16
assert tensors["diffusion_model.scale.diff"].item() == pytest.approx(2.0)
assert tensors["diffusion_model.lowrank.lora_down.weight"].shape[0] == 2
assert tensors["diffusion_model.lowrank.lora_down.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_up.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.alpha"].dtype == torch.float32
default_output = merger.merge_multi_loras(
[str(first), str(second)], [2.0, -1.0], "concatenate", "cpu",
torch.float16, "merged_default", base_model_path=str(base), verbose=False,
)
default_tensors = load_file(default_output)
assert "diffusion_model.norm.diff" not in default_tensors
assert "diffusion_model.matrix.diff" in default_tensors
assert "diffusion_model.scale.diff" in default_tensors
assert "diffusion_model.lowrank.lora_down.weight" in default_tensors
def test_multi_merge_rejects_direct_shape_mismatch(monkeypatch, tmp_path, modules):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / "shape_first.safetensors"
second = tmp_path / "shape_second.safetensors"
save_file({"diffusion_model.bad.diff": torch.ones(2)}, str(first))
save_file({"diffusion_model.bad.diff": torch.ones(3)}, str(second))
with pytest.raises(ValueError, match="Direct LoRA shape mismatch"):
merger.merge_multi_loras(
[str(first), str(second)], [1.0, 1.0], "weighted_sum", "cpu",
torch.float16, "bad_shapes", verbose=False, include_1d_diffs=True,
)
@pytest.mark.parametrize("enhanced", [False, True])
def test_dare_variants_emit_direct_layers(monkeypatch, tmp_path, modules, enhanced):
_, merger = modules
_patch_output(monkeypatch, merger, tmp_path)
first = tmp_path / f"dare_first_{enhanced}.safetensors"
second = tmp_path / f"dare_second_{enhanced}.safetensors"
save_file({
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.float32),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.float32),
}, str(first))
save_file({
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.float32),
"diffusion_model.lowrank.lora_down.weight": torch.ones((1, 2), dtype=torch.bfloat16),
"diffusion_model.lowrank.lora_up.weight": torch.ones((2, 1), dtype=torch.bfloat16),
}, str(second))
if enhanced:
output = merger.merge_multi_loras_dare_enhanced(
[str(first), str(second)], [1.0, 1.0], 1.0, 1.0, 0.0, 0.0, 7,
"cpu", torch.float16, "enhanced", verbose=False, include_1d_diffs=True,
)
else:
output = merger.merge_multi_loras_dare(
[str(first), str(second)], [1.0, 1.0], 0.0, 0.0, 7,
"cpu", torch.float16, "dare", verbose=False, include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["diffusion_model.norm.diff"], torch.tensor([2.0, 3.0]))
assert tensors["diffusion_model.norm.diff"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_down.weight"].dtype == torch.float32
assert tensors["diffusion_model.lowrank.lora_up.weight"].dtype == torch.float32
if enhanced:
default_output = merger.merge_multi_loras_dare_enhanced(
[str(first), str(second)], [1.0, 1.0], 1.0, 1.0, 0.0, 0.0, 7,
"cpu", torch.float16, "enhanced_default", verbose=False,
)
else:
default_output = merger.merge_multi_loras_dare(
[str(first), str(second)], [1.0, 1.0], 0.0, 0.0, 7,
"cpu", torch.float16, "dare_default", verbose=False,
)
default_tensors = load_file(default_output)
assert "diffusion_model.norm.diff" not in default_tensors
assert "diffusion_model.lowrank.lora_down.weight" in default_tensors
def test_merge_to_model_applies_generic_direct_keys_and_preserves_dtype(
monkeypatch, tmp_path, modules, capsys
):
resize, _ = modules
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "base.safetensors"
adapter = tmp_path / "adapter.safetensors"
save_file({
"model.diffusion_model.block.weight": torch.zeros((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.norm.weight": torch.zeros(2, dtype=torch.bfloat16),
"model.diffusion_model.block.bias": torch.zeros(2, dtype=torch.bfloat16),
"model.diffusion_model.scale": torch.tensor(1.0, dtype=torch.bfloat16),
"model.diffusion_model.lin": torch.zeros((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.bad.weight": torch.ones((2, 2), dtype=torch.bfloat16),
"model.diffusion_model.untouched": torch.tensor(1e-8, dtype=torch.float32),
}, str(base))
save_file({
"diffusion_model.block.diff": torch.full((2, 2), 1e-8, dtype=torch.float32),
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.block.diff_b": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.scale.diff": torch.tensor(0.5, dtype=torch.bfloat16),
"diffusion_model.lin.diff": torch.ones((2, 2), dtype=torch.float32),
"diffusion_model.bad.diff": torch.ones(3, dtype=torch.float32),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16, "merged_model",
verbose=False, include_1d_diffs=True,
)
tensors = load_file(output)
torch.testing.assert_close(tensors["model.diffusion_model.norm.weight"], torch.tensor([0.5, 1.0]))
torch.testing.assert_close(tensors["model.diffusion_model.block.bias"], torch.tensor([1.5, 2.0]))
assert tensors["model.diffusion_model.scale"].item() == pytest.approx(1.25)
torch.testing.assert_close(tensors["model.diffusion_model.lin"], torch.full((2, 2), 0.5))
torch.testing.assert_close(tensors["model.diffusion_model.bad.weight"].float(), torch.ones((2, 2)))
assert "Shape mismatch" in capsys.readouterr().out
assert tensors["model.diffusion_model.block.weight"].dtype == torch.float32
assert tensors["model.diffusion_model.norm.weight"].dtype == torch.float32
assert tensors["model.diffusion_model.block.bias"].dtype == torch.float32
assert tensors["model.diffusion_model.scale"].dtype == torch.float16
assert tensors["model.diffusion_model.lin"].dtype == torch.float32
assert tensors["model.diffusion_model.untouched"].dtype == torch.float32
default_output = resize.merge_loras_to_model(
[str(adapter)], [0.5], str(base), "cpu", torch.float16, "merged_model_default",
verbose=False,
)
default_tensors = load_file(default_output)
torch.testing.assert_close(
default_tensors["model.diffusion_model.norm.weight"].float(), torch.zeros(2)
)
torch.testing.assert_close(
default_tensors["model.diffusion_model.block.bias"].float(), torch.zeros(2)
)
assert default_tensors["model.diffusion_model.scale"].item() == pytest.approx(1.25)
torch.testing.assert_close(
default_tensors["model.diffusion_model.lin"], torch.full((2, 2), 0.5)
)
def _generic_params(output_filename, override_dtype, include_1d_diffs=False):
return {
"mismatch_mode": "skip",
"alignment_mode": "pad/crop",
"alpha": 0.5,
"beta": 0.0,
"gamma": 0.5,
"delta": 2.0,
"epsilon": 0.01,
"zeta": 0.0,
"seed": 0,
"output_filename": output_filename,
"save_dtype": "bf16",
"override_dtype": override_dtype,
"device": "cpu",
"dtype": torch.float32,
"exclude_patterns": "",
"discard_patterns": "",
"glob_patterns": False,
"lazy_load": True,
"force_clear_cache": False,
"include_1d_diffs": include_1d_diffs,
}
def test_generic_two_and_three_mergers_enforce_direct_dtype_contract(
monkeypatch, tmp_path, generic_modules
):
generic, operations = generic_modules
paths = {}
for name, tensors in {
"a": {
"diffusion_model.norm.diff": torch.tensor([1.0, 2.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.ones((2, 2), dtype=torch.float32),
},
"b": {
"diffusion_model.norm.diff": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
"diffusion_model.matrix.diff": torch.full((2, 2), 3.0, dtype=torch.bfloat16),
"diffusion_model.secondary_only.diff": torch.ones(1, dtype=torch.float32),
},
"c": {
"diffusion_model.norm.diff": torch.tensor([0.5, 1.0], dtype=torch.float32),
"diffusion_model.matrix.diff": torch.full((2, 2), 0.5, dtype=torch.bfloat16),
},
}.items():
path = tmp_path / f"{name}.safetensors"
save_file(tensors, str(path))
paths[name] = str(path)
monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths.get(name))
monkeypatch.setattr(generic.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None)
two_params = _generic_params("generic_two", override_dtype=True)
two_params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"},
"Weight-Sum",
operations.TWO_MODEL_MODES,
two_params,
"loras",
)
two = load_file(str(tmp_path / "generic_two.safetensors"))
torch.testing.assert_close(
two["diffusion_model.norm.diff"].float(), torch.tensor([1.0, 2.0])
)
assert two["diffusion_model.norm.diff"].dtype == torch.bfloat16
assert two["diffusion_model.matrix.diff"].dtype == torch.bfloat16
assert "diffusion_model.secondary_only.diff" not in two
two_enabled_params = _generic_params(
"generic_two_enabled", override_dtype=True, include_1d_diffs=True
)
two_enabled_params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"},
"Weight-Sum",
operations.TWO_MODEL_MODES,
two_enabled_params,
"loras",
)
two_enabled = load_file(str(tmp_path / "generic_two_enabled.safetensors"))
torch.testing.assert_close(
two_enabled["diffusion_model.norm.diff"], torch.tensor([2.0, 3.0])
)
assert two_enabled["diffusion_model.norm.diff"].dtype == torch.float32
three_params = _generic_params("generic_three", override_dtype=False)
three_params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"},
"Add-Difference",
operations.THREE_MODEL_MODES,
three_params,
"loras",
)
three = load_file(str(tmp_path / "generic_three.safetensors"))
torch.testing.assert_close(
three["diffusion_model.norm.diff"].float(), torch.tensor([1.0, 2.0])
)
assert three["diffusion_model.norm.diff"].dtype == torch.bfloat16
assert three["diffusion_model.matrix.diff"].dtype == torch.float32
three_enabled_params = _generic_params(
"generic_three_enabled", override_dtype=False, include_1d_diffs=True
)
three_enabled_params.update({"model_a": "a", "model_b": "b", "model_c": "c"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b", "model_c": "c"},
"Add-Difference",
operations.THREE_MODEL_MODES,
three_enabled_params,
"loras",
)
three_enabled = load_file(str(tmp_path / "generic_three_enabled.safetensors"))
torch.testing.assert_close(
three_enabled["diffusion_model.norm.diff"], torch.tensor([2.25, 3.5])
)
assert three_enabled["diffusion_model.norm.diff"].dtype == torch.float32
tests/test_lora_quantization_guard.py
+390
View File
@@ -0,0 +1,390 @@
import importlib
import json
import logging
import struct
import sys
import types
from pathlib import Path
import pytest
import torch
from safetensors.torch import load_file, save_file
REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "modelutils_quant_guard_tests"
def _load_module(name):
if PACKAGE_NAME not in sys.modules:
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(REPO_ROOT)]
sys.modules[PACKAGE_NAME] = package
return importlib.import_module(f"{PACKAGE_NAME}.{name}")
@pytest.fixture(scope="module")
def modules():
return {
"guard": _load_module("nodes.quantization_guard"),
"extract": _load_module("nodes.lora_extract_svd"),
"resize": _load_module("nodes.lora_resize"),
"multi": _load_module("nodes.lora_merger"),
"generic": _load_module("nodes.merger"),
"operations": _load_module("nodes.merger_ops"),
}
def _patch_output(monkeypatch, module, output_dir):
monkeypatch.setattr(module.folder_paths, "get_folder_paths", lambda _: [str(output_dir)])
monkeypatch.setattr(module, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "cleanup_after_operation", lambda: None)
def _inspect_file(path, modules, label="input"):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
return modules["guard"].inspect_low_bit_input(handler, label, "Guard Test")
finally:
handler.__exit__(None, None, None)
def _round_trip_preserved(path, output, modules):
handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
try:
with modules["resize"].IncrementalSafetensorsWriter(str(output)) as writer:
modules["guard"].write_preserved_tensor(writer, "low", handler)
finally:
handler.__exit__(None, None, None)
return load_file(str(output))["low"]
def _tensor_payload(path, key="low"):
with open(path, "rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
header = json.loads(source.read(header_size).decode("utf-8"))
start, end = header[key]["data_offsets"]
source.seek(8 + header_size + start)
return source.read(end - start)
@pytest.mark.parametrize("dtype", [torch.int8, torch.uint8])
def test_int8_uint8_are_detected_and_warn_once(tmp_path, modules, caplog, dtype):
path = tmp_path / f"{dtype}.safetensors"
save_file({"low": torch.arange(4, dtype=dtype), "normal": torch.ones(1)}, str(path))
with caplog.at_level(logging.WARNING):
detected = _inspect_file(path, modules, "threshold input")
assert detected == {"low"}
records = [record for record in caplog.records if "threshold input" in record.message]
assert len(records) == 1
assert "1 isolated low-bit tensor" in records[0].message
assert "low" in records[0].message
def test_exact_two_tensor_threshold(tmp_path, modules):
for count in range(4):
path = tmp_path / f"threshold_{count}.safetensors"
tensors = {f"low_{i}": torch.ones(2, dtype=torch.uint8) for i in range(count)}
tensors["normal"] = torch.ones(1)
save_file(tensors, str(path))
if count <= 2:
assert len(_inspect_file(path, modules)) == count
else:
with pytest.raises(ValueError, match="3 low-bit tensors.*casted or quantized"):
_inspect_file(path, modules)
@pytest.mark.parametrize(
("tensors", "metadata"),
[
({"layer.comfy_quant": torch.tensor([123], dtype=torch.uint8)}, None),
({"normal": torch.ones(1)}, {"_quantization_metadata": "{}"}),
({"scaled_fp8": torch.ones(2)}, None),
],
)
def test_recognized_comfy_quantization_hard_errors(tmp_path, modules, tensors, metadata):
path = tmp_path / f"quant_{len(list(tmp_path.iterdir()))}.safetensors"
save_file(tensors, str(path), metadata=metadata)
with pytest.raises(ValueError, match="ComfyUI quantization metadata.*unsupported"):
_inspect_file(path, modules, "recognized input")
def test_supported_fp8_storage_round_trip(tmp_path, modules):
candidates = [
dtype
for name in (
"float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz",
"float8_e8m0fnu",
)
if (dtype := getattr(torch, name, None)) is not None
]
supported = []
for index, dtype in enumerate(candidates):
path = tmp_path / f"low_float_{index}.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
continue
supported.append(dtype)
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / f"preserved_{index}.safetensors", modules)
assert preserved.dtype == dtype
torch.testing.assert_close(preserved.float(), loaded.float())
if candidates and not supported:
pytest.skip("Installed safetensors cannot write any installed torch FP8 dtype")
def test_fp4_storage_when_supported(tmp_path, modules):
dtype = getattr(torch, "float4_e2m1fn_x2", None)
if dtype is None:
pytest.skip("Installed PyTorch has no torch.float4_e2m1fn_x2")
path = tmp_path / "fp4.safetensors"
try:
save_file({"low": torch.zeros(4, dtype=dtype)}, str(path))
loaded = load_file(str(path))["low"]
except (KeyError, RuntimeError, TypeError, ValueError):
pytest.skip("Installed safetensors cannot write torch.float4_e2m1fn_x2")
assert loaded.dtype == dtype
assert _inspect_file(path, modules) == {"low"}
preserved = _round_trip_preserved(path, tmp_path / "preserved_fp4.safetensors", modules)
assert preserved.dtype == dtype
assert _tensor_payload(path) == _tensor_payload(tmp_path / "preserved_fp4.safetensors")
def test_realistic_packed_uint8_comfy_sidecar_is_rejected(tmp_path, modules):
path = tmp_path / "packed_comfy.safetensors"
save_file({
"layer.weight": torch.arange(8, dtype=torch.uint8).reshape(2, 4),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
"layer.weight_scale": torch.ones(2),
}, str(path))
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
_inspect_file(path, modules)
def test_resize_preserves_complete_low_bit_layer_and_resizes_normal(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "resize_quant.safetensors"
guarded_down = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"diffusion_model.guarded.lora_down.weight": guarded_down,
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 2)),
"diffusion_model.guarded.alpha": torch.tensor(2.0),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.lora_down.weight": torch.ones((2, 2)),
"diffusion_model.normal.lora_up.weight": torch.ones((2, 2)),
"isolated_aux": torch.tensor([7], dtype=torch.int8),
}, str(source))
output = resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, "resize_guarded", verbose=False
)
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_down.weight"], guarded_down)
assert result["diffusion_model.guarded.lora_down.weight"].dtype == torch.uint8
assert result["diffusion_model.guarded.lora_up.weight"].dtype == torch.float32
assert result["diffusion_model.guarded.dora_scale"].dtype == torch.float32
assert result["diffusion_model.normal.lora_down.weight"].shape[0] == 1
assert result["isolated_aux"].dtype == torch.int8
def test_extract_excludes_low_bit_candidate_and_extracts_float_layer(
monkeypatch, tmp_path, modules
):
extract = modules["extract"]
monkeypatch.setattr(extract, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(extract, "cleanup_after_operation", lambda: None)
model_a = tmp_path / "extract_a.safetensors"
model_b = tmp_path / "extract_b.safetensors"
output = tmp_path / "extracted.safetensors"
save_file({
"model.diffusion_model.normal.weight": torch.eye(2),
"model.diffusion_model.quant.weight": torch.ones((2, 2), dtype=torch.uint8),
}, str(model_a))
save_file({
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
"model.diffusion_model.quant.weight": torch.zeros((2, 2), dtype=torch.uint8),
}, str(model_b))
extract.extract_lora_from_files(
str(model_a), str(model_b), "fixed", 1, 1, "cpu", "fp32", str(output),
force_clear_cache=False,
)
result = load_file(str(output))
assert any("normal" in key for key in result)
assert not any("quant" in key for key in result)
def _run_multi_variant(module, variant, paths, output_name):
common = (paths, [1.0, 1.0])
if variant == "standard":
return module.merge_multi_loras(
*common, "weighted_sum", "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
if variant == "dare":
return module.merge_multi_loras_dare(
*common, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
return module.merge_multi_loras_dare_enhanced(
*common, 1.0, 1.0, 0.0, 0.0, 11, "cpu", torch.float16, output_name,
verbose=False, include_1d_diffs=True,
)
@pytest.mark.parametrize("variant", ["standard", "dare", "enhanced"])
def test_multi_variants_preserve_earliest_complete_affected_layer(
monkeypatch, tmp_path, modules, variant
):
multi = modules["multi"]
_patch_output(monkeypatch, multi, tmp_path)
first = tmp_path / f"{variant}_first.safetensors"
second = tmp_path / f"{variant}_second.safetensors"
first_down = torch.tensor([[1.0, 2.0]])
save_file({
"diffusion_model.guarded.lora_down.weight": first_down,
"diffusion_model.guarded.lora_up.weight": torch.tensor([[3.0], [4.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([0.125, 0.25]),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(first))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 2]], dtype=torch.uint8),
"diffusion_model.guarded.lora_up.weight": torch.tensor([[9.0], [9.0]]),
"diffusion_model.guarded.dora_scale": torch.tensor([9.0, 9.0]),
"diffusion_model.normal.diff": torch.full((2, 2), 3.0),
}, str(second))
output = _run_multi_variant(multi, variant, [str(first), str(second)], f"guard_{variant}")
result = load_file(output)
torch.testing.assert_close(result["diffusion_model.guarded.lora_down.weight"], first_down)
torch.testing.assert_close(
result["diffusion_model.guarded.lora_up.weight"], torch.tensor([[3.0], [4.0]])
)
torch.testing.assert_close(
result["diffusion_model.guarded.dora_scale"], torch.tensor([0.125, 0.25])
)
assert "diffusion_model.normal.diff" in result
def test_hard_error_happens_before_resize_output_creation(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "three_low.safetensors"
save_file({f"low_{i}": torch.ones(1, dtype=torch.uint8) for i in range(3)}, str(source))
output = tmp_path / "must_not_exist.safetensors"
with pytest.raises(ValueError, match="3 low-bit tensors"):
resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
)
assert not output.exists()
def test_recognized_quantization_creates_no_resize_output(monkeypatch, tmp_path, modules):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
source = tmp_path / "recognized_quant.safetensors"
save_file({
"layer.weight": torch.ones((2, 2)),
"layer.comfy_quant": torch.tensor(list(b'{"format":"int8"}'), dtype=torch.uint8),
}, str(source))
output = tmp_path / "recognized_must_not_exist.safetensors"
with pytest.raises(ValueError, match="ComfyUI quantization metadata"):
resize.resize_lora_file(
str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
)
assert not output.exists()
def _generic_params(output_filename):
return {
"mismatch_mode": "skip", "alignment_mode": "pad/crop", "alpha": 0.5,
"beta": 0.0, "gamma": 0.5, "delta": 2.0, "epsilon": 0.01, "zeta": 0.0,
"seed": 0, "output_filename": output_filename, "save_dtype": "bf16",
"override_dtype": True, "device": "cpu", "dtype": torch.float32,
"exclude_patterns": "", "discard_patterns": "", "glob_patterns": False,
"lazy_load": True, "force_clear_cache": False, "include_1d_diffs": True,
}
def test_generic_merge_preserves_model_a_when_either_source_key_is_low_bit(
monkeypatch, tmp_path, modules
):
generic = modules["generic"]
operations = modules["operations"]
model_a = tmp_path / "generic_guard_a.safetensors"
model_b = tmp_path / "generic_guard_b.safetensors"
a_tensors = {
"a_low": torch.tensor([[1, 2]], dtype=torch.uint8),
"b_low": torch.tensor([[0.125, 0.25]], dtype=torch.float32),
"normal": torch.ones((1, 2)),
}
save_file(a_tensors, str(model_a))
save_file({
"a_low": torch.tensor([[9.0, 9.0]]),
"b_low": torch.tensor([[3, 4]], dtype=torch.int8),
"normal": torch.full((1, 2), 3.0),
}, str(model_b))
paths = {"a": str(model_a), "b": str(model_b)}
monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths[name])
monkeypatch.setattr(generic.folder_paths, "get_folder_paths", lambda _: [str(tmp_path)])
monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None)
monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None)
params = _generic_params("generic_guarded")
params.update({"model_a": "a", "model_b": "b"})
generic.MergerLogic.execute_merge(
{"model_a": "a", "model_b": "b"}, "Weight-Sum",
operations.TWO_MODEL_MODES, params, "loras",
)
result = load_file(str(tmp_path / "generic_guarded.safetensors"))
torch.testing.assert_close(result["a_low"], a_tensors["a_low"])
torch.testing.assert_close(result["b_low"], a_tensors["b_low"])
assert result["a_low"].dtype == torch.uint8
assert result["b_low"].dtype == torch.float32
torch.testing.assert_close(result["normal"].float(), torch.full((1, 2), 2.0))
def test_merge_to_model_preserves_low_bit_base_and_skips_low_bit_adapter_layer(
monkeypatch, tmp_path, modules
):
resize = modules["resize"]
_patch_output(monkeypatch, resize, tmp_path)
base = tmp_path / "merge_guard_base.safetensors"
adapter = tmp_path / "merge_guard_adapter.safetensors"
base_quant = torch.tensor([[1, 2], [3, 4]], dtype=torch.uint8)
save_file({
"model.diffusion_model.base_quant.weight": base_quant,
"model.diffusion_model.guarded.weight": torch.zeros((2, 2)),
"model.diffusion_model.normal.weight": torch.zeros((2, 2)),
}, str(base))
save_file({
"diffusion_model.guarded.lora_down.weight": torch.tensor([[1, 1]], dtype=torch.int8),
"diffusion_model.guarded.lora_up.weight": torch.ones((2, 1)),
"diffusion_model.guarded.diff": torch.full((2, 2), 5.0),
"diffusion_model.normal.diff": torch.ones((2, 2)),
}, str(adapter))
output = resize.merge_loras_to_model(
[str(adapter)], [1.0], str(base), "cpu", torch.float16, "merge_guarded",
verbose=False, include_1d_diffs=True,
)
result = load_file(output)
torch.testing.assert_close(result["model.diffusion_model.base_quant.weight"], base_quant)
assert result["model.diffusion_model.base_quant.weight"].dtype == torch.uint8
torch.testing.assert_close(
result["model.diffusion_model.guarded.weight"].float(), torch.zeros((2, 2))
)
torch.testing.assert_close(
result["model.diffusion_model.normal.weight"].float(), torch.ones((2, 2))
)