447 lines
19 KiB
Python
447 lines
19 KiB
Python
"""Inference-time LoRA loading for the engine: any common LoRA file format ->
|
|
per-module low-rank deltas, applied either as forward hooks (dynamic strength,
|
|
removable) or merged into the weights (quantized weights re-quantized with
|
|
stochastic rounding so small deltas survive).
|
|
|
|
Independent of the training network classes on purpose.
|
|
|
|
Supported key formats (auto-detected per key):
|
|
peft / diffusers transformer.blocks.0.attn.to_q.lora_A.weight / lora_B
|
|
comfy diffusion_model.<...>.lora_A|lora_down (+ .alpha)
|
|
kohya / lycoris lora_unet_<flat_name>.lora_down.weight / lora_up + alpha,
|
|
lora_te_ / lora_te1_ / lora_te2_ for text encoders
|
|
full diffs <module>.diff (weight delta), <module>.diff_b (bias delta)
|
|
lokr <module>.lokr_w1 | lokr_w1_a/lokr_w1_b (+lokr_t1),
|
|
lokr_w2 | lokr_w2_a/lokr_w2_b (+lokr_t2), alpha;
|
|
toolkit peft-style dotted keys or old lycoris_<flat> keys
|
|
The holder's convert_lora_weights_before_load hook runs first, so archs with
|
|
their own conversions (anima, ltx2, hidream_o1) keep working.
|
|
"""
|
|
|
|
import os
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
|
|
_A_SUFFIXES = (".lora_A.weight", ".lora_down.weight", ".lora_A.default.weight")
|
|
_B_SUFFIXES = (".lora_B.weight", ".lora_up.weight", ".lora_B.default.weight")
|
|
_KRON_PARTS = ("lokr_w1", "lokr_w1_a", "lokr_w1_b", "lokr_t1", "lokr_w2", "lokr_w2_a", "lokr_w2_b", "lokr_t2")
|
|
_TE_PREFIXES = {
|
|
"lora_te1_": 0, "lora_te_": 0, "lora_te2_": 1, "lora_te3_": 2,
|
|
"text_encoder.": 0, "text_encoder_2.": 1, "text_encoder_3.": 2,
|
|
"te.": 0, "te1.": 0, "te2.": 1,
|
|
}
|
|
_DIT_PREFIXES = (
|
|
"model.diffusion_model.", "diffusion_model.", "transformer.", "unet.",
|
|
"lora_unet_", "lora_transformer_", "lycoris_", "base_model.model.",
|
|
)
|
|
|
|
|
|
def _cp_weight(t, wa, wb):
|
|
# tucker core t (r1, r2, kh, kw) with wa (r1, p), wb (r2, q) -> (p, q, kh, kw)
|
|
return torch.einsum("ijkl,ip,jr->prkl", t, wa, wb)
|
|
|
|
|
|
def _kron_factors(k: dict):
|
|
"""Rebuild the two kron factors (w1: (a, b), w2: (c, d[, kh, kw])) from stored parts."""
|
|
if "lokr_w1" in k:
|
|
w1 = k["lokr_w1"]
|
|
elif "lokr_t1" in k:
|
|
w1 = _cp_weight(k["lokr_t1"], k["lokr_w1_a"], k["lokr_w1_b"])
|
|
else:
|
|
w1 = k["lokr_w1_a"] @ k["lokr_w1_b"]
|
|
if "lokr_w2" in k:
|
|
w2 = k["lokr_w2"]
|
|
elif "lokr_t2" in k:
|
|
w2 = _cp_weight(k["lokr_t2"], k["lokr_w2_a"], k["lokr_w2_b"])
|
|
else:
|
|
w2 = k["lokr_w2_a"] @ k["lokr_w2_b"]
|
|
return w1, w2
|
|
|
|
|
|
def _kron_apply(x: torch.Tensor, k: dict) -> torch.Tensor:
|
|
"""x @ kron(w1, w2).T without materializing the full delta (linear only).
|
|
x (..., b*d) viewed as (..., b, d) -> (..., a, c) -> (..., a*c)."""
|
|
w1, w1a, w1b = k.get("lokr_w1"), k.get("lokr_w1_a"), k.get("lokr_w1_b")
|
|
w2, w2a, w2b = k.get("lokr_w2"), k.get("lokr_w2_a"), k.get("lokr_w2_b")
|
|
b = (w1 if w1 is not None else w1b).shape[1]
|
|
d = (w2 if w2 is not None else w2b).shape[1]
|
|
t = x.reshape(*x.shape[:-1], b, d)
|
|
if w2 is not None:
|
|
t = t @ w2.t()
|
|
else:
|
|
t = (t @ w2b.t()) @ w2a.t()
|
|
if w1 is not None:
|
|
t = torch.einsum("...bc,ab->...ac", t, w1)
|
|
else:
|
|
t = torch.einsum("...bc,rb->...rc", t, w1b)
|
|
t = torch.einsum("...rc,ar->...ac", t, w1a)
|
|
return t.reshape(*x.shape[:-1], -1)
|
|
|
|
|
|
class LoRAEntry:
|
|
def __init__(self, module: torch.nn.Module, name: str):
|
|
self.module = module
|
|
self.name = name
|
|
self.A: Optional[torch.Tensor] = None # (r, in)
|
|
self.B: Optional[torch.Tensor] = None # (out, r)
|
|
self.alpha: Optional[float] = None
|
|
self.diff: Optional[torch.Tensor] = None
|
|
self.diff_b: Optional[torch.Tensor] = None
|
|
self.kron: Optional[Dict[str, torch.Tensor]] = None # lokr parts by name
|
|
|
|
@property
|
|
def scale(self) -> float:
|
|
if self.alpha is None:
|
|
return 1.0
|
|
if self.A is not None:
|
|
r = self.A.shape[0]
|
|
return float(self.alpha) / r if r else 1.0
|
|
if self.kron is not None:
|
|
# rank lives in the _b factor; both factors full -> scale 1
|
|
rb = self.kron.get("lokr_w1_b", self.kron.get("lokr_w2_b"))
|
|
if rb is None:
|
|
return 1.0
|
|
return float(self.alpha) / rb.shape[0]
|
|
return 1.0
|
|
|
|
@property
|
|
def kron_is_linear(self) -> bool:
|
|
return self.kron is not None and all(t.dim() == 2 for t in self.kron.values())
|
|
|
|
def kron_shape(self) -> Optional[Tuple[int, int]]:
|
|
"""(out, in) of the rebuilt kron delta, linear parts only."""
|
|
if not self.kron_is_linear:
|
|
return None
|
|
k = self.kron
|
|
a = (k["lokr_w1"] if "lokr_w1" in k else k["lokr_w1_a"]).shape[0]
|
|
b = (k["lokr_w1"] if "lokr_w1" in k else k["lokr_w1_b"]).shape[1]
|
|
c = (k["lokr_w2"] if "lokr_w2" in k else k["lokr_w2_a"]).shape[0]
|
|
d = (k["lokr_w2"] if "lokr_w2" in k else k["lokr_w2_b"]).shape[1]
|
|
return a * c, b * d
|
|
|
|
def delta(self, strength: float, device, dtype=torch.float32) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
|
|
"""(weight delta, bias delta) at this strength, full precision."""
|
|
dw = None
|
|
if self.A is not None and self.B is not None:
|
|
A = self.A.to(device=device, dtype=dtype)
|
|
B = self.B.to(device=device, dtype=dtype)
|
|
if A.dim() == 4: # conv lora: (r, in, kh, kw) @ (out, r, 1, 1)
|
|
dw = torch.einsum("or,rikl->oikl", B.flatten(1), A) * (self.scale * strength)
|
|
else:
|
|
dw = (B @ A) * (self.scale * strength)
|
|
if self.kron is not None:
|
|
w1, w2 = _kron_factors({n: t.to(device=device, dtype=dtype) for n, t in self.kron.items()})
|
|
if w2.dim() == 4:
|
|
w1 = w1[:, :, None, None]
|
|
dk = torch.kron(w1, w2.contiguous()) * (self.scale * strength)
|
|
dw = dk if dw is None else dw + dk
|
|
if self.diff is not None:
|
|
d = self.diff.to(device=device, dtype=dtype) * strength
|
|
dw = d if dw is None else dw + d
|
|
db = self.diff_b.to(device=device, dtype=dtype) * strength if self.diff_b is not None else None
|
|
return dw, db
|
|
|
|
|
|
class InferenceLoRA:
|
|
def __init__(self, path: str, strength: float = 1.0, name: Optional[str] = None):
|
|
self.path = path
|
|
self.local_path: Optional[str] = None # set by load(): path resolved on disk
|
|
self.strength = float(strength)
|
|
self.name = name or os.path.splitext(os.path.basename(path))[0]
|
|
self.entries: List[LoRAEntry] = []
|
|
self.unmatched: List[str] = []
|
|
self._hooks: List = []
|
|
|
|
# ---- loading / key resolution ----
|
|
def load(self, holder, status_fn=None) -> "InferenceLoRA":
|
|
from toolkit.models.v2.resolver import resolve_lora_file
|
|
|
|
# a spec path may be a hub reference (org/repo/file.safetensors); it is
|
|
# searched for under the models folder before anything is downloaded.
|
|
# self.path stays as given: it is the spec identity the engine keys its
|
|
# loaded stack on.
|
|
self.local_path = resolve_lora_file(self.path, status_fn=status_fn)
|
|
sd = load_file(self.local_path)
|
|
convert = getattr(holder, "convert_lora_weights_before_load", None)
|
|
if callable(convert):
|
|
try:
|
|
sd = convert(sd)
|
|
except Exception:
|
|
pass
|
|
roots = self._roots(holder)
|
|
grouped: Dict[Tuple[int, str], dict] = {}
|
|
for key, tensor in sd.items():
|
|
base, part = self._split(key)
|
|
if part is None:
|
|
continue
|
|
grouped.setdefault(base, {})[part] = tensor
|
|
self.entries = []
|
|
self.unmatched = []
|
|
for base, parts in grouped.items():
|
|
module, name = self._resolve(base, roots)
|
|
if module is None:
|
|
self.unmatched.append(base)
|
|
continue
|
|
e = LoRAEntry(module, name)
|
|
e.A = parts.get("A")
|
|
e.B = parts.get("B")
|
|
e.diff = parts.get("diff")
|
|
e.diff_b = parts.get("diff_b")
|
|
kron = {n: parts[n] for n in _KRON_PARTS if n in parts}
|
|
if kron:
|
|
has_w1 = "lokr_w1" in kron or ("lokr_w1_a" in kron and "lokr_w1_b" in kron)
|
|
has_w2 = "lokr_w2" in kron or ("lokr_w2_a" in kron and "lokr_w2_b" in kron)
|
|
if not (has_w1 and has_w2):
|
|
self.unmatched.append(f"{base} (incomplete lokr parts {sorted(kron)})")
|
|
continue
|
|
e.kron = kron
|
|
alpha = parts.get("alpha")
|
|
if alpha is not None:
|
|
e.alpha = float(alpha.flatten()[0])
|
|
w_shape = getattr(getattr(module, "weight", None), "shape", None)
|
|
if e.A is not None and e.B is not None and e.A.dim() == 2:
|
|
if w_shape is not None and (e.B.shape[0] != w_shape[0] or e.A.shape[1] != w_shape[1]):
|
|
self.unmatched.append(f"{base} (shape {tuple(e.B.shape[0:1]) + tuple(e.A.shape[1:])} vs {tuple(w_shape)})")
|
|
continue
|
|
ks = e.kron_shape()
|
|
if ks is not None and w_shape is not None and len(w_shape) == 2 and tuple(ks) != tuple(w_shape):
|
|
self.unmatched.append(f"{base} (lokr shape {ks} vs {tuple(w_shape)})")
|
|
continue
|
|
if e.A is None and e.kron is None and e.diff is None and e.diff_b is None:
|
|
continue
|
|
self.entries.append(e)
|
|
return self
|
|
|
|
@staticmethod
|
|
def _split(key: str):
|
|
for suf in _A_SUFFIXES:
|
|
if key.endswith(suf):
|
|
return key[: -len(suf)], "A"
|
|
for suf in _B_SUFFIXES:
|
|
if key.endswith(suf):
|
|
return key[: -len(suf)], "B"
|
|
if key.endswith(".alpha"):
|
|
return key[: -len(".alpha")], "alpha"
|
|
if key.endswith(".diff"):
|
|
return key[: -len(".diff")], "diff"
|
|
if key.endswith(".diff_b"):
|
|
return key[: -len(".diff_b")], "diff_b"
|
|
for part in _KRON_PARTS:
|
|
if key.endswith("." + part):
|
|
return key[: -len(part) - 1], part
|
|
return key, None
|
|
|
|
@staticmethod
|
|
def _roots(holder) -> List[Tuple[str, torch.nn.Module, dict, dict]]:
|
|
"""(kind, root module, dotted-name map, flat-name map) for the DiT and TEs."""
|
|
roots = []
|
|
dit = getattr(holder, "model", None)
|
|
if isinstance(dit, torch.nn.Module):
|
|
names = dict(dit.named_modules())
|
|
roots.append(("dit", dit, names, {n.replace(".", "_"): n for n in names}))
|
|
tes = getattr(holder, "text_encoder", None)
|
|
tes = tes if isinstance(tes, list) else [tes]
|
|
for i, te in enumerate(tes):
|
|
if isinstance(te, torch.nn.Module) and not type(te).__name__.startswith("Fake"):
|
|
names = dict(te.named_modules())
|
|
roots.append((f"te{i}", te, names, {n.replace(".", "_"): n for n in names}))
|
|
return roots
|
|
|
|
@classmethod
|
|
def _resolve(cls, base: str, roots):
|
|
te_index = None
|
|
rest = base
|
|
for pref, idx in _TE_PREFIXES.items():
|
|
if base.startswith(pref):
|
|
te_index = idx
|
|
rest = base[len(pref):]
|
|
break
|
|
if te_index is None:
|
|
for pref in _DIT_PREFIXES:
|
|
if base.startswith(pref):
|
|
rest = base[len(pref):]
|
|
break
|
|
candidates = [r for r in roots if (r[0] == f"te{te_index}" if te_index is not None else r[0] == "dit")]
|
|
if not candidates:
|
|
candidates = roots
|
|
for _, root, names, flat in candidates:
|
|
if rest in names:
|
|
return names[rest], rest
|
|
flat_key = rest.replace(".", "_")
|
|
if flat_key in flat:
|
|
n = flat[flat_key]
|
|
return names[n], n
|
|
return None, None
|
|
|
|
# ---- hook mode ----
|
|
def attach(self):
|
|
"""Forward hooks that add strength * delta(x) to each module's output."""
|
|
self.detach()
|
|
for e in self.entries:
|
|
hookable = (e.A is not None and e.A.dim() == 2) or e.kron_is_linear \
|
|
or e.diff is not None or e.diff_b is not None
|
|
if not hookable:
|
|
continue
|
|
lora = self
|
|
|
|
def make(entry: LoRAEntry):
|
|
cache = {}
|
|
|
|
def hook(module, inputs, output):
|
|
if lora.strength == 0:
|
|
return output
|
|
x = inputs[0] if inputs else None
|
|
if x is None:
|
|
return output
|
|
key = (x.device, x.dtype)
|
|
if key not in cache:
|
|
cache.clear()
|
|
cache[key] = (
|
|
entry.A.to(device=x.device, dtype=x.dtype) if entry.A is not None and entry.A.dim() == 2 else None,
|
|
entry.B.to(device=x.device, dtype=x.dtype) if entry.B is not None and entry.B.dim() == 2 else None,
|
|
entry.diff.to(device=x.device, dtype=x.dtype) if entry.diff is not None and entry.diff.dim() == 2 else None,
|
|
entry.diff_b.to(device=x.device, dtype=x.dtype) if entry.diff_b is not None else None,
|
|
{n: t.to(device=x.device, dtype=x.dtype) for n, t in entry.kron.items()} if entry.kron_is_linear else None,
|
|
)
|
|
A, B, diff, diff_b, kron = cache[key]
|
|
add = None
|
|
if A is not None and B is not None:
|
|
add = ((x @ A.t()) @ B.t()) * (entry.scale * lora.strength)
|
|
if kron is not None:
|
|
dk = _kron_apply(x, kron) * (entry.scale * lora.strength)
|
|
add = dk if add is None else add + dk
|
|
if diff is not None:
|
|
d = (x @ diff.t()) * lora.strength
|
|
add = d if add is None else add + d
|
|
if diff_b is not None:
|
|
add = (diff_b * lora.strength) if add is None else add + diff_b * lora.strength
|
|
if add is None:
|
|
return output
|
|
if isinstance(output, tuple):
|
|
return (output[0] + add.to(output[0].dtype),) + tuple(output[1:])
|
|
return output + add.to(output.dtype)
|
|
|
|
return hook
|
|
|
|
self._hooks.append(e.module.register_forward_hook(make(e)))
|
|
return self
|
|
|
|
def detach(self):
|
|
for h in self._hooks:
|
|
h.remove()
|
|
self._hooks = []
|
|
|
|
# ---- merge mode ----
|
|
@torch.no_grad()
|
|
def merge(self, stochastic: bool = True) -> int:
|
|
"""Add strength * delta into each module's weight. Quantized (Ostris)
|
|
weights: dequantize -> add -> requantize on the stored grid with
|
|
stochastic rounding (int8 backends), nearest otherwise. Plain weights
|
|
in bf16/fp16 on cuda: stochastic rounding via copy_stochastic."""
|
|
from toolkit.optimizers.optimizer_utils import copy_stochastic
|
|
|
|
merged = 0
|
|
for e in self.entries:
|
|
m = e.module
|
|
w = getattr(m, "weight", None)
|
|
if w is None:
|
|
continue
|
|
device = w.device
|
|
dw, db = e.delta(self.strength, device)
|
|
if dw is not None:
|
|
if getattr(m, "is_ostris_quantized", False):
|
|
fp = m.dequantize_weight().to(torch.float32) + dw.reshape(w.shape)
|
|
q = m.ostris_quantizer
|
|
rq_codes = getattr(q, "requantize_codes_", None)
|
|
if rq_codes is not None and stochastic:
|
|
try:
|
|
rq_codes(m, fp, stochastic=True)
|
|
except TypeError:
|
|
rq_codes(m, fp)
|
|
elif rq_codes is not None:
|
|
rq_codes(m, fp)
|
|
else:
|
|
m.requantize_(fp)
|
|
elif isinstance(w, torch.nn.Parameter) or isinstance(w, torch.Tensor):
|
|
fp = w.data.to(torch.float32) + dw.reshape(w.shape)
|
|
if stochastic and w.is_cuda and w.dtype in (torch.bfloat16, torch.float16):
|
|
copy_stochastic(w.data, fp)
|
|
else:
|
|
w.data.copy_(fp.to(w.dtype))
|
|
else:
|
|
continue
|
|
if db is not None and getattr(m, "bias", None) is not None:
|
|
m.bias.data.add_(db.to(m.bias.dtype))
|
|
merged += 1
|
|
m_marker = getattr(self, "_marker", None)
|
|
return merged
|
|
|
|
def summary(self) -> dict:
|
|
return {
|
|
"name": self.name,
|
|
"path": self.path,
|
|
"strength": self.strength,
|
|
"modules": len(self.entries),
|
|
"unmatched": len(self.unmatched),
|
|
}
|
|
|
|
|
|
class LoRAStack:
|
|
"""The LoRAs applied to one holder, in one mode."""
|
|
|
|
def __init__(self, holder, mode: str = "hook"):
|
|
self.holder = holder
|
|
self.mode = mode
|
|
self.loras: List[InferenceLoRA] = []
|
|
|
|
def key(self) -> str:
|
|
return "|".join(f"{l.path}@{l.strength:.4f}" for l in self.loras) + f"#{self.mode}"
|
|
|
|
def load(self, specs: List[dict], status_fn=None):
|
|
for spec in specs:
|
|
lora = InferenceLoRA(spec["path"], spec.get("strength", 1.0), spec.get("name")).load(
|
|
self.holder, status_fn=status_fn
|
|
)
|
|
if status_fn:
|
|
status_fn(
|
|
f"LoRA {lora.name}: {len(lora.entries)} modules"
|
|
+ (f", {len(lora.unmatched)} unmatched" if lora.unmatched else "")
|
|
)
|
|
if not lora.entries:
|
|
raise ValueError(f"LoRA {lora.name} matched no modules of this model (unmatched: {lora.unmatched[:3]})")
|
|
self.loras.append(lora)
|
|
return self
|
|
|
|
def apply(self, status_fn=None):
|
|
touched = set()
|
|
for lora in self.loras:
|
|
if self.mode == "merge":
|
|
n = lora.merge(stochastic=True)
|
|
if status_fn:
|
|
status_fn(f"Merged LoRA {lora.name} into {n} modules (stochastic rounding)")
|
|
else:
|
|
lora.attach()
|
|
for e in lora.entries:
|
|
touched.add(id(e.module))
|
|
return touched
|
|
|
|
def remove(self):
|
|
"""Hook mode: detach every hook. Merge mode: cannot un-merge; the
|
|
engine reloads the affected components instead."""
|
|
for lora in self.loras:
|
|
lora.detach()
|
|
self.loras = []
|
|
|
|
def set_strengths(self, specs: List[dict]) -> bool:
|
|
"""Hook mode only: update strengths in place for the same paths."""
|
|
if self.mode != "hook" or len(specs) != len(self.loras):
|
|
return False
|
|
for lora, spec in zip(self.loras, specs):
|
|
if spec["path"] != lora.path:
|
|
return False
|
|
for lora, spec in zip(self.loras, specs):
|
|
lora.strength = float(spec.get("strength", 1.0))
|
|
return True
|