From ba9e26a99db9d38534bc2b446c6027e024cf23be Mon Sep 17 00:00:00 2001
From: bmaltais <bernard@ducourier.com>
Date: Thu, 9 Mar 2023 07:49:50 -0500
Subject: [PATCH] LoHa support

---
 README.md             |  1 +
 library/common_gui.py |  5 ++++-
 lora_gui.py           | 28 +++++++++++++++++++++-------
 requirements.txt      |  2 +-
 4 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 24ad0b7..07c775e 100644
--- a/README.md
+++ b/README.md
@@ -177,6 +177,7 @@ This will store your a backup file with your current locally installed pip packa
 ## Change History
 
 * 2023/03/05 (v21.2.0):
+    - Add support for LoRA LoHa type. See https://github.com/KohakuBlueleaf/LyCORIS for more detais.
     - Fix issue https://github.com/bmaltais/kohya_ss/issues/335
     - Add option to print LoRA trainer command without executing it
     - Add support for samples during trainin via a new `Sample images config` accordion in the `Training parameters` tab.
diff --git a/library/common_gui.py b/library/common_gui.py
index b22594f..477d346 100644
--- a/library/common_gui.py
+++ b/library/common_gui.py
@@ -58,7 +58,10 @@ def update_my_data(my_data):
                 my_data[key] = int(value)
             else:
                 my_data[key] = -1
-
+                
+    if my_data.get('LoRA_type', 'Standard') == 'LoCon':
+        my_data['LoRA_type'] = 'LyCORIS/LoCon'
+        
     return my_data
 
 
diff --git a/lora_gui.py b/lora_gui.py
index d175b30..23eae18 100644
--- a/lora_gui.py
+++ b/lora_gui.py
@@ -474,17 +474,29 @@ def train_model(
         run_cmd += f' --save_model_as={save_model_as}'
     if not float(prior_loss_weight) == 1.0:
         run_cmd += f' --prior_loss_weight={prior_loss_weight}'
-    if LoRA_type == 'LoCon':
+    if LoRA_type == 'LoCon' or LoRA_type == 'LyCORIS/LoCon':
         try:
-            import locon.locon_kohya
+            import lycoris
         except ModuleNotFoundError:
             print(
-                "\033[1;31mError:\033[0m The required module 'locon' is not installed. Please install by running \033[33mupgrade.ps1\033[0m before running this program."
+                "\033[1;31mError:\033[0m The required module 'lycoris_lora' is not installed. Please install by running \033[33mupgrade.ps1\033[0m before running this program."
             )
             return
-        run_cmd += f' --network_module=locon.locon_kohya'
+        run_cmd += f' --network_module=lycoris.kohya'
         run_cmd += (
-            f' --network_args "conv_dim={conv_dim}" "conv_alpha={conv_alpha}"'
+            f' --network_args "conv_dim={conv_dim}" "conv_alpha={conv_alpha}" "algo=lora"'
+        )
+    if LoRA_type == 'LyCORIS/LoHa':
+        try:
+            import lycoris
+        except ModuleNotFoundError:
+            print(
+                "\033[1;31mError:\033[0m The required module 'lycoris_lora' is not installed. Please install by running \033[33mupgrade.ps1\033[0m before running this program."
+            )
+            return
+        run_cmd += f' --network_module=lycoris.kohya'
+        run_cmd += (
+            f' --network_args "conv_dim={conv_dim}" "conv_alpha={conv_alpha}" "algo=loha"'
         )
     if LoRA_type == 'Kohya LoCon':
         run_cmd += f' --network_module=networks.lora'
@@ -700,7 +712,9 @@ def lora_tab(
                 label='LoRA type',
                 choices=[
                     'Kohya LoCon',
-                    'LoCon',
+                    # 'LoCon',
+                    'LyCORIS/LoCon',
+                    'LyCORIS/LoHa',
                     'Standard',
                 ],
                 value='Standard',
@@ -787,7 +801,7 @@ def lora_tab(
         # Show of hide LoCon conv settings depending on LoRA type selection
         def LoRA_type_change(LoRA_type):
             print('LoRA type changed...')
-            if LoRA_type == 'LoCon' or LoRA_type == 'Kohya LoCon':
+            if LoRA_type == 'LoCon' or LoRA_type == 'Kohya LoCon' or LoRA_type == 'LyCORIS/LoHa' or LoRA_type == 'LyCORIS/LoCon':
                 return gr.Group.update(visible=True)
             else:
                 return gr.Group.update(visible=False)
diff --git a/requirements.txt b/requirements.txt
index 27bde40..479c7f6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -25,6 +25,6 @@ timm==0.6.12
 huggingface-hub==0.12.0
 tensorflow==2.10.1
 # For locon support
-locon==0.0.2
+lycoris_lora==0.0.5
 # for kohya_ss library
 .
\ No newline at end of file