Add a lora loader
This commit is contained in:
Zuellni
@@ -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).<br>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.<br><br>To load a LoRA specify its directory in `lora_dir`. It should contain `adapter_model.bin` and `adapter_config.json`.<br><br>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).<br><br>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).
|
||||
|
||||
|
||||
|
||||
@@ -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,)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user