diff --git a/one_click.py b/one_click.py index cccb0dc9..2abd9ebc 100644 --- a/one_click.py +++ b/one_click.py @@ -17,6 +17,7 @@ import sys # Define the required versions TORCH_VERSION = "2.6.0" +TORCH_VERSION_BLACKWELL = "2.7.0" TORCHVISION_VERSION = "0.21.0" TORCHAUDIO_VERSION = "2.6.0" PYTHON_VERSION = "3.11" @@ -119,12 +120,13 @@ def get_gpu_choice(): 'B': 'AMD - Linux/macOS only, requires ROCm 6.2.4', 'C': 'Apple M Series', 'D': 'Intel Arc (beta)', + 'E': 'NVIDIA - CUDA 12.8 - RTX 50XX BLACKWELL', 'N': 'CPU mode' }, ) # Convert choice to GPU name - gpu_choice = {"A": "NVIDIA", "B": "AMD", "C": "APPLE", "D": "INTEL", "N": "NONE"}[choice] + gpu_choice = {"A": "NVIDIA", "B": "AMD", "C": "APPLE", "D": "INTEL", "E": "NVIDIA_BLACKWELL", "N": "NONE"}[choice] # Save choice to state state['gpu_choice'] = gpu_choice @@ -137,8 +139,13 @@ def get_pytorch_install_command(gpu_choice): """Get PyTorch installation command based on GPU choice""" base_cmd = f"python -m pip install torch=={TORCH_VERSION} torchvision=={TORCHVISION_VERSION} torchaudio=={TORCHAUDIO_VERSION} " + if gpu_choice == "NVIDIA_BLACKWELL": + base_cmd = f"python -m pip install torch=={TORCH_VERSION_BLACKWELL} torchvision torchaudio " + if gpu_choice == "NVIDIA": return base_cmd + "--index-url https://download.pytorch.org/whl/cu124" + elif gpu_choice == "NVIDIA_BLACKWELL": + return base_cmd + "--index-url https://download.pytorch.org/whl/cu128" elif gpu_choice == "AMD": return base_cmd + "--index-url https://download.pytorch.org/whl/rocm6.2.4" elif gpu_choice in ["APPLE", "NONE"]: @@ -156,8 +163,13 @@ def get_pytorch_update_command(gpu_choice): """Get PyTorch update command based on GPU choice""" base_cmd = f"python -m pip install --upgrade torch=={TORCH_VERSION} torchvision=={TORCHVISION_VERSION} torchaudio=={TORCHAUDIO_VERSION}" + if gpu_choice == "NVIDIA_BLACKWELL": + base_cmd = f"python -m pip install --upgrade torch=={TORCH_VERSION_BLACKWELL} torchvision torchaudio" + if gpu_choice == "NVIDIA": return f"{base_cmd} --index-url https://download.pytorch.org/whl/cu124" + elif gpu_choice == "NVIDIA_BLACKWELL": + return f"{base_cmd} --index-url https://download.pytorch.org/whl/cu128" elif gpu_choice == "AMD": return f"{base_cmd} --index-url https://download.pytorch.org/whl/rocm6.2.4" elif gpu_choice in ["APPLE", "NONE"]: @@ -181,6 +193,8 @@ def get_requirements_file(gpu_choice): file_name = f"requirements_cpu_only{'_noavx2' if not cpu_has_avx2() else ''}.txt" elif gpu_choice == "NVIDIA": file_name = f"requirements{'_noavx2' if not cpu_has_avx2() else ''}.txt" + elif gpu_choice == "NVIDIA_BLACKWELL": + file_name = f"requirements_blackwell.txt" else: raise ValueError(f"Unknown GPU choice: {gpu_choice}") @@ -328,6 +342,9 @@ def install_webui(): # Handle CUDA version display elif any((is_windows(), is_linux())) and gpu_choice == "NVIDIA": print("CUDA: 12.4") + + elif any((is_windows(), is_linux())) and gpu_choice == "NVIDIA_BLACKWELL": + print("CUDA: 12.8") # No PyTorch for AMD on Windows (?) elif is_windows() and gpu_choice == "AMD": @@ -350,7 +367,7 @@ def install_webui(): update_requirements(initial_installation=True, pull=False) -def update_requirements(initial_installation=False, pull=True): +def update_requirements(initial_installation=False, pull=True, isUpgradeBlackwell=False): # Create .git directory if missing if not os.path.exists(os.path.join(script_dir, ".git")): run_cmd( @@ -369,6 +386,8 @@ def update_requirements(initial_installation=False, pull=True): wheels_changed = True gpu_choice = get_gpu_choice() + if isUpgradeBlackwell: + gpu_choice = "NVIDIA_BLACKWELL" requirements_file = get_requirements_file(gpu_choice) if pull: @@ -448,8 +467,11 @@ def update_requirements(initial_installation=False, pull=True): run_cmd(f"python -m pip uninstall -y {package_name}", environment=True) print(f"Uninstalled {package_name}") - # Install/update the project requirements - run_cmd("python -m pip install -r temp_requirements.txt --upgrade", assert_success=True, environment=True) + if gpu_choice == "NVIDIA_BLACKWELL": + # Install/update the project requirements + run_cmd(f"python -m pip install -r {requirements_file} --upgrade", assert_success=True, environment=True) + else: + run_cmd("python -m pip install -r temp_requirements.txt --upgrade", assert_success=True, environment=True) # Clean up os.remove('temp_requirements.txt') @@ -485,6 +507,7 @@ if __name__ == "__main__": 'A': 'Update the web UI', 'B': 'Install/update extensions requirements', 'C': 'Revert local changes to repository files with \"git reset --hard\"', + 'D': 'Upgrade to NVIDIA BLACKWELL', 'N': 'Nothing (exit)' }, ) @@ -508,6 +531,8 @@ if __name__ == "__main__": update_requirements(pull=False) elif choice == 'C': run_cmd("git reset --hard", assert_success=True, environment=True) + elif choice == 'D': + update_requirements(isUpgradeBlackwell=True) elif choice == 'N': sys.exit() else: diff --git a/requirements/full/requirements_blackwell.txt b/requirements/full/requirements_blackwell.txt new file mode 100644 index 00000000..4a83e158 --- /dev/null +++ b/requirements/full/requirements_blackwell.txt @@ -0,0 +1,46 @@ +accelerate==1.5.* +beautifulsoup4==4.13.4 +bitsandbytes==0.45.* +colorama +datasets +duckduckgo_search==8.0.2 +einops +fastapi==0.112.4 +gradio==4.37.* +jinja2==3.1.6 +markdown +numpy +pandas +peft==0.15.* +Pillow>=9.5.0 +psutil +pydantic==2.8.2 +PyPDF2==3.0.1 +pyyaml +requests +rich +safetensors==0.5.* +scipy +sentencepiece +tensorboard +transformers +tqdm +wandb +triton-windows; platform_system == "Windows" +triton; platform_system == "Linux" + +# API +flask_cloudflared==0.0.14 +sse-starlette==1.6.5 +tiktoken + +# CUDA wheels +https://github.com/oobabooga/llama-cpp-binaries/releases/download/v0.15.0/llama_cpp_binaries-0.15.0+cu124-py3-none-win_amd64.whl; platform_system == "Windows" and python_version == "3.11" +https://github.com/oobabooga/llama-cpp-binaries/releases/download/v0.15.0/llama_cpp_binaries-0.15.0+cu124-py3-none-linux_x86_64.whl; platform_system == "Linux" and platform_machine == "x86_64" and python_version == "3.11" +https://github.com/turboderp-org/exllamav3/releases/download/v0.0.2/exllamav3-0.0.2+cu128.torch2.7.0-cp311-cp311-win_amd64.whl; platform_system == "Windows" and python_version == "3.11" +https://github.com/turboderp-org/exllamav3/releases/download/v0.0.2/exllamav3-0.0.2+cu128.torch2.7.0-cp311-cp311-linux_x86_64.whl; platform_system == "Linux" and platform_machine == "x86_64" and python_version == "3.11" +https://github.com/turboderp-org/exllamav2/releases/download/v0.3.0/exllamav2-0.3.0+cu128.torch2.7.0-cp311-cp311-win_amd64.whl; platform_system == "Windows" and python_version == "3.11" +https://github.com/turboderp-org/exllamav2/releases/download/v0.3.0/exllamav2-0.3.0+cu128.torch2.7.0-cp311-cp311-linux_x86_64.whl; platform_system == "Linux" and platform_machine == "x86_64" and python_version == "3.11" +https://github.com/turboderp-org/exllamav2/releases/download/v0.3.0/exllamav2-0.3.0-py3-none-any.whl; platform_system == "Linux" and platform_machine != "x86_64" +https://github.com/kingbri1/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu128torch2.7.0cxx11abiFALSE-cp311-cp311-win_amd64.whl; platform_system == "Windows" and python_version == "3.11" +https://github.com/kingbri1/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu128torch2.7.0cxx11abiFALSE-cp311-cp311-linux_x86_64.whl; platform_system == "Linux" and platform_machine == "x86_64" and python_version == "3.11"