Balance Attention

Here's a subtle problem that ruins merges nobody can quite explain: two models can have exactly the same architecture and still arrange their attention heads in incompatible orders. Fine-tunes reorder heads all the time, and when you add two tensors whose heads don't correspond, you're not merging - you're averaging misaligned noise and hoping it averages out. Balance Attention fixes the alignment before the addition happens.

How it works

This is the pack's implementation of weight-balancing via permutation. Before merging, the node finds the optimal permutation of each attention head in one model that best aligns it with the other, so that after permuting, corresponding heads actually correspond. It's the same family of ideas as the rebasin/SD-Merge techniques that the paper-following part of the community uses to fix merges where naive averaging produces a mush. Concretely the implementation here is registered for the SDXL (sdxl-sgm) layout - it extracts the attention projections (Q, K, V, O) from a recipe, computes alignment, and permutes before any downstream merge node runs.

Input:

a (delta|param|weight) - the recipe whose attention is being balanced. Usually one of the two models in a planned merge.

Output: one recipe (MECHA_RECIPE) with permuted attention keys, ready to feed into a Weighted Sum or Add Difference alongside the other model.

It's lazy and streaming like the rest of the pack - no full load, no RAM spike - and it exposes the standard merge_checkpointing toggle to cache the permutation work on CPU in fp16 when the input branch is stable.

When it's worth the trouble

The honest take: for two checkpoints from the same base, most merges are fine without it, and you can get a long way on plain weighted sums. But the failure mode it fixes is real and specific - a merge where the two parents look similar in isolation but blend into a washed-out, featureless result that no amount of alpha fiddling rescues. If you've ever had a merge that's "just wrong" despite sane weights, permutation alignment is one of the few genuinely principled fixes, and this node is the easiest way to get at it without writing permutation code yourself. The research community's rebasin-style merging is one of the better-grounded "free lunch" ideas in this hobby - the papers it comes from show permuted merges beating plain averaging on the same models.

Install & gotchas

ComfyUI Manager → search mecha → "Mecha Merge Node Pack", or:

cd ComfyUI/custom_nodes
git clone https://github.com/ljleb/comfy-mecha.git
pip install -r comfy-mecha/requirements.txt

Restart after. Dependency is sd-mecha==1.1.7.

The catch is architecture scope: balancing is implemented for the SDXL attention layout, so don't expect it to align SD1.5 or Flux models - those will fall back to pass-through behavior. And permuting attention costs compute; with merge_checkpointing off, every workflow execution re-runs it, so flip that toggle on once your recipe is stable.

Name	Type	Default	Description
a (delta\|param\|weight)	MECHA_RECIPE		—
merge_checkpointingopt	BOOLEAN	false	Speeds up an entire branch of a merge graph that does not change often in exchange of memory. - true: store the first output of this recipe node on cpu memory in fp16. On subsequent workflow executions, as long as the inputs do not change, the cached keys are returned after being cast to the original device and dtype. - false: do not store the output. The recipe and its inputs will re-execute on subsequent workflow executions. Note that the memory used to checkpoint the output is distinct from the cache feature. In general, you probably want to either use this or a cache unit, but not both at the same time because the memory adds up. The difference between merge checkpointing and cache is that merge checkpointing completely re-merges from scratch if any input changes. Merge checkpointing is also generally much faster than cache in the fast path.

How it works

When it's worth the trouble

Install & gotchas

Inputs (2)

Outputs (1)