From d70b8d9048357f514ddbecd389eab47af5447166 Mon Sep 17 00:00:00 2001 From: mongolu Date: Sat, 23 Sep 2023 04:43:11 +0300 Subject: [PATCH] Added two ENVs in webui.py for docker (#111) --- one_click.py | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/one_click.py b/one_click.py index d88f88fb..3d092c26 100644 --- a/one_click.py +++ b/one_click.py @@ -98,19 +98,24 @@ def run_cmd(cmd, assert_success=False, environment=False, capture_output=False, def install_webui(): - print("What is your GPU") - print() - print("A) NVIDIA") - print("B) AMD (Linux/MacOS only. Requires ROCm SDK 5.4.2/5.4.3 on Linux)") - print("C) Apple M Series") - print("D) None (I want to run models in CPU mode)") - print() - # Select your GPU, or choose to run in CPU mode - choice = input("Input> ").upper() - while choice not in ['A', 'B', 'C', 'D']: - print("Invalid choice. Please try again.") + 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("What is your GPU") + print() + print("A) NVIDIA") + print("B) AMD (Linux/MacOS only. Requires ROCm SDK 5.4.2/5.4.3 on Linux)") + print("C) Apple M Series") + print("D) None (I want to run models in CPU mode)") + print() + choice = input("Input> ").upper() + while choice not in ['A', 'B', 'C', 'D']: + print("Invalid choice. Please try again.") + choice = input("Input> ").upper() + if choice == "D": print_big_message("Once the installation ends, make sure to open CMD_FLAGS.txt with\na text editor and add the --cpu flag.") @@ -261,6 +266,10 @@ if __name__ == "__main__": install_webui() os.chdir(script_dir) + if os.environ.get("LAUNCH_AFTER_INSTALL", "").lower() in ("no", "n", "false", "0", "f", "off"): + print_big_message("Install finished successfully and will now exit due to LAUNCH_AFTER_INSTALL.") + sys.exit() + # Check if a model has been downloaded yet if len([item for item in glob.glob('models/*') if not item.endswith(('.txt', '.yaml'))]) == 0: print_big_message("WARNING: You haven't downloaded any model yet.\nOnce the web UI launches, head over to the \"Model\" tab and download one.")