v2.0.0-beta.3

Set PROMPTCONTROL_ENABLE_CACHE_HACK=1 in your environment to enable

README.md

@@ -4,10 +4,7 @@ Control LoRA and prompt scheduling, advanced text encoding, regional prompting,
 
 ## Prompt Control v2
 
-Prompt control has been almost completely rewritten. It now uses ComfyUI's lazy execution to build graphs from the text prompt at runtime. This has some advantages:
-
-- ComfyUI will not re-run unchanged parts of generated graphs. This is especially useful for two-pass workflows where previously you'd be forced to re-run the first sampling pass even with filtering. That is no longer the case and it does the right thing.
-- The generated graph is often exactly equivalent to a manually built workflow using native ComfyUI nodes. There are no more weird sampling hooks that could cause problems with other nodes
+Prompt control has been almost completely rewritten. It now uses ComfyUI's lazy execution to build graphs from the text prompt at runtime. The generated graph is often exactly equivalent to a manually built workflow using native ComfyUI nodes. There are no more weird sampling hooks that could cause problems with other nodes
 
 Prompt Control also comes with `PCTextEncode`, which provides advanced text encoding with many additional features compared to ComfyUI's base `CLIPTextEncode`.
 
@@ -168,4 +165,8 @@ The parameters affect how the masked and unmasked prompts are combined to produc
 
 # Known issues
 
-- None at the moment
+- ComfyUI's caching mechanism has an issue that makes it unnecessarily invalidate caches for certain inputs; you'll still get some benefit from the lazy nodes, but changing inputs that shouldn't affect downstream nodes (especially if using filtering) will still cause them to be recomputed because ComfyUI doesn't realize the inputs haven't changed.
+
+If you want to enable a hack to fix this, set `PROMPTCONTROL_ENABLE_CACHE_HACK=1` in your environment. Unset it to disable.
+
+It's a purely optional performance optimization that allows Prompt Control nodes to override their cache keys in a way that should not interfere with other nodes. Note that the optimization only works if the text input to the lazy nodes is a constant (so either directly on the node or from a primitive); outputs from other nodes can't be optimized.

__init__.py

@@ -15,7 +15,7 @@ log = logging.getLogger("comfyui-prompt-control")
 log.propagate = False
 if not log.handlers:
     h = logging.StreamHandler(sys.stdout)
-    h.setFormatter(logging.Formatter("[%(levelname)s] PromptControl: %(message)s"))
+    h.setFormatter(logging.Formatter("[PromptControl] %(levelname)s: %(message)s"))
     log.addHandler(h)
 
 if os.environ.get("PROMPTCONTROL_DEBUG"):
@@ -23,6 +23,8 @@ if os.environ.get("PROMPTCONTROL_DEBUG"):
 else:
     log.setLevel(logging.INFO)
 
+cache_hack = importlib.import_module(".prompt_control.cache_hack", package=__name__)
+cache_hack.init()
 
 NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}

prompt_control/cache_hack.py

@@ -0,0 +1,47 @@
+import comfy_execution.caching
+from comfy_execution.graph_utils import is_link
+import nodes
+from os import environ
+
+import logging
+
+log = logging.getLogger("comfyui-prompt-control")
+
+
+include_unique_id_in_input = comfy_execution.caching.include_unique_id_in_input
+
+
+def promptcontrol_get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
+    if not dynprompt.has_node(node_id):
+        # This node doesn't exist -- we can't cache it.
+        return [float("NaN")]
+    node = dynprompt.get_node(node_id)
+    class_type = node["class_type"]
+    class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+    inputs = node["inputs"]
+    if hasattr(class_def, "CACHE_KEY"):
+        inputs = getattr(class_def, "CACHE_KEY")(inputs)
+    signature = [class_type, self.is_changed_cache.get(node_id)]
+    if (
+        self.include_node_id_in_input()
+        or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT)
+        or include_unique_id_in_input(class_type)
+    ):
+        signature.append(node_id)
+    for key in sorted(inputs.keys()):
+        if is_link(inputs[key]):
+            (ancestor_id, ancestor_socket) = inputs[key]
+            ancestor_index = ancestor_order_mapping[ancestor_id]
+            signature.append((key, ("ANCESTOR", ancestor_index, ancestor_socket)))
+        else:
+            signature.append((key, inputs[key]))
+    return signature
+
+
+def init():
+    if environ.get("PROMPTCONTROL_ENABLE_CACHE_HACK") != "1":
+        return
+    log.warning("Enabling Prompt Control cache hack")
+    comfy_execution.caching.CacheKeySetInputSignature.get_immediate_node_signature = (
+        promptcontrol_get_immediate_node_signature
+    )

prompt_control/nodes_lazy.py

@@ -1,6 +1,6 @@
 import logging
 from .parser import parse_prompt_schedules
-from comfy_execution.graph_utils import GraphBuilder
+from comfy_execution.graph_utils import GraphBuilder, is_link
 
 from .prompts import get_function
 
@@ -9,6 +9,13 @@ log = logging.getLogger("comfyui-prompt-control")
 from .utils import consolidate_schedule, find_nonscheduled_loras
 
 
+def cache_key_hack(inputs):
+    out = inputs.copy()
+    if not is_link(inputs["text"]):
+        out["text"] = cache_key_from_inputs(**inputs)
+    return out
+
+
 def create_lora_loader_nodes(graph, model, clip, loras):
     for path, info in loras.items():
         log.info("Creating LoraLoader for %s", path)
@@ -24,6 +31,7 @@ def create_lora_loader_nodes(graph, model, clip, loras):
 
 
 def create_hook_nodes_for_lora(graph, path, info, existing_node, start_pct, end_pct):
+    prev_keyframe = None
     next_keyframe = None
     if not existing_node:
         log.debug("Creating hook for %s", path)
@@ -106,13 +114,13 @@ def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_h
             n.set_input("hooks_B", h.out(0))
             res = n
         res = res.out(0)
-    if apply_hooks:
-        n = graph.node("SetClipHooks")
-        n.set_input("clip", clip)
-        n.set_input("hooks", res)
-        n.set_input("apply_to_conds", True)
-        n.set_input("schedule_clip", True)
-        clip = n.out(0)
+        if apply_hooks:
+            n = graph.node("SetClipHooks")
+            n.set_input("clip", clip)
+            n.set_input("hooks", res)
+            n.set_input("apply_to_conds", True)
+            n.set_input("schedule_clip", True)
+            clip = n.out(0)
 
     r = graph.finalize()
 
@@ -125,6 +133,8 @@ def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_h
 
 
 class PCLazyLoraLoaderAdvanced:
+    CACHE_KEY = cache_key_hack
+
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -132,9 +142,9 @@ class PCLazyLoraLoaderAdvanced:
                 "text": ("STRING", {"multiline": True}),
                 "model": ("MODEL", {"rawLink": True}),
                 "clip": ("CLIP", {"rawLink": True}),
-                "apply_hooks": ("BOOLEAN", {"default": True}),
             },
             "optional": {
+                "apply_hooks": ("BOOLEAN", {"default": True}),
                 "tags": ("STRING", {"default": ""}),
                 "start": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 0.0, "step": 0.01}),
                 "end": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 1.0, "step": 0.01}),
@@ -147,13 +157,15 @@ class PCLazyLoraLoaderAdvanced:
     CATEGORY = "promptcontrol"
     FUNCTION = "apply"
 
-    def apply(self, model, clip, text, apply_hooks, unique_id, tags="", start=0.0, end=1.0):
-        schedule = parse_prompt_schedules(text).with_filters(filters=tags, start=start, end=end)
+    def apply(self, model, clip, text, unique_id, apply_hooks=True, tags="", start=0.0, end=1.0):
+        schedule = parse_prompt_schedules(text, filters=tags, start=start, end=end)
         graph = GraphBuilder(f"PCLazyLoraLoaderAdvanced-{unique_id}")
         return build_lora_schedule(graph, schedule, model, clip, apply_hooks=apply_hooks, return_hooks=True)
 
 
 class PCLazyLoraLoader:
+    CACHE_KEY = cache_key_hack
+
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -210,16 +222,23 @@ def build_scheduled_prompts(graph, schedules, clip):
         combiner.set_input("conditioning_2", othernode.out(0))
         node = combiner
 
-    return {"result": (node.out(0),), "expand": graph.finalize()}
+    g = graph.finalize()
+
+    return {"result": (node.out(0),), "expand": g}
+
+
+def cache_key_from_inputs(text, tags="", start=0.0, end=1.0, **kwargs):
+    schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end)
+    return [(pct, s["prompt"]) for pct, s in schedules]
 
 
 class PCLazyTextEncode:
+    CACHE_KEY = cache_key_hack
+
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {"clip": ("CLIP", {"rawLink": True}), "text": ("STRING", {"multiline": True})},
-            # "optional": {"defaults": ("SCHEDULE_DEFAULTS",)},
-            "hidden": {"unique_id": "UNIQUE_ID"},
         }
 
     RETURN_TYPES = ("CONDITIONING",)
@@ -227,13 +246,15 @@ class PCLazyTextEncode:
     CATEGORY = "promptcontrol"
     FUNCTION = "apply"
 
-    def apply(self, clip, text, unique_id):
+    def apply(self, clip, text):
         schedules = parse_prompt_schedules(text)
-        graph = GraphBuilder(f"PCEncodeLazy-{unique_id}")
+        graph = GraphBuilder()
         return build_scheduled_prompts(graph, schedules, clip)
 
 
 class PCLazyTextEncodeAdvanced:
+    CACHE_KEY = cache_key_hack
+
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -251,7 +272,7 @@ class PCLazyTextEncodeAdvanced:
     FUNCTION = "apply"
 
     def apply(self, clip, text, unique_id, tags="", start=0.1, end=1.0):
-        schedules = parse_prompt_schedules(text).with_filters(start=start, end=end, filters=tags)
+        schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end)
         graph = GraphBuilder(f"PCLazyTextEncodeAdvanced-{unique_id}")
         return build_scheduled_prompts(graph, schedules, clip)
 

prompt_control/parser.py

@@ -5,6 +5,8 @@ from math import ceil
 logging.basicConfig()
 log = logging.getLogger("comfyui-prompt-control")
 
+from functools import lru_cache
+
 if lark.__version__ == "0.12.0":
     x = "Your lark package reports an ancient version (0.12.0) and will not work. If you have the 'lark-parser' package in your Python environment, remove that and *reinstall* lark!"
     log.error(x)
@@ -323,5 +325,6 @@ class PromptSchedule(object):
         return len(self.parsed_prompt) - 1, self.parsed_prompt[-1]
 
 
-def parse_prompt_schedules(prompt):
-    return PromptSchedule(prompt)
+@lru_cache
+def parse_prompt_schedules(prompt, **kwargs):
+    return PromptSchedule(prompt, **kwargs)

prompt_control/utils.py

@@ -35,7 +35,6 @@ def find_nonscheduled_loras(consolidated_schedule):
     if not consolidated_schedule:
         return {}
     last_end, candidate_loras = consolidated_schedule[0]
-    print(candidate_loras)
     to_remove = set()
     for candidate, weights in candidate_loras.items():
         for end, loras in consolidated_schedule[1:]:

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-prompt-control"
 description = "Nodes for convenient prompt editing, making many common operations prompt-controllable"
-version = "2.0.0-beta.2"
+version = "2.0.0-beta.3"
 license = { file = "LICENSE" }
 # some lark versions older than 1.1.9 apparently have a bug that breaks things, see https://github.com/asagi4/comfyui-prompt-control/issues/35
 dependencies = ["lark >= 1.1.9"]