Vito
55fa2c4870
* fix(usage): count cached re-runs of runtime-linked prompt text (#347) A PromptManager-family node whose text comes from a runtime link (PreviewAny, an LLM, a list node) counted only the job that executed it; later jobs served from ComfyUI's output cache counted nothing. Each execution now records a receipt (digests of its STRING outputs and hashes of the text it stores) per run key and queue item. On execution_cached the receipt whose digests equal the cache entry's text is replayed once for the job, with the node's role and what the queue hook already saw recomputed from the job's own graph. The executing node and the replay count through one shared own_counts. Unmatched cache entries count nothing, as before; a node that executes after its replay in the same job skips its own count. No schema change, no backfill. * fix(usage): never replay a node twice in one job; dedupe receipts once complete (#347) A repeated execution_cached id or message for the same job now counts the node once. Receipts replace equal older ones only when their job finishes (or the next job starts), so a list still recording [A] of [A, B] no longer drops an older [A] receipt. Node-level tests check the receipt each execution records. * test(comfyui): count #347's cached-run cases in the execution suite pin (#347) The pin listed 19 execution-text-link cases per project; #347 adds five node kinds and eight one-renderer cases, so the suite lists 32 per project and 118 in total with #332's image-attribution case from main.
342 lines
10 KiB
Python
342 lines
10 KiB
Python
"""Deterministic nodes and queue controls installed only in an owned test host."""
|
|
|
|
import asyncio
|
|
import os
|
|
import threading
|
|
import time
|
|
|
|
import nodes
|
|
import execution
|
|
import torch
|
|
|
|
from aiohttp import web
|
|
import comfy.model_management
|
|
from server import PromptServer
|
|
|
|
_GATE = threading.Event()
|
|
_GATE.set()
|
|
_ENTERED = threading.Event()
|
|
|
|
|
|
@PromptServer.instance.routes.post("/pm_execution_fixture/gate/{action}")
|
|
async def control_gate(request: web.Request) -> web.Response:
|
|
action = request.match_info["action"]
|
|
if action == "hold":
|
|
_ENTERED.clear()
|
|
_GATE.clear()
|
|
elif action == "release":
|
|
_GATE.set()
|
|
else:
|
|
raise web.HTTPBadRequest()
|
|
return web.json_response({"held": not _GATE.is_set()})
|
|
|
|
|
|
@PromptServer.instance.routes.get("/pm_execution_fixture/gate")
|
|
async def gate_status(request: web.Request) -> web.Response:
|
|
return web.json_response({"entered": _ENTERED.is_set(), "held": not _GATE.is_set()})
|
|
|
|
|
|
class PMHistoryCached:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {"required": {"value": ("STRING", {"default": "deterministic"})}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "testing/execution-history"
|
|
OUTPUT_NODE = True
|
|
|
|
def execute(self, value: str) -> dict:
|
|
return {"ui": {"text": [value]}, "result": (value,)}
|
|
|
|
|
|
class PMHistoryFailure(PMHistoryCached):
|
|
def execute(self, value: str) -> dict:
|
|
raise ValueError(value)
|
|
|
|
|
|
class PMHistoryGate(PMHistoryCached):
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs: object) -> float:
|
|
return float("nan")
|
|
|
|
def execute(self, value: str) -> dict:
|
|
_ENTERED.set()
|
|
deadline = time.monotonic() + 30
|
|
while not _GATE.wait(0.05):
|
|
comfy.model_management.throw_exception_if_processing_interrupted()
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError("Owned execution fixture gate timed out")
|
|
return super().execute(value)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PMHistoryCached": PMHistoryCached,
|
|
"PMHistoryFailure": PMHistoryFailure,
|
|
"PMHistoryGate": PMHistoryGate,
|
|
}
|
|
|
|
# Only external inference is replaced; PM classes and host execution remain real.
|
|
|
|
|
|
class _Clip:
|
|
def tokenize(self, text: str) -> str:
|
|
return text
|
|
|
|
def encode_from_tokens_scheduled(self, tokens: str) -> list:
|
|
return [[tokens, {}]]
|
|
|
|
|
|
class PMHistoryInputs:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {"required": {}}
|
|
|
|
RETURN_TYPES = ("CLIP", "MODEL", "LATENT")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "testing/execution-history"
|
|
|
|
def execute(self) -> tuple:
|
|
return (_Clip(), object(), {"samples": "owned deterministic latent"})
|
|
|
|
|
|
class PMHistoryOutput:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {"value": ("LATENT",)},
|
|
"optional": {"mode": (["normal", "failure", "gate"],)},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "execute"
|
|
CATEGORY = "testing/execution-history"
|
|
OUTPUT_NODE = True
|
|
|
|
def execute(self, value: dict, mode: str = "normal") -> tuple:
|
|
if mode == "failure":
|
|
raise ValueError("owned downstream failure")
|
|
if mode == "gate":
|
|
PMHistoryGate().execute("owned downstream gate")
|
|
return ()
|
|
|
|
|
|
def _sample(self: object, **kwargs: object) -> tuple:
|
|
return (kwargs["latent_image"],)
|
|
|
|
|
|
nodes.KSampler.sample = _sample
|
|
NODE_CLASS_MAPPINGS.update(
|
|
{"PMHistoryInputs": PMHistoryInputs, "PMHistoryOutput": PMHistoryOutput}
|
|
)
|
|
|
|
|
|
class PMHistoryList(PMHistoryCached):
|
|
OUTPUT_IS_LIST = (True,)
|
|
|
|
def execute(self, value: str) -> tuple:
|
|
return ([value, value],)
|
|
|
|
|
|
class PMHistoryTransform:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {"required": {"conditioning": ("CONDITIONING",)}}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "testing/execution-history"
|
|
|
|
def execute(self, conditioning: list) -> tuple:
|
|
return (conditioning,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS.update(
|
|
{"PMHistoryList": PMHistoryList, "PMHistoryTransform": PMHistoryTransform}
|
|
)
|
|
|
|
|
|
class PMHistoryImages:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"value": ("INT", {"default": 0, "min": 0, "max": 255}),
|
|
"batch": ("INT", {"default": 1, "min": 1, "max": 2}),
|
|
},
|
|
"optional": {"latent": ("LATENT",)},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "testing/execution-history"
|
|
|
|
def execute(self, value: int, batch: int, latent: dict | None = None) -> tuple:
|
|
return (torch.full((batch, 8, 8, 3), value / 255.0, device="cpu"),)
|
|
|
|
|
|
class PMHistoryImageOutput:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"mode": (["failure", "gate"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "execute"
|
|
CATEGORY = "testing/execution-history"
|
|
OUTPUT_NODE = True
|
|
|
|
def execute(self, images: torch.Tensor, mode: str) -> tuple:
|
|
if mode == "failure":
|
|
raise ValueError("owned failure after real image output")
|
|
PMHistoryGate().execute("owned gate after real image output")
|
|
return ()
|
|
|
|
|
|
NODE_CLASS_MAPPINGS.update(
|
|
{"PMHistoryImages": PMHistoryImages, "PMHistoryImageOutput": PMHistoryImageOutput}
|
|
)
|
|
|
|
# #347: text whose value (and cache signature) follows a value the test sets, so
|
|
# one queued graph can serve A, then B, then A again (each cached under LRU).
|
|
_CYCLED = {"value": "cycled"}
|
|
|
|
|
|
class PMCycledText:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict:
|
|
return {"required": {}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "testing/execution-history"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs: object) -> str:
|
|
return _CYCLED["value"]
|
|
|
|
def execute(self) -> tuple:
|
|
return (_CYCLED["value"],)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS["PMCycledText"] = PMCycledText
|
|
|
|
|
|
@PromptServer.instance.routes.post("/pm_execution_fixture/cycled")
|
|
async def cycled_text(request: web.Request) -> web.Response:
|
|
value = (await request.json())["value"]
|
|
if type(value) is not str:
|
|
raise web.HTTPBadRequest()
|
|
_CYCLED["value"] = value
|
|
return web.json_response({"value": value})
|
|
|
|
|
|
@PromptServer.instance.routes.post("/pm_execution_fixture/metadata")
|
|
async def metadata(request: web.Request) -> web.Response:
|
|
disabled = (await request.json())["disabled"]
|
|
if type(disabled) is not bool:
|
|
raise web.HTTPBadRequest()
|
|
nodes.args.disable_metadata = disabled
|
|
return web.json_response({"disabled": disabled})
|
|
|
|
|
|
_INJECTION_ORIGINALS = {}
|
|
|
|
|
|
@PromptServer.instance.routes.post("/pm_execution_fixture/injection")
|
|
async def injection(request: web.Request) -> web.Response:
|
|
suffix = (await request.json())["suffix"]
|
|
if not isinstance(suffix, str) or len(suffix) > 128:
|
|
raise web.HTTPBadRequest()
|
|
for name in ("PromptManager", "PromptManagerText"):
|
|
cls = nodes.NODE_CLASS_MAPPINGS[name]
|
|
if name not in _INJECTION_ORIGINALS:
|
|
_INJECTION_ORIGINALS[name] = cls._inject_lora_trigger_words
|
|
cls._inject_lora_trigger_words = (
|
|
(lambda self, text: text + suffix) if suffix else _INJECTION_ORIGINALS[name]
|
|
)
|
|
return web.json_response({"configured": True})
|
|
|
|
|
|
# Apply a requested host cache configuration on its own worker before execution.
|
|
# Retain only the mode scalar; never retain the worker/executor/cache objects.
|
|
|
|
_CACHE_MODE = None
|
|
_EXECUTE = execution.PromptExecutor.execute
|
|
|
|
|
|
def _execute_with_cache_mode(self: object, *args: object, **kwargs: object) -> object:
|
|
global _CACHE_MODE
|
|
if _CACHE_MODE is not None:
|
|
self.cache_type = execution.CacheType[_CACHE_MODE]
|
|
self.cache_args = {**self.cache_args, "lru": 100, "ram": 0}
|
|
self.reset()
|
|
_CACHE_MODE = None
|
|
return _EXECUTE(self, *args, **kwargs)
|
|
|
|
|
|
execution.PromptExecutor.execute = _execute_with_cache_mode
|
|
|
|
|
|
@PromptServer.instance.routes.post("/pm_execution_fixture/cache")
|
|
async def cache_mode(request: web.Request) -> web.Response:
|
|
global _CACHE_MODE
|
|
mode = (await request.json())["mode"]
|
|
if mode not in ("CLASSIC", "LRU", "RAM_PRESSURE", "NONE"):
|
|
raise web.HTTPBadRequest()
|
|
_CACHE_MODE = mode
|
|
return web.json_response({"configured": True})
|
|
|
|
|
|
def _plugin_module(suffix: str):
|
|
"""A loaded PromptManager module (ComfyUI names the package by its path)."""
|
|
import sys
|
|
|
|
for name, module in list(sys.modules.items()):
|
|
if name.endswith("promptmanager." + suffix):
|
|
return module
|
|
raise web.HTTPServiceUnavailable(text=f"PromptManager {suffix} not loaded")
|
|
|
|
|
|
@PromptServer.instance.routes.post("/pm_execution_fixture/monitor")
|
|
async def image_monitor(request: web.Request) -> web.Response:
|
|
"""Start or stop PromptManager's image monitor on the host output folder.
|
|
|
|
The managed host's config disables monitoring; only the attribution cases
|
|
(#342) turn it on, and they turn it off again.
|
|
"""
|
|
enabled = (await request.json())["enabled"]
|
|
if type(enabled) is not bool:
|
|
raise web.HTTPBadRequest()
|
|
import folder_paths
|
|
|
|
monitor_module = _plugin_module("utils.image_monitor")
|
|
config = _plugin_module("py.config").GalleryConfig
|
|
monitor = monitor_module._monitor_instance
|
|
if monitor is None:
|
|
database = _plugin_module("database.operations").PromptDatabase()
|
|
tracker = _plugin_module("utils.prompt_tracker").get_prompt_tracker(database)
|
|
monitor = monitor_module.get_image_monitor(database, tracker)
|
|
if enabled:
|
|
config.MONITORING_ENABLED = True
|
|
monitor.start_monitoring([folder_paths.get_output_directory()])
|
|
else:
|
|
# Joins the image worker and the observer (one shared 5 s deadline),
|
|
# so it runs on a worker thread, not on the event loop.
|
|
await asyncio.to_thread(monitor.stop_monitoring)
|
|
config.MONITORING_ENABLED = False
|
|
return web.json_response({"monitoring": monitor.is_monitoring})
|
|
|
|
|
|
@PromptServer.instance.routes.get("/pm_execution_fixture/monitor")
|
|
async def image_monitor_state(request: web.Request) -> web.Response:
|
|
"""File names of keyless images waiting for their run's recorded outputs."""
|
|
monitor = _plugin_module("utils.image_monitor")._monitor_instance
|
|
handler = getattr(monitor, "handler", None)
|
|
waiting = handler.awaiting_history() if handler is not None else []
|
|
return web.json_response({"awaiting": [os.path.basename(p) for p in waiting]})
|