diff --git a/README.md b/README.md index 78c6bd3..d6b6bf2 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@ # ComfyUI ExLlama Nodes -A simple prompt generator for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) utilizing [ExLlama](https://github.com/turboderp/exllama). +A simple prompt generator for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) utilizing [ExLlamaV2](https://github.com/turboderp/exllamav2). ## Installation Clone the repository to `custom_nodes` in your ComfyUI directory and install dependencies: @@ -8,18 +8,17 @@ git clone https://github.com/Zuellni/ComfyUI-ExLlama-Nodes python -m pip install -r requirements.txt ``` -If you see any ExLlama errors while loading, install it manually from [here](https://github.com/jllllll/exllama/releases/latest).
-For example, on Windows with Python 3.10 and PyTorch CUDA 11.7: +If you see any ExLlama errors while loading, install it manually from [here](https://github.com/turboderp/exllamav2/releases/latest).
+For example, on Windows with Python 3.10 and CUDA 11.7: ``` -python -m pip install https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-0.0.17+cu117-cp310-cp310-win_amd64.whl +python -m pip install https://github.com/turboderp/exllamav2/releases/download/v0.0.4/exllamav2-0.0.4+cu117-cp310-cp310-win_amd64.whl ``` ## Nodes Name | Description :--- | :--- -Loader | Used to load 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 and place them in an empty directory, then specify the path in `model_dir`. The `model.safetensors` file won't work on its own.

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. -LoRA | Used to load LoRAs. The directory should contain `adapter_model.bin` and `adapter_config.json`.
LoRA parameter count has to match the model. -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. +Loader | Used to load 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 and place them in an empty directory, then specify the path in `model_dir`. The `model.safetensors` file won't work on its own.

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 the model to 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, so the outputs may differ even with the same seed. Previewer | Displays generated outputs in the UI and appends them to workflow metadata. ## Workflow diff --git a/__init__.py b/__init__.py index b287d9a..eb7005d 100644 --- a/__init__.py +++ b/__init__.py @@ -1,15 +1,13 @@ -from .nodes import Generator, Loader, Lora, Previewer +from .nodes import Generator, Loader, Previewer NODE_CLASS_MAPPINGS = { "ZuellniExLlamaLoader": Loader, - "ZuellniExLlamaLora": Lora, "ZuellniExLlamaGenerator": Generator, "ZuellniExLlamaPreviewer": Previewer, } NODE_DISPLAY_NAME_MAPPINGS = { "ZuellniExLlamaLoader": "ExLlama Loader", - "ZuellniExLlamaLora": "ExLlama LoRA", "ZuellniExLlamaGenerator": "ExLlama Generator", "ZuellniExLlamaPreviewer": "ExLlama Previewer", } diff --git a/nodes.py b/nodes.py index 9015773..45f1194 100644 --- a/nodes.py +++ b/nodes.py @@ -1,37 +1,7 @@ -from pathlib import Path -from platform import sys - import torch -from colorama import Fore from comfy.utils import ProgressBar - -if not torch.cuda.is_available(): - raise Exception(f"\n{Fore.RED}No CUDA found. ExLlama doesn't support CPU mode.{Fore.RESET}") - -cuda = "" - -if torch.version.hip: - cuda = "rocm" + str(torch.version.hip) -else: - cuda = "cu" + str(torch.version.cuda).replace(".", "") - -pckg = f"{cuda}-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 Exception( - f"\n{Fore.RED}ExLlama not installed. Get {Fore.CYAN}{pckg}{Fore.RED} from" - f"\n{Fore.MAGENTA}https://github.com/jllllll/exllama/releases/latest{Fore.RESET}" - ) -except ImportError: - raise Exception( - f"\n{Fore.RED}Wrong ExLlama version installed. Get {Fore.CYAN}{pckg}{Fore.RED} from" - f"\n{Fore.MAGENTA}https://github.com/jllllll/exllama/releases/latest{Fore.RESET}" - ) +from exllamav2 import ExLlamaV2, ExLlamaV2Cache, ExLlamaV2Config, ExLlamaV2Tokenizer +from exllamav2.generator import ExLlamaV2Sampler, ExLlamaV2StreamingGenerator class Generator: @@ -39,20 +9,17 @@ class Generator: def INPUT_TYPES(cls): return { "required": { - "model": ("GPTQ",), + "model": ("EXLLAMA_MODEL",), "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}), "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01}), - "typical_p": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "typical": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), "penalty": ("FLOAT", {"default": 1.15, "min": 1.0, "max": 2.0, "step": 0.01}), "seed": ("INT", {"default": 0, "min": 0, "max": 2**64 - 1}), "prompt": ("STRING", {"default": "", "multiline": True}), }, - "optional": { - "lora": ("LORA",), - }, } CATEGORY = "Zuellni/ExLlama" @@ -68,46 +35,43 @@ class Generator: temperature, top_k, top_p, - typical_p, + typical, penalty, seed, prompt, - lora=None, ): - progress = ProgressBar(max_tokens) - prompt = prompt.strip() - torch.manual_seed(seed) - if not prompt: return ("",) - settings = ExLlamaAltGenerator.Settings() - settings.temperature = temperature - settings.top_k = top_k - settings.top_p = top_p - settings.typical = typical_p - settings.token_repetition_penalty_max = penalty - settings.lora = lora - + torch.manual_seed(seed) + progress = ProgressBar(max_tokens) + input = model.tokenizer.encode(prompt) stop_conditions = [model.tokenizer.eos_token_id] if stop_on_newline: stop_conditions += [model.tokenizer.newline_token_id] - model.begin_stream(prompt, stop_conditions, max_tokens, settings) + settings = ExLlamaV2Sampler.Settings() + settings.temperature = temperature + settings.top_k = top_k + settings.top_p = top_p + settings.typical = typical + settings.token_repetition_penalty = penalty + + model.set_stop_conditions(stop_conditions) + model.begin_stream(input, settings) eos = False + tokens = 0 text = "" - while not eos: - chunk, eos = model.stream() + while not eos and tokens < max_tokens: + chunk, eos, _ = model.stream() progress.update(1) text += chunk + tokens += 1 progress.update_absolute(max_tokens) - text = text.strip() - print(f"[{Fore.CYAN}ExLlama{Fore.RESET}]: {text}\n") - - return (text,) + return (text.strip(),) class Loader: @@ -123,45 +87,23 @@ class Loader: CATEGORY = "Zuellni/ExLlama" FUNCTION = "load" RETURN_NAMES = ("MODEL",) - RETURN_TYPES = ("GPTQ",) + RETURN_TYPES = ("EXLLAMA_MODEL",) def load(self, model_dir, max_seq_len): - model_dir = Path(model_dir).expanduser() - config = ExLlamaConfig(str(model_dir / "config.json")) - config.model_path = model_dir.glob("*.safetensors") + config = ExLlamaV2Config() config.max_seq_len = max_seq_len + config.model_dir = model_dir + config.prepare() - model = ExLlama(config) - cache = ExLlamaCache(model) - tokenizer = ExLlamaTokenizer(str(model_dir / "tokenizer.model")) - generator = ExLlamaAltGenerator(model, tokenizer, cache) + model = ExLlamaV2(config) + model.load() + tokenizer = ExLlamaV2Tokenizer(config) + cache = ExLlamaV2Cache(model) + generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer) return (generator,) -class Lora: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "model": ("GPTQ",), - "lora_dir": ("STRING", {"default": ""}), - }, - } - - CATEGORY = "Zuellni/ExLlama" - FUNCTION = "load" - RETURN_TYPES = ("LORA",) - - def load(self, model, lora_dir): - lora_dir = Path(lora_dir).expanduser() - lora_config = str(lora_dir / "adapter_config.json") - lora_model = str(lora_dir / "adapter_model.bin") - lora = ExLlamaLora(model.model, lora_config, lora_model) - - return (lora,) - - class Previewer: @classmethod def INPUT_TYPES(cls): @@ -178,8 +120,7 @@ class Previewer: CATEGORY = "Zuellni/ExLlama" FUNCTION = "preview" OUTPUT_NODE = True - RETURN_NAMES = ("TEXT",) - RETURN_TYPES = ("STRING",) + RETURN_TYPES = () def preview(self, text, info=None, id=None): if id and info and "workflow" in info: @@ -189,4 +130,4 @@ class Previewer: if node: node["widgets_values"] = [text] - return {"ui": {"text": [text]}, "result": (text,)} + return {"ui": {"text": [text]}} diff --git a/requirements.txt b/requirements.txt index 3447dfc..23d148c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,4 @@ -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+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" -colorama -sentencepiece +https://github.com/turboderp/exllamav2/releases/download/v0.0.4/exllamav2-0.0.4+cu121-cp311-cp311-win_amd64.whl; platform_system == "Windows" and python_version == "3.11" +https://github.com/turboderp/exllamav2/releases/download/v0.0.4/exllamav2-0.0.4+cu118-cp310-cp310-win_amd64.whl; platform_system == "Windows" and python_version == "3.10" +https://github.com/turboderp/exllamav2/releases/download/v0.0.4/exllamav2-0.0.4+cu121-cp311-cp311-linux_x86_64.whl; platform_system == "Linux" and python_version == "3.11" +https://github.com/turboderp/exllamav2/releases/download/v0.0.4/exllamav2-0.0.4+cu118-cp310-cp310-linux_x86_64.whl; platform_system == "Linux" and python_version == "3.10"