Now things should actually work

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 $(ARGS)
 
 test_graph:
 	PYTHONPATH=../../  pytest tests/test_graph.py $(ARGS)

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/nodes_lazy.py

@@ -3,20 +3,15 @@ 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 .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):

prompt_control/nodes_tools.py

@@ -2,7 +2,8 @@ import logging
 
 from comfy_api.latest import io
 
-from .parser import expand_macros, parse_prompt_schedules
+from .macros import expand_macros
+from .parser import parse_prompt_schedules
 
 log = logging.getLogger("comfyui-prompt-control")
 

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") != "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  # noqa
+else:
+    from .parser_lark import parse_prompt_schedules  # noqa

prompt_control/parser_lark.py

@@ -0,0 +1,360 @@
+# 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
+
+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 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/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"
+version = "3.0.0-beta.2"
 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"
 

requirements.txt

@@ -0,0 +1 @@
+# Nothing for now

tests/test_parser.py

@@ -2,9 +2,7 @@ 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
+from prompt_control.macros import expand_macros
 
 
 def lora_dict(*loras):
@@ -27,9 +25,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)