v3.0.0-beta.5

This one might go away or change

.github/workflows/tests.yml

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

.github/workflows/tests_with_comfy.yml

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

Makefile

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

README.md

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

__init__.py

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

doc/attention_couple.md

@@ -1,7 +1,5 @@
 # Attention Couple
 
-NOTE: This is still considered an experimental feature, so the syntax may change.
-
 Attention Couple is an attention-based implementation of regional prompting. it is faster and often more flexible than latent-based masking.
 
 The implementation is based on the one by [pamparamm](https://github.com/pamparamm/ComfyUI-ppm.git), modified to use ComfyUI's hook system. This enables it to work with prompt scheduling.
@@ -12,6 +10,11 @@ As a consequence of this, however, you can also use `COUPLE` in your negative pr
 
 To enable batching negative prompts, run your positive and negative prompt through the `PPCAttentionCoupleBatchNegative` node. This will make the outputs identical to pamparamm's implementation and will also improve performance. It will fall back to the default behaviour in cases where batching can't be done, so it should always be safe to use.
 
+## Anima
+
+There is a **very experimental** port of pamparamm's Anima support for Attention Couple in Prompt Control. Because ComfyUI lacks the built-in schedulable hooks required, you must first patch your model with `PC: Anima Attention Couple Model Patch` in addition to using `COUPLE` as usual.
+
+The code was hacked together with minimal thought, so expect bugs and misbehaviour. The port is also currently *not* compatible with NegPIP.
 
 ## Syntax
 

doc/basic.md

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

doc/macros.md

@@ -58,3 +58,40 @@ a "$1" b "$2"
 a "" b "$2"
 a "A" b "$2"
 ```
+
+## SEG: Split your prompt into named segments
+
+Syntax: `SEG(segment_name)`
+
+To help with organizing prompts, you can use the `SEG` function. For example:
+```
+This is a comic
+Top panel: $CAT. $SEG3
+Bottom panel: $DOG
+
+SEG(DOG)
+A dog chasing its
+  tail in a living room.
+SEG(CAT)
+
+a sleeping cat
+
+SEG
+The cat has orange fur with white stripes
+```
+This produces:
+```
+This is a comic
+Top panel: a sleeping cat. The cat has orange fur with white stripes
+Bottom panel: A dog chasing its
+  tail in a living room.
+```
+Unlike macros, SEG is processed *after* scheduling syntax.
+
+In this case, the first section before any `SEG` becomes the *template* and any text after a `SEG` call becomes part of that segment. Whitespace is stripped from the start and end of segments and the template.
+
+In the template, you can refer to segments by either their index (starting from 1) or the given name, prefixed with a `$SEG`, so in this example, `$SEG1` is the same as `$PANEL1`
+
+Segments can also refer to each other. Recursion will terminate, but produces weird outputs.
+
+Naming segments is optional, in which case you will have to refer to it by its index.

prompt_control/anima_couple.py

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

prompt_control/attention_couple_ppm.py

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

prompt_control/macros.py

@@ -4,11 +4,52 @@ from __future__ import annotations
 import logging
 import re
 
-from .utils import find_closing_paren, get_function
+from .utils import find_closing_paren, get_function, split_by_function
 
 log = logging.getLogger("comfyui-prompt-control")
 
 
+def substitute_template(template, segments, do_subs):
+    def _substitute(template, segments, stack):
+        name = ""
+        if "$" in template:
+            for name, value in sorted(segments):
+                value = substitute_var(value, name, "")
+                if name not in stack:
+                    stack.add(name)
+                    value = _substitute(value, segments, stack)
+                    stack.remove(name)
+                template = substitute_var(template, name, value)
+        if do_subs and name not in stack:
+            template = expand_subs(template)
+        return template
+
+    return _substitute(template, segments, set())
+
+
+def expand_segs(text, do_subs=True):
+    template, segments = split_by_function(text, "SEG", defaults=[""], require_args=True)
+    named_segs = [(f.args[0].strip() or f"SEG{i + 1}", c.strip()) for i, (c, f) in enumerate(segments)]
+
+    new_text = substitute_template(template, named_segs, do_subs).strip()
+    if new_text != text.strip():
+        log.debug("Template expanded to: %s", new_text)
+    return new_text
+
+
+def expand_subs(text):
+    text, subs = get_function(text, "SUB", defaults=None)
+    subs = [spec.strip() for f in subs for spec in f.args[0].split(";")]
+    for spec in subs:
+        if len(spec) <= 3 or spec[0] != "s":
+            log.warning("Invalid SUB spec ignored: '%s'", spec)
+            continue
+        splitchar = spec[1]
+        search, replace, *_ = spec[2:].split(splitchar)
+        text = re.sub(search, replace, text)
+    return text
+
+
 def parse_search(search):
     arg_start = search.find("(")
     args = ""
@@ -59,6 +100,13 @@ def expand_macros(text):
     return res
 
 
+def substitute_var(text, name, replace, boundary=r"\b"):
+    if f"${name}" not in text:
+        return text
+    name = re.escape(str(name))
+    return re.sub(rf"\${name}{boundary}", replace, text)
+
+
 def substitute_defcall(text, search, replace):
     name, default_args = search
     text, defns = get_function(text, name, defaults=None, placeholder=f"DEFNCALL{name}", require_args=False)
@@ -70,11 +118,12 @@ def substitute_defcall(text, search, replace):
         if parameters:
             paramvals = [x.strip() for x in parameters[0].split(";")]
         r = replace
+        end_re = r"(?![0-9])"
         for i, v in enumerate(paramvals):
-            r = re.sub(rf"\${i + 1}\b", v, r)
+            r = substitute_var(r, i + 1, v, boundary=end_re)
 
         for i, v in enumerate(default_args):
-            r = re.sub(rf"\${i + 1}\b", v, r)
+            r = substitute_var(r, i + 1, v, boundary=end_re)
 
         text = text.replace(ph, r)
     return text

prompt_control/nodes_anima.py

@@ -0,0 +1,71 @@
+# Adapted from ComfyUI-ppm into hook form
+
+
+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
+
+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):
+        self.block.to("cuda")
+        return cosmos_attention_forward_couple(self.fn, *args, **kwargs)
+
+
+class PCAnimaAttnCouplePatch(io.ComfyNode):
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        return io.Schema(
+            node_id="PCAnimaAttnCouplePatch",
+            display_name="PC: Anima attention Couple Model Patch",
+            category="promptcontrol/experimental",
+            inputs=[
+                io.Model.Input("model"),
+            ],
+            outputs=[
+                io.Model.Output(),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model: ModelPatcher) -> io.NodeOutput:
+        model_type = type(model.model)
+        m = model
+
+        if issubclass(model_type, Anima):
+            m = model.clone()
+            anima_model = model.get_model_object("diffusion_model")
+            m.add_wrapper_with_key(
+                comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL,
+                cls.__name__,
+                anima_forward_wrapper,
+            )
+            m.add_wrapper_with_key(
+                comfy.patcher_extension.WrappersMP.SAMPLER_SAMPLE,
+                cls.__name__,
+                anima_sample_wrapper,
+            )
+
+            for block_name, b in (
+                (n, b) for n, b in anima_model.named_modules() if "cross_attn" in n and isinstance(b, CosmosAttention)
+            ):
+                attn_forward_prev = m.get_model_object(f"diffusion_model.{block_name}.forward")
+                m.add_object_patch(f"diffusion_model.{block_name}.forward", CoupleForward(attn_forward_prev, b))
+
+        return io.NodeOutput(m)
+
+
+NODES = [PCAnimaAttnCouplePatch]

prompt_control/nodes_base.py

@@ -2,6 +2,7 @@ import logging
 
 from comfy_api.latest import io
 
+from .macros import expand_segs
 from .prompts import encode_prompt
 
 log = logging.getLogger("comfyui-prompt-control")
@@ -25,10 +26,11 @@ class PCTextEncodeWithRange(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, clip, text, start=0.0, end=1.0) -> io.NodeOutput:  # ty: ignore[invalid-method-override]
+    def execute(cls, clip, text, start=0.0, end=1.0) -> io.NodeOutput:
         log.debug("PCTextEncode: Encoding '%s'", text)
         defaults = clip.patcher.model_options.get("x-promptcontrol.defaults", {})
         masks = clip.patcher.model_options.get("x-promptcontrol.masks", None)
+        text = expand_segs(text)
         out = encode_prompt(clip, text, start, end, defaults, masks)
         return io.NodeOutput(out)
 
@@ -49,7 +51,7 @@ class PCTextEncode(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, clip, text) -> io.NodeOutput:  # ty: ignore[invalid-method-override]
+    def execute(cls, clip, text) -> io.NodeOutput:
         # Use the WithRange node for the range 0.0, 1.0
         return PCTextEncodeWithRange.execute(clip, text, 0.0, 1.0)
 

prompt_control/nodes_hooks.py

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

prompt_control/nodes_tools.py

@@ -2,7 +2,9 @@ import logging
 
 from comfy_api.latest import io
 
-from .macros import expand_macros
+from .macros import expand_macros as macroexpand
+from .macros import expand_segs as segexpand
+from .macros import expand_subs as subexpand
 from .parser import parse_prompt_schedules
 
 log = logging.getLogger("comfyui-prompt-control")
@@ -132,22 +134,32 @@ class PCExtractScheduledPrompt(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="PCExtractScheduledPrompt",
-            display_name="PC: Extract Scheduled Prompt",
+            display_name="PC: Show Prompt",
             category="promptcontrol/tools",
             description="Parses the input prompt and returns the prompt scheduled at the specified point",
             inputs=[
                 io.String.Input("text", multiline=True),
                 io.Float.Input("at", min=0.0, max=1.0, default=1.0, step=0.01),
                 io.String.Input("tags", default="", optional=True),
+                io.Boolean.Input("expand_segs", default=False, optional=True),
+                io.Boolean.Input("expand_subs", default=False, optional=True),
+                io.Boolean.Input("expand_macros", default=False, optional=True),
             ],
             outputs=[io.String.Output()],
+            search_aliases=["extract scheduled prompt"],
         )
 
     @classmethod
-    def execute(cls, text, at, tags="") -> io.NodeOutput:
+    def execute(cls, text, at, tags="", expand_segs=False, expand_subs=False, expand_macros=False) -> io.NodeOutput:
+        if expand_macros:
+            text = macroexpand(text)
         schedule = parse_prompt_schedules(text, filters=tags)
-        _, entry = schedule.at_step(at, total_steps=1)
+        _, entry = schedule.at_step(at)
         prompt_text = entry.get("prompt", "")
+        if expand_segs:
+            prompt_text = segexpand(prompt_text, do_subs=expand_subs)
+        if expand_subs:
+            prompt_text = subexpand(prompt_text)
         return io.NodeOutput(prompt_text)
 
 
@@ -167,7 +179,7 @@ class PCMacroExpand(io.ComfyNode):
 
     @classmethod
     def execute(cls, text) -> io.NodeOutput:
-        return io.NodeOutput(expand_macros(text))
+        return io.NodeOutput(macroexpand(text))
 
 
 NODES = [

prompt_control/parser.py

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

prompt_control/parser_parsy.py

@@ -352,7 +352,7 @@ escape = (string("\\") >> char_from("\\[]:#") | string(r"\(") | string(r"\)")).m
 emphasis = seq(lpar, (prompt | col).at_least(0), rpar)
 sign = string("+") | string("-")
 number = (
-    (sign.optional("") + (digit.at_least(1) + string(".") * 1 + digit.many() | digit.at_least(1)).concat())
+    (sign.optional("") + (digit.many() + string(".") * 1 + digit.many() | digit.at_least(1)).concat())
     .concat()
     .map(float)
 )

prompt_control/prompts.py

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

pyproject.toml

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

requirements.txt

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

tests/test_encode.py

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

tests/test_macros.py

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

tests/test_parser.py

@@ -2,8 +2,6 @@ import os
 
 import pytest
 
-from prompt_control.macros import expand_macros
-
 
 def lora_dict(*loras):
     return {lora: {"weight": unet, "weight_clip": te} for lora, unet, te in loras}
@@ -223,23 +221,6 @@ def test_def(parse):
     p2 = parse("[(test):(test:0.7):0.7] [(test):(test:0.5):0.5]")
     assert p.parsed_prompt == p2.parsed_prompt
 
-    p = expand_macros("DEF(X(a;b)=$1 $2 $3 d)X(A) X(A;B;C)")
-    assert p == "A b $3 d A B C d"
-
-    p = expand_macros("DEF(MACRO()=[empty:$1:$2])MACRO MACRO(;) MACRO(;0.5) MACRO(a;0.5)")
-    assert p == "[empty::$2] [empty::] [empty::0.5] [empty:a:0.5]"
-
-    p = expand_macros("DEF(X=$1)DEF(Y()=$1)[X Y][X() Y()][X(1) Y(1)]")
-    assert p == "[$1 ][ ][1 1]"
-
-    p = parse("DEF(test(1)=prompt $1)DEF(test2((a); (test))=[$1:$2:0.5])test test2")
-    p2 = parse("prompt 1 [(a):(prompt 1):0.5]")
-    assert p.parsed_prompt == p2.parsed_prompt
-
-    with pytest.raises(ValueError) as c:
-        expand_macros("DEF(X=recurse Y) DEF(Y=recurse X) X")
-    assert "Unable to resolve DEFs" in str(c.value)
-
 
 @pytest.mark.parametrize(
     "text, cases",
@@ -361,6 +342,12 @@ def test_cornercase_corrected(parse):
     assert p.parsed_prompt[1:] == p2.parsed_prompt
 
 
+def test_floats(parse):
+    p = parse("[a:b:0.5] [c:d:e:0.2,0.7] <lora:test:-0.3>")
+    p2 = parse("[a:b:.5] [c:d:e:.2,.7] <lora:test:-.3>")
+    assert p.parsed_prompt == p2.parsed_prompt
+
+
 def test_alternating_lora(parse):
     p4 = parse("[cat|[dog:wolf<lora:canine:1>:0.5]:0.2]")
     for i, (text, *_loras) in enumerate(

tests/test_workflow.json

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

tests/test_workflow.py

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