From 6e78068c1a14b1b3a5b6024606806bbfbdfcf1f4 Mon Sep 17 00:00:00 2001 From: Zuellni <123005779+Zuellni@users.noreply.github.com> Date: Sat, 7 Oct 2023 10:56:18 +0200 Subject: [PATCH] Split nodes into files, allow loading each separately without cloning the repo, some other minor stuff, install exllamav2 pip package by default --- README.md | 5 +- __init__.py | 22 ++---- exllama.py | 122 ++++++++++++++++++++++++++++++++ nodes.py | 157 ------------------------------------------ requirements.txt | 5 +- scripts.js => text.js | 0 text.py | 67 ++++++++++++++++++ 7 files changed, 198 insertions(+), 180 deletions(-) create mode 100644 exllama.py delete mode 100644 nodes.py rename scripts.js => text.js (100%) create mode 100644 text.py diff --git a/README.md b/README.md index 6d9b238..de48d90 100644 --- a/README.md +++ b/README.md @@ -8,10 +8,7 @@ git clone https://github.com/Zuellni/ComfyUI-ExLlama-Nodes pip install -r requirements.txt ``` -If you see any ExLlama-related 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: -``` -pip install https://github.com/turboderp/exllamav2/releases/download/v0.0.4/exllamav2-0.0.4+cu117-cp310-cp310-win_amd64.whl -``` +If you see any ExLlama-related errors while loading, install it manually following the instructions from [here](https://github.com/turboderp/exllamav2#installation). ## Nodes Name | Description diff --git a/__init__.py b/__init__.py index 8def457..280ef52 100644 --- a/__init__.py +++ b/__init__.py @@ -1,17 +1,9 @@ -from .nodes import Generator, Loader, Previewer, Replacer - -NODE_CLASS_MAPPINGS = { - "ZuellniExLlamaLoader": Loader, - "ZuellniExLlamaGenerator": Generator, - "ZuellniTextPreviewer": Previewer, - "ZuellniTextReplacer": Replacer, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "ZuellniExLlamaLoader": "ExLlama Loader", - "ZuellniExLlamaGenerator": "ExLlama Generator", - "ZuellniTextPreviewer": "Preview Text", - "ZuellniTextReplacer": "Replace Text", -} +from . import exllama, text +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} WEB_DIRECTORY = "." + +for module in (exllama, text): + NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS) + NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS) diff --git a/exllama.py b/exllama.py new file mode 100644 index 0000000..6b780d2 --- /dev/null +++ b/exllama.py @@ -0,0 +1,122 @@ +from time import time + +import torch +from comfy.utils import ProgressBar +from exllamav2 import ExLlamaV2, ExLlamaV2Cache, ExLlamaV2Config, ExLlamaV2Tokenizer +from exllamav2.generator import ExLlamaV2Sampler, ExLlamaV2StreamingGenerator + + +class Loader: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_dir": ("STRING", {"default": ""}), + "max_seq_len": ("INT", {"default": 2048, "min": 1, "max": 8192}), + }, + } + + CATEGORY = "Zuellni/ExLlama" + FUNCTION = "load" + RETURN_NAMES = ("MODEL",) + RETURN_TYPES = ("EL_MODEL",) + + def load(self, model_dir, max_seq_len): + config = ExLlamaV2Config() + config.model_dir = model_dir + config.prepare() + config.max_seq_len = max_seq_len + + model = ExLlamaV2(config) + model.load() + + cache = ExLlamaV2Cache(model) + tokenizer = ExLlamaV2Tokenizer(config) + generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer) + settings = ExLlamaV2Sampler.Settings() + + return ((tokenizer, generator, settings),) + + +class Generator: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("EL_MODEL",), + "stop_on_newline": ((False, True),), + "max_tokens": ("INT", {"default": 128, "min": 1, "max": 8192}), + "temperature": ("FLOAT", {"default": 0.7, "max": 2, "step": 0.01}), + "top_k": ("INT", {"default": 20, "max": 200}), + "top_p": ("FLOAT", {"default": 0.9, "max": 1, "step": 0.01}), + "typical": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}), + "penalty": ("FLOAT", {"default": 1.15, "min": 1, "max": 2, "step": 0.01}), + "seed": ("INT", {"max": 2**64 - 1}), + "text": ("STRING", {"multiline": True}), + }, + } + + CATEGORY = "Zuellni/ExLlama" + FUNCTION = "generate" + RETURN_NAMES = ("TEXT",) + RETURN_TYPES = ("STRING",) + + def generate( + self, + model, + stop_on_newline, + max_tokens, + temperature, + top_k, + top_p, + typical, + penalty, + seed, + text, + ): + tokenizer, generator, settings = model + progress = ProgressBar(max_tokens) + + if not text: + return ("",) + + prompt = tokenizer.encode(text) + stop_conditions = [tokenizer.eos_token_id] + stop_on_newline and stop_conditions.append(tokenizer.newline_token_id) + generator.set_stop_conditions(stop_conditions) + + settings.temperature = temperature + settings.top_k = top_k + settings.top_p = top_p + settings.typical = typical + settings.token_repetition_penalty = penalty + + torch.manual_seed(seed) + generator.begin_stream(prompt, settings) + start = time() + eos = False + output = "" + tokens = 0 + + while not eos and tokens < max_tokens: + chunk, eos, _ = generator.stream() + progress.update(1) + output += chunk + tokens += 1 + + total = round(time() - start, 2) + speed = round(tokens / total, 2) + + print(f"Output generated in {total} seconds ({tokens} tokens, {speed} tokens/s)") + return (output.strip(),) + + +NODE_CLASS_MAPPINGS = { + "ZuellniExLlamaLoader": Loader, + "ZuellniExLlamaGenerator": Generator, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ZuellniExLlamaLoader": "Loader", + "ZuellniExLlamaGenerator": "Generator", +} diff --git a/nodes.py b/nodes.py deleted file mode 100644 index 8d530a5..0000000 --- a/nodes.py +++ /dev/null @@ -1,157 +0,0 @@ -import torch -from comfy.utils import ProgressBar -from exllamav2 import ExLlamaV2, ExLlamaV2Cache, ExLlamaV2Config, ExLlamaV2Tokenizer -from exllamav2.generator import ExLlamaV2Sampler, ExLlamaV2StreamingGenerator - - -class Loader: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "model_dir": ("STRING", {"default": ""}), - "max_seq_len": ("INT", {"default": 2048, "min": 1, "max": 8192}), - }, - } - - CATEGORY = "Zuellni/ExLlama" - FUNCTION = "load" - RETURN_NAMES = ("MODEL",) - RETURN_TYPES = ("EXLLAMA_MODEL",) - - def load(self, model_dir, max_seq_len): - config = ExLlamaV2Config() - config.model_dir = model_dir - config.prepare() - config.max_seq_len = max_seq_len - - model = ExLlamaV2(config) - model.load() - - tokenizer = ExLlamaV2Tokenizer(config) - cache = ExLlamaV2Cache(model) - generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer) - - return (generator,) - - -class Generator: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "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, "max": 2, "step": 0.01}), - "top_k": ("INT", {"default": 20, "min": 0, "max": 200}), - "top_p": ("FLOAT", {"default": 0.9, "min": 0, "max": 1, "step": 0.01}), - "typical": ("FLOAT", {"default": 1, "min": 0, "max": 1, "step": 0.01}), - "penalty": ("FLOAT", {"default": 1.15, "min": 1, "max": 2, "step": 0.01}), - "seed": ("INT", {"default": 0, "min": 0, "max": 2**64 - 1}), - "text": ("STRING", {"default": "", "multiline": True}), - }, - } - - CATEGORY = "Zuellni/ExLlama" - FUNCTION = "generate" - RETURN_NAMES = ("TEXT",) - RETURN_TYPES = ("STRING",) - - def generate( - self, - model, - stop_on_newline, - max_tokens, - temperature, - top_k, - top_p, - typical, - penalty, - seed, - text, - ): - torch.manual_seed(seed) - progress = ProgressBar(max_tokens) - prompt = model.tokenizer.encode(text) - stop_conditions = [model.tokenizer.eos_token_id] - - if stop_on_newline: - stop_conditions += [model.tokenizer.newline_token_id] - - 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(prompt, settings) - eos = False - tokens = 0 - text = "" - - while not eos and tokens < max_tokens: - chunk, eos, _ = model.stream() - progress.update(1) - text += chunk - tokens += 1 - - return (text.strip(),) - - -class Previewer: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "text": ("STRING", {"forceInput": True}), - }, - "hidden": { - "info": "EXTRA_PNGINFO", - "id": "UNIQUE_ID", - }, - } - - CATEGORY = "Zuellni/Text" - FUNCTION = "preview" - OUTPUT_NODE = True - RETURN_TYPES = () - - def preview(self, text, info=None, id=None): - if id and info and "workflow" in info: - nodes = info["workflow"]["nodes"] - node = next((n for n in nodes if str(n["id"]) == id), None) - - if node: - node["widgets_values"] = [text] - - return {"ui": {"text": [text]}} - - -class Replacer: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "text": ("STRING", {"default": "", "multiline": True}), - }, - "optional": { - "a": ("STRING", {"forceInput": True, "multiline": True}), - "b": ("STRING", {"forceInput": True, "multiline": True}), - "c": ("STRING", {"forceInput": True, "multiline": True}), - "d": ("STRING", {"forceInput": True, "multiline": True}), - } - } - - CATEGORY = "Zuellni/Text" - FUNCTION = "replace" - RETURN_NAMES = ("TEXT",) - RETURN_TYPES = ("STRING",) - - def replace(self, text, **vars): - for key, value, in vars.items(): - text = text.replace(f"[{key}]", value) - - return (text,) diff --git a/requirements.txt b/requirements.txt index 4c9e0f9..7aea264 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1 @@ -https://github.com/turboderp/exllamav2/releases/download/v0.0.5/exllamav2-0.0.5+cu121-cp311-cp311-win_amd64.whl; platform_system == "Windows" and python_version == "3.11" -https://github.com/turboderp/exllamav2/releases/download/v0.0.5/exllamav2-0.0.5+cu118-cp310-cp310-win_amd64.whl; platform_system == "Windows" and python_version == "3.10" -https://github.com/turboderp/exllamav2/releases/download/v0.0.5/exllamav2-0.0.5+cu121-cp311-cp311-linux_x86_64.whl; platform_system == "Linux" and python_version == "3.11" -https://github.com/turboderp/exllamav2/releases/download/v0.0.5/exllamav2-0.0.5+cu118-cp310-cp310-linux_x86_64.whl; platform_system == "Linux" and python_version == "3.10" +exllamav2 diff --git a/scripts.js b/text.js similarity index 100% rename from scripts.js rename to text.js diff --git a/text.py b/text.py new file mode 100644 index 0000000..6488c82 --- /dev/null +++ b/text.py @@ -0,0 +1,67 @@ +class Previewer: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"forceInput": True}), + }, + "hidden": { + "info": "EXTRA_PNGINFO", + "id": "UNIQUE_ID", + }, + } + + CATEGORY = "Zuellni/Text" + FUNCTION = "preview" + OUTPUT_NODE = True + RETURN_TYPES = () + + def preview(self, text, info=None, id=None): + if id and info and "workflow" in info: + nodes = info["workflow"]["nodes"] + node = next((n for n in nodes if str(n["id"]) == id), None) + + if node: + node["widgets_values"] = [text] + + return {"ui": {"text": [text]}} + + +class Replacer: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"multiline": True}), + }, + "optional": { + "a": ("STRING", {"forceInput": True, "multiline": True}), + "b": ("STRING", {"forceInput": True, "multiline": True}), + "c": ("STRING", {"forceInput": True, "multiline": True}), + "d": ("STRING", {"forceInput": True, "multiline": True}), + }, + } + + CATEGORY = "Zuellni/Text" + FUNCTION = "replace" + RETURN_NAMES = ("TEXT",) + RETURN_TYPES = ("STRING",) + + def replace(self, text, **vars): + for key, value in vars.items(): + text = text.replace(f"[{key}]", value) + + return (text,) + + +NODE_CLASS_MAPPINGS = { + "ZuellniTextPreviewer": Previewer, + "ZuellniTextReplacer": Replacer, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ZuellniTextPreviewer": "Preview Text", + "ZuellniTextReplacer": "Replace Text", +} + +WEB_DIRECTORY = "."