From 0bd03f4daee2ed3cf8b1abdbfea68a894ff83c3f Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Mon, 24 Jun 2024 12:33:30 +0200
Subject: [PATCH] Clean up imports

---
 README.md  |  2 +-
 exllama.py | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index a4577ec..eadcaa9 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,7 @@ git clone https://github.com/Zuellni/ComfyUI-ExLlama-Nodes custom_nodes/ComfyUI-
 pip install -r custom_nodes/ComfyUI-ExLlamaV2-Nodes/requirements.txt
 ```
 
-Use wheels for [ExLlamaV2](https://github.com/turboderp/exllamav2/releases/latest) and [Flash Attention](https://github.com/bdashore3/flash-attention/releases/latest) on Windows:
+Use wheels for [ExLlamaV2](https://github.com/turboderp/exllamav2/releases/latest) and [FlashAttention](https://github.com/bdashore3/flash-attention/releases/latest) on Windows:
 ```
 pip install exllamav2-X.X.X+cuXXX.torch2.X.X-cp3XX-cp3XX-win_amd64.whl
 pip install flash_attn-X.X.X+cuXXX.torch2.X.X-cp3XX-cp3XX-win_amd64.whl
diff --git a/exllama.py b/exllama.py
index 37ef347..11fb1a0 100644
--- a/exllama.py
+++ b/exllama.py
@@ -4,8 +4,20 @@ import random
 from pathlib import Path
 from time import time
 
-from exllamav2 import *
-from exllamav2.generator import *
+from exllamav2 import (
+    ExLlamaV2,
+    ExLlamaV2Cache,
+    ExLlamaV2Cache_Q4,
+    ExLlamaV2Cache_Q6,
+    ExLlamaV2Cache_Q8,
+    ExLlamaV2Config,
+    ExLlamaV2Tokenizer,
+)
+from exllamav2.generator import (
+    ExLlamaV2DynamicGenerator,
+    ExLlamaV2DynamicJob,
+    ExLlamaV2Sampler,
+)
 
 from comfy.model_management import soft_empty_cache, unload_all_models
 from comfy.utils import ProgressBar