Added gradient accumulation finally

toolkit/config_modules.py

@@ -52,6 +52,7 @@ class LormModuleSettingsConfig:
 class LoRMConfig:
     def __init__(self, **kwargs):
         self.extract_mode: str = kwargs.get('extract_mode', 'ratio')
+        self.do_conv: bool = kwargs.get('do_conv', False)
         self.extract_mode_param: dict = kwargs.get('extract_mode_param', 0.25)
         self.parameter_threshold: int = kwargs.get('parameter_threshold', 0)
         module_settings = kwargs.get('module_settings', [])
@@ -110,6 +111,8 @@ class NetworkConfig:
             # set linear to arbitrary values so it makes them
             self.linear = 4
             self.rank = 4
+            if self.lorm_config.do_conv:
+                self.conv = 4
 
 
 AdapterTypes = Literal['t2i', 'ip', 'ip+']
@@ -177,6 +180,10 @@ class TrainConfig:
         self.noise_multiplier = kwargs.get('noise_multiplier', 1.0)
         self.img_multiplier = kwargs.get('img_multiplier', 1.0)
 
+        # set to -1 to accumulate gradients for entire epoch
+        # warning, only do this with a small dataset or you will run out of memory
+        self.gradient_accumulation_steps = kwargs.get('gradient_accumulation_steps', 1)
+
         # short long captions will double your batch size. This only works when a dataset is
         # prepared with a json caption file that has both short and long captions in it. It will
         # Double up every image and run it through with both short and long captions. The idea

toolkit/lorm.py

@@ -191,6 +191,8 @@ def extract_conv(
     if lora_rank >= out_ch / 2:
         lora_rank = int(out_ch / 2)
         print(f"rank is higher than it should be")
+        # print(f"Skipping layer as determined rank is too high")
+        # return None, None, None, None
         # return weight, 'full'
 
     U = U[:, :lora_rank]
@@ -243,6 +245,8 @@ def extract_linear(
         # print(f"rank is higher than it should be")
         lora_rank = int(out_ch / 2)
         # return weight, 'full'
+        # print(f"Skipping layer as determined rank is too high")
+        # return None, None, None, None
 
     U = U[:, :lora_rank]
     S = S[:lora_rank]
@@ -358,6 +362,8 @@ def convert_diffusers_unet_to_lorm(
                             mode_param=extract_mode_param,
                             device=child_module.weight.device,
                         )
+                        if down_weight is None:
+                            continue
                         down_weight = down_weight.to(dtype=dtype)
                         up_weight = up_weight.to(dtype=dtype)
                         bias_weight = None
@@ -398,6 +404,8 @@ def convert_diffusers_unet_to_lorm(
                             mode_param=extract_mode_param,
                             device=child_module.weight.device,
                         )
+                        if down_weight is None:
+                            continue
                         down_weight = down_weight.to(dtype=dtype)
                         up_weight = up_weight.to(dtype=dtype)
                         bias_weight = None