Allow users to init their LoKr with perturbed normal w2

helpers/configuration/cmd_args.py

@@ -319,6 +319,13 @@ def parse_cmdline_args(input_args=None):
         default="config/lycoris_config.json",
         help=("The location for the JSON file of the Lycoris configuration."),
     )
+    parser.add_argument(
+        "--init_lokr_norm",
+        type=float,
+        required=False,
+        default=None,
+        help=("Setting this turns on perturbed normal initialization of the LyCORIS LoKr PEFT layers. A good value is between 1e-4 and 1e-2."),
+    )
     parser.add_argument(
         "--controlnet",
         action="store_true",

helpers/training/peft_init.py

@@ -0,0 +1,24 @@
+import torch
+
+
+def approximate_normal_tensor(inp, target, scale=1.0):
+    tensor = torch.randn_like(target)
+    desired_norm = inp.norm()
+    desired_mean = inp.mean()
+    desired_std = inp.std()
+
+    current_norm = tensor.norm()
+    tensor = tensor * (desired_norm / current_norm)
+    current_std = tensor.std()
+    tensor = tensor * (desired_std / current_std)
+    tensor = tensor - tensor.mean() + desired_mean
+    tensor.mul_(scale)
+
+    target.copy_(tensor)
+
+
+def init_lokr_network_with_perturbed_normal(lycoris, scale=1e-3):
+    with torch.no_grad():
+        for lora in lycoris.loras:
+            lora.lokr_w1.fill_(1.0)
+            approximate_normal_tensor(lora.org_weight, lora.lokr_w2, scale=scale)

helpers/training/trainer.py

@@ -50,6 +50,7 @@
     segmented_timestep_selection,
 )
 from helpers.training.min_snr_gamma import compute_snr
+from helpers.training.peft_init import init_lokr_network_with_perturbed_standard
 from accelerate.logging import get_logger
 from diffusers.models.embeddings import get_2d_rotary_pos_embed
 from helpers.models.smoldit import get_resize_crop_region_for_grid
@@ -794,6 +795,12 @@ def init_trainable_peft_adapter(self):
                     **self.lycoris_config,
                 )
 
+                if self.config.init_lycoris_lokr_perturbed_normal is not None:
+                    init_lokr_network_with_perturbed_standard(
+                        self.lycoris_wrapped_network,
+                        scale=self.config.init_lokr_norm,
+                    )
+
             self.lycoris_wrapped_network.apply_to()
             setattr(
                 self.accelerator,