150 lines
5.1 KiB
Python
150 lines
5.1 KiB
Python
# "Design-focused node suite for ComfyUI."
|
|
# Copyright 2026 Augment Studio
|
|
# Augmentstudio.app
|
|
|
|
"""
|
|
Augment Executor Override
|
|
─────────────────────────────
|
|
Bypasses ComfyUI's caching for trigger-driven workflows.
|
|
|
|
What it does:
|
|
1. Monkeypatches the execute() function in ComfyUI's execution module
|
|
2. Any node class with ALWAYS_EXECUTE = True skips the output cache — always re-runs
|
|
3. Any node that accepts/outputs TRIGGER type is auto-detected and also skips cache
|
|
4. Everything else executes normally — zero impact on other workflows
|
|
|
|
No fork needed. No CLI flags. Just drop and go.
|
|
"""
|
|
|
|
import logging
|
|
import execution as comfy_execution
|
|
from comfy_execution.caching import BasicCache, CacheKeySetInputSignature, HierarchicalCache
|
|
|
|
logger = logging.getLogger("Augment.Executor")
|
|
|
|
# ── Store reference to the original execute function ─────────────
|
|
_original_execute = comfy_execution.execute
|
|
|
|
|
|
def _should_skip_cache(class_def):
|
|
"""
|
|
Determine if a node class should bypass the output cache.
|
|
|
|
Returns True if:
|
|
- Node has ALWAYS_EXECUTE = True
|
|
- Node accepts or returns TRIGGER type
|
|
- Node has IS_CHANGED that returns float("NaN") (legacy force-run pattern)
|
|
"""
|
|
# Explicit opt-in
|
|
if getattr(class_def, "ALWAYS_EXECUTE", False):
|
|
return True
|
|
|
|
# Check if node deals with TRIGGER types (input or output)
|
|
try:
|
|
input_types = class_def.INPUT_TYPES()
|
|
for category in ("required", "optional"):
|
|
if category in input_types:
|
|
for key, val in input_types[category].items():
|
|
input_type = val[0] if isinstance(val, tuple) else val
|
|
if input_type == "TRIGGER":
|
|
return True
|
|
except Exception:
|
|
pass
|
|
|
|
# Check return types for TRIGGER
|
|
if hasattr(class_def, "RETURN_TYPES"):
|
|
for rt in class_def.RETURN_TYPES:
|
|
if rt == "TRIGGER":
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
async def patched_execute(server, dynprompt, caches, current_item, extra_data,
|
|
executed, prompt_id, execution_list,
|
|
pending_subgraph_results, pending_async_nodes, ui_outputs):
|
|
"""
|
|
Drop-in replacement for execution.execute().
|
|
|
|
For nodes that should skip cache:
|
|
- Evict any cached output BEFORE the original execute() runs
|
|
- This forces the original code past the cache-hit early return (line 419-428)
|
|
- The node executes fresh every time
|
|
- After execution, evict again so next run also re-executes
|
|
|
|
For all other nodes: passthrough to original, zero overhead.
|
|
"""
|
|
import nodes
|
|
|
|
unique_id = current_item
|
|
node_info = dynprompt.get_node(unique_id)
|
|
class_type = node_info["class_type"]
|
|
class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
|
|
|
|
skip_cache = _should_skip_cache(class_def)
|
|
|
|
if skip_cache:
|
|
# Evict cached output so the original execute() can't short-circuit
|
|
try:
|
|
caches.outputs.delete(unique_id)
|
|
except (AttributeError, KeyError):
|
|
# NullCache or key doesn't exist — fine
|
|
pass
|
|
|
|
# Run the original execute
|
|
result = await _original_execute(
|
|
server, dynprompt, caches, current_item, extra_data,
|
|
executed, prompt_id, execution_list,
|
|
pending_subgraph_results, pending_async_nodes, ui_outputs
|
|
)
|
|
|
|
if skip_cache:
|
|
# Evict AFTER execution too, so next prompt also re-runs this node
|
|
try:
|
|
caches.outputs.delete(unique_id)
|
|
except (AttributeError, KeyError):
|
|
pass
|
|
|
|
return result
|
|
|
|
|
|
# ── Also patch the cache's set_prompt to skip signature checks for our nodes ──
|
|
|
|
_original_cache_set_prompt = None
|
|
|
|
def _patch_cache_set_prompt():
|
|
"""
|
|
Patch CacheSet's set_prompt behavior so IsChangedCache
|
|
never marks our trigger nodes as 'unchanged'.
|
|
"""
|
|
global _original_cache_set_prompt
|
|
|
|
original_is_changed_get = comfy_execution.IsChangedCache.get
|
|
|
|
async def patched_is_changed_get(self, node_id):
|
|
"""Force trigger/always-execute nodes to always report as changed."""
|
|
import nodes
|
|
node = self.dynprompt.get_node(node_id)
|
|
class_type = node["class_type"]
|
|
class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
|
|
|
|
if _should_skip_cache(class_def):
|
|
# Return NaN — ComfyUI treats this as "always changed"
|
|
self.is_changed[node_id] = float("NaN")
|
|
return float("NaN")
|
|
|
|
return await original_is_changed_get(self, node_id)
|
|
|
|
comfy_execution.IsChangedCache.get = patched_is_changed_get
|
|
|
|
|
|
# ── Apply patches on import ─────────────────────────────────────
|
|
|
|
def _install():
|
|
"""Monkeypatch ComfyUI's execution module. Called once on import."""
|
|
comfy_execution.execute = patched_execute
|
|
_patch_cache_set_prompt()
|
|
logger.info("Augment executor override installed — trigger nodes will always re-execute")
|
|
|
|
_install()
|