Load works, unload infinite recursion

extensions-builtin/Lora/lora_patches.py

@@ -20,6 +20,11 @@ class LoraPatches:
     def apply(self):
         if self.active or shared.opts.lora_force_diffusers:
             return
+        try:
+            import bitsandbytes
+            self.Linear4bit_forward = patches.patch(__name__, bitsandbytes.nn.Linear4bit, 'forward', networks.network_Linear4bit_forward)
+        except:
+            pass
         if "Model" in shared.opts.optimum_quanto_weights or "Text Encoder" in shared.opts.optimum_quanto_weights:
             from optimum import quanto # pylint: disable=no-name-in-module
             self.QLinear_forward = patches.patch(__name__, quanto.nn.QLinear, 'forward', networks.network_QLinear_forward) # pylint: disable=attribute-defined-outside-init

extensions-builtin/Lora/network_lora.py

@@ -55,12 +55,13 @@ class NetworkModuleLora(network.NetworkModule):
         return module
 
     def calc_updown(self, target): # pylint: disable=W0237
-        up = self.up_model.weight.to(target.device, dtype=target.dtype)
-        down = self.down_model.weight.to(target.device, dtype=target.dtype)
+        target_dtype = target.dtype if target.dtype != torch.uint8 else self.up_model.weight.dtype
+        up = self.up_model.weight.to(target.device, dtype=target_dtype)
+        down = self.down_model.weight.to(target.device, dtype=target_dtype)
         output_shape = [up.size(0), down.size(1)]
         if self.mid_model is not None:
             # cp-decomposition
-            mid = self.mid_model.weight.to(target.device, dtype=target.dtype)
+            mid = self.mid_model.weight.to(target.device, dtype=target_dtype)
             updown = lyco_helpers.rebuild_cp_decomposition(up, down, mid)
             output_shape += mid.shape[2:]
         else:

extensions-builtin/Lora/networks.py

@@ -329,16 +329,16 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
         if current_names != ():
             raise RuntimeError("no backup weights found and current weights are not unchanged")
         if isinstance(self, torch.nn.MultiheadAttention):
-            weights_backup = (self.in_proj_weight.to(devices.cpu, copy=True), self.out_proj.weight.to(devices.cpu, copy=True))
+            weights_backup = (self.in_proj_weight.clone().to(devices.cpu), self.out_proj.weight.clone().to(devices.cpu))
         else:
-            weights_backup = self.weight.to(devices.cpu, copy=True)
+            weights_backup = self.weight.clone().to(devices.cpu)
         self.network_weights_backup = weights_backup
     bias_backup = getattr(self, "network_bias_backup", None)
     if bias_backup is None:
         if isinstance(self, torch.nn.MultiheadAttention) and self.out_proj.bias is not None:
-            bias_backup = self.out_proj.bias.to(devices.cpu, copy=True)
+            bias_backup = self.out_proj.bias.clone().to(devices.cpu)
         elif getattr(self, 'bias', None) is not None:
-            bias_backup = self.bias.to(devices.cpu, copy=True)
+            bias_backup = self.bias.clone().to(devices.cpu)
         else:
             bias_backup = None
         self.network_bias_backup = bias_backup
@@ -356,7 +356,14 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
                         if len(weight.shape) == 4 and weight.shape[1] == 9:
                             # inpainting model. zero pad updown to make channel[1]  4 to 9
                             updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) # pylint: disable=not-callable
-                        self.weight = torch.nn.Parameter(weight + updown)
+                        if getattr(self.weight, "quant_type", None) == "nf4":
+                            import bitsandbytes
+                            device = self.weight.device
+                            weight = bitsandbytes.functional.dequantize_4bit(self.weight, quant_state=self.weight.quant_state, quant_type=self.weight.quant_type, blocksize=self.weight.blocksize)
+                            self.weight = bitsandbytes.nn.Params4bit(weight + updown)
+                            self.weight.to(device)
+                        else:
+                            self.weight = torch.nn.Parameter(weight + updown)
                         if hasattr(self, "qweight") and hasattr(self, "freeze"):
                             self.freeze()
                         if ex_bias is not None and hasattr(self, 'bias'):
@@ -446,6 +453,14 @@ def network_Linear_load_state_dict(self, *args, **kwargs):
     return originals.Linear_load_state_dict(self, *args, **kwargs)
 
 
+def network_Linear4bit_forward(self, input): # pylint: disable=W0622
+    network_apply_weights(self)
+    return originals.Linear4bit_forward(self, input)
+#
+# def network_Linear4bit_load_state_dict(self, *args, **kwargs):
+#     network_reset_cached_weight(self)
+#     return originals.Linear4bit_load_state_dict(self, *args, **kwargs)
+
 def network_Conv2d_forward(self, input): # pylint: disable=W0622
     network_apply_weights(self)
     return originals.Conv2d_forward(self, input)