nodes/alisson-anjos-ComfyUI-BFSNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

LoRA Surgery: detect the block axis from the file, and list every layer (1.44.0)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
alisson-anjos
2026-09-27 12:30:03 +00:00
co-authored by Claude Opus 5
parent 5073cf00ae
commit e72dd2569a
3 changed files with 116 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bfs_lora_surgery.py
+65 -12
View File
@@ -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:
pyproject.toml
+1 -1
View File
@@ -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",
web/js/bfs_lora_surgery.js
+50 -4
View File
@@ -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; } }, "->"),
])),
])),
]))),
]);
}