v3.0.0-beta.10

It seems to be fixed now. Schedules appear to be properly lazy.

README.md

@@ -11,6 +11,9 @@ A `Basic Text to Image` template is included with the extension, and can be load
 > The parser was rewritten using parsy. It is intended to have the same behaviour as the old parser, but is **significantly** faster.
 > Please report any bugs or incompatibilities you find.
 
+## Notable changes
+
+- `PC: Schedule Prompt` now strips surrounding whitespace by default, which may change some prompts. Add `NOSTRIP()` to your prompt to restore previous behaviour.
 
 ## What can it do?
 
@@ -23,6 +26,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 +56,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
@@ -80,6 +80,4 @@ This node configures `PCTextEncode` default values for some functions by attachi
 
 # Known issues
 
-- ComfyUI's caching mechanism has an issue that makes it unnecessarily invalidate caches for certain inputs; you'll still get some benefit from the lazy nodes, but changing inputs that shouldn't affect downstream nodes (especially if using filtering) will still cause them to be recomputed because ComfyUI doesn't realize the inputs haven't changed.
-
 - Cutoff does not work with models that use non-CLIP text encoders, like Flux. This might be fixable, but it's uncertain if cutoff even makes sense for those models.

doc/attention_couple.md

@@ -24,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:
@@ -31,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,41 @@
+# 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.
+
+> [!NOTE]
+> When using NODE, you lose access to *all* special syntax provided by `PCTextEncode`. Only SEGs, macros and scheduling will continue to work since those are processed at graph expansion time before the text prompt is passed into the node.
+
+## 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.

doc/schedules.md

@@ -14,6 +14,9 @@ Besides the syntax documented below, the [basic syntax](/doc/basic.md) and [prom
 a [large::0.1] [cat|dog:0.05] [<lora:somelora:0.5:0.6>::0.5]
 [in a park:in space:0.4]
 ```
+## Note on whitespace
+`PC: Schedule Prompt` will strip leading and following whitespace from the prompt automatically. If you really want whitespace in your prompt, include `NOSTRIP()` in your prompt.
+
 ## Comments and escaping
 
 In schedules, any text on a line following a `#` is considered a comment and removed, including the `#` character.

prompt_control/anima_couple.py

@@ -1,6 +1,7 @@
 # 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
@@ -14,18 +15,6 @@ from comfy.samplers import process_conds
 COND = 0
 UNCOND = 1
 
-COND_NEGPIP_MASK_KEY = "c_ppm_negpip_mask"
-NEGPIP_MASKS_COUPLE_KEY = "ppm_couple_negpip_masks"
-NEGPIP_MASK_KEY = "ppm_negpip_mask"
-
-
-class CoupleForward:
-    def __init__(self, fn):
-        self.fn = fn
-
-    def __call__(self, *args, **kwargs):
-        return cosmos_attention_forward_couple(self.fn, *args, **kwargs)
-
 
 def reshape_mask(mask: torch.Tensor, size: tuple[int, int], bs: int, num_tokens: int) -> torch.Tensor:
     num_conds = mask.shape[0]
@@ -42,7 +31,7 @@ def wrap_forwards(anima_model):
         (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 = CoupleForward(b.forward)
+        b.forward = partial(cosmos_attention_forward_couple, b.forward)
     return backups
 
 
@@ -70,17 +59,10 @@ def anima_sample_wrapper(executor, *args, **kwargs):
             seed,
             latent_shapes=[latent_image.shape],
         )
-        conds_p = [
+        return [
             c["model_conds"]["c_crossattn"].cond * pc_conds[i][1].get("strength", 1.0)
             for i, c in enumerate(conds["positive"])
         ]
-        # TODO: weight?
-        negpip_masks_couple = []
-        if all(COND_NEGPIP_MASK_KEY in cond["model_conds"] for cond in conds["positive"]):
-            negpip_masks_couple = [
-                cond["model_conds"][COND_NEGPIP_MASK_KEY].cond.to(device) for cond in conds["positive"]
-            ]
-        return conds_p, negpip_masks_couple
 
     extra_options["model_options"]["transformer_options"]["pc_process_conds"] = pc_process_conds
     return executor(*args, **kwargs)
@@ -94,9 +76,7 @@ def anima_forward_wrapper(executor: WrapperExecutor, *args, **kwargs):
     transformer_options: dict = kwargs.get("transformer_options", {}).copy()
     pc = transformer_options.get("pc_couple")
     if pc and "processed_conds" not in pc:
-        conds, negpip_masks = transformer_options["pc_process_conds"](pc["conds"])
-        pc["processed_conds"] = conds
-        transformer_options[COND_NEGPIP_MASK_KEY] = negpip_masks if negpip_masks else None
+        pc["processed_conds"] = transformer_options["pc_process_conds"](pc["conds"])
     patch_spatial = anima_model.patch_spatial
 
     activations_shape = list(x.shape)
@@ -127,7 +107,6 @@ def cosmos_attention_forward_couple(_forward: Callable, x, context, rope_emb, tr
 
     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"]
@@ -136,9 +115,6 @@ def cosmos_attention_forward_couple(_forward: Callable, x, context, rope_emb, tr
     num_chunks = len(cond_or_uncond)
     bs = x.shape[0] // num_chunks
 
-    n = transformer_options.get(NEGPIP_MASK_KEY)
-    negpip_masks = transformer_options.get(NEGPIP_MASKS_COUPLE_KEY)
-    has_negpip = n is not None and negpip_masks is not None
     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)
@@ -147,16 +123,7 @@ def cosmos_attention_forward_couple(_forward: Callable, x, context, rope_emb, tr
         dim=0,
     )
 
-    if has_negpip:
-        n_chunks = n.chunk(num_chunks, dim=0)
-        num_tokens_n: list[int] = [mask.shape[1] for mask in negpip_masks]
-        lcm_tokens_n = lcm(*(num_tokens_n + [n.shape[1]]))
-        conds_n_tensor = torch.cat(
-            [mask.repeat(bs, lcm_tokens_n // num_tokens_n[i], 1) for i, mask in enumerate(negpip_masks)],
-            dim=0,
-        )
-
-    xs, cs, ns = [], [], []
+    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)
@@ -169,20 +136,9 @@ def cosmos_attention_forward_couple(_forward: Callable, x, context, rope_emb, tr
             cs.append(torch.cat([c_target, conds_c_tensor], dim=0))
             cond_or_uncond_couple.extend(itertools.repeat(COND, num_conds))
 
-        if has_negpip:
-            n_target = n_chunks[i].repeat(1, lcm_tokens_n // n.shape[1], 1)
-            if cond_type == UNCOND:
-                ns.append(x_target)
-            else:
-                ns.append(torch.cat([n_target, conds_n_tensor], dim=0))
-
     xs = torch.cat(xs, dim=0)
     cs = torch.cat(cs, dim=0)
 
-    if has_negpip:
-        ns = torch.cat(ns, dim=0)
-        transformer_options[NEGPIP_MASK_KEY] = ns
-
     out = _forward(xs, cs, rope_emb, transformer_options)
 
     size = tuple(transformer_options["activations_shape"][-2:])

prompt_control/attention_couple_ppm.py

@@ -52,7 +52,7 @@ class Proxy:
         return self
 
     def __call__(self, *args, **kwargs):
-        return self.function(*args, *kwargs)
+        return self.function(*args, **kwargs)
 
 
 class AttentionCoupleHook(TransformerOptionsHook):

prompt_control/macros.py

@@ -69,8 +69,12 @@ def parse_search(search):
     return name, args
 
 
-def expand_macros(text):
-    text, defs = get_function(text, "DEF", defaults=None)
+def expand_macros(text, defs=None):
+    silent = False
+    if defs is None:
+        text, defs = get_function(text, "DEF", defaults=None)
+    else:
+        silent = True
     res = text
     prevres = text
     replacements = []
@@ -95,7 +99,8 @@ def expand_macros(text):
         prevres = res
     if res.strip() != text.strip():
         res = res.strip()
-        log.debug("DEFs expanded to: %s", res)
+        if not silent:
+            log.debug("DEFs expanded to: %s", res)
     return res
 
 
@@ -108,11 +113,8 @@ def substitute_var(text, name, replace, boundary=r"\b"):
 
 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)
-    for i, d in enumerate(defns):
-        ph = d.placeholder
-        assert ph is not None, "This is a bug"
-        parameters = d.args
+
+    def run_macro(*parameters):
         paramvals = []
         if parameters:
             paramvals = [x.strip() for x in parameters[0].split(";")]
@@ -123,6 +125,7 @@ def substitute_defcall(text, search, replace):
 
         for i, v in enumerate(default_args):
             r = substitute_var(r, i + 1, v, boundary=end_re)
+        return r
 
-        text = text.replace(ph, r)
+    text, _ = get_function(text, name, defaults=None, processor=run_macro, require_args=False)
     return text

prompt_control/nodes_anima.py

@@ -3,7 +3,6 @@
 
 import comfy.model_management
 import comfy.patcher_extension
-from comfy.ldm.cosmos.predict2 import Attention as CosmosAttention
 from comfy.model_base import Anima
 from comfy.model_patcher import ModelPatcher
 from comfy_api.latest import io
@@ -11,23 +10,9 @@ from comfy_api.latest import io
 from .anima_couple import (
     anima_forward_wrapper,
     anima_sample_wrapper,
-    cosmos_attention_forward_couple,
 )
 
 
-class CoupleForward:
-    def __init__(self, fn, block):
-        self.fn = fn
-        self.block = block
-
-    def __call__(self, *args, **kwargs):
-        transformer_options = kwargs["transformer_options"]
-        pc = transformer_options.get("pc_couple")
-        if pc:
-            self.block.to("cuda")
-        return cosmos_attention_forward_couple(self.fn, *args, **kwargs)
-
-
 class PCAnimaAttnCouplePatch(io.ComfyNode):
     @classmethod
     def define_schema(cls) -> io.Schema:

prompt_control/nodes_base.py

@@ -3,7 +3,7 @@ import logging
 from comfy_api.latest import io
 
 from .macros import expand_segs
-from .prompts import encode_prompt
+from .prompts import encode_prompt, hook_te
 
 log = logging.getLogger("comfyui-prompt-control")
 
@@ -56,7 +56,31 @@ class PCTextEncode(io.ComfyNode):
         return PCTextEncodeWithRange.execute(clip, text, 0.0, 1.0)
 
 
-NODES = [
-    PCTextEncodeWithRange,
-    PCTextEncode,
-]
+class PCHookEncoderModsInternal(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="PCHookTextEncoderModsInternal",
+            display_name="PC: Apply Text Encoder Mods",
+            category="promptcontrol",
+            description="Apply TE modifications (internal)",
+            is_experimental=True,
+            is_dev_only=True,
+            inputs=[
+                io.Clip.Input("clip"),
+                io.String.Input("te_names"),
+                io.String.Input("style"),
+                io.String.Input("normalization"),
+                io.Custom("PC_EXTRA_DATA").Input("extra", optional=True),
+            ],
+            outputs=[io.Clip.Output()],
+        )
+
+    @classmethod
+    def execute(cls, clip, te_names, style, normalization, extra) -> io.NodeOutput:
+        te_names = [x.strip() for x in te_names.split(",")]
+        clip = hook_te(clip, te_names, style, normalization, extra)
+        return io.NodeOutput(clip)
+
+
+NODES = [PCTextEncodeWithRange, PCTextEncode, PCHookEncoderModsInternal]

prompt_control/nodes_lazy.py

@@ -8,8 +8,9 @@ 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
+from .utils import consolidate_schedule, find_nonscheduled_loras, get_function, split_by_function
 
 log = logging.getLogger("comfyui-prompt-control")
 
@@ -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)
@@ -190,25 +192,106 @@ class PCLazyLoraLoader(io.ComfyNode):
         return io.NodeOutput(*no.args[:2], expand=no.expand)
 
 
+def parse_extra_inputs(args, defaults):
+    params = {}
+    if not args.strip():
+        return defaults + [{}]
+    defaults = defaults[:]
+    defaults.append("")
+    for i, v in enumerate(args.split(",", maxsplit=len(defaults) - 1)):
+        defaults[i] = v
+        # We should strip extra whitespace so that people don't have to worry about functions.
+    magic_spec = defaults[-1]
+    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:
+            params[name.strip()] = json.loads(jsondata.strip())
+        except ValueError as e:
+            raise ValueError(f"Invalid JSON input: '{jsondata}'") from e
+    return [x.strip() for x in defaults[:-1]] + [params]
+
+
+def make_node(graph, p, clip, strip):
+    p, classnames = get_function(p, "NODE", defaults=None)
+    p, filters = get_function(p, "FILTER", defaults=None)
+    args = ""
+    if len(classnames) > 1:
+        log.warning("You have more than one NODE call in your prompt. Only the first one will be used")
+    if classnames:
+        args = classnames[0].args[0]
+        if not args.strip():
+            raise ValueError("NODE can't be empty!")
+    classname, paramname, extras = parse_extra_inputs(args, ["PCTextEncode", "text"])
+    # We should strip extra whitespace so that people don't have to worry about functions.
+    node = graph.node(classname.strip())
+    node.set_input("clip", clip)
+    node.set_input(paramname.strip(), p.strip() if strip else p)
+    for e, v in extras.items():
+        node.set_input(e, v)
+
+    for f in filters:
+        classname, paramname, extras = parse_extra_inputs(f.args[0], ["", "conditioning"])
+        if not classname:
+            raise ValueError("FILTER requires a Node class name")
+        extras[paramname] = node.out(0)
+        node = graph.node(classname)
+        for e, v in extras.items():
+            node.set_input(e, v)
+
+    return node
+
+
+def build_prompt(graph, prompt, clip, start=None, end=None):
+    p = prompt
+    strip = "NOSTRIP()" not in p
+    p = p.replace("NOSTRIP()", "")
+    # Need to explicitly expand SEGs here *before* NODE is processed
+    p = expand_segs(p)
+    p, combines = split_by_function(p, "COMBINE")
+    current_cond = make_node(graph, p, clip, strip)
+    for text, f in combines:
+        classname, param1, param2, extra = parse_extra_inputs(f.args[0], ["", "conditioning_1", "conditioning_2"])
+        if classname.strip() == "":
+            raise ValueError("Can't use COMBINE without a class name")
+        combiner = graph.node(classname.strip())
+        c2 = make_node(graph, text, clip, strip)
+        extra[param1] = current_cond.out(0)
+        extra[param2] = c2.out(0)
+        for e, v in extra.items():
+            combiner.set_input(e, v)
+        current_cond = combiner
+
+    node = current_cond
+    if start is not None and end is not None:
+        node = graph.node("ConditioningSetTimestepRange")
+        node.set_input("conditioning", current_cond.out(0))
+        node.set_input("start", start)
+        node.set_input("end", end)
+
+    return node
+
+
 def build_scheduled_prompts(graph, schedules, clip):
     nodes = []
     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"
-        if classnames:
-            classname, paramname = classnames[0].args
-        node = graph.node(classname)
-        node.set_input("clip", clip)
-        node.set_input(paramname, p)
-        timestep = graph.node("ConditioningSetTimestepRange")
-        timestep.set_input("conditioning", node.out(0))
-        timestep.set_input("start", start_pct)
-        timestep.set_input("end", end_pct)
-        nodes.append(timestep)
+        node = build_prompt(graph, p, clip, start_pct, end_pct)
+        nodes.append(node)
         start_pct = end_pct
+
     node = nodes[0]
     for othernode in nodes[1:]:
         combiner = graph.node("ConditioningCombine")
@@ -272,9 +355,49 @@ class PCLazyTextEncode(io.ComfyNode):
         return PCLazyTextEncodeAdvanced.execute(clip, text)
 
 
+predefined_macros = get_function(
+    """
+DEF(AND=COMBINE(ConditioningCombine, conditioning_1, conditioning_2))
+DEF(CAT=COMBINE(ConditioningConcat, conditioning_to, conditioning_from))
+DEF(AVG(0.5)=COMBINE(ConditioningAverage, conditioning_from, conditioning_to, conditioning_to_strength $1))
+""",
+    "DEF",
+    defaults=None,
+)
+
+
+class PCLazyTextEncodeSingle(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="PCLazyTextEncodeSingle",
+            display_name="PC: Prompt (without scheduling)",
+            is_experimental=True,
+            is_dev_only=True,
+            enable_expand=True,
+            category="promptcontrol",
+            inputs=[
+                io.Clip.Input("clip", raw_link=True),
+                io.String.Input("text", multiline=True, default=""),
+            ],
+            outputs=[
+                io.Conditioning.Output("conditioning"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, clip, text):
+        graph = GraphBuilder()
+        text = expand_macros(text, predefined_macros)
+        node = build_prompt(graph, text, clip)
+        g = graph.finalize()
+        return io.NodeOutput(node.out(0), expand=g)
+
+
 NODES = [
     PCLazyTextEncode,
     PCLazyTextEncodeAdvanced,
+    PCLazyTextEncodeSingle,
     PCLazyLoraLoader,
     PCLazyLoraLoaderAdvanced,
 ]

prompt_control/nodes_tools.py

@@ -1,3 +1,4 @@
+import json
 import logging
 
 from comfy_api.latest import io
@@ -5,6 +6,7 @@ from comfy_api.latest import io
 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")
@@ -182,6 +184,70 @@ class PCMacroExpand(io.ComfyNode):
         return io.NodeOutput(macroexpand(text))
 
 
+class PCLinkHelper(io.ComfyNode):
+    # a-z
+    NAMES = [chr(97 + i) for i in range(26)]
+
+    @classmethod
+    def define_schema(cls):
+        t1 = io.Autogrow.TemplateNames(io.AnyType.Input("link", raw_link=True), min=0, names=cls.NAMES)
+        t2 = io.Autogrow.TemplateNames(
+            io.AnyType.Input("value", lazy=True), min=0, names=[f"var{i + 1}" for i in range(50)]
+        )
+        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 -> $z with JSON link values.",
+            is_experimental=True,
+            inputs=[
+                io.Autogrow.Input("links", template=t1),
+                io.Autogrow.Input(
+                    "vars",
+                    template=t2,
+                ),
+                io.String.Input(
+                    "template",
+                    tooltip="The variables $a to $z will be replaced in this text with their corresponding input's JSON link value",
+                    placeholder="In this text you can refer to the input links as $a, $b etc. and the var inputs as either $var1 or $json1 etc. (the latter will be rendered through Python's json.dumps function which will cause strings to be quoted)",
+                    multiline=True,
+                ),
+            ],
+            outputs=[io.String.Output()],
+        )
+
+    # This requires https://github.com/Comfy-Org/ComfyUI/pull/15103 to work properly
+    # Without that PR, all inputs will be evaluated non-lazily
+    @classmethod
+    def check_lazy_status(cls, template, links, vars):
+        r = []
+        for name, (v, input_name) in vars.items():
+            if v is None and f"${name}" in template or v is None and f"$json{name[3:]}" in template:
+                r.append(input_name)
+        return r
+
+    @classmethod
+    def execute(cls, template, links, vars) -> io.NodeOutput:
+        text = template
+        for k in cls.NAMES:
+            v = "__EMPTY__"
+            if k in links:
+                # Replace : with \: to avoid breaking scheduling syntax when linking subgraphs. Any function that consumes this should replace \: with :
+                v = json.dumps(links[k]).replace(":", r"\:")
+            text = substitute_var(text, k, v)
+        for i in range(50):
+            v = "__EMPTY__"
+            k = f"var{i + 1}"
+            if k in vars:
+                v = vars[k]
+            text = substitute_var(text, k, str(v))
+            if f"$json{i + 1}" in text:
+                v = v if v == "__EMPTY__" else json.dumps(v)
+                text = substitute_var(text, f"json{i + 1}", v)
+
+        return io.NodeOutput(text)
+
+
 NODES = [
     PCSetPCTextEncodeSettings,
     PCAddMaskToCLIP,
@@ -189,4 +255,5 @@ NODES = [
     PCSetLogLevel,
     PCExtractScheduledPrompt,
     PCMacroExpand,
+    PCLinkHelper,
 ]

prompt_control/parser.py

@@ -1,10 +1,397 @@
-import logging
-import os
+from __future__ import annotations
 
-log = logging.getLogger("comfyui-prompt-control")
+import itertools as it
+from dataclasses import dataclass
+from math import ceil
+from typing import Any, TypeAlias
 
-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
+from typing_extensions import override
+
+from .macros import expand_macros
+from .parsy import any_char, char_from, digit, eof, forward_declaration, generate, regex, seq, string, success
+
+FOREVER = float("inf")
+
+EvalResult: TypeAlias = tuple[float, str, list["LoRA"]]
+
+
+def merge_until(i: EvalResult, minimum: float):
+    until, p, loras = i
+    until = min(until, minimum)
+    return until, p, loras
+
+
+def batched(iterable, n, *, strict=False):
+    # batched('ABCDEFG', 2) → AB CD EF G
+    if n < 1:
+        raise ValueError("n must be at least one")
+    iterator = iter(iterable)
+    while batch := tuple(it.islice(iterator, n)):
+        if strict and len(batch) != n:
+            raise ValueError("batched(): incomplete batch")
+        yield batch
+
+
+EvalResult: TypeAlias = tuple[float, str, list["LoRA"]]
+
+
+class Expression:
+    def eval(self, step: float, tags: list[str]) -> EvalResult:
+        return (FOREVER, "", [])
+
+    def required_steps(self, max_steps: float) -> set[float]:
+        return set()
+
+
+@dataclass
+class Text(Expression):
+    string: str
+
+    @override
+    def eval(self, step: float, tags: list[str]) -> EvalResult:
+        assert isinstance(self.string, str)
+        return FOREVER, self.string, []
+
+
+@dataclass
+class Alternate(Expression):
+    prompts: list[Expression]
+    step: float = 0.1
+
+    @override
+    def eval(self, step: float, tags: list[str]) -> EvalResult:
+        SCALE = 10_000
+        step = max(step, self.step)
+        position = (step * SCALE) / (self.step * SCALE)
+        idx = (ceil(position) - 1) % len(self.prompts)
+
+        r = self.prompts[max(0, idx)].eval(step, tags)
+        r = merge_until(r, max(self.step, ceil(position) * self.step))
+        return r
+
+    @override
+    def required_steps(self, max_steps: float):
+        r = set()
+        for x in self.prompts:
+            r.update(x.required_steps(max_steps))
+        r.update(set(x / 100 for x in range(0, int(max_steps * 100), int(self.step * 100))))
+        return r
+
+
+@dataclass
+class Sequence(Expression):
+    prompts: list[tuple[Expression, float]]
+
+    @override
+    def eval(self, step: float, tags: list[str]) -> EvalResult:
+        item = Text("")
+        found_step = FOREVER
+        for prompt, switch_step in self.prompts:
+            if step <= switch_step:
+                found_step = switch_step
+                item = prompt
+                break
+
+        return merge_until(item.eval(step, tags), found_step)
+
+    @override
+    def required_steps(self, max_steps: float):
+        return set(step for _, step in self.prompts if step <= max_steps)
+
+
+@dataclass
+class Schedule(Expression):
+    before: Prompt
+    during: Prompt
+    after: Prompt
+    start: float
+    end: float
+    tag: str | None
+
+    def tag_matches(self, tags: list[str]):
+        return self.tag in tags
+
+    @override
+    def eval(self, step: float, tags: list[str]) -> EvalResult:
+        if self.tag is not None and not self.tag_matches(tags):
+            return self.before.eval(step, tags)
+        if self.tag_matches(tags):
+            return self.during.eval(step, tags)
+
+        if step <= self.start:
+            return merge_until(self.before.eval(step, tags), self.start)
+        if self.start < step <= self.end:
+            return merge_until(self.during.eval(step, tags), self.end)
+        if step > self.end:
+            return self.after.eval(step, tags)
+        raise AssertionError("How are you here?")
+
+    @override
+    def required_steps(self, max_steps: float):
+        r = set()
+        if self.start < max_steps:
+            r.add(self.start)
+        if self.end < max_steps:
+            r.add(self.end)
+        r.update(self.before.required_steps(max_steps))
+        r.update(self.during.required_steps(max_steps))
+        r.update(self.after.required_steps(max_steps))
+        return r
+
+
+@dataclass
+class Prompt(Expression):
+    data: list[Expression]
+
+    @override
+    def eval(self, step: float, tags: list[str]) -> EvalResult:
+        evals = [x.eval(step, tags) for x in self.data]
+        text = "".join(x[1] for x in evals)
+        untils = [x[0] for x in evals]
+        loras = []
+        for x in evals:
+            loras.extend(x[2])
+        until = FOREVER if not untils else min(untils)
+        return until, text, loras
+
+    @override
+    def required_steps(self, max_steps):
+        r = set()
+        for x in self.data:
+            r.update(x.required_steps(max_steps))
+        return r
+
+
+@dataclass
+class LoRA(Expression):
+    filename: str
+    w_model: float = 1.0
+    w_te: float = 1.0
+
+    def eval(self, step: float, tags: list[str]) -> EvalResult:
+        return FOREVER, "", [self]
+
+
+def find_weight_at(weights: list[tuple[float, float]], step: float, until: float):
+    res_w = 0
+    for this, next in zip(weights, it.chain(weights[1:], [(0, FOREVER)]), strict=False):
+        w, start = this
+        _, next_start = next
+        if start > step or next_start < step:
+            until = min(until, start)
+            continue
+        res_w = w
+    return until, res_w
+
+
+@dataclass
+class LoRACTL(Expression):
+    filename: str
+    w_model: list[tuple[float, float]]
+    w_te: list[tuple[float, float]]
+
+    def eval(self, step: float, tags: list[str]) -> EvalResult:
+        until, w1 = find_weight_at(self.w_model, step, FOREVER)
+        until, w2 = find_weight_at(self.w_te, step, until)
+        lora = []
+        if w1 != 0 or w1 != 0:
+            lora = [LoRA(self.filename, w1, w2)]
+
+        return until, "", lora
+
+    def required_steps(self, max_steps):
+        r = set(x[1] for x in self.w_model)
+        r.update(set(x[1] for x in self.w_te))
+        return r
+
+
+def combine_arglist(prompts, start_end) -> Schedule:
+    a, b, c = prompts
+    start_or_tag, end = start_end
+    empty = Prompt([])
+    start = start_or_tag
+    # Handle [a:b:TAG]
+    if isinstance(start_or_tag, str):
+        if b is None:
+            before = empty
+            during = a  # [a:TAG] produces a when tag is active
+        else:
+            before, during = a, b  # [a:b:TAG] changes from a to b when tag is active
+        return Schedule(before, during, empty, start=0.0, end=FOREVER, tag=start_or_tag)
+    during = before = after = empty
+    if end is not None:
+        if b is None:  # [a:0,0.5] == [:a:0,0.5]
+            during = a
+            before = after = empty
+        elif c is None:  # [a:b:0,0.5]
+            before = empty
+            during = a
+            after = b
+        else:
+            before, during, after = a, b, c
+    else:
+        end = FOREVER
+        if b is None:  # [a:0.5] == [::a:0.5,0.5]
+            before = empty
+            during = a
+            after = a
+        else:
+            before = a
+            during = b
+            after = b
+            # c always gets ignored
+    start = float(start)  # for typechecking
+    return Schedule(before, during, after, start, end, tag=None)
+
+
+def token(s: str):
+    return string(s).map(Text)
+
+
+def combine_prompt(*prompts):
+    p = prompts
+    if len(p) == 1:
+        p = p[0]
+    if isinstance(p, Prompt):
+        p = p.data[0] if len(p.data) == 1 else combine_prompt(*p.data)
+    if isinstance(p, Expression):
+        return p
+    p = [combine_prompt(x) for x in p]
+    return Prompt(p)
+
+
+@dataclass
+class PromptSchedule:
+    parse_tree: Expression
+    filters: list[str]
+    start: float
+    end: float
+    num_steps: int
+
+    def at_step(self, step: float) -> tuple[float, dict[str, Any]]:
+        max_step = self.num_steps or 1.0
+        if max_step > 1 and step < 1:
+            step = step * max_step
+        until, p, lora_list = self.parse_tree.eval(step, self.filters)
+        loras = {}
+        for lora in lora_list:
+            d = loras.get(lora.filename, {})
+            d["weight"] = d.get("weight", 0) + lora.w_model
+            d["weight_clip"] = d.get("weight_clip", 0) + lora.w_te
+            loras[lora.filename] = d
+        if max_step > 0 and until > 1:
+            # TODO: better logic for this?
+            until = min(until / max_step, 1.0)
+        return (min(max_step, round(until, 2)), {"prompt": p, "loras": loras})
+
+    def with_filters(self, filters: str | None = None, start: float | None = None, end: float | None = None):
+        return PromptSchedule(
+            self.parse_tree,
+            self.filters if filters is None else parse_filters(filters),
+            self.start if start is None else start,
+            self.end if end is None else end,
+            self.num_steps,
+        )
+
+    def clone(self):
+        return self.with_filters()
+
+    def __iter__(self):
+        return (x for x in self.parsed_prompt if x[0] != 0)
+
+    @property
+    def parsed_prompt(self):
+        max_step = self.num_steps or 1.0
+        required_steps = self.parse_tree.required_steps(max_step).union({max_step})
+
+        prompts = list(sorted((self.at_step(step) for step in required_steps), key=lambda x: x[0]))
+        res = []
+        prev_end = -1
+        for end_at, p in prompts:
+            if end_at < self.start:
+                continue
+            elif end_at < self.end and prev_end < end_at:
+                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
+
+        if len(res) == 0:
+            res = [[1.0], prompts[-1][1]]
+
+        return res
+
+
+def lora_weights(p):
+    @generate
+    def parser():
+        w_model = yield col >> p
+        w_te = yield (col >> p).optional(w_model)
+        return [w_model, w_te]
+
+    return parser.desc("lora_weights")
+
+
+prompt = forward_declaration()
+empty = Text("")
+comma = token(",")
+col = token(":")
+lsq = token("[")
+rsq = token("]")
+lpar = token("(")
+rpar = token(")")
+tag = regex(r"[A-Z_]+")
+non_special = regex(r"[^:\[\]()|\\<>#]+").map(Text)
+filename = regex(r"[^:<>]+")
+
+comment = string("#") >> any_char.until(eof | char_from("\n")) >> success(empty)
+escape = (string("\\") >> char_from("\\[]:#") | string(r"\(") | string(r"\)")).map(Text)
+emphasis = seq(lpar, (prompt | col).at_least(0), rpar)
+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())
+arglist = seq((opt_prompt << col).optional() * 3, step_range)
+schedule = lsq >> arglist.combine(combine_arglist) << rsq
+
+alternate = (lsq >> seq(prompt.sep_by(string("|"), min=1), (col >> number).optional(0.1)) << rsq).combine(Alternate)
+sequence = (lsq >> string("SEQ") >> seq(col >> opt_prompt << col, number).at_least(1) << rsq).map(Sequence)
+bracketed = seq(lsq, prompt.at_least(0), rsq) | sequence | schedule | alternate
+lora = (string("<lora:") >> filename * 1 + lora_weights(number) << string(">")).combine(LoRA)
+ctlweight = seq(number, (string("@") >> number).optional(0)).sep_by(comma, min=1)
+loractl = (string("<loractl:") >> filename * 1 + lora_weights(ctlweight) << string(">")).combine(LoRACTL)
+emb = (string("<emb:") >> filename << string(">")).map(lambda f: Text(f"embedding:{f}"))
+
+expr = (
+    escape
+    | comment
+    | non_special
+    | bracketed
+    | emphasis.combine(combine_prompt)
+    | lora
+    | loractl
+    | emb
+    | char_from("<>").map(Text)
+)
+prompt_ = expr.at_least(1).combine(combine_prompt)
+prompt.become(prompt_)
+# Treat any character that isn't valid prompt syntax as just text
+all = (prompt | any_char.map(Text)).at_least(0).combine(combine_prompt)
+
+
+def parse_filters(filters: str):
+    return [x.strip().upper() for x in filters.split(",") if x.strip()]
+
+
+def parse(text):
+    return combine_prompt(all.parse(text))
+
+
+def parse_prompt_schedules(text, filters="", start=0, end=1.0, num_steps=0):
+    return PromptSchedule(parse(expand_macros(text.strip())), parse_filters(filters), start, end, num_steps)

prompt_control/parser_lark.py

@@ -1,359 +0,0 @@
-# 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

@@ -1,389 +0,0 @@
-from __future__ import annotations
-
-import itertools as it
-from dataclasses import dataclass
-from math import ceil
-from typing import Any, TypeAlias
-
-from typing_extensions import override
-
-from .macros import expand_macros
-from .parsy import any_char, char_from, digit, eof, forward_declaration, generate, regex, seq, string, success
-
-FOREVER = float("inf")
-
-EvalResult: TypeAlias = tuple[float, str, list["LoRA"]]
-
-
-def merge_until(i: EvalResult, minimum: float):
-    until, p, loras = i
-    until = min(until, minimum)
-    return until, p, loras
-
-
-def batched(iterable, n, *, strict=False):
-    # batched('ABCDEFG', 2) → AB CD EF G
-    if n < 1:
-        raise ValueError("n must be at least one")
-    iterator = iter(iterable)
-    while batch := tuple(it.islice(iterator, n)):
-        if strict and len(batch) != n:
-            raise ValueError("batched(): incomplete batch")
-        yield batch
-
-
-EvalResult: TypeAlias = tuple[float, str, list["LoRA"]]
-
-
-class Expression:
-    def eval(self, step: float, tags: list[str]) -> EvalResult:
-        return (FOREVER, "", [])
-
-    def required_steps(self, max_steps: float) -> set[float]:
-        return set()
-
-
-@dataclass
-class Text(Expression):
-    string: str
-
-    @override
-    def eval(self, step: float, tags: list[str]) -> EvalResult:
-        assert isinstance(self.string, str)
-        return FOREVER, self.string, []
-
-
-@dataclass
-class Alternate(Expression):
-    prompts: list[Expression]
-    step: float = 0.1
-
-    @override
-    def eval(self, step: float, tags: list[str]) -> EvalResult:
-        SCALE = 10_000
-        step = max(step, self.step)
-        position = (step * SCALE) / (self.step * SCALE)
-        idx = (ceil(position) - 1) % len(self.prompts)
-
-        r = self.prompts[max(0, idx)].eval(step, tags)
-        r = merge_until(r, max(self.step, ceil(position) * self.step))
-        return r
-
-    @override
-    def required_steps(self, max_steps: float):
-        r = set()
-        for x in self.prompts:
-            r.update(x.required_steps(max_steps))
-        r.update(set(x / 100 for x in range(0, int(max_steps * 100), int(self.step * 100))))
-        return r
-
-
-@dataclass
-class Sequence(Expression):
-    prompts: list[tuple[Expression, float]]
-
-    @override
-    def eval(self, step: float, tags: list[str]) -> EvalResult:
-        item = Text("")
-        found_step = FOREVER
-        for prompt, switch_step in self.prompts:
-            if step <= switch_step:
-                found_step = switch_step
-                item = prompt
-                break
-
-        return merge_until(item.eval(step, tags), found_step)
-
-    @override
-    def required_steps(self, max_steps: float):
-        return set(step for _, step in self.prompts if step <= max_steps)
-
-
-@dataclass
-class Schedule(Expression):
-    before: Prompt
-    during: Prompt
-    after: Prompt
-    start: float
-    end: float
-    tag: str | None
-
-    def tag_matches(self, tags: list[str]):
-        return self.tag in tags
-
-    @override
-    def eval(self, step: float, tags: list[str]) -> EvalResult:
-        if self.tag is not None and not self.tag_matches(tags):
-            return self.before.eval(step, tags)
-        if self.tag_matches(tags):
-            return self.during.eval(step, tags)
-
-        if step <= self.start:
-            return merge_until(self.before.eval(step, tags), self.start)
-        if self.start < step <= self.end:
-            return merge_until(self.during.eval(step, tags), self.end)
-        if step > self.end:
-            return self.after.eval(step, tags)
-        raise AssertionError("How are you here?")
-
-    @override
-    def required_steps(self, max_steps: float):
-        r = set()
-        if self.tag is not None:
-            return r
-        if self.start < max_steps:
-            r.add(self.start)
-        if self.end < max_steps:
-            r.add(self.end)
-        r.update(self.before.required_steps(max_steps))
-        r.update(self.during.required_steps(max_steps))
-        r.update(self.after.required_steps(max_steps))
-        return r
-
-
-@dataclass
-class Prompt(Expression):
-    data: list[Expression]
-
-    @override
-    def eval(self, step: float, tags: list[str]) -> EvalResult:
-        evals = [x.eval(step, tags) for x in self.data]
-        text = "".join(x[1] for x in evals)
-        untils = [x[0] for x in evals]
-        loras = []
-        for x in evals:
-            loras.extend(x[2])
-        until = FOREVER if not untils else min(untils)
-        return until, text, loras
-
-    @override
-    def required_steps(self, max_steps):
-        r = set()
-        for x in self.data:
-            r.update(x.required_steps(max_steps))
-        return r
-
-
-@dataclass
-class LoRA(Expression):
-    filename: str
-    w_model: float = 1.0
-    w_te: float = 1.0
-
-    def eval(self, step: float, tags: list[str]) -> EvalResult:
-        return FOREVER, "", [self]
-
-
-def find_weight_at(weights: list[tuple[float, float]], step: float, until: float):
-    res_w = 0
-    for this, next in zip(weights, it.chain(weights[1:], [(0, FOREVER)]), strict=False):
-        w, start = this
-        _, next_start = next
-        if start > step or next_start < step:
-            until = min(until, start)
-            continue
-        res_w = w
-    return until, res_w
-
-
-@dataclass
-class LoRACTL(Expression):
-    filename: str
-    w_model: list[tuple[float, float]]
-    w_te: list[tuple[float, float]]
-
-    def eval(self, step: float, tags: list[str]) -> EvalResult:
-        until, w1 = find_weight_at(self.w_model, step, FOREVER)
-        until, w2 = find_weight_at(self.w_te, step, until)
-        lora = []
-        if w1 != 0 or w1 != 0:
-            lora = [LoRA(self.filename, w1, w2)]
-
-        return until, "", lora
-
-    def required_steps(self, max_steps):
-        r = set(x[1] for x in self.w_model)
-        r.update(set(x[1] for x in self.w_te))
-        return r
-
-
-def combine_arglist(prompts, start_end) -> Schedule:
-    a, b, c = prompts
-    start_or_tag, end = start_end
-    empty = Prompt([])
-    start = start_or_tag
-    # Handle [a:b:TAG]
-    if isinstance(start_or_tag, str):
-        if b is None:
-            before = empty
-            during = a  # [a:TAG] produces a when tag is active
-        else:
-            before, during = a, b  # [a:b:TAG] changes from a to b when tag is active
-        return Schedule(before, during, empty, start=0.0, end=FOREVER, tag=start_or_tag)
-    during = before = after = empty
-    if end is not None:
-        if b is None:  # [a:0,0.5] == [:a:0,0.5]
-            during = a
-            before = after = empty
-        elif c is None:  # [a:b:0,0.5]
-            before = empty
-            during = a
-            after = b
-        else:
-            before, during, after = a, b, c
-    else:
-        end = FOREVER
-        if b is None:  # [a:0.5] == [::a:0.5,0.5]
-            before = empty
-            during = a
-            after = a
-        else:
-            before = a
-            during = b
-            after = b
-            # c always gets ignored
-    start = float(start)  # for typechecking
-    return Schedule(before, during, after, start, end, tag=None)
-
-
-def token(s: str):
-    return string(s).map(Text)
-
-
-def combine_prompt(*prompts):
-    p = prompts
-    if len(p) == 1:
-        p = p[0]
-    if isinstance(p, Prompt):
-        p = p.data[0] if len(p.data) == 1 else combine_prompt(*p.data)
-    if isinstance(p, Expression):
-        return p
-    p = [combine_prompt(x) for x in p]
-    return Prompt(p)
-
-
-@dataclass
-class PromptSchedule:
-    parse_tree: Expression
-    filters: list[str]
-    start: float
-    end: float
-    num_steps: int
-
-    def at_step(self, step: float) -> tuple[float, dict[str, Any]]:
-        max_step = self.num_steps or 1.0
-        if max_step > 1 and step < 1:
-            step = step * max_step
-        until, p, lora_list = self.parse_tree.eval(step, self.filters)
-        loras = {}
-        for lora in lora_list:
-            d = loras.get(lora.filename, {})
-            d["weight"] = d.get("weight", 0) + lora.w_model
-            d["weight_clip"] = d.get("weight_clip", 0) + lora.w_te
-            loras[lora.filename] = d
-        if max_step > 0 and until > 1:
-            # TODO: better logic for this?
-            until = min(until / max_step, 1.0)
-        return (min(max_step, round(until, 2)), {"prompt": p, "loras": loras})
-
-    def with_filters(self, filters: str | None = None, start: float | None = None, end: float | None = None):
-        return PromptSchedule(
-            self.parse_tree,
-            self.filters if filters is None else parse_filters(filters),
-            self.start if start is None else start,
-            self.end if end is None else end,
-            self.num_steps,
-        )
-
-    def clone(self):
-        return self.with_filters()
-
-    def __iter__(self):
-        return (x for x in self.parsed_prompt if x[0] != 0)
-
-    @property
-    def parsed_prompt(self):
-        max_step = self.num_steps or 1.0
-        required_steps = self.parse_tree.required_steps(max_step).union({max_step})
-
-        prompts = list(sorted((self.at_step(step) for step in required_steps), key=lambda x: x[0]))
-        res = []
-        prev_end = -1
-        for end_at, p in prompts:
-            if end_at < self.start:
-                continue
-            elif end_at < self.end and prev_end < end_at:
-                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
-
-        if len(res) == 0:
-            res = [[1.0], prompts[-1][1]]
-
-        return res
-
-
-def lora_weights(p):
-    @generate
-    def parser():
-        w_model = yield col >> p
-        w_te = yield (col >> p).optional(w_model)
-        return [w_model, w_te]
-
-    return parser.desc("lora_weights")
-
-
-prompt = forward_declaration()
-empty = Text("")
-comma = token(",")
-col = token(":")
-lsq = token("[")
-rsq = token("]")
-lpar = token("(")
-rpar = token(")")
-tag = regex(r"[A-Z_]+")
-non_special = regex(r"[^:\[\]()|\\<>#]+").map(Text)
-filename = regex(r"[^:<>]+")
-
-comment = string("#") >> any_char.until(eof | char_from("\n")) >> success(empty)
-escape = (string("\\") >> char_from("\\[]:#") | string(r"\(") | string(r"\)")).map(Text)
-emphasis = seq(lpar, (prompt | col).at_least(0), rpar)
-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())
-arglist = seq((opt_prompt << col).optional() * 3, step_range)
-schedule = lsq >> arglist.combine(combine_arglist) << rsq
-
-alternate = (lsq >> seq(prompt.sep_by(string("|"), min=1), (col >> number).optional(0.1)) << rsq).combine(Alternate)
-sequence = (lsq >> string("SEQ") >> seq(col >> opt_prompt << col, number).at_least(1) << rsq).map(Sequence)
-bracketed = seq(lsq, prompt.at_least(0), rsq) | sequence | schedule | alternate
-lora = (string("<lora:") >> filename * 1 + lora_weights(number) << string(">")).combine(LoRA)
-ctlweight = seq(number, (string("@") >> number).optional(0)).sep_by(comma, min=1)
-loractl = (string("<loractl:") >> filename * 1 + lora_weights(ctlweight) << string(">")).combine(LoRACTL)
-emb = (string("<emb:") >> filename << string(">")).map(lambda f: Text(f"embedding:{f}"))
-
-expr = escape | comment | non_special | bracketed | emphasis.combine(combine_prompt) | lora | loractl | emb
-prompt_ = expr.at_least(1).combine(combine_prompt)
-prompt.become(prompt_)
-# Treat any character that isn't valid prompt syntax as just text
-all = (prompt | any_char.map(Text)).at_least(0).combine(combine_prompt)
-
-
-def parse_filters(filters: str):
-    return [x.strip().upper() for x in filters.split(",") if x.strip()]
-
-
-def parse(text):
-    return combine_prompt(all.parse(text))
-
-
-def parse_prompt_schedules(text, filters="", start=0, end=1.0, num_steps=0):
-    return PromptSchedule(parse(expand_macros(text.strip())), parse_filters(filters), start, end, num_steps)

prompt_control/utils.py

@@ -3,10 +3,9 @@ from __future__ import annotations
 import copy
 import logging
 import re
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from dataclasses import dataclass
 from pathlib import Path
-from timeit import default_timer as timer
 from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar
 
 if TYPE_CHECKING:
@@ -21,7 +20,6 @@ class FunctionSpec:
     name: str
     args: FunctionArgs
     position: int
-    placeholder: str | None
 
 
 # Allow testing
@@ -35,22 +33,6 @@ except ImportError:
 
 log = logging.getLogger("comfyui-prompt-control")
 
-TIME_CONTEXT = []
-
-
-def push_timer(name):
-    TIME_CONTEXT.append((name, timer()))
-    name = ":".join(n[0] for n in TIME_CONTEXT)
-    log.info("START: %s", name)
-
-
-def pop_timer():
-    global TIME_CONTEXT
-    name = ":".join(n[0] for n in TIME_CONTEXT)
-    _, start = TIME_CONTEXT.pop()
-    end = timer()
-    log.info("%s: took %s seconds", name, round(end - start, 3))
-
 
 def flatten(x):
     if type(x) in [str, tuple, int, type(None)] or isinstance(x, dict) and "type" in x:
@@ -159,6 +141,10 @@ def find_function_spans(
         if text[at_paren:after_first_paren] == "(":
             end = find_closing_paren(text, after_first_paren)
             if end < 0:
+                # Unclosed paren: skip past this match so the loop terminates
+                idx += match.end()
+                text = text[match.end() :]
+                match = rex.search(text)
                 continue
             args = parse_strings(text[after_first_paren:end], defaults)
             end += 1
@@ -172,7 +158,11 @@ def find_function_spans(
 
 
 def get_function(
-    text: str, func: str, defaults: list[str] | None, placeholder: str = "", require_args: bool = True
+    text: str,
+    func: str,
+    defaults: list[str] | None,
+    processor: Callable[..., str] | None = None,
+    require_args: bool = True,
 ) -> tuple[str, list[FunctionSpec]]:
     spans = [x.span() for x in re.finditer(r'".+?"', text)]
     instances = []
@@ -181,14 +171,13 @@ def get_function(
     current = 0
     skipped = 0
     for start, end, funcname, args in find_function_spans(text, func, require_args, defaults):
-        ph = None
         if spans_include(spans, start, end):
             continue
-        if placeholder:
-            ph = f"\0{placeholder}{count}\0"
-        instances.append(FunctionSpec(funcname, args, start - skipped, ph))
+        instances.append(FunctionSpec(funcname, args, start - skipped))
         skipped += end - start
-        chunks.append(text[current:start] + (ph or ""))
+        chunks.append(text[current:start])
+        if processor:
+            chunks.append(processor(*args))
         current = end
         count += 1
     chunks.append(text[current:])
@@ -290,6 +279,10 @@ def lora_name_to_file(name: str) -> str | None:
     search = [f for f in filenames if all(p in f for p in parts)]
     if len(search) == 1:
         return search[0]
+    elif len(search) > 1:
+        if len(search) > 4:
+            search[4] = "..."
+        log.warning("Ignored LoRA search 's%'; matched more than one file: %s", name, ", ".join(search[:5]))
 
     return None
 

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.5"
+version = "3.0.0-beta.10"
 license = { file = "LICENSE" }
 
 requires-python = ">= 3.10"

tests/test_graph.py

@@ -460,6 +460,102 @@ def test_textencode_lora_with_schedule():
     }
 
 
+def test_textencode_custom():
+    r = te("NODE(CLIPTextEncode)simple [test:0.1,0.5] $p SEG(p) 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\: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: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: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: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 +678,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

@@ -30,6 +30,11 @@ def test_macro_recursion():
     assert "Unable to resolve DEFs" in str(c.value)
 
 
+def test_parsing_cornercase():
+    r = expand_macros("This should not get stuck DEF(")
+    assert r == "This should not get stuck DEF("
+
+
 @pytest.mark.parametrize(
     "input, output",
     [

tests/test_parser.py

@@ -1,7 +1,7 @@
-import os
-
 import pytest
 
+from prompt_control.parser import parse_prompt_schedules as parse
+
 
 def lora_dict(*loras):
     return {lora: {"weight": unet, "weight_clip": te} for lora, unet, te in loras}
@@ -19,18 +19,9 @@ def assert_prompt(p, at, until, text, *loras):
     assert prompts_match(p.at_step(at), prompt(until, text, *loras))
 
 
-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)
+params.append(parse)
 
 
 @pytest.fixture(scope="module", autouse=True, params=params)
@@ -342,6 +333,12 @@ def test_cornercase_corrected(parse):
     assert p.parsed_prompt[1:] == p2.parsed_prompt
 
 
+def test_ltgt_in_schedule(parse):
+    p = parse("This should [<parse> correctly:be <Picture 1>:0.1]<lora:test:1>")
+    assert_prompt(p, 0.1, 0.1, "This should <parse> correctly", ("test", 1.0, 1.0))
+    assert_prompt(p, 0.15, 1.0, "This should be <Picture 1>", ("test", 1.0, 1.0))
+
+
 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>")
@@ -364,3 +361,9 @@ def test_alternating_nested(parse):
     for i, x in enumerate(catdogtigers):
         step = round((i * 0.1) + 0.1, 2)
         assert_prompt(p3, step, step, x)
+
+
+def test_alternating_with_tags(parse):
+    p1 = parse("[[a|b]:HR]", filters="HR")
+    p2 = parse("[a|b]")
+    assert p1.parsed_prompt == p2.parsed_prompt

tests/test_utils.py

@@ -0,0 +1,7 @@
+from prompt_control import utils
+
+
+def test_smart_split():
+    assert utils.smarter_split(",", "foo,bar") == ["foo", "bar"]
+    assert utils.smarter_split(",", "(foo,bar),zonk") == ["(foo,bar)", "zonk"]
+    assert utils.smarter_split(",", r"\(foo,bar),zonk") == [r"\(foo", "bar)", "zonk"]