|
|
import logging |
|
|
from typing import Optional |
|
|
|
|
|
import torch |
|
|
import comfy.model_management |
|
|
from .base import ( |
|
|
WeightAdapterBase, |
|
|
WeightAdapterTrainBase, |
|
|
weight_decompose, |
|
|
factorization, |
|
|
) |
|
|
|
|
|
|
|
|
class LokrDiff(WeightAdapterTrainBase): |
|
|
def __init__(self, weights): |
|
|
super().__init__() |
|
|
(lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2, dora_scale) = weights |
|
|
self.use_tucker = False |
|
|
if lokr_w1_a is not None: |
|
|
_, rank_a = lokr_w1_a.shape[0], lokr_w1_a.shape[1] |
|
|
rank_a, _ = lokr_w1_b.shape[0], lokr_w1_b.shape[1] |
|
|
self.lokr_w1_a = torch.nn.Parameter(lokr_w1_a) |
|
|
self.lokr_w1_b = torch.nn.Parameter(lokr_w1_b) |
|
|
self.w1_rebuild = True |
|
|
self.ranka = rank_a |
|
|
|
|
|
if lokr_w2_a is not None: |
|
|
_, rank_b = lokr_w2_a.shape[0], lokr_w2_a.shape[1] |
|
|
rank_b, _ = lokr_w2_b.shape[0], lokr_w2_b.shape[1] |
|
|
self.lokr_w2_a = torch.nn.Parameter(lokr_w2_a) |
|
|
self.lokr_w2_b = torch.nn.Parameter(lokr_w2_b) |
|
|
if lokr_t2 is not None: |
|
|
self.use_tucker = True |
|
|
self.lokr_t2 = torch.nn.Parameter(lokr_t2) |
|
|
self.w2_rebuild = True |
|
|
self.rankb = rank_b |
|
|
|
|
|
if lokr_w1 is not None: |
|
|
self.lokr_w1 = torch.nn.Parameter(lokr_w1) |
|
|
self.w1_rebuild = False |
|
|
|
|
|
if lokr_w2 is not None: |
|
|
self.lokr_w2 = torch.nn.Parameter(lokr_w2) |
|
|
self.w2_rebuild = False |
|
|
|
|
|
self.alpha = torch.nn.Parameter(torch.tensor(alpha), requires_grad=False) |
|
|
|
|
|
@property |
|
|
def w1(self): |
|
|
if self.w1_rebuild: |
|
|
return (self.lokr_w1_a @ self.lokr_w1_b) * (self.alpha / self.ranka) |
|
|
else: |
|
|
return self.lokr_w1 |
|
|
|
|
|
@property |
|
|
def w2(self): |
|
|
if self.w2_rebuild: |
|
|
if self.use_tucker: |
|
|
w2 = torch.einsum( |
|
|
'i j k l, j r, i p -> p r k l', |
|
|
self.lokr_t2, |
|
|
self.lokr_w2_b, |
|
|
self.lokr_w2_a |
|
|
) |
|
|
else: |
|
|
w2 = self.lokr_w2_a @ self.lokr_w2_b |
|
|
return w2 * (self.alpha / self.rankb) |
|
|
else: |
|
|
return self.lokr_w2 |
|
|
|
|
|
def __call__(self, w): |
|
|
diff = torch.kron(self.w1, self.w2) |
|
|
return w + diff.reshape(w.shape).to(w) |
|
|
|
|
|
def passive_memory_usage(self): |
|
|
return sum(param.numel() * param.element_size() for param in self.parameters()) |
|
|
|
|
|
|
|
|
class LoKrAdapter(WeightAdapterBase): |
|
|
name = "lokr" |
|
|
|
|
|
def __init__(self, loaded_keys, weights): |
|
|
self.loaded_keys = loaded_keys |
|
|
self.weights = weights |
|
|
|
|
|
@classmethod |
|
|
def create_train(cls, weight, rank=1, alpha=1.0): |
|
|
out_dim = weight.shape[0] |
|
|
in_dim = weight.shape[1:].numel() |
|
|
out1, out2 = factorization(out_dim, rank) |
|
|
in1, in2 = factorization(in_dim, rank) |
|
|
mat1 = torch.empty(out1, in1, device=weight.device, dtype=weight.dtype) |
|
|
mat2 = torch.empty(out2, in2, device=weight.device, dtype=weight.dtype) |
|
|
torch.nn.init.kaiming_uniform_(mat2, a=5**0.5) |
|
|
torch.nn.init.constant_(mat1, 0.0) |
|
|
return LokrDiff( |
|
|
(mat1, mat2, alpha, None, None, None, None, None, None) |
|
|
) |
|
|
|
|
|
def to_train(self): |
|
|
return LokrDiff(self.weights) |
|
|
|
|
|
@classmethod |
|
|
def load( |
|
|
cls, |
|
|
x: str, |
|
|
lora: dict[str, torch.Tensor], |
|
|
alpha: float, |
|
|
dora_scale: torch.Tensor, |
|
|
loaded_keys: set[str] = None, |
|
|
) -> Optional["LoKrAdapter"]: |
|
|
if loaded_keys is None: |
|
|
loaded_keys = set() |
|
|
lokr_w1_name = "{}.lokr_w1".format(x) |
|
|
lokr_w2_name = "{}.lokr_w2".format(x) |
|
|
lokr_w1_a_name = "{}.lokr_w1_a".format(x) |
|
|
lokr_w1_b_name = "{}.lokr_w1_b".format(x) |
|
|
lokr_t2_name = "{}.lokr_t2".format(x) |
|
|
lokr_w2_a_name = "{}.lokr_w2_a".format(x) |
|
|
lokr_w2_b_name = "{}.lokr_w2_b".format(x) |
|
|
|
|
|
lokr_w1 = None |
|
|
if lokr_w1_name in lora.keys(): |
|
|
lokr_w1 = lora[lokr_w1_name] |
|
|
loaded_keys.add(lokr_w1_name) |
|
|
|
|
|
lokr_w2 = None |
|
|
if lokr_w2_name in lora.keys(): |
|
|
lokr_w2 = lora[lokr_w2_name] |
|
|
loaded_keys.add(lokr_w2_name) |
|
|
|
|
|
lokr_w1_a = None |
|
|
if lokr_w1_a_name in lora.keys(): |
|
|
lokr_w1_a = lora[lokr_w1_a_name] |
|
|
loaded_keys.add(lokr_w1_a_name) |
|
|
|
|
|
lokr_w1_b = None |
|
|
if lokr_w1_b_name in lora.keys(): |
|
|
lokr_w1_b = lora[lokr_w1_b_name] |
|
|
loaded_keys.add(lokr_w1_b_name) |
|
|
|
|
|
lokr_w2_a = None |
|
|
if lokr_w2_a_name in lora.keys(): |
|
|
lokr_w2_a = lora[lokr_w2_a_name] |
|
|
loaded_keys.add(lokr_w2_a_name) |
|
|
|
|
|
lokr_w2_b = None |
|
|
if lokr_w2_b_name in lora.keys(): |
|
|
lokr_w2_b = lora[lokr_w2_b_name] |
|
|
loaded_keys.add(lokr_w2_b_name) |
|
|
|
|
|
lokr_t2 = None |
|
|
if lokr_t2_name in lora.keys(): |
|
|
lokr_t2 = lora[lokr_t2_name] |
|
|
loaded_keys.add(lokr_t2_name) |
|
|
|
|
|
if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None): |
|
|
weights = (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2, dora_scale) |
|
|
return cls(loaded_keys, weights) |
|
|
else: |
|
|
return None |
|
|
|
|
|
def calculate_weight( |
|
|
self, |
|
|
weight, |
|
|
key, |
|
|
strength, |
|
|
strength_model, |
|
|
offset, |
|
|
function, |
|
|
intermediate_dtype=torch.float32, |
|
|
original_weight=None, |
|
|
): |
|
|
v = self.weights |
|
|
w1 = v[0] |
|
|
w2 = v[1] |
|
|
w1_a = v[3] |
|
|
w1_b = v[4] |
|
|
w2_a = v[5] |
|
|
w2_b = v[6] |
|
|
t2 = v[7] |
|
|
dora_scale = v[8] |
|
|
dim = None |
|
|
|
|
|
if w1 is None: |
|
|
dim = w1_b.shape[0] |
|
|
w1 = torch.mm(comfy.model_management.cast_to_device(w1_a, weight.device, intermediate_dtype), |
|
|
comfy.model_management.cast_to_device(w1_b, weight.device, intermediate_dtype)) |
|
|
else: |
|
|
w1 = comfy.model_management.cast_to_device(w1, weight.device, intermediate_dtype) |
|
|
|
|
|
if w2 is None: |
|
|
dim = w2_b.shape[0] |
|
|
if t2 is None: |
|
|
w2 = torch.mm(comfy.model_management.cast_to_device(w2_a, weight.device, intermediate_dtype), |
|
|
comfy.model_management.cast_to_device(w2_b, weight.device, intermediate_dtype)) |
|
|
else: |
|
|
w2 = torch.einsum('i j k l, j r, i p -> p r k l', |
|
|
comfy.model_management.cast_to_device(t2, weight.device, intermediate_dtype), |
|
|
comfy.model_management.cast_to_device(w2_b, weight.device, intermediate_dtype), |
|
|
comfy.model_management.cast_to_device(w2_a, weight.device, intermediate_dtype)) |
|
|
else: |
|
|
w2 = comfy.model_management.cast_to_device(w2, weight.device, intermediate_dtype) |
|
|
|
|
|
if len(w2.shape) == 4: |
|
|
w1 = w1.unsqueeze(2).unsqueeze(2) |
|
|
if v[2] is not None and dim is not None: |
|
|
alpha = v[2] / dim |
|
|
else: |
|
|
alpha = 1.0 |
|
|
|
|
|
try: |
|
|
lora_diff = torch.kron(w1, w2).reshape(weight.shape) |
|
|
if dora_scale is not None: |
|
|
weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength, intermediate_dtype, function) |
|
|
else: |
|
|
weight += function(((strength * alpha) * lora_diff).type(weight.dtype)) |
|
|
except Exception as e: |
|
|
logging.error("ERROR {} {} {}".format(self.name, key, e)) |
|
|
return weight |
|
|
|
