diff --git a/gguf/gguf.py b/gguf/gguf.py index 9e93e67..d835211 100644 --- a/gguf/gguf.py +++ b/gguf/gguf.py @@ -34,9 +34,7 @@ def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modul key = "diffusion_model." + module_prefix + "weight" patch = patches.get(key, []) - lora_diffs = None - lora_strengths = None - lora_alphas = None + lora_diffs = lora_strengths = lora_alphas = None if len(patch) != 0: lora_diffs = [p[1].weights for p in patch] lora_strengths = [p[0] for p in patch] @@ -85,6 +83,7 @@ class GGUFLinear(nn.Linear): def forward(self, inputs): weight = dequantize_without_compile(self.weight) + temp_weight = weight.to(torch.float32) weight = weight.to(self.compute_dtype) bias = self.bias.to(self.compute_dtype) if self.bias is not None else None @@ -93,12 +92,12 @@ class GGUFLinear(nn.Linear): for lora_diff, lora_strength, lora_alpha in zip(self.lora_diffs, self.lora_strengths, self.lora_alphas): # Calculate the diff for this patch patch_diff = torch.mm( - lora_diff[0].flatten(start_dim=1).to(weight.device), - lora_diff[1].flatten(start_dim=1).to(weight.device) + lora_diff[0].flatten(start_dim=1).to(temp_weight.device), + lora_diff[1].flatten(start_dim=1).to(temp_weight.device) ).reshape(weight.shape) # Apply the patch with its strength - weight = weight + ((lora_strength * lora_alpha) * patch_diff).to(self.compute_dtype) + weight = (temp_weight + ((lora_strength * lora_alpha) * patch_diff)).to(self.compute_dtype) output = torch.nn.functional.linear(inputs, weight, bias) return output \ No newline at end of file