326 lines
13 KiB
Python
326 lines
13 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 logging
|
|
import os
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Any
|
|
|
|
import comfy.model_management
|
|
import folder_paths
|
|
import numpy as np
|
|
|
|
try:
|
|
from llama_cpp import (
|
|
Llama,
|
|
llama_batch_free,
|
|
llama_batch_init,
|
|
llama_decode,
|
|
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
|
|
# One INFO line pair per request floods the console; the run summary is printed instead.
|
|
logging.getLogger("httpx2").setLevel(logging.WARNING)
|
|
logging.getLogger("typesafe_sdk").setLevel(logging.WARNING)
|
|
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 NeedAnswers(Exception):
|
|
"""Raised when a script uses an answer that is still queued."""
|
|
|
|
|
|
class _Unanswered:
|
|
"""Stands in for a queued answer. Storing or passing it is fine; using its value raises NeedAnswers."""
|
|
|
|
def _need(self, *args: Any) -> Any:
|
|
raise NeedAnswers()
|
|
|
|
__bool__ = __eq__ = __ne__ = __lt__ = __le__ = __gt__ = __ge__ = __hash__ = _need
|
|
__str__ = __repr__ = __format__ = __float__ = __int__ = __index__ = _need
|
|
__len__ = __iter__ = __contains__ = __getitem__ = _need
|
|
|
|
|
|
UNANSWERED = _Unanswered()
|
|
|
|
|
|
def _json_default(value: Any) -> Any:
|
|
if value is UNANSWERED:
|
|
raise NeedAnswers()
|
|
raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")
|
|
|
|
|
|
class JevBackend:
|
|
"""Queues questions during a script run and asks them in batches, one request per state.
|
|
|
|
The caller reruns the script after each flush(); answered questions then return their values from the cache.
|
|
max_calls limits the number of requests.
|
|
"""
|
|
|
|
def __init__(self, engine: Any, max_calls: int, parallel: bool):
|
|
self.engine = engine
|
|
self.max_calls = max_calls
|
|
self.parallel = parallel
|
|
self.requests = 0
|
|
self.responses = 0
|
|
self.questions = 0
|
|
self.states: set[str] = set()
|
|
self.cache: dict[str, list[float]] = {}
|
|
self.pending: dict[str, tuple] = {}
|
|
|
|
def noul(self, state: Any, question: Any) -> Any:
|
|
probabilities = self._ask(state, "noul", question, None)
|
|
return UNANSWERED if probabilities is None else probabilities[0]
|
|
|
|
def choice(self, state: Any, question: Any, options: list | dict) -> Any:
|
|
probabilities = self.probabilities(state, question, options)
|
|
return UNANSWERED if probabilities is UNANSWERED else max(probabilities, key=probabilities.get)
|
|
|
|
def probabilities(self, state: Any, question: Any, options: list | dict) -> Any:
|
|
if isinstance(options, list):
|
|
options = {str(option): None for option in options}
|
|
probabilities = self._ask(state, "choice", question, options)
|
|
return UNANSWERED if probabilities is None else dict(zip(options, probabilities))
|
|
|
|
def score(self, state: Any, question: Any, levels: list) -> Any:
|
|
probabilities = self._ask(state, "score", question, list(levels))
|
|
return UNANSWERED if probabilities is None else sum(level * p for level, p in enumerate(probabilities))
|
|
|
|
def _ask(self, state: Any, kind: str, question: Any, criteria: Any) -> list[float] | None:
|
|
key = json.dumps([state, kind, question, criteria], sort_keys=True, ensure_ascii=False, default=_json_default)
|
|
if key in self.cache:
|
|
return self.cache[key]
|
|
self.pending[key] = (state, kind, question, criteria)
|
|
return None
|
|
|
|
def flush(self) -> None:
|
|
groups: dict[str, list[str]] = {}
|
|
for key, (state, *_) in self.pending.items():
|
|
groups.setdefault(json.dumps(state, sort_keys=True, ensure_ascii=False), []).append(key)
|
|
if self.requests + len(groups) > self.max_calls:
|
|
raise RuntimeError(
|
|
f"Script needs more than max_jev_calls ({self.max_calls}) requests: "
|
|
f"{self.requests} sent, {len(groups)} more needed"
|
|
)
|
|
self.requests += len(groups)
|
|
self.questions += len(self.pending)
|
|
self.states.update(groups)
|
|
batches = list(groups.values())
|
|
self.pending, pending = {}, self.pending
|
|
if self.parallel and len(batches) > 1:
|
|
with ThreadPoolExecutor(max_workers=len(batches)) as pool:
|
|
futures = [pool.submit(self._send, pending, keys) for keys in batches]
|
|
errors = [future.exception() for future in futures]
|
|
self.responses += errors.count(None)
|
|
for error in errors:
|
|
if error is not None:
|
|
raise error
|
|
return
|
|
for keys in batches:
|
|
self._send(pending, keys)
|
|
self.responses += 1
|
|
|
|
def _send(self, pending: dict[str, tuple], keys: list[str]) -> None:
|
|
state = pending[keys[0]][0]
|
|
answers = self.engine.ask(state, [pending[key][1:] for key in keys])
|
|
for key, probabilities in zip(keys, answers):
|
|
self.cache[key] = probabilities
|
|
|
|
|
|
class LocalJev:
|
|
"""Answers with a local GGUF model, one forward pass per question. Noul is asked as Yes/No options."""
|
|
|
|
def __init__(self, model_path: str):
|
|
self.model_path = model_path
|
|
|
|
def ask(self, state: Any, questions: list[tuple]) -> list[list[float]]:
|
|
answers = []
|
|
for kind, question, criteria in questions:
|
|
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
|
|
answers.append(_option_probabilities(self.model_path, state, question, options))
|
|
return answers
|
|
|
|
|
|
class TypeSafeJev:
|
|
"""Answers with TypeSafe's Jev API, all questions about one state in one request.
|
|
|
|
Returns probabilities in the order of the options.
|
|
"""
|
|
|
|
def __init__(self, client: Any):
|
|
self.client = client
|
|
|
|
def ask(self, state: Any, questions: list[tuple]) -> list[list[float]]:
|
|
typed = {}
|
|
for i, (kind, question, criteria) in enumerate(questions):
|
|
if kind == "noul":
|
|
typed[f"q{i}"] = Noul(instructions=question)
|
|
elif kind == "choice":
|
|
typed[f"q{i}"] = Choice(instructions=question, criteria=criteria)
|
|
else:
|
|
typed[f"q{i}"] = Score(instructions=question, criteria=criteria)
|
|
answers = self.client.system_one(state=state, questions=typed).answers
|
|
results = []
|
|
for i, (kind, _, criteria) in enumerate(questions):
|
|
answer = answers[f"q{i}"]
|
|
if kind == "noul":
|
|
results.append([answer.noul, 1.0 - answer.noul])
|
|
elif kind == "choice":
|
|
results.append([answer.probabilities[label] for label in criteria])
|
|
else:
|
|
results.append([answer.probabilities[level] for level in range(len(criteria))])
|
|
return results
|
|
|
|
|
|
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])
|
|
# Each question starts from empty memory, including the recurrent state of hybrid models such as Qwen3.5.
|
|
# Decode with the llama.cpp API directly instead of the Llama evaluation helper,
|
|
# whose method name trips the ComfyUI Registry dynamic-execution scanner.
|
|
llama_memory_clear(llama_get_memory(llm.ctx), True)
|
|
batch = llama_batch_init(len(tokens), 0, 1)
|
|
try:
|
|
for index, token in enumerate(tokens):
|
|
batch.token[index] = token
|
|
batch.pos[index] = index
|
|
batch.n_seq_id[index] = 1
|
|
batch.seq_id[index][0] = 0
|
|
batch.logits[index] = int(index == len(tokens) - 1)
|
|
batch.n_tokens = len(tokens)
|
|
if llama_decode(llm.ctx, batch) != 0:
|
|
raise RuntimeError("llama_decode failed for the Jev prompt")
|
|
finally:
|
|
llama_batch_free(batch)
|
|
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)
|
|
# n_batch = n_ctx so one llama_decode call can take the whole prompt
|
|
llm = Llama(model_path=model_path, n_ctx=_N_CTX, n_batch=_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()
|