diff --git a/README.md b/README.md index 029106d..ad77dfb 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,5 @@ # ExLlama nodes for ComfyUI -A simple prompt generator for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) utilizing [ExLlama](https://github.com/turboderp/exllama). -Outputs are printed in the console, if you'd like to see them in the UI you can use the `Show Text` node from [here](https://github.com/pythongosssss/ComfyUI-Custom-Scripts). +A simple prompt generator for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) utilizing [ExLlama](https://github.com/turboderp/exllama). ## Installation Clone the repository to `custom_nodes` in your ComfyUI directory and install dependencies: diff --git a/__init__.py b/__init__.py index 44290f5..6fd7ec2 100644 --- a/__init__.py +++ b/__init__.py @@ -9,3 +9,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ZuellniExLlamaGenerator": "ExLlama Generator", "ZuellniExLlamaLoader": "ExLlama Loader", } + +WEB_DIRECTORY = "." diff --git a/nodes.py b/nodes.py index 544178c..703df5c 100644 --- a/nodes.py +++ b/nodes.py @@ -32,6 +32,11 @@ class Generator: RETURN_TYPES = ("STRING",) def generate(self, model, max_tokens, temperature, top_k, top_p, typical_p, penalty, seed, prompt): + prompt = prompt.strip() + + if not prompt: + return ("",) + settings = ExLlamaAltGenerator.Settings() settings.temperature = temperature settings.top_k = top_k @@ -41,7 +46,7 @@ class Generator: torch.manual_seed(seed) stop_conditions = [model.tokenizer.eos_token_id, model.tokenizer.newline_token_id] - model.begin_stream(prompt.strip(), stop_conditions, max_tokens, settings) + model.begin_stream(prompt, stop_conditions, max_tokens, settings) progress = ProgressBar(max_tokens) eos = False @@ -49,12 +54,11 @@ class Generator: while not eos: chunk, eos = model.stream() - text += chunk progress.update(1) + text += chunk text = text.strip() - print("\n[\033[94mExLlama\033[0m]: " + text, end="\n\n") - return (text,) + return {"ui": {"text": text}, "result": (text,)} class Loader: diff --git a/scripts.js b/scripts.js new file mode 100644 index 0000000..ce3771e --- /dev/null +++ b/scripts.js @@ -0,0 +1,49 @@ +import { app } from "../../../scripts/app.js"; +import { ComfyWidgets } from "../../../scripts/widgets.js"; + +app.registerExtension({ + name: "Zuellni.ExLlama.Generator", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "ZuellniExLlamaGenerator") { + const onExecuted = nodeType.prototype.onExecuted; + + nodeType.prototype.onExecuted = function (message) { + onExecuted?.apply(this, arguments); + + if (this.widgets) { + const pos = this.widgets.findIndex((w) => w.name === "text"); + + if (pos !== -1) { + for (let i = pos; i < this.widgets.length; i++) { + this.widgets[i].onRemove?.(); + } + + this.widgets.length = pos; + } + + const string = ["STRING", { multiline: true }] + const widget = ComfyWidgets["STRING"](this, "text", string, app).widget; + + widget.inputEl.readOnly = true; + widget.inputEl.style.opacity = 0.7; + widget.value = message.text.join(""); + + requestAnimationFrame(() => { + const size = this.computeSize(); + + if (size[0] < this.size[0]) { + size[0] = this.size[0]; + } + + if (size[1] < this.size[1]) { + size[1] = this.size[1]; + } + + this.onResize?.(size); + app.graph.setDirtyCanvas(true, false); + }); + } + }; + } + }, +});