From 43be1be598e40a8dd2ad0db84ae8ceba8fbbb3a8 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Thu, 12 Oct 2023 21:02:44 -0700
Subject: [PATCH] Manually install CUDA runtime libraries

---
 README.md    | 11 ++++++-----
 one_click.py |  4 ++++
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index a4c807b9..f834e6c0 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,3 @@
-**Breaking change: WebUI now uses PyTorch 2.1.**
-
-* For one-click installer users: If you encounter problems after updating, rerun the update script. If issues persist, delete the `installer_files` folder and use the start script to reinstall requirements.
-* For manual installations, update PyTorch with the [provided command](https://github.com/oobabooga/text-generation-webui/#2-install-pytorch).
-
 # Text generation web UI
 
 A Gradio web UI for Large Language Models.
@@ -100,6 +95,12 @@ conda activate textgen
 
 The up-to-date commands can be found here: https://pytorch.org/get-started/locally/. 
 
+For NVIDIA, you may also need to manually install the CUDA runtime libraries:
+
+```
+conda install -y -c "nvidia/label/cuda-11.8.0" cuda-runtime
+```
+
 #### 3. Install the web UI
 
 ```
diff --git a/one_click.py b/one_click.py
index 8e1f06dc..2da76055 100644
--- a/one_click.py
+++ b/one_click.py
@@ -187,6 +187,10 @@ def install_webui():
     # Install Git and then 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":
+        run_cmd("conda install -y -c \"nvidia/label/cuda-11.8.0\" cuda-runtime", assert_success=True, environment=True)
+
     # Install the webui requirements
     update_requirements(initial_installation=True)