From e72dd2569a99e2012a5cd6abc59e00298eab075c Mon Sep 17 00:00:00 2001 From: alisson-anjos Date: Sun, 27 Sep 2026 12:30:03 +0000 Subject: [PATCH] LoRA Surgery: detect the block axis from the file, and list every layer (1.44.0) The parser took the last numeric component of a key as the block index. That is wrong whenever a name carries more than one number: transformer_blocks.12.attn1.to_q would be read as block 1 of a family ending in attn, and a nested layout such as down_blocks.0.attentions.1.transformer_blocks.0.ff.net.0.proj collapsed into a family called ff.net. The block axis is now chosen per path template as the numeric slot that varies most across the whole LoRA, which is information a single key does not have. Numbers outside that slot stay in the module name, so attn1 and attn2 stay distinct rather than folding together. The grouped tree is a summary, and a summary can hide things, so the panel now also has an "every layer" view: the full module list with a filter box, each row showing its norm and rank, and a + that writes a rule matching exactly that path. The structure endpoint returns that list alongside the grouped one. Checked against the body-swap recipe: same 136 of 192 modules kept, 56 dropped, 96 scaled as before the change, and attn.to_out.0 is still treated as attention rather than as block 0. Co-Authored-By: Claude Opus 5 --- bfs_lora_surgery.py | 77 ++++++++++++++++++++++++++++++++------ pyproject.toml | 2 +- web/js/bfs_lora_surgery.js | 54 ++++++++++++++++++++++++-- 3 files changed, 116 insertions(+), 17 deletions(-) diff --git a/bfs_lora_surgery.py b/bfs_lora_surgery.py index 22694a4..b4abe28 100644 --- a/bfs_lora_surgery.py +++ b/bfs_lora_surgery.py @@ -50,19 +50,64 @@ def _strip_suffix(key: str) -> str | None: return None -def _split_block(path: str) -> tuple[str, int | None, str]: - """Return (family, block_index, module_type) using the LAST numeric path component. +_NUM_RE = re.compile(r"\d+") - ``transformer_blocks.12.attn.to_q`` -> ("transformer_blocks", 12, "attn.to_q") - Keys with no numeric component come back with block ``None``. + +def _template(path: str) -> tuple[str, list[int], list[tuple[int, int]]]: + """``a.12.attn1.q`` -> ("a.{}.attn{}.q", [12, 1], spans).""" + nums, spans, out, last = [], [], [], 0 + for m in _NUM_RE.finditer(path): + out.append(path[last:m.start()]) + out.append("{}") + nums.append(int(m.group())) + spans.append((m.start(), m.end())) + last = m.end() + out.append(path[last:]) + return "".join(out), nums, spans + + +def choose_block_axis(paths: list[str]) -> dict[str, int]: + """For each path template, decide which numeric slot is the block index. + + Picking the last number breaks on names like ``transformer_blocks.12.attn1.to_q``, where + the last one is the ``1`` of ``attn1``, and on nested layouts such as + ``down_blocks.0.attentions.1.transformer_blocks.0.ff.net.0.proj``. The block axis is the + slot that varies most across the whole LoRA, so decide it from the file, not from one key. """ - matches = list(_BLOCK_RE.finditer(path)) - if not matches: + seen: dict[str, list[set]] = {} + for path in paths: + tpl, nums, _ = _template(path) + slots = seen.setdefault(tpl, [set() for _ in nums]) + for i, v in enumerate(nums): + if i < len(slots): + slots[i].add(v) + axis: dict[str, int] = {} + for tpl, slots in seen.items(): + if not slots: + continue + counts = [len(x) for x in slots] + best = max(range(len(counts)), key=lambda i: (counts[i], -i)) + axis[tpl] = best if counts[best] > 1 else -1 + return axis + + +def _split_block(path: str, axis: dict[str, int] | None = None) -> tuple[str, int | None, str]: + """Return (family, block_index, module_type). + + ``transformer_blocks.12.attn.to_q`` -> ("transformer_blocks", 12, "attn.to_q"). + Numbers that are not the block axis stay in the name, so ``attn1`` and ``attn2`` remain + distinct module types instead of being folded together. + """ + tpl, nums, spans = _template(path) + if not nums: return path, None, "" - m = matches[-1] - family = path[: m.start()].rstrip("._") - module = path[m.end():].lstrip("._") - return family, int(m.group(1)), module + idx = axis.get(tpl, -1) if axis is not None else (len(nums) - 1) + if idx < 0 or idx >= len(nums): + return path.rstrip("._"), None, "" + start, end = spans[idx] + family = path[:start].rstrip("._") + module = path[end:].lstrip("._") + return family, nums[idx], module def read_structure(lora_path: str, with_norms: bool = True) -> dict[str, Any]: @@ -90,11 +135,15 @@ def read_structure(lora_path: str, with_norms: bool = True) -> dict[str, Any]: norms: dict[str, float] = {} if with_norms and modules: norms = _module_norms(lora_path, modules) + axis = choose_block_axis(list(modules)) families: dict[str, dict[str, Any]] = {} loose: list[dict[str, Any]] = [] + all_modules: list[dict[str, Any]] = [] for base, entry in modules.items(): - family, block, mtype = _split_block(base) + family, block, mtype = _split_block(base, axis) + all_modules.append({"path": base, "family": family, "block": block, "type": mtype, + "rank": entry["rank"], "norm": norms.get(base)}) if block is None: loose.append({"path": base, "rank": entry["rank"], "norm": norms.get(base)}) continue @@ -125,11 +174,13 @@ def read_structure(lora_path: str, with_norms: bool = True) -> dict[str, Any]: }) out_families.sort(key=lambda f: -sum(t["count"] for t in f["types"])) + all_modules.sort(key=lambda m: (m["family"], m["block"] if m["block"] is not None else -1, m["type"])) return { "file": os.path.basename(lora_path), "total_modules": len(modules), "families": out_families, "loose": sorted(loose, key=lambda x: x["path"])[:64], + "modules": all_modules, "per_module_norms": norms, } @@ -294,6 +345,8 @@ def apply_rules(sd: dict[str, torch.Tensor], rules: list[dict[str, Any]]) -> tup """Return (new state dict, stats). Scale 0 drops the module entirely.""" if not rules: return sd, {"kept": None, "dropped": 0, "scaled": 0} + bases = {b for b in (_strip_suffix(k) for k in sd) if b} + axis = choose_block_axis(sorted(bases)) out: dict[str, torch.Tensor] = {} dropped = scaled = kept = 0 seen: dict[str, float] = {} @@ -303,7 +356,7 @@ def apply_rules(sd: dict[str, torch.Tensor], rules: list[dict[str, Any]]) -> tup out[key] = tensor continue if base not in seen: - family, block, mtype = _split_block(base) + family, block, mtype = _split_block(base, axis) seen[base] = resolve_scale(rules, family, block, mtype, full_path=base) s = seen[base] if s == 0.0: diff --git a/pyproject.toml b/pyproject.toml index 3e04bd9..9177438 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "bfsnodes" description = "Custom ComfyUI nodes for reference-guided video: the BFS Sampler (guide-driven IC-LoRAs with auto-configured crop, mask inpainting and temporal chunking in one node), LTX Multiple Controls and Identity Transfer (IPT2V reference + source-phase RoPE), LTXV Edit Anything, Head-Swap/Anime2Real Bernini conditioning, the AMV Guide Builder, MiniMax-H3 single-frame VAE loaders with tile control, and LoRA Surgery (read any LoRA's layers and blocks, then scale or drop them live to find which modules carry a defect, and save the result as a new LoRA)." -version = "1.43.2" +version = "1.44.0" license = {file = "LICENSE"} dependencies = [ "torch", diff --git a/web/js/bfs_lora_surgery.js b/web/js/bfs_lora_surgery.js index 5eead13..0214c15 100644 --- a/web/js/bfs_lora_surgery.js +++ b/web/js/bfs_lora_surgery.js @@ -51,6 +51,13 @@ function styles() { font-size:9px;display:flex;align-items:center;justify-content:center;cursor:pointer;color:#a8a8b2} .bfsls .blk.sel{background:#2b5fd9;border-color:#4a7ef0;color:#fff} .bfsls .err{color:#ff8a8a} +.bfsls .mod{display:grid;grid-template-columns:1fr 62px 34px 26px;gap:6px;align-items:center; + padding:3px 6px;border-bottom:1px solid #26262c;font-size:11px} +.bfsls .mod:hover{background:#26262c} +.bfsls .modlist{max-height:260px;overflow:auto;border:1px solid #33333c;border-radius:6px} +.bfsls .tab{padding:3px 10px;border-radius:5px 5px 0 0;border:1px solid #33333c;border-bottom:none; + background:#1f1f25;cursor:pointer;font-size:11px} +.bfsls .tab.on{background:#2b2b33;color:#fff} .bfsls .sec{border-top:1px solid #2c2c34;margin-top:8px;padding-top:6px} .bfsls details.help{border:1px solid #33333c;border-radius:6px;background:#1c1c22;margin-bottom:8px} .bfsls details.help>summary{cursor:pointer;padding:5px 8px;color:#9fb8ff;list-style:none;user-select:none} @@ -74,6 +81,8 @@ function Panel(props, expose) { const selBlocks = ref(new Set()); const scale = ref(0); const dragging = ref(false); + const view = ref("grouped"); // "grouped" | "all" + const filter = ref(""); let dragAdds = true; // Stored shape is {rules, ui}. A bare array is still accepted, both because older @@ -220,7 +229,7 @@ function Panel(props, expose) { ]), h("p", {}, [h("b", {}, "How to use it")]), h("ol", {}, [ - h("li", {}, "Pick a LoRA. The Layers list below shows the families, block counts and module types found in that specific file."), + h("li", {}, ["Pick a LoRA. Layers below shows what that file contains, grouped by type, or switch to ", h("code", {}, "every layer"), " to see and filter each module individually."]), h("li", {}, ["The bar next to each type is ", h("code", {}, "||dW||"), ", how large that group's update is. Big bars are where training invested."]), h("li", {}, ["Click ", h("code", {}, "->"), " on a type to select it, then drag across the block ruler to pick a range."]), @@ -331,8 +340,45 @@ function Panel(props, expose) { [h("button", { onClick: () => { rules.value = []; writeRules(); } }, "clear all")]) : h("div", { class: "muted" }, "No rules, the LoRA is applied unchanged."), - h("div", { class: "sec" }, [h("h4", {}, "Layers")]), - ...(struct.value?.families || []).map(fam => + h("div", { class: "sec" }, [ + h("div", { class: "row", style: "justify-content:space-between" }, [ + h("h4", {}, `Layers (${struct.value?.total_modules ?? 0} modules)`), + h("div", { class: "row", style: "gap:2px" }, [ + h("span", { class: "tab" + (view.value === "grouped" ? " on" : ""), + onClick: () => { view.value = "grouped"; } }, "grouped"), + h("span", { class: "tab" + (view.value === "all" ? " on" : ""), + onClick: () => { view.value = "all"; } }, "every layer"), + ]), + ]), + ]), + view.value === "all" + ? h("div", {}, [ + h("input", { type: "text", placeholder: "filter, e.g. gate_up or blocks.1 or attn", + value: filter.value, style: "margin-bottom:4px", + onInput: e => { filter.value = e.target.value; } }), + (() => { + const all = struct.value?.modules || []; + const q = filter.value.trim().toLowerCase(); + const hits = q ? all.filter(m => m.path.toLowerCase().includes(q)) : all; + const shown = hits.slice(0, 400); + return h("div", {}, [ + h("div", { class: "muted", style: "margin-bottom:3px" }, + `${hits.length} of ${all.length} modules${hits.length > shown.length ? `, showing first ${shown.length}` : ""}`), + h("div", { class: "modlist" }, shown.map(m => h("div", { class: "mod" }, [ + h("div", { class: "name", title: m.path }, m.path), + h("span", { class: "muted" }, m.norm != null ? m.norm.toFixed(3) : ""), + h("span", { class: "muted" }, `r${m.rank ?? "?"}`), + h("button", { title: "add a rule for exactly this module", + onClick: () => { + const esc = m.path.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); + rules.value.push({ enabled: true, match: { regex: `^${esc}$` }, scale: Number(scale.value) }); + writeRules(); + } }, "+"), + ]))), + ]); + })(), + ]) + : h("div", {}, (struct.value?.families || []).map(fam => h("details", { class: "fam", open: fam.name === selFamily.value }, [ h("summary", { onClick: () => { selFamily.value = fam.name; } }, [ h("span", { class: "name" }, fam.name || "(root)"), @@ -347,7 +393,7 @@ function Panel(props, expose) { h("span", { class: "muted" }, `r${t.rank ?? "?"}`), h("button", { title: "select this type", onClick: () => { selFamily.value = fam.name; selType.value = t.name; } }, "->"), ])), - ])), + ]))), ]); }