diff --git a/toolkit/inference_lora.py b/toolkit/inference_lora.py index 9abc324..f406ab2 100644 --- a/toolkit/inference_lora.py +++ b/toolkit/inference_lora.py @@ -11,6 +11,9 @@ Supported key formats (auto-detected per key): kohya / lycoris lora_unet_.lora_down.weight / lora_up + alpha, lora_te_ / lora_te1_ / lora_te2_ for text encoders full diffs .diff (weight delta), .diff_b (bias delta) + lokr .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_ 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