Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor(lora): extract the lokr operand rebuild

Browse Source 
The seventeen lines that rebuild w1 and w2 from whatever the file stored
were copied into all three lokr variants, character for character. They
move to the base class; each variant keeps only the part that differs,
which is how it addresses the product.

The base class keeps its conv branch, which the two chunk variants
deliberately lack: those address 2-d fused weights.
This commit is contained in:
CalamitousFelicitousness
2026-08-30 05:10:34 +01:00
parent 2ae31ce46f
commit 5538b19390
1 changed files with 9 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/lora/network_lokr.py
+9 -36
View File
@@ -32,7 +32,8 @@ class NetworkModuleLokr(network.NetworkModule): # pylint: disable=abstract-metho
self.dim = self.w2b.shape[0] if self.w2b is not None else self.dim
self.t2 = weights.w.get("lokr_t2")
def calc_updown(self, target):
def rebuild_operands(self, target):
"""The two Kronecker operands on the target's device and dtype, each either stored whole or rebuilt from its factors."""
if self.w1 is not None:
w1 = self.w1.to(target.device, dtype=target.dtype)
else:
@@ -50,8 +51,12 @@ class NetworkModuleLokr(network.NetworkModule): # pylint: disable=abstract-metho
w2a = self.w2a.to(target.device, dtype=target.dtype)
w2b = self.w2b.to(target.device, dtype=target.dtype)
w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b)
return w1, w2
def calc_updown(self, target):
w1, w2 = self.rebuild_operands(target)
output_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)]
if len(target.shape) == 4:
if len(target.shape) == 4: # a conv target keeps its own shape; the chunk variants below only ever address 2-d fused weights
output_shape = target.shape
updown = make_kron(output_shape, w1, w2)
return self.finalize_updown(updown, target, output_shape)
@@ -70,23 +75,7 @@ class NetworkModuleLokrChunk(NetworkModuleLokr):
self.num_chunks = num_chunks
def calc_updown(self, target):
if self.w1 is not None:
w1 = self.w1.to(target.device, dtype=target.dtype)
else:
w1a = self.w1a.to(target.device, dtype=target.dtype)
w1b = self.w1b.to(target.device, dtype=target.dtype)
w1 = w1a @ w1b
if self.w2 is not None:
w2 = self.w2.to(target.device, dtype=target.dtype)
elif self.t2 is None:
w2a = self.w2a.to(target.device, dtype=target.dtype)
w2b = self.w2b.to(target.device, dtype=target.dtype)
w2 = w2a @ w2b
else:
t2 = self.t2.to(target.device, dtype=target.dtype)
w2a = self.w2a.to(target.device, dtype=target.dtype)
w2b = self.w2b.to(target.device, dtype=target.dtype)
w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b)
w1, w2 = self.rebuild_operands(target)
full_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)]
updown = make_kron(full_shape, w1, w2)
updown = torch.chunk(updown, self.num_chunks, dim=0)[self.chunk_index]
@@ -109,23 +98,7 @@ class NetworkModuleLokrSliceChunk(NetworkModuleLokr):
self.end_row = end_row
def calc_updown(self, target):
if self.w1 is not None:
w1 = self.w1.to(target.device, dtype=target.dtype)
else:
w1a = self.w1a.to(target.device, dtype=target.dtype)
w1b = self.w1b.to(target.device, dtype=target.dtype)
w1 = w1a @ w1b
if self.w2 is not None:
w2 = self.w2.to(target.device, dtype=target.dtype)
elif self.t2 is None:
w2a = self.w2a.to(target.device, dtype=target.dtype)
w2b = self.w2b.to(target.device, dtype=target.dtype)
w2 = w2a @ w2b
else:
t2 = self.t2.to(target.device, dtype=target.dtype)
w2a = self.w2a.to(target.device, dtype=target.dtype)
w2b = self.w2b.to(target.device, dtype=target.dtype)
w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b)
w1, w2 = self.rebuild_operands(target)
full_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)]
updown = make_kron(full_shape, w1, w2)
updown = updown[self.start_row:self.end_row]