From b6f3170fc3c298f08824e1425ca991f66ad850c0 Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Wed, 20 Sep 2023 20:51:30 +0200
Subject: [PATCH] Simplify the installation process a bit (hopefully)

---
 README.md        | 13 +++----------
 nodes.py         | 27 ++++++++++++++++++++++-----
 requirements.txt |  8 ++++++++
 3 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index 27bfce3..0b63303 100644
--- a/README.md
+++ b/README.md
@@ -8,17 +8,10 @@ git clone https://github.com/Zuellni/ComfyUI-ExLlama-Nodes
 python -m pip install -r requirements.txt
 ```
 
-Next, install the latest pre-built ExLlama wheel from https://github.com/jllllll/exllama/releases/latest.  
-Choose the version matching your system, Python, and PyTorch CUDA/ROCm.
-
-Example for Windows with Python 3.10 and CUDA 11.8, which should match the portable ComfyUI build:
+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). For example:
 ```
-python -m pip install https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu118-cp310-cp310-win_amd64.whl
-```
-
-You can use the following command to check what you have:
-```
-python -c "import platform; import torch; print(f'Python: {platform.python_version()}, CUDA: {torch.version.cuda}')"
+python -m pip install https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu121-cp311-cp311-win_amd64.whl
 ```
 
 ## Nodes
diff --git a/nodes.py b/nodes.py
index 03e583e..1444f56 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1,12 +1,29 @@
 from pathlib import Path
+from platform import sys
 
 import torch
+from colorama import Fore
 from comfy.model_management import soft_empty_cache
 from comfy.utils import ProgressBar
-from exllama.alt_generator import ExLlamaAltGenerator
-from exllama.lora import ExLlamaLora
-from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig
-from exllama.tokenizer import ExLlamaTokenizer
+
+cu = "cu" + torch.version.cuda.replace(".", "")
+cp = f"cp{sys.version_info.major}{sys.version_info.minor}"
+
+try:
+    from exllama.alt_generator import ExLlamaAltGenerator
+    from exllama.lora import ExLlamaLora
+    from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig
+    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"{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"{Fore.MAGENTA}https://github.com/jllllll/exllama/releases/latest{Fore.RESET}"
+    )
 
 
 class Generator:
@@ -86,7 +103,7 @@ class Generator:
 
         progress.update_absolute(max_tokens)
         text = text.strip()
-        print(text)
+        print(Fore.CYAN + text + Fore.RESET + "\n")
 
         return (text,)
 
diff --git a/requirements.txt b/requirements.txt
index 080a0d3..1d8a331 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,9 @@
+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+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"
 sentencepiece