fix: support multiplier, no forward pass hook (fce86ab7) · Commits · github_fork / Stable Diffusion Webui

extensions-builtin/Lora/network_oft.py

+33 −10

Original line number	Diff line number	Diff line
		@@ -32,21 +32,27 @@ class NetworkModuleOFT(network.NetworkModule):
		self.org_module: list[torch.Module] = [self.sd_module]
		self.org_weight = self.org_module[0].weight.to(self.org_module[0].weight.device, copy=True)
		#self.org_weight = self.org_module[0].weight.to(devices.cpu, copy=True)
		self.R = self.get_weight(self.oft_blocks)
		init_multiplier = self.multiplier() * self.calc_scale()
		self.last_multiplier = init_multiplier
		self.R = self.get_weight(self.oft_blocks, init_multiplier)

		self.merged_weight = self.merge_weight()
		self.apply_to()
		self.merged = False

		# weights_backup = getattr(self.org_module[0], 'network_weights_backup', None)
		# if weights_backup is None:
		# self.org_module[0].network_weights_backup = self.org_weight


		def merge_weight(self):
		org_sd = self.org_module[0].state_dict()
		#org_sd = self.org_module[0].state_dict()
		R = self.R.to(self.org_weight.device, dtype=self.org_weight.dtype)
		if self.org_weight.dim() == 4:
		weight = torch.einsum("oihw, op -> pihw", self.org_weight, R)
		else:
		weight = torch.einsum("oi, op -> pi", self.org_weight, R)
		org_sd['weight'] = weight
		#org_sd['weight'] = weight
		# replace weight
		#self.org_module[0].load_state_dict(org_sd)
		return weight
		@@ -74,6 +80,7 @@ class NetworkModuleOFT(network.NetworkModule):
		self.org_module[0].register_forward_hook(self.forward_hook)

		def get_weight(self, oft_blocks, multiplier=None):
		multiplier = multiplier.to(oft_blocks.device, dtype=oft_blocks.dtype)
		constraint = self.constraint.to(oft_blocks.device, dtype=oft_blocks.dtype)
		block_Q = oft_blocks - oft_blocks.transpose(1, 2)
		norm_Q = torch.norm(block_Q.flatten())
		@@ -81,9 +88,9 @@ class NetworkModuleOFT(network.NetworkModule):
		block_Q = block_Q * ((new_norm_Q + 1e-8) / (norm_Q + 1e-8))
		m_I = torch.eye(self.block_size, device=oft_blocks.device).unsqueeze(0).repeat(self.num_blocks, 1, 1)
		block_R = torch.matmul(m_I + block_Q, (m_I - block_Q).inverse())
		#block_R_weighted = multiplier * block_R + (1 - multiplier) * I
		#R = torch.block_diag(*block_R_weighted)
		R = torch.block_diag(*block_R)
		block_R_weighted = multiplier * block_R + (1 - multiplier) * m_I
		R = torch.block_diag(*block_R_weighted)
		#R = torch.block_diag(*block_R)

		return R

		@@ -93,6 +100,8 @@ class NetworkModuleOFT(network.NetworkModule):
		#R = self.R.to(orig_weight.device, dtype=orig_weight.dtype)
		##self.R = R

		#R = self.R.to(orig_weight.device, dtype=orig_weight.dtype)
		##self.R = R
		#if orig_weight.dim() == 4:
		# weight = torch.einsum("oihw, op -> pihw", orig_weight, R)
		#else:
		@@ -103,19 +112,33 @@ class NetworkModuleOFT(network.NetworkModule):
		updown = torch.zeros_like(orig_weight, device=orig_weight.device, dtype=orig_weight.dtype)
		#updown = orig_weight
		output_shape = orig_weight.shape
		#orig_weight = self.merged_weight.to(orig_weight.device, dtype=orig_weight.dtype)
		orig_weight = self.merged_weight.to(orig_weight.device, dtype=orig_weight.dtype)
		#output_shape = self.oft_blocks.shape

		return self.finalize_updown(updown, orig_weight, output_shape)

		def pre_forward_hook(self, module, input):
		if not self.merged:
		multiplier = self.multiplier() * self.calc_scale()
		if not multiplier==self.last_multiplier or not self.merged:

		#if multiplier != self.last_multiplier or not self.merged:
		self.R = self.get_weight(self.oft_blocks, multiplier)
		self.last_multiplier = multiplier
		self.merged_weight = self.merge_weight()
		self.replace_weight(self.merged_weight)
		#elif not self.merged:
		# self.replace_weight(self.merged_weight)


		def forward_hook(self, module, args, output):
		if self.merged:
		pass
		#output = output * self.multiplier() * self.calc_scale()
		#if len(args) > 0:
		# y = args[0]
		# output = output + y
		#return output
		#if self.merged:
		# pass
		#self.restore_weight()
		#print(f'Forward hook in {self.network_key} called')

Admin message