v3.0.0-beta.7

This is extremely cursed, but you can do prompt scheduling with video models using this.

For example:

NODE(UC_AdvancedMiniMaxH3ImageToVideo, prompt, vae 161:152 0; reference_images.reference_image_1 [114:201:0.3] 0; width 120 0; height 120 1; length 176 1)

This changes the reference image parameter mid-sampling.

Syntax for the last parameter is a ;-separated list of paramname node_id output_slot tuples. The given slot from the node will then be plugged in as the named parameter)

Export your workflow in API format to find the node IDs.

This will probably change later.

README.md

@@ -23,6 +23,7 @@ A `Basic Text to Image` template is included with the extension, and can be load
   - Different weight interpretation types (ComfyUI, A1111, compel, etc.)
   - Prompt masking with an implementation of [cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff).
 - Organize complicated prompts with [segments and prompt macros](/doc/macros.md).
+- [Schedule your own encoder nodes](/doc/node_function.md), allowing prompt control of eg. video or audio models with non-text inputs.
 
 All features are fully schedulable unless otherwise stated. See the [scheduling syntax documentation](doc/schedules.md) to get started.
 
@@ -52,10 +53,6 @@ If you run into problems, update ComfyUI first.
 
 for example, if you first encode `[cat:dog:0.1]` and later change that to `[cat:dog:0.5]`, no re-encoding takes place.
 
-for added fun, put `NODE(NodeClassName, textinputname)` in a prompt to generate a graph using **any other node** that's compatible. The node can't have required parameters besides a single CLIP parameter (which must be named `clip`) and the text prompt, and it must return a `CONDITIONING` as its first return value. The "default" values are `PCTextEncode` and `text`.
-
-For example, if you for some reason do not want the advanced features of `PCTextEncode`, use `NODE(CLIPTextEncode)` in the prompt and you'll still get scheduling with ComfyUI's regular TE node.
-
 The advanced node enables filtering the prompt for multi-pass workflows.
 
 ## PCLazyLoraLoader and PCLazyLoraLoaderAdvanced

__init__.py

@@ -30,7 +30,7 @@ if "PYTEST_CURRENT_TEST" not in os.environ:
         h = logging.StreamHandler(sys.stdout)
         h.setFormatter(logging.Formatter("[PromptControl] %(levelname)s: %(message)s"))
         log.addHandler(h)
-    for node in ["base", "hooks", "tools", "lazy"]:
+    for node in ["base", "hooks", "tools", "lazy", "anima"]:
         mod = importlib.import_module(f".prompt_control.nodes_{node}", package=__name__)
         v3_modules.append(mod)
 

doc/attention_couple.md

@@ -1,7 +1,5 @@
 # Attention Couple
 
-NOTE: This is still considered an experimental feature, so the syntax may change.
-
 Attention Couple is an attention-based implementation of regional prompting. it is faster and often more flexible than latent-based masking.
 
 The implementation is based on the one by [pamparamm](https://github.com/pamparamm/ComfyUI-ppm.git), modified to use ComfyUI's hook system. This enables it to work with prompt scheduling.
@@ -12,6 +10,11 @@ As a consequence of this, however, you can also use `COUPLE` in your negative pr
 
 To enable batching negative prompts, run your positive and negative prompt through the `PPCAttentionCoupleBatchNegative` node. This will make the outputs identical to pamparamm's implementation and will also improve performance. It will fall back to the default behaviour in cases where batching can't be done, so it should always be safe to use.
 
+## Anima
+
+There is a **very experimental** port of pamparamm's Anima support for Attention Couple in Prompt Control. Because ComfyUI lacks the built-in schedulable hooks required, you must first patch your model with `PC: Anima Attention Couple Model Patch` in addition to using `COUPLE` as usual.
+
+The code was hacked together with minimal thought, so expect bugs and misbehaviour. The port is also currently *not* compatible with NegPIP.
 
 ## Syntax
 
@@ -21,6 +24,13 @@ See also the [regional prompting documentation](/doc/regional_prompts.md) for in
 
 You can use `COUPLE` to attach attention-coupled prompts to a base prompt:
 
+For example:
+```
+dog FILL() COUPLE(0.5 1) cat
+```
+
+The full syntax looks as follows (to use `IMASK` you need to attach a custom mask)
+
 `base_prompt COUPLE MASK(0 0.5) coupled prompt 1 with mask COUPLE IMASK(0) coupled prompt 2 with custom mask`
 
 as a shortcut, `COUPLE(maskparams)` is expanded to `COUPLE MASK(maskparams)`, so the above prompt can also be written as:
@@ -28,15 +38,17 @@ as a shortcut, `COUPLE(maskparams)` is expanded to `COUPLE MASK(maskparams)`, so
 `base_prompt COUPLE(0 0.5) coupled prompt 1 with mask COUPLE IMASK(0) coupled prompt 2 with custom mask`
 
 Behaviour:
-- If no mask is specified, an implicit `MASK()` is assumed.
+- If no mask is specified, an implicit `MASK()` is assumed, meaning that the prompt affects the entire image.
 
-- For the base prompt, you can also use `FILL()` to automatically mask all parts not masked by coupled prompts
+- For the base prompt, you can use `FILL()` to automatically mask all parts not masked by other coupled prompts
 
-- If the base prompt has weight set to zero (ie. ´:0` at the end), then the first coupled prompt with non-zero weight becomes the base prompt.
+- If the base prompt has weight set to zero (ie. ´:0` at the end), then the first coupled prompt with non-zero weight becomes the base prompt:
 
-For example:
 ```
-dog FILL() COUPLE(0.5 1) cat
+disabled prompt :0 COUPLE new base prompt COUPLE coupled prompt
 ```
 
-Note that because the generation still sees and diffuses the full latent, attention coupling is not guaranteed to perfectly limit the effect of your prompt to the masked area.
+You can also schedule the weight normally: `prompt :[1:0:0.35]`
+
+> ![NOTE]
+> Note that because the generation still sees and diffuses the full latent, attention coupling is not guaranteed to perfectly limit the effect of your prompt to the masked area.

doc/macros.md

@@ -86,7 +86,9 @@ Top panel: a sleeping cat. The cat has orange fur with white stripes
 Bottom panel: A dog chasing its
   tail in a living room.
 ```
-Unlike macros, SEG is processed *after* scheduling syntax.
+
+> [!NOTE]
+> Unlike macros, SEGs are processed *after* scheduling syntax has been expanded, except in the LoRA loader (this may change later, but requires a bit of refactoring)
 
 In this case, the first section before any `SEG` becomes the *template* and any text after a `SEG` call becomes part of that segment. Whitespace is stripped from the start and end of segments and the template.
 

doc/node_function.md

@@ -0,0 +1,38 @@
+# The NODE function
+
+The `NODE` function allows you to use any other text encoding node within `PC: Schedule Prompt`, replacing the default `PCTextEncode` and allowing for example video model scheduling.
+
+## Basic usage
+
+Use `NODE(NodeClassName, textinputname)` in a prompt to generate a graph using any node that's compatible. The requirements are as follows:
+- The node must have a CLIP parameter (which must be named `clip`)
+- It must have a text field
+- It must return a `CONDITIONING` as its first return value.
+
+For example, if you for some reason do not want the advanced features of `PCTextEncode`, use `NODE(CLIPTextEncode)` in the prompt and you'll still get scheduling with ComfyUI's regular TE node.
+
+The default parameters are `PCTextEncode` and `text`.
+
+## Advanced Usage with arbitrary parameters
+
+Advanced usage of `NODE` can be complicated. For an example, see [The H3 workflow](/example_workflows/Prompt%20Control%20with%20MiniMax%20H3.json?raw=1). You can also find it in the template library.
+
+The full synopsis of the function is `NODE(NodeClassName, textinputname, arg_spec)` where `arg_spec` is a semicolon-separated list of `parameter_name json_value` pairs. In raw form, it looks like this:
+```
+NODE(MiniMaxH3ImageToVideo, prompt, vae ["1", 0]; width 1024; height 1024; first_frame ["2", 0])
+```
+
+The names and inputs must match the ComfyUI API format which **may differ from frontend names**. You can export your workflow in API format and inspect it to see how inputs are passed in to nodes.
+
+The arrays are literal ComfyUI node links, meaning the `vae` parameter is taken from node ID "1" first output and `first_frame` from node ID "2" first output.
+
+The values are arbitrary JSON literals, meaning that you can also pass in constant values. To pass in literal strings for example, you need to use quotes `"like this"`.
+
+This is intended to be used with the helper node `PC: NODE Input Helper`, which can be used to pass arbitrary parameters (named `$a` to `$n`) to the encoder. The recommended pattern is to put something like the following:
+```
+SEG(node)
+NODE(MiniMaxH3ImageToVideo, prompt, vae $a; width $b; height $c; length $d; first_frame $e)
+```
+to the helper and then concatenate it at the end of your prompt (use whitespace as a separator). You can then trigger the node with `$node` in your prompt. (see [documentation](/doc/macros.md) for `SEG`)
+
+The helper will replace the parameters with the correct ComfyUI link values.

prompt_control/anima_couple.py

@@ -0,0 +1,167 @@
+# Adapted from https://github.com/pamparamm/ComfyUI-ppm
+import itertools
+from collections.abc import Callable
+from functools import partial
+from math import lcm
+
+import torch
+import torch.nn.functional as F
+from comfy.ldm.anima.model import Anima as AnimaDIT
+from comfy.ldm.cosmos.predict2 import Attention as CosmosAttention
+from comfy.patcher_extension import WrapperExecutor
+from comfy.sampler_helpers import convert_cond
+from comfy.samplers import process_conds
+
+COND = 0
+UNCOND = 1
+
+
+def reshape_mask(mask: torch.Tensor, size: tuple[int, int], bs: int, num_tokens: int) -> torch.Tensor:
+    num_conds = mask.shape[0]
+
+    mask_downsample = F.interpolate(mask, size=size, mode="nearest")
+    mask_downsample_reshaped = mask_downsample.view(num_conds, num_tokens, 1).repeat_interleave(bs, dim=0)
+
+    return mask_downsample_reshaped
+
+
+def wrap_forwards(anima_model):
+    backups = {}
+    for block_name, b in (
+        (n, b) for n, b in anima_model.named_modules() if "cross_attn" in n and isinstance(b, CosmosAttention)
+    ):
+        backups[block_name] = b.forward
+        b.forward = partial(cosmos_attention_forward_couple, b.forward)
+    return backups
+
+
+def unwrap_forwards(anima_model, backups):
+    for block_name, b in (
+        (n, b) for n, b in anima_model.named_modules() if "cross_attn" in n and isinstance(b, CosmosAttention)
+    ):
+        b.forward = backups[block_name]
+
+
+def anima_sample_wrapper(executor, *args, **kwargs):
+    guider, _, extra_options, _, noise, latent_image, denoise_mask, *_ = args
+    seed = extra_options["seed"]
+    device = "cuda"  # TODO: fix
+
+    def pc_process_conds(pc_conds):
+        conds = [convert_cond([c])[0] for c in pc_conds]
+        conds = process_conds(
+            guider.inner_model,
+            noise,
+            {"positive": conds},
+            device,
+            latent_image,
+            denoise_mask,
+            seed,
+            latent_shapes=[latent_image.shape],
+        )
+        return [
+            c["model_conds"]["c_crossattn"].cond * pc_conds[i][1].get("strength", 1.0)
+            for i, c in enumerate(conds["positive"])
+        ]
+
+    extra_options["model_options"]["transformer_options"]["pc_process_conds"] = pc_process_conds
+    return executor(*args, **kwargs)
+
+
+def anima_forward_wrapper(executor: WrapperExecutor, *args, **kwargs):
+    """Model wrapper does something with activation shapes?"""
+    anima_model: AnimaDIT = executor.class_obj  # type: ignore
+
+    x: torch.Tensor = args[0]
+    transformer_options: dict = kwargs.get("transformer_options", {}).copy()
+    pc = transformer_options.get("pc_couple")
+    if pc and "processed_conds" not in pc:
+        pc["processed_conds"] = transformer_options["pc_process_conds"](pc["conds"])
+    patch_spatial = anima_model.patch_spatial
+
+    activations_shape = list(x.shape)
+    activations_shape[-2] = activations_shape[-2] // patch_spatial
+    activations_shape[-1] = activations_shape[-1] // patch_spatial
+
+    transformer_options["activations_shape"] = activations_shape
+    kwargs["transformer_options"] = transformer_options
+
+    b = {}
+    if pc:
+        b = wrap_forwards(anima_model)
+    r = executor(*args, **kwargs)
+    if pc:
+        unwrap_forwards(anima_model, b)
+    return r
+
+
+def cosmos_attention_forward_couple(_forward: Callable, x, context, rope_emb, transformer_options):
+    """attention block wrapper"""
+    if "pc_couple" not in transformer_options:
+        return _forward(x, context, rope_emb, transformer_options)
+    c: torch.Tensor = context
+    # FIXME: base cond weight
+    # c = args["processed_conds"][0]
+
+    args = transformer_options["pc_couple"]
+
+    mask = args["mask"]
+    conds = args["processed_conds"][1:]
+    num_conds = len(conds) + 1
+    num_tokens_c: list[int] = [c.shape[1] for c in conds]
+    cond_or_uncond = transformer_options["cond_or_uncond"]
+    cond_or_uncond_couple = []
+
+    num_chunks = len(cond_or_uncond)
+    bs = x.shape[0] // num_chunks
+
+    x_chunks = x.chunk(num_chunks, dim=0)
+    c_chunks = c.chunk(num_chunks, dim=0)
+    lcm_tokens_c = lcm(c.shape[1], *num_tokens_c)
+    conds_c_tensor = torch.cat(
+        [cond.repeat(bs, lcm_tokens_c // num_tokens_c[i], 1) for i, cond in enumerate(conds)],
+        dim=0,
+    )
+
+    xs, cs = [], []
+    for i, cond_type in enumerate(cond_or_uncond):
+        x_target = x_chunks[i]
+        c_target = c_chunks[i].repeat(1, lcm_tokens_c // c.shape[1], 1)
+        if cond_type == UNCOND:
+            xs.append(x_target)
+            cs.append(c_target)
+            cond_or_uncond_couple.append(UNCOND)
+        else:
+            xs.append(x_target.repeat(num_conds, 1, 1))
+            cs.append(torch.cat([c_target, conds_c_tensor], dim=0))
+            cond_or_uncond_couple.extend(itertools.repeat(COND, num_conds))
+
+    xs = torch.cat(xs, dim=0)
+    cs = torch.cat(cs, dim=0)
+
+    out = _forward(xs, cs, rope_emb, transformer_options)
+
+    size = tuple(transformer_options["activations_shape"][-2:])
+    num_tokens = out.shape[1]
+    mask_downsample = reshape_mask(mask, size, bs, num_tokens)
+
+    outputs = []
+    cond_outputs = []
+    i_cond = 0
+
+    for i, cond_type in enumerate(cond_or_uncond_couple):
+        pos, next_pos = i * bs, (i + 1) * bs
+
+        if cond_type == UNCOND:
+            outputs.append(out[pos:next_pos])
+        else:
+            pos_cond, next_pos_cond = i_cond * bs, (i_cond + 1) * bs
+            masked_output = out[pos:next_pos] * mask_downsample[pos_cond:next_pos_cond]
+            cond_outputs.append(masked_output)
+            i_cond += 1
+
+    if len(cond_outputs) > 0:
+        cond_output = torch.stack(cond_outputs).sum(0)
+        outputs.append(cond_output)
+
+    return torch.cat(outputs, dim=0)

prompt_control/attention_couple_ppm.py

@@ -63,21 +63,23 @@ class AttentionCoupleHook(TransformerOptionsHook):
     def __init__(self):
         super().__init__(hook_scope=EnumHookScope.HookedOnly)
 
-        self.transformers_dict = {
+        self.transformers_dict: dict[str, Any] = {
             "patches": {
                 "attn2_output_patch": [Proxy(self.attn2_output_patch)],
                 "attn2_patch": [Proxy(self.attn2_patch)],
-            }
+            },
+            "pc_couple": {},
         }
+
         self.has_negpip = False
-        # calculate later. All clones must refer to the same kv dict
-        # The type is here to shut up the type checker
-        self.kv: dict[str, list] = {"k": None, "v": None}
+        # The list will be calculated later. All clones must refer to the same kv dict
+        self.kv: dict[str, list] = {"k": None, "v": None}  # type: ignore
 
     def initialize_regions(self, base_cond, conds, fill):
         self.num_conds = len(conds) + 1
         self.base_strength = base_cond[1].get("strength", 1.0)
         self.strengths: list[float] = [cond[1].get("strength", 1.0) for cond in conds]
+        self.comfy_conds = [base_cond] + conds
         self.conds: list[torch.Tensor] = [base_cond[0]] + [cond[0] for cond in conds]
         base_mask = base_cond[1].get("mask", None)
         masks = [cond[1].get("mask") * cond[1].get("mask_strength") for cond in conds]
@@ -117,6 +119,11 @@ class AttentionCoupleHook(TransformerOptionsHook):
         self.mask = mask / mask.sum(dim=0, keepdim=True)
 
     def on_apply_hooks(self, model: ModelPatcher, transformer_options: dict[str, Any]):
+        self.transformers_dict["pc_couple"] = {
+            "conds": self.comfy_conds,
+            "num_conds": self.num_conds,
+            "mask": self.mask,
+        }
         if self.kv["k"] is None:
             self.has_negpip = model.model_options.get("ppm_negpip", False)
             log.debug("AttentionCouple has_negpip=%s", self.has_negpip)

prompt_control/macros.py

@@ -9,30 +9,47 @@ from .utils import find_closing_paren, get_function, split_by_function
 log = logging.getLogger("comfyui-prompt-control")
 
 
-def substitute_template(template, segments):
+def substitute_template(template, segments, do_subs):
     def _substitute(template, segments, stack):
-        for name, value in sorted(segments):
-            value = substitute_var(value, name, "")
-            if name not in stack:
-                stack.add(name)
-                value = _substitute(value, segments, stack)
-                stack.remove(name)
-            template = substitute_var(template, name, value)
+        name = ""
+        if "$" in template:
+            for name, value in sorted(segments):
+                value = substitute_var(value, name, "")
+                if name not in stack:
+                    stack.add(name)
+                    value = _substitute(value, segments, stack)
+                    stack.remove(name)
+                template = substitute_var(template, name, value)
+        if do_subs and name not in stack:
+            template = expand_subs(template)
         return template
 
     return _substitute(template, segments, set())
 
 
-def expand_segs(text):
+def expand_segs(text, do_subs=True):
     template, segments = split_by_function(text, "SEG", defaults=[""], require_args=True)
     named_segs = [(f.args[0].strip() or f"SEG{i + 1}", c.strip()) for i, (c, f) in enumerate(segments)]
 
-    new_text = substitute_template(template, named_segs).strip()
+    new_text = substitute_template(template, named_segs, do_subs).strip()
     if new_text != text.strip():
         log.debug("Template expanded to: %s", new_text)
     return new_text
 
 
+def expand_subs(text):
+    text, subs = get_function(text, "SUB", defaults=None)
+    subs = [spec.strip() for f in subs for spec in f.args[0].split(";")]
+    for spec in subs:
+        if len(spec) <= 3 or spec[0] != "s":
+            log.warning("Invalid SUB spec ignored: '%s'", spec)
+            continue
+        splitchar = spec[1]
+        search, replace, *_ = spec[2:].split(splitchar)
+        text = re.sub(search, replace, text)
+    return text
+
+
 def parse_search(search):
     arg_start = search.find("(")
     args = ""
@@ -71,7 +88,6 @@ def expand_macros(text):
         iterations += 1
         if iterations > 10:
             raise ValueError("Unable to resolve DEFs, make sure there are no cycles!")
-            return text
         for search, replace in replacements:
             res = substitute_defcall(res, search, replace)
         if res == prevres:
@@ -79,13 +95,15 @@ def expand_macros(text):
         prevres = res
     if res.strip() != text.strip():
         res = res.strip()
-        log.info("DEFs expanded to: %s", res)
+        log.debug("DEFs expanded to: %s", res)
     return res
 
 
-def substitute_var(text, name, replace):
+def substitute_var(text, name, replace, boundary=r"\b"):
+    if f"${name}" not in text:
+        return text
     name = re.escape(str(name))
-    return re.sub(rf"\${name}\b", replace, text)
+    return re.sub(rf"\${name}{boundary}", replace, text)
 
 
 def substitute_defcall(text, search, replace):
@@ -99,11 +117,12 @@ def substitute_defcall(text, search, replace):
         if parameters:
             paramvals = [x.strip() for x in parameters[0].split(";")]
         r = replace
+        end_re = r"(?![0-9])"
         for i, v in enumerate(paramvals):
-            r = substitute_var(r, i + 1, v)
+            r = substitute_var(r, i + 1, v, boundary=end_re)
 
         for i, v in enumerate(default_args):
-            r = substitute_var(r, i + 1, v)
+            r = substitute_var(r, i + 1, v, boundary=end_re)
 
         text = text.replace(ph, r)
     return text

prompt_control/nodes_anima.py

@@ -0,0 +1,52 @@
+# Adapted from ComfyUI-ppm into hook form
+
+
+import comfy.model_management
+import comfy.patcher_extension
+from comfy.model_base import Anima
+from comfy.model_patcher import ModelPatcher
+from comfy_api.latest import io
+
+from .anima_couple import (
+    anima_forward_wrapper,
+    anima_sample_wrapper,
+)
+
+
+class PCAnimaAttnCouplePatch(io.ComfyNode):
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        return io.Schema(
+            node_id="PCAnimaAttnCouplePatch",
+            display_name="PC: Anima attention Couple Model Patch",
+            category="promptcontrol/experimental",
+            inputs=[
+                io.Model.Input("model"),
+            ],
+            outputs=[
+                io.Model.Output(),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model: ModelPatcher) -> io.NodeOutput:
+        model_type = type(model.model)
+        m = model
+
+        if issubclass(model_type, Anima):
+            m = model.clone()
+            m.add_wrapper_with_key(
+                comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL,
+                cls.__name__,
+                anima_forward_wrapper,
+            )
+            m.add_wrapper_with_key(
+                comfy.patcher_extension.WrappersMP.SAMPLER_SAMPLE,
+                cls.__name__,
+                anima_sample_wrapper,
+            )
+
+        return io.NodeOutput(m)
+
+
+NODES = [PCAnimaAttnCouplePatch]

prompt_control/nodes_base.py

@@ -26,7 +26,7 @@ class PCTextEncodeWithRange(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, clip, text, start=0.0, end=1.0) -> io.NodeOutput:  # ty: ignore[invalid-method-override]
+    def execute(cls, clip, text, start=0.0, end=1.0) -> io.NodeOutput:
         log.debug("PCTextEncode: Encoding '%s'", text)
         defaults = clip.patcher.model_options.get("x-promptcontrol.defaults", {})
         masks = clip.patcher.model_options.get("x-promptcontrol.masks", None)
@@ -51,7 +51,7 @@ class PCTextEncode(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, clip, text) -> io.NodeOutput:  # ty: ignore[invalid-method-override]
+    def execute(cls, clip, text) -> io.NodeOutput:
         # Use the WithRange node for the range 0.0, 1.0
         return PCTextEncodeWithRange.execute(clip, text, 0.0, 1.0)
 

prompt_control/nodes_hooks.py

@@ -29,7 +29,7 @@ class PCLoraHooksFromText(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, text) -> io.NodeOutput:  # ty: ignore[invalid-method-override]
+    def execute(cls, text) -> io.NodeOutput:
         prompt_schedule = parse_prompt_schedules(text)
         consolidated = consolidate_schedule(prompt_schedule)
         hooks = lora_hooks_from_schedule(consolidated, {})
@@ -104,7 +104,7 @@ class PCAttentionCoupleBatchNegative(io.ComfyNode):
 
     @classmethod
     @override
-    def execute(cls, positive, negative) -> io.NodeOutput:  # ty: ignore[invalid-method-override]
+    def execute(cls, positive, negative) -> io.NodeOutput:
         if len(negative) != 1:
             log.warning("Batching scheduled negatives is not supported yet")
             return io.NodeOutput(positive, negative)

prompt_control/nodes_lazy.py

@@ -8,6 +8,7 @@ from comfy_api.latest import io
 from comfy_execution.graph import ExecutionBlocker
 from comfy_execution.graph_utils import GraphBuilder
 
+from .macros import expand_macros, expand_segs
 from .parser import parse_prompt_schedules
 from .utils import consolidate_schedule, find_nonscheduled_loras, get_function
 
@@ -158,6 +159,7 @@ class PCLazyLoraLoaderAdvanced(io.ComfyNode):
 
     @classmethod
     def execute(cls, model=None, clip=None, text="", apply_hooks=True, tags="", start=0.0, end=1.0, num_steps=0):
+        text = expand_segs(expand_macros(text))
         schedule = parse_prompt_schedules(text, filters=tags, start=start, end=end, num_steps=num_steps)
         graph = GraphBuilder()
         r = build_lora_schedule(graph, schedule, model, clip, apply_hooks=apply_hooks)
@@ -195,14 +197,36 @@ def build_scheduled_prompts(graph, schedules, clip):
     start_pct = 0.0
     for end_pct, c in schedules:
         p = c["prompt"]
-        p, classnames = get_function(p, "NODE", ["PCTextEncode", "text"])
-        classname = "PCTextEncode"
-        paramname = "text"
+        p, classnames = get_function(p, "NODE", defaults=None)
+        realargs = ["PCTextEncode", "text", ""]
         if classnames:
-            classname, paramname = classnames[0].args
-        node = graph.node(classname)
+            args = classnames[0].args[0]
+            if not args.strip():
+                raise ValueError("NODE can't be empty!")
+            for i, v in enumerate(args.split(",", maxsplit=2)):
+                realargs[i] = v
+        classname, paramname, magic_spec = realargs
+        node = graph.node(classname.strip())
         node.set_input("clip", clip)
-        node.set_input(paramname, p)
+        node.set_input(paramname.strip(), p)
+        magic_spec.replace(r"\;", "__ESCAPED_SEMICOLON__")
+        extra_inputs = magic_spec.split(";") if magic_spec.strip() else []
+        for e in extra_inputs:
+            e = e.strip()
+            if not e:
+                continue
+            e = e.replace("__ESCAPED_SEMICOLON__", ";")
+            name, jsondata = e.split(maxsplit=1)
+            jsondata = jsondata.strip()
+            if not jsondata.strip():
+                continue
+            # From helper node:
+            if jsondata == "__EMPTY__":
+                continue
+            try:
+                node.set_input(name.strip(), json.loads(jsondata.strip()))
+            except ValueError as e:
+                raise ValueError(f"Invalid JSON input: '{jsondata}'") from e
         timestep = graph.node("ConditioningSetTimestepRange")
         timestep.set_input("conditioning", node.out(0))
         timestep.set_input("start", start_pct)

prompt_control/nodes_tools.py

@@ -1,8 +1,12 @@
+import json
 import logging
 
 from comfy_api.latest import io
 
-from .macros import expand_macros
+from .macros import expand_macros as macroexpand
+from .macros import expand_segs as segexpand
+from .macros import expand_subs as subexpand
+from .macros import substitute_var
 from .parser import parse_prompt_schedules
 
 log = logging.getLogger("comfyui-prompt-control")
@@ -132,22 +136,32 @@ class PCExtractScheduledPrompt(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="PCExtractScheduledPrompt",
-            display_name="PC: Extract Scheduled Prompt",
+            display_name="PC: Show Prompt",
             category="promptcontrol/tools",
             description="Parses the input prompt and returns the prompt scheduled at the specified point",
             inputs=[
                 io.String.Input("text", multiline=True),
                 io.Float.Input("at", min=0.0, max=1.0, default=1.0, step=0.01),
                 io.String.Input("tags", default="", optional=True),
+                io.Boolean.Input("expand_segs", default=False, optional=True),
+                io.Boolean.Input("expand_subs", default=False, optional=True),
+                io.Boolean.Input("expand_macros", default=False, optional=True),
             ],
             outputs=[io.String.Output()],
+            search_aliases=["extract scheduled prompt"],
         )
 
     @classmethod
-    def execute(cls, text, at, tags="") -> io.NodeOutput:
+    def execute(cls, text, at, tags="", expand_segs=False, expand_subs=False, expand_macros=False) -> io.NodeOutput:
+        if expand_macros:
+            text = macroexpand(text)
         schedule = parse_prompt_schedules(text, filters=tags)
-        _, entry = schedule.at_step(at, total_steps=1)
+        _, entry = schedule.at_step(at)
         prompt_text = entry.get("prompt", "")
+        if expand_segs:
+            prompt_text = segexpand(prompt_text, do_subs=expand_subs)
+        if expand_subs:
+            prompt_text = subexpand(prompt_text)
         return io.NodeOutput(prompt_text)
 
 
@@ -167,7 +181,33 @@ class PCMacroExpand(io.ComfyNode):
 
     @classmethod
     def execute(cls, text) -> io.NodeOutput:
-        return io.NodeOutput(expand_macros(text))
+        return io.NodeOutput(macroexpand(text))
+
+
+class PCLinkHelper(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        inputs: list = [io.String.Input("text", multiline=True)]
+        for x in "abcdefghijklmn":
+            inputs.append(io.AnyType.Input(x, optional=True, lazy=True, extra_dict={"rawLink": True}))
+        return io.Schema(
+            node_id="PCNODELinkHelper",
+            display_name="PC: Extra argument helper for NODE",
+            category="promptcontrol/tools",
+            description="Takes in arbitrary inputs and renders them as NODE-compatible values, replacing $a -> $n with JSON link values",
+            is_experimental=True,
+            inputs=inputs,
+            outputs=[io.String.Output()],
+        )
+
+    @classmethod
+    def execute(cls, text, **vars) -> io.NodeOutput:
+        for k in "abcdefghijklmn":
+            v = "__EMPTY__"
+            if k in vars:
+                v = json.dumps(vars[k])
+            text = substitute_var(text, k, v)
+        return io.NodeOutput(text)
 
 
 NODES = [
@@ -177,4 +217,5 @@ NODES = [
     PCSetLogLevel,
     PCExtractScheduledPrompt,
     PCMacroExpand,
+    PCLinkHelper,
 ]

prompt_control/parser.py

@@ -4,7 +4,7 @@ import os
 log = logging.getLogger("comfyui-prompt-control")
 
 if os.environ.get("PC_USE_OLD_PARSER", "0") != "1":
-    log.info("Using new parser implementation. Set PC_USE_OLD_PARSER=1 to use old parser instead")
     from .parser_parsy import parse_prompt_schedules  # noqa
 else:
+    log.warning("Using old Lark parser (UNSUPPORTED)")
     from .parser_lark import parse_prompt_schedules  # noqa

prompt_control/prompts.py

@@ -604,7 +604,7 @@ def encode_prompt(clip, text, start_pct, end_pct, defaults, masks):
         return f"MASK({args[0]})"
 
     for prompt in prompts:
-        text, noise_w, generator = get_noise(text)
+        prompt, noise_w, generator = get_noise(prompt)
         base_prompt, attn_couple_prompts = split_by_function(prompt, "COUPLE", defaults=None, require_args=False)
 
         prompts = [base_prompt] + [couple_mask(f.args) + chunk for (chunk, f) in attn_couple_prompts]

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-prompt-control"
 description = "Nodes for prompt editing and LoRA scheduling, advanced regional prompting (including attention masking) and advanced prompt encoding, all controlled through your text prompt. Feature keywords: comfyui-prompt-control, schedule, macros, attention couple, loractl, A1111"
-version = "3.0.0-beta.4"
+version = "3.0.0-beta.7"
 license = { file = "LICENSE" }
 
 requires-python = ">= 3.10"

tests/test_encode.py

@@ -73,6 +73,16 @@ def tensors_equal(t1, t2):
     npt.assert_equal(t1.detach().numpy(), t2.detach().numpy())
 
 
+def cond_neq(c1, c2, key=None, key_assert=None):
+    ok = False
+    try:
+        cond_equal(c1, c2, key=key, key_assert=key_assert)
+    except AssertionError:
+        ok = True
+    if not ok:
+        raise ValueError("Tensors should not be equal")
+
+
 def cond_equal(c1, c2, key=None, key_assert=None):
     assert len(c1) == len(c2)
     for i in range(len(c1)):
@@ -241,3 +251,15 @@ class TestPCTextEncode:
             (c2,) = run(pc_text_encode, clip, "test COUPLE MASK(0 0.2, 0.5) prompt1")
             cond_equal(c, c2)
             cond_equal(c, c2, "hooks", compare_hookgroup_mask)
+
+    def test_noise_weight0(self, text_encoder_clips, pc_text_encode, node_class_objs):
+        for _k, clip in text_encoder_clips:
+            (c1,) = run(pc_text_encode, clip, "test")
+            (c2,) = run(pc_text_encode, clip, "test NOISE(0, 0)")
+            cond_equal(c1, c2)
+
+    def test_noise(self, text_encoder_clips, pc_text_encode, node_class_objs):
+        for _k, clip in text_encoder_clips:
+            (c1,) = run(pc_text_encode, clip, "test")
+            (c2,) = run(pc_text_encode, clip, "test NOISE(1, 0)")
+            cond_neq(c1, c2)

tests/test_graph.py

@@ -460,6 +460,107 @@ def test_textencode_lora_with_schedule():
     }
 
 
+def test_textencode_custom():
+    r = te("NODE(CLIPTextEncode)simple [test:0.1,0.5] prompt")
+    assert r == {
+        "result": (["UID.0.0.8", 0],),
+        "expand": {
+            "UID.0.0.1": {
+                "class_type": "CLIPTextEncode",
+                "inputs": {"clip": [0, 0], "text": "simple  prompt"},
+            },
+            "UID.0.0.2": {
+                "class_type": "ConditioningSetTimestepRange",
+                "inputs": {"conditioning": ["UID.0.0.1", 0], "start": 0.0, "end": 0.1},
+            },
+            "UID.0.0.3": {
+                "class_type": "CLIPTextEncode",
+                "inputs": {"clip": [0, 0], "text": "simple test prompt"},
+            },
+            "UID.0.0.4": {
+                "class_type": "ConditioningSetTimestepRange",
+                "inputs": {"conditioning": ["UID.0.0.3", 0], "start": 0.1, "end": 0.5},
+            },
+            "UID.0.0.5": {
+                "class_type": "CLIPTextEncode",
+                "inputs": {"clip": [0, 0], "text": "simple  prompt"},
+            },
+            "UID.0.0.6": {
+                "class_type": "ConditioningSetTimestepRange",
+                "inputs": {"conditioning": ["UID.0.0.5", 0], "start": 0.5, "end": 1.0},
+            },
+            "UID.0.0.7": {
+                "class_type": "ConditioningCombine",
+                "inputs": {"conditioning_1": ["UID.0.0.2", 0], "conditioning_2": ["UID.0.0.4", 0]},
+            },
+            "UID.0.0.8": {
+                "class_type": "ConditioningCombine",
+                "inputs": {"conditioning_1": ["UID.0.0.7", 0], "conditioning_2": ["UID.0.0.6", 0]},
+            },
+        },
+    }
+
+
+def test_textencode_custom_extra():
+    r = te(
+        'NODE(CustomTextEncode, prompt, image ["1", 0]; option "test"; float [10.0:__EMPTY__:0.5])simple [test:0.1,0.5] prompt'
+    )
+    assert r == {
+        "result": (["UID.0.0.8", 0],),
+        "expand": {
+            "UID.0.0.1": {
+                "class_type": "CustomTextEncode",
+                "inputs": {
+                    "clip": [0, 0],
+                    "prompt": "simple  prompt",
+                    "image": ["1", 0],
+                    "option": "test",
+                    "float": 10.0,
+                },
+            },
+            "UID.0.0.2": {
+                "class_type": "ConditioningSetTimestepRange",
+                "inputs": {"conditioning": ["UID.0.0.1", 0], "start": 0.0, "end": 0.1},
+            },
+            "UID.0.0.3": {
+                "class_type": "CustomTextEncode",
+                "inputs": {
+                    "clip": [0, 0],
+                    "prompt": "simple test prompt",
+                    "image": ["1", 0],
+                    "option": "test",
+                    "float": 10.0,
+                },
+            },
+            "UID.0.0.4": {
+                "class_type": "ConditioningSetTimestepRange",
+                "inputs": {"conditioning": ["UID.0.0.3", 0], "start": 0.1, "end": 0.5},
+            },
+            "UID.0.0.5": {
+                "class_type": "CustomTextEncode",
+                "inputs": {
+                    "clip": [0, 0],
+                    "prompt": "simple  prompt",
+                    "image": ["1", 0],
+                    "option": "test"
+                },
+            },
+            "UID.0.0.6": {
+                "class_type": "ConditioningSetTimestepRange",
+                "inputs": {"conditioning": ["UID.0.0.5", 0], "start": 0.5, "end": 1.0},
+            },
+            "UID.0.0.7": {
+                "class_type": "ConditioningCombine",
+                "inputs": {"conditioning_1": ["UID.0.0.2", 0], "conditioning_2": ["UID.0.0.4", 0]},
+            },
+            "UID.0.0.8": {
+                "class_type": "ConditioningCombine",
+                "inputs": {"conditioning_1": ["UID.0.0.7", 0], "conditioning_2": ["UID.0.0.6", 0]},
+            },
+        },
+    }
+
+
 def test_loraloader_empty(monkeypatch, caplog):
     result = loraloader("prompt here <lora:nonexistent:1.0:0.5>")["expand"]
     result_adv = loraloader("prompt here <lora:nonexistent:1.0:0.5>", adv=True)["expand"]
@@ -582,3 +683,9 @@ def test_loraloader_adv_start():
 def test_loraloader_end_zero():
     result2 = loraloader("prompt [<lora:test:0.5>:0.5]", adv=True, end=0.5)["expand"]
     assert result2 == {}
+
+
+def test_loraloader_segs():
+    result = loraloader("prompt [<lora:test:0.5>:0.5]")["expand"]
+    result2 = loraloader("prompt [$lora:0.5]\nSEG(lora)<lora:test:0.5>\nSEG(lora2)<lora:ignored:1>")["expand"]
+    assert result == result2

tests/test_macros.py

@@ -14,6 +14,10 @@ from prompt_control.macros import expand_macros, expand_segs
             "[empty::$2] [empty::] [empty::0.5] [empty:a:0.5]",
         ),
         ("DEF(X=$1)DEF(Y()=$1)[X Y][X() Y()][X(1) Y(1)]", "[$1 ][ ][1 1]"),
+        (
+            "DEF(C_ANIMAL=cat)DEF(D_ANIMAL=dog)DEF(IT=It is a $1_ANIMAL $10_ANIMAL)IT(D) IT(C)",
+            "It is a dog $10_ANIMAL It is a cat $10_ANIMAL",
+        ),
     ],
 )
 def test_basic_macro(text, result):