diff --git a/exllama.py b/exllama.py index 1857faa..052c69e 100644 --- a/exllama.py +++ b/exllama.py @@ -18,6 +18,7 @@ from exllamav2.generator import ( ExLlamaV2DynamicJob, ExLlamaV2Sampler, ) +from jinja2 import Template from comfy.model_management import soft_empty_cache, unload_all_models from comfy.utils import ProgressBar @@ -45,7 +46,7 @@ class Loader: return { "required": { "model": (models, {"default": default}), - "cache_bits": (caches, {"default": 16}), + "cache_bits": (caches, {"default": 4}), "fast_tensors": ("BOOLEAN", {"default": True}), "flash_attention": ("BOOLEAN", {"default": True}), "max_seq_len": ("INT", {"default": 2048, "max": 2**20, "step": 256}), @@ -67,6 +68,7 @@ class Loader: def setup(self, model, cache_bits, fast_tensors, flash_attention, max_seq_len): self.unload() self.cache_bits = cache_bits + self.config = ExLlamaV2Config(__class__._MODELS[model]) self.config.fasttensors = fast_tensors self.config.no_flash_attn = not flash_attention @@ -78,24 +80,22 @@ class Loader: self.config.max_input_len = max_seq_len self.config.max_attention_len = max_seq_len**2 + self.tokenizer = ExLlamaV2Tokenizer(self.config) return (self,) def load(self): if ( hasattr(self, "model") and hasattr(self, "cache") - and hasattr(self, "tokenizer") and hasattr(self, "generator") and self.model and self.cache - and self.tokenizer and self.generator ): return self.model = ExLlamaV2(self.config) self.cache = __class__._CACHES[self.cache_bits](self.model) - self.tokenizer = ExLlamaV2Tokenizer(self.config) progress = ProgressBar(len(self.model.modules)) self.model.load_autosplit(self.cache, callback=lambda _, __: progress.update(1)) @@ -113,37 +113,154 @@ class Loader: self.model = None self.cache = None - self.tokenizer = None self.generator = None gc.collect() soft_empty_cache() +class Formatter: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("EXL_MODEL",), + "messages": ("EXL_MESSAGES",), + "add_assistant_role": ("BOOLEAN", {"default": True}), + }, + } + + CATEGORY = _CATEGORY + FUNCTION = "format" + RETURN_NAMES = ("TEXT",) + RETURN_TYPES = ("STRING",) + + def raise_exception(self, message): + raise Exception(message) + + def render(self, template, messages, add_assistant_role): + return ( + template.render( + add_generation_prompt=add_assistant_role, + raise_exception=self.raise_exception, + messages=messages, + bos_token="", + ), + ) + + def format(self, model, messages, add_assistant_role): + template = model.tokenizer.tokenizer_config_dict["chat_template"] + template = Template(template) + + try: + return self.render(template, messages, add_assistant_role) + except: + system = None + merged = [] + + for message in messages: + if message["role"] == "system": + system = {"role": "user", "content": message["content"]} + merged.append(system) + elif system and message["role"] == "user": + index = merged.index(system) + merged[index]["content"] += "\n" + message["content"] + system = None + else: + merged.append(message) + system = None + + return self.render(template, merged, add_assistant_role) + + +class Tokenizer: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("EXL_MODEL",), + "text": ("STRING", {"forceInput": True, "multiline": True}), + "add_bos_token": ("BOOLEAN", {"default": True}), + "encode_special_tokens": ("BOOLEAN", {"default": True}), + }, + } + + CATEGORY = _CATEGORY + FUNCTION = "tokenize" + RETURN_NAMES = ("TOKENS",) + RETURN_TYPES = ("EXL_TOKENS",) + + def tokenize(self, model, text, add_bos_token, encode_special_tokens): + return ( + model.tokenizer.encode( + text=text, + add_bos=add_bos_token, + encode_special_tokens=encode_special_tokens, + ), + ) + + +class Settings: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "temperature": ("FLOAT", {"default": 1, "max": 10, "step": 0.01}), + "penalty": ("FLOAT", {"default": 1, "min": 1, "max": 10, "step": 0.01}), + "top_k": ("INT", {"default": 1, "max": 1000}), + "top_p": ("FLOAT", {"max": 1, "step": 0.01}), + "top_a": ("FLOAT", {"max": 1, "step": 0.01}), + "min_p": ("FLOAT", {"max": 1, "step": 0.01}), + "tfs": ("FLOAT", {"max": 1, "step": 0.01}), + "typical": ("FLOAT", {"max": 1, "step": 0.01}), + "temperature_last": ("BOOLEAN", {"default": True}), + }, + } + + CATEGORY = _CATEGORY + FUNCTION = "set" + RETURN_NAMES = ("SETTINGS",) + RETURN_TYPES = ("EXL_SETTINGS",) + + def set( + self, + temperature, + penalty, + top_k, + top_p, + top_a, + min_p, + tfs, + typical, + temperature_last, + ): + settings = ExLlamaV2Sampler.Settings() + settings.temperature = temperature + settings.token_repetition_penalty = penalty + settings.top_k = top_k + settings.top_p = top_p + settings.top_a = top_a + settings.min_p = min_p + settings.tfs = tfs + settings.typical = typical + settings.temperature_last = temperature_last + return (settings,) + + class Generator: @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("EXL_MODEL",), + "tokens": ("EXL_TOKENS",), "unload": ("BOOLEAN", {"default": False}), - "stop_conditions": ("STRING", {"default": r'["\n"]'}), + "stop_conditions": ("STRING", {"default": r'"\n"'}), "max_tokens": ("INT", {"default": 128, "max": 2**20}), - "temperature": ("FLOAT", {"default": 1, "max": 5, "step": 0.01}), - "top_k": ("INT", {"max": 200}), - "top_p": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}), - "typical_p": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}), - "min_p": ("FLOAT", {"max": 1, "step": 0.01}), - "top_a": ("FLOAT", {"max": 1, "step": 0.01}), - "repetition_penalty": ("FLOAT", {"default": 1, "min": 1, "max": 3, "step": 0.01}), - "temperature_last": ("BOOLEAN", {"default": True}), "seed": ("INT", {"max": 2**64 - 1}), - "text": ("STRING", {"multiline": True}), - }, - "hidden": { - "info": "EXTRA_PNGINFO", - "id": "UNIQUE_ID", }, + "optional": {"settings": ("EXL_SETTINGS",)}, + "hidden": {"info": "EXTRA_PNGINFO", "id": "UNIQUE_ID"}, } CATEGORY = _CATEGORY @@ -154,57 +271,41 @@ class Generator: def generate( self, model, + tokens, unload, stop_conditions, max_tokens, - temperature, - top_k, - top_p, - typical_p, - min_p, - top_a, - repetition_penalty, - temperature_last, seed, - text, + settings=None, info=None, id=None, ): - if not text.strip(): - return ("",) - if unload: unload_all_models() model.unload() model.load() random.seed(seed) - input = model.tokenizer.encode(text, encode_special_tokens=True) - input_len = input.shape[-1] - max_len = model.config.max_seq_len - input_len + tokens_len = tokens.shape[-1] + max_len = model.config.max_seq_len - tokens_len stop = [model.tokenizer.eos_token_id] if not max_tokens or max_tokens > max_len: max_tokens = max_len if stop_conditions.strip(): - stop_conditions = json.loads(stop_conditions) + stop_conditions = json.loads(f"[{stop_conditions}]") stop.extend(stop_conditions) - settings = ExLlamaV2Sampler.Settings() - settings.temperature = temperature - settings.top_k = top_k - settings.top_p = top_p - settings.typical = typical_p - settings.min_p = min_p - settings.top_a = top_a - settings.token_repetition_penalty = repetition_penalty - settings.temperature_last = temperature_last + if not settings: + settings = ExLlamaV2Sampler.Settings() + settings.greedy() job = ExLlamaV2DynamicJob( - input_ids=input, + input_ids=tokens, max_new_tokens=max_tokens, stop_conditions=stop, + gen_settings=settings, ) progress = ProgressBar(max_tokens) @@ -212,7 +313,7 @@ class Generator: start = time() eos = False chunks = [] - tokens = 0 + count = 0 while not eos: for response in model.generator.iterate(): @@ -221,15 +322,15 @@ class Generator: eos = response["eos"] chunks.append(chunk) progress.update(1) - tokens += 1 + count += 1 output = "".join(chunks).strip() total = round(time() - start, 2) - speed = round(tokens / total, 2) + speed = round(count / total, 2) print( f"Output generated in {total} seconds", - f"({input_len} context, {tokens} tokens, {speed}t/s)", + f"({tokens_len} context, {count} tokens, {speed}t/s)", ) if unload: @@ -247,10 +348,16 @@ class Generator: NODE_CLASS_MAPPINGS = { f"{_MAPPING}Loader": Loader, + f"{_MAPPING}Formatter": Formatter, + f"{_MAPPING}Tokenizer": Tokenizer, + f"{_MAPPING}Settings": Settings, f"{_MAPPING}Generator": Generator, } NODE_DISPLAY_NAME_MAPPINGS = { f"{_MAPPING}Loader": "Loader", + f"{_MAPPING}Formatter": "Formatter", + f"{_MAPPING}Tokenizer": "Tokenizer", + f"{_MAPPING}Settings": "Settings", f"{_MAPPING}Generator": "Generator", } diff --git a/text.js b/text.js index 68097f3..40d7af2 100644 --- a/text.js +++ b/text.js @@ -4,7 +4,7 @@ import { ComfyWidgets } from "../../../scripts/widgets.js"; app.registerExtension({ name: "ZuellniText", async beforeRegisterNodeDef(nodeType, nodeData, app) { - if (nodeData.name === "ZuellniTextPreviewer") { + if (nodeData.name === "ZuellniTextPreview") { const onExecuted = nodeType.prototype.onExecuted; nodeType.prototype.onExecuted = function(message) { diff --git a/text.py b/text.py index 346211b..0841c08 100644 --- a/text.py +++ b/text.py @@ -2,7 +2,27 @@ _CATEGORY = "Zuellni/Text" _MAPPING = "ZuellniText" -class Previewer: +class Message: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "role": (("system", "user", "assistant"),), + "content": ("STRING", {"multiline": True}), + }, + "optional": {"messages": ("EXL_MESSAGES",)}, + } + + CATEGORY = _CATEGORY + FUNCTION = "append" + RETURN_NAMES = ("MESSAGES",) + RETURN_TYPES = ("EXL_MESSAGES",) + + def append(self, role, content, messages=[]): + return (messages + [{"role": role, "content": content}],) + + +class Preview: @classmethod def INPUT_TYPES(cls): return {"required": {"text": ("STRING", {"forceInput": True})}} @@ -16,7 +36,7 @@ class Previewer: return {"ui": {"text": [text]}} -class Replacer: +class Replace: @classmethod def INPUT_TYPES(cls): return { @@ -43,13 +63,15 @@ class Replacer: NODE_CLASS_MAPPINGS = { - f"{_MAPPING}Previewer": Previewer, - f"{_MAPPING}Replacer": Replacer, + f"{_MAPPING}Message": Message, + f"{_MAPPING}Preview": Preview, + f"{_MAPPING}Replace": Replace, } NODE_DISPLAY_NAME_MAPPINGS = { - f"{_MAPPING}Previewer": "Preview", - f"{_MAPPING}Replacer": "Replace", + f"{_MAPPING}Message": "Message", + f"{_MAPPING}Preview": "Preview", + f"{_MAPPING}Replace": "Replace", } WEB_DIRECTORY = "."