v3.0.0-beta.6

This avoids using stale block references and seems to fix LoRAs

.github/workflows/tests.yml

@@ -19,5 +19,5 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version: '3.11'
-      - run: pip install pytest typing-extensions -r requirements.txt
+      - run: pip install pytest typing-extensions
       - run: PYTHONPATH=ComfyUI pytest tests/test_parser.py

.github/workflows/tests_with_comfy.yml

@@ -31,7 +31,7 @@ jobs:
       - name: install-torch
         run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
       - name: install ComfyUI
-        run: pip install pytest typing-extensions -r requirements.txt -r ComfyUI/requirements.txt
+        run: pip install pytest typing-extensions -r ComfyUI/requirements.txt
       - name: Download clip_l.safetensors
         run: curl -LO https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors
       - name: Force Comfy to use the CPU

Makefile

@@ -12,7 +12,7 @@ format:
 	ruff format
 
 test:
-	PYTHONPATH=../../ pytest tests/test_parser.py $(ARGS)
+	PYTHONPATH=../../ pytest tests/test_parser.py tests/test_cutout.py tests/test_macros.py $(ARGS)
 
 test_graph:
 	PYTHONPATH=../../  pytest tests/test_graph.py $(ARGS)
@@ -20,6 +20,9 @@ test_graph:
 test_encode:
 	PYTHONPATH=../../  pytest tests/test_encode.py $(ARGS)
 
+test_workflow:
+	PYTHONPATH=../../  pytest tests/test_workflow.py $(ARGS)
+
 test_encode_both:
 	TEST_TE="clip_l t5" PYTHONPATH=../../  pytest tests/test_encode.py $(ARGS)
 

README.md

@@ -1,6 +1,6 @@
 # ComfyUI prompt control
 
-Control LoRA and prompt scheduling, advanced text encoding, regional prompting, and much more, through your text prompt. Generates dynamic graphs that are literally identical to handcrafted noodle soup.
+Control LoRA and prompt scheduling, advanced text encoding, regional prompting, and much more, through your text prompt. Prompt Control generates dynamic graphs that are literally identical to handcrafted noodle soup, condensing complicated workflows with dozens of nodes into simple text prompts.
 
 Prompt Control comes with `PCTextEncode`, which provides advanced text encoding with many additional features compared to ComfyUI's base `CLIPTextEncode`.
 
@@ -14,17 +14,15 @@ A `Basic Text to Image` template is included with the extension, and can be load
 
 ## What can it do?
 
-You can use text prompts to control the following:
-
 - A1111-style prompt scheduling and filtering without noodle soup.
-- LoRA loading and [scheduling](/doc/schedules.md) via the prompt, using ComfyUI's hook system
+- LoRA loading and [scheduling](/doc/schedules.md) using ComfyUI's built-in hook system.
 - Masking, composition and area control ([regional prompting](/doc/regional_prompts.md)) with an implementation of [Attention Couple](/doc/attention_couple.md), also fully schedulable.
 - [Advanced prompt encoding](/doc/basic.md)
-  - Per-encoder prompts for models with multiple text encoders, such as SDXL and Flux
+  - Per-encoder prompts for models with multiple text encoders, such as SDXL and Flux.
   - Prompt combinators like `BREAK`, as well as `CAT`, `AVG()` and `AND` corresponding to ComfyUI's `ConditioningConcat`, `ConditioningAverage` and `ConditioningCombine` nodes.
   - Different weight interpretation types (ComfyUI, A1111, compel, etc.)
-  - Prompt masking with an implementation of [cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff)
-- Simple [prompt macros](/doc/macros.md) with `DEF`
+  - Prompt masking with an implementation of [cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff).
+- Organize complicated prompts with [segments and prompt macros](/doc/macros.md).
 
 All features are fully schedulable unless otherwise stated. See the [scheduling syntax documentation](doc/schedules.md) to get started.
 

__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 follow (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/basic.md

@@ -31,6 +31,7 @@ Prompt operators are processed in the following order, meaning that all features
 
 - DEF macros are expanded
 - Scheduling is expanded, and for each scheduled prompt:
+  - SEGs are processed and the template is expanded
   - The prompt is split by AND, and for each:
     - Prompts are split by COUPLE. and for each:
       - Most functions (like MASK) and cutoffs are evaluated

doc/macros.md

@@ -58,3 +58,42 @@ a "$1" b "$2"
 a "" b "$2"
 a "A" b "$2"
 ```
+
+## SEG: Split your prompt into named segments
+
+Syntax: `SEG(segment_name)`
+
+To help with organizing prompts, you can use the `SEG` function. For example:
+```
+This is a comic
+Top panel: $CAT. $SEG3
+Bottom panel: $DOG
+
+SEG(DOG)
+A dog chasing its
+  tail in a living room.
+SEG(CAT)
+
+a sleeping cat
+
+SEG
+The cat has orange fur with white stripes
+```
+This produces:
+```
+This is a comic
+Top panel: a sleeping cat. The cat has orange fur with white stripes
+Bottom panel: A dog chasing its
+  tail in a living room.
+```
+
+> [!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.
+
+In the template, you can refer to segments by either their index (starting from 1) or the given name, prefixed with a `$SEG`, so in this example, `$SEG1` is the same as `$PANEL1`
+
+Segments can also refer to each other. Recursion will terminate, but produces weird outputs.
+
+Naming segments is optional, in which case you will have to refer to it by its index.

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,20 +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
-        self.kv = {"k": [], "v": []}
+        # 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]
@@ -116,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/cutoff_parser.py

@@ -1,58 +1,25 @@
-from typing import TypeAlias
+import re
 
-import lark
+from .utils import parse_args
 
-from .parser import flatten
-
-cut_parser = lark.Lark(
-    r"""
-!start: (cut | prompt | /[][:()]/+)*
-prompt: (PLAIN | WHITESPACE)+
-cut: "[CUT:" prompt ":" prompt [":" NUMBER  [ ":" NUMBER [":" NUMBER [ ":" PLAIN ] ] ] ]"]"
-WHITESPACE: /\s+/
-PLAIN: /([^\[\]:])+/
-%import common.SIGNED_NUMBER -> NUMBER
-"""
-)
+CUTOFF_RE = re.compile(r"\[CUT:((.*?):(.*?))\]")
 
 
-class CutTransform(lark.Transformer):
-    def __default__(self, data, children, meta):
-        return children
+def noop(x):
+    return x
 
-    def NUMBER(self, args):
-        return float(args)
 
-    def cut(self, args):
-        prompt, cutout, weight, strict_mask, start_from_masked, mask_token = args
-
-        # prompts and cutouts are always sequences of str
-        return (
-            "".join(prompt),
-            "".join(cutout),
-            weight,
-            strict_mask,
-            start_from_masked,
-            mask_token,
+def parse_cuts(string):
+    text = CUTOFF_RE.sub(r"\2", string)
+    cutoffs = CUTOFF_RE.findall(string)
+    cs = []
+    for x, *_ in cutoffs:
+        p = x.split(":")
+        args = parse_args(
+            p, [(str, ""), (str, ""), (float, 0), (float, None), (float, None), (noop, None)], strip=False
         )
-
-    def start(self, args):
-        prompt = []
-        cuts = []
-        for a in flatten(args):
-            if isinstance(a, str):
-                prompt.append(a)
-            else:
-                prompt.append(a[0])
-                cuts.append(a)
-        return "".join(prompt), cuts
-
-    def PLAIN(self, args: str) -> str:
-        return str(args)
-
-
-CutResult: TypeAlias = tuple[str, str, float, float, float, str]
-
-
-def parse_cuts(text: str) -> tuple[str, CutResult]:
-    return CutTransform().transform(cut_parser.parse(text))
+        args = tuple(args)
+        if not args[0] or not args[1] or (args[5] is not None and not args[5].strip()):
+            raise ValueError(f"Invalid CUT spec: [CUT:{x}]")
+        cs.append(args)
+    return text, cs

prompt_control/macros.py

@@ -4,12 +4,52 @@ from __future__ import annotations
 import logging
 import re
 
-from .utils import find_closing_paren, get_function
+from .utils import find_closing_paren, get_function, split_by_function
 
-logging.basicConfig()
 log = logging.getLogger("comfyui-prompt-control")
 
 
+def substitute_template(template, segments, do_subs):
+    def _substitute(template, segments, stack):
+        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, 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, 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 = ""
@@ -60,6 +100,13 @@ def expand_macros(text):
     return res
 
 
+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}{boundary}", replace, text)
+
+
 def substitute_defcall(text, search, replace):
     name, default_args = search
     text, defns = get_function(text, name, defaults=None, placeholder=f"DEFNCALL{name}", require_args=False)
@@ -71,11 +118,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 = re.sub(rf"\${i + 1}\b", v, r)
+            r = substitute_var(r, i + 1, v, boundary=end_re)
 
         for i, v in enumerate(default_args):
-            r = re.sub(rf"\${i + 1}\b", v, r)
+            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

@@ -2,6 +2,7 @@ import logging
 
 from comfy_api.latest import io
 
+from .macros import expand_segs
 from .prompts import encode_prompt
 
 log = logging.getLogger("comfyui-prompt-control")
@@ -25,10 +26,11 @@ 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)
+        text = expand_segs(text)
         out = encode_prompt(clip, text, start, end, defaults, masks)
         return io.NodeOutput(out)
 
@@ -49,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

@@ -3,20 +3,16 @@ from __future__ import annotations
 
 import json
 import logging
-import os
 
 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
 
 log = logging.getLogger("comfyui-prompt-control")
-if os.environ.get("PC_USE_OLD_PARSER", "0") != "0":
-    log.info("Using new parser implementation. Set PC_USE_OLD_PARSER=1 to use old parser instead")
-    from .parser_parsy import parse_prompt_schedules as parse_prompt_schedules
-else:
-    from .parser import parse_prompt_schedules
 
 
 def create_lora_loader_nodes(graph, model, clip, loras):
@@ -163,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)

prompt_control/nodes_tools.py

@@ -2,7 +2,10 @@ import logging
 
 from comfy_api.latest import io
 
-from .parser import expand_macros, parse_prompt_schedules
+from .macros import expand_macros as macroexpand
+from .macros import expand_segs as segexpand
+from .macros import expand_subs as subexpand
+from .parser import parse_prompt_schedules
 
 log = logging.getLogger("comfyui-prompt-control")
 
@@ -131,22 +134,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)
 
 
@@ -166,7 +179,7 @@ class PCMacroExpand(io.ComfyNode):
 
     @classmethod
     def execute(cls, text) -> io.NodeOutput:
-        return io.NodeOutput(expand_macros(text))
+        return io.NodeOutput(macroexpand(text))
 
 
 NODES = [

prompt_control/parser.py

@@ -1,367 +1,10 @@
-# vim: sw=4 ts=4
-from __future__ import annotations
-
 import logging
-from functools import lru_cache
-from math import ceil
+import os
 
-import lark
-
-from .macros import expand_macros
-
-logging.basicConfig()
 log = logging.getLogger("comfyui-prompt-control")
 
-if lark.__version__ == "0.12.0":
-    from sys import executable
-
-    x = "\n".join(
-        [
-            "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!",
-            f"{executable} -m pip uninstall lark-parser lark",
-            f"{executable} -m pip install lark",
-        ]
-    )
-    log.error(x)
-    raise ImportError(x)
-
-
-ESCAPES = [
-    ("XxPCBackslashESCAPExX", "\\"),
-    ("XxPCColonESCAPExX", ":"),
-    ("XxPCCommentESCAPExX", "#"),
-]
-
-
-def escape_specials(string: str) -> str:
-    for ph, c in ESCAPES:
-        string = string.replace(rf"\{c}", ph)
-    return string
-
-
-def restore_escaped(string: str) -> str:
-    for ph, c in ESCAPES:
-        string = string.replace(ph, c)
-    return string
-
-
-def remove_comments(string: str) -> str:
-    r = []
-    for line in string.split("\n"):
-        comment = line.find("#")
-        if comment >= 0:
-            r.append(line[:comment])
-        else:
-            r.append(line)
-    return "\n".join(r)
-
-
-prompt_parser = lark.Lark(
-    r"""
-!start: (prompt | /[][():|]/+)*
-prompt: (emphasized | embedding | scheduled | alternate | sequence | loraspec | PLAIN | | /\\:/ | /</ | />/ | WHITESPACE)+
-!emphasized: "(" prompt? ")"
-        | "(" prompt ":" prompt ")"
-        | "[" prompt "]"
-promptlist: ([prompt] ":")~1..3
-scheduled: "[" promptlist _WS? NUMBER ["," NUMBER] "]"
-        | "[" promptlist _WS? TAG "]"
-sequence.5:  "[SEQ" ":" [prompt] ":" NUMBER (":" [prompt] ":" NUMBER)* "]"
-alternate: "[" [prompt] ("|" [prompt])+ [":" NUMBER] "]"
-loraspec.99: "<lora:" FILENAME lora_weights [lora_block_weights] ">"
-lora_weights.1: (":" _WS? NUMBER)~1..2
-lora_block_weights.-1: ":" PLAIN
-embedding.100: "<emb:" FILENAME ">"
-WHITESPACE: /\s+/
-_WS: WHITESPACE
-PLAIN: /([^<>\\\[\]():|]|\\.)+/
-FILENAME: /[^<>:]+/
-TAG: /[A-Z_]+/
-%import common.SIGNED_NUMBER -> NUMBER
-""",
-    lexer="dynamic",
-)
-
-
-def flatten(x):
-    if type(x) in [str, tuple, int, type(None)] or isinstance(x, dict) and "type" in x:
-        yield x
-    else:
-        for g in x:
-            yield from flatten(g)
-
-
-def clamp(a, b, c):
-    """clamp b between a and c"""
-    return min(max(a, b), c)
-
-
-def get_steps(tree, num_steps):
-    res = [num_steps or 100]
-
-    def tostep(s):
-        steps = num_steps or 100
-        if "." in str(s) or not num_steps:
-            w = float(s)
-            value = w * steps
-        else:
-            w = int(s)
-            value = w
-
-        if w > 1 and not num_steps:
-            log.warning(
-                "You haven't configured the number of steps for Prompt Control to use, %s will be clipped to 1.0", w
-            )
-            value = steps
-
-        return int(clamp(0, value, steps))
-
-    class CollectSteps(lark.Visitor):
-        def scheduled(self, tree):
-            i = tree.children[-1]
-            if i and i.type == "TAG":
-                return
-            for i in [-1, -2]:
-                if tree.children[i] is not None:
-                    tree.children[i] = tostep(tree.children[i])
-                    res.append(tree.children[i])
-
-        def interp_steps(self, tree):
-            tree.children[-1] = tostep(tree.children[-1] or 0.1)
-            for i, _ in enumerate(tree.children[:-1]):
-                tree.children[i] = tostep(tree.children[i])
-
-            res.extend(tree.children[:-1])
-
-        def sequence(self, tree):
-            steps = tree.children[1::2]
-            for i, _ in enumerate(steps):
-                w = tostep(tree.children[i * 2 + 1])
-                tree.children[i * 2 + 1] = w
-                res.append(w)
-
-        def alternate(self, tree):
-            step_size = tostep(round(float(tree.children[-1] or 0.1), 2))
-            tree.children[-1] = step_size
-            res.extend([x for x in range(step_size, num_steps or 100, step_size)])
-
-    CollectSteps().visit(tree)
-
-    return sorted(set(res))
-
-
-def at_step(step, filters, tree):
-    class AtStep(lark.Transformer):
-        def scheduled(self, args):
-            before = None
-            during = None
-            after = None
-            when_end = None
-            pl, when, *rest = args
-            if rest:
-                when_end = rest[0]
-
-            pl = list(pl)
-            if len(pl) == 1:
-                (during,) = pl  # [after:0.5] == [::after:0.5,0.5]
-                if when_end is None:
-                    when_end = when
-                    after = during
-            elif len(pl) == 2:
-                during, after = pl  # [during:after:0.5] = [before::after:0.5,0.5]
-                if when_end is None:
-                    when_end = when
-                    before = during
-            else:
-                before, during, after = pl  # [before:during:after:0.5,0.8]
-
-            if isinstance(when, str):
-                return before or "" if when not in filters else after or ""
-
-            if when_end is None:
-                when_end = 1000_000
-
-            if step <= when:
-                return before or ""
-            if when < step <= when_end:
-                return during or ""
-            else:
-                return after or ""
-
-        def sequence(self, args):
-            previous_step = 0.0
-            prompts = args[::2]
-            steps = args[1::2]
-            for s, p in zip(steps, prompts, strict=False):
-                if s >= step and step >= previous_step:
-                    previous_step = step
-                    return p or ""
-                else:
-                    previous_step = s
-            return ""
-
-        def alternate(self, args):
-            step_size = args[-1]
-            idx = ceil(step / step_size)
-            return args[(idx - 1) % (len(args) - 1)] or ""
-
-        def start(self, args):
-            prompt = []
-            loraspecs = {}
-            args = flatten(args)
-            for a in args:
-                if isinstance(a, str):
-                    prompt.append(a)
-                elif isinstance(a, tuple):
-                    # sum identical specs together
-                    n = a[0]
-                    # if clip weight is not provided, use unet weight
-                    w, w_clip = a[1][0], a[1][1 % len(a[1])]
-                    e = loraspecs.get(n, {})
-                    loraspecs[n] = {
-                        "weight": round(e.get("weight", 0.0) + w, 2),
-                        "weight_clip": round(e.get("weight_clip", 0.0) + w_clip, 2),
-                    }
-                    lbw = a[2]
-                    if lbw:
-                        loraspecs[n]["lbw"] = lbw
-                    if loraspecs[n]["weight"] == 0 and loraspecs[n]["weight_clip"] == 0 and not lbw:
-                        del loraspecs[n]
-                else:
-                    pass
-            p = "".join(prompt)
-            return {"prompt": p, "loras": loraspecs}
-
-        def PLAIN(self, args):
-            return restore_escaped(args)
-
-        def FILENAME(self, value):
-            return str(value)
-
-        def embedding(self, args):
-            return "embedding:" + str(args[0])
-
-        def lora_weights(self, args):
-            return [float(str(a)) for a in args]
-
-        def lora_block_weights(self, args):
-            vals = args[0].split(";")
-            r = {}
-            for v in vals:
-                x = v.split("=", 2)
-                if len(x) != 2:
-                    continue
-                k, v = x[0].strip().upper(), x[1].strip()
-                r[k] = v
-            return r
-
-        def loraspec(self, args):
-            name = args[0]
-            params = args[1]
-            lbw = args[2]
-
-            return name, params, lbw
-
-        def __default__(self, data, children, meta):
-            return children
-
-    return AtStep().transform(tree)
-
-
-class PromptSchedule:
-    # 0 num_steps means unconfigured
-    def __init__(self, prompt, filters="", start=0.0, end=1.0, num_steps=0):
-        self.filters = filters
-        self.start = start
-        self.end = end
-        self.num_steps = num_steps
-        # placeholder is restored on parse
-        self.prompt = remove_comments(escape_specials(prompt.strip()))
-        self.defaults = {}
-        self.loaded_loras = {}
-
-        self.parsed_prompt = self._parse(num_steps)
-
-    def __iter__(self):
-        # Filter out zero, it's only useful for interpolation
-        return (x for x in self.parsed_prompt if x[0] != 0)
-
-    def _parse(self, num_steps):
-        filters = [x.strip() for x in self.filters.upper().split(",")]
-        try:
-            parsed = []
-            tree = prompt_parser.parse(self.prompt)
-            steps = get_steps(tree, num_steps=num_steps)
-
-            def f(x):
-                return round(x / (num_steps or 100), 2)
-
-            for t in steps:
-                p = at_step(t, filters, tree)
-                parsed.append([f(t), p])
-
-        except lark.exceptions.LarkError as e:
-            log.error("Prompt editing parse error: %s", e)
-            parsed = [[1.0, {"prompt": self.prompt, "loras": {}}]]
-            raise
-
-        # Tag filtering may return redundant prompts, so filter them out here
-        res = []
-        prev_end = -1
-
-        for end_at, p in parsed:
-            if end_at < self.start:
-                continue
-            elif end_at <= self.end:
-                res.append([end_at, p])
-                prev_end = end_at
-            elif end_at > self.end and prev_end < self.end:
-                res.append([end_at, p])
-                break
-
-        # Always use the last prompt if everything was filtered
-        if len(res) == 0:
-            res = [[1.0, parsed[-1][1]]]
-
-        final = [res[0]]
-
-        # Clean up duplicates
-        for p in res[1:]:
-            if p[1] != final[-1][1]:
-                final.append(p)
-            else:
-                final[-1][0] = p[0]
-        return final
-
-    def clone(self):
-        return self.with_filters()
-
-    def with_filters(self, filters=None, start=None, end=None, defaults=None):
-        def ifspecified(x, defval):
-            return x if x is not None else defval
-
-        p = PromptSchedule(
-            self.prompt,
-            filters=ifspecified(filters, self.filters),
-            start=ifspecified(start, self.start),
-            end=ifspecified(end, self.end),
-            num_steps=self.num_steps,
-        )
-        return p
-
-    def at_step(self, step, total_steps=1):
-        _, x = self.at_step_idx(step, total_steps)
-        return x
-
-    def at_step_idx(self, step, total_steps=1):
-        for i, x in enumerate(self.parsed_prompt):
-            if x[0] * total_steps >= step:
-                return i, x
-        return len(self.parsed_prompt) - 1, self.parsed_prompt[-1]
-
-
-@lru_cache
-def parse_prompt_schedules(prompt, **kwargs):
-    prompt = expand_macros(prompt)
-    return PromptSchedule(prompt, **kwargs)
+if os.environ.get("PC_USE_OLD_PARSER", "0") != "1":
+    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/parser_lark.py

@@ -0,0 +1,359 @@
+# vim: sw=4 ts=4
+from __future__ import annotations
+
+import logging
+from functools import lru_cache
+from math import ceil
+
+import lark
+
+from .macros import expand_macros
+from .utils import flatten
+
+log = logging.getLogger("comfyui-prompt-control")
+
+if lark.__version__ == "0.12.0":
+    from sys import executable
+
+    x = "\n".join(
+        [
+            "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!",
+            f"{executable} -m pip uninstall lark-parser lark",
+            f"{executable} -m pip install lark",
+        ]
+    )
+    log.error(x)
+    raise ImportError(x)
+
+
+ESCAPES = [
+    ("XxPCBackslashESCAPExX", "\\"),
+    ("XxPCColonESCAPExX", ":"),
+    ("XxPCCommentESCAPExX", "#"),
+]
+
+
+def escape_specials(string: str) -> str:
+    for ph, c in ESCAPES:
+        string = string.replace(rf"\{c}", ph)
+    return string
+
+
+def restore_escaped(string: str) -> str:
+    for ph, c in ESCAPES:
+        string = string.replace(ph, c)
+    return string
+
+
+def remove_comments(string: str) -> str:
+    r = []
+    for line in string.split("\n"):
+        comment = line.find("#")
+        if comment >= 0:
+            r.append(line[:comment])
+        else:
+            r.append(line)
+    return "\n".join(r)
+
+
+prompt_parser = lark.Lark(
+    r"""
+!start: (prompt | /[][():|]/+)*
+prompt: (emphasized | embedding | scheduled | alternate | sequence | loraspec | PLAIN | | /\\:/ | /</ | />/ | WHITESPACE)+
+!emphasized: "(" prompt? ")"
+        | "(" prompt ":" prompt ")"
+        | "[" prompt "]"
+promptlist: ([prompt] ":")~1..3
+scheduled: "[" promptlist _WS? NUMBER ["," NUMBER] "]"
+        | "[" promptlist _WS? TAG "]"
+sequence.5:  "[SEQ" ":" [prompt] ":" NUMBER (":" [prompt] ":" NUMBER)* "]"
+alternate: "[" [prompt] ("|" [prompt])+ [":" NUMBER] "]"
+loraspec.99: "<lora:" FILENAME lora_weights [lora_block_weights] ">"
+lora_weights.1: (":" _WS? NUMBER)~1..2
+lora_block_weights.-1: ":" PLAIN
+embedding.100: "<emb:" FILENAME ">"
+WHITESPACE: /\s+/
+_WS: WHITESPACE
+PLAIN: /([^<>\\\[\]():|]|\\.)+/
+FILENAME: /[^<>:]+/
+TAG: /[A-Z_]+/
+%import common.SIGNED_NUMBER -> NUMBER
+""",
+    lexer="dynamic",
+)
+
+
+def clamp(a, b, c):
+    """clamp b between a and c"""
+    return min(max(a, b), c)
+
+
+def get_steps(tree, num_steps):
+    res = [num_steps or 100]
+
+    def tostep(s):
+        steps = num_steps or 100
+        if "." in str(s) or not num_steps:
+            w = float(s)
+            value = w * steps
+        else:
+            w = int(s)
+            value = w
+
+        if w > 1 and not num_steps:
+            log.warning(
+                "You haven't configured the number of steps for Prompt Control to use, %s will be clipped to 1.0", w
+            )
+            value = steps
+
+        return int(clamp(0, value, steps))
+
+    class CollectSteps(lark.Visitor):
+        def scheduled(self, tree):
+            i = tree.children[-1]
+            if i and i.type == "TAG":
+                return
+            for i in [-1, -2]:
+                if tree.children[i] is not None:
+                    tree.children[i] = tostep(tree.children[i])
+                    res.append(tree.children[i])
+
+        def interp_steps(self, tree):
+            tree.children[-1] = tostep(tree.children[-1] or 0.1)
+            for i, _ in enumerate(tree.children[:-1]):
+                tree.children[i] = tostep(tree.children[i])
+
+            res.extend(tree.children[:-1])
+
+        def sequence(self, tree):
+            steps = tree.children[1::2]
+            for i, _ in enumerate(steps):
+                w = tostep(tree.children[i * 2 + 1])
+                tree.children[i * 2 + 1] = w
+                res.append(w)
+
+        def alternate(self, tree):
+            step_size = tostep(round(float(tree.children[-1] or 0.1), 2))
+            tree.children[-1] = step_size
+            res.extend([x for x in range(step_size, num_steps or 100, step_size)])
+
+    CollectSteps().visit(tree)
+
+    return sorted(set(res))
+
+
+def at_step(step, filters, tree):
+    class AtStep(lark.Transformer):
+        def scheduled(self, args):
+            before = None
+            during = None
+            after = None
+            when_end = None
+            pl, when, *rest = args
+            if rest:
+                when_end = rest[0]
+
+            pl = list(pl)
+            if len(pl) == 1:
+                (during,) = pl  # [after:0.5] == [::after:0.5,0.5]
+                if when_end is None:
+                    when_end = when
+                    after = during
+            elif len(pl) == 2:
+                during, after = pl  # [during:after:0.5] = [before::after:0.5,0.5]
+                if when_end is None:
+                    when_end = when
+                    before = during
+            else:
+                before, during, after = pl  # [before:during:after:0.5,0.8]
+
+            if isinstance(when, str):
+                return before or "" if when not in filters else after or ""
+
+            if when_end is None:
+                when_end = 1000_000
+
+            if step <= when:
+                return before or ""
+            if when < step <= when_end:
+                return during or ""
+            else:
+                return after or ""
+
+        def sequence(self, args):
+            previous_step = 0.0
+            prompts = args[::2]
+            steps = args[1::2]
+            for s, p in zip(steps, prompts, strict=False):
+                if s >= step and step >= previous_step:
+                    previous_step = step
+                    return p or ""
+                else:
+                    previous_step = s
+            return ""
+
+        def alternate(self, args):
+            step_size = args[-1]
+            idx = ceil(step / step_size)
+            return args[(idx - 1) % (len(args) - 1)] or ""
+
+        def start(self, args):
+            prompt = []
+            loraspecs = {}
+            args = flatten(args)
+            for a in args:
+                if isinstance(a, str):
+                    prompt.append(a)
+                elif isinstance(a, tuple):
+                    # sum identical specs together
+                    n = a[0]
+                    # if clip weight is not provided, use unet weight
+                    w, w_clip = a[1][0], a[1][1 % len(a[1])]
+                    e = loraspecs.get(n, {})
+                    loraspecs[n] = {
+                        "weight": round(e.get("weight", 0.0) + w, 2),
+                        "weight_clip": round(e.get("weight_clip", 0.0) + w_clip, 2),
+                    }
+                    lbw = a[2]
+                    if lbw:
+                        loraspecs[n]["lbw"] = lbw
+                    if loraspecs[n]["weight"] == 0 and loraspecs[n]["weight_clip"] == 0 and not lbw:
+                        del loraspecs[n]
+                else:
+                    pass
+            p = "".join(prompt)
+            return {"prompt": p, "loras": loraspecs}
+
+        def PLAIN(self, args):
+            return restore_escaped(args)
+
+        def FILENAME(self, value):
+            return str(value)
+
+        def embedding(self, args):
+            return "embedding:" + str(args[0])
+
+        def lora_weights(self, args):
+            return [float(str(a)) for a in args]
+
+        def lora_block_weights(self, args):
+            vals = args[0].split(";")
+            r = {}
+            for v in vals:
+                x = v.split("=", 2)
+                if len(x) != 2:
+                    continue
+                k, v = x[0].strip().upper(), x[1].strip()
+                r[k] = v
+            return r
+
+        def loraspec(self, args):
+            name = args[0]
+            params = args[1]
+            lbw = args[2]
+
+            return name, params, lbw
+
+        def __default__(self, data, children, meta):
+            return children
+
+    return AtStep().transform(tree)
+
+
+class PromptSchedule:
+    # 0 num_steps means unconfigured
+    def __init__(self, prompt, filters="", start=0.0, end=1.0, num_steps=0):
+        self.filters = filters
+        self.start = start
+        self.end = end
+        self.num_steps = num_steps
+        # placeholder is restored on parse
+        self.prompt = remove_comments(escape_specials(prompt.strip()))
+        self.defaults = {}
+        self.loaded_loras = {}
+
+        self.parsed_prompt = self._parse(num_steps)
+
+    def __iter__(self):
+        # Filter out zero, it's only useful for interpolation
+        return (x for x in self.parsed_prompt if x[0] != 0)
+
+    def _parse(self, num_steps):
+        filters = [x.strip() for x in self.filters.upper().split(",")]
+        try:
+            parsed = []
+            tree = prompt_parser.parse(self.prompt)
+            steps = get_steps(tree, num_steps=num_steps)
+
+            def f(x):
+                return round(x / (num_steps or 100), 2)
+
+            for t in steps:
+                p = at_step(t, filters, tree)
+                parsed.append([f(t), p])
+
+        except lark.exceptions.LarkError as e:
+            log.error("Prompt editing parse error: %s", e)
+            parsed = [[1.0, {"prompt": self.prompt, "loras": {}}]]
+            raise
+
+        # Tag filtering may return redundant prompts, so filter them out here
+        res = []
+        prev_end = -1
+
+        for end_at, p in parsed:
+            if end_at < self.start:
+                continue
+            elif end_at <= self.end:
+                res.append([end_at, p])
+                prev_end = end_at
+            elif end_at > self.end and prev_end < self.end:
+                res.append([end_at, p])
+                break
+
+        # Always use the last prompt if everything was filtered
+        if len(res) == 0:
+            res = [[1.0, parsed[-1][1]]]
+
+        final = [res[0]]
+
+        # Clean up duplicates
+        for p in res[1:]:
+            if p[1] != final[-1][1]:
+                final.append(p)
+            else:
+                final[-1][0] = p[0]
+        return final
+
+    def clone(self):
+        return self.with_filters()
+
+    def with_filters(self, filters=None, start=None, end=None, defaults=None):
+        def ifspecified(x, defval):
+            return x if x is not None else defval
+
+        p = PromptSchedule(
+            self.prompt,
+            filters=ifspecified(filters, self.filters),
+            start=ifspecified(start, self.start),
+            end=ifspecified(end, self.end),
+            num_steps=self.num_steps,
+        )
+        return p
+
+    def at_step(self, step, total_steps=1):
+        _, x = self.at_step_idx(step, total_steps)
+        return x
+
+    def at_step_idx(self, step, total_steps=1):
+        for i, x in enumerate(self.parsed_prompt):
+            if x[0] * total_steps >= step:
+                return i, x
+        return len(self.parsed_prompt) - 1, self.parsed_prompt[-1]
+
+
+@lru_cache
+def parse_prompt_schedules(prompt, **kwargs):
+    prompt = expand_macros(prompt)
+    return PromptSchedule(prompt, **kwargs)

prompt_control/parser_parsy.py

@@ -348,9 +348,14 @@ non_special = regex(r"[^:\[\]()|\\<>#]+").map(Text)
 filename = regex(r"[^:<>]+")
 
 comment = string("#") >> any_char.until(eof | char_from("\n")) >> success(empty)
-escape = (string("\\") >> char_from("\\[]:#")).map(Text)
+escape = (string("\\") >> char_from("\\[]:#") | string(r"\(") | string(r"\)")).map(Text)
 emphasis = seq(lpar, (prompt | col).at_least(0), rpar)
-number = (digit.at_least(1) + string(".") * 1 + digit.many() | digit.at_least(1)).concat().map(float)
+sign = string("+") | string("-")
+number = (
+    (sign.optional("") + (digit.many() + string(".") * 1 + digit.many() | digit.at_least(1)).concat())
+    .concat()
+    .map(float)
+)
 
 opt_prompt = prompt.optional(empty)
 step_range = seq(number | tag, (comma >> number).optional())

prompt_control/prompts.py

@@ -241,7 +241,7 @@ def encode_prompt_segment(
     can_break = {}
     for k in empty:
         tokenizer = getattr(clip.tokenizer, f"clip_{k}", getattr(clip.tokenizer, k, None))
-        can_break[k] = tokenizer and tokenizer.pad_to_max_length
+        can_break[k] = tokenizer and getattr(tokenizer, "pad_to_max_length", False)
 
     clip = hook_te(clip, empty.keys(), style, normalization, extra)
 
@@ -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]

prompt_control/utils.py

@@ -35,6 +35,14 @@ except ImportError:
 log = logging.getLogger("comfyui-prompt-control")
 
 
+def flatten(x):
+    if type(x) in [str, tuple, int, type(None)] or isinstance(x, dict) and "type" in x:
+        yield x
+    else:
+        for g in x:
+            yield from flatten(g)
+
+
 def call_node(cls, *args, **kwargs):
     if hasattr(cls, "execute"):
         # v3 node

pyproject.toml

@@ -1,10 +1,8 @@
 [project]
 name = "comfyui-prompt-control"
-description = "Provides nodes for prompt editing and LoRA scheduling, advanced regional prompting (including attention masking) and more, all controlled through your text prompt"
-version = "3.0.0-beta.1"
+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.6"
 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"]
 
 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_macros.py

@@ -0,0 +1,72 @@
+from textwrap import dedent
+
+import pytest
+
+from prompt_control.macros import expand_macros, expand_segs
+
+
+@pytest.mark.parametrize(
+    "text, result",
+    [
+        ("DEF(X(a;b)=$1 $2 $3 d)X(A) X(A;B;C)", "A b $3 d A B C d"),
+        (
+            "DEF(MACRO()=[empty:$1:$2])MACRO MACRO(;) MACRO(;0.5) MACRO(a;0.5)",
+            "[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):
+    assert expand_macros(text) == result
+
+
+def test_macro_recursion():
+    with pytest.raises(ValueError) as c:
+        expand_macros("DEF(X=recurse Y) DEF(Y=recurse X) X")
+    assert "Unable to resolve DEFs" in str(c.value)
+
+
+@pytest.mark.parametrize(
+    "input, output",
+    [
+        (
+            """\
+      A red $b and
+      a blue $a
+      SEG(a)
+      cat
+      SEG(b)
+
+      dog
+      SEG(c)""",
+            "A red dog and\na blue cat",
+        ),
+        (
+            """\
+    $a and $b
+    SEG(a)
+    cat, $b
+    SEG(b)
+    dog, $c
+    SEG(c)
+    tiger
+    """,
+            "cat, dog, tiger and dog, tiger",
+        ),
+        (
+            """\
+    $a
+    SEG(a)
+    a $b
+    SEG(b)
+    b $a""",
+            "a b a b $a",
+        ),
+    ],
+)
+def test_segments(input, output):
+    assert expand_segs(dedent(input)) == output

tests/test_parser.py

@@ -2,10 +2,6 @@ import os
 
 import pytest
 
-from prompt_control.parser import expand_macros
-from prompt_control.parser import parse_prompt_schedules as old_parse  # noqa
-from prompt_control.parser_parsy import parse_prompt_schedules as new_parse  # noqa
-
 
 def lora_dict(*loras):
     return {lora: {"weight": unet, "weight_clip": te} for lora, unet, te in loras}
@@ -27,9 +23,13 @@ parsers_to_test = os.environ.get("PC_PARSERS_TO_TEST", "new").split()
 
 params = []
 if "old" in parsers_to_test:
+    from prompt_control.parser_lark import parse_prompt_schedules as old_parse  # noqa
+
     params.append(old_parse)
 
 if "new" in parsers_to_test:
+    from prompt_control.parser_parsy import parse_prompt_schedules as new_parse  # noqa
+
     params.append(new_parse)
 
 
@@ -111,9 +111,9 @@ def test_basic_ok(parse):
 
 @pytest.mark.parametrize("step", [0, 0.5, 1])
 def test_lora(step, parse):
-    p = parse("This is a (lora:0.6) (prompt) with [no scheduling] features <lora:foo:0.5> <lora:bar:0.5:1.0>")
+    p = parse("This is a (lora:0.6) (prompt) with [no scheduling] features <lora:foo:0.5> <lora:bar:0.5:-1.0>")
     expected = prompt(
-        1.0, "This is a (lora:0.6) (prompt) with [no scheduling] features  ", ("foo", 0.5, 0.5), ("bar", 0.5, 1.0)
+        1.0, "This is a (lora:0.6) (prompt) with [no scheduling] features  ", ("foo", 0.5, 0.5), ("bar", 0.5, -1.0)
     )
     assert prompts_match(p.at_step(step), expected)
 
@@ -221,23 +221,6 @@ def test_def(parse):
     p2 = parse("[(test):(test:0.7):0.7] [(test):(test:0.5):0.5]")
     assert p.parsed_prompt == p2.parsed_prompt
 
-    p = expand_macros("DEF(X(a;b)=$1 $2 $3 d)X(A) X(A;B;C)")
-    assert p == "A b $3 d A B C d"
-
-    p = expand_macros("DEF(MACRO()=[empty:$1:$2])MACRO MACRO(;) MACRO(;0.5) MACRO(a;0.5)")
-    assert p == "[empty::$2] [empty::] [empty::0.5] [empty:a:0.5]"
-
-    p = expand_macros("DEF(X=$1)DEF(Y()=$1)[X Y][X() Y()][X(1) Y(1)]")
-    assert p == "[$1 ][ ][1 1]"
-
-    p = parse("DEF(test(1)=prompt $1)DEF(test2((a); (test))=[$1:$2:0.5])test test2")
-    p2 = parse("prompt 1 [(a):(prompt 1):0.5]")
-    assert p.parsed_prompt == p2.parsed_prompt
-
-    with pytest.raises(ValueError) as c:
-        expand_macros("DEF(X=recurse Y) DEF(Y=recurse X) X")
-    assert "Unable to resolve DEFs" in str(c.value)
-
 
 @pytest.mark.parametrize(
     "text, cases",
@@ -249,6 +232,7 @@ def test_def(parse):
         ),
         (r"[a\:b\\:c:0.5]", [(0.0, 0.5, "a:b\\"), (0.55, 1, r"c")]),
         (r"[a:\#b:0.5]", [(0.0, 0.5, "a"), (0.55, 1, "#b")]),
+        (r"[a:b \(test\):0.2]", [(0, 0.2, r"a"), (0.25, 1, r"b \(test\)")]),
     ],
 )
 def test_escapes(text, cases, parse):
@@ -358,6 +342,12 @@ def test_cornercase_corrected(parse):
     assert p.parsed_prompt[1:] == p2.parsed_prompt
 
 
+def test_floats(parse):
+    p = parse("[a:b:0.5] [c:d:e:0.2,0.7] <lora:test:-0.3>")
+    p2 = parse("[a:b:.5] [c:d:e:.2,.7] <lora:test:-.3>")
+    assert p.parsed_prompt == p2.parsed_prompt
+
+
 def test_alternating_lora(parse):
     p4 = parse("[cat|[dog:wolf<lora:canine:1>:0.5]:0.2]")
     for i, (text, *_loras) in enumerate(

tests/test_workflow.json

@@ -0,0 +1,139 @@
+{
+  "1": {
+    "inputs": {
+      "text": "positive prompt",
+      "clip": [
+        "4",
+        1
+      ]
+    },
+    "class_type": "PCLazyTextEncode",
+    "_meta": {
+      "title": "PC: Schedule prompt"
+    }
+  },
+  "2": {
+    "inputs": {
+      "ckpt_name": "$TEST_CHECKPOINT"
+    },
+    "class_type": "CheckpointLoaderSimple",
+    "_meta": {
+      "title": "Load Checkpoint"
+    }
+  },
+  "3": {
+    "inputs": {
+      "seed": 0,
+      "steps": 8,
+      "cfg": 3,
+      "sampler_name": "euler",
+      "scheduler": "simple",
+      "denoise": 1,
+      "model": [
+        "4",
+        0
+      ],
+      "positive": [
+        "9",
+        0
+      ],
+      "negative": [
+        "9",
+        1
+      ],
+      "latent_image": [
+        "5",
+        0
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "KSampler"
+    }
+  },
+  "4": {
+    "inputs": {
+      "text": "<lora:$TEST_LORA:1>",
+      "model": [
+        "2",
+        0
+      ],
+      "clip": [
+        "2",
+        1
+      ]
+    },
+    "class_type": "PCLazyLoraLoader",
+    "_meta": {
+      "title": "PC: Schedule LoRAs"
+    }
+  },
+  "5": {
+    "inputs": {
+      "width": 1024,
+      "height": 1024,
+      "batch_size": 1
+    },
+    "class_type": "EmptyLatentImage",
+    "_meta": {
+      "title": "Empty Latent Image"
+    }
+  },
+  "6": {
+    "inputs": {
+      "samples": [
+        "3",
+        0
+      ],
+      "vae": [
+        "2",
+        2
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    }
+  },
+  "7": {
+    "inputs": {
+      "images": [
+        "6",
+        0
+      ]
+    },
+    "class_type": "PreviewImage",
+    "_meta": {
+      "title": "Preview Image"
+    }
+  },
+  "8": {
+    "inputs": {
+      "text": "worst quality,",
+      "clip": [
+        "4",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "9": {
+    "inputs": {
+      "positive": [
+        "1",
+        0
+      ],
+      "negative": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "PCAttentionCoupleBatchNegative",
+    "_meta": {
+      "title": "PC: Attention Couple (batch negative)"
+    }
+  }
+}

tests/test_workflow.py

@@ -0,0 +1,40 @@
+import json
+import os
+import uuid
+from time import sleep
+
+import pytest
+import requests
+
+
+@pytest.fixture(scope="module", autouse=True)
+def workflow(request):
+    with open(str(request.path).replace(".py", ".json")) as f:
+        data = f.read()
+    data = data.replace("$TEST_CHECKPOINT", os.environ["PC_TEST_CHECKPOINT"])
+    data = data.replace("$TEST_LORA", os.environ["PC_TEST_LORA"])
+    return json.loads(data)
+
+
+def assert_prompt(url, p):
+    timeout = 60
+    r = requests.post(f"{url}/prompt", json={"prompt": p, "client_id": str(uuid.uuid4())}).json()
+    prompt_id = r["prompt_id"]
+    r = {"status": "pending"}
+    while r["status"] in ["pending", "in_progress"]:
+        sleep(1)
+        assert timeout > 0
+        timeout -= 1
+        r = requests.get(f"{url}/api/jobs/{prompt_id}").json()
+    assert r["status"] == "completed"
+
+
+@pytest.fixture
+def comfyui():
+    return os.environ.get("PC_TEST_COMFYUI", "http://localhost:8188")
+
+
+def test_workflow(workflow, comfyui):
+    prompt = "DEF(blue=green)a blue dog and a cat sitting [COUPLE(0 0.5, 0 1) red (cat,:1.3) COUPLE(0.5 1, 0 1) (blue:1.2) dog,:0.1]"
+    workflow["1"]["inputs"]["text"] = prompt
+    assert_prompt(comfyui, workflow)