Replace TypeSafe API with local GGUF decisions via llama.cpp in MultiOutputScript (Jev)

2026-09-26 17:54:57 +09:00
parent 288275c100
commit a977b97d29
4 changed files with 155 additions and 53 deletions
@@ -32,9 +32,6 @@ env/
 *.swo
 *~

-# API Keys
-api_key.txt
-
 # Model files (too large for git)
 *.gguf
 *.safetensors
@@ -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/<your-installation-folder>/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`
@@ -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(
@@ -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