From a977b97d2991c8b0d03084fb0851a2e44a66fa17 Mon Sep 17 00:00:00 2001 From: kantan-kanto Date: Sat, 26 Sep 2026 17:54:57 +0900 Subject: [PATCH] Replace TypeSafe API with local GGUF decisions via llama.cpp in MultiOutputScript (Jev) --- .gitignore | 3 - README.md | 31 +++++------ __init__.py | 29 +++++----- jev_backend.py | 145 +++++++++++++++++++++++++++++++++++++++++++------ 4 files changed, 155 insertions(+), 53 deletions(-) diff --git a/.gitignore b/.gitignore index b8749dd..0c81f1c 100644 --- a/.gitignore +++ b/.gitignore @@ -32,9 +32,6 @@ env/ *.swo *~ -# API Keys -api_key.txt - # Model files (too large for git) *.gguf *.safetensors diff --git a/README.md b/README.md index e14f6a1..b6a04c0 100644 --- a/README.md +++ b/README.md @@ -77,27 +77,22 @@ ov1, ov2 = (512, 384) if w >= h else (384, 512) ### `MultiOutputScript (Jev)` **Category:** `utils` -Same inputs, outputs, and script rules as `MultiOutputScript`, plus a `jev` namespace that asks [Jev](https://typesafe.ai) (TypeSafe AI's decision model) typed questions about text. +Same inputs, outputs, and script rules as `MultiOutputScript`, plus a `jev` namespace that answers typed questions about text with a local GGUF language model. It works like [Jev](https://typesafe.ai) (TypeSafe AI's decision model) but runs on your machine: the decision is read from the model's next-token probabilities for lettered options in one forward pass, following [SemIf (OpenJev)](https://github.com/TheoLeeCJ/SemIf-OpenJev). No text is generated and nothing is sent over the network. Setup: -1. Install the SDK in the Python environment that runs ComfyUI: - ```bash - pip install typesafe-sdk - ``` -2. Create `api_key.txt` in the actual installation directory of this custom node and add your TypeSafe API key (single line, no quotes): - ``` - ComfyUI/custom_nodes//api_key.txt - ``` - The key is read from this file on every run; the `TYPESAFE_API_KEY` environment variable is not used. `api_key.txt` is listed in `.gitignore`; never commit it. +- Install [llama-cpp-python](https://github.com/abetlen/llama-cpp-python) with GPU support for your platform in the Python environment that runs ComfyUI. +- Put an instruction-tuned GGUF model with a chat template in `ComfyUI/models/LLM` (or `models/text_encoders`). Tested with Qwen3.5-9B Q8_0 and Gemma-4-E4B Q8. -Extra input: -- `max_jev_calls` (INT, default 8): maximum Jev API requests per run. Identical questions within a run are asked only once. +Extra inputs: +- `model`: GGUF model to use (`mmproj-*` files are hidden). +- `max_jev_calls` (INT, default 8): maximum model evaluations per run. Identical questions within a run are evaluated only once. +- `keep_model_loaded` (BOOLEAN, default true): keep the model in memory between runs. Turn it off to free VRAM after each run. -Functions (`state` is a string, dict, or list; Jev accepts text only): +Functions (`state` is a string, dict, or list; text only): - `jev.yes(state, question[, threshold])` → `bool` (yes-probability ≥ `threshold`, default `0.5`) - `jev.noul(state, question)` → yes-probability (`0.0`–`1.0`). Compare it with a threshold; using it directly as a condition raises an error. -- `jev.choice(state, question, options)` → the most likely option (`str`). `options` is a list of labels, or a dict of labels to descriptions. -- `jev.score(state, question, levels)` → expected level (`float`, `0` to `len(levels) - 1`) +- `jev.choice(state, question, options)` → the most likely option (`str`). `options` is a list of 2–16 labels, or a dict of labels to descriptions. +- `jev.score(state, question, levels)` → expected level (`float`, `0` to `len(levels) - 1`) for 2–16 ordered level descriptions ```python # it1: a prompt text @@ -111,9 +106,9 @@ ov2, ov3 = (384, 512) if kind == "portrait" else (512, 384) ``` Notes: -- This node sends `state` and questions to the TypeSafe API over the internet. Do not use it with text you cannot share externally. -- Results are probabilistic; test thresholds on your own inputs. -- `MultiOutputScript` itself never makes network requests; `jev` is not available there. +- Each question takes one forward pass (about 1–1.5 s with a 4B–9B model on GPU); the first question also loads the model. +- Probabilities are uncalibrated and depend on the model; test thresholds on your own inputs. +- `MultiOutputScript` itself never loads a model; `jev` is not available there. ### `centi` **Category:** `utils` diff --git a/__init__.py b/__init__.py index 6dd5ffa..9ae7f9f 100644 --- a/__init__.py +++ b/__init__.py @@ -14,16 +14,14 @@ import ast import datetime import math import random -from pathlib import Path from typing import Any, Dict -from .jev_backend import JevBackend, TypeSafeClient +from .jev_backend import JevBackend, list_local_models, unload_local_model _DEFAULT_MAX_STEPS = 10000 _DEFAULT_MAX_CALL_DEPTH = 32 _MAX_RANGE_SIZE = 10000 -_API_KEY_PATH = Path(__file__).resolve().parent / "api_key.txt" _BLOCKED_AST_NODES = ( ast.Import, @@ -875,13 +873,17 @@ class MultiOutputScriptJev(MultiOutputScript): "kind = jev.choice(it1, \"What orientation suits this prompt?\", [\"portrait\", \"landscape\"])\n" "ov1, ov2 = (384, 512) if kind == \"portrait\" else (512, 384)\n" ) + types["required"]["model"] = (list(list_local_models()),) types["required"]["max_jev_calls"] = ("INT", {"default": 8, "min": 1, "max": 64}) + types["required"]["keep_model_loaded"] = ("BOOLEAN", {"default": True}) return types def run( self, code: str, + model: str, max_jev_calls: int, + keep_model_loaded: bool, in_text_1: Any = None, in_text_2: Any = None, in_text_3: Any = None, @@ -889,21 +891,13 @@ class MultiOutputScriptJev(MultiOutputScript): in_value_2: Any = None, in_value_3: Any = None, ): - if TypeSafeClient is None: - raise RuntimeError("typesafe-sdk is not installed. Run: pip install typesafe-sdk") - if not _API_KEY_PATH.exists(): - raise EnvironmentError( - f"API key file not found: {_API_KEY_PATH}\nPlease create this file with your TypeSafe API key." - ) - api_key = _API_KEY_PATH.read_text(encoding="utf-8").strip() - if not api_key: - raise EnvironmentError( - f'API key is not set in "{_API_KEY_PATH}"\nPlease add your TypeSafe API key to this file.' - ) - with TypeSafeClient(api_key=api_key) as client: + model_path = list_local_models().get(model) + if model_path is None: + raise ValueError(f"GGUF model not found under models/LLM or models/text_encoders: {model}") + try: return _run_script( code, - JevBackend(client, max_jev_calls), + JevBackend(model_path, max_jev_calls), in_text_1, in_text_2, in_text_3, @@ -911,6 +905,9 @@ class MultiOutputScriptJev(MultiOutputScript): in_value_2, in_value_3, ) + finally: + if not keep_model_loaded: + unload_local_model() def _run_script( diff --git a/jev_backend.py b/jev_backend.py index acf0658..6bb640d 100644 --- a/jev_backend.py +++ b/jev_backend.py @@ -5,45 +5,158 @@ # it under the terms of the GNU General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. +# +# The decision readout is adapted from SemIf (formerly OpenJev), +# Copyright (c) 2026 TheoLeeCJ, MIT License: https://github.com/TheoLeeCJ/SemIf-OpenJev from __future__ import annotations import json +import os from typing import Any +import comfy.model_management +import folder_paths +import numpy as np + try: - from typesafe_sdk import Choice, Noul, Score, TypeSafeClient + from llama_cpp import Llama, llama_get_memory, llama_memory_clear + from llama_cpp.llama_chat_format import Jinja2ChatFormatter except ImportError: - TypeSafeClient = None + Llama = None + + +_LETTERS = "ABCDEFGHIJKLMNOP" +_DIRECT_SYSTEM = ( + "Apply the supplied criterion to the supplied evidence. Choose exactly one listed option. " + "Respond with only its uppercase letter, with no explanation or reasoning." +) +_N_CTX = 8192 + +_loaded_model: tuple[str, Any, Any] | None = None class JevBackend: - """Asks Jev one question per call. Scoped to a single script run.""" + """Caches answers and limits model calls for a single script run.""" - def __init__(self, client: Any, max_calls: int): - self.client = client + def __init__(self, model_path: str, max_calls: int): + self.model_path = model_path self.max_calls = max_calls self.calls = 0 self.cache: dict[str, Any] = {} - def noul(self, state: Any, question: str) -> float: - return self._ask(state, Noul(instructions=question)).noul + def noul(self, state: Any, question: Any) -> float: + return self._ask(state, question, ["Yes", "No"])[0] - def choice(self, state: Any, question: str, options: list | dict) -> str: + def choice(self, state: Any, question: Any, options: list | dict) -> str: if isinstance(options, list): options = {str(option): None for option in options} - return self._ask(state, Choice(instructions=question, criteria=options)).choice + labels = list(options) + descriptions = [label if description is None else f"{label}: {description}" for label, description in options.items()] + probabilities = self._ask(state, question, descriptions) + return labels[probabilities.index(max(probabilities))] - def score(self, state: Any, question: str, levels: list) -> float: - return self._ask(state, Score(instructions=question, criteria=levels)).score + def score(self, state: Any, question: Any, levels: list) -> float: + probabilities = self._ask(state, question, list(levels)) + return sum(level * p for level, p in enumerate(probabilities)) - def _ask(self, state: Any, question: Any) -> Any: - key = json.dumps([state, question.model_dump()], sort_keys=True, ensure_ascii=False) + def _ask(self, state: Any, question: Any, options: list) -> list[float]: + key = json.dumps([state, question, options], sort_keys=True, ensure_ascii=False) if key in self.cache: return self.cache[key] if self.calls >= self.max_calls: raise RuntimeError(f"Script exceeded max_jev_calls ({self.max_calls})") self.calls += 1 - answer = self.client.system_one(state=state, questions={"q": question}).answers["q"] - self.cache[key] = answer - return answer + probabilities = _option_probabilities(self.model_path, state, question, options) + self.cache[key] = probabilities + return probabilities + + +def _option_probabilities(model_path: str, state: Any, question: Any, options: list) -> list[float]: + """Reads one decision from the next-token logits of the option letters, like SemIf's direct path.""" + if not 2 <= len(options) <= len(_LETTERS): + raise ValueError(f"Jev questions need 2-{len(_LETTERS)} options, got {len(options)}") + llm, formatter = load_local_model(model_path) + payload = { + "evidence": state, + "criterion": question, + "options": [{"letter": letter, "description": option} for letter, option in zip(_LETTERS, options)], + } + messages = [ + {"role": "system", "content": _DIRECT_SYSTEM}, + {"role": "user", "content": json.dumps(payload, ensure_ascii=False)}, + ] + prompt = formatter(messages=messages, enable_thinking=False).prompt + tokens = llm.tokenize(prompt.encode("utf-8"), add_bos=False, special=True) + if len(tokens) >= llm.n_ctx(): + raise ValueError(f"Jev prompt is {len(tokens)} tokens, over the {llm.n_ctx()} token context") + slots = [] + for letter in _LETTERS[: len(options)]: + encoded = llm.tokenize(letter.encode("utf-8"), add_bos=False, special=False) + if len(encoded) != 1: + raise ValueError(f"Model tokenizer does not encode {letter!r} as one token") + slots.append(encoded[0]) + # reset() alone keeps the recurrent state of hybrid models such as Qwen3.5 + llama_memory_clear(llama_get_memory(llm.ctx), True) + llm.reset() + llm.eval(tokens) + logits = llm.scores[0, slots].astype(np.float64) + weights = np.exp(logits - logits.max()) + return (weights / weights.sum()).tolist() + + +def _model_dirs() -> list[str]: + dirs = [] + for key in ("LLM", "llm", "text_encoders"): + try: + dirs.extend(folder_paths.get_folder_paths(key)) + except KeyError: + pass + dirs.append(os.path.join(folder_paths.models_dir, "LLM")) + return list(dict.fromkeys(os.path.normpath(d) for d in dirs)) + + +def list_local_models() -> dict[str, str]: + """GGUF language models under models/LLM and models/text_encoders, keyed by display name.""" + models = {} + for root in _model_dirs(): + if not os.path.isdir(root): + continue + for current_dir, _, files in os.walk(root): + for file_name in files: + if file_name.endswith(".gguf") and not file_name.startswith("mmproj"): + path = os.path.join(current_dir, file_name) + models[os.path.relpath(path, folder_paths.models_dir).replace("\\", "/")] = path + return dict(sorted(models.items(), key=lambda item: item[0].lower())) + + +def load_local_model(model_path: str) -> tuple[Any, Any]: + global _loaded_model + if _loaded_model is not None and _loaded_model[0] == model_path: + return _loaded_model[1], _loaded_model[2] + if Llama is None: + raise RuntimeError("llama-cpp-python is not installed. See the README for installation.") + unload_local_model() + device = comfy.model_management.get_torch_device() + comfy.model_management.free_memory(os.path.getsize(model_path) * 1.2, device) + llm = Llama(model_path=model_path, n_ctx=_N_CTX, n_gpu_layers=-1, verbose=False) + template = llm.metadata.get("tokenizer.chat_template") + if not template: + llm.close() + raise ValueError(f"GGUF has no chat template: {model_path}") + bos, eos = llm.token_bos(), llm.token_eos() + formatter = Jinja2ChatFormatter( + template=template, + bos_token=llm.detokenize([bos], special=True).decode("utf-8") if bos != -1 else "", + eos_token=llm.detokenize([eos], special=True).decode("utf-8") if eos != -1 else "", + ) + _loaded_model = (model_path, llm, formatter) + return llm, formatter + + +def unload_local_model() -> None: + global _loaded_model + if _loaded_model is None: + return + _loaded_model[1].close() + _loaded_model = None