evan/ai-toolkit
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add lokr support to inference engine

Browse Source
This commit is contained in:
Jaret Burkett committed 2026-09-12 14:17:44 -06:00
1 parent ca097e63da
commit 737ee20733
1 file changed
+106 -10
Download Patch File Download Diff File Expand all files Collapse all files
toolkit/inference_lora.py
+106 -10
View File
@@ -11,6 +11,9 @@ Supported key formats (auto-detected per key):
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.
"""
@@ -23,6 +26,7 @@ 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,
@@ -30,10 +34,52 @@ _TE_PREFIXES = {
}
_DIT_PREFIXES = (
"model.diffusion_model.", "diffusion_model.", "transformer.", "unet.",
"lora_unet_", "lora_transformer_", "base_model.model.",
"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
@@ -43,13 +89,37 @@ class LoRAEntry:
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.A is None or self.alpha is None:
if self.alpha is None:
return 1.0
r = self.A.shape[0]
return float(self.alpha) / r if r else 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."""
@@ -61,6 +131,12 @@ class LoRAEntry:
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
@@ -107,15 +183,27 @@ class InferenceLoRA:
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:
w_shape = getattr(getattr(module, "weight", None), "shape", None)
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
if e.A is None and e.diff is None and e.diff_b is None:
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
@@ -134,6 +222,9 @@ class InferenceLoRA:
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
@@ -183,9 +274,10 @@ class InferenceLoRA:
"""Forward hooks that add strength * delta(x) to each module's output."""
self.detach()
for e in self.entries:
if e.A is None or e.A.dim() != 2:
if e.diff is None and e.diff_b is None:
continue
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):
@@ -205,11 +297,15 @@ class InferenceLoRA:
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 = cache[key]
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