v2.0.0-rc.3

Broken by moving the LCM calculation outside the loop

See #108

.github/workflows/publish.yml

@@ -7,15 +7,23 @@ on:
     paths:
       - "pyproject.toml"
 
+permissions:
+  issues: write
+
 jobs:
+  tests:
+    uses: ./.github/workflows/tests.yml
+  tests_with_comfy:
+    uses: ./.github/workflows/tests_with_comfy.yml
   publish-node:
     name: Publish Custom Node to registry
     runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'asagi4' }}
+    needs: [tests, tests_with_comfy]
     steps:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
+        uses: Comfy-Org/publish-node-action@v1
         with:
-          ## Add your own personal access token to your Github Repository secrets and reference it here.
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.github/workflows/tests.yml

@@ -0,0 +1,18 @@
+name: Run parser tests
+on:
+  - workflow_call
+  - workflow_dispatch
+  - push
+
+jobs:
+  run-parser-tests:
+    name: Run parser tests
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - run: pip install -r requirements.txt
+      - run: python -m prompt_control.test_parser

.github/workflows/tests_with_comfy.yml

@@ -0,0 +1,28 @@
+name: Run tests requiring ComfyUI
+on:
+  workflow_call:
+  workflow_dispatch:
+  push:
+    paths:
+      - prompt_control/nodes_lazy.py
+      - prompt_control/utils.py
+
+
+jobs:
+  run-graph-tests:
+    name: Run graph tests
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Check out ComfyUI
+        uses: actions/checkout@v4
+        with:
+          repository: comfyanonymous/ComfyUI
+          path: ComfyUI
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+      - run: pip install -r requirements.txt -r ComfyUI/requirements.txt
+      - run: PYTHONPATH=ComfyUI python -m prompt_control.test_graph

Makefile

@@ -1,8 +1,14 @@
-all: format check
+all: format check test
 	@echo "Done"
 check:
 	find . -name "*.py" | xargs pyflakes
 format:
 	find . -name "*.py" | xargs black -l 120
 
+test:
+	python -m prompt_control.test_parser
+
+test_graph:
+	PYTHONPATH=../../ python -m prompt_control.test_graph
+
 .PHONY: check format all

README.md

@@ -2,51 +2,51 @@
 
 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.
 
+Prompt Control comes with `PCTextEncode`, which provides advanced text encoding with many additional features compared to ComfyUI's base `CLIPTextEncode`.
+
+A `Basic Text to Image` template is included with the extension, and can be loaded from ComfyUI's template library.
+
+## What can it do?
+
+You can use text prompts to control the following:
+
+- Prompt scheduling and filtering without noodle soup.
+- LoRA loading and scheduling via ComfyUI's hook system
+- Masking, composition and area control (regional prompting) with an implementation of Attention Couple, also fully schedulable.
+- Per-encoder prompts for models with multiple text encoders, such as SDXL and Flux
+- Prompt operations like `BREAK` and `AND`
+- Different weight interpretation types (ComfyUI, A1111, compel, etc.)
+- Prompt masking with an implementation of [cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff)
+- Simple prompt macros with `DEF`
+- And a bunch more
+
+All features are fully schedulable unless otherwise stated. See the [syntax documentation](doc/syntax.md) for details on how to use each feature.
+
+If you find prompt scheduling inconvenient for some reason, `PCTextEncode` can be used as a drop-in replacement for `CLIPTextEncode` to get everything else.
+
+[This workflow](example_workflows/Workflow%20Comparison.json?raw=1) shows LoRA scheduling and prompt editing and compares it with the same prompt implemented with built-in ComfyUI nodes. You can also find it in the template library.
+
+## Compatibility
+
+Prompt Control uses graph generation,  and tries to delegate functionality to core ComfyUI wherever possible, implementing any hooks and patches in a way that is maximally compatible. This means that it should just work in most cases, even with models and nodes not explicitly supported.
+
+If you encounter issues as a user or if you're a node developer and Prompt Control somehow breaks something, feel free to file a bug report.
+
 ## Prompt Control v2
 
 Prompt control has been almost completely rewritten. It now uses ComfyUI's lazy execution to build graphs from the text prompt at runtime. The generated graph is often exactly equivalent to a manually built workflow using native ComfyUI nodes. There are no more weird sampling hooks that could cause problems with other nodes
 
-Prompt Control also comes with `PCTextEncode`, which provides advanced text encoding with many additional features compared to ComfyUI's base `CLIPTextEncode`.
-
 ### Removed features
 
 - Prompt interpolation syntax; it was too cumbersome to maintain
 - LoRA block weight integration; ditto, for now.
 
-
-### Is it stable now?
-
-Unless I run into bugs or significant annoyances that require changing the interface, it probably won't change too much, but until I tag 2.0, everything can change.
-
 ### Everything broke, where are the old nodes?
 
 If you really need them, you can install the [legacy nodes](https://github.com/asagi4/comfyui-prompt-control-legacy). However, I will not fix bugs in those nodes, and I strongly recommend just migrating your workflows to the new nodes.
 
 You can have both installed at the same time; none of the nodes conflict.
 
-## What can it do?
-
-See [features](#features) below. Things you can control via the prompt:
-- Prompt editing and filtering without noodle soup
-- LoRA loading and scheduling via ComfyUI's hook system
-- Masking, composition and area control (regional prompting)
-- Prompt operations like `BREAK` and `AND`
-- Weight interpretation types (comfy, A1111, etc.)
-- Prompt masking with [cutoff](#cutoff)
-- And a bunch more
-
-See the [syntax documentation](doc/syntax.md)
-
-If you find prompt scheduling inconvenient for some reason, `PCTextEncode` can be used as a drop-in replacement for `CLIPTextEncode` to get everything else.
-
-
-[This workflow](workflows/example-lazy.json?raw=1) shows LoRA scheduling and prompt editing and compares it with the same prompt implemented with built-in ComfyUI nodes.
-
-[Here](workflows/example-2pass.json?raw=1) is a two-pass workflow illustrating more features, including custom masks and filtering.
-
-The tools in this repository combine well with the macro and wildcard functionality in [comfyui-utility-nodes](https://github.com/asagi4/comfyui-utility-nodes)
-
-
 ## Requirements
 
 For LoRA scheduling to work, you'll need at least version 0.3.7 of ComfyUI (0.3.36 of ComfyUI desktop).
@@ -96,75 +96,6 @@ This node attaches masks to a `CLIP` model so that they can be referred to when
 
 This node configures `PCTextEncode` default values for some functions by attaching the information to a `CLIP` model.
 
-# Features
-## Scheduling and LoRA loading
-
-Prompt control provides a way to easily schedule different prompts and control LoRA loading.
-
-See the [syntax documentation](doc/syntax.md)
-
-### Note on how schedules work
-
-ComfyUI does not use the step number to determine whether to apply conds; instead, it uses the sampler's timestep value which is affected by the scheduler you're using. This means that when the sampler scheduler isn't linear, the schedules generated by prompt control will not be either.
-
-## Advanced CLIP encoding
-
-If you use `PCTextEncode`, advanced encodings are available automatically. Thanks to BlenderNeko for the original code.
-
-Use the syntax `STYLE(weight_interpretation, normalization)` in a prompt to affect how prompts are interpreted.
-
-The weight interpretations available are:
-  - comfy (default)
-  - comfy++
-  - compel
-  - down_weight
-  - A1111
-  - perp
-
-Normalizations are:
-  - none (default)
-  - length
-  - mean
-
-The normalization calculations are independent operations and you can combine them with `+`, eg `STYLE(A1111, length+mean)` or `STYLE(comfy, mean+length)`, or even something silly like `STYLE(perp, mean+length+mean+length)`
-
-The style can be specified separately for each AND:ed prompt, but the first prompt is special; later prompts will "inherit" it as default. For example:
-
-```
-STYLE(A1111) a (red:1.1) cat with (brown:0.9) spots and a long tail AND an (old:0.5) dog AND a (green:1.4) (balloon:1.1)
-```
-will interpret everything as A1111, but
-```
-a (red:1.1) cat with (brown:0.9) spots and a long tail AND STYLE(A1111) an (old:0.5) dog AND a (green:1.4) (balloon:1.1)
-```
-Will interpret the first one using the default ComfyUI behaviour, the second prompt with A1111 and the last prompt with the default again
-
-For things (ie. the code imports) to work, the nodes must be cloned in a directory named exactly `ComfyUI_ADV_CLIP_emb`.
-
-## Cutoff
-
-NOTE: Cutoff syntax might change at some point; it's pretty clunky.
-
-`PCTextEncode` reimplements cutoff from [ComfyUI Cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff).
-
-The syntax is
-```
-a group of animals, [CUT:white cat:white], [CUT:brown dog:brown:0.5:1.0:1.0:_]
-```
-You should read the prompt as `a group of animals, white cat, brown dog`, but CUT causes the tokens in `target_tokens` to be masked off from the base prompt in `region_text`, so that their effect can be isolated, and you're less likely to get brown cats or white dogs.
-
-Target tokens are treated individually, separated by space, for example, `[CUT:green apple, red apple, green leaf:green apple]` will mask *both* greens and the apple, giving you `+ +, red +, + leaf`. To mask out just `green apple`, use `[CUT:green apple, red apple:green_apple]` which will result in a masked prompt of `+ +, red apple`. Escape `_` with a `\`.
-
-the parameters in the `CUT` section are `region_text:target_tokens:weight;strict_mask:start_from_masked:padding_token` of which only the first two are required. The default values are `weight=1.0`, `strict_mask=1.0` `start_from_masked=1.0`, `padding_token=+`
-
-If `strict_mask`, `start_from_masked` or `padding_token` are specified in more than one CUT, the *last* one becomes the default for any CUTs afterwards that do not explicitly set the parameters. For example, in:
-
-`[CUT:white cat:white:0.5] and [CUT:black parrot, flying:black:1.0:0.5] and [CUT:green apple:green]`
-
-`white cat` will a weight of 0.5, and 1.0 for all parameters, and `black parrot` and `green apple` will *both* have a `strict_mask` parameter of 0.5.
-
-The parameters affect how the masked and unmasked prompts are combined to produce the final embedding. Just play around with them.
-
 # 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.
@@ -172,3 +103,5 @@ The parameters affect how the masked and unmasked prompts are combined to produc
 If you want to enable a hack to fix this, set `PROMPTCONTROL_ENABLE_CACHE_HACK=1` in your environment. Unset it to disable.
 
 It's a purely optional performance optimization that allows Prompt Control nodes to override their cache keys in a way that should not interfere with other nodes. Note that the optimization only works if the text input to the lazy nodes is a constant (so either directly on the node or from a primitive); outputs from other nodes can't be optimized.
+
+- 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.

__init__.py

@@ -30,21 +30,8 @@ NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}
 
 nodes = ["base", "lazy", "tools"]
-optional_nodes = ["attnmask"]
-if importlib.util.find_spec("comfy.hooks"):
-    nodes.extend(["hooks"])
-else:
-    log.error("Your ComfyUI version is too old, can't import comfy.hooks. Update your installation.")
 
 for node in nodes:
     mod = importlib.import_module(f".prompt_control.nodes_{node}", package=__name__)
     NODE_CLASS_MAPPINGS.update(mod.NODE_CLASS_MAPPINGS)
     NODE_DISPLAY_NAME_MAPPINGS.update(mod.NODE_DISPLAY_NAME_MAPPINGS)
-
-for node in optional_nodes:
-    try:
-        mod = importlib.import_module(f".prompt_control.nodes_{node}", package=__name__)
-        NODE_CLASS_MAPPINGS.update(mod.NODE_CLASS_MAPPINGS)
-        NODE_DISPLAY_NAME_MAPPINGS.update(mod.NODE_DISPLAY_NAME_MAPPINGS)
-    except ImportError:
-        log.info(f"Could not import optional nodes: {node}; continuing anyway")

doc/syntax.md

@@ -1,6 +1,8 @@
-# Scheduling syntax
+# Prompt Control Syntax
 
-Syntax is like A1111 for now, but only fractions are supported for steps. LoRAs are scheduled by including them in a scheduling expression.
+If you're viewing this on GitHub, I recommend opening the outline by clicking the button in the top right corner of the text view (it is annoyingly easy to miss).
+
+Scheduling syntax is similar to A1111, but only fractions are supported for steps. LoRAs are scheduled by including them in a scheduling expression.
 
 ```
 a [large::0.1] [cat|dog:0.05] [<lora:somelora:0.5:0.6>::0.5]
@@ -12,7 +14,7 @@ a [large::0.1] [cat|dog:0.05] [<lora:somelora:0.5:0.6>::0.5]
 There are two forms of scheduled prompts.
 
 ### Basic scheduling expressions
-Basic expressions take the form `[before:after:X]` where `X` is the switch point, a decimal number between 0.0 and 1.0 inclusive, representing 0 to 100% of timesteps.
+Basic expressions take the form `[before:after:X]` where `X` is the switch point, a decimal number between 0.0 and 1.0 inclusive, representing 0 to 100% of timesteps. Either prompt can also be empty.
 For example:
 ```
 a [red:blue:0.5] cat
@@ -24,13 +26,24 @@ a [red:[blue::0.7]:0.5] cat
 
 switches from `a red cat` to `a blue cat` at 0.5 and to `a cat` at 0.7
 
-
-**Note:** As a special case, `[cat:0.5]` is like `[:cat:0.5]` meaning it switches from empty to `cat` at 0.5. Currently, `[:cat:0.5]` doesn't actually parse correctly, so you **must** use the shortcut form
+For convenience `[cat:0.5]` is equivalent to `[:cat:0.5]` meaning it switches from empty to `cat` at 0.5.
 
 ### Range expressions
 
-You can also use `a [during:after:0.3,0.7]` as a shortcut. The prompt be `a` until 0.3, `a during` until 0.7, and then `a after`. This form is equivalent to `[[during:after:0.7]:0.3]`
-For convenience, `[during:0.1,0.4]` is equivalent to `[during::0.1,0.4]`
+The most general form of a schedule is a range expression: For example, in `[before:during:after:0.3,0.7]`, The prompt be `a before` until 0.3, `a during` until 0.7, and then `a after`. This form is equivalent to `[before:[during:after:0.7]:0.3]`
+
+For convenience, `[during:0.1,0.4]` is equivalent to `[:during::0.1,0.4]` and `[during:after:0.1,0.4]` is equivalent to `[:during:after:0.1,0.4]`.
+
+`[before:during:after:0.1]` is the same as `[before:during:after:0.1,1.0]` which is same as `[before:during:0.1]`
+
+
+### Using step numbers with the Advanced nodes
+
+If you provide a non-zero value to `num_steps` to the `Advanced` versions of the scheduling nodes, you will be able to use step numbers in prompts.
+
+For now, a value between 0 and 1.0 will be interpreted as a percentage if it contains a ., and as an absolute step otherwise.
+
+This is just syntactic sugar. Behind the scenes, the values are converted to percentages and have normal ComfyUI scheduling behaviour.
 
 ## Tag selection
 Using the `FilterSchedule` node, in addition to step percentages, you can use a *tag* to select part of an input:
@@ -47,6 +60,7 @@ a [black:blue:X] [cat:dog:Y] [walking:running:Z] in space
 ```
 with `tags` `x,z` would result in the prompt `a blue cat running in space`
 
+The three prompt form `[a:b:c:TAG]` is parsed, but ignores `b` and is equivalent to `[a:c:TAG]`.
 
 ## LoRA Scheduling
 When using the lazy graph building nodes, LoRAs can be scheduled by referring to them in a scheduling expression, like so:
@@ -79,7 +93,6 @@ Might be useful with Jinja templating (see https://github.com/asagi4/comfyui-uti
 ```
 generates a LoRA schedule based on a sinewave
 
-
 # Basic prompt syntax
 
 This syntax is also available in outside scheduled prompts, where applicable.
@@ -90,41 +103,95 @@ The A111-style syntax `<lora:loraname:weight>` can be used to load LoRAs via the
 
 ## Combining prompts, A1111-style
 
-- The keyword `BREAK` causes the prompt to be tokenized in separate chunks, which results in each chunk being individually padded to the text encoder's maximum token length. This is mostly equivalent to the `ConditioningConcat` node.
+### BREAK
+The keyword `BREAK` causes the prompt to be tokenized in separate chunks, which results in each chunk being individually padded to the text encoder's maximum token length. This is mostly equivalent to the `ConditioningConcat` node.
 
-`AND` can be used to combine prompts. You can also use a weight at the end. It does a weighted sum of each prompt,
+### AND
+
+`AND` can be used to create "prompt segments". By default, it works as if you had combined the different prompts with `ConditioningCombine`. 
+
+It is also used with regional prompting to separate different prompts; see `MASK` and `ATTN` below.
+
+Prompts can have a weight at the end:
 ```
 cat :1 AND dog :2
 ```
-The weight defaults to 1 and are normalized so that `a:2 AND b:2` is equal to `a AND b`. `AND` is processed after schedule parsing, so you can change the weight mid-prompt: `cat:[1:2:0.5] AND dog`
+`AND` is processed after schedule parsing, so you can change the weight mid-prompt: `cat:[1:2:0.5] AND dog`
 
+The weight defaults to 1. If a prompt's weight is set to 0, it's **skipped entirely.** This can be useful when scheduling to completely disable a prompt:
 
-## Functions
+```
+cat [\:0::0.5] AND dog
+```
+Note that the `:` needs to be escaped with a `\` or it will be interpreted as scheduling syntax.
 
-There are some "functions" that can be included in a prompt to do various things. 
+# Functions
 
-Functions have the form `FUNCNAME(param1, param2, ...)`. How parameters are interpreted is up to the function.
-Note: Whitespace is *not* stripped from string parameters by default. Commas can be escaped with `\,`
+There are some "functions" that can be included in a prompt to affect how it is interpreted.
 
-Like `AND`, these functions are parsed after regular scheduling syntax has been expanded, allowing things like `[AREA:MASK:0.3](...)`, in case that's somehow useful.
+Functions have the form `FUNCNAME(param1, param2, ...)`. How parameters are interpreted is up to the function. 
 
-### SDXL
+In general, function parameters will have default values that are used if the parameter is left empty.
+
+Note: Whitespace is usually *not* stripped from string parameters by default. Commas can be escaped with `\,`
+
+Like `AND`, functions are parsed after regular scheduling syntax has been expanded, allowing things like `[AREA:MASK:0.3](...)`, in case that's somehow useful.
+
+### STYLE: Configure prompt weighting (also known as "Advanced CLIP Encode")
+
+Use the syntax `STYLE(weight_interpretation, normalization)` in a prompt to affect how prompts are interpreted.
+
+The weight interpretations available are:
+  - comfy (default)
+  - comfy++
+  - compel
+  - down_weight
+  - A1111
+  - perp
+
+Normalizations are:
+  - none (default)
+  - length
+  - mean
+
+The normalization calculations are independent operations and you can combine them with `+`, eg `STYLE(A1111, length+mean)` or `STYLE(comfy, mean+length)`, or even something silly like `STYLE(perp, mean+length+mean+length)`
+
+The style can be specified separately for each AND:ed prompt, but the first prompt is special; later prompts will "inherit" it as default. For example:
+
+```
+STYLE(A1111) a (red:1.1) cat with (brown:0.9) spots and a long tail AND an (old:0.5) dog AND a (green:1.4) (balloon:1.1)
+```
+will interpret everything as A1111, but
+```
+a (red:1.1) cat with (brown:0.9) spots and a long tail AND STYLE(A1111) an (old:0.5) dog AND a (green:1.4) (balloon:1.1)
+```
+Will interpret the first one using the default ComfyUI behaviour, the second prompt with A1111 and the last prompt with the default again
+
+### SDXL: Configure SDXL prompting parameters
 
 The nodes do not treat SDXL models specially, but there are some utilities that enable SDXL specific functionality.
 
 You can use the function `SDXL(width height, target_width target_height, crop_w crop_h)` to set SDXL prompt parameters. `SDXL()` is equivalent to `SDXL(1024 1024, 1024 1024, 0 0)` unless the default values have been overridden by `PCScheduleSettings`.
 
-To set the `clip_l` prompt, as with `CLIPTextEncodeSDXL`, use the function `CLIP_L(prompt text goes here)`.
+### TE: Per-encoder prompts for multi-encoder models
+
+You can specify per-encoder prompts using the `TE` function. The syntax is as follows:
+`TE(encoder_name=prompt)`. Whitespace surrounding the prompt and encoder name are ignored.
+
+For example:
+```
+TE(l=cat) TE(g = (dog:1.1)) TE(t5xxl=tiger)
+```
+The keys to use depend on what key ComfyUI uses for the encoder; for example `l` for CLIP L, `g` for CLIP G, and `t5xxl` for T5 XXL (Flux text encoder).
+
+Use `TE(help)` to print a help text listing available keys.
 
 Things to note:
-- Multiple instances of `CLIP_L` are joined with a space. That is, `CLIP_L(foo)CLIP_L(bar)` is the same as `CLIP_L(foo bar)`
-- Using `BREAK` isn't supported in it; it'll just parse as the plain word BREAK.
-- similarly, `AND` inside `CLIP_L` does not do anything sensible; `CLIP_L(foo AND bar)` will parse as two prompts `CLIP_L(foo` and `bar)`
-- `CLIP_L` and `SDXL` have no effect on SD 1.5.
-- The rest of the prompt becomes the `clip_g` prompt.
-- If there is no `CLIP_L` or `SDXL`, the prompts will work as with `CLIPTextEncode`.
+- If you set a prompt with `TE`, it will override the prompt outside the function for the specified text encoder.
+- Multiple instances of `TE` are joined with a space. That is, `TE(l=foo)TE(l=bar)` is the same as `TE(l=foo bar)`
+- `AND` inside `TE` does not do anything sensible; `TE(l=foo AND bar)` will parse as two prompts `TE(foo` and `bar)`. `BREAK`, `SHIFT` and `SHUFFLE` do work, however
 
-### SHUFFLE and SHIFT
+### SHUFFLE and SHIFT: Create prompt permutations
 
 Default parameters: `SHUFFLE(seed=0, separator=,, joiner=,)`, `SHIFT(steps=0, separator=,, joiner=,)`
 
@@ -149,23 +216,30 @@ For example:
 Whitespace is *not* stripped and may also be used as a joiner or separator
 - `SHIFT(1,, ) cat,dog` results in `dog cat`
 
-### NOISE
+### NOISE: Add noise to a prompt
 
 The function `NOISE(weight, seed)` adds some random noise into the prompt. The seed is optional, and if not specified, the global RNG is used. `weight` should be between 0 and 1.
 
+
+## Regional prompting
+
+See also [Attention Couple](#attention-couple) below
+
 ### MASK, IMASK and AREA
 
 You can use `MASK(x1 x2, y1 y2, weight, op)` to specify a region mask for a prompt. The values are specified as a percentage with a float between `0` and `1`, or as absolute pixel values (these can't be mixed). `1` will be interpreted as a percentage instead of a pixel value.
 
+Multiple `MASK` or `IMASK` calls will be composited together using ComfyUI's `MaskComposite` node, using `op` as the `operation` parameter (defaulting to `multiply`).
+
 Similarly, you can use `AREA(x1 x2, y1 y2, weight)` to specify an area for the prompt (see ComfyUI's area composition examples). The area is calculated by ComfyUI relative to your latent size.
 
-#### Custom masks: IMASK and `PCAddMaskToCLIP`
+### Custom masks: IMASK and `PCAddMaskToCLIP`
 
 You can attach custom masks to a `CLIP` with the `PC: Attach Mask` nodes and then refer to those masks in the prompt using `IMASK(index, weight, op)`. Indexing starts from zero, so 0 is the first attached mask etc. `PCSCheduleAddMasks` ignores empty inputs, so if you only add a mask to the `mask4` input, it will still have index 0.
 
 Applying the nodes multiple times *appends* masks rather than overriding existing ones, so if you need more than 4, you can just use it more than once.
 
-#### Behaviour of masks
+### Behaviour of masks
 If multiple `MASK`s are specified, they are combined together with ComfyUI's `MaskComposite` node, with `op` specifying the operation to use (default `multiply`). In this case, the combined mask weight can be set with `MASKW(weight)` (defaults to 1.0).
 
 Masks assume a size of `(512, 512)`, unless overridden with `PC: Configure PCTextEncode` and pixel values will be relative to that. ComfyUI will scale the mask to match the image resolution. You can change it manually by using `MASK_SIZE(width, height)` anywhere in the prompt,
@@ -178,7 +252,7 @@ Note that because the default values are percentages, `MASK(0 256, 64 512)` is v
 
 Masking does not affect LoRA scheduling unless you set unet weights to 0 for a LoRA.
 
-### FEATHER
+### FEATHER: Mask operations
 
 When you use `MASK` or `IMASK`, you can also call `FEATHER(left top right bottom)` to apply feathering using ComfyUI's `FeatherMask` node. The values are in pixels and default to `0`.
 
@@ -193,17 +267,120 @@ gives you a mask that is a combination of 1, 2 and 3, where 1 and 3 are feathere
 
 The order of the `FEATHER` and `MASK` calls doesn't matter; you can have `FEATHER` before `MASK` or even interleave them.
 
+## Cutoff
+
+NOTE: Cutoff syntax might change at some point; it's pretty clunky.
+
+`PCTextEncode` reimplements cutoff from [ComfyUI Cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff).
+
+The syntax is
+```
+a group of animals, [CUT:white cat:white], [CUT:brown dog:brown:0.5:1.0:1.0:_]
+```
+You should read the prompt as `a group of animals, white cat, brown dog`, but CUT causes the tokens in `target_tokens` to be masked off from the base prompt in `region_text`, so that their effect can be isolated, and you're less likely to get brown cats or white dogs.
+
+Target tokens are treated individually, separated by space, for example, `[CUT:green apple, red apple, green leaf:green apple]` will mask *both* greens and the apple, giving you `+ +, red +, + leaf`. To mask out just `green apple`, use `[CUT:green apple, red apple:green_apple]` which will result in a masked prompt of `+ +, red apple`. Escape `_` with a `\`.
+
+the parameters in the `CUT` section are `region_text:target_tokens:weight;strict_mask:start_from_masked:padding_token` of which only the first two are required. The default values are `weight=1.0`, `strict_mask=1.0` `start_from_masked=1.0`, `padding_token=+`
+
+If `strict_mask`, `start_from_masked` or `padding_token` are specified in more than one CUT, the *last* one becomes the default for any CUTs afterwards that do not explicitly set the parameters. For example, in:
+
+`[CUT:white cat:white:0.5] and [CUT:black parrot, flying:black:1.0:0.5] and [CUT:green apple:green]`
+
+`white cat` will a weight of 0.5, and 1.0 for all parameters, and `black parrot` and `green apple` will *both* have a `strict_mask` parameter of 0.5.
+
+The parameters affect how the masked and unmasked prompts are combined to produce the final embedding. Just play around with them.
+
 ## Miscellaneous
 - `<emb:xyz>` is alternative syntax for `embedding:xyz` to work around a syntax conflict with `[embedding:xyz:0.5]` which is parsed as a schedule that switches from `embedding` to `xyz`.
 
 # Experimental features
 
-Experimental features are unstable and may disappear or break without warning.
+Experimental features are unstable and may disappear or change without warning.
 
-## Attention masking
+## DEF: Lightweight prompt macros
 
-Use `ATTN()` in combination with `MASK()` or `IMASK()` to enable attention masking. Currently, it's pretty slow and only works with SDXL. You need to have a recent enough version of ComfyUI for this to work.
+You can define "prompt macros" by using `DEF`:
+```
+DEF(MYMACRO=this is a prompt)
+[(MYMACRO:0.6):(MYMACRO:1.1):0.5]
+```
+is equivalent to
+```
+[(this is a prompt:0.5):(this is a prompt:1.1):0.5]
+```
+
+It's also possible to give parameters to a macro:
+```
+DEF(MYMACRO=[(prompt $1:$2):(prompt $1:$3):$4])
+MYMACRO(test; 1.1; 0.7; 0.2)
+```
+gives
+```
+[(prompt test:1.1):(prompt test:0.7):0.2]
+```
+in this form, the variables $N (where N is any number corresponding to a positional parameter) will be replaced with the given parameter. The parameters must be separated with a semicolon.
+
+You can also optionally specify default values:
+
+```
+DEF(MACRO(example; 0; 1)=[$1:$2,$3])
+MACRO MACRO(test; 0.2)
+```
+gives
+```
+[example:0,1] [test:0.2,1]
+```
+
+Note that unspecified parameters will not be substituted:
+```
+DEF(mything=a $1 b $2)
+mything
+mything(A)
+```
+gives
+```
+a $1 b $2
+a A b $2
+```
+
+Macros are expanded before any other parsing takes place. The expansion continues until no further changes occur. Recursion will raise an error.
+
+## Attention Couple
+
+Attention Couple is an attention-based implementation of regional prompting. it can often be faster and more flexible than latent-based masking.
+
+The implementation is based on the one by [pamparamm](https://github.com/pamparamm/ComfyUI-ppm.git), but modified to use ComfyUI's hook system. This enables it to work with prompt scheduling.
+
+The implementation produces slightly different results from Pamparamm's implementation because ComfyUI will only run the hook for conds that have it attached, unlike the ModelPatcher based implementation which has special logic to avoid messing up negative prompts with attention masks. It's also slightly slower because ComfyUI can't batch cond and uncond calculations while the hook is in use.
+
+As a consequence of this, however, you can also use `ATTN()` in your negative prompt, and it will work correctly.
+
+### ATTN: Trigger Attention Couple
+
+Use `ATTN()` to mark a prompt to be used with Attention Couple. `ATTN()` needs to be combined with either `MASK()` or `IMASK()` to work correctly.
+
+If no mask is specified, an implicit `MASK()` is assumed.
+
+For attention masking to take effect, you need at least two prompt segments with the `ATTN()` marker (separated with `AND`). A single prompt with `ATTN()` will simply ignore the marker.
+
+For the first prompt (and the first prompt only) you can also use `FILL()` to automatically mask all parts not masked by other prompt segments.
+
+For example:
+```
+dog FILL() ATTN() AND cat MASK(0.5 1) ATTN()
+```
+
+If typing `ATTN() MASK()` feels bothersome, try the following macro:
+```
+DEF(AM=ATTN() MASK($1))
+```
+and then use it like `MASK`: `AM(0 1, 0.5 1)`
 
 ## TE_WEIGHT
 
-For models using multiple text encoders, you can set weights per TE using the syntax `TE_WEIGHT(clipname=weight, clipname2=weight2, ...)` where `clipname` is one of `g`, `l`, or `t5xxl`. For example with SDXL, try `TE_WEIGHT(g=0.25, l=0.75`). The weights are applied as a multiplier to the TE output.
+For models using multiple text encoders, you can set weights per TE using the syntax `TE_WEIGHT(clipname=weight, clipname2=weight2, ...)` where `clipname` is one of the encoder names printed by `TE(help)`. For example with SDXL, try `TE_WEIGHT(g=0.25, l=0.75)`.
+
+The weights are applied as a multiplier to the TE output. You can also override pooled output multipliers using eg. `l_pooled`.
+
+To set a default value for all encoders, use `TE_WEIGHT(all=weight)`

example_workflows/Basic Text to Image.json

@@ -0,0 +1,672 @@
+{
+  "id": "e820c2fb-9502-45b7-a864-684757dddcdf",
+  "revision": 0,
+  "last_node_id": 18,
+  "last_link_id": 20,
+  "nodes": [
+    {
+      "id": 1,
+      "type": "CheckpointLoaderSimple",
+      "pos": [
+        -135,
+        -930
+      ],
+      "size": [
+        315,
+        98
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "MODEL",
+          "type": "MODEL",
+          "slot_index": 0,
+          "links": [
+            2
+          ]
+        },
+        {
+          "name": "CLIP",
+          "type": "CLIP",
+          "slot_index": 1,
+          "links": [
+            3
+          ]
+        },
+        {
+          "name": "VAE",
+          "type": "VAE",
+          "slot_index": 2,
+          "links": [
+            18
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.3.18",
+        "Node name for S&R": "CheckpointLoaderSimple"
+      },
+      "widgets_values": [
+        "NoobAI-XL-Vpred-v1.0.safetensors"
+      ]
+    },
+    {
+      "id": 2,
+      "type": "PCLazyTextEncode",
+      "pos": [
+        555,
+        -720
+      ],
+      "size": [
+        252,
+        78
+      ],
+      "flags": {
+        "collapsed": true
+      },
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "clip",
+          "type": "CLIP",
+          "link": 5
+        },
+        {
+          "name": "text",
+          "type": "STRING",
+          "widget": {
+            "name": "text"
+          },
+          "link": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "CONDITIONING",
+          "type": "CONDITIONING",
+          "slot_index": 0,
+          "links": [
+            12
+          ]
+        }
+      ],
+      "title": "PC: Schedule Prompt (positive)",
+      "properties": {
+        "cnr_id": "comfyui-prompt-control",
+        "ver": "2.0.0-beta.7",
+        "Node name for S&R": "PCLazyTextEncode"
+      },
+      "widgets_values": [
+        "STYLE(A1111) 1girl, [painting \\(medium\\), realistic,::0.2] fennec fox girl, animal ear fluff, [[purple:white pupils, purple:0.2] eyes:sparkling eyes:0.85], cargo pants, long sleeves, cardigan, winter, snow, steaming cup, coffee mug, [thermos,:0.1] [long hair,:0.25] [BREAK:0.3]\n[(masterpiece, best quality, newest, very awa,):0.1], night sky, full moon, star \\(sky\\),"
+      ],
+      "color": "#232",
+      "bgcolor": "#353"
+    },
+    {
+      "id": 3,
+      "type": "PCLazyLoraLoader",
+      "pos": [
+        257.5,
+        -745
+      ],
+      "size": [
+        210,
+        78
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "model",
+          "shape": 7,
+          "type": "MODEL",
+          "link": 2
+        },
+        {
+          "name": "clip",
+          "shape": 7,
+          "type": "CLIP",
+          "link": 3
+        },
+        {
+          "name": "text",
+          "type": "STRING",
+          "widget": {
+            "name": "text"
+          },
+          "link": 6
+        }
+      ],
+      "outputs": [
+        {
+          "name": "MODEL",
+          "type": "MODEL",
+          "slot_index": 0,
+          "links": [
+            17
+          ]
+        },
+        {
+          "name": "CLIP",
+          "type": "CLIP",
+          "slot_index": 1,
+          "links": [
+            5,
+            9
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfyui-prompt-control",
+        "ver": "2.0.0-beta.7",
+        "Node name for S&R": "PCLazyLoraLoader"
+      },
+      "widgets_values": [
+        "STYLE(A1111) 1girl, [painting \\(medium\\), realistic,::0.2] fennec fox girl, animal ear fluff, [[purple:white pupils, purple:0.2] eyes:sparkling eyes:0.85], cargo pants, long sleeves, cardigan, winter, snow, steaming cup, coffee mug, [thermos,:0.1] [long hair,:0.25] [BREAK:0.3]\n[(masterpiece, best quality, newest, very awa,):0.1], night sky, full moon, star \\(sky\\),"
+      ],
+      "color": "#223",
+      "bgcolor": "#335"
+    },
+    {
+      "id": 4,
+      "type": "KSampler",
+      "pos": [
+        930,
+        -780
+      ],
+      "size": [
+        315,
+        474
+      ],
+      "flags": {},
+      "order": 9,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "model",
+          "type": "MODEL",
+          "link": 17
+        },
+        {
+          "name": "positive",
+          "type": "CONDITIONING",
+          "link": 12
+        },
+        {
+          "name": "negative",
+          "type": "CONDITIONING",
+          "link": 13
+        },
+        {
+          "name": "latent_image",
+          "type": "LATENT",
+          "link": 14
+        }
+      ],
+      "outputs": [
+        {
+          "name": "LATENT",
+          "type": "LATENT",
+          "slot_index": 0,
+          "links": [
+            15
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.3.18",
+        "Node name for S&R": "KSampler"
+      },
+      "widgets_values": [
+        2,
+        "fixed",
+        25,
+        1.4000000000000001,
+        "euler_cfg_pp",
+        "simple",
+        1
+      ]
+    },
+    {
+      "id": 5,
+      "type": "PrimitiveNode",
+      "pos": [
+        -270,
+        -780
+      ],
+      "size": [
+        495,
+        225
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "STRING",
+          "type": "STRING",
+          "widget": {
+            "name": "text"
+          },
+          "links": [
+            6,
+            7
+          ]
+        }
+      ],
+      "title": "Positive prompt (with LoRAs)",
+      "properties": {
+        "Run widget replace on values": false
+      },
+      "widgets_values": [
+        "STYLE(A1111) 1girl, [painting \\(medium\\), realistic,::0.2] fennec fox girl, animal ear fluff, [[purple:white pupils, purple:0.2] eyes:sparkling eyes:0.85], cargo pants, long sleeves, cardigan, winter, snow, steaming cup, coffee mug, [thermos,:0.1] [long hair,:0.25] [BREAK:0.3]\n[(masterpiece, best quality, newest, very awa,):0.1], night sky, full moon, star \\(sky\\),"
+      ],
+      "color": "#232",
+      "bgcolor": "#353"
+    },
+    {
+      "id": 6,
+      "type": "PrimitiveNode",
+      "pos": [
+        -270,
+        -510
+      ],
+      "size": [
+        480,
+        225
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "STRING",
+          "type": "STRING",
+          "widget": {
+            "name": "text"
+          },
+          "slot_index": 0,
+          "links": [
+            8
+          ]
+        }
+      ],
+      "title": "Negative prompt",
+      "properties": {
+        "Run widget replace on values": false
+      },
+      "widgets_values": [
+        "chibi, [bad hands,low quality, worst quality,:0.05], simple background, blurry, sketch, unfinished, [holding two cups,no pupils,:0.1]"
+      ],
+      "color": "#322",
+      "bgcolor": "#533"
+    },
+    {
+      "id": 7,
+      "type": "PCLazyTextEncode",
+      "pos": [
+        555,
+        -675
+      ],
+      "size": [
+        252,
+        78
+      ],
+      "flags": {
+        "collapsed": true
+      },
+      "order": 8,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "clip",
+          "type": "CLIP",
+          "link": 9
+        },
+        {
+          "name": "text",
+          "type": "STRING",
+          "widget": {
+            "name": "text"
+          },
+          "link": 8
+        }
+      ],
+      "outputs": [
+        {
+          "name": "CONDITIONING",
+          "type": "CONDITIONING",
+          "slot_index": 0,
+          "links": [
+            13
+          ]
+        }
+      ],
+      "title": "PC: Schedule Prompt (negative)",
+      "properties": {
+        "cnr_id": "comfyui-prompt-control",
+        "ver": "2.0.0-beta.7",
+        "Node name for S&R": "PCLazyTextEncode"
+      },
+      "widgets_values": [
+        "chibi, [bad hands,low quality, worst quality,:0.05], simple background, blurry, sketch, unfinished, [holding two cups,no pupils,:0.1]"
+      ],
+      "color": "#322",
+      "bgcolor": "#533"
+    },
+    {
+      "id": 9,
+      "type": "EmptyLatentImage",
+      "pos": [
+        525,
+        -615
+      ],
+      "size": [
+        315,
+        106
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "LATENT",
+          "type": "LATENT",
+          "links": [
+            14
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.3.18",
+        "Node name for S&R": "EmptyLatentImage"
+      },
+      "widgets_values": [
+        896,
+        1152,
+        1
+      ],
+      "color": "#432",
+      "bgcolor": "#653"
+    },
+    {
+      "id": 10,
+      "type": "VAEDecode",
+      "pos": [
+        1290,
+        -780
+      ],
+      "size": [
+        210,
+        46
+      ],
+      "flags": {},
+      "order": 10,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "samples",
+          "type": "LATENT",
+          "link": 15
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": 18
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "slot_index": 0,
+          "links": [
+            20
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.3.18",
+        "Node name for S&R": "VAEDecode"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 13,
+      "type": "MarkdownNote",
+      "pos": [
+        240,
+        -615
+      ],
+      "size": [
+        240,
+        105
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [],
+      "properties": {},
+      "widgets_values": [
+        "If you do not need LoRA scheduling, you can simply skip this node."
+      ],
+      "color": "#432",
+      "bgcolor": "#653"
+    },
+    {
+      "id": 15,
+      "type": "MarkdownNote",
+      "pos": [
+        240,
+        -450
+      ],
+      "size": [
+        600,
+        210
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [],
+      "properties": {},
+      "widgets_values": [
+        "`PC: Schedule prompt` will expand into instances of `PCTextEncode`. `PC: Schedule LoRAs` will expand into the required `LoRALoader`s and `CLIP` hooks required to schedule LoRAs in the prompt.\n\nYou can pass the same prompt to both nodes; `PC: Schedule Prompt` will simply ignore any `<lora:xyz:1>` elements, so they will not affect the prompt.\nSee the [full syntax available in the prompts](https://github.com/asagi4/comfyui-prompt-control/blob/master/doc/syntax.md) on GitHub"
+      ],
+      "color": "#432",
+      "bgcolor": "#653"
+    },
+    {
+      "id": 18,
+      "type": "SaveImage",
+      "pos": [
+        1290,
+        -690
+      ],
+      "size": [
+        270,
+        405
+      ],
+      "flags": {},
+      "order": 11,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 20
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.3.18"
+      },
+      "widgets_values": [
+        "PromptControl"
+      ]
+    }
+  ],
+  "links": [
+    [
+      2,
+      1,
+      0,
+      3,
+      0,
+      "MODEL"
+    ],
+    [
+      3,
+      1,
+      1,
+      3,
+      1,
+      "CLIP"
+    ],
+    [
+      5,
+      3,
+      1,
+      2,
+      0,
+      "CLIP"
+    ],
+    [
+      6,
+      5,
+      0,
+      3,
+      2,
+      "STRING"
+    ],
+    [
+      7,
+      5,
+      0,
+      2,
+      1,
+      "STRING"
+    ],
+    [
+      8,
+      6,
+      0,
+      7,
+      1,
+      "STRING"
+    ],
+    [
+      9,
+      3,
+      1,
+      7,
+      0,
+      "CLIP"
+    ],
+    [
+      12,
+      2,
+      0,
+      4,
+      1,
+      "CONDITIONING"
+    ],
+    [
+      13,
+      7,
+      0,
+      4,
+      2,
+      "CONDITIONING"
+    ],
+    [
+      14,
+      9,
+      0,
+      4,
+      3,
+      "LATENT"
+    ],
+    [
+      15,
+      4,
+      0,
+      10,
+      0,
+      "LATENT"
+    ],
+    [
+      17,
+      3,
+      0,
+      4,
+      0,
+      "MODEL"
+    ],
+    [
+      20,
+      10,
+      0,
+      18,
+      0,
+      "IMAGE"
+    ],
+    [
+      18,
+      1,
+      2,
+      10,
+      1,
+      "VAE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.8,
+      "offset": [
+        591.75,
+        1235
+      ]
+    },
+    "linkExtensions": [
+      {
+        "id": 18,
+        "parentId": 1
+      }
+    ],
+    "reroutes": [
+      {
+        "id": 1,
+        "pos": [
+          1273.75,
+          -879.5
+        ],
+        "linkIds": [
+          18
+        ]
+      }
+    ]
+  },
+  "version": 0.4,
+  "models": [{
+    "name": "NoobAI-XL-Vpred-v1.0.safetensors",
+    "url": "https://huggingface.co/Laxhar/noobai-XL-Vpred-1.0/resolve/main/NoobAI-XL-Vpred-v1.0.safetensors",
+    "directory": "checkpoints"
+  }]
+}

prompt_control/attention_couple_ppm.py

@@ -0,0 +1,150 @@
+# Lifted from https://github.com/pamparamm/ComfyUI-ppm/blob/c3e6b673ee2d424405dcb99aeed89f21943c89ac/nodes_ppm/attention_couple_ppm.py
+# Original implementation by laksjdjf, hako-mikan, Haoming02 licensed under GPL-3.0
+# https://github.com/laksjdjf/cgem156-ComfyUI/blob/1f5533f7f31345bafe4b833cbee15a3c4ad74167/scripts/attention_couple/node.py
+# https://github.com/Haoming02/sd-forge-couple/blob/e8e258e982a8d149ba59a4bc43b945467604311c/scripts/attention_couple.py
+import math
+
+import torch
+import torch.nn.functional as F
+from comfy.hooks import TransformerOptionsHook, HookGroup, EnumHookScope, set_hooks_for_conditioning
+
+from comfy.model_patcher import ModelPatcher
+
+import logging
+
+log = logging.getLogger("comfyui-prompt-control")
+
+
+def set_cond_attnmask(base_cond, extra_conds, fill=False):
+    hook = AttentionCoupleHook(base_cond[0], extra_conds, fill=fill)
+    group = HookGroup()
+    group.add(hook)
+    return set_hooks_for_conditioning(base_cond, hooks=group)
+
+
+def lcm_for_list(numbers):
+    current_lcm = numbers[0]
+    for number in numbers[1:]:
+        current_lcm = math.lcm(current_lcm, number)
+    return current_lcm
+
+
+class Proxy:
+    def __init__(self, function):
+        self.function = function
+
+    def to(self, *args, **kwargs):
+        self.function.__self__.to(*args, **kwargs)
+        return self
+
+    def __call__(self, *args, **kwargs):
+        return self.function(*args, *kwargs)
+
+
+class AttentionCoupleHook(TransformerOptionsHook):
+    def __init__(self, base_cond, conds, fill):
+        super().__init__(hook_scope=EnumHookScope.HookedOnly)
+        self.transformers_dict = {
+            "patches": {
+                "attn2_output_patch": [Proxy(self.attn2_output_patch)],
+                "attn2_patch": [Proxy(self.attn2_patch)],
+            }
+        }
+
+        self.num_conds = len(conds) + 1
+        self.base_strength = base_cond[1].pop("strength", 1.0)
+        self.strengths = [cond[1].get("strength", 1.0) for cond in conds]
+        self.conds: list[torch.Tensor] = [base_cond[0]] + [cond[0] for cond in conds]
+        base_mask = base_cond[1].pop("mask", None)
+        masks = [cond[1].pop("mask") * cond[1].pop("mask_strength") for cond in conds]
+
+        if base_mask is None and not fill:
+            raise ValueError("You must specify a base mask when fill=False")
+        elif base_mask is None:
+            sum = torch.stack(masks, dim=0).sum(dim=0)
+            base_mask = torch.zeros_like(sum)
+            base_mask[sum <= 0] = 1.0
+        mask = [base_mask] + masks
+        mask = torch.stack(mask, dim=0)
+        if mask.sum(dim=0).min() <= 0 and not fill:
+            raise ValueError("Masks contain non-filled areas")
+
+        self.mask = mask / mask.sum(dim=0, keepdim=True)
+        # calculate later
+        self.conds_k_tensor = None
+        self.conds_v_tensor = None
+
+    def on_apply_hooks(self, model: ModelPatcher, transformer_options: dict[str]):
+        if self.conds_k_tensor is None:
+            attn_patches = model.model_options["transformer_options"].get("patches", {}).get("attn2_patch", [])
+            has_negpip = any("negpip_attn" in i.__name__ for i in attn_patches)
+            log.debug("AttentionCouple has_negpip=%s", has_negpip)
+
+            conds_kv = (
+                [(cond[:, 0::2], cond[:, 1::2]) for cond in self.conds]
+                if has_negpip
+                else [(cond, cond) for cond in self.conds]
+            )
+
+            num_tokens_k = [cond[0].shape[1] for cond in conds_kv]
+            num_tokens_v = [cond[1].shape[1] for cond in conds_kv]
+
+            lcm_tokens_k = lcm_for_list(num_tokens_k)
+            lcm_tokens_v = lcm_for_list(num_tokens_v)
+            # Skip the base cond here, which is always first
+            self.conds_k_tensor = torch.cat(
+                [
+                    cond[0].repeat(1, lcm_tokens_k // num_tokens_k[i + 1], 1) * self.strengths[i]
+                    for i, cond in enumerate(conds_kv[1:])
+                ],
+                dim=0,
+            )
+            if has_negpip:
+                self.conds_v_tensor = torch.cat(
+                    [
+                        cond[1].repeat(1, lcm_tokens_v // num_tokens_v[i + 1], 1) * self.strengths[i]
+                        for i, cond in enumerate(conds_kv[1:])
+                    ],
+                    dim=0,
+                )
+            else:
+                self.conds_v_tensor = self.conds_k_tensor
+
+        return super().on_apply_hooks(model, transformer_options)
+
+    def to(self, *args, **kwargs):
+        self.conds = [c.to(*args, **kwargs) for c in self.conds]
+        self.mask = self.mask.to(*args, **kwargs)
+        return self
+
+    def attn2_patch(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, extra_options):
+        cond_or_uncond = extra_options["cond_or_uncond"]
+        num_chunks = len(cond_or_uncond)  # should always be 1
+        bs = q.shape[0] // num_chunks
+
+        conds_k_tensor = self.conds_k_tensor.expand(bs, *self.conds_k_tensor.shape[1:])
+        conds_v_tensor = self.conds_v_tensor.expand(bs, *self.conds_v_tensor.shape[1:])
+
+        q = q.repeat(self.num_conds, 1, 1)
+        k = k.repeat(1, self.conds_k_tensor.shape[1] // k.shape[1], 1)
+        v = v.repeat(1, self.conds_v_tensor.shape[1] // v.shape[1], 1)
+
+        k = torch.cat([k * self.base_strength, conds_k_tensor], dim=0)
+        v = torch.cat([v * self.base_strength, conds_v_tensor], dim=0)
+
+        return q, k, v
+
+    def attn2_output_patch(self, out, extra_options):
+        # out has been extended to shape [num_conds*batch_size, TOKENS, N]
+        # out is [b1c1 b1c2 ... b1cN, b2c1 b2c2 ... b2cn, ...]
+        num_conds = self.mask.shape[0]
+        bs = out.shape[0] // num_conds
+        num_tokens = out.shape[1]
+        mask_size = extra_options["activations_shape"][-2:]
+        mask_downsample = F.interpolate(self.mask, size=mask_size, mode="nearest")
+        mask_downsample = mask_downsample.view(num_conds, num_tokens, 1).repeat_interleave(bs, dim=0)
+
+        # cond_outputs is [num_conds*bs, tokens, N], output needs to be [bs, tokens, N]
+        cond_outputs = out * mask_downsample
+        cond_output = cond_outputs.view(num_conds, bs, out.shape[1], out.shape[2]).sum(0)
+        return cond_output

prompt_control/nodes_attnmask.py

@@ -1,79 +0,0 @@
-import logging
-
-log = logging.getLogger("comfyui-prompt-control")
-
-from comfy.hooks import TransformerOptionsHook, HookGroup, EnumHookScope
-from comfy.ldm.modules.attention import optimized_attention
-import torch.nn.functional as F
-import torch
-from math import sqrt
-
-
-class MaskedAttn2:
-    def __init__(self, mask):
-        self.mask = mask
-
-    def __call__(self, q, k, v, extra_options):
-        mask = self.mask
-        orig_shape = extra_options["original_shape"]
-        _, _, oh, ow = orig_shape
-        seq_len = q.shape[1]
-        mask_h = oh / sqrt(oh * ow / seq_len)
-        mask_h = int(mask_h) + int((seq_len % int(mask_h)) != 0)
-        mask_w = seq_len // mask_h
-        r = optimized_attention(q, k, v, extra_options["n_heads"])
-        mask = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="nearest").squeeze(1)
-        mask = mask.view(mask.shape[0], -1, 1).repeat(1, 1, r.shape[2])
-
-        return mask * r
-
-
-def create_attention_hook(mask):
-    attn_replacements = {}
-    mask = mask.detach().to(device="cuda", dtype=torch.float16)
-
-    masked_attention = MaskedAttn2(mask)
-
-    for id in [4, 5, 7, 8]:  # id of input_blocks that have cross attention
-        block_indices = range(2) if id in [4, 5] else range(10)  # transformer_depth
-        for index in block_indices:
-            k = ("input", id, index)
-            attn_replacements[k] = masked_attention
-    for id in range(6):  # id of output_blocks that have cross attention
-        block_indices = range(2) if id in [3, 4, 5] else range(10)  # transformer_depth
-        for index in block_indices:
-            k = ("output", id, index)
-            attn_replacements[k] = masked_attention
-    for index in range(10):
-        k = ("middle", 1, index)
-        attn_replacements[k] = masked_attention
-
-    hook = TransformerOptionsHook(
-        transformers_dict={"patches_replace": {"attn2": attn_replacements}}, hook_scope=EnumHookScope.HookedOnly
-    )
-    group = HookGroup()
-    group.add(hook)
-
-    return group
-
-
-class AttentionMaskHookExperimental:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {"mask": ("MASK",)},
-        }
-
-    RETURN_TYPES = ("HOOKS",)
-    CATEGORY = "promptcontrol/_testing"
-    FUNCTION = "apply"
-    EXPERIMENTAL = True
-    DESCRIPTION = "Experimental attention masking hook. For testing only"
-
-    def apply(self, mask):
-        return (create_attention_hook(mask),)
-
-
-NODE_CLASS_MAPPINGS = {"AttentionMaskHookExperimental": AttentionMaskHookExperimental}
-
-NODE_DISPLAY_NAME_MAPPINGS = {}

prompt_control/nodes_lazy.py

@@ -2,7 +2,9 @@ import logging
 from .parser import parse_prompt_schedules
 from comfy_execution.graph_utils import GraphBuilder, is_link
 
-from .prompts import get_function
+from comfy_execution.graph import ExecutionBlocker
+
+from .utils import get_function
 
 log = logging.getLogger("comfyui-prompt-control")
 
@@ -85,11 +87,15 @@ def create_hook_nodes_for_lora(graph, path, info, existing_node, start_pct, end_
     return hook_node, next_keyframe
 
 
-def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_hooks=True):
+def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True):
     # This gets rid of non-existent LoRAs
     consolidated = consolidate_schedule(schedule)
-    non_scheduled = find_nonscheduled_loras(consolidated)
-    model, clip = create_lora_loader_nodes(graph, model, clip, non_scheduled)
+    if model is not None:
+        non_scheduled = find_nonscheduled_loras(consolidated)
+        model, clip = create_lora_loader_nodes(graph, model, clip, non_scheduled)
+    else:
+        non_scheduled = {}
+        model = ExecutionBlocker("No model provided to PCLazyLoRALoader or PCLazyLoRALoaderAdvanced")
 
     hook_nodes = {}
     start_pct = 0.0
@@ -124,7 +130,7 @@ def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_h
             n.set_input("hooks_B", h.out(0))
             res = n
         res = res.out(0)
-        if apply_hooks:
+        if clip is not None and apply_hooks:
             n = graph.node("SetClipHooks")
             n.set_input("clip", clip)
             n.set_input("hooks", res)
@@ -132,13 +138,12 @@ def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_h
             n.set_input("schedule_clip", True)
             clip = n.out(0)
 
+    if clip is None:
+        clip = ExecutionBlocker("No clip model provided to PCLazyLoRALoader or PCLazyLoRALoaderAdvanced")
     r = graph.finalize()
     log.debug("LazyLoraLoader built graph: %s", json.dumps(r))
 
-    if return_hooks:
-        ret = (model, clip, res)
-    else:
-        ret = (model, clip)
+    ret = (model, clip, res)
 
     return {"result": ret, "expand": r}
 
@@ -149,16 +154,15 @@ class PCLazyLoraLoaderAdvanced:
     @classmethod
     def INPUT_TYPES(s):
         return {
-            "required": {
-                "text": ("STRING", {"multiline": True}),
+            "optional": {
                 "model": ("MODEL", {"rawLink": True}),
                 "clip": ("CLIP", {"rawLink": True}),
-            },
-            "optional": {
+                "text": ("STRING", {"multiline": True, "default": ""}),
                 "apply_hooks": ("BOOLEAN", {"default": True}),
                 "tags": ("STRING", {"default": ""}),
                 "start": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 0.0, "step": 0.01}),
                 "end": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 1.0, "step": 0.01}),
+                "num_steps": ("INT", {"min": 0, "max": 10000, "default": 0, "step": 1}),
             },
             "hidden": {"unique_id": "UNIQUE_ID"},
         }
@@ -168,22 +172,23 @@ class PCLazyLoraLoaderAdvanced:
     CATEGORY = "promptcontrol"
     FUNCTION = "apply"
 
-    def apply(self, model, clip, text, unique_id, apply_hooks=True, tags="", start=0.0, end=1.0):
-        schedule = parse_prompt_schedules(text, filters=tags, start=start, end=end)
-        graph = GraphBuilder(f"PCLazyLoraLoaderAdvanced-{unique_id}")
-        return build_lora_schedule(graph, schedule, model, clip, apply_hooks=apply_hooks, return_hooks=True)
+    def apply(
+        self, unique_id, model=None, clip=None, text="", apply_hooks=True, tags="", start=0.0, end=1.0, num_steps=0
+    ):
+        schedule = parse_prompt_schedules(text, filters=tags, start=start, end=end, num_steps=num_steps)
+        graph = GraphBuilder(f"{unique_id}-")
+        r = build_lora_schedule(graph, schedule, model, clip, apply_hooks=apply_hooks)
+        return r
 
 
-class PCLazyLoraLoader:
-    CACHE_KEY = cache_key_lora
-
+class PCLazyLoraLoader(PCLazyLoraLoaderAdvanced):
     @classmethod
     def INPUT_TYPES(s):
         return {
-            "required": {
+            "optional": {
                 "model": ("MODEL", {"rawLink": True}),
                 "clip": ("CLIP", {"rawLink": True}),
-                "text": ("STRING", {"multiline": True}),
+                "text": ("STRING", {"multiline": True, "default": ""}),
             },
             "hidden": {"unique_id": "UNIQUE_ID"},
         }
@@ -192,14 +197,12 @@ class PCLazyLoraLoader:
         "MODEL",
         "CLIP",
     )
-    OUTPUT_TOOLTIPS = ("Returns a model and clip with LoRAs scheduled",)
     CATEGORY = "promptcontrol"
-    FUNCTION = "apply"
 
-    def apply(self, model, clip, text, unique_id):
-        graph = GraphBuilder(f"PCLazyLoraLoader-{unique_id}")
-        schedule = parse_prompt_schedules(text)
-        return build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_hooks=False)
+    def apply(self, *args, **kwargs):
+        r = super().apply(*args, **kwargs)
+        r["result"] = r["result"][:2]
+        return r
 
 
 def build_scheduled_prompts(graph, schedules, clip):
@@ -235,31 +238,11 @@ def build_scheduled_prompts(graph, schedules, clip):
     return {"result": (node.out(0),), "expand": g}
 
 
-def cache_key_from_inputs(cachekey, text, tags="", start=0.0, end=1.0, **kwargs):
-    schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end)
+def cache_key_from_inputs(cachekey, text, tags="", start=0.0, end=1.0, num_steps=0, **kwargs):
+    schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end, num_steps=num_steps)
     return [(pct, s[cachekey]) for pct, s in schedules]
 
 
-class PCLazyTextEncode:
-    CACHE_KEY = cache_key_prompt
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {"clip": ("CLIP", {"rawLink": True}), "text": ("STRING", {"multiline": True})},
-        }
-
-    RETURN_TYPES = ("CONDITIONING",)
-    OUTPUT_TOOLTIPS = ("A fully encoded and scheduled conditioning",)
-    CATEGORY = "promptcontrol"
-    FUNCTION = "apply"
-
-    def apply(self, clip, text):
-        schedules = parse_prompt_schedules(text)
-        graph = GraphBuilder()
-        return build_scheduled_prompts(graph, schedules, clip)
-
-
 class PCLazyTextEncodeAdvanced:
     CACHE_KEY = cache_key_prompt
 
@@ -271,6 +254,7 @@ class PCLazyTextEncodeAdvanced:
                 "tags": ("STRING", {"default": ""}),
                 "start": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 0.0, "step": 0.01}),
                 "end": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 1.0, "step": 0.01}),
+                "num_steps": ("INT", {"min": 0, "max": 10000, "default": 0, "step": 1}),
             },
             "hidden": {"unique_id": "UNIQUE_ID"},
         }
@@ -279,12 +263,23 @@ class PCLazyTextEncodeAdvanced:
     CATEGORY = "promptcontrol"
     FUNCTION = "apply"
 
-    def apply(self, clip, text, unique_id, tags="", start=0.1, end=1.0):
-        schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end)
-        graph = GraphBuilder(f"PCLazyTextEncodeAdvanced-{unique_id}")
+    def apply(self, clip, text, unique_id, tags="", start=0.0, end=1.0, num_steps=0):
+        schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end, num_steps=num_steps)
+        graph = GraphBuilder(f"{unique_id}-")
         return build_scheduled_prompts(graph, schedules, clip)
 
 
+class PCLazyTextEncode(PCLazyTextEncodeAdvanced):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {"clip": ("CLIP", {"rawLink": True}), "text": ("STRING", {"multiline": True})},
+            "hidden": {"unique_id": "UNIQUE_ID"},
+        }
+
+    CATEGORY = "promptcontrol"
+
+
 NODE_CLASS_MAPPINGS = {
     "PCLazyTextEncode": PCLazyTextEncode,
     "PCLazyTextEncodeAdvanced": PCLazyTextEncodeAdvanced,

prompt_control/nodes_tools.py

@@ -1,9 +1,71 @@
 import logging
 from .parser import parse_prompt_schedules
+from .nodes_lazy import NODE_CLASS_MAPPINGS as LAZY_NODES
+import json
+import folder_paths
+from pathlib import Path
+from comfy_execution.graph_utils import is_link
 
 log = logging.getLogger("comfyui-prompt-control")
 
 
+class PCSaveExpandedWorkflow:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "any": ("*", {}),
+            },
+            "hidden": {
+                "prompt": "DYNPROMPT",
+            },
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(self, input_types):
+        return True
+
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+    CATEGORY = "promptcontrol/tools"
+    DESCRIPTION = "Saves the current expanded dynamic prompt into a JSON file"
+
+    FUNCTION = "apply"
+
+    def apply(self, any, prompt):
+        full_output_folder, filename, counter, subfolder, prefix = folder_paths.get_save_image_path(
+            "pc_workflow_debug", self.output_dir
+        )
+        p = {}
+        input_replace_map = {}
+        for node in prompt.all_node_ids():
+            n = prompt.get_node(node)
+            t = n["class_type"]
+            if t in LAZY_NODES:
+                expanded_prompt = LAZY_NODES[t]().apply(**n["inputs"], unique_id=node)
+                for k in expanded_prompt["expand"]:
+                    p[k] = expanded_prompt["expand"][k]
+                for i, _ in enumerate(expanded_prompt["result"]):
+                    input_replace_map[(node, i)] = [k, i]
+            else:
+                p[node] = n
+        for k in p:
+            for ik in p[k]["inputs"]:
+                x = p[k]["inputs"][ik]
+                if is_link(x) and tuple(x) in input_replace_map:
+                    p[k]["inputs"][ik] = input_replace_map[tuple(x)]
+        file = f"{filename}_{counter:05}_.json"
+        full_path = Path(full_output_folder) / file
+        with open(full_path, "w") as f:
+            log.info(f"Saving workflow to {full_path}")
+            json.dump(p, f)
+
+        return ()
+
+
 class PCSetLogLevel:
     @classmethod
     def INPUT_TYPES(s):
@@ -81,7 +143,6 @@ class PCSetPCTextEncodeSettings:
         return {
             "required": {"clip": ("CLIP",)},
             "optional": {
-                "steps": ("INT", {"default": 0, "min": 0, "max": 10000}),
                 "mask_width": ("INT", {"default": 512, "min": 64, "max": 4096 * 4}),
                 "mask_height": ("INT", {"default": 512, "min": 64, "max": 4096 * 4}),
                 "sdxl_width": ("INT", {"default": 1024, "min": 0, "max": 4096 * 4}),
@@ -101,7 +162,6 @@ class PCSetPCTextEncodeSettings:
     def apply(
         self,
         clip,
-        steps=0,
         mask_width=512,
         mask_height=512,
         sdxl_width=1024,
@@ -112,7 +172,6 @@ class PCSetPCTextEncodeSettings:
         sdxl_crop_h=0,
     ):
         settings = {
-            "steps": steps,
             "mask_width": mask_width,
             "mask_height": mask_height,
             "sdxl_width": sdxl_width,
@@ -156,6 +215,7 @@ NODE_CLASS_MAPPINGS = {
     "PCAddMaskToCLIPMany": PCAddMaskToCLIPMany,
     "PCSetLogLevel": PCSetLogLevel,
     "PCExtractScheduledPrompt": PCExtractScheduledPrompt,
+    "PCSaveExpandedWorkflow": PCSaveExpandedWorkflow,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -164,4 +224,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "PCAddMaskToCLIPMany": "PC: Attach Mask (multi)",
     "PCSetLogLevel": "PC: Configure Logging (for debug)",
     "PCExtractScheduledPrompt": "PC: Extract Scheduled Prompt",
+    "PCSaveExpandedWorkflow": "PC: Save Expanded Workflow (for debug)",
 }

prompt_control/parser.py

@@ -1,3 +1,4 @@
+# vim: sw=4 ts=4
 import lark
 import logging
 from math import ceil
@@ -7,7 +8,7 @@ log = logging.getLogger("comfyui-prompt-control")
 import re
 
 from functools import lru_cache
-from .utils import get_function
+from .utils import get_function, find_closing_paren
 
 if lark.__version__ == "0.12.0":
     from sys import executable
@@ -30,9 +31,10 @@ prompt: (emphasized | embedding | scheduled | alternate | sequence | loraspec |
 !emphasized: "(" prompt? ")"
         | "(" prompt ":" prompt ")"
         | "[" prompt "]"
-scheduled: "[" [prompt ":"] [prompt] ":" _WS? NUMBER ["," NUMBER] "]"
-        | "[" [prompt ":"] [prompt] ":" _WS? TAG "]"
-sequence:  "[SEQ" ":" [prompt] ":" NUMBER (":" [prompt] ":" NUMBER)+ "]"
+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
@@ -90,7 +92,7 @@ def parse_cuts(text):
 
 
 def flatten(x):
-    if type(x) in [str, tuple] or isinstance(x, dict) and "type" in 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:
@@ -102,13 +104,25 @@ def clamp(a, b, c):
     return min(max(a, b), c)
 
 
-def get_steps(tree):
-    res = [100]
+def get_steps(tree, num_steps):
+    res = [num_steps or 100]
 
     def tostep(s):
-        w = float(s) * 100
-        w = int(clamp(0, w, 100))
-        return w
+        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):
@@ -130,15 +144,14 @@ def get_steps(tree):
         def sequence(self, tree):
             steps = tree.children[1::2]
             for i, steps in enumerate(steps):
-                w = float(tree.children[i * 2 + 1]) * 100
-                tree.children[i * 2 + 1] = clamp(0, w, 100)
+                w = tostep(tree.children[i * 2 + 1])
+                tree.children[i * 2 + 1] = w
                 res.append(w)
 
         def alternate(self, tree):
-            step_size = int(round(float(tree.children[-1] or 0.1), 2) * 100)
-            step_size = clamp(1, step_size, 100)
+            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, 100, step_size)])
+            res.extend([x for x in range(step_size, num_steps or 100, step_size)])
 
     CollectSteps().visit(tree)
 
@@ -148,31 +161,38 @@ def get_steps(tree):
 def at_step(step, filters, tree):
     class AtStep(lark.Transformer):
         def scheduled(self, args):
+            before = None
+            during = None
+            after = None
             when_end = None
-            before, after, when, *rest = args
-            if isinstance(when, str):
-                return before or "" if when not in filters else after or ""
-
+            pl, when, *rest = args
             if rest:
                 when_end = rest[0]
 
-            if when_end is not None and step <= when and before is not None:
-                return ""
+            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 when_end is not None and (step > when and step <= when_end):
-                # handle [a:0,1]
-                if before is None:
-                    return after or ""
-                return before or ""
+            if isinstance(when, str):
+                return before or "" if when not in filters else after or ""
 
-            if when_end is not None and step >= when_end:
-                # handle [a:0,1]
-                if before is None:
-                    return ""
-                return 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 ""
 
@@ -258,29 +278,31 @@ def at_step(step, filters, tree):
 
 
 class PromptSchedule(object):
-    def __init__(self, prompt, filters="", start=0.0, end=1.0):
+    # 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
         self.prompt = prompt.strip()
         self.defaults = {}
         self.loaded_loras = {}
 
-        self.parsed_prompt = self._parse()
+        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):
+    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)
+            steps = get_steps(tree, num_steps=num_steps)
 
             def f(x):
-                return round(x / 100, 2)
+                return round(x / (num_steps or 100), 2)
 
             for t in steps:
                 p = at_step(t, filters, tree)
@@ -289,6 +311,7 @@ class PromptSchedule(object):
         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 = []
@@ -308,7 +331,15 @@ class PromptSchedule(object):
         if len(res) == 0:
             res = [[1.0, parsed[-1][1]]]
 
-        return res
+        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()
@@ -322,6 +353,7 @@ class PromptSchedule(object):
             filters=ifspecified(filters, self.filters),
             start=ifspecified(start, self.start),
             end=ifspecified(end, self.end),
+            num_steps=self.num_steps,
         )
         return p
 
@@ -336,34 +368,80 @@ class PromptSchedule(object):
         return len(self.parsed_prompt) - 1, self.parsed_prompt[-1]
 
 
-def replace_defs(text):
+def parse_search(search):
+    arg_start = search.find("(")
+    args = ""
+    name = search.strip()
+    if arg_start > 0:
+        arg_end = find_closing_paren(search, arg_start)
+        name = search[:arg_start].strip()
+        args = search[arg_start + 1 : arg_end - 1]
+
+    if not name:
+        return None
+    args = args.strip()
+    if args:
+        args = [a.strip() for a in args.split(";")]
+    else:
+        args = []
+    return name, args
+
+
+def replace_def(text):
     text, defs = get_function(text, "DEF", defaults=None)
     res = text
     prevres = text
     replacements = []
     for d in defs:
         r = d.split("=", 1)
-        if len(r) != 2 or not r[0].strip():
+        search = parse_search(r[0].strip())
+        if not search or len(r) != 2:
             log.warning("Ignoring invalid DEF(%s)", d)
             continue
-        replacements.append((r[0].strip(), r[1].strip()))
+        replacements.append((search, r[1].strip()))
     iterations = 0
     while True:
         iterations += 1
         if iterations > 10:
-            log.error("Unable to resolve DEFs, make sure there are no cycles!")
+            raise ValueError("Unable to resolve DEFs, make sure there are no cycles!")
             return text
         for search, replace in replacements:
-            res = re.sub(rf"\b{re.escape(search)}\b", replace, res)
+            res = substitute_defcall(res, search, replace)
+            res = substitute_def(res, search, replace)
         if res == prevres:
             break
         prevres = res
-    if res != text:
+    if res.strip() != text.strip():
+        res = res.strip()
         log.info("DEFs expanded to: %s", res)
     return res
 
 
+def substitute_def(text, search, replace):
+    search, default_args = search
+    for i, v in enumerate(default_args):
+        replace = re.sub(rf"\${i+1}\b", v, replace)
+    return re.sub(rf"\b{re.escape(search)}\b", replace, text)
+
+
+def substitute_defcall(text, search, replace):
+    name, default_args = search
+    text, defns = get_function(text, name, defaults=None, placeholder=f"DEFNCALL{search}")
+    for i, defn in enumerate(defns):
+        ph = f"\0DEFNCALL{search}{i}\0"
+        paramvals = [x.strip() for x in defn.split(";")]
+        r = replace
+        for i, v in enumerate(paramvals):
+            r = re.sub(rf"\${i+1}\b", v, r)
+
+        for i, v in enumerate(default_args):
+            r = re.sub(rf"\${i+1}\b", v, r)
+
+        text = text.replace(ph, r)
+    return text
+
+
 @lru_cache
 def parse_prompt_schedules(prompt, **kwargs):
-    prompt = replace_defs(prompt)
+    prompt = replace_def(prompt)
     return PromptSchedule(prompt, **kwargs)

prompt_control/prompts.py

@@ -9,20 +9,7 @@ from .adv_encode import advanced_encode_from_tokens
 from .cutoff import process_cuts
 from .parser import parse_cuts
 
-try:
-    from .nodes_attnmask import create_attention_hook
-    from comfy.hooks import set_hooks_for_conditioning
-
-    def set_cond_attnmask(cond, mask):
-        hook = create_attention_hook(mask)
-        return set_hooks_for_conditioning(cond, hooks=hook)
-
-except ImportError:
-
-    def set_cond_attnmask(cond, mask):
-        log.info("Attention masking is not available")
-        return cond
-
+from .attention_couple_ppm import set_cond_attnmask
 
 log = logging.getLogger("comfyui-prompt-control")
 
@@ -138,6 +125,28 @@ def fix_word_ids(tokens):
     return tokens
 
 
+def tokenize_chunks(clip, text, need_word_ids):
+    chunks = re.split(r"\bBREAK\b", text)
+    token_chunks = []
+    for c in chunks:
+        c, shuffles = get_function(c.strip(), "(SHIFT|SHUFFLE)", ["0", "default", "default"], return_func_name=True)
+        r = c
+        for s in shuffles:
+            r = shuffle_chunk(s, r)
+        if r != c:
+            log.info("Shuffled prompt chunk to %s", r)
+            c = r
+        t = clip.tokenize(c, return_word_ids=need_word_ids)
+        token_chunks.append(t)
+    tokens = token_chunks[0]
+
+    for key in tokens:
+        for c in token_chunks[1:]:
+            tokens[key].extend(c[key])
+
+    return tokens
+
+
 def encode_prompt_segment(
     clip,
     text,
@@ -157,37 +166,46 @@ def encode_prompt_segment(
 
     # defaults=None means there is no argument parsing at all
     text, l_prompts = get_function(text, "CLIP_L", defaults=None)
-    chunks = re.split(r"\bBREAK\b", text)
-    token_chunks = []
+    text, te_prompts = get_function(text, "TE", defaults=None)
     need_word_ids = True
-    for c in chunks:
-        c, shuffles = get_function(c.strip(), "(SHIFT|SHUFFLE)", ["0", "default", "default"], return_func_name=True)
-        r = c
-        for s in shuffles:
-            r = shuffle_chunk(s, r)
-        if r != c:
-            log.info("Shuffled prompt chunk to %s", r)
-            c = r
-        t = clip.tokenize(c, return_word_ids=need_word_ids)
-        token_chunks.append(t)
-    tokens = token_chunks[0]
+    tokens = tokenize_chunks(clip, text, need_word_ids)
 
-    for key in tokens:
-        for c in token_chunks[1:]:
-            tokens[key].extend(c[key])
+    per_te_prompts = {}
+    if l_prompts:
+        log.warning("Note: CLIP_L is deprecated. Use TE(l=prompt) instead")
+        per_te_prompts["l"] = l_prompts
 
-    # Non-SDXL has only "l"
-    if "g" in tokens and l_prompts:
-        text_l = " ".join(l_prompts)
-        log.info("Encoded SDXL CLIP_L prompt: %s", text_l)
-        tokens["l"] = clip.tokenize(text_l, return_word_ids=need_word_ids)["l"]
+    for prompt in te_prompts:
+        if prompt.strip() == "help":
+            log.info("Encoders available for TE: %s", ", ".join(tokens.keys()))
+            continue
+        params = prompt.split("=", 1)
+        if len(params) != 2:
+            log.warning("Invalid TE call, ignoring: %s", prompt)
+            continue
+        te = params[0].strip()
+        prompt = params[1].strip()
+        if te not in tokens:
+            log.warning("Invalid TE call, no TE with key '%s', ignoring: %s", te)
+            log.info("Encoders available for TE: %s", ", ".join(tokens.keys()))
+            continue
+        l = per_te_prompts.get(te, [])
+        l.append(prompt)
+        per_te_prompts[te] = l
 
-    if "g" in tokens and "l" in tokens and len(tokens["l"]) != len(tokens["g"]):
-        empty = clip.tokenize("", return_word_ids=need_word_ids)
-        while len(tokens["l"]) < len(tokens["g"]):
-            tokens["l"] += empty["l"]
-        while len(tokens["l"]) > len(tokens["g"]):
-            tokens["g"] += empty["g"]
+    if per_te_prompts:
+        for key in per_te_prompts:
+            prompt = " ".join(per_te_prompts[key])
+            tokens[key] = tokenize_chunks(clip, prompt, need_word_ids)[key]
+            log.info("Encoded prompt with TE '%s': %s", key, prompt)
+
+    maxlen = max(len(tokens[k]) for k in tokens)
+    empty = None
+    for k in tokens:
+        while len(tokens[k]) < maxlen:
+            if empty is None:
+                empty = clip.tokenize("", return_word_ids=need_word_ids)
+            tokens[k] += empty[k]
 
     tokens = fix_word_ids(tokens)
 
@@ -211,21 +229,28 @@ def apply_weights(output, te_name, spec):
     if te_name.startswith("clip_"):
         te_name = te_name[5:]
 
+    default = spec.get("all", None)
+
     if isinstance(output, tuple):
         out, pooled = output
-        if te_name in spec:
-            log.info("Weighting %s output by %s", te_name, spec[te_name])
-            out = out * spec[te_name]
         pkey = te_name + "_pooled"
-        if pkey in spec:
-            log.info("Weighting %s pooled output by %s", te_name, spec[pkey])
-            pooled = pooled * spec[pkey]
+        if te_name in spec or pkey in spec or default is not None:
+            w = spec.get(te_name, default)
+            pooled_w = spec.get(pkey, w)
+            if w is None:
+                w = 1.0
+            if pooled_w is None:
+                pooled_w = 1.0
+            log.info("Weighting %s output by %s, pooled by %s", te_name, w, pooled_w)
+            out = out * w
+            pooled = pooled * pooled_w
 
         return out, pooled
     else:
-        if te_name in spec:
-            log.info("Weighting %s output by %s", te_name, spec[te_name])
-            output = output * spec[te_name]
+        if te_name in spec or default is not None:
+            w = spec.get(te_name, default)
+            log.info("Weighting %s output by %s", te_name, w)
+            output = output * w
         return output
 
 
@@ -444,19 +469,25 @@ def encode_prompt(clip, text, start_pct, end_pct, defaults, masks):
     conds = []
     # TODO: is this still needed?
     # scale = sum(abs(weight(p)[0]) for p in prompts if not ("AREA(" in p or "MASK(" in p))
+    attnmasked_prompts = []
+    fill = False
     for prompt in prompts:
-        attn = False
+        attn_couple = False
+        prompt_has_fill = False
         if "ATTN()" in prompt:
             prompt = prompt.replace("ATTN()", "")
-            attn = True
-            log.info("Using attention masking for prompt segment")
+            attn_couple = True
+        if "FILL()" in prompt:
+            prompt = prompt.replace("FILL()", "")
+            prompt_has_fill = True
         prompt, mask, mask_weight = get_mask(prompt, mask_size, masks)
-        w, opts, prompt = weight(prompt)
         text, noise_w, generator = get_noise(text)
-        if not w:
-            continue
         prompt, area = get_area(prompt)
         prompt, local_sdxl_opts = get_sdxl(prompt, defaults)
+        # Get weight last so other syntax doesn't interfere with it
+        w, opts, prompt = weight(prompt)
+        if not w:
+            continue
         settings = {"prompt": prompt}
         settings["strength"] = w
         settings.update(sdxl_opts)
@@ -471,12 +502,42 @@ def encode_prompt(clip, text, start_pct, end_pct, defaults, masks):
 
         settings["start_percent"] = start_pct
         settings["end_percent"] = end_pct
-        x = encode_prompt_segment(clip, prompt, settings, style, normalization)
-        if attn and mask is not None:
-            mask = settings.pop("mask")
-            strength = settings.pop("mask_strength")
-            x = set_cond_attnmask(x, mask * strength)
 
-        conds.extend(x)
+        x = encode_prompt_segment(clip, prompt, settings, style, normalization)
+        if attn_couple:
+            if prompt_has_fill:
+                if attnmasked_prompts:
+                    log.warning("FILL() can only be used for the first prompt, ignoring")
+                elif mask is not None:
+                    log.warning("MASK() and FILL() can't be used together, ignoring FILL()")
+                else:
+                    fill = True
+            log.info("Using attention masking for prompt segment")
+            attnmasked_prompts.extend(x)
+        else:
+            conds.extend(x)
+
+    def ensure_mask(c):
+        if "mask" not in c[1]:
+            _, mask, _ = get_mask("MASK()", mask_size, masks)
+            c[1]["mask"] = mask
+            c[1]["mask_strength"] = 1.0
+        return c
+
+    if attnmasked_prompts:
+        base_cond = attnmasked_prompts[0]
+        if not fill:
+            ensure_mask(base_cond)
+        # else, set_cond_attnmask will have the base mask fill any unspecified areas
+        base_cond = [base_cond]
+        if len(attnmasked_prompts) > 1:
+            base_cond = set_cond_attnmask(
+                base_cond,
+                [ensure_mask(c) for c in attnmasked_prompts[1:]],
+                fill=fill,
+            )
+        else:
+            log.warning("You must specify at least two prompt segments with ATTN() for attention couple to work")
+        conds.extend(base_cond)
 
     return conds

prompt_control/test_graph.py

@@ -0,0 +1,215 @@
+import unittest
+import unittest.mock as mock
+import logging
+
+log = logging.getLogger("comfyui-prompt-control")
+
+
+def find_file(name):
+    names = {"test": "test.safetensors", "other": "some/other.safetensors"}
+    return names.get(name)
+
+
+def apply(cls, text, **kwargs):
+    model = [0, 1]
+    clip = [0, 0]
+    return cls().apply(unique_id="UID", model=model, clip=clip, text=text, **kwargs)
+
+
+@mock.patch("prompt_control.utils.lora_name_to_file", find_file)
+@mock.patch("torch.cuda.current_device", lambda: "cpu")
+class GraphTests(unittest.TestCase):
+    maxDiff = 4096
+
+    def test_textencode(self):
+        clip = [0, 0]
+        from .nodes_lazy import PCLazyTextEncode, PCLazyTextEncodeAdvanced
+
+        for p in ["test", "[test:0.2] test", "[test[test::0.5]]<lora:test:1>"]:
+            r1 = PCLazyTextEncode().apply(clip, p, "UID")
+            r2 = PCLazyTextEncodeAdvanced().apply(clip, p, "UID")
+            self.assertEqual(r1, r2)
+
+        r = PCLazyTextEncode().apply(clip, "test<lora:test:1>", "UID")
+        self.assertEqual(
+            r,
+            {
+                "result": (["UID-2", 0],),
+                "expand": {
+                    "UID-1": {"class_type": "PCTextEncode", "inputs": {"clip": [0, 0], "text": "test"}},
+                    "UID-2": {
+                        "class_type": "ConditioningSetTimestepRange",
+                        "inputs": {"conditioning": ["UID-1", 0], "start": 0.0, "end": 1.0},
+                    },
+                },
+            },
+        )
+        r = PCLazyTextEncode().apply(clip, "simple [test:0.1,0.5] prompt<lora:test:1>", "UID")
+        self.assertEqual(
+            r,
+            {
+                "result": (["UID-8", 0],),
+                "expand": {
+                    "UID-1": {"class_type": "PCTextEncode", "inputs": {"clip": [0, 0], "text": "simple  prompt"}},
+                    "UID-2": {
+                        "class_type": "ConditioningSetTimestepRange",
+                        "inputs": {"conditioning": ["UID-1", 0], "start": 0.0, "end": 0.1},
+                    },
+                    "UID-3": {"class_type": "PCTextEncode", "inputs": {"clip": [0, 0], "text": "simple test prompt"}},
+                    "UID-4": {
+                        "class_type": "ConditioningSetTimestepRange",
+                        "inputs": {"conditioning": ["UID-3", 0], "start": 0.1, "end": 0.5},
+                    },
+                    "UID-5": {"class_type": "PCTextEncode", "inputs": {"clip": [0, 0], "text": "simple  prompt"}},
+                    "UID-6": {
+                        "class_type": "ConditioningSetTimestepRange",
+                        "inputs": {"conditioning": ["UID-5", 0], "start": 0.5, "end": 1.0},
+                    },
+                    "UID-7": {
+                        "class_type": "ConditioningCombine",
+                        "inputs": {"conditioning_1": ["UID-2", 0], "conditioning_2": ["UID-4", 0]},
+                    },
+                    "UID-8": {
+                        "class_type": "ConditioningCombine",
+                        "inputs": {"conditioning_1": ["UID-7", 0], "conditioning_2": ["UID-6", 0]},
+                    },
+                },
+            },
+        )
+
+    @mock.patch("prompt_control.utils.lora_name_to_file", find_file)
+    def test_loraloader(self):
+        from .nodes_lazy import PCLazyLoraLoader, PCLazyLoraLoaderAdvanced
+
+        model = [0, 1]
+        clip = [0, 0]
+        with self.assertLogs(log, level="WARNING") as cm:
+            result = apply(PCLazyLoraLoader, "prompt here <lora:nonexistent:1.0:0.5>")["expand"]
+            result_adv = apply(PCLazyLoraLoaderAdvanced, "prompt here <lora:nonexistent:1.0:0.5>")["expand"]
+        self.assertIn("LoRA 'nonexistent' not found", cm.output[0])
+        self.assertEqual(result, {})
+        self.assertEqual(result_adv, {})
+
+        result = apply(PCLazyLoraLoader, "<lora:test:1>")["expand"]
+        result2 = apply(PCLazyLoraLoader, "prompt here <lora:test:1.0:0.5><lora:test:0:0.5>")["expand"]
+        result3 = apply(PCLazyLoraLoaderAdvanced, "prompt here <lora:test:1.0:0.5><lora:test:0:0.5>")["expand"]
+        self.assertEqual(result, result2)
+        self.assertEqual(result2, result3)
+        self.assertEqual(
+            result,
+            {
+                "UID-1": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": [0, 1],
+                        "clip": [0, 0],
+                        "strength_model": 1.0,
+                        "strength_clip": 1.0,
+                        "lora_name": "test.safetensors",
+                    },
+                }
+            },
+        )
+        result = apply(PCLazyLoraLoader, "<lora:test:1><lora:other:0.5>")["expand"]
+        self.assertEqual(
+            result,
+            {
+                "UID-1": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": [0, 1],
+                        "clip": [0, 0],
+                        "strength_model": 1.0,
+                        "strength_clip": 1.0,
+                        "lora_name": "test.safetensors",
+                    },
+                },
+                "UID-2": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": ["UID-1", 0],
+                        "clip": ["UID-1", 1],
+                        "strength_model": 0.5,
+                        "strength_clip": 0.5,
+                        "lora_name": "some/other.safetensors",
+                    },
+                },
+            },
+        )
+
+        result = apply(PCLazyLoraLoader, "prompt here <lora:test:1.0:0.5>")["expand"]
+        self.assertEqual(
+            result,
+            {
+                "UID-1": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": [0, 1],
+                        "clip": [0, 0],
+                        "strength_model": 1.0,
+                        "strength_clip": 0.5,
+                        "lora_name": "test.safetensors",
+                    },
+                }
+            },
+        )
+
+        result = apply(PCLazyLoraLoader, "prompt [<lora:test:0.5>:0.5]")["expand"]
+        result2 = apply(PCLazyLoraLoaderAdvanced, "prompt [<lora:test:0.5>:0.5]")["expand"]
+        self.assertEqual(result, result2)
+        expected = {
+            "UID-1": {
+                "class_type": "CreateHookLora",
+                "inputs": {"lora_name": "test.safetensors", "strength_model": 0.5, "strength_clip": 0.5},
+            },
+            "UID-2": {
+                "class_type": "CreateHookKeyframe",
+                "inputs": {"strength_mult": 0.0, "start_percent": 0.0},
+            },
+            "UID-3": {
+                "class_type": "CreateHookKeyframe",
+                "inputs": {
+                    "start_percent": 0.5,
+                    "prev_hook_kf": ["UID-2", 0],
+                    "strength_mult": 1.0,
+                },
+            },
+            "UID-4": {
+                "class_type": "SetHookKeyframes",
+                "inputs": {"hooks": ["UID-1", 0], "hook_kf": ["UID-3", 0]},
+            },
+            "UID-5": {
+                "class_type": "SetClipHooks",
+                "inputs": {
+                    "clip": [0, 0],
+                    "hooks": ["UID-4", 0],
+                    "apply_to_conds": True,
+                    "schedule_clip": True,
+                },
+            },
+        }
+        self.assertEqual(result, expected)
+        result2 = apply(PCLazyLoraLoaderAdvanced, "prompt [<lora:test:0.5>:0.5]", start=0.6)["expand"]
+        self.assertEqual(
+            result2,
+            {
+                "UID-1": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": [0, 1],
+                        "clip": [0, 0],
+                        "strength_model": 0.5,
+                        "strength_clip": 0.5,
+                        "lora_name": "test.safetensors",
+                    },
+                }
+            },
+        )
+        result2 = PCLazyLoraLoaderAdvanced().apply(model, clip, "prompt [<lora:test:0.5>:0.5]", "UID", end=0.5)[
+            "expand"
+        ]
+        self.assertEqual(result2, {})
+
+
+if __name__ == "__main__":
+    unittest.main()

prompt_control/test_parser.py

@@ -0,0 +1,202 @@
+import unittest
+from .parser import parse_prompt_schedules as parse
+
+
+def prompt(until, text, *loras):
+    loras = {lora: {"weight": unet, "weight_clip": te} for lora, unet, te in loras}
+    return [until, {"prompt": text, "loras": loras}]
+
+
+class TestParser(unittest.TestCase):
+    def assertPrompt(self, p, at, until, text, *loras):
+        self.assertEqual(p.at_step(at), prompt(until, text, *loras))
+
+    def test_no_scheduling(self):
+        p = parse("This is a (basic:0.6) (prompt) with [no scheduling] features")
+        expected = prompt(1.0, "This is a (basic:0.6) (prompt) with [no scheduling] features")
+        self.assertEqual(p.at_step(0), expected)
+        self.assertEqual(p.at_step(0.5), expected)
+        self.assertEqual(p.at_step(1), expected)
+
+    def test_equivalences(self):
+        eqs = [parse(p) for p in ["[a:0.1]", "[:a:0.1]", "[:a:0,0.1]", "[:a::0.1,1.0]", "[:a::0.1]"]]
+        for p in eqs[1:]:
+            self.assertEqual(eqs[0].parsed_prompt, p.parsed_prompt)
+
+        eqs = [parse(p) for p in ["[before:during:after:0.1]", "[before:during:after:0.1,1.0]", "[before:during:0.1]"]]
+        for p in eqs[1:]:
+            self.assertEqual(eqs[0].parsed_prompt, p.parsed_prompt)
+
+        eqs = [parse(p) for p in ["[a:0.1,0.5]", "[[a:0.1]::0.5]", "[:a::0.1,0.5]", "[a::0.1,0.5]"]]
+        for p in eqs[1:]:
+            self.assertEqual(eqs[0].parsed_prompt, p.parsed_prompt)
+
+        eqs = [parse(p) for p in ["[a:b:0.5]", "[a::b:0.5,0.5]"]]
+        for p in eqs[1:]:
+            self.assertEqual(eqs[0].parsed_prompt, p.parsed_prompt)
+
+        eqs = [parse(p) for p in ["[a::0.5]", "[a:::0.5,0.5]"]]
+        for p in eqs[1:]:
+            self.assertEqual(eqs[0].parsed_prompt, p.parsed_prompt)
+
+    def test_basic(self):
+        p = parse(
+            "This is a (basic:0.6) (prompt) with (very [[simple]:(basic:0.6):0.5]:1.1) [features::0.8][ and this is ignored:1]"
+        )
+        self.assertPrompt(p, 0, 0.5, "This is a (basic:0.6) (prompt) with (very [simple]:1.1) features")
+        self.assertPrompt(p, 0.5, 0.5, "This is a (basic:0.6) (prompt) with (very [simple]:1.1) features")
+        self.assertPrompt(p, 0.7, 0.8, "This is a (basic:0.6) (prompt) with (very (basic:0.6):1.1) features")
+        self.assertPrompt(p, 1.0, 1.0, "This is a (basic:0.6) (prompt) with (very (basic:0.6):1.1) ")
+
+    def test_lora(self):
+        p = parse("This is a (lora:0.6) (prompt) with [no scheduling] features <lora:foo:0.5> <lora:bar:0.5:1.0>")
+        expected = prompt(
+            1.0, "This is a (lora:0.6) (prompt) with [no scheduling] features  ", ("foo", 0.5, 0.5), ("bar", 0.5, 1.0)
+        )
+        self.assertEqual(p.at_step(0), expected)
+        self.assertEqual(p.at_step(0.5), expected)
+        self.assertEqual(p.at_step(1), expected)
+
+    def test_scheduled_lora(self):
+        p = parse(
+            "This is a (lora:0.6) (prompt) with [scheduling] features [<lora:foo:0.5>:<lora:bar:0.5:0.2>:0.3] <lora:bar:0.5:1.0>"
+        )
+        self.assertPrompt(
+            p,
+            0.1,
+            0.3,
+            "This is a (lora:0.6) (prompt) with [scheduling] features  ",
+            ("foo", 0.5, 0.5),
+            ("bar", 0.5, 1.0),
+        )
+        self.assertPrompt(p, 0.5, 1.0, "This is a (lora:0.6) (prompt) with [scheduling] features  ", ("bar", 1.0, 1.2))
+
+    def test_seq(self):
+        p = parse("This is a sequence of [SEQ:a:0.2::0.5:c:0.8][SEQ: and x:0.8]")
+        p2 = parse("This is a sequence of [[a:[c:0.5]:0.2]::0.8][ and x::0.8]")
+        prompts = {
+            0.2: "This is a sequence of a and x",
+            0.5: "This is a sequence of  and x",
+            0.8: "This is a sequence of c and x",
+            1.0: "This is a sequence of ",
+        }
+        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+        for k, v in prompts.items():
+            self.assertPrompt(p, k, k, v)
+
+    def test_shortcuts_scheduling(self):
+        p = parse("A schedule [a:0.1,0.7] b")
+        p2 = parse("A schedule [[a:0.1]::0.7] b")
+        p3 = parse("A schedule [a:b:0.5,0.8]")
+        p4 = parse("A schedule [[a:0.5]:b:0.8]")
+        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+        self.assertEqual(p3.parsed_prompt, p4.parsed_prompt)
+
+    def test_range(self):
+        p = parse("test [excluded::excluded2:0.1,0.4] test")
+        self.assertPrompt(p, 0, 0.1, "test excluded test")
+        self.assertPrompt(p, 0.2, 0.4, "test  test")
+        self.assertPrompt(p, 0.45, 1.0, "test excluded2 test")
+        p = parse("test [[:included::0.2,0.8]|[excluded::excluded2:0.4,0.9]:0.1] test")
+        self.assertPrompt(p, 0, 0.1, "test  test")
+        self.assertPrompt(p, 0.25, 0.3, "test included test")
+        self.assertPrompt(p, 0.15, 0.2, "test excluded test")
+        self.assertPrompt(p, 0.25, 0.3, "test included test")
+        self.assertPrompt(p, 0.55, 0.6, "test  test")
+        self.assertPrompt(p, 0.95, 1.0, "test excluded2 test")
+
+    def test_nested(self):
+        p = parse(
+            "This [prompt is [SEQ:[crazy:weird:0.2] stuff:0.5:<lora:cool:1>:0.7:nesting:1.0]:completely ignored with tags:HR]"
+        )
+        prompts = {
+            0.2: (0.2, "This prompt is crazy stuff"),
+            0.3: (0.5, "This prompt is weird stuff"),
+            0.5: (0.5, "This prompt is weird stuff"),
+            0.8: (1.0, "This prompt is nesting"),
+        }
+        for k in prompts:
+            self.assertEqual(p.at_step(k), [prompts[k][0], {"prompt": prompts[k][1], "loras": {}}])
+
+        self.assertPrompt(p, 0.6, 0.7, "This prompt is ", ("cool", 1.0, 1.0))
+        self.assertPrompt(p, 0.7, 0.7, "This prompt is ", ("cool", 1.0, 1.0))
+        p2 = p.with_filters(filters="hr, xyz")
+
+        self.assertEqual(p2.at_step(0), p2.at_step(1))
+
+    def test_def(self):
+        p = parse("DEF(X=0.5) [a:b:X] DEF(test = [c:X]) test test")
+        prompts = {
+            0.2: (0.5, "a   "),
+            0.6: (1.0, "b  c c"),
+        }
+        for k, v in prompts.items():
+            self.assertPrompt(p, k, v[0], v[1])
+
+        p = parse("DEF(X=[($1):($1:$2):$2])X(test;0.7)")
+        p2 = parse("[(test):(test:0.7):0.7]")
+        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+
+        p = parse("DEF(X=[($1):($1:$2):$2])DEF(Y=X(test;$1))Y(0.7) Y(0.5)")
+        p2 = parse("[(test):(test:0.7):0.7] [(test):(test:0.5):0.5]")
+        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+
+        p = parse("DEF(X(a;b)=$1 $2 $3 d)X(A) X(A;B;C)")
+        p2 = parse("A b $3 d A B C d")
+        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+
+        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]")
+        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+
+        with self.assertRaises(ValueError) as c:
+            parse("DEF(X=recurse Y) DEF(Y=recurse X) X")
+        self.assertTrue("Unable to resolve DEFs" in str(c.exception))
+
+    def test_misc(self):
+        p = parse("[[a:c:0.5]:0.7]")
+        p2 = parse("[:[a:c:0.5]:0.7]")
+        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+
+        p = parse("test [[a:[b<lora:test:0.5>:0.6]:0.5]:HR]")
+        p2 = parse("test [:[a:[:b<lora:test:0.5>:0.6]:0.5]:HR]")
+        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+
+        pf = p.with_filters(filters="hr")
+        self.assertEqual(pf.parsed_prompt, p2.with_filters(filters="hr").parsed_prompt)
+        self.assertPrompt(pf, 0, 0.5, "test a")
+        self.assertPrompt(pf, 0.55, 0.6, "test ")
+        self.assertPrompt(pf, 0.8, 1.0, "test b", ("test", 0.5, 0.5))
+
+        p = parse("[:[<lora:test:1>:c:0.5]:0.3]")
+        self.assertPrompt(p, 0, 0.3, "")
+        self.assertPrompt(p, 0.4, 0.5, "", ("test", 1.0, 1.0))
+        self.assertPrompt(p, 1.0, 1.0, "c")
+
+        p = parse("an [<emb:foo>:<emb:bar>:0.5]")
+        prompts = {
+            0.2: (0.5, "an embedding:foo"),
+            0.8: (1.0, "an embedding:bar"),
+        }
+        for k, v in prompts.items():
+            self.assertPrompt(p, k, v[0], v[1])
+
+    def test_alternating(self):
+        p = parse("[cat|dog|tiger]")
+        p2 = parse("[cat|dog|tiger:0.1]")
+        p3 = parse("[cat|[dog|wolf]|tiger]")
+        p4 = parse("[cat|[dog:wolf<lora:canine:1>:0.5]:0.2]")
+
+        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+        for i, x in enumerate(["cat", "wolf", "tiger", "cat", "dog", "tiger", "cat", "wolf", "tiger", "cat"]):
+            step = round((i * 0.1) + 0.1, 2)
+            self.assertPrompt(p3, step, step, x)
+
+        for i, x in enumerate([["cat"], ["dog"], ["cat"], ["wolf", ("canine", 1.0, 1.0)], ["cat"]]):
+            step = round((i * 0.2) + 0.2, 2)
+            self.assertPrompt(p4, step, step, *x)
+        self.assertPrompt(p4, 0.7, 0.8, "wolf", ("canine", 1.0, 1.0))
+
+
+if __name__ == "__main__":
+    unittest.main()

prompt_control/utils.py

@@ -87,10 +87,11 @@ def find_closing_paren(text, start):
     return len(text)
 
 
-def get_function(text, func, defaults, return_func_name=False):
+def get_function(text, func, defaults, return_func_name=False, placeholder=""):
     rex = re.compile(rf"\b{func}\(", re.MULTILINE)
     instances = []
     match = rex.search(text)
+    count = 0
     while match:
         # Match start, content start
         start, after_first_paren = match.span()
@@ -102,8 +103,12 @@ def get_function(text, func, defaults, return_func_name=False):
         else:
             instances.append(args)
 
-        text = text[:start] + text[end + 1 :]
+        if placeholder:
+            text = text[:start] + f"\0{placeholder}{count}\0" + text[end + 1 :]
+        else:
+            text = text[:start] + text[end + 1 :]
         match = rex.search(text)
+        count += 1
     return text, instances
 
 

pyproject.toml

@@ -1,14 +1,13 @@
 [project]
 name = "comfyui-prompt-control"
 description = "Nodes for convenient prompt editing, making many common operations prompt-controllable"
-version = "2.0.0-beta.5"
+version = "2.0.0-rc.3"
 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"]
 
 [project.urls]
 Repository = "https://github.com/asagi4/comfyui-prompt-control"
-#  Used by Comfy Registry https://comfyregistry.org
 
 [tool.comfy]
 PublisherId = "asagi4"