Add support for LoRA-FA

bmaltais · Aug 15, 2023 · d311bba · d311bba
1 parent 859f71b
commit d311bba
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -520,9 +520,11 @@ If you come across a `FileNotFoundError`, it is likely due to an installation is
 
 * 2023/08/05 (v21.8.8)
   - Fix issue with aiofiles: https://github.com/bmaltais/kohya_ss/issues/1359
-  - Merge sd-scripts updates as of Aug 11 2023
+  - Merge sd-scripts updates as of Aug 18 2023
   - Add new blip2 caption processor tool
   - Add dataset preparation tab to appropriate trainers
+  - Add GUI support for new block_lr lora network parameter
+  - Add support for experimental LoRA-FA network
 * 2023/08/05 (v21.8.7)
   - Add manual captioning option. Thanks to https://github.com/channelcat for this great contribution. (https://github.com/bmaltais/kohya_ss/pull/1352)
   - Added support for `v_pred_like_loss` to the advanced training tab
diff --git a/lora_gui.py b/lora_gui.py
@@ -812,6 +812,39 @@ def train_model(
 
         if network_args:
             run_cmd += f' --network_args{network_args}'
+
+    if LoRA_type in ['LoRA-FA',]:
+        kohya_lora_var_list = [
+            'down_lr_weight',
+            'mid_lr_weight',
+            'up_lr_weight',
+            'block_lr_zero_threshold',
+            'block_dims',
+            'block_alphas',
+            'block_lr',
+            'conv_block_dims',
+            'conv_block_alphas',
+            'rank_dropout',
+            'module_dropout',
+        ]
+
+        run_cmd += f' --network_module=networks.lora_fa'
+        kohya_lora_vars = {
+            key: value
+            for key, value in vars().items()
+            if key in kohya_lora_var_list and value
+        }
+
+        network_args = ''
+        if LoRA_type == 'Kohya LoCon':
+            network_args += f' conv_dim="{conv_dim}" conv_alpha="{conv_alpha}"'
+
+        for key, value in kohya_lora_vars.items():
+            if value:
+                network_args += f' {key}="{value}"'
+
+        if network_args:
+            run_cmd += f' --network_args{network_args}'
 
     if LoRA_type in ['Kohya DyLoRA']:
         kohya_lora_var_list = [
@@ -1061,6 +1094,7 @@ def list_presets(path):
                         choices=[
                             'Kohya DyLoRA',
                             'Kohya LoCon',
+                            'LoRA-FA',
                             'LyCORIS/DyLoRA',
                             'LyCORIS/iA3',
                             'LyCORIS/LoCon',
@@ -1228,6 +1262,7 @@ def update_LoRA_settings(LoRA_type):
                                 {
                                     'Kohya DyLoRA',
                                     'Kohya LoCon',
+                                    'LoRA-FA',
                                     'LyCORIS/DyLoRA',
                                     'LyCORIS/LoCon',
                                     'LyCORIS/LoHa',
@@ -1241,6 +1276,7 @@ def update_LoRA_settings(LoRA_type):
                                     'LoCon',
                                     'Kohya DyLoRA',
                                     'Kohya LoCon',
+                                    'LoRA-FA',
                                     'LyCORIS/DyLoRA',
                                     'LyCORIS/LoHa',
                                     'LyCORIS/LoKr',
@@ -1249,29 +1285,29 @@ def update_LoRA_settings(LoRA_type):
                                 gr.Row,
                             ),
                             'kohya_advanced_lora': (
-                                {'Standard', 'Kohya DyLoRA', 'Kohya LoCon'},
+                                {'Standard', 'Kohya DyLoRA', 'Kohya LoCon', 'LoRA-FA',},
                                 gr.Row,
                             ),
                             'kohya_dylora': (
                                 {'Kohya DyLoRA', 'LyCORIS/DyLoRA'},
                                 gr.Row,
                             ),
                             'lora_network_weights': (
-                                {'Standard', 'LoCon', 'Kohya DyLoRA', 'Kohya LoCon','LyCORIS/DyLoRA',
+                                {'Standard', 'LoCon', 'Kohya DyLoRA', 'Kohya LoCon', 'LoRA-FA', 'LyCORIS/DyLoRA',
                                     'LyCORIS/LoHa',
                                     'LyCORIS/LoCon',
                                     'LyCORIS/LoKr',},
                                 gr.Textbox,
                             ),
                             'lora_network_weights_file': (
-                                {'Standard', 'LoCon', 'Kohya DyLoRA', 'Kohya LoCon','LyCORIS/DyLoRA',
+                                {'Standard', 'LoCon', 'Kohya DyLoRA', 'Kohya LoCon', 'LoRA-FA', 'LyCORIS/DyLoRA',
                                     'LyCORIS/LoHa',
                                     'LyCORIS/LoCon',
                                     'LyCORIS/LoKr',},
                                 gr.Button,
                             ),
                             'dim_from_weights': (
-                                {'Standard', 'LoCon', 'Kohya DyLoRA', 'Kohya LoCon','LyCORIS/DyLoRA',
+                                {'Standard', 'LoCon', 'Kohya DyLoRA', 'Kohya LoCon', 'LoRA-FA', 'LyCORIS/DyLoRA',
                                     'LyCORIS/LoHa',
                                     'LyCORIS/LoCon',
                                     'LyCORIS/LoKr',},
@@ -1294,6 +1330,7 @@ def update_LoRA_settings(LoRA_type):
                                     'LoCon',
                                     'Kohya DyLoRA',
                                     'Kohya LoCon',
+                                    'LoRA-FA',
                                     'LyCORIS/DyLoRA',
                                     'LyCORIS/LoHa',
                                     'LyCORIS/LoCon',
@@ -1307,6 +1344,7 @@ def update_LoRA_settings(LoRA_type):
                                     'LoCon',
                                     'Kohya DyLoRA',
                                     'Kohya LoCon',
+                                    'LoRA-FA',
                                     'LyCORIS/DyLoRA',
                                     'LyCORIS/LoHa',
                                     'LyCORIS/LoCon',
@@ -1316,12 +1354,12 @@ def update_LoRA_settings(LoRA_type):
                                 gr.Slider,
                             ),
                             'rank_dropout': (
-                                {'LoCon', 'Kohya DyLoRA', 'Kohya LoCon',
+                                {'LoCon', 'Kohya DyLoRA', 'Kohya LoCon', 'LoRA-FA',
                                     'Standard',},
                                 gr.Slider,
                             ),
                             'module_dropout': (
-                                {'LoCon', 'Kohya DyLoRA', 'Kohya LoCon',
+                                {'LoCon', 'Kohya DyLoRA', 'Kohya LoCon', 'LoRA-FA',
                                     'Standard',},
                                 gr.Slider,
                             ),