Skip to content

Commit

Permalink
Various one-click installer improvements (oobabooga#4994)
Browse files Browse the repository at this point in the history
---------

Co-authored-by: oobabooga <[email protected]>
  • Loading branch information
matthewraaff and oobabooga authored Jan 5, 2024
1 parent c9d8145 commit c9c31f7
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 29 deletions.
58 changes: 29 additions & 29 deletions one_click.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,15 +185,28 @@ def install_webui():
print("Invalid choice. Please try again.")
choice = input("Input> ").upper()

if choice == "N":
print_big_message("Once the installation ends, make sure to open CMD_FLAGS.txt with\na text editor and add the --cpu flag.")
gpu_choice_to_name = {
"A": "NVIDIA",
"B": "AMD",
"C": "APPLE",
"D": "INTEL",
"N": "NONE"
}

selected_gpu = gpu_choice_to_name[choice]

if selected_gpu == "NONE":
with open(cmd_flags_path, 'r+') as cmd_flags_file:
if "--cpu" not in cmd_flags_file.read():
print_big_message("Adding the --cpu flag to CMD_FLAGS.txt.")
cmd_flags_file.write("\n--cpu")

# Find the proper Pytorch installation command
install_git = "conda install -y -k ninja git"
install_pytorch = "python -m pip install torch torchvision torchaudio"

use_cuda118 = "N"
if any((is_windows(), is_linux())) and choice == "A":
if any((is_windows(), is_linux())) and selected_gpu == "NVIDIA":
if "USE_CUDA118" in os.environ:
use_cuda118 = "Y" if os.environ.get("USE_CUDA118", "").lower() in ("yes", "y", "true", "1", "t", "on") else "N"
else:
Expand All @@ -203,29 +216,30 @@ def install_webui():
while use_cuda118 not in 'YN':
print("Invalid choice. Please try again.")
use_cuda118 = input("Input> ").upper().strip('"\'').strip()

if use_cuda118 == 'Y':
print("CUDA: 11.8")
else:
print("CUDA: 12.1")

install_pytorch = f"python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/{'cu121' if use_cuda118 == 'N' else 'cu118'}"
elif not is_macos() and choice == "B":
elif not is_macos() and selected_gpu == "AMD":
if is_linux():
install_pytorch = "python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6"
else:
print("AMD GPUs are only supported on Linux. Exiting...")
sys.exit(1)
elif is_linux() and (choice == "C" or choice == "N"):
elif is_linux() and selected_gpu in ["APPLE", "NONE"]:
install_pytorch = "python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu"
elif choice == "D":
elif selected_gpu == "INTEL":
install_pytorch = "python -m pip install torch==2.1.0a0 torchvision==0.16.0a0 intel_extension_for_pytorch==2.1.10+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/"

# Install Git and then Pytorch
print_big_message("Installing PyTorch.")
run_cmd(f"{install_git} && {install_pytorch} && python -m pip install py-cpuinfo==9.0.0", assert_success=True, environment=True)

# Install CUDA libraries (this wasn't necessary for Pytorch before...)
if choice == "A":
if selected_gpu == "NVIDIA":
print_big_message("Installing the CUDA runtime libraries.")
run_cmd(f"conda install -y -c \"nvidia/label/{'cuda-12.1.1' if use_cuda118 == 'N' else 'cuda-11.8.0'}\" cuda-runtime", assert_success=True, environment=True)

Expand Down Expand Up @@ -279,29 +293,19 @@ def update_requirements(initial_installation=False):
is_cuda118 = '+cu118' in torver # 2.1.0+cu118
is_cuda117 = '+cu117' in torver # 2.0.1+cu117
is_rocm = '+rocm' in torver # 2.0.1+rocm5.4.2
is_intel = '+cxx11' in torver # 2.0.1a0+cxx11.abi
# is_intel = '+cxx11' in torver # 2.0.1a0+cxx11.abi
is_cpu = '+cpu' in torver # 2.0.1+cpu

if is_rocm:
if cpu_has_avx2():
requirements_file = "requirements_amd.txt"
else:
requirements_file = "requirements_amd_noavx2.txt"
base_requirements = "requirements_amd" + ("_noavx2" if not cpu_has_avx2() else "") + ".txt"
elif is_cpu:
if cpu_has_avx2():
requirements_file = "requirements_cpu_only.txt"
else:
requirements_file = "requirements_cpu_only_noavx2.txt"
base_requirements = "requirements_cpu_only" + ("_noavx2" if not cpu_has_avx2() else "") + ".txt"
elif is_macos():
if is_x86_64():
requirements_file = "requirements_apple_intel.txt"
else:
requirements_file = "requirements_apple_silicon.txt"
base_requirements = "requirements_apple_" + ("intel" if is_x86_64() else "silicon") + ".txt"
else:
if cpu_has_avx2():
requirements_file = "requirements.txt"
else:
requirements_file = "requirements_noavx2.txt"
base_requirements = "requirements" + ("_noavx2" if not cpu_has_avx2() else "") + ".txt"

requirements_file = base_requirements

print_big_message(f"Installing webui requirements from file: {requirements_file}")
print(f"TORCH: {torver}\n")
Expand Down Expand Up @@ -346,10 +350,6 @@ def update_requirements(initial_installation=False):
clear_cache()


def download_model():
run_cmd("python download-model.py", environment=True)


def launch_webui():
run_cmd(f"python server.py {flags}", environment=True)

Expand Down Expand Up @@ -378,7 +378,7 @@ def launch_webui():
if '--model-dir' in flags:
# Splits on ' ' or '=' while maintaining spaces within quotes
flags_list = re.split(' +(?=(?:[^\"]*\"[^\"]*\")*[^\"]*$)|=', flags)
model_dir = [flags_list[(flags_list.index(flag)+1)] for flag in flags_list if flag == '--model-dir'][0].strip('"\'')
model_dir = [flags_list[(flags_list.index(flag) + 1)] for flag in flags_list if flag == '--model-dir'][0].strip('"\'')
else:
model_dir = 'models'

Expand Down
3 changes: 3 additions & 0 deletions setup.cfg
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
[pycodestyle]
max-line-length = 120
ignore = E402, E501, E722

0 comments on commit c9c31f7

Please sign in to comment.