From e421b49575cf895ae3b765bb6b498e76dd20a826 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Fri, 4 Sep 2026 16:34:51 +0200 Subject: [PATCH] feat(lora): map alpha normalization keys --- nodes/lora_merger.py | 89 +---------------- nodes/lora_resize.py | 131 ++++++++++++++++++++++++- tests/test_lora_resize_factor_space.py | 44 +++++++++ 3 files changed, 175 insertions(+), 89 deletions(-) diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index 7614bbb..6c4f3ba 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -10,9 +10,10 @@ import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -from typing import List, Dict, Tuple, Optional +from typing import List, Optional from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation from .lora_resize import ( + build_lora_layer_map, canonical_lora_key, detect_lora_format, layer_has_companions, @@ -29,88 +30,6 @@ from .artifact_paths import canonical_model_artifact_path from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned -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 or "b_norm" in block_keys: - candidate = f"{core}.bias" - target = normalized_base.get(candidate) or normalized_base.get(candidate.replace(".", "_")) - elif any(role in block_keys for role in ("diff", "w_norm", "set_weight")): - 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 or "b_norm" in block_keys) and normalized_target.endswith(".bias"): - output_core = normalized_target[:-5] - elif any(role in block_keys for role in ("diff", "w_norm", "set_weight")) and normalized_target.endswith(".weight"): - output_core = normalized_target[:-7] - elif any(role in block_keys for role in ("diff", "w_norm", "set_weight")): - 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 any( - role in block_keys for role in ("diff", "diff_b", "w_norm", "b_norm", "set_weight") - ): - layer_map.setdefault(block_name, []).append((idx, block_name)) - return layer_map - - def _logical_output_dtypes(indices, lora_infos, handlers): pair_dtypes = [] direct_dtypes = {role: [] for role in _DIRECT_MERGE_ROLES} @@ -727,9 +646,9 @@ def _run_multi_lora_merge( f"format={info['format']['format']}" ) - layer_map = _build_layer_map(lora_infos, base_model_path) - validate_canonical_blocks(layer_map, operation) + layer_map = build_lora_layer_map(lora_infos, base_model_path) _ensure_guarded_layers_mapped(layer_map, lora_infos) + validate_canonical_blocks(layer_map, operation) units, stream_keys = _build_merge_work_units( layer_map, lora_infos, handlers, include_1d_diffs ) diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index d1f4c53..1655ff2 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -57,6 +57,14 @@ LORA_AUX_SUFFIXES = ( CANONICAL_DOWN_SUFFIX = ".lora_A.weight" CANONICAL_UP_SUFFIX = ".lora_B.weight" +BASE_MODEL_KEY_PREFIXES = ( + "model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net.", +) +LORA_BLOCK_PREFIXES = ( + "lora_unet_", "lora_transformer_", "lora_te1_", "lora_te2_", "lora_te_", + "lycoris_", "diffusion_model.", "transformer.", "unet.", +) + def canonical_lora_block_name(block_name: str) -> str: """Return the preferred generic ComfyUI diffusion-model LoRA root.""" @@ -110,6 +118,77 @@ def canonical_lora_key(block_name: str, role: str) -> str: return f"{root}{suffixes[role]}" +def _strip_lora_prefix(value: str, prefixes) -> str: + for prefix in prefixes: + if value.startswith(prefix): + return value[len(prefix):] + return value + + +def build_lora_layer_map( + lora_infos: List[Dict], base_model_path: Optional[str] +) -> Dict[str, List[Tuple[int, str]]]: + """Map parsed LoRA blocks to a reference model using existing merge rules.""" + 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_lora_prefix(base_key, BASE_MODEL_KEY_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 = _strip_lora_prefix(block_name, LORA_BLOCK_PREFIXES) + if core.endswith(".lora"): + core = core[:-5] + 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 or "b_norm" in block_keys: + candidate = f"{core}.bias" + target = normalized_base.get(candidate) or normalized_base.get(candidate.replace(".", "_")) + elif any(role in block_keys for role in ("diff", "w_norm", "set_weight")): + 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_lora_prefix(target, BASE_MODEL_KEY_PREFIXES) + if ("diff_b" in block_keys or "b_norm" in block_keys) and normalized_target.endswith(".bias"): + output_core = normalized_target[:-5] + elif any(role in block_keys for role in ("diff", "w_norm", "set_weight")) and normalized_target.endswith(".weight"): + output_core = normalized_target[:-7] + elif any(role in block_keys for role in ("diff", "w_norm", "set_weight")): + 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)) + + for idx, info in enumerate(lora_infos): + for block_name, block_keys in info["pairs"].items(): + if (idx, block_name) not in mapped and any( + role in block_keys for role in ("diff", "diff_b", "w_norm", "b_norm", "set_weight") + ): + layer_map.setdefault(block_name, []).append((idx, block_name)) + return layer_map + + def layer_tensor_keys(layer: Dict[str, str]) -> Dict[str, str]: roles = { "down", "up", "alpha", "mid", "reshape", "dora_scale", @@ -637,6 +716,7 @@ def normalize_lora_alpha_file( lora_path: str, output_filename: str, verbose: bool = True, + reference_model_path: Optional[str] = None, ) -> str: """Materialize LoRA alpha scaling and remove alpha tensors.""" handler = MemoryEfficientSafeOpen(lora_path, low_memory=True) @@ -647,6 +727,21 @@ def normalize_lora_alpha_file( metadata = (handler.metadata() or {}).copy() all_keys = handler.keys() pairs, _ = parse_lora_layers(all_keys) + output_keys = {} + mapped_blocks = {} + if reference_model_path: + layer_map = build_lora_layer_map( + [{"pairs": pairs}], reference_model_path + ) + validate_canonical_blocks(layer_map, "LoRA Alpha Normalize") + mapped_blocks = { + source_block: output_block + for output_block, sources in layer_map.items() + for _, source_block in sources + } + for block_name, output_block in mapped_blocks.items(): + for role, source_key in layer_tensor_keys(pairs[block_name]).items(): + output_keys[source_key] = canonical_lora_key(output_block, role) work_units = [] normalized_up_keys = set() @@ -697,7 +792,13 @@ def normalize_lora_alpha_file( for key in all_keys: if key in all_alpha_keys or key in normalized_up_keys: continue - write_preserved_tensor(writer, key, handler, force_raw=True) + write_preserved_tensor( + writer, + key, + handler, + output_key=output_keys.get(key), + force_raw=True, + ) with torch.no_grad(): iterator = stream_work_units( @@ -725,7 +826,10 @@ def normalize_lora_alpha_file( .to(dtype=up.dtype) .contiguous() ) - writer.write_batch([(up_key, output.cpu().contiguous())]) + writer.write_batch([( + output_keys.get(up_key, up_key), + output.cpu().contiguous(), + )]) finally: del output del up @@ -737,6 +841,10 @@ def normalize_lora_alpha_file( print( f"[LoRA Alpha Normalize] Removed {len(work_units)} alpha tensors" ) + if reference_model_path: + print( + f"[LoRA Alpha Normalize] Reference-mapped {len(mapped_blocks)} LoRA layers" + ) print(f"[LoRA Alpha Normalize] Saved to {output_path}") return output_path finally: @@ -1028,15 +1136,30 @@ class LoRANormalizeAlpha(io.ComfyNode): default="normalized_lora", tooltip="Output filename without extension, written under ComfyUI's LoRA directory.", ), + io.Combo.Input( + "reference_model", + options=["None", *folder_paths.get_filename_list("diffusion_models")], + default="None", + tooltip="Optional diffusion model used to normalize flattened LoRA layer names to its exact model paths.", + ), ], outputs=[io.AnyType.Output(display_name="output_path")], is_output_node=True, ) @classmethod - def execute(cls, lora_name, output_filename) -> io.NodeOutput: + def execute(cls, lora_name, output_filename, reference_model="None") -> io.NodeOutput: lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) - normalize_lora_alpha_file(lora_path, output_filename) + reference_path = None + if reference_model != "None": + reference_path = folder_paths.get_full_path_or_raise( + "diffusion_models", reference_model + ) + normalize_lora_alpha_file( + lora_path, + output_filename, + reference_model_path=reference_path, + ) _, output_name = canonical_model_artifact_path("loras", output_filename) return io.NodeOutput(output_name) diff --git a/tests/test_lora_resize_factor_space.py b/tests/test_lora_resize_factor_space.py index be11372..8eea70f 100644 --- a/tests/test_lora_resize_factor_space.py +++ b/tests/test_lora_resize_factor_space.py @@ -274,6 +274,50 @@ def test_normalize_alpha_rejects_alpha_free_input(monkeypatch, tmp_path, resize) ) +def test_normalize_alpha_maps_flattened_names_with_reference( + monkeypatch, tmp_path, resize +): + source = tmp_path / "flattened_alpha.safetensors" + reference = tmp_path / "reference.safetensors" + down = torch.eye(2, dtype=torch.float16) + up = torch.tensor([[2.0, 4.0], [6.0, 8.0]], dtype=torch.float16) + _write_uel( + source, + { + "lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight": down, + "lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight": up, + "lora_unet_double_blocks_0_img_attn_qkv.alpha": torch.tensor(1.0), + }, + ) + _write_uel( + reference, + {"diffusion_model.double_blocks.0.img_attn.qkv.weight": torch.eye(2)}, + ) + monkeypatch.setattr(resize.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(resize, "cleanup_after_operation", lambda: None) + + output = resize.normalize_lora_alpha_file( + str(source), + "reference_mapped_alpha", + verbose=False, + reference_model_path=str(reference), + ) + + output_tensors = _read_uel(output) + assert set(output_tensors) == { + "diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight", + "diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight", + } + assert torch.equal( + output_tensors["diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight"], + down, + ) + assert torch.equal( + output_tensors["diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight"], + up * 0.5, + ) + + def test_resize_does_not_invent_alpha_for_peft_pair(monkeypatch, tmp_path, resize): source = tmp_path / "peft_no_alpha.safetensors" _write_uel(