Save embeddings as their trigger to match auto and comfy style loading. Also, FINALLY found why gradients were wonkey and fixed it. The root problem is dropping out of network state before backward pass.

toolkit/lora_special.py

@@ -50,10 +50,7 @@ class LoRAModule(ToolkitModuleMixin, torch.nn.Module):
             **kwargs
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
-        super().__init__(
-            org_module=org_module,
-            parent=parent
-        )
+        super().__init__()
         self.lora_name = lora_name
         self.scalar = torch.tensor(1.0)
 

toolkit/lycoris_special.py

@@ -36,12 +36,7 @@ class LoConSpecialModule(ToolkitModuleMixin, LoConModule):
         # call super of super
         torch.nn.Module.__init__(self)
         # call super of
-        super().__init__(
-            org_module=org_module,
-            call_super_init=False,
-            parent=parent,
-            **kwargs
-        )
+        super().__init__(call_super_init=False)
         self.lora_name = lora_name
         self.lora_dim = lora_dim
         self.cp = False

toolkit/network_mixins.py

@@ -52,18 +52,9 @@ class ToolkitModuleMixin:
     ):
         if call_super_init:
             super().__init__(*args, **kwargs)
-        self.tk_orig_module: torch.nn.Module = kwargs.get('org_module', None)
-        self.tk_orig_parent = kwargs.get('parent', None)
         self.is_checkpointing = False
         self.is_normalizing = False
         self.normalize_scaler = 1.0
-        # see if is conv or linear
-        self.is_conv = False
-        self.is_linear = False
-        if self.tk_orig_module.__class__.__name__ in LINEAR_MODULES:
-            self.is_linear = True
-        elif self.tk_orig_module.__class__.__name__ in CONV_MODULES:
-            self.is_conv = True
         self._multiplier: Union[float, list, torch.Tensor] = 1.0
 
     # this allows us to set different multipliers on a per item in a batch basis
@@ -140,10 +131,6 @@ class ToolkitModuleMixin:
         lora_output_batch_size = lora_output.size(0)
         multiplier_batch_size = multiplier.size(0)
         if lora_output_batch_size != multiplier_batch_size:
-            print(
-                f"Warning: lora_output_batch_size {lora_output_batch_size} != multiplier_batch_size {multiplier_batch_size}")
-            # doing cfg
-            # should be 1 for if total batch size was 1
             num_interleaves = (lora_output_batch_size // 2) // multiplier_batch_size
             multiplier = multiplier.repeat_interleave(num_interleaves)
         # multiplier = 1.0