v3.0.0-beta.6

This avoids using stale block references and seems to fix LoRAs

__init__.py

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

doc/attention_couple.md

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

doc/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.
 

prompt_control/anima_couple.py

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

prompt_control/attention_couple_ppm.py

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

prompt_control/macros.py

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

prompt_control/nodes_anima.py

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

prompt_control/nodes_base.py

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

prompt_control/nodes_hooks.py

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

prompt_control/nodes_lazy.py

@@ -8,6 +8,7 @@ from comfy_api.latest import io
 from comfy_execution.graph import ExecutionBlocker
 from comfy_execution.graph_utils import GraphBuilder
 
+from .macros import expand_macros, expand_segs
 from .parser import parse_prompt_schedules
 from .utils import consolidate_schedule, find_nonscheduled_loras, get_function
 
@@ -158,6 +159,7 @@ class PCLazyLoraLoaderAdvanced(io.ComfyNode):
 
     @classmethod
     def execute(cls, model=None, clip=None, text="", apply_hooks=True, tags="", start=0.0, end=1.0, num_steps=0):
+        text = expand_segs(expand_macros(text))
         schedule = parse_prompt_schedules(text, filters=tags, start=start, end=end, num_steps=num_steps)
         graph = GraphBuilder()
         r = build_lora_schedule(graph, schedule, model, clip, apply_hooks=apply_hooks)

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

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

pyproject.toml

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

tests/test_encode.py

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

tests/test_macros.py

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