From 1f3bd2f7d40a135df8d75503c6ec3a7f6e703109 Mon Sep 17 00:00:00 2001 From: kantan-kanto Date: Sat, 3 Oct 2026 14:02:41 +0900 Subject: [PATCH] Batch Jev questions per state and limit max_jev_calls by requests --- __init__.py | 49 ++++++++++++--- jev_backend.py | 167 ++++++++++++++++++++++++++++++++++++------------- 2 files changed, 161 insertions(+), 55 deletions(-) diff --git a/__init__.py b/__init__.py index 6900134..55d0a8a 100644 --- a/__init__.py +++ b/__init__.py @@ -16,8 +16,10 @@ from pathlib import Path from typing import Any, Dict from .jev_backend import ( + UNANSWERED, JevBackend, LocalJev, + NeedAnswers, TypeSafeClient, TypeSafeJev, list_local_models, @@ -194,11 +196,13 @@ class _SafeScriptInterpreter: def __init__( self, variables: Dict[str, Any], + now: datetime.datetime, max_steps: int = _DEFAULT_MAX_STEPS, max_call_depth: int = _DEFAULT_MAX_CALL_DEPTH, jev: JevBackend | None = None, ): self.frames = [variables] + self.now = now self.jev = jev self.max_steps = max_steps self.max_call_depth = max_call_depth @@ -392,6 +396,9 @@ class _SafeScriptInterpreter: if name == "int": return int(*args) if name == "float": + # float() of a queued probability stays queued so storing it does not end the batch. + if len(args) == 1 and args[0] is UNANSWERED: + return UNANSWERED return float(*args) if name == "str": return str(*args) @@ -500,11 +507,11 @@ class _SafeScriptInterpreter: if namespace == "datetime.datetime" and name == "now": if args: raise ValueError("datetime.datetime.now() expects no arguments") - return datetime.datetime.now() + return self.now if namespace == "datetime.date" and name == "today": if args: raise ValueError("datetime.date.today() expects no arguments") - return datetime.date.today() + return self.now.date() if namespace == "math" and name in _SAFE_MATH_FUNCTIONS: return self._call_math(name, args) if namespace == "jev": @@ -516,11 +523,13 @@ class _SafeScriptInterpreter: if len(args) not in (2, 3): raise ValueError("jev.yes() expects 2 or 3 arguments") threshold = args[2] if len(args) == 3 else 0.5 - return self.jev.noul(args[0], args[1]) >= threshold + p = self.jev.noul(args[0], args[1]) + return p if p is UNANSWERED else p >= threshold if name == "noul": if len(args) != 2: raise ValueError("jev.noul() expects 2 arguments") - return _JevProbability(self.jev.noul(*args)) + p = self.jev.noul(*args) + return p if p is UNANSWERED else _JevProbability(p) if name == "choice": if len(args) != 3: raise ValueError("jev.choice() expects 3 arguments") @@ -680,6 +689,8 @@ class _SafeScriptInterpreter: return left > right if isinstance(op, ast.GtE): return left >= right + if isinstance(op, (ast.Is, ast.IsNot)) and (left is UNANSWERED or right is UNANSWERED): + raise NeedAnswers() if isinstance(op, ast.Is): return left is right if isinstance(op, ast.IsNot): @@ -868,7 +879,7 @@ class MultiOutputScript: in_value_3: Any = None, ): return _run_script( - code, None, in_text_1, in_text_2, in_text_3, in_value_1, in_value_2, in_value_3 + code, None, datetime.datetime.now(), in_text_1, in_text_2, in_text_3, in_value_1, in_value_2, in_value_3 ) @@ -926,21 +937,38 @@ class MultiOutputScriptJev(MultiOutputScript): inputs = (in_text_1, in_text_2, in_text_3, in_value_1, in_value_2, in_value_3) if model == _TYPESAFE_API: with _typesafe_client() as client: - return _run_jev_script(code, JevBackend(TypeSafeJev(client), max_jev_calls), inputs) + return _run_jev_script(code, JevBackend(TypeSafeJev(client), max_jev_calls, True), inputs) 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_jev_script(code, JevBackend(LocalJev(model_path), max_jev_calls), inputs) + return _run_jev_script(code, JevBackend(LocalJev(model_path), max_jev_calls, False), inputs) finally: unload_local_model() def _run_jev_script(code: str, jev: JevBackend, inputs: tuple): + """Reruns the script until it finishes, asking the questions queued before each stop in one batch. + + Every rerun starts from the same random state and clock, so it asks the same questions as the run before. + A script error while questions are queued may come from a queued answer, so those are asked and the script reruns. + """ + random_state = random.getstate() + now = datetime.datetime.now() try: - return _run_script(code, jev, *inputs) + while True: + random.setstate(random_state) + try: + return _run_script(code, jev, now, *inputs) + except Exception: + if not jev.pending: + raise + jev.flush() finally: - print(f"[ComfyUI-ScriptFlow] Jev requests: {jev.calls}, responses: {jev.responses}") + print( + f"[ComfyUI-ScriptFlow] Jev requests: {jev.requests}, states: {len(jev.states)}, " + f"questions: {jev.questions}, responses: {jev.responses}" + ) def _typesafe_client() -> Any: @@ -961,6 +989,7 @@ def _typesafe_client() -> Any: def _run_script( code: str, jev: JevBackend | None, + now: datetime.datetime, in_text_1: Any, in_text_2: Any, in_text_3: Any, @@ -990,7 +1019,7 @@ def _run_script( "ov3": None, } - locals_dict = _SafeScriptInterpreter(locals_dict, jev=jev).run(code) + locals_dict = _SafeScriptInterpreter(locals_dict, now, jev=jev).run(code) _validate_outputs(locals_dict) diff --git a/jev_backend.py b/jev_backend.py index eac0c57..8a04c82 100644 --- a/jev_backend.py +++ b/jev_backend.py @@ -15,6 +15,7 @@ import gc import json import logging import os +from concurrent.futures import ThreadPoolExecutor from typing import Any import comfy.model_management @@ -38,7 +39,7 @@ except ImportError: 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("httpx").setLevel(logging.WARNING) + logging.getLogger("httpx2").setLevel(logging.WARNING) logging.getLogger("typesafe_sdk").setLevel(logging.WARNING) except ImportError: TypeSafeClient = None @@ -54,79 +55,155 @@ _N_CTX = 8192 _loaded_model: tuple[str, Any, Any] | None = None -class JevBackend: - """Caches answers and limits model calls for a single script run.""" +class NeedAnswers(Exception): + """Raised when a script uses an answer that is still queued.""" - def __init__(self, engine: Any, max_calls: int): + +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.calls = 0 + self.parallel = parallel + self.requests = 0 self.responses = 0 - self.cache: dict[str, Any] = {} + 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) -> float: - return self._ask(state, "noul", question, None)[0] + 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) -> str: + def choice(self, state: Any, question: Any, options: list | dict) -> Any: probabilities = self.probabilities(state, question, options) - return max(probabilities, key=probabilities.get) + return UNANSWERED if probabilities is UNANSWERED else max(probabilities, key=probabilities.get) - def probabilities(self, state: Any, question: Any, options: list | dict) -> dict[str, float]: + def probabilities(self, state: Any, question: Any, options: list | dict) -> Any: if isinstance(options, list): options = {str(option): None for option in options} - return dict(zip(options, self._ask(state, "choice", question, 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) -> float: + def score(self, state: Any, question: Any, levels: list) -> Any: probabilities = self._ask(state, "score", question, list(levels)) - return sum(level * p for level, p in enumerate(probabilities)) + 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]: - key = json.dumps([state, kind, question, criteria], sort_keys=True, ensure_ascii=False) + 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] - 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.responses += 1 - self.cache[key] = probabilities - return probabilities + 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. Noul is asked as Yes/No options.""" + """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, 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) + 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. Returns probabilities in the order of the options.""" + """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, 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 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]: