Merge pull request bmaltais#363 from bmaltais/LoHa-Support

Lo ha support

README.md

@@ -110,15 +110,28 @@ pip install --use-pep517 --upgrade -r requirements.txt
 
 Once the commands have completed successfully you should be ready to use the new version.
 
-## Launching the GUI
+## Launching the GUI using gui.bat or gui.ps1
 
-To run the GUI, simply use this command:
+The script can be run with several optional command line arguments:
 
-```
-.\gui.ps1
-```
+--listen: the IP address to listen on for connections to Gradio.
+--username: a username for authentication.
+--password: a password for authentication.
+--server_port: the port to run the server listener on.
+--inbrowser: opens the Gradio UI in a web browser.
+--share: shares the Gradio UI.
+
+These command line arguments can be passed to the UI function as keyword arguments. To launch the Gradio UI, run the script in a terminal with the desired command line arguments, for example:
+
+`gui.ps1 --listen 127.0.0.1 --server_port 7860 --inbrowser --share`
 
-or you can also do:
+or
+
+`gui.bat --listen 127.0.0.1 --server_port 7860 --inbrowser --share`
+
+## Launching the GUI using kohya_gui.py
+
+To run the GUI, simply use this command:
 
 ```
 .\venv\Scripts\activate
@@ -176,6 +189,10 @@ This will store your a backup file with your current locally installed pip packa
 
 ## Change History
 
+* 2023.03/12 (v21.2.3):
+    - Add validation that all requirements are met before starting the GUI.
+* 2023/03/11 (v21.2.2):
+    - Add support for LoRA LoHa type. See https://github.com/KohakuBlueleaf/LyCORIS for more details.
 * 2023/03/10 (v21.2.1):
     - Update to latest sd-script code
     - Add support for SVD based LoRA merge

gui.bat

@@ -1,6 +1,12 @@
 @echo off
 
-call venv\Scripts\activate.bat
-python.exe kohya_gui.py %*
+:: Activate the virtual environment
+call .\venv\Scripts\activate.bat
 
-pause
+:: Validate the requirements and store the exit code
+python.exe .\tools\validate_requirements.py
+
+:: If the exit code is 0, run the kohya_gui.py script with the command-line arguments
+if %errorlevel% equ 0 (
+    python.exe kohya_gui.py %*
+)

gui.ps1

@@ -1,2 +1,10 @@
-.\venv\Scripts\activate
-python.exe kohya_gui.py $args
+# Activate the virtual environment
+& .\venv\Scripts\activate
+
+# Validate the requirements and store the exit code
+python.exe .\tools\validate_requirements.py
+
+# If the exit code is 0, run the kohya_gui.py script with the command-line arguments
+if ($LASTEXITCODE -eq 0) {
+    python.exe kohya_gui.py $args
+}

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
 
 

lora_gui.py

@@ -480,17 +480,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'
@@ -707,7 +719,9 @@ def lora_tab(
                 label='LoRA type',
                 choices=[
                     'Kohya LoCon',
-                    'LoCon',
+                    # 'LoCon',
+                    'LyCORIS/LoCon',
+                    'LyCORIS/LoHa',
                     'Standard',
                 ],
                 value='Standard',
@@ -782,19 +796,19 @@ def lora_tab(
                 maximum=512,
                 value=1,
                 step=1,
-                label='LoCon Convolution Rank (Dimension)',
+                label='Convolution Rank (Dimension)',
             )
             conv_alpha = gr.Slider(
                 minimum=1,
                 maximum=512,
                 value=1,
                 step=1,
-                label='LoCon Convolution Alpha',
+                label='Convolution Alpha',
             )
         # 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)

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.9
 # for kohya_ss library
 .

tools/validate_requirements.py

@@ -0,0 +1,41 @@
+import sys
+import pkg_resources
+
+print("Validating that requirements are satisfied.")
+
+# Load the requirements from the requirements.txt file
+with open('requirements.txt') as f:
+    requirements = f.readlines()
+
+# Check each requirement against the installed packages
+missing_requirements = []
+wrong_version_requirements = []
+for requirement in requirements:
+    requirement = requirement.strip()
+    if requirement == ".":
+        # Skip the current requirement if it is a dot (.)
+        continue
+    try:
+        pkg_resources.require(requirement)
+    except pkg_resources.DistributionNotFound:
+        missing_requirements.append(requirement)
+    except pkg_resources.VersionConflict as e:
+        wrong_version_requirements.append((requirement, str(e.req), e.dist.version))
+
+# If there are any missing or wrong version requirements, print an error message and exit with a non-zero exit code
+if missing_requirements or wrong_version_requirements:
+    if missing_requirements:
+        print("Error: The following packages are missing:")
+        for requirement in missing_requirements:
+            print(f" - {requirement}")
+    if wrong_version_requirements:
+        print("Error: The following packages have the wrong version:")
+        for requirement, expected_version, actual_version in wrong_version_requirements:
+            print(f" - {requirement} (expected version {expected_version}, found version {actual_version})")
+    print('\nRun \033[33mupgrade.ps1\033[0m or \033[33mpip install -U -r requirements.txt\033[0m to resolve the missing requirements listed above...')
+
+    sys.exit(1)
+
+# All requirements satisfied
+print("All requirements satisfied.")
+sys.exit(0)