diff --git a/README.md b/README.md
index 7e01392..5e9ed70 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,7 @@ git clone https://github.com/Zuellni/ComfyUI-ExLlama-Nodes
python -m pip install -r requirements.txt
```
-Install the latest pre-built ExLlama wheel from [here](https://github.com/jllllll/exllama/releases/latest).
+Next, install the latest pre-built ExLlama wheel from [here](https://github.com/jllllll/exllama/releases/latest).
Choose the version matching your platform, Python, and PyTorch CUDA/ROCm.
Example for Windows with Python 3.10 and CUDA 11.8 which should match the portable ComfyUI build:
@@ -17,26 +17,13 @@ python -m pip install https://github.com/jllllll/exllama/releases/download/0.0.1
```
## Nodes
-Comes with the following nodes:
-
-### Loader
-Used to load 4-bit GPTQ Llama/2 models. You can find a lot of them over at [Hugging Face](https://huggingface.co/TheBloke).
-
-You should either clone the model repository or download all the files in it manually, then point to the directory in `model_dir`. The `model.safetensors` file on its own is not enough to work.
-
-ExLlama allocates [memory](https://github.com/turboderp/exllama/issues/259) according to `max_seq_len`. Lowering it is a good way to save on GPU RAM.
-It's currently not possible to [offload](https://github.com/turboderp/exllama/issues/177) the models to CPU RAM.
-
-### Generator
-Generates a `string` based on the given `prompt` for use with other nodes.
-Default values correspond to the `simple-1` preset from [text-generation-webui](https://github.com/oobabooga/text-generation-webui).
-ExLlama isn't [deterministic](https://github.com/turboderp/exllama/issues/201), so the outputs may differ even with the same seed.
-
-### Previewer
-Displays generated outputs in the UI.
+Name | Description
+:--- | :---
+Loader | Loads 4-bit GPTQ Llama/2 models. You can find a lot of them on [Hugging Face](https://huggingface.co/TheBloke).
Clone the model repository or download all the files in it to an empty directory, then point to it in `model_dir`. `model.safetensors` won't work on its own.
To load a LoRA specify its directory in `lora_dir`. It should contain `adapter_model.bin` and `adapter_config.json`.
ExLlama allocates memory based on `max_seq_len`. Lowering it is a good way to save on VRAM. It's currently not possible to [offload](https://github.com/turboderp/exllama/issues/177) the model to RAM.
+Generator | Returns a `string` based on the given `prompt` for use with other nodes. Default values correspond to the `simple-1` preset from [text-generation-webui](https://github.com/oobabooga/text-generation-webui).
ExLlama isn't [deterministic](https://github.com/turboderp/exllama/issues/201), so the outputs may differ even with the same seed.
+Previewer | Displays generated outputs in the UI.
## Workflow
-Can be opened directly in ComfyUI.
-Model used: [MythoLogic-Mini-7B](https://huggingface.co/TheBloke/MythoLogic-Mini-7B-GPTQ).
+Can be opened directly in ComfyUI. Model used: [MythoLogic-Mini-7B](https://huggingface.co/TheBloke/MythoLogic-Mini-7B-GPTQ).
diff --git a/nodes.py b/nodes.py
index e0e0c53..a1ca2a4 100644
--- a/nodes.py
+++ b/nodes.py
@@ -4,6 +4,7 @@ import torch
from comfy.utils import ProgressBar
from comfy.model_management import soft_empty_cache
from exllama.alt_generator import ExLlamaAltGenerator
+from exllama.lora import ExLlamaLora
from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig
from exllama.tokenizer import ExLlamaTokenizer
@@ -14,6 +15,7 @@ class Generator:
return {
"required": {
"model": ("GPTQ",),
+ "stop_on_newline": ([False, True], {"default": False}),
"max_tokens": ("INT", {"default": 128, "min": 1, "max": 8192}),
"temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01}),
"top_k": ("INT", {"default": 20, "min": 0, "max": 200}),
@@ -30,8 +32,10 @@ class Generator:
RETURN_NAMES = ("TEXT",)
RETURN_TYPES = ("STRING",)
- def generate(self, model, max_tokens, temperature, top_k, top_p, typical_p, penalty, seed, prompt):
+ def generate(self, model, stop_on_newline, max_tokens, temperature, top_k, top_p, typical_p, penalty, seed, prompt):
+ progress = ProgressBar(max_tokens)
prompt = prompt.strip()
+ torch.manual_seed(seed)
if not prompt:
return ("",)
@@ -42,12 +46,12 @@ class Generator:
settings.top_p = top_p
settings.typical = typical_p
settings.token_repetition_penalty_max = penalty
+ stop_conditions = [model.tokenizer.eos_token_id]
+
+ if stop_on_newline:
+ stop_conditions.append(model.tokenizer.newline_token_id)
- torch.manual_seed(seed)
- stop_conditions = [model.tokenizer.eos_token_id, model.tokenizer.newline_token_id]
model.begin_stream(prompt, stop_conditions, max_tokens, settings)
-
- progress = ProgressBar(max_tokens)
eos = False
text = ""
@@ -56,6 +60,7 @@ class Generator:
progress.update(1)
text += chunk
+ progress.update_absolute(max_tokens)
text = text.strip()
print(text)
return (text,)
@@ -67,6 +72,7 @@ class Loader:
return {
"required": {
"model_dir": ("STRING", {"default": ""}),
+ "lora_dir": ("STRING", {"default": ""}),
"max_seq_len": ("INT", {"default": 2048, "min": 1, "max": 8192}),
},
}
@@ -76,20 +82,26 @@ class Loader:
RETURN_NAMES = ("MODEL",)
RETURN_TYPES = ("GPTQ",)
- def load(self, model_dir, max_seq_len):
- model_dir = Path(model_dir).expanduser()
+ def load(self, model_dir, lora_dir, max_seq_len):
soft_empty_cache()
+ model_dir = Path(model_dir).expanduser()
config = ExLlamaConfig(model_dir / "config.json")
config.model_path = model_dir.glob("*.safetensors")
config.max_seq_len = max_seq_len
model = ExLlama(config)
cache = ExLlamaCache(model)
-
- # sentencepiece requires a string
tokenizer = ExLlamaTokenizer(str(model_dir / "tokenizer.model"))
generator = ExLlamaAltGenerator(model, tokenizer, cache)
+
+ if lora_dir:
+ lora_dir = Path(lora_dir).expanduser()
+ lora_config = lora_dir / "adapter_config.json"
+ lora_model = str(lora_dir / "adapter_model.bin")
+ lora = ExLlamaLora(model, lora_config, lora_model)
+ generator.lora = lora
+
return (generator,)