mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
feat(flux2): native BOFT (butterfly-OFT) support
Discriminated from OFT by 4-D oft_blocks shape (boft_m, block_num, block_size, block_size), mirroring LyCORIS algo_check. The cascade of log_2(block_num)+1 Cayley rotations is applied via the unflatten/transpose/flatten reshape sequence in NetworkModuleBOFT, porting LyCORIS make_weight verbatim.
This commit is contained in:
CalamitousFelicitousness
@@ -0,0 +1,119 @@
|
||||
"""BOFT (Butterfly-OFT) — cascade of butterfly orthogonal factors.
|
||||
|
||||
Saves with the same ``oft_blocks`` key as OFT but as a 4-D tensor
|
||||
``(boft_m, block_num, block_size, block_size)``. The caller in
|
||||
:func:`pipelines.flux.flux2_lora.try_load_oft` discriminates BOFT from
|
||||
OFT by ``oft_blocks.ndim == 4``. Math ported from
|
||||
``KohakuBlueleaf/LyCORIS/lycoris/modules/boft.py``.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import modules.lora.network as network
|
||||
|
||||
|
||||
class ModuleTypeBOFT(network.ModuleType):
|
||||
def create_module(self, net: network.Network, weights: network.NetworkWeights):
|
||||
ob = weights.w.get("oft_blocks")
|
||||
if ob is not None and ob.ndim == 4:
|
||||
return NetworkModuleBOFT(net, weights)
|
||||
return None
|
||||
|
||||
|
||||
class NetworkModuleBOFT(network.NetworkModule): # pylint: disable=abstract-method
|
||||
"""Butterfly-OFT module: cascade of orthogonal factors.
|
||||
|
||||
Constructor signature mirrors :class:`NetworkModuleOFT` so it slots into
|
||||
the same ``finalize_updown`` pipeline. The ``boft_m``/``block_num``/
|
||||
``block_size`` triple is read from the saved tensor's shape rather than
|
||||
re-derived via :func:`butterfly_factor`, which keeps loading deterministic
|
||||
even if the upstream factorization heuristic changes.
|
||||
"""
|
||||
|
||||
def __init__(self, net: network.Network, weights: network.NetworkWeights):
|
||||
super().__init__(net, weights)
|
||||
self.org_module: list[torch.nn.Module] = [self.sd_module]
|
||||
self.scale = 1.0
|
||||
|
||||
# 4-D oft_blocks: (boft_m, block_num, block_size, block_size)
|
||||
self.oft_blocks = weights.w["oft_blocks"]
|
||||
self.alpha = weights.w["alpha"]
|
||||
self.rescale = weights.w.get("rescale")
|
||||
self.boft_m = self.oft_blocks.shape[0]
|
||||
self.block_num = self.oft_blocks.shape[1]
|
||||
self.block_size = self.oft_blocks.shape[2]
|
||||
self.boft_b = self.block_size
|
||||
|
||||
# Resolve out_dim from the host module — matches NetworkModuleOFT's
|
||||
# discrimination so Linear/Conv2d hosts both work.
|
||||
is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear]
|
||||
is_conv = type(self.sd_module) in [torch.nn.Conv2d]
|
||||
if is_linear:
|
||||
self.out_dim = self.sd_module.out_features
|
||||
elif is_conv:
|
||||
self.out_dim = self.sd_module.out_channels
|
||||
else:
|
||||
self.out_dim = self.block_num * self.block_size
|
||||
|
||||
# constraint scales with out_dim per LyCORIS BOFT init
|
||||
self.constraint = float(self.alpha) * self.out_dim if self.alpha is not None else 0.0
|
||||
|
||||
def _get_r(self, target: torch.Tensor):
|
||||
"""Compute the per-stage Cayley rotations.
|
||||
|
||||
Returns a tensor of shape ``(boft_m, block_num, block_size, block_size)``
|
||||
where each ``r[i]`` is a stack of ``block_num`` orthogonal matrices
|
||||
derived from the i-th butterfly factor via Cayley's parameterization
|
||||
of SO(n): ``R = (I + Q)(I - Q)^-1`` for skew-symmetric ``Q``.
|
||||
"""
|
||||
eye = torch.eye(self.block_size, device=target.device, dtype=target.dtype)
|
||||
oft_blocks = self.oft_blocks.to(target.device, dtype=target.dtype)
|
||||
q = oft_blocks - oft_blocks.transpose(-1, -2)
|
||||
if self.constraint > 0:
|
||||
q_norm = torch.norm(q) + 1e-8
|
||||
if q_norm > self.constraint:
|
||||
q = q * self.constraint / q_norm
|
||||
# Inverse needs fp32 to be numerically well-behaved across all dtypes;
|
||||
# cast back to target dtype after.
|
||||
r = (eye + q) @ (eye - q).float().inverse().to(target.dtype)
|
||||
return r
|
||||
|
||||
def _make_weight(self, target: torch.Tensor):
|
||||
"""Apply the butterfly cascade to ``target`` and return the transformed weight.
|
||||
|
||||
Direct port of :meth:`ButterflyOFTModule.make_weight` (LyCORIS
|
||||
boft.py:158-191) for the merge-mode (no-bypass) path. ``target`` is the
|
||||
host weight; iteratively reshape to expose the per-stage block layout,
|
||||
einsum-multiply by the stage rotation, then reshape back. The reshape
|
||||
recipe at each stage is what makes the rotations interleave across
|
||||
butterfly partitions, giving the algorithm its O(d log d) parameter
|
||||
density.
|
||||
"""
|
||||
m = self.boft_m
|
||||
b = self.boft_b
|
||||
r_b = b // 2
|
||||
r = self._get_r(target)
|
||||
inp = target
|
||||
|
||||
for i in range(m):
|
||||
bi = r[i]
|
||||
g = 2
|
||||
k = 2 ** i * r_b
|
||||
inp = (
|
||||
inp.unflatten(0, (-1, g, k))
|
||||
.transpose(1, 2)
|
||||
.flatten(0, 2)
|
||||
.unflatten(0, (-1, b))
|
||||
)
|
||||
inp = torch.einsum("b i j, b j ... -> b i ...", bi, inp)
|
||||
inp = (
|
||||
inp.flatten(0, 1).unflatten(0, (-1, k, g)).transpose(1, 2).flatten(0, 2)
|
||||
)
|
||||
|
||||
if self.rescale is not None:
|
||||
inp = inp * self.rescale.to(target.device, dtype=target.dtype)
|
||||
return inp
|
||||
|
||||
def calc_updown(self, target: torch.Tensor):
|
||||
merged = self._make_weight(target)
|
||||
updown = merged - target
|
||||
return self.finalize_updown(updown, target, target.shape)
|
||||
@@ -32,6 +32,23 @@ Per-family fused-QKV handling:
|
||||
warning.
|
||||
- Norm: targets 1-D LayerNorm/RMSNorm parameters; never fused.
|
||||
|
||||
LyCORIS algorithm coverage relative to upstream
|
||||
``KohakuBlueleaf/LyCORIS/lycoris/modules/``:
|
||||
|
||||
- Native: LoRA, LoKR, LoHA, OFT, BOFT, IA3, GLoRA, Norm, Full.
|
||||
- Saved as standard LoRA: LoCon and DyLoRA. Both ``custom_state_dict``
|
||||
outputs collapse to ``lora_up.weight``/``lora_down.weight``/``alpha``
|
||||
(LoCon bakes its ``scalar`` into ``lora_up``; DyLoRA concats its
|
||||
per-block slabs into a max-rank matrix), so ``try_load_lora`` loads
|
||||
them losslessly relative to upstream's own export.
|
||||
- Deferred: TLoRA. The file saves only ``q_layer.weight`` /
|
||||
``p_layer.weight`` / ``lambda_layer`` / ``alpha``; the base SVD
|
||||
reference (``base_q`` / ``base_p`` / ``base_lambda``) that the
|
||||
delta math subtracts is unsaved by upstream design and the
|
||||
``sig_type`` selection mode is unrecoverable from the file, so
|
||||
any loader has a silent-correctness gap for ``sig_type != 'principal'``.
|
||||
Files fail cleanly with "not loaded".
|
||||
|
||||
Diffusers-PEFT fallback (used when ``lora_force_diffusers`` is on) is preserved
|
||||
via :func:`apply_patch`, which monkey-patches ``Flux2LoraLoaderMixin.lora_state_dict``
|
||||
to inject the ``diffusion_model.`` prefix for bare-BFL keys and bake kohya
|
||||
@@ -44,7 +61,7 @@ import torch
|
||||
from modules import shared, sd_models
|
||||
from modules.logger import log
|
||||
from modules.lora import (
|
||||
network, network_lora, network_lokr, network_hada, network_oft,
|
||||
network, network_lora, network_lokr, network_hada, network_oft, network_boft,
|
||||
network_ia3, network_glora, network_norm, network_full, lora_convert,
|
||||
)
|
||||
from modules.lora import lora_common as l
|
||||
@@ -542,15 +559,24 @@ def try_load_loha(name, network_on_disk, lora_scale):
|
||||
|
||||
|
||||
def try_load_oft(name, network_on_disk, lora_scale):
|
||||
"""Load a Flux2/Klein OFT (Orthogonal Fine-Tuning) adapter as native modules.
|
||||
"""Load a Flux2/Klein OFT or BOFT adapter as native modules.
|
||||
|
||||
Both kohya (``oft_blocks`` + alpha-as-constraint) and LyCORIS
|
||||
(``oft_diag``) layouts are recognized via :class:`NetworkModuleOFT`.
|
||||
Fused QKV in double_blocks is skipped with a warning: an OFT block
|
||||
structure is tied to the target module's ``out_features``, so a per-Q/K/V
|
||||
split would require re-deriving the rotation per chunk and is not a
|
||||
drop-in. Single-block linear1 (a single fused diffusers module) and all
|
||||
non-QKV double-block targets work fully.
|
||||
Both algorithms share the ``oft_blocks`` save key and are discriminated
|
||||
by tensor dimensionality, mirroring LyCORIS's own ``algo_check``:
|
||||
|
||||
- **OFT** — 3-D ``(num_blocks, block_size, block_size)``. Both kohya
|
||||
(``oft_blocks`` + alpha-as-constraint) and LyCORIS (``oft_diag``)
|
||||
layouts route through :class:`NetworkModuleOFT`.
|
||||
- **BOFT** — 4-D ``(boft_m, block_num, block_size, block_size)``,
|
||||
a cascade of butterfly factors. Routes through
|
||||
:class:`NetworkModuleBOFT` which ports the butterfly-cascade
|
||||
``make_weight`` from LyCORIS boft.py.
|
||||
|
||||
Fused QKV in double_blocks is skipped with a warning for both: an OFT
|
||||
block structure (and BOFT's per-stage block partition) is tied to the
|
||||
target module's ``out_features``, so a per-Q/K/V split would require
|
||||
re-deriving the rotations per chunk. Single-block ``linear1`` (a single
|
||||
fused diffusers module) and all non-QKV double-block targets work fully.
|
||||
"""
|
||||
t0 = time.time()
|
||||
state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
|
||||
@@ -566,9 +592,10 @@ def try_load_oft(name, network_on_disk, lora_scale):
|
||||
for (prefix, base), w in groups.items():
|
||||
if not ('oft_blocks' in w or 'oft_diag' in w):
|
||||
continue
|
||||
is_boft = 'oft_blocks' in w and w['oft_blocks'].ndim == 4
|
||||
targets = resolve_targets(prefix, base)
|
||||
if any(t[1] is not None for t in targets):
|
||||
log.warning(f'Network load: type=OFT name="{name}" key={base} fused QKV skipped (unsupported)')
|
||||
log.warning(f'Network load: type={"BOFT" if is_boft else "OFT"} name="{name}" key={base} fused QKV skipped (unsupported)')
|
||||
skipped += 1
|
||||
continue
|
||||
for diffusers_path, _, _ in targets:
|
||||
@@ -578,7 +605,10 @@ def try_load_oft(name, network_on_disk, lora_scale):
|
||||
unmapped += 1
|
||||
continue
|
||||
nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module)
|
||||
net.modules[network_key] = network_oft.NetworkModuleOFT(net, nw)
|
||||
if is_boft:
|
||||
net.modules[network_key] = network_boft.NetworkModuleBOFT(net, nw)
|
||||
else:
|
||||
net.modules[network_key] = network_oft.NetworkModuleOFT(net, nw)
|
||||
|
||||
return finalize_network(net, name, 'OFT', lora_scale, t0, unmapped=unmapped, skipped=skipped)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user