diff --git a/modules/lora/network_boft.py b/modules/lora/network_boft.py new file mode 100644 index 000000000..eccdb30d6 --- /dev/null +++ b/modules/lora/network_boft.py @@ -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) diff --git a/pipelines/flux/flux2_lora.py b/pipelines/flux/flux2_lora.py index 5689f9481..5148a8cfe 100644 --- a/pipelines/flux/flux2_lora.py +++ b/pipelines/flux/flux2_lora.py @@ -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)