From e069a4d10b0823e4eea524d9ecf1ae2ece110643 Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Thu, 21 Sep 2023 00:33:40 +0200
Subject: [PATCH] Fix and simplify requirements

---
 README.md        |  3 +--
 nodes.py         |  6 +++---
 requirements.txt | 10 +++-------
 3 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 6cc54ae..6d5a37d 100644
--- a/README.md
+++ b/README.md
@@ -8,8 +8,7 @@ git clone https://github.com/Zuellni/ComfyUI-ExLlama-Nodes
 python -m pip install -r requirements.txt
 ```
 
-ExLlama for CUDA 11.8 is installed by default. It should be compatible with the latest portable build of ComfyUI.  
-If you see an error while loading, you should manually install the version matching your system from [here](https://github.com/jllllll/exllama/releases/latest).  
+If you see any errors related to ExLlama while loading, you should manually install the wheel matching your system from [here](https://github.com/jllllll/exllama/releases/latest).  
 For example, on Windows with Python 3.11 and PyTorch CUDA 12.1, you would install:
 ```
 python -m pip install https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu121-cp311-cp311-win_amd64.whl
diff --git a/nodes.py b/nodes.py
index 1444f56..a92f26c 100644
--- a/nodes.py
+++ b/nodes.py
@@ -16,12 +16,12 @@ try:
     from exllama.tokenizer import ExLlamaTokenizer
 except ModuleNotFoundError:
     raise ModuleNotFoundError(
-        f"\n{Fore.RED}ExLlama package not installed. Get {Fore.CYAN}{cu}-{cp}{Fore.RED} from:\n"
+        f"\n{Fore.RED}ExLlama not installed. Get {Fore.CYAN}{cu}-{cp}{Fore.RED} from:\n"
         f"{Fore.MAGENTA}https://github.com/jllllll/exllama/releases/latest{Fore.RESET}"
     )
 except ImportError:
     raise ImportError(
-        f"\n{Fore.RED}Wrong version of ExLlama installed. Get {Fore.CYAN}{cu}-{cp}{Fore.RED} from:\n"
+        f"\n{Fore.RED}Wrong ExLlama version installed. Get {Fore.CYAN}{cu}-{cp}{Fore.RED} from:\n"
         f"{Fore.MAGENTA}https://github.com/jllllll/exllama/releases/latest{Fore.RESET}"
     )
 
@@ -103,7 +103,7 @@ class Generator:
 
         progress.update_absolute(max_tokens)
         text = text.strip()
-        print(Fore.CYAN + text + Fore.RESET + "\n")
+        print(f"{Fore.GREEN}ExLlama: {Fore.CYAN}{text}{Fore.RESET}\n")
 
         return (text,)
 
diff --git a/requirements.txt b/requirements.txt
index 1d8a331..9b53740 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,9 +1,5 @@
-https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp311-cp311-win_amd64.whl; platform_system == "Windows" and python_version == "3.11"
+https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu121-cp311-cp311-win_amd64.whl; platform_system == "Windows" and python_version == "3.11"
 https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp310-cp310-win_amd64.whl; platform_system == "Windows" and python_version == "3.10"
-https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp39-cp39-win_amd64.whl; platform_system == "Windows" and python_version == "3.9"
-https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp38-cp38-win_amd64.whl; platform_system == "Windows" and python_version == "3.8"
-https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp311-cp311-win_amd64.whl; platform_system == "Linux" and python_version == "3.11"
-https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp310-cp310-win_amd64.whl; platform_system == "Linux" and python_version == "3.10"
-https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp39-cp39-win_amd64.whl; platform_system == "Linux" and python_version == "3.9"
-https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp38-cp38-win_amd64.whl; platform_system == "Linux" and python_version == "3.8"
+https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu121-cp311-cp311-linux_x86_64.whl; platform_system == "Linux" and python_version == "3.11"
+https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp310-cp310-linux_x86_64.whl; platform_system == "Linux" and python_version == "3.10"
 sentencepiece