From 5538b193909b7cd7501e99f4f2bdc1cd98937c56 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 30 Aug 2026 05:10:34 +0100 Subject: [PATCH] refactor(lora): extract the lokr operand rebuild 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. --- modules/lora/network_lokr.py | 45 ++++++++---------------------------- 1 file changed, 9 insertions(+), 36 deletions(-) diff --git a/modules/lora/network_lokr.py b/modules/lora/network_lokr.py index 096d7f568..0f2328a9e 100644 --- a/modules/lora/network_lokr.py +++ b/modules/lora/network_lokr.py @@ -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]