222 lines
8.8 KiB
Python
222 lines
8.8 KiB
Python
# ComfyUI-ScriptFlow
|
|
# Copyright (C) 2026 kantan-kanto (https://github.com/kantan-kanto)
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# 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 gc
|
|
import json
|
|
import os
|
|
from typing import Any
|
|
|
|
import comfy.model_management
|
|
import folder_paths
|
|
import numpy as np
|
|
|
|
try:
|
|
from llama_cpp import Llama, llama_get_logits_ith, llama_get_memory, llama_memory_clear
|
|
from llama_cpp.llama_chat_format import Jinja2ChatFormatter
|
|
except ImportError:
|
|
Llama = None
|
|
|
|
try:
|
|
from typesafe_sdk import Choice, Noul, Score, TypeSafeClient
|
|
except ImportError:
|
|
TypeSafeClient = 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:
|
|
"""Caches answers and limits model calls for a single script run."""
|
|
|
|
def __init__(self, engine: Any, max_calls: int):
|
|
self.engine = engine
|
|
self.max_calls = max_calls
|
|
self.calls = 0
|
|
self.cache: dict[str, Any] = {}
|
|
|
|
def noul(self, state: Any, question: Any) -> float:
|
|
return self._ask(state, "noul", question, None)[0]
|
|
|
|
def choice(self, state: Any, question: Any, options: list | dict) -> str:
|
|
probabilities = self.probabilities(state, question, options)
|
|
return max(probabilities, key=probabilities.get)
|
|
|
|
def probabilities(self, state: Any, question: Any, options: list | dict) -> dict[str, float]:
|
|
if isinstance(options, list):
|
|
options = {str(option): None for option in options}
|
|
return dict(zip(options, self._ask(state, "choice", question, options)))
|
|
|
|
def score(self, state: Any, question: Any, levels: list) -> float:
|
|
probabilities = self._ask(state, "score", question, list(levels))
|
|
return sum(level * p for level, p in enumerate(probabilities))
|
|
|
|
def _ask(self, state: Any, kind: str, question: Any, criteria: Any) -> list[float]:
|
|
key = json.dumps([state, kind, question, criteria], 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
|
|
probabilities = self.engine.ask(state, kind, question, criteria)
|
|
self.cache[key] = probabilities
|
|
return probabilities
|
|
|
|
|
|
class LocalJev:
|
|
"""Answers with a local GGUF model. Noul is asked as Yes/No options."""
|
|
|
|
def __init__(self, model_path: str):
|
|
self.model_path = model_path
|
|
|
|
def ask(self, state: Any, kind: str, question: Any, criteria: Any) -> list[float]:
|
|
if kind == "noul":
|
|
options = ["Yes", "No"]
|
|
elif kind == "choice":
|
|
options = [label if description is None else f"{label}: {description}" for label, description in criteria.items()]
|
|
else:
|
|
options = criteria
|
|
return _option_probabilities(self.model_path, state, question, options)
|
|
|
|
|
|
class TypeSafeJev:
|
|
"""Answers with TypeSafe's Jev API. Returns probabilities in the order of the options."""
|
|
|
|
def __init__(self, client: Any):
|
|
self.client = client
|
|
|
|
def ask(self, state: Any, kind: str, question: Any, criteria: Any) -> list[float]:
|
|
if kind == "noul":
|
|
p = self._answer(state, Noul(instructions=question)).noul
|
|
return [p, 1.0 - p]
|
|
if kind == "choice":
|
|
probabilities = self._answer(state, Choice(instructions=question, criteria=criteria)).probabilities
|
|
return [probabilities[label] for label in criteria]
|
|
probabilities = self._answer(state, Score(instructions=question, criteria=criteria)).probabilities
|
|
return [probabilities[level] for level in range(len(criteria))]
|
|
|
|
def _answer(self, state: Any, question: Any) -> Any:
|
|
return self.client.system_one(state=state, questions={"q": question}).answers["q"]
|
|
|
|
|
|
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)
|
|
# Llama.scores is not filled without logits_all in upstream llama-cpp-python, so read the context directly
|
|
vocabulary = np.ctypeslib.as_array(llama_get_logits_ith(llm.ctx, -1), shape=(llm.n_vocab(),))
|
|
logits = vocabulary[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[_display_name(path)] = path
|
|
return dict(sorted(models.items(), key=lambda item: item[0].lower()))
|
|
|
|
|
|
def _display_name(path: str) -> str:
|
|
# models registered through extra_model_paths.yaml may be outside models_dir or on another drive
|
|
try:
|
|
rel = os.path.relpath(path, folder_paths.models_dir)
|
|
except ValueError:
|
|
return path
|
|
return path if rel.startswith("..") else rel.replace("\\", "/")
|
|
|
|
|
|
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:
|
|
# Same cleanup as ComfyUI-LLM-Session's Unload LLM Model: drop the model, collect, then empty the device cache
|
|
global _loaded_model
|
|
if _loaded_model is None:
|
|
return
|
|
llm = _loaded_model[1]
|
|
_loaded_model = None
|
|
llm.close()
|
|
del llm
|
|
gc.collect()
|
|
comfy.model_management.soft_empty_cache()
|