Split some nodes, added a format node that uses tokenizer config template

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)
+        tokens_len = tokens.shape[-1]
-        input_len = input.shape[-1]
+        max_len = model.config.max_seq_len - tokens_len
-        max_len = model.config.max_seq_len - input_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()
+        if not settings:
-        settings.temperature = temperature
+            settings = ExLlamaV2Sampler.Settings()
-        settings.top_k = top_k
+            settings.greedy()
-        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
 
         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",
 }

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) {

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}Message": Message,
-    f"{_MAPPING}Replacer": Replacer,
+    f"{_MAPPING}Preview": Preview,
+    f"{_MAPPING}Replace": Replace,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    f"{_MAPPING}Previewer": "Preview",
+    f"{_MAPPING}Message": "Message",
-    f"{_MAPPING}Replacer": "Replace",
+    f"{_MAPPING}Preview": "Preview",
+    f"{_MAPPING}Replace": "Replace",
 }
 
 WEB_DIRECTORY = "."