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]