Merge 34faa8975e into dc8ed6dbe7

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"
@@ -207,7 +208,7 @@ def get_user_choice(question, options_dict):
     return choice
 
 
-def update_pytorch_and_python():
+def update_pytorch_and_python(selected_gpu=""):
     print_big_message("Checking for PyTorch updates.")
 
     # Update the Python version. Left here for future reference in case this becomes necessary.
@@ -219,8 +220,13 @@ def update_pytorch_and_python():
     torver = torch_version()
     base_cmd = f"python -m pip install --upgrade torch=={TORCH_VERSION} torchvision=={TORCHVISION_VERSION} torchaudio=={TORCHAUDIO_VERSION}"
 
-    if "+cu" in torver:
+    if selected_gpu == "NVIDIA_BLACKWELL":
-        install_cmd = f"{base_cmd} --index-url https://download.pytorch.org/whl/cu124"
+        base_cmd = f"python -m pip install --upgrade torch=={TORCH_VERSION_BLACKWELL} torchvision torchaudio"
+
+    if "+cu" in torver and selected_gpu == "NVIDIA_BLACKWELL":
+        install_cmd = f"{base_cmd} --index-url https://download.pytorch.org/whl/cu128"
+    elif "+cu" in torver:
+        install_cmd = f"{base_cmd} --index-url https://download.pytorch.org/whl/cu124" 
     elif "+rocm" in torver:
         install_cmd = f"{base_cmd} --index-url https://download.pytorch.org/whl/rocm6.2.4"
     elif "+cpu" in torver:
@@ -276,6 +282,7 @@ def install_webui():
                 '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'
             },
         )
@@ -286,6 +293,7 @@ def install_webui():
         "B": "AMD",
         "C": "APPLE",
         "D": "INTEL",
+        "E": "NVIDIA_BLACKWELL",
         "N": "NONE"
     }
 
@@ -302,6 +310,9 @@ def install_webui():
     # Handle CUDA version display
     elif any((is_windows(), is_linux())) and selected_gpu == "NVIDIA":
         print("CUDA: 12.4")
+    
+    elif any((is_windows(), is_linux())) and selected_gpu == "NVIDIA_BLACKWELL":
+        print("CUDA: 12.8")
 
     # No PyTorch for AMD on Windows (?)
     elif is_windows() and selected_gpu == "AMD":
@@ -311,8 +322,13 @@ def install_webui():
     # Find the Pytorch installation command
     install_pytorch = f"python -m pip install torch=={TORCH_VERSION} torchvision=={TORCHVISION_VERSION} torchaudio=={TORCHAUDIO_VERSION} "
 
+    if selected_gpu == "NVIDIA_BLACKWELL":
+        install_pytorch = f"python -m pip install torch=={TORCH_VERSION_BLACKWELL} torchvision torchaudio "
+
     if selected_gpu == "NVIDIA":
         install_pytorch += "--index-url https://download.pytorch.org/whl/cu124"
+    elif selected_gpu == "NVIDIA_BLACKWELL":
+        install_pytorch += "--index-url https://download.pytorch.org/whl/cu128"
     elif selected_gpu == "AMD":
         install_pytorch += "--index-url https://download.pytorch.org/whl/rocm6.2.4"
     elif selected_gpu in ["APPLE", "NONE"]:
@@ -335,10 +351,10 @@ def install_webui():
         run_cmd("conda install -y libuv", environment=True)
 
     # Install the webui requirements
-    update_requirements(initial_installation=True, pull=False)
+    update_requirements(initial_installation=True, pull=False, selected_gpu=selected_gpu)
 
 
-def update_requirements(initial_installation=False, pull=True):
+def update_requirements(initial_installation=False, pull=True, selected_gpu=""):
     # Create .git directory if missing
     if not os.path.exists(os.path.join(script_dir, ".git")):
         run_cmd(
@@ -358,6 +374,8 @@ def update_requirements(initial_installation=False, pull=True):
         file_name = f"requirements_cpu_only{'_noavx2' if not cpu_has_avx2() else ''}.txt"
     elif is_macos():
         file_name = f"requirements_apple_{'intel' if is_x86_64() else 'silicon'}.txt"
+    elif selected_gpu == "NVIDIA_BLACKWELL":
+        file_name = f"requirements_blackwell.txt"
     else:
         file_name = f"requirements{'_noavx2' if not cpu_has_avx2() else ''}.txt"
 
@@ -431,7 +449,7 @@ def update_requirements(initial_installation=False, pull=True):
 
     # Update PyTorch
     if not initial_installation:
-        update_pytorch_and_python()
+        update_pytorch_and_python(selected_gpu=selected_gpu)
         torver = torch_version()
         clean_outdated_pytorch_cuda_dependencies()
 
@@ -455,8 +473,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
+    if selected_gpu == "NVIDIA_BLACKWELL":
-    run_cmd("python -m pip install -r temp_requirements.txt --upgrade", assert_success=True, environment=True)
+        # 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')
@@ -492,6 +513,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)'
                 },
             )
@@ -515,6 +537,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(selected_gpu='NVIDIA_BLACKWELL')
             elif choice == 'N':
                 sys.exit()
     else:

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"