From 79c2c2debe1f393b83e06d518cb78afeeb105935 Mon Sep 17 00:00:00 2001
From: bmaltais <bernard@ducourier.com>
Date: Sun, 12 Mar 2023 10:11:41 -0400
Subject: [PATCH] Add validation that all requirements are met

---
 README.md                      | 29 ++++++++++++++++++------
 gui.bat                        | 12 +++++++---
 gui.ps1                        | 12 ++++++++--
 tools/validate_requirements.py | 41 ++++++++++++++++++++++++++++++++++
 4 files changed, 82 insertions(+), 12 deletions(-)
 create mode 100644 tools/validate_requirements.py

diff --git a/README.md b/README.md
index fc49070..393464f 100644
--- a/README.md
+++ b/README.md
@@ -110,16 +110,29 @@ 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
+
+The script can be run with several optional command line arguments:
+
+--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
+
+`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:
 
-```
-.\gui.ps1
-```
-
-or you can also do:
-
 ```
 .\venv\Scripts\activate
 
@@ -176,6 +189,8 @@ 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 detais.
 * 2023/03/10 (v21.2.1):
diff --git a/gui.bat b/gui.bat
index a7e8d08..d1f2864 100644
--- a/gui.bat
+++ b/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
\ No newline at end of file
+:: 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 %*
+)
\ No newline at end of file
diff --git a/gui.ps1 b/gui.ps1
index fde5028..81b43a0 100644
--- a/gui.ps1
+++ b/gui.ps1
@@ -1,2 +1,10 @@
-.\venv\Scripts\activate
-python.exe kohya_gui.py $args
\ No newline at end of file
+# 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
+}
\ No newline at end of file
diff --git a/tools/validate_requirements.py b/tools/validate_requirements.py
new file mode 100644
index 0000000..9af2ce4
--- /dev/null
+++ b/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)