sleepdeep/llm-deploy
0
1import argparse2import glob3import hashlib4import os5import platform6import re7import signal8import site9import subprocess10import sys11 12script_dir = os.getcwd()13conda_env_path = os.path.join(script_dir, "installer_files", "env")14 15# Remove the '# ' from the following lines as needed for your AMD GPU on Linux16# os.environ["ROCM_PATH"] = '/opt/rocm'17# os.environ["HSA_OVERRIDE_GFX_VERSION"] = '10.3.0'18# os.environ["HCC_AMDGPU_TARGET"] = 'gfx1030'19 20# Command-line flags21cmd_flags_path = os.path.join(script_dir, "CMD_FLAGS.txt")22if os.path.exists(cmd_flags_path):23 with open(cmd_flags_path, 'r') as f:24 CMD_FLAGS = ' '.join(line.strip().rstrip('\\').strip() for line in f if line.strip().rstrip('\\').strip() and not line.strip().startswith('#'))25else:26 CMD_FLAGS = ''27 28flags = f"{' '.join([flag for flag in sys.argv[1:] if flag != '--update'])} {CMD_FLAGS}"29 30 31def signal_handler(sig, frame):32 sys.exit(0)33 34 35signal.signal(signal.SIGINT, signal_handler)36 37 38def is_linux():39 return sys.platform.startswith("linux")40 41 42def is_windows():43 return sys.platform.startswith("win")44 45 46def is_macos():47 return sys.platform.startswith("darwin")48 49 50def is_x86_64():51 return platform.machine() == "x86_64"52 53 54def cpu_has_avx2():55 try:56 import cpuinfo57 58 info = cpuinfo.get_cpu_info()59 if 'avx2' in info['flags']:60 return True61 else:62 return False63 except:64 return True65 66 67def cpu_has_amx():68 try:69 import cpuinfo70 71 info = cpuinfo.get_cpu_info()72 if 'amx' in info['flags']:73 return True74 else:75 return False76 except:77 return True78 79 80def torch_version():81 site_packages_path = None82 for sitedir in site.getsitepackages():83 if "site-packages" in sitedir and conda_env_path in sitedir:84 site_packages_path = sitedir85 break86 87 if site_packages_path:88 torch_version_file = open(os.path.join(site_packages_path, 'torch', 'version.py')).read().splitlines()89 torver = [line for line in torch_version_file if '__version__' in line][0].split('__version__ = ')[1].strip("'")90 else:91 from torch import __version__ as torver92 93 return torver94 95 96def is_installed():97 site_packages_path = None98 for sitedir in site.getsitepackages():99 if "site-packages" in sitedir and conda_env_path in sitedir:100 site_packages_path = sitedir101 break102 103 if site_packages_path:104 return os.path.isfile(os.path.join(site_packages_path, 'torch', '__init__.py'))105 else:106 return os.path.isdir(conda_env_path)107 108 109def check_env():110 # If we have access to conda, we are probably in an environment111 conda_exist = run_cmd("conda", environment=True, capture_output=True).returncode == 0112 if not conda_exist:113 print("Conda is not installed. Exiting...")114 sys.exit(1)115 116 # Ensure this is a new environment and not the base environment117 if os.environ["CONDA_DEFAULT_ENV"] == "base":118 print("Create an environment for this project and activate it. Exiting...")119 sys.exit(1)120 121 122def clear_cache():123 run_cmd("conda clean -a -y", environment=True)124 run_cmd("python -m pip cache purge", environment=True)125 126 127def print_big_message(message):128 message = message.strip()129 lines = message.split('\n')130 print("\n\n*******************************************************************")131 for line in lines:132 if line.strip() != '':133 print("*", line)134 135 print("*******************************************************************\n\n")136 137 138def calculate_file_hash(file_path):139 p = os.path.join(script_dir, file_path)140 if os.path.isfile(p):141 with open(p, 'rb') as f:142 return hashlib.sha256(f.read()).hexdigest()143 else:144 return ''145 146 147def run_cmd(cmd, assert_success=False, environment=False, capture_output=False, env=None):148 # Use the conda environment149 if environment:150 if is_windows():151 conda_bat_path = os.path.join(script_dir, "installer_files", "conda", "condabin", "conda.bat")152 cmd = "\"" + conda_bat_path + "\" activate \"" + conda_env_path + "\" >nul && " + cmd153 else:154 conda_sh_path = os.path.join(script_dir, "installer_files", "conda", "etc", "profile.d", "conda.sh")155 cmd = ". \"" + conda_sh_path + "\" && conda activate \"" + conda_env_path + "\" && " + cmd156 157 # Run shell commands158 result = subprocess.run(cmd, shell=True, capture_output=capture_output, env=env)159 160 # Assert the command ran successfully161 if assert_success and result.returncode != 0:162 print("Command '" + cmd + "' failed with exit status code '" + str(result.returncode) + "'.\n\nExiting now.\nTry running the start/update script again.")163 sys.exit(1)164 165 return result166 167 168def install_webui():169 # Select your GPU, or choose to run in CPU mode170 if "GPU_CHOICE" in os.environ:171 choice = os.environ["GPU_CHOICE"].upper()172 print_big_message(f"Selected GPU choice \"{choice}\" based on the GPU_CHOICE environment variable.")173 else:174 print()175 print("What is your GPU?")176 print()177 print("A) NVIDIA")178 print("B) AMD (Linux/MacOS only. Requires ROCm SDK 5.6 on Linux)")179 print("C) Apple M Series")180 print("D) Intel Arc (IPEX)")181 print("N) None (I want to run models in CPU mode)")182 print()183 184 choice = input("Input> ").upper()185 while choice not in 'ABCDN':186 print("Invalid choice. Please try again.")187 choice = input("Input> ").upper()188 189 gpu_choice_to_name = {190 "A": "NVIDIA",191 "B": "AMD",192 "C": "APPLE",193 "D": "INTEL",194 "N": "NONE"195 }196 197 selected_gpu = gpu_choice_to_name[choice]198 199 if selected_gpu == "NONE":200 with open(cmd_flags_path, 'r+') as cmd_flags_file:201 if "--cpu" not in cmd_flags_file.read():202 print_big_message("Adding the --cpu flag to CMD_FLAGS.txt.")203 cmd_flags_file.write("\n--cpu")204 205 # Find the proper Pytorch installation command206 install_git = "conda install -y -k ninja git"207 install_pytorch = "python -m pip install torch==2.1.* torchvision==0.16.* torchaudio==2.1.* "208 209 use_cuda118 = "N"210 if any((is_windows(), is_linux())) and selected_gpu == "NVIDIA":211 if "USE_CUDA118" in os.environ:212 use_cuda118 = "Y" if os.environ.get("USE_CUDA118", "").lower() in ("yes", "y", "true", "1", "t", "on") else "N"213 else:214 # Ask for CUDA version if using NVIDIA215 print("\nDo you want to use CUDA 11.8 instead of 12.1? Only choose this option if your GPU is very old (Kepler or older).\nFor RTX and GTX series GPUs, say \"N\". If unsure, say \"N\".\n")216 use_cuda118 = input("Input (Y/N)> ").upper().strip('"\'').strip()217 while use_cuda118 not in 'YN':218 print("Invalid choice. Please try again.")219 use_cuda118 = input("Input> ").upper().strip('"\'').strip()220 221 if use_cuda118 == 'Y':222 print("CUDA: 11.8")223 install_pytorch += "--index-url https://download.pytorch.org/whl/cu118"224 else:225 print("CUDA: 12.1")226 install_pytorch += "--index-url https://download.pytorch.org/whl/cu121"227 elif not is_macos() and selected_gpu == "AMD":228 if is_linux():229 install_pytorch += "--index-url https://download.pytorch.org/whl/rocm5.6"230 else:231 print("AMD GPUs are only supported on Linux. Exiting...")232 sys.exit(1)233 elif is_linux() and selected_gpu in ["APPLE", "NONE"]:234 install_pytorch += "--index-url https://download.pytorch.org/whl/cpu"235 elif selected_gpu == "INTEL":236 install_pytorch = "python -m pip install torch==2.1.0a0 torchvision==0.16.0a0 torchaudio==2.1.0a0 intel-extension-for-pytorch==2.1.10 --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/"237 238 # Install Git and then Pytorch239 print_big_message("Installing PyTorch.")240 run_cmd(f"{install_git} && {install_pytorch} && python -m pip install py-cpuinfo==9.0.0", assert_success=True, environment=True)241 242 # Install CUDA libraries (this wasn't necessary for Pytorch before...)243 if selected_gpu == "NVIDIA":244 print_big_message("Installing the CUDA runtime libraries.")245 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)246 247 if selected_gpu == "INTEL":248 # Install oneAPI dependencies via conda249 print_big_message("Installing Intel oneAPI runtime libraries.")250 run_cmd("conda install -y -c intel dpcpp-cpp-rt=2024.0 mkl-dpcpp=2024.0")251 # Install libuv required by Intel-patched torch252 run_cmd("conda install -y libuv")253 254 # Install the webui requirements255 update_requirements(initial_installation=True)256 257 258def update_requirements(initial_installation=False):259 # Create .git directory if missing260 if not os.path.isdir(os.path.join(script_dir, ".git")):261 git_creation_cmd = 'git init -b main && git remote add origin https://github.com/oobabooga/text-generation-webui && git fetch && git symbolic-ref refs/remotes/origin/HEAD refs/remotes/origin/main && git reset --hard origin/main && git branch --set-upstream-to=origin/main'262 run_cmd(git_creation_cmd, environment=True, assert_success=True)263 264 files_to_check = [265 'start_linux.sh', 'start_macos.sh', 'start_windows.bat', 'start_wsl.bat',266 'update_linux.sh', 'update_macos.sh', 'update_windows.bat', 'update_wsl.bat',267 'one_click.py'268 ]269 270 before_pull_hashes = {file_name: calculate_file_hash(file_name) for file_name in files_to_check}271 run_cmd("git pull --autostash", assert_success=True, environment=True)272 after_pull_hashes = {file_name: calculate_file_hash(file_name) for file_name in files_to_check}273 274 # Check for differences in installation file hashes275 for file_name in files_to_check:276 if before_pull_hashes[file_name] != after_pull_hashes[file_name]:277 print_big_message(f"File '{file_name}' was updated during 'git pull'. Please run the script again.")278 exit(1)279 280 # Extensions requirements are installed only during the initial install by default.281 # That can be changed with the INSTALL_EXTENSIONS environment variable.282 install = initial_installation283 if "INSTALL_EXTENSIONS" in os.environ:284 install = os.environ["INSTALL_EXTENSIONS"].lower() in ("yes", "y", "true", "1", "t", "on")285 286 if install:287 print_big_message("Installing extensions requirements.")288 skip = ['superbooga', 'superboogav2', 'coqui_tts'] # Fail to install on Windows289 extensions = [foldername for foldername in os.listdir('extensions') if os.path.isfile(os.path.join('extensions', foldername, 'requirements.txt'))]290 extensions = [x for x in extensions if x not in skip]291 for i, extension in enumerate(extensions):292 print(f"\n\n--- [{i+1}/{len(extensions)}]: {extension}\n\n")293 extension_req_path = os.path.join("extensions", extension, "requirements.txt")294 run_cmd("python -m pip install -r " + extension_req_path + " --upgrade", assert_success=False, environment=True)295 elif initial_installation:296 print_big_message("Will not install extensions due to INSTALL_EXTENSIONS environment variable.")297 298 # Detect the Python and PyTorch versions299 torver = torch_version()300 is_cuda = '+cu' in torver301 is_cuda118 = '+cu118' in torver # 2.1.0+cu118302 is_cuda117 = '+cu117' in torver # 2.0.1+cu117303 is_rocm = '+rocm' in torver # 2.0.1+rocm5.4.2304 is_intel = '+cxx11' in torver # 2.0.1a0+cxx11.abi305 is_cpu = '+cpu' in torver # 2.0.1+cpu306 307 if is_rocm:308 base_requirements = "requirements_amd" + ("_noavx2" if not cpu_has_avx2() else "") + ".txt"309 elif is_cpu or is_intel:310 base_requirements = "requirements_cpu_only" + ("_noavx2" if not cpu_has_avx2() else "") + ".txt"311 elif is_macos():312 base_requirements = "requirements_apple_" + ("intel" if is_x86_64() else "silicon") + ".txt"313 else:314 base_requirements = "requirements" + ("_noavx2" if not cpu_has_avx2() else "") + ".txt"315 316 requirements_file = base_requirements317 318 print_big_message(f"Installing webui requirements from file: {requirements_file}")319 print(f"TORCH: {torver}\n")320 321 # Prepare the requirements file322 textgen_requirements = open(requirements_file).read().splitlines()323 if is_cuda117:324 textgen_requirements = [req.replace('+cu121', '+cu117').replace('+cu122', '+cu117').replace('torch2.1', 'torch2.0') for req in textgen_requirements]325 elif is_cuda118:326 textgen_requirements = [req.replace('+cu121', '+cu118').replace('+cu122', '+cu118') for req in textgen_requirements]327 if is_windows() and (is_cuda117 or is_cuda118): # No flash-attention on Windows for CUDA 11328 textgen_requirements = [req for req in textgen_requirements if 'jllllll/flash-attention' not in req]329 330 with open('temp_requirements.txt', 'w') as file:331 file.write('\n'.join(textgen_requirements))332 333 # Workaround for git+ packages not updating properly.334 git_requirements = [req for req in textgen_requirements if req.startswith("git+")]335 for req in git_requirements:336 url = req.replace("git+", "")337 package_name = url.split("/")[-1].split("@")[0].rstrip(".git")338 run_cmd("python -m pip uninstall -y " + package_name, environment=True)339 print(f"Uninstalled {package_name}")340 341 # Make sure that API requirements are installed (temporary)342 extension_req_path = os.path.join("extensions", "openai", "requirements.txt")343 if os.path.exists(extension_req_path):344 run_cmd("python -m pip install -r " + extension_req_path + " --upgrade", environment=True)345 346 # Install/update the project requirements347 run_cmd("python -m pip install -r temp_requirements.txt --upgrade", assert_success=True, environment=True)348 os.remove('temp_requirements.txt')349 350 # Check for '+cu' or '+rocm' in version string to determine if torch uses CUDA or ROCm. Check for pytorch-cuda as well for backwards compatibility351 if not any((is_cuda, is_rocm)) and run_cmd("conda list -f pytorch-cuda | grep pytorch-cuda", environment=True, capture_output=True).returncode == 1:352 clear_cache()353 return354 355 if not os.path.exists("repositories/"):356 os.mkdir("repositories")357 358 clear_cache()359 360 361def launch_webui():362 run_cmd(f"python server.py {flags}", environment=True)363 364 365if __name__ == "__main__":366 # Verifies we are in a conda environment367 check_env()368 369 parser = argparse.ArgumentParser(add_help=False)370 parser.add_argument('--update', action='store_true', help='Update the web UI.')371 args, _ = parser.parse_known_args()372 373 if args.update:374 update_requirements()375 else:376 # If webui has already been installed, skip and run377 if not is_installed():378 install_webui()379 os.chdir(script_dir)380 381 if os.environ.get("LAUNCH_AFTER_INSTALL", "").lower() in ("no", "n", "false", "0", "f", "off"):382 print_big_message("Install finished successfully and will now exit due to LAUNCH_AFTER_INSTALL.")383 sys.exit()384 385 # Check if a model has been downloaded yet386 if '--model-dir' in flags:387 # Splits on ' ' or '=' while maintaining spaces within quotes388 flags_list = re.split(' +(?=(?:[^\"]*\"[^\"]*\")*[^\"]*$)|=', flags)389 model_dir = [flags_list[(flags_list.index(flag) + 1)] for flag in flags_list if flag == '--model-dir'][0].strip('"\'')390 else:391 model_dir = 'models'392 393 if len([item for item in glob.glob(f'{model_dir}/*') if not item.endswith(('.txt', '.yaml'))]) == 0:394 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.")395 396 # Workaround for llama-cpp-python loading paths in CUDA env vars even if they do not exist397 conda_path_bin = os.path.join(conda_env_path, "bin")398 if not os.path.exists(conda_path_bin):399 os.mkdir(conda_path_bin)400 401 # Launch the webui402 launch_webui()403