Create an update wizard (#5623)

README.md

@@ -32,7 +32,7 @@ To restart the web UI in the future, just run the `start_` script again. This sc
 
 The script accepts command-line flags. Alternatively, you can edit the `CMD_FLAGS.txt` file with a text editor and add your flags there.
 
-To get updates in the future, run `update_linux.sh`, `update_windows.bat`, `update_macos.sh`, or `update_wsl.bat`.
+To get updates in the future, run `update_wizard_linux.sh`, `update_wizard_windows.bat`, `update_wizard_macos.sh`, or `update_wizard_wsl.bat`.
 
 <details>
 <summary>
@@ -45,7 +45,7 @@ The script uses Miniconda to set up a Conda environment in the `installer_files`
 
 If you ever need to install something manually in the `installer_files` environment, you can launch an interactive shell using the cmd script: `cmd_linux.sh`, `cmd_windows.bat`, `cmd_macos.sh`, or `cmd_wsl.bat`.
 
-* There is no need to run any of those scripts (`start_`, `update_`, or `cmd_`) as admin/root.
+* There is no need to run any of those scripts (`start_`, `update_wizard_`, or `cmd_`) as admin/root.
 * To install the requirements for extensions, you can use the `extensions_reqs` script for your OS. At the end, this script will install the main requirements for the project to make sure that they take precedence in case of version conflicts.
 * For additional instructions about AMD and WSL setup, consult [the documentation](https://github.com/oobabooga/text-generation-webui/wiki).
 * For automated installation, you can use the `GPU_CHOICE`, `USE_CUDA118`, and `LAUNCH_AFTER_INSTALL` environment variables. For instance: `GPU_CHOICE=A USE_CUDA118=FALSE LAUNCH_AFTER_INSTALL=FALSE ./start_linux.sh`.

one_click.py

@@ -32,7 +32,7 @@
 else:
     CMD_FLAGS = ''
 
-flags = f"{' '.join([flag for flag in sys.argv[1:] if flag != '--update'])} {CMD_FLAGS}"
+flags = f"{' '.join([flag for flag in sys.argv[1:] if flag != '--update-wizard'])} {CMD_FLAGS}"
 
 
 def signal_handler(sig, frame):
@@ -200,27 +200,41 @@ def run_cmd(cmd, assert_success=False, environment=False, capture_output=False,
     return result
 
 
+def get_user_choice(question, options_dict):
+    print()
+    print(question)
+    print()
+
+    for key, value in options_dict.items():
+        print(f"{key}) {value}")
+
+    print()
+
+    choice = input("Input> ").upper()
+    while choice not in options_dict.keys():
+        print("Invalid choice. Please try again.")
+        choice = input("Input> ").upper()
+
+    return choice
+
+
 def install_webui():
 
     # Ask the user for the GPU vendor
     if "GPU_CHOICE" in os.environ:
         choice = os.environ["GPU_CHOICE"].upper()
         print_big_message(f"Selected GPU choice \"{choice}\" based on the GPU_CHOICE environment variable.")
     else:
-        print()
-        print("What is your GPU?")
-        print()
-        print("A) NVIDIA")
-        print("B) AMD (Linux/MacOS only. Requires ROCm SDK 5.6 on Linux)")
-        print("C) Apple M Series")
-        print("D) Intel Arc (IPEX)")
-        print("N) None (I want to run models in CPU mode)")
-        print()
-
-        choice = input("Input> ").upper()
-        while choice not in 'ABCDN':
-            print("Invalid choice. Please try again.")
-            choice = input("Input> ").upper()
+        choice = get_user_choice(
+            "What is your GPU?",
+            {
+                'A': 'NVIDIA',
+                'B': 'AMD (Linux/MacOS only. Requires ROCm SDK 5.6 on Linux)',
+                'C': 'Apple M Series',
+                'D': 'Intel Arc (IPEX)',
+                'N': 'None (I want to run models in CPU mode)'
+            },
+        )
 
     gpu_choice_to_name = {
         "A": "NVIDIA",
@@ -395,15 +409,29 @@ def launch_webui():
     check_env()
 
     parser = argparse.ArgumentParser(add_help=False)
-    parser.add_argument('--update', action='store_true', help='Update the web UI.')
-    parser.add_argument('--install-extensions', action='store_true', help='Install extensions requirements.')
+    parser.add_argument('--update-wizard', action='store_true', help='Launch a menu with update options.')
     args, _ = parser.parse_known_args()
 
-    if args.update:
-        update_requirements()
-    elif args.install_extensions:
-        install_extensions_requirements()
-        update_requirements()
+    if args.update_wizard:
+        choice = get_user_choice(
+            "What would you like to do?",
+            {
+                'A': 'Update the web UI',
+                'B': 'Install/update extensions requirements',
+                'C': 'Revert local changes to repository files with \"git reset --hard\"',
+                'N': 'Nothing (exit).'
+            },
+        )
+
+        if choice == 'A':
+            update_requirements()
+        elif choice == 'B':
+            install_extensions_requirements()
+            update_requirements()
+        elif choice == 'C':
+            run_cmd("git reset --hard", assert_success=True, environment=True)
+        elif choice == 'N':
+            sys.exit()
     else:
         if not is_installed():
             install_webui()

update_macos.sh → update_wizard_linux.sh

@@ -23,4 +23,4 @@ source "$CONDA_ROOT_PREFIX/etc/profile.d/conda.sh" # otherwise conda complains a
 conda activate "$INSTALL_ENV_DIR"
 
 # update installer env
-python one_click.py --update && echo -e "\nDone!"
+python one_click.py --update-wizard && echo -e "\nDone!"

update_linux.sh → update_wizard_macos.sh

@@ -23,4 +23,4 @@ source "$CONDA_ROOT_PREFIX/etc/profile.d/conda.sh" # otherwise conda complains a
 conda activate "$INSTALL_ENV_DIR"
 
 # update installer env
-python one_click.py --update && echo -e "\nDone!"
+python one_click.py --update-wizard && echo -e "\nDone!"

update_windows.bat → update_wizard_windows.sh

@@ -28,7 +28,7 @@ set "CUDA_HOME=%CUDA_PATH%"
 call "%CONDA_ROOT_PREFIX%\condabin\conda.bat" activate "%INSTALL_ENV_DIR%" || ( echo. && echo Miniconda hook not found. && goto end )
 
 @rem update installer env
-call python one_click.py --update && (
+call python one_click.py --update-wizard && (
     echo.
     echo Done!
 )

update_wsl.bat → update_wizard_wsl.sh

@@ -5,7 +5,7 @@ cd /D "%~dp0"
 set PATH=%PATH%;%SystemRoot%\system32
 
 @rem sed -i 's/\x0D$//' ./wsl.sh converts newlines to unix format in the wsl script   calling wsl.sh with 'update' will run updater
-call wsl -e bash -lic "sed -i 's/\x0D$//' ./wsl.sh; source ./wsl.sh update"
+call wsl -e bash -lic "sed -i 's/\x0D$//' ./wsl.sh; source ./wsl.sh update-wizard"
 
 :end
 pause

wsl.sh

@@ -110,7 +110,6 @@ fi
 
 # setup installer env   update env if called with 'wsl.sh update'
 case "$1" in
-("update") python one_click.py --update;;
-("install-extensions") python one_click.py --install-extensions;;
+("update-wizard") python one_click.py --update-wizard;;
 (*) python one_click.py $@;;
 esac