forestcalled/text-generation-webui
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 return torver93 94 95def is_installed():96 site_packages_path = None97 for sitedir in site.getsitepackages():98 if "site-packages" in sitedir and conda_env_path in sitedir:99 site_packages_path = sitedir100 break101 102 if site_packages_path:103 return os.path.isfile(os.path.join(site_packages_path, 'torch', '__init__.py'))104 else:105 return os.path.isdir(conda_env_path)106 107 108def check_env():109 # If we have access to conda, we are probably in an environment110 conda_exist = run_cmd("conda", environment=True, capture_output=True).returncode == 0111 if not conda_exist:112 print("Conda is not installed. Exiting...")113 sys.exit(1)114 115 # Ensure this is a new environment and not the base environment116 if os.environ["CONDA_DEFAULT_ENV"] == "base":117 print("Create an environment for this project and activate it. Exiting...")118 sys.exit(1)119 120 121def clear_cache():122 run_cmd("conda clean -a -y", environment=True)123 run_cmd("python -m pip cache purge", environment=True)124 125 126def print_big_message(message):127 message = message.strip()128 lines = message.split('\n')129 print("\n\n*******************************************************************")130 for line in lines:131 if line.strip() != '':132 print("*", line)133 134 print("*******************************************************************\n\n")135 136 137def calculate_file_hash(file_path):138 p = os.path.join(script_dir, file_path)139 if os.path.isfile(p):140 with open(p, 'rb') as f:141 return hashlib.sha256(f.read()).hexdigest()142 else:143 return ''144 145 146def run_cmd(cmd, assert_success=False, environment=False, capture_output=False, env=None):147 # Use the conda environment148 if environment:149 if is_windows():150 conda_bat_path = os.path.join(script_dir, "installer_files", "conda", "condabin", "conda.bat")151 cmd = "\"" + conda_bat_path + "\" activate \"" + conda_env_path + "\" >nul && " + cmd152 else:153 conda_sh_path = os.path.join(script_dir, "installer_files", "conda", "etc", "profile.d", "conda.sh")154 cmd = ". \"" + conda_sh_path + "\" && conda activate \"" + conda_env_path + "\" && " + cmd155 156 # Run shell commands157 result = subprocess.run(cmd, shell=True, capture_output=capture_output, env=env)158 159 # Assert the command ran successfully160 if assert_success and result.returncode != 0:161 print("Command '" + cmd + "' failed with exit status code '" + str(result.returncode) + "'.\n\nExiting now.\nTry running the start/update script again.")162 sys.exit(1)163 164 return result165 166 167def install_webui():168 # Select your GPU, or choose to run in CPU mode169 if "GPU_CHOICE" in os.environ:170 choice = os.environ["GPU_CHOICE"].upper()171 print_big_message(f"Selected GPU choice \"{choice}\" based on the GPU_CHOICE environment variable.")172 else:173 print()174 print("What is your GPU?")175 print()176 print("A) NVIDIA")177 print("B) AMD (Linux/MacOS only. Requires ROCm SDK 5.6 on Linux)")178 print("C) Apple M Series")179 print("D) Intel Arc (IPEX)")180 print("N) None (I want to run models in CPU mode)")181 print()182 183 choice = input("Input> ").upper()184 while choice not in 'ABCDN':185 print("Invalid choice. Please try again.")186 choice = input("Input> ").upper()187 188 if choice == "N":189 print_big_message("Once the installation ends, make sure to open CMD_FLAGS.txt with\na text editor and add the --cpu flag.")190 191 # Find the proper Pytorch installation command192 install_git = "conda install -y -k ninja git"193 install_pytorch = "python -m pip install torch torchvision torchaudio"194 195 use_cuda118 = "N"196 if any((is_windows(), is_linux())) and choice == "A":197 if "USE_CUDA118" in os.environ:198 use_cuda118 = "Y" if os.environ.get("USE_CUDA118", "").lower() in ("yes", "y", "true", "1", "t", "on") else "N"199 else:200 # Ask for CUDA version if using NVIDIA201 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")202 use_cuda118 = input("Input (Y/N)> ").upper().strip('"\'').strip()203 while use_cuda118 not in 'YN':204 print("Invalid choice. Please try again.")205 use_cuda118 = input("Input> ").upper().strip('"\'').strip()206 if use_cuda118 == 'Y':207 print("CUDA: 11.8")208 else:209 print("CUDA: 12.1")210 211 install_pytorch = f"python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/{'cu121' if use_cuda118 == 'N' else 'cu118'}"212 elif not is_macos() and choice == "B":213 if is_linux():214 install_pytorch = "python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6"215 else:216 print("AMD GPUs are only supported on Linux. Exiting...")217 sys.exit(1)218 elif is_linux() and (choice == "C" or choice == "N"):219 install_pytorch = "python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu"220 elif choice == "D":221 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/"222 223 # Install Git and then Pytorch224 print_big_message("Installing PyTorch.")225 run_cmd(f"{install_git} && {install_pytorch} && python -m pip install py-cpuinfo==9.0.0", assert_success=True, environment=True)226 227 # Install CUDA libraries (this wasn't necessary for Pytorch before...)228 if choice == "A":229 print_big_message("Installing the CUDA runtime libraries.")230 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)231 232 # Install the webui requirements233 update_requirements(initial_installation=True)234 235 236def update_requirements(initial_installation=False):237 # Create .git directory if missing238 if not os.path.isdir(os.path.join(script_dir, ".git")):239 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'240 run_cmd(git_creation_cmd, environment=True, assert_success=True)241 242 files_to_check = [243 'start_linux.sh', 'start_macos.sh', 'start_windows.bat', 'start_wsl.bat',244 'update_linux.sh', 'update_macos.sh', 'update_windows.bat', 'update_wsl.bat',245 'one_click.py'246 ]247 248 before_pull_hashes = {file_name: calculate_file_hash(file_name) for file_name in files_to_check}249 run_cmd("git pull --autostash", assert_success=True, environment=True)250 after_pull_hashes = {file_name: calculate_file_hash(file_name) for file_name in files_to_check}251 252 # Check for differences in installation file hashes253 for file_name in files_to_check:254 if before_pull_hashes[file_name] != after_pull_hashes[file_name]:255 print_big_message(f"File '{file_name}' was updated during 'git pull'. Please run the script again.")256 exit(1)257 258 # Extensions requirements are installed only during the initial install by default.259 # That can be changed with the INSTALL_EXTENSIONS environment variable.260 install = initial_installation261 if "INSTALL_EXTENSIONS" in os.environ:262 install = os.environ["INSTALL_EXTENSIONS"].lower() in ("yes", "y", "true", "1", "t", "on")263 264 if install:265 print_big_message("Installing extensions requirements.")266 skip = ['superbooga', 'superboogav2', 'coqui_tts'] # Fail to install on Windows267 extensions = [foldername for foldername in os.listdir('extensions') if os.path.isfile(os.path.join('extensions', foldername, 'requirements.txt'))]268 extensions = [x for x in extensions if x not in skip]269 for i, extension in enumerate(extensions):270 print(f"\n\n--- [{i+1}/{len(extensions)}]: {extension}\n\n")271 extension_req_path = os.path.join("extensions", extension, "requirements.txt")272 run_cmd("python -m pip install -r " + extension_req_path + " --upgrade", assert_success=False, environment=True)273 elif initial_installation:274 print_big_message("Will not install extensions due to INSTALL_EXTENSIONS environment variable.")275 276 # Detect the Python and PyTorch versions277 torver = torch_version()278 is_cuda = '+cu' in torver279 is_cuda118 = '+cu118' in torver # 2.1.0+cu118280 is_cuda117 = '+cu117' in torver # 2.0.1+cu117281 is_rocm = '+rocm' in torver # 2.0.1+rocm5.4.2282 is_intel = '+cxx11' in torver # 2.0.1a0+cxx11.abi283 is_cpu = '+cpu' in torver # 2.0.1+cpu284 285 if is_rocm:286 if cpu_has_avx2():287 requirements_file = "requirements_amd.txt"288 else:289 requirements_file = "requirements_amd_noavx2.txt"290 elif is_cpu:291 if cpu_has_avx2():292 requirements_file = "requirements_cpu_only.txt"293 else:294 requirements_file = "requirements_cpu_only_noavx2.txt"295 elif is_macos():296 if is_x86_64():297 requirements_file = "requirements_apple_intel.txt"298 else:299 requirements_file = "requirements_apple_silicon.txt"300 else:301 if cpu_has_avx2():302 requirements_file = "requirements.txt"303 else:304 requirements_file = "requirements_noavx2.txt"305 306 print_big_message(f"Installing webui requirements from file: {requirements_file}")307 print(f"TORCH: {torver}\n")308 309 # Prepare the requirements file310 textgen_requirements = open(requirements_file).read().splitlines()311 if is_cuda117:312 textgen_requirements = [req.replace('+cu121', '+cu117').replace('+cu122', '+cu117').replace('torch2.1', 'torch2.0') for req in textgen_requirements]313 elif is_cuda118:314 textgen_requirements = [req.replace('+cu121', '+cu118').replace('+cu122', '+cu118') for req in textgen_requirements]315 if is_windows() and (is_cuda117 or is_cuda118): # No flash-attention on Windows for CUDA 11316 textgen_requirements = [req for req in textgen_requirements if 'jllllll/flash-attention' not in req]317 318 with open('temp_requirements.txt', 'w') as file:319 file.write('\n'.join(textgen_requirements))320 321 # Workaround for git+ packages not updating properly.322 git_requirements = [req for req in textgen_requirements if req.startswith("git+")]323 for req in git_requirements:324 url = req.replace("git+", "")325 package_name = url.split("/")[-1].split("@")[0].rstrip(".git")326 run_cmd("python -m pip uninstall -y " + package_name, environment=True)327 print(f"Uninstalled {package_name}")328 329 # Make sure that API requirements are installed (temporary)330 extension_req_path = os.path.join("extensions", "openai", "requirements.txt")331 if os.path.exists(extension_req_path):332 run_cmd("python -m pip install -r " + extension_req_path + " --upgrade", environment=True)333 334 # Install/update the project requirements335 run_cmd("python -m pip install -r temp_requirements.txt --upgrade", assert_success=True, environment=True)336 os.remove('temp_requirements.txt')337 338 # Check for '+cu' or '+rocm' in version string to determine if torch uses CUDA or ROCm. Check for pytorch-cuda as well for backwards compatibility339 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:340 clear_cache()341 return342 343 if not os.path.exists("repositories/"):344 os.mkdir("repositories")345 346 os.chdir("repositories")347 348 # Install or update ExLlama as needed349 if not os.path.exists("exllama/"):350 run_cmd("git clone https://github.com/turboderp/exllama.git", environment=True)351 else:352 os.chdir("exllama")353 run_cmd("git pull", environment=True)354 os.chdir("..")355 356 if is_linux():357 # Fix JIT compile issue with ExLlama in Linux/WSL358 if not os.path.exists(f"{conda_env_path}/lib64"):359 run_cmd(f'ln -s "{conda_env_path}/lib" "{conda_env_path}/lib64"', environment=True)360 361 # On some Linux distributions, g++ may not exist or be the wrong version to compile GPTQ-for-LLaMa362 gxx_output = run_cmd("g++ -dumpfullversion -dumpversion", environment=True, capture_output=True)363 if gxx_output.returncode != 0 or int(gxx_output.stdout.strip().split(b".")[0]) > 11:364 # Install the correct version of g++365 run_cmd("conda install -y -k conda-forge::gxx_linux-64=11.2.0", environment=True)366 367 clear_cache()368 369 370def download_model():371 run_cmd("python download-model.py", environment=True)372 373 374def launch_webui():375 run_cmd(f"python server.py {flags}", environment=True)376 377 378if __name__ == "__main__":379 # Verifies we are in a conda environment380 check_env()381 382 parser = argparse.ArgumentParser(add_help=False)383 parser.add_argument('--update', action='store_true', help='Update the web UI.')384 args, _ = parser.parse_known_args()385 386 if args.update:387 update_requirements()388 else:389 # If webui has already been installed, skip and run390 if not is_installed():391 install_webui()392 os.chdir(script_dir)393 394 if os.environ.get("LAUNCH_AFTER_INSTALL", "").lower() in ("no", "n", "false", "0", "f", "off"):395 print_big_message("Install finished successfully and will now exit due to LAUNCH_AFTER_INSTALL.")396 sys.exit()397 398 # Check if a model has been downloaded yet399 if '--model-dir' in flags:400 # Splits on ' ' or '=' while maintaining spaces within quotes401 flags_list = re.split(' +(?=(?:[^\"]*\"[^\"]*\")*[^\"]*$)|=', flags)402 model_dir = [flags_list[(flags_list.index(flag)+1)] for flag in flags_list if flag == '--model-dir'][0].strip('"\'')403 else:404 model_dir = 'models'405 406 if len([item for item in glob.glob(f'{model_dir}/*') if not item.endswith(('.txt', '.yaml'))]) == 0:407 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.")408 409 # Workaround for llama-cpp-python loading paths in CUDA env vars even if they do not exist410 conda_path_bin = os.path.join(conda_env_path, "bin")411 if not os.path.exists(conda_path_bin):412 os.mkdir(conda_path_bin)413 414 # Launch the webui415 launch_webui()416