v2.1.0

This should prevent things from breaking, but it needs a bit of testing.

.github/ISSUE_TEMPLATE/bug_report.md

@@ -20,5 +20,7 @@ A clear and concise description of what the bug is.
 Information needed to trigger the problem.
 If possible, attach a workflow to reproduce the problem
 
+If a workflow works, but isn't producing the correct output, please enable debug logging with the `PCSetLogLevel` node (from `promptcontrol/tools`) and run your workflow with debug logging enabled, and copy the outputs here.
+
 **Expected behavior**
 A description of what you expected to happen.

.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,42 @@
+name: Run tests requiring ComfyUI
+on:
+  workflow_call:
+  workflow_dispatch:
+  push:
+    paths:
+      - prompt_control/adv_encode.py
+      - prompt_control/attention_couple_ppm.py
+      - prompt_control/nodes_lazy.py
+      - prompt_control/prompts.py
+      - prompt_control/parser.py
+      - prompt_control/utils.py
+
+
+jobs:
+  run-graph-tests:
+    name: Run tests requiring ComfyUI
+    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'
+          cache: pip
+      - name: install-torch
+        run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+      - name: install ComfyUI
+        run: pip install -r requirements.txt -r ComfyUI/requirements.txt
+      - name: Download clip_l.safetensors
+        run: curl -LO https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors
+      - name: Force Comfy to use the CPU
+        run: sed -i "s/^cpu_state = CPUState.GPU/cpu_state = CPUState.CPU/g" ComfyUI/comfy/model_management.py
+      - name: Run graph tests
+        run: PYTHONPATH=ComfyUI python -m prompt_control.test_graph
+      - name: Run encoder tests (clip_l only)
+        run: PYTHONPATH=ComfyUI python -m prompt_control.test_encode

Makefile

@@ -1,8 +1,25 @@
-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
+
+test_encode:
+	PYTHONPATH=../../ python -m prompt_control.test_encode --verbose
+
+test_encode_both:
+	TEST_TE="clip_l t5" PYTHONPATH=../../ python -m prompt_control.test_encode --verbose
+
+test_heavy: test_graph test_encode_both
+
+manual_test:
+	PYTHONPATH=../../ python -im prompt_control.manual_test
+
 .PHONY: check format all

README.md

@@ -1,54 +1,36 @@
 # ComfyUI prompt control
 
-Nodes for LoRA and prompt scheduling that make basic operations in ComfyUI completely prompt-controllable.
+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 v2
+Prompt Control comes with `PCTextEncode`, which provides advanced text encoding with many additional features compared to ComfyUI's base `CLIPTextEncode`.
 
-Prompt control has been almost completely rewritten. It now uses ComfyUI's lazy execution to build graphs from the text prompt at runtime. This has some advantages:
-
-- ComfyUI will not re-run unchanged parts of generated graphs. This is especially useful for two-pass workflows where previously you'd be forced to re-run the first sampling pass even with filtering. That is no longer the case and it does the right thing.
-- 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, probably.
-
-### Everything broke, where are the old nodes?
-
-If you really need them, check out the [legacy branch](https://github.com/asagi4/comfyui-prompt-control/tree/legacy). However, I will not fix bugs in that branch, 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.
+A `Basic Text to Image` template is included with the extension, and can be loaded from ComfyUI's template library.
 
 ## 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
+You can use text prompts to control the following:
 
-See the [syntax documentation](doc/syntax.md)
+- A1111-style prompt scheduling and filtering without noodle soup.
+- LoRA loading and [scheduling](/doc/schedules.md) via the prompt, using ComfyUI's hook system
+- Masking, composition and area control ([regional prompting](/doc/regional_prompts.md)) with an implementation of [Attention Couple](/doc/attention_couple.md), also fully schedulable.
+- [Advanced prompt encoding](/doc/basic.md)
+  - Per-encoder prompts for models with multiple text encoders, such as SDXL and Flux
+  - Prompt combinators like `BREAK`, as well as `CAT`, `AVG()` and `AND` corresponding to ComfyUI's `ConditioningConcat`, `ConditioningAverage` and `ConditioningCombine` nodes.
+  - Different weight interpretation types (ComfyUI, A1111, compel, etc.)
+  - Prompt masking with an implementation of [cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff)
+- Simple [prompt macros](/doc/macros.md) with `DEF`
+
+All features are fully schedulable unless otherwise stated. See the [scheduling syntax documentation](doc/schedules.md) to get started.
 
 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.
 
-[This example 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.
+## Compatibility
 
-[This example workflow](workflows/example.json?raw=1) implements 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)
+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.
 
 ## Requirements
 
@@ -65,23 +47,25 @@ Then restart ComfyUI afterwards.
 
 # Core nodes
 
-## PCLazyTextEncode and PCLazyTextEncodeAvanced
+**Note**: The documentation refers to the nodes with their internal names for consistency. The display name may change, but ComfyUI's search will always find the nodes with the internal name. `PCLazyTextEncode` and `PCLazyLoraLoader` are the main ones you'll want to use, also known as `PC: Schedule Prompt` and `PC: Schedule LoRas`.
+
+## PCLazyTextEncode and PCLazyTextEncodeAdvanced
 
 `PCLazyTextEncode` uses ComfyUI's lazy graph execution mechanism to generate a graph of `PCTextEncode` and `SetConditioningTimestepRange` nodes from a prompt with schedules. This has the advantage that if a part of the schedule doesn't change, ComfyUI's caching mechanism allows you to avoid re-encoding the non-changed part.
 
 for example, if you first encode `[cat:dog:0.1]` and later change that to `[cat:dog:0.5]`, no re-encoding takes place.
 
-for added fun, put `NODE(NodeClassName, paramname)` in a prompt to generate a graph using **any other node** that's compatible. The node can't have required parameters besides a single CLIP parameter (which must be named `clip`) and the text prompt, and it must return a `CONDITIONING` as its first return value.
+for added fun, put `NODE(NodeClassName, textinputname)` in a prompt to generate a graph using **any other node** that's compatible. The node can't have required parameters besides a single CLIP parameter (which must be named `clip`) and the text prompt, and it must return a `CONDITIONING` as its first return value. The "default" values are `PCTextEncode` and `text`.
+
+For example, if you for some reason do not want the advanced features of `PCTextEncode`, use `NODE(CLIPTextEncode)` in the prompt and you'll still get scheduling with ComfyUI's regular TE node.
 
 The advanced node enables filtering the prompt for multi-pass workflows.
 
 ## PCLazyLoraLoader and PCLazyLoraLoaderAdvanced
 
-This node reads LoRA expressions from the scheduled prompt and constructs a graph of `LoraLoader`s and `CreateHookLora`s as necessary to provide the necessary LoRA scheduling.
+This node reads LoRA expressions from the scheduled prompt and constructs a graph of `LoraLoader`s and `CreateHookLora`s as necessary to provide the necessary LoRA scheduling. Just use it in place of a `LoRALoader` and use the output normally.
 
-If you have `apply_hooks` set to true, you **do not** need to apply the `HOOKS`  output to a CLIP model separately; it's provided in case you want to use it elsewhere.
-
-The advanced node enables filtering the prompt for multi-pass workflows.
+The Advanced node gives you access to the generated hooks. If you have `apply_hooks` set to true, you **do not** need to apply the `HOOKS`  output to a CLIP model separately; it's provided in case you want to use it elsewhere. The advanced node also enables filtering the prompt for multi-pass workflows.
 
 ## PCTextEncode
 
@@ -97,75 +81,8 @@ 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
 
-- None at the moment
+- ComfyUI's caching mechanism has an issue that makes it unnecessarily invalidate caches for certain inputs; you'll still get some benefit from the lazy nodes, but changing inputs that shouldn't affect downstream nodes (especially if using filtering) will still cause them to be recomputed because ComfyUI doesn't realize the inputs haven't changed.
+
+- Cutoff does not work with models that use non-CLIP text encoders, like Flux. This might be fixable, but it's uncertain if cutoff even makes sense for those models.

__init__.py

@@ -1,3 +1,10 @@
+"""
+@author: asagi4
+@title: ComfyUI Prompt Control
+@nickname: ComfyUI Prompt Control
+@description: 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.
+"""
+
 import os
 import sys
 import logging
@@ -8,7 +15,7 @@ log = logging.getLogger("comfyui-prompt-control")
 log.propagate = False
 if not log.handlers:
     h = logging.StreamHandler(sys.stdout)
-    h.setFormatter(logging.Formatter("[%(levelname)s] PromptControl: %(message)s"))
+    h.setFormatter(logging.Formatter("[PromptControl] %(levelname)s: %(message)s"))
     log.addHandler(h)
 
 if os.environ.get("PROMPTCONTROL_DEBUG"):
@@ -16,17 +23,12 @@ if os.environ.get("PROMPTCONTROL_DEBUG"):
 else:
     log.setLevel(logging.INFO)
 
-
 NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}
 
-nodes = ["base", "lazy", "tools"]
-if importlib.util.find_spec("comfy.hooks"):
-    nodes.append("hooks")
-else:
-    log.warning(
-        "Your ComfyUI version is too old, can't import comfy.hooks for PCEncodeSchedule and PCLoraHooksFromSchedule. Update your installation."
-    )
+WEB_DIRECTORY = "web"
+
+nodes = ["base", "lazy", "tools", "hooks"]
 
 for node in nodes:
     mod = importlib.import_module(f".prompt_control.nodes_{node}", package=__name__)

doc/attention_couple.md

@@ -0,0 +1,40 @@
+# Attention Couple
+
+NOTE: This is still considered an experimental feature, so the syntax may change.
+
+Attention Couple is an attention-based implementation of regional prompting. it is faster and often more flexible than latent-based masking.
+
+The implementation is based on the one by [pamparamm](https://github.com/pamparamm/ComfyUI-ppm.git), modified to use ComfyUI's hook system. This enables it to work with prompt scheduling.
+
+By default, the implementation produces slightly different results from Pamparamm's implementation because ComfyUI will only run the hook for conds that have it attached and can't batch negative conditionings.
+
+As a consequence of this, however, you can also use `COUPLE` in your negative prompt, and it will work correctly.
+
+To enable batching negative prompts, run your positive and negative prompt through the `PPCAttentionCoupleBatchNegative` node. This will make the outputs identical to pamparamm's implementation and will also improve performance. It will fall back to the default behaviour in cases where batching can't be done, so it should always be safe to use.
+
+
+## Syntax
+
+See also the main syntax documentation for `MASK` etc.
+
+### COUPLE: Trigger Attention Couple
+
+You can use `COUPLE` to attach attention-coupled prompts to a base prompt:
+
+`base_prompt COUPLE MASK(0 0.5) coupled prompt 1 with mask COUPLE IMASK(0) coupled prompt 2 with custom mask`
+
+as a shortcut, `COUPLE(maskparams)` is expanded to `COUPLE MASK(maskparams)`, so the above prompt can also be written as:
+
+`base_prompt COUPLE(0 0.5) coupled prompt 1 with mask COUPLE IMASK(0) coupled prompt 2 with custom mask`
+
+Behaviour:
+- If no mask is specified, an implicit `MASK()` is assumed.
+
+- For the base prompt, you can also use `FILL()` to automatically mask all parts not masked by coupled prompts
+
+- If the base prompt has weight set to zero (ie. ´:0` at the end), then the first coupled prompt with non-zero weight becomes the base prompt.
+
+For example:
+```
+dog FILL() COUPLE(0.5 1) cat
+```

doc/basic.md

@@ -0,0 +1,206 @@
+# Basic Prompt Syntax
+
+The syntax below documents the features of `PCTextEncode`
+
+## Combining prompts
+
+### 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, see `MASK` and `COUPLE` below.
+
+Prompts can have a weight at the end:
+```
+cat :1 AND dog :2
+```
+`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:
+
+```
+cat [\:0::0.5] AND dog
+```
+Note that the `:` needs to be escaped with a `\` or it will be interpreted as scheduling syntax.
+
+## Note about processing order
+
+Prompt operators are processed in the following order, meaning that all features "below" another can be affected by the feature above it. That is, `BREAK` can go inside a `TE()` call, but not `AND` or `CAT`.
+
+- DEF macros are expanded
+- Scheduling is expanded, and for each scheduled prompt:
+  - The prompt is split by AND, and for each:
+    - Prompts are split by COUPLE. and for each:
+      - Most functions (like MASK) and cutoffs are evaluated
+      - prompts are split by `AVG()` or CAT
+        - the TE() function is evaluated to set per-encoder prompts
+        - BREAK is evaluated
+        - Everything else
+      - Prompts are combined with `ConditioningAverage` (for `AVG`) or `ConditioningConcat` (for `CAT`)
+    - If coupled prompts exist, the base cond is set up for attention coupling and returned
+  - Prompts split with `AND` are combined with `ConditioningCombine`
+- Each scheduled prompt is restricted to its effective range with  `ConditioningSetTimestepRange`
+
+## Functions
+
+There are some "functions" that can be included in a prompt to affect how it is interpreted.
+
+Functions have the form `FUNCNAME(param1, param2, ...)`. How parameters are interpreted is up to the function. 
+
+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.
+
+### BREAK
+The keyword `BREAK` causes the prompt to be tokenized in separate chunks, padding each chunk to the text encoder's maximum size before encoding.
+
+For some text encoders (like t5), this operation doesn't really make sense and BREAKs are simply ignored.
+
+### CAT
+
+`CAT` encodes each prompt separately before concatenating the resulting tensors into a single conditioning. It behaves identically to ComfyUI's `ConditioningConcat`.
+
+### AVG()
+
+`prompt1 AVG(weight) prompt2` encodes prompt1 and prompt2 separately, and then combines them using `ConditioningAverage`. The default for `weight` is `0.5`.
+
+`AVG` is processed before `BREAK` but after `AND`
+
+`p1 AVG() p2 AVG() p3` combines `p1` and `p2` first, then combines the result with `p3`.
+
+## Prompt weighting (also known as "Advanced CLIP Encode")
+
+### STYLE
+
+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`.
+
+### 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:
+- 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` and `BREAK` are processed before `TE`, so they do not do anything sensible; `TE(l=foo AND bar)` will parse as two prompts `TE(foo` and `bar)`. `SHIFT`, `SHUFFLE` and `OLDBREAK` do work, however.
+
+### SHUFFLE and SHIFT: Create prompt permutations
+
+Default parameters: `SHUFFLE(seed=0, separator=,, joiner=,)`, `SHIFT(steps=0, separator=,, joiner=,)`
+
+`SHIFT` moves elements to the left by `steps`. The default is 0 so `SHIFT()` does nothing
+`SHUFFLE` generates a random permutation with `seed` as its seed.
+
+These functions are applied to each prompt chunk **after** `BREAK`, `AND` etc. have been parsed. The prompt is split by `separator`, the operation is applied, and it's then joined back by `joiner`.
+
+Multiple instances of these functions are applied in the order they appear in the prompt.
+
+**NOTE** To avoid breaking emphasis syntax, the functions ignore any separators inside parentheses
+
+For example:
+- `SHIFT(1) cat, dog, tiger, mouse` does a shift and results in `dog, tiger, mouse, cat`. (whitespace may vary)
+- `SHIFT(1,;) cat, dog ; tiger, mouse` results in `tiger, mouse, cat, dog`
+- `SHUFFLE() cat, dog, tiger, mouse` results in `cat, dog, mouse, tiger`
+- `SHUFFLE() SHIFT(1) cat, dog, tiger, mouse` results in `dog, mouse, tiger, cat`
+
+- `SHIFT(1) cat,dog BREAK tiger,mouse` results in `dog,cat BREAK tiger,mouse`
+- `SHIFT(1) cat, dog AND SHIFT(1) tiger, mouse` results in `dog, cat BREAK mouse, tiger`
+
+Whitespace is *not* stripped and may also be used as a joiner or separator
+- `SHIFT(1,, ) cat,dog` results in `dog cat`
+
+### NOISE: Add noise to a prompt
+
+The function `NOISE(weight, seed)` adds some random noise into the cond tensor. The seed is optional, and if not specified, the global RNG is used. `weight` should be between 0 and 1.
+
+The usefulness of this is questionable, but it wasn't difficult to implement, so here it is.
+
+## Regional prompting
+
+See [Regional prompting](/doc/regional_prompting.md)
+
+## 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
+
+> [!WARN]
+> These features are may change or disappear without warning
+
+## COUPLE: Attention couple
+
+See [here](/doc/attention_couple.md)
+
+## 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 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)`

doc/macros.md

@@ -0,0 +1,60 @@
+## DEF: Lightweight prompt macros
+
+You can define "prompt macros" by using `DEF`.  Macros are expanded before any other parsing takes place. The expansion continues until no further changes occur. Recursion will raise an error.
+
+`PCLazyTextEncode` and `PCLazyLoraLoader` expand macros, but `PCTextEncode` **does not**. If you need to expand macros for a single prompt, use `PCMacroExpand`
+
+```
+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]
+```
+### Macro parameters
+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, and can be empty.
+
+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]
+```
+
+```
+DEF(MACRO() = [a:$1:0.5])
+```
+sets the default value of `$1` to an empty string.
+
+### Unspecified parameters in macros
+
+Unspecified parameters (either via defaults or explicitly given) will not be substituted. Compare:
+
+```
+DEF(mything=a "$1" b "$2")
+mything
+mything()
+mything(A)
+```
+
+gives
+
+```
+a "$1" b "$2"
+a "" b "$2"
+a "A" b "$2"
+```

doc/regional_prompts.md

@@ -0,0 +1,63 @@
+# Regional prompting
+
+This section documents the masking functionality of `PCTextEncode`
+
+See also [Attention Couple](/doc/attention_couple.md)
+
+Remember that when using the lazy nodes, prompt scheduling applies to masks as well, so you can change or enable/disable regional prompts at any point during sampling.
+
+## Behaviour
+
+For each prompt separated by `AND`, you can specify either latent masks or an area.
+
+- When masked, ComfyUI generates the model output using the **full latent** as the input, and then applies the mask to the output before adding it to your latent for the next step.
+- When an area is specified, ComfyUI generates a separate model output using the **part of the latent specified by the area** and then composites it into the full latent afterwards.
+- You can have *both* an AREA and a MASK specified, in which case the mask is applied to the latent specified by the AREA.
+
+For example, consider a 1024 by 1024 (width x height) generation:
+
+- `cat MASK(0 0.5, 0 1) AND dog MASK(0.5 1, 0 1)` generates two outputs at 1024x1024 for "dog" and "cat", then masks half of them off and adds the results together. The following step still see both the dog and the cat from the previous step, so they may blend slightly.
+
+- `cat AREA(0 0.5, 0 1) AND dog AREA(0.5 1, 0 1)` generates two completely separate outputs at **512**x1024 and then composites them together into the 1024x1024 latent. Because the areas do not overlap, the generation for `cat` will not see the output of `dog` and vice versa in subsequent steps as long as the area restriction is in effect.
+
+## 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`
+
+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 multiple 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,
+
+These are handled per `AND`-ed prompt, so in `prompt1 AND MASK(...) prompt2`, the mask will only affect prompt2.
+
+The default values are `MASK(0 1, 0 1, 1)` and you can omit unnecessary ones, that is, `MASK(0 0.5, 0.3)` is `MASK(0 0.5, 0.3 1, 1)`
+
+Note that because the default values are percentages, `MASK(0 256, 64 512)` is valid, but `MASK(0 200)` will raise an error.
+
+Masking does not affect LoRA scheduling unless you set unet weights to 0 for a LoRA.
+
+## 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`.
+
+If multiple masks are used, `FEATHER` is applied *before compositing* in the order they appear in the prompt, and any leftovers are applied to the combined mask. If you want to skip feathering a mask while compositing, just use `FEATHER()` with no arguments.
+
+For example:
+```
+MASK(1) MASK(2) MASK(3) FEATHER(1) FEATHER() FEATHER(3) weirdmask FEATHER(4)
+```
+
+gives you a mask that is a combination of 1, 2 and 3, where 1 and 3 are feathered before compositing and then `FEATHER(4)` is applied to the composite.
+
+The order of the `FEATHER` and `MASK` calls doesn't matter; you can have `FEATHER` before `MASK` or even interleave them.

doc/schedules.md

@@ -0,0 +1,117 @@
+# Prompt Schedule Syntax
+
+> [!TIP]
+> 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).
+
+> [!NOTE]
+> The syntax documented in this section is only available with the `PC: Schedule Prompt` and `PC: Schedule LoRAs` nodes and their advanced variants.
+
+Scheduling syntax is available with is similar to A1111, but only fractions are supported for steps. LoRAs are scheduled by including them in a scheduling expression.
+
+Besides the syntax documented below, the [basic syntax](/doc/basic.md) and [prompt macro](/doc/macros.md) features are also automatically available.
+
+```
+a [large::0.1] [cat|dog:0.05] [<lora:somelora:0.5:0.6>::0.5]
+[in a park:in space:0.4]
+```
+## Comments and escaping
+
+In schedules, any text on a line following a `#` is considered a comment and removed, including the `#` character.
+You can escape the following characters in places where they would otherwise conflict with syntax:
+
+- `#` with `\#`
+- `:` with `\:`
+- `\` with `\\`
+
+Escaping is only required if it would otherwise be considered syntax, that is `\o/` will be interpreted literally and the `\` does not need to be escaped, but in `[embedding:a:0.5]` you would need to escape the `:`.
+
+## Scheduled prompts
+
+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. Either prompt can also be empty.
+For example:
+```
+a [red:blue:0.5] cat
+```
+switches from `a red cat` to `a blue cat` at 0.5. `before` and `after` can be arbitrary prompts (`after` can also be empty), including other scheduling expressions, allowing nesting:
+```
+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
+
+For convenience `[cat:0.5]` is equivalent to `[:cat:0.5]` meaning it switches from empty to `cat` at 0.5.
+
+### Range expressions
+
+The most general form of a schedule is a range expression: For example, in `prompt [before:during:after:0.3,0.7]`, The prompt be `prompt before` until 0.3, `prompt during` until 0.7, and then `prompt after`. This form is equivalent to `prompt [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:
+```
+a large [dog:cat<lora:catlora:0.5>:SECOND_PASS]
+```
+Set the `tags` parameter in the `FilterSchedule` node to filter the prompt. If the tag matches any tag `tags` (comma-separated), the second option is returned (`cat`, in this case, with the LoRA). Otherwise, the first option is chosen (`dog`, without LoRA).
+
+the values in `tags` are case-insensitive, but the tags in the input **must** be uppercase A-Z and underscores only, or they won't be recognized. That is, `[dog:cat:hr]` will not work.
+
+For example, a prompt
+```
+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:
+
+`<lora:fulllora:1> [<lora:partialora:1>::0.5]`
+
+This will schedule `fulllora` for the entire duration of the prompt and `partiallora` until half of sampling is complete.
+
+You can refer to LoRAs by using the filename without extension and subdirectories will also be searched. For example, `<lora:cats:1>`. will match both `cats.safetensors` and `sd15/animals/cats.safetensors`. If there are multiple LoRAs with the same name, the first match will be loaded.
+
+Alternatively, the name can include the full directory path relative to ComfyUI's search paths, without extension: `<lora:XL/sdxllora:0.5>`. In this case, the *full* path must match.
+
+You can also give the exact path (including the extension) as shown in `LoRALoader`.
+
+If no match is found, the node will try to replace spaces with underscores and search again. That is, `<lora:cats and dogs:1>` will find `cats_and_dogs.safetensors`. This helps with some autocompletion scripts that replace underscores with spaces.
+
+Finally, if none of the above produce a match, the search term will be split by whitespace and files that contain all of the parts in any order will be considered. If this returns only a single match, it will be loaded. For example, consider LoRAs:
+
+- `xl/red_cats.safetensors`
+- `flux/blue_cats.safetensors`
+- `flux/red_cats.safetensors`
+
+Then `<lora:cats xl:1>` would match the red cats LoRA, but `cats flux` would be ambiguous and not match.
+
+## Alternating
+
+Alternating syntax is `[a|b:pct_steps]`, causing the prompt to alternate every `pct_steps`. `pct_steps` defaults to 0.1 if not specified. You can also have more than two options.
+
+
+## Sequences
+
+The syntax `[SEQ:a:N1:b:N2:c:N3]` is shorthand for `[a:[b:[c::N3]:N2]:N1]` ie. it switches from `a` to `b` to `c` to nothing at the specified points in sequence.
+
+Might be useful with Jinja templating (see https://github.com/asagi4/comfyui-utility-nodes). For example:
+```
+[SEQ<% for x in steps(0.1, 0.9, 0.1) %>:<lora:test:<= sin(x*pi) + 0.1 =>>:<= x =><% endfor %>]
+```
+generates a LoRA schedule based on a sinewave

doc/syntax.md

@@ -1,197 +0,0 @@
-# Scheduling syntax
-
-Syntax is like A1111 for now, 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]
-[in a park:in space:0.4]
-```
-
-## Scheduled prompts
-
-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.
-For example:
-```
-a [red:blue:0.5] cat
-```
-switches from `a red cat` to `a blue cat` at 0.5. `before` and `after` can be arbitrary prompts (`after` can also be empty), including other scheduling expressions, allowing nesting:
-```
-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
-
-### 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]`
-
-## Tag selection
-Using the `FilterSchedule` node, in addition to step percentages, you can use a *tag* to select part of an input:
-```
-a large [dog:cat<lora:catlora:0.5>:SECOND_PASS]
-```
-Set the `tags` parameter in the `FilterSchedule` node to filter the prompt. If the tag matches any tag `tags` (comma-separated), the second option is returned (`cat`, in this case, with the LoRA). Otherwise, the first option is chosen (`dog`, without LoRA).
-
-the values in `tags` are case-insensitive, but the tags in the input **must** be uppercase A-Z and underscores only, or they won't be recognized. That is, `[dog:cat:hr]` will not work.
-
-For example, a prompt
-```
-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`
-
-
-## LoRA Scheduling
-When using the lazy graph building nodes, LoRAs can be scheduled by referring to them in a scheduling expression, like so:
-
-`<lora:fulllora:1> [<lora:partialora:1>::0.5]`
-
-This will schedule `fulllora` for the entire duration of the prompt and `partiallora` until half of sampling is complete.
-
-You can refer to LoRAs by using the filename without extension and subdirectories will also be searched. For example, `<lora:cats:1>`. will match both `cats.safetensors` and `sd15/animals/cats.safetensors`. If there are multiple LoRAs with the same name, the first match will be loaded.
-
-Alternatively, the name can include the full directory path relative to ComfyUI's search paths, without extension: `<lora:XL/sdxllora:0.5>`. In this case, the *full* path must match.
-
-If no match is found, the node will try to replace spaces with underscores and search again. That is, `<lora:cats and dogs:1>` will find `cats_and_dogs.safetensors`. This helps with some autocompletion scripts that replace underscores with spaces.
-
-Finally, you can give the exact path (including the extension) as shown in `LoRALoader`.
-
-
-## Alternating
-
-Alternating syntax is `[a|b:pct_steps]`, causing the prompt to alternate every `pct_steps`. `pct_steps` defaults to 0.1 if not specified. You can also have more than two options.
-
-
-## Sequences
-
-The syntax `[SEQ:a:N1:b:N2:c:N3]` is shorthand for `[a:[b:[c::N3]:N2]:N1]` ie. it switches from `a` to `b` to `c` to nothing at the specified points in sequence.
-
-Might be useful with Jinja templating (see https://github.com/asagi4/comfyui-utility-nodes). For example:
-```
-[SEQ<% for x in steps(0.1, 0.9, 0.1) %>:<lora:test:<= sin(x*pi) + 0.1 =>>:<= x =><% endfor %>]
-```
-generates a LoRA schedule based on a sinewave
-
-
-# Basic prompt syntax
-
-This syntax is also available in outside scheduled prompts, where applicable.
-
-## LoRA loading
-
-The A111-style syntax `<lora:loraname:weight>` can be used to load LoRAs via the prompt. See LoRA scheduling above.
-
-## 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.
-
-`AND` can be used to combine prompts. You can also use a weight at the end. It does a weighted sum of each prompt,
-```
-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`
-
-
-## Functions
-
-There are some "functions" that can be included in a prompt to do various things. 
-
-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 `\,`
-
-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.
-
-### SDXL
-
-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)`.
-
-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`.
-
-### SHUFFLE and SHIFT
-
-Default parameters: `SHUFFLE(seed=0, separator=,, joiner=,)`, `SHIFT(steps=0, separator=,, joiner=,)`
-
-`SHIFT` moves elements to the left by `steps`. The default is 0 so `SHIFT()` does nothing
-`SHUFFLE` generates a random permutation with `seed` as its seed.
-
-These functions are applied to each prompt chunk **after** `BREAK`, `AND` etc. have been parsed. The prompt is split by `separator`, the operation is applied, and it's then joined back by `joiner`.
-
-Multiple instances of these functions are applied in the order they appear in the prompt.
-
-**NOTE:** These functions are *not* smart about syntax and will break emphasis if the separator occurs inside parentheses. I might fix this at some point, but for now, keep this in mind.
-
-For example:
-- `SHIFT(1) cat, dog, tiger, mouse` does a shift and results in `dog, tiger, mouse, cat`. (whitespace may vary)
-- `SHIFT(1,;) cat, dog ; tiger, mouse` results in `tiger, mouse, cat, dog`
-- `SHUFFLE() cat, dog, tiger, mouse` results in `cat, dog, mouse, tiger`
-- `SHUFFLE() SHIFT(1) cat, dog, tiger, mouse` results in `dog, mouse, tiger, cat`
-
-- `SHIFT(1) cat,dog BREAK tiger,mouse` results in `dog,cat BREAK tiger,mouse`
-- `SHIFT(1) cat, dog AND SHIFT(1) tiger, mouse` results in `dog, cat BREAK mouse, tiger`
-
-Whitespace is *not* stripped and may also be used as a joiner or separator
-- `SHIFT(1,, ) cat,dog` results in `dog cat`
-
-### NOISE
-
-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.
-
-### 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.
-
-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`
-
-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
-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,
-
-These are handled per `AND`-ed prompt, so in `prompt1 AND MASK(...) prompt2`, the mask will only affect prompt2.
-
-The default values are `MASK(0 1, 0 1, 1)` and you can omit unnecessary ones, that is, `MASK(0 0.5, 0.3)` is `MASK(0 0.5, 0.3 1, 1)`
-
-Note that because the default values are percentages, `MASK(0 256, 64 512)` is valid, but `MASK(0 200)` will raise an error.
-
-Masking does not affect LoRA scheduling unless you set unet weights to 0 for a LoRA.
-
-### FEATHER
-
-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`.
-
-If multiple masks are used, `FEATHER` is applied *before compositing* in the order they appear in the prompt, and any leftovers are applied to the combined mask. If you want to skip feathering a mask while compositing, just use `FEATHER()` with no arguments.
-
-For example:
-```
-MASK(1) MASK(2) MASK(3) FEATHER(1) FEATHER() FEATHER(3) weirdmask FEATHER(4)
-```
-
-gives you a mask that is a combination of 1, 2 and 3, where 1 and 3 are feathered before compositing and then `FEATHER(4)` is applied to the composite.
-
-The order of the `FEATHER` and `MASK` calls doesn't matter; you can have `FEATHER` before `MASK` or even interleave 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`.

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

@@ -1,7 +1,17 @@
 import torch
 import numpy as np
+from math import copysign
+import logging
 import itertools
 
+log = logging.getLogger("comfyui-prompt-control")
+
+
+def _norm_mag(w, n):
+    d = w - 1
+    return 1 + np.sign(d) * np.sqrt(np.abs(d) ** 2 / n)
+    # return  np.sign(w) * np.sqrt(np.abs(w)**2 / n)
+
 
 def _grouper(n, iterable):
     it = iter(iterable)
@@ -12,29 +22,22 @@ def _grouper(n, iterable):
         yield chunk
 
 
-def _norm_mag(w, n):
-    d = w - 1
-    return 1 + np.sign(d) * np.sqrt(np.abs(d) ** 2 / n)
-    # return  np.sign(w) * np.sqrt(np.abs(w)**2 / n)
+def batched_clip_encode(tokens, length, encode_func, num_chunks):
+    embs = []
+    for e in _grouper(32, tokens):
+        enc, pooled = encode_func(e)
+        enc = enc.reshape((len(e), length, -1))
+        embs.append(enc)
+
+    embs = torch.cat(embs)
+    embs = embs.reshape((len(tokens) // num_chunks, length * num_chunks, -1))
+    return embs
 
 
 def weights_like(weights, emb):
     return torch.tensor(weights, dtype=emb.dtype, device=emb.device).reshape(1, -1, 1).expand(emb.shape)
 
 
-def divide_length(word_ids, weights):
-    sums = dict(zip(*np.unique(word_ids, return_counts=True)))
-    sums[0] = 1
-    weights = [[_norm_mag(w, sums[id]) if id != 0 else 1.0 for w, id in zip(x, y)] for x, y in zip(weights, word_ids)]
-    return weights
-
-
-def shift_mean_weight(word_ids, weights):
-    delta = 1 - np.mean([w for x, y in zip(weights, word_ids) for w, id in zip(x, y) if id != 0])
-    weights = [[w if id == 0 else w + delta for w, id in zip(x, y)] for x, y in zip(weights, word_ids)]
-    return weights
-
-
 def scale_to_norm(weights, word_ids, w_max):
     top = np.max(weights)
     w_max = min(top, w_max)
@@ -48,61 +51,6 @@ def mask_word_id(tokens, word_ids, target_id, mask_token):
     return (new_tokens, mask)
 
 
-def batched_clip_encode(tokens, length, encode_func, num_chunks):
-    embs = []
-    for e in _grouper(32, tokens):
-        enc, pooled = encode_func(e)
-        enc = enc.reshape((len(e), length, -1))
-        embs.append(enc)
-    embs = torch.cat(embs)
-    embs = embs.reshape((len(tokens) // num_chunks, length * num_chunks, -1))
-    return embs
-
-
-def from_masked(tokens, weights, word_ids, base_emb, length, encode_func, m_token=266):
-    pooled_base = base_emb[0, length - 1 : length, :]
-    wids, inds = np.unique(np.array(word_ids).reshape(-1), return_index=True)
-    weight_dict = dict((id, w) for id, w in zip(wids, np.array(weights).reshape(-1)[inds]) if w != 1.0)
-
-    if len(weight_dict) == 0:
-        return torch.zeros_like(base_emb), base_emb[0, length - 1 : length, :]
-
-    weight_tensor = weights_like(weights, base_emb)
-
-    # m_token = (clip.tokenizer.end_token, 1.0) if  clip.tokenizer.pad_with_end else (0,1.0)
-    # TODO: find most suitable masking token here
-    m_token = (m_token, 1.0)
-
-    ws = []
-    masked_tokens = []
-    masks = []
-
-    # create prompts
-    for id, w in weight_dict.items():
-        masked, m = mask_word_id(tokens, word_ids, id, m_token)
-        masked_tokens.extend(masked)
-        masks.append(weights_like(m, base_emb))
-
-        ws.append(w)
-
-    # batch process prompts
-    embs = batched_clip_encode(masked_tokens, length, encode_func, len(tokens))
-    masks = torch.cat(masks)
-
-    embs = base_emb.expand(embs.shape) - embs
-    pooled = embs[0, length - 1 : length, :]
-
-    embs *= masks
-    embs = embs.sum(axis=0, keepdim=True)
-
-    pooled_start = pooled_base.expand(len(ws), -1)
-    ws = torch.tensor(ws).reshape(-1, 1).expand(pooled_start.shape)
-    pooled = (pooled - pooled_start) * (ws - 1)
-    pooled = pooled.mean(axis=0, keepdim=True)
-
-    return ((weight_tensor - 1) * embs), pooled_base + pooled
-
-
 def mask_inds(tokens, inds, mask_token):
     clip_len = len(tokens[0])
     inds_set = set(inds)
@@ -112,34 +60,6 @@ def mask_inds(tokens, inds, mask_token):
     return new_tokens
 
 
-def down_weight(tokens, weights, word_ids, base_emb, length, encode_func, m_token=266):
-    w, w_inv = np.unique(weights, return_inverse=True)
-
-    if np.sum(w < 1) == 0:
-        return base_emb, tokens, base_emb[0, length - 1 : length, :]
-    # m_token = (clip.tokenizer.end_token, 1.0) if  clip.tokenizer.pad_with_end else (0,1.0)
-    # using the comma token as a masking token seems to work better than aos tokens for SD 1.x
-    m_token = (m_token, 1.0)
-
-    masked_tokens = []
-
-    masked_current = tokens
-    for i in range(len(w)):
-        if w[i] >= 1:
-            continue
-        masked_current = mask_inds(masked_current, np.where(w_inv == i)[0], m_token)
-        masked_tokens.extend(masked_current)
-
-    embs = batched_clip_encode(masked_tokens, length, encode_func, len(tokens))
-    embs = torch.cat([base_emb, embs])
-    w = w[w <= 1.0]
-    w_mix = np.diff([0] + w.tolist())
-    w_mix = torch.tensor(w_mix, dtype=embs.dtype, device=embs.device).reshape((-1, 1, 1))
-
-    weighted_emb = (w_mix * embs).sum(axis=0, keepdim=True)
-    return weighted_emb, masked_current, weighted_emb[0, length - 1 : length, :]
-
-
 def scale_emb_to_mag(base_emb, weighted_emb):
     norm_base = torch.linalg.norm(base_emb)
     norm_weighted = torch.linalg.norm(weighted_emb)
@@ -147,12 +67,6 @@ def scale_emb_to_mag(base_emb, weighted_emb):
     return embeddings_final
 
 
-def recover_dist(base_emb, weighted_emb):
-    fixed_std = (base_emb.std() / weighted_emb.std()) * (weighted_emb - weighted_emb.mean())
-    embeddings_final = fixed_std + (base_emb.mean() - fixed_std.mean())
-    return embeddings_final
-
-
 def perp_weight(weights, unweighted_embs, empty_embs):
     unweighted, unweighted_pooled = unweighted_embs
     zero, zero_pooled = empty_embs
@@ -171,72 +85,294 @@ def perp_weight(weights, unweighted_embs, empty_embs):
     result[~over1] = (unweighted - (1 - weights) * perp)[~over1]
     result[weights == 0.0] = zero[weights == 0.0]
 
+    # Not sure if this is an implementation bug or if this just doesn't make sense with T5
+    nans = result.isnan()
+    if nans.any():
+        log.warning("perp weight returned NaNs (known to happen with T5), replacing with 0")
+        result[nans] = 0.0
+
     return result, unweighted_pooled
 
 
+def style_comfy(encoder, tokens, **kwargs):
+    tokens = encoder.without_word_ids(tokens)
+    return encoder.encode_fn(tokens)
+
+
+def style_a1111(encoder, tokens, **kwargs):
+    base_emb, pooled = encoder.base_emb(tokens)
+    weighted_emb = base_emb * weights_like(encoder.weights(tokens), base_emb)
+    weighted_emb = (base_emb.mean() / weighted_emb.mean()) * weighted_emb  # renormalize
+    return weighted_emb, pooled
+
+
+def style_compel(encoder, tokens, **kwargs):
+    pos_tokens = encoder.weighted_with(tokens, lambda w: w if w > 1.0 else 1.0)
+    weighted_emb, pooled = encoder.encode_fn(pos_tokens)
+    weighted_emb, _, pooled = encoder.down_weight(
+        pos_tokens, encoder.weights(tokens), encoder.word_ids(tokens), weighted_emb, pooled
+    )
+    return weighted_emb, pooled
+
+
+def style_comfypp(encoder, tokens, **kwargs):
+    unweighted_tokens = encoder.unweighted(tokens)
+    base_emb, pooled_base = encoder.base_emb(tokens)
+    weighted_emb, tokens_down, _ = encoder.down_weight(
+        unweighted_tokens, encoder.weights(tokens), encoder.word_ids(tokens), base_emb, pooled_base
+    )
+    weights = encoder.weights(encoder.weighted_with(tokens, lambda w: w if w > 1.0 else 1.0))
+    embs, pooled = encoder.from_masked(
+        unweighted_tokens,
+        weights,
+        encoder.word_ids(tokens),
+        base_emb,
+        pooled_base,
+    )
+    weighted_emb += embs
+
+    return weighted_emb, pooled
+
+
+def style_downweight(encoder, tokens, **kwargs):
+    weights = scale_to_norm(encoder.weights(tokens), encoder.word_ids(tokens), encoder.w_max)
+    base_emb, pooled_base = encoder.base_emb(tokens)
+    weighted_emb, _, pooled = encoder.down_weight(
+        encoder.unweighted(tokens), weights, encoder.word_ids(tokens), base_emb, pooled_base
+    )
+
+    return weighted_emb, pooled
+
+
+def style_perp(encoder, tokens, **kwargs):
+    zero_emb, zero_pooled = encoder.encode_fn(encoder.tokenizer.tokenize_with_weights(""))
+    base_emb, pooled = encoder.base_emb(tokens)
+    return perp_weight(encoder.weights(tokens), (base_emb, pooled), (zero_emb, zero_pooled))
+
+
+def apply_negpip(encoder, emb, pooled, **kwargs):
+    original_tokens = kwargs["original_tokens"]
+    emb_negpip = torch.empty_like(emb).repeat(1, 2, 1)
+    emb_negpip[:, 0::2, :] = emb
+    emb_negpip[:, 1::2, :] = emb * weights_like(encoder.signs(original_tokens), emb)
+    return emb_negpip, pooled
+
+
+def norm_length(encoder, tokens, **kwargs):
+    word_ids = encoder.word_ids(tokens)
+    sums = dict(zip(*np.unique(word_ids, return_counts=True)))
+    sums[0] = 1
+    tokens = [[(t, _norm_mag(w, sums[id]) if id != 0 else 1.0, id) for (t, w, id) in x] for x in tokens]
+    return tokens
+
+
+def norm_mean(encoder, tokens, **kwargs):
+    weights = encoder.weights(tokens)
+    word_ids = encoder.word_ids(tokens)
+    delta = 1 - np.mean([w for x, y in zip(weights, word_ids) for w, id in zip(x, y) if id != 0])
+    tokens = [[(t, w if id == 0 else w + delta, id) for (t, w, id) in x] for x in tokens]
+    return tokens
+
+
+def norm_none(encoder, tokens, **kwargs):
+    return tokens
+
+
+class AdvancedEncoder:
+    STYLES = {
+        "A1111": style_a1111,
+        "comfy": style_comfy,
+        "comfy++": style_comfypp,
+        "compel": style_compel,
+        "down_weight": style_downweight,
+        "perp": style_perp,
+    }
+    NORMALIZATION_OPS = {
+        "none": norm_none,
+        "length": norm_length,
+        "mean": norm_mean,
+    }
+
+    @classmethod
+    def add_encoder(cls, name, fn):
+        cls.STYLES[name] = fn
+
+    def add_normalization_op(cls, name, fn):
+        cls.NORMALIZATION_OPS[name] = fn
+
+    @classmethod
+    def weighted_with(cls, tokens, fn=id, word_ids=True):
+        w = ([(t, fn(w), id) for t, w, id in x] for x in tokens)
+        if not word_ids:
+            w = cls.without_word_ids(w)
+        return list(w)
+
+    @classmethod
+    def unweighted(cls, tokens, word_ids=False):
+        return cls.weighted_with(tokens, fn=lambda w: 1.0, word_ids=word_ids)
+
+    @classmethod
+    def tokens_only(cls, tokens):
+        return list([t[0] for t in x] for x in tokens)
+
+    @classmethod
+    def weights(cls, tokens):
+        return list([t[1] for t in x] for x in tokens)
+
+    @classmethod
+    def word_ids(cls, tokens):
+        return list([t[2] for t in x] for x in tokens)
+
+    @classmethod
+    def signs(cls, tokens):
+        return list([copysign(1, t[1]) for t in x] for x in tokens)
+
+    @classmethod
+    def without_word_ids(cls, tokens):
+        return list([(t, w) for t, w, _ in x] for x in tokens)
+
+    def __init__(self, encode_fn, style, normalization, tokenizer, m_token="+", w_max=1.0, **extra_args):
+        self.encode_fn = encode_fn
+        self.preprocessors = []
+        self.postprocessors = []
+        self.tokenizer = tokenizer
+        self.extra_args = extra_args
+        self.m_token = tokenizer.tokenize_with_weights(m_token)[0][tokenizer.tokens_start]
+        self.max_length = tokenizer.max_length if tokenizer.pad_to_max_length else None
+        self.w_max = w_max
+
+        if style == "comfy++" and not self.max_length:
+            log.warning("comfy++ does not work with tokenizer %s, using default weighting", tokenizer)
+            style = "comfy"
+
+        norms = normalization.split("+")
+        assert style in self.STYLES, f"Invalid weight interpretation: {style}"
+        self.weight_fn = self.STYLES[style]
+        for n in norms:
+            n = n.strip()
+            assert n in self.NORMALIZATION_OPS, f"Invalid normalization: {normalization}"
+            self.preprocessors.append(self.NORMALIZATION_OPS[n])
+
+        negpip = extra_args.get("has_negpip")
+        if negpip:
+
+            def _encode(t):
+                emb, pooled = encode_fn(t)
+                return emb[:, 0::2, :], pooled
+
+            self.encode_fn = _encode
+            self.preprocessors.insert(0, lambda encoder, tokens, **kwargs: encoder.weighted_with(tokens, abs))
+            self.postprocessors.insert(0, apply_negpip)
+
+    def base_emb(self, tokens):
+        unweighted = self.unweighted(tokens)
+        return self.encode_fn(unweighted)
+
+    def down_weight(self, tokens, weights, word_ids, base_emb, pooled_base):
+        w, w_inv = np.unique(weights, return_inverse=True)
+
+        if np.sum(w < 1) == 0:
+            return (
+                base_emb,
+                tokens,
+                (
+                    base_emb[0, self.max_length - 1 : self.max_length, :]
+                    if (pooled_base is not None and self.max_length)
+                    else None
+                ),
+            )
+
+        masked_current = tokens
+        emblist = [base_emb]
+        for i in range(len(w)):
+            if w[i] >= 1:
+                continue
+            masked_current = mask_inds(masked_current, np.where(w_inv == i)[0], self.m_token)
+            masked, _ = self.encode_fn(masked_current)
+            emblist.append(masked)
+
+        embs = torch.cat(emblist)
+        w = w[w <= 1.0]
+        w_mix = np.diff([0] + w.tolist())
+        w_mix = torch.tensor(w_mix, dtype=embs.dtype, device=embs.device).reshape((-1, 1, 1))
+
+        weighted_emb = (w_mix * embs).sum(axis=0, keepdim=True)
+        pooled = pooled_base
+        if pooled is not None and self.max_length:
+            pooled = weighted_emb[0, self.max_length - 1 : self.max_length, :]
+        return weighted_emb, masked_current, pooled
+
+    def from_masked(self, tokens, weights, word_ids, base_emb, pooled_base):
+        wids, inds = np.unique(np.array(word_ids).reshape(-1), return_index=True)
+        weight_dict = dict((id, w) for id, w in zip(wids, np.array(weights).reshape(-1)[inds]) if w != 1.0)
+
+        if len(weight_dict) == 0:
+            return torch.zeros_like(base_emb), torch.zeros_like(pooled_base) if pooled_base is not None else None
+
+        weight_tensor = weights_like(weights, base_emb)
+
+        ws = []
+        masked_tokens = []
+        masks = []
+
+        # create prompts
+        for id, w in weight_dict.items():
+            masked, m = mask_word_id(tokens, word_ids, id, self.m_token)
+            masks.append(weights_like(m, base_emb))
+            masked_tokens.extend(masked)
+
+            ws.append(w)
+
+        # TODO: figure out how to get rid of this
+        embs = batched_clip_encode(masked_tokens, self.max_length, self.encode_fn, len(tokens))
+        masks = torch.cat(masks)
+
+        embs = base_emb.expand(embs.shape) - embs
+        if pooled_base is not None and self.max_length:
+            pooled = embs[0, self.max_length - 1 : self.max_length, :]
+            pooled_start = pooled_base.expand(len(ws), -1)
+            ws = torch.tensor(ws).reshape(-1, 1).expand(pooled_start.shape)
+            pooled = (pooled - pooled_start) * (ws - 1)
+            pooled = pooled.mean(axis=0, keepdim=True)
+            pooled = pooled_base + pooled
+
+        if embs.shape[0] != masks.shape[0]:
+            embs = embs.repeat(masks.shape[0], 1, 1)
+        embs *= masks
+        embs = embs.sum(axis=0, keepdim=True)
+
+        return ((weight_tensor - 1) * embs), pooled
+
+    def __call__(self, tokens, apply_to_pooled=False, return_pooled=False):
+        normalized_tokens = tokens
+        for op in self.preprocessors:
+            normalized_tokens = op(self, normalized_tokens)
+
+        emb, pooled = self.weight_fn(self, normalized_tokens, original_tokens=tokens)
+
+        for fn in self.postprocessors:
+            emb, pooled = fn(self, emb, pooled, tokens=tokens, original_tokens=tokens)
+
+        if return_pooled:
+            if not apply_to_pooled:
+                _, pooled = self.base_emb(tokens)
+            return emb, pooled
+        return emb, None
+
+
 def advanced_encode_from_tokens(
     tokenized,
     token_normalization,
     weight_interpretation,
     encode_func,
-    m_token=266,
-    length=77,
+    m_token="+",
     w_max=1.0,
     return_pooled=False,
     apply_to_pooled=False,
-    **extra_args
+    tokenizer=None,
+    **extra_args,
 ):
-    tokens = [[t for t, _, _ in x] for x in tokenized]
-    weights = [[w for _, w, _ in x] for x in tokenized]
-    word_ids = [[wid for _, _, wid in x] for x in tokenized]
-
-    for op in token_normalization.split("+"):
-        op = op.strip()
-        if op == "length":
-            # distribute down/up weights over word lengths
-            weights = divide_length(word_ids, weights)
-        if op == "mean":
-            weights = shift_mean_weight(word_ids, weights)
-
-    pooled = None
-
-    if weight_interpretation == "comfy":
-        weighted_tokens = [[(t, w) for t, w in zip(x, y)] for x, y in zip(tokens, weights)]
-        weighted_emb, pooled_base = encode_func(weighted_tokens)
-        pooled = pooled_base
-    else:
-        unweighted_tokens = [[(t, 1.0) for t, _, _ in x] for x in tokenized]
-        base_emb, pooled_base = encode_func(unweighted_tokens)
-
-    if weight_interpretation == "A1111":
-        weighted_emb = base_emb * weights_like(weights, base_emb)  # from_zero
-        weighted_emb = (base_emb.mean() / weighted_emb.mean()) * weighted_emb  # renormalize
-        pooled = pooled_base
-
-    if weight_interpretation == "compel":
-        pos_tokens = [[(t, w) if w >= 1.0 else (t, 1.0) for t, w in zip(x, y)] for x, y in zip(tokens, weights)]
-        weighted_emb, _ = encode_func(pos_tokens)
-        weighted_emb, _, pooled = down_weight(pos_tokens, weights, word_ids, weighted_emb, length, encode_func)
-
-    if weight_interpretation == "comfy++":
-        weighted_emb, tokens_down, _ = down_weight(unweighted_tokens, weights, word_ids, base_emb, length, encode_func)
-        weights = [[w if w > 1.0 else 1.0 for w in x] for x in weights]
-        # unweighted_tokens = [[(t,1.0) for t, _,_ in x] for x in tokens_down]
-        embs, pooled = from_masked(unweighted_tokens, weights, word_ids, base_emb, length, encode_func)
-        weighted_emb += embs
-
-    if weight_interpretation == "down_weight":
-        weights = scale_to_norm(weights, word_ids, w_max)
-        weighted_emb, _, pooled = down_weight(unweighted_tokens, weights, word_ids, base_emb, length, encode_func)
-
-    if weight_interpretation == "perp":
-        weighted_emb, pooled = perp_weight(
-            weights, (base_emb, pooled_base), encode_func(extra_args["tokenizer"].tokenize_with_weights(""))
-        )
-
-    if return_pooled:
-        if apply_to_pooled:
-            return weighted_emb, pooled
-        else:
-            return weighted_emb, pooled_base
-    return weighted_emb, None
+    enc = AdvancedEncoder(
+        encode_func, weight_interpretation, token_normalization, tokenizer, m_token, w_max, **extra_args
+    )
+    return enc(tokenized, return_pooled=return_pooled, apply_to_pooled=apply_to_pooled)

prompt_control/attention_couple_ppm.py

@@ -0,0 +1,248 @@
+# 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 itertools
+import logging
+import math
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from comfy.hooks import EnumHookScope, HookGroup, TransformerOptionsHook, set_hooks_for_conditioning
+from comfy.model_patcher import ModelPatcher
+
+log = logging.getLogger("comfyui-prompt-control")
+
+
+def set_cond_attnmask(base_cond, extra_conds, fill=False):
+    hook = AttentionCoupleHook()
+    c = [base_cond[0][0], base_cond[0][1].copy()]
+    # hook uses these, remove them to avoid doing latent masking
+    c[1].pop("mask", None)
+    c[1].pop("strength", None)
+    c[1].pop("mask_strength", None)
+    c = [c]
+    c.extend(base_cond[1:])
+
+    hook.initialize_regions(base_cond[0], extra_conds, fill=fill)
+    group = HookGroup()
+    group.add(hook)
+
+    return set_hooks_for_conditioning(c, hooks=group, append_hooks=True)
+
+
+def get_mask(mask, batch_size, num_tokens, extra_options):
+    activations_shape = extra_options["activations_shape"]
+    size = activations_shape[-2:]
+
+    num_conds = mask.shape[0]
+    mask_downsample = F.interpolate(mask, size=size, mode="nearest")
+    mask_downsample_reshaped = mask_downsample.view(num_conds, num_tokens, 1).repeat_interleave(batch_size, dim=0)
+
+    return mask_downsample_reshaped
+
+
+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):
+    COND_UNCOND_COUPLE_OPTION = "cond_or_uncond_hook_couple"
+    COND = 0
+    UNCOND = 1
+
+    def __init__(self):
+        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.has_negpip = False
+
+        # calculate later. All clones must refer to the same kv dict
+        self.kv = {"k": None, "v": None}
+
+    def initialize_regions(self, base_cond, conds, fill):
+        self.num_conds = len(conds) + 1
+        self.base_strength = base_cond[1].get("strength", 1.0)
+        self.strengths = [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].get("mask", None)
+        masks = [cond[1].get("mask") * cond[1].get("mask_strength") for cond in conds]
+        if len(masks) < 1:
+            raise ValueError("Attention Couple hook makes no sense without masked conds")
+
+        if any(m is None for m in masks):
+            raise ValueError("All conds given to Attention Couple must have masks")
+
+        if any(m.shape != masks[0].shape for m in masks) or (
+            base_mask is not None and base_mask.shape != masks[0].shape
+        ):
+            largest_shape = max(m.shape for m in masks)
+            if base_mask is not None:
+                largest_shape = max(largest_shape, base_mask.shape)
+            log.warning("Attention Couple: Masks are irregularly shaped, resizing them all to match the largest")
+            for i in range(len(masks)):
+                masks[i] = F.interpolate(masks[i].unsqueeze(1), size=largest_shape[1:], mode="nearest-exact").squeeze(1)
+
+            if base_mask is not None:
+                base_mask = F.interpolate(base_mask.unsqueeze(1), size=largest_shape[1:], mode="nearest-exact").squeeze(
+                    1
+                )
+
+        if base_mask is None:
+            if not fill:
+                raise ValueError("You must specify a base mask when fill=False")
+            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)
+
+    def on_apply_hooks(self, model: ModelPatcher, transformer_options: dict[str, Any]):
+        if self.kv["k"] is None:
+            self.has_negpip = model.model_options.get("ppm_negpip", False)
+            log.debug("AttentionCouple has_negpip=%s", self.has_negpip)
+
+            # Skip the base cond here, which is always first
+            if self.has_negpip:
+                self.kv["k"] = [cond[:, 0::2] for cond in self.conds[1:]]
+                self.kv["v"] = [cond[:, 1::2] for cond in self.conds[1:]]
+            else:
+                self.kv["k"] = self.kv["v"] = self.conds[1:]
+
+        return super().on_apply_hooks(model, transformer_options)
+
+    def clone(self):
+        c: AttentionCoupleHook = super().clone()
+        c.mask = self.mask
+        c.conds = self.conds
+        c.kv = self.kv
+        c.has_negpip = self.has_negpip
+        c.base_strength = self.base_strength
+        c.strengths = self.strengths
+        c.num_conds = self.num_conds
+        return c
+
+    def to(self, *args, **kwargs):
+        self.conds = [c.to(*args, **kwargs) for c in self.conds]
+        self.mask = self.mask.to(*args, **kwargs)
+        if self.kv["k"] is not None:
+            self.kv["k"] = [c.to(*args, **kwargs) for c in self.kv["k"]]
+            self.kv["v"] = [c.to(*args, **kwargs) for c in self.kv["v"]]
+        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"]
+        cond_or_uncond_couple = extra_options[self.COND_UNCOND_COUPLE_OPTION] = list(cond_or_uncond)
+        num_chunks = len(cond_or_uncond)
+
+        # Cloning messes up the device sometimes
+        if self.kv["k"][0].device != k.device:
+            self.to(k)
+
+        conds_k = self.kv["k"]
+        conds_v = self.kv["v"]
+
+        lcm_tokens_k = math.lcm(k.shape[1], *(cond.shape[1] for cond in conds_k))
+        lcm_tokens_v = math.lcm(v.shape[1], *(cond.shape[1] for cond in conds_v))
+        q_chunks = q.chunk(num_chunks, dim=0)
+        k_chunks = k.chunk(num_chunks, dim=0)
+        v_chunks = v.chunk(num_chunks, dim=0)
+
+        bs = q.shape[0] // num_chunks
+
+        conds_k_tensor = conds_v_tensor = torch.cat(
+            [cond.repeat(bs, lcm_tokens_k // cond.shape[1], 1) * self.strengths[i] for i, cond in enumerate(conds_k)],
+            dim=0,
+        )
+        if self.has_negpip:
+            conds_v_tensor = torch.cat(
+                [
+                    cond.repeat(bs, lcm_tokens_v // cond.shape[1], 1) * self.strengths[i]
+                    for i, cond in enumerate(conds_v)
+                ],
+                dim=0,
+            )
+
+        qs, ks, vs = [], [], []
+        cond_or_uncond_couple.clear()
+
+        for i, cond_type in enumerate(cond_or_uncond):
+            q_target = q_chunks[i]
+            k_target = k_chunks[i].repeat(1, lcm_tokens_k // k.shape[1], 1)
+            v_target = v_chunks[i].repeat(1, lcm_tokens_v // v.shape[1], 1)
+            if cond_type == self.UNCOND:
+                qs.append(q_target)
+                ks.append(k_target)
+                vs.append(v_target)
+                cond_or_uncond_couple.append(self.UNCOND)
+            else:
+                qs.append(q_target.repeat(self.num_conds, 1, 1))
+                ks.append(
+                    torch.cat(
+                        [
+                            k_target * self.base_strength,
+                            conds_k_tensor,
+                        ],
+                        dim=0,
+                    )
+                )
+                vs.append(
+                    torch.cat(
+                        [
+                            v_target * self.base_strength,
+                            conds_v_tensor,
+                        ],
+                        dim=0,
+                    )
+                )
+                cond_or_uncond_couple.extend(itertools.repeat(self.COND, self.num_conds))
+
+        q = torch.cat(qs, dim=0)
+        k = torch.cat(ks, dim=0)
+        v = torch.cat(vs, dim=0)
+
+        return q, k, v
+
+    def attn2_output_patch(self, out, extra_options):
+        cond_or_uncond = extra_options[self.COND_UNCOND_COUPLE_OPTION]
+        bs = out.shape[0] // len(cond_or_uncond)
+        mask_downsample = get_mask(self.mask, bs, out.shape[1], extra_options)
+        outputs = []
+        cond_outputs = []
+        i_cond = 0
+        for i, cond_type in enumerate(cond_or_uncond):
+            pos, next_pos = i * bs, (i + 1) * bs
+
+            if cond_type == self.UNCOND:
+                outputs.append(out[pos:next_pos])
+            else:
+                pos_cond, next_pos_cond = i_cond * bs, (i_cond + 1) * bs
+                masked_output = out[pos:next_pos] * mask_downsample[pos_cond:next_pos_cond]
+                cond_outputs.append(masked_output)
+                i_cond += 1
+
+        if len(cond_outputs) > 0:
+            cond_output = torch.stack(cond_outputs).sum(0)
+            outputs.append(cond_output)
+
+        return torch.cat(outputs, dim=0)

prompt_control/cache_hack.py

@@ -0,0 +1,47 @@
+import comfy_execution.caching
+from comfy_execution.graph_utils import is_link
+import nodes
+from os import environ
+
+import logging
+
+log = logging.getLogger("comfyui-prompt-control")
+
+
+include_unique_id_in_input = comfy_execution.caching.include_unique_id_in_input
+
+
+def promptcontrol_get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
+    if not dynprompt.has_node(node_id):
+        # This node doesn't exist -- we can't cache it.
+        return [float("NaN")]
+    node = dynprompt.get_node(node_id)
+    class_type = node["class_type"]
+    class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+    inputs = node["inputs"]
+    if hasattr(class_def, "CACHE_KEY"):
+        inputs = getattr(class_def, "CACHE_KEY")(inputs)
+    signature = [class_type, self.is_changed_cache.get(node_id)]
+    if (
+        self.include_node_id_in_input()
+        or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT)
+        or include_unique_id_in_input(class_type)
+    ):
+        signature.append(node_id)
+    for key in sorted(inputs.keys()):
+        if is_link(inputs[key]):
+            (ancestor_id, ancestor_socket) = inputs[key]
+            ancestor_index = ancestor_order_mapping[ancestor_id]
+            signature.append((key, ("ANCESTOR", ancestor_index, ancestor_socket)))
+        else:
+            signature.append((key, inputs[key]))
+    return signature
+
+
+def init():
+    if environ.get("PROMPTCONTROL_ENABLE_CACHE_HACK") != "1":
+        return
+    log.warning("Enabling Prompt Control cache hack")
+    comfy_execution.caching.CacheKeySetInputSignature.get_immediate_node_signature = (
+        promptcontrol_get_immediate_node_signature
+    )

prompt_control/cutoff.py

@@ -209,6 +209,9 @@ def encode_regions(clip_regions, encode, tokenizer):
         debug_tokens("region", region_prompt, tokenizer)
         region_emb, _ = encode(region_prompt)
         region_emb -= base_embedding_start
+        # NegPiP support:
+        if region_emb.shape[1] == 2 * region_masking.shape[1]:
+            region_masking = torch.repeat_interleave(region_masking, 2, dim=1)
         region_emb *= region_masking
 
         region_embeddings.append(region_emb)
@@ -217,6 +220,10 @@ def encode_regions(clip_regions, encode, tokenizer):
     embeddings_final_mask = torch.tensor(
         global_region_mask, dtype=base_embedding_full.dtype, device=base_embedding_full.device
     ).unsqueeze(-1)
+    # NegPiP support:
+    if region_embeddings.shape[1] == 2 * embeddings_final_mask.shape[1]:
+        embeddings_final_mask = torch.repeat_interleave(embeddings_final_mask, 2, dim=1)
+
     embeddings_final = base_embedding_start * embeddings_final_mask + base_embedding_outer * (1 - embeddings_final_mask)
     embeddings_final += region_embeddings
     return embeddings_final, pool

prompt_control/manual_test.py

@@ -0,0 +1,46 @@
+import main
+import nodes
+import prompt_control.adv_encode
+
+(l,) = nodes.CLIPLoader.load_clip(None, "clip_l.safetensors")
+(t5,) = nodes.CLIPLoader.load_clip(None, "t5base.safetensors")
+
+id(main)  # get rid of warning
+
+
+def adv(t, text, style="A1111", norm="none", new=True, **kwargs):
+    c = t.tokenize(text, return_word_ids=True)
+    if new:
+        style = "new+" + style
+    if t is t5:
+        te = t.patcher.model.t5base.encode_token_weights
+        token = t.tokenizer.clip_t5base
+        tok = c["t5base"]
+    else:
+        te = t.patcher.model.clip_l.encode_token_weights
+        token = t.tokenizer.clip_l
+        tok = c["l"]
+    return prompt_control.adv_encode.advanced_encode_from_tokens(tok, norm, style, te, tokenizer=token)
+
+
+def adv_all(t, text, styles=[], **kwargs):
+    r = []
+    for s in styles or prompt_control.adv_encode.AdvancedEncoder.STYLES:
+        print("Testing", s, kwargs)
+        r.append([s, adv(t, text, style=s, **kwargs)])
+    return r
+
+
+def replacenan(t):
+    t[t.isnan()] = 42.123321
+    return t
+
+
+def adv_equal(t, text, **kwargs):
+    old = adv_all(t, text, new=False, **kwargs)
+    new = adv_all(t, text, new=True, **kwargs)
+    r = {}
+    for i, o in enumerate(old):
+        n = new[i]
+        r[n[0]] = (replacenan(n[1][0]) == replacenan(o[1][0])).all()
+    return r

prompt_control/nodes_base.py

@@ -4,6 +4,29 @@ from .prompts import encode_prompt
 log = logging.getLogger("comfyui-prompt-control")
 
 
+class PCTextEncodeWithRange:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {"clip": ("CLIP",), "text": ("STRING", {"multiline": True})},
+            "optional": {
+                "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}),
+            },
+        }
+
+    RETURN_TYPES = ("CONDITIONING",)
+    CATEGORY = "promptcontrol/tools"
+    FUNCTION = "apply"
+    DESCRIPTION = "Like PCTextEncode, but if you know the range you need for a prompt, can be slightly more efficient when you have LoRAs scheduled on a CLIP model"
+
+    def apply(self, clip, text, start=0.0, end=1.0):
+        log.debug("PCTextEncode: Encoding '%s'", text)
+        defaults = clip.patcher.model_options.get("x-promptcontrol.defaults", {})
+        masks = clip.patcher.model_options.get("x-promptcontrol.masks", None)
+        return (encode_prompt(clip, text, start, end, defaults, masks),)
+
+
 class PCTextEncode:
     @classmethod
     def INPUT_TYPES(s):
@@ -14,17 +37,15 @@ class PCTextEncode:
     RETURN_TYPES = ("CONDITIONING",)
     CATEGORY = "promptcontrol"
     FUNCTION = "apply"
+    DESCRIPTION = "Encodes a prompt with extra goodies from Prompt Control. This node does *not* support scheduling"
 
     def apply(self, clip, text):
-        defaults = clip.patcher.model_options.get("x-promptcontrol.defaults", {})
-        masks = clip.patcher.model_options.get("x-promptcontrol.masks", None)
-        return (encode_prompt(clip, text, 0, 1.0, defaults, masks),)
+        return PCTextEncodeWithRange.apply(self, clip, text, 0.0, 1.0)
 
 
-NODE_CLASS_MAPPINGS = {
-    "PCTextEncode": PCTextEncode,
-}
+NODE_CLASS_MAPPINGS = {"PCTextEncode": PCTextEncode, "PCTextEncodeWithRange": PCTextEncodeWithRange}
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "PCTextEncode": "PC Text Encode (no scheduling)",
+    "PCTextEncode": "PC: Text Encode (no scheduling)",
+    "PCTextEncodeWithRange": "PC: Text Encode with Range (no scheduling)",
 }

prompt_control/nodes_hooks.py

@@ -1,9 +1,13 @@
 import logging
-import comfy.utils
+
 import comfy.hooks
+import comfy.utils
 import folder_paths
-from .utils import consolidate_schedule
+from comfy.comfy_types.node_typing import IO, ComfyNodeABC, InputTypeDict
+
+from .attention_couple_ppm import AttentionCoupleHook
 from .parser import parse_prompt_schedules
+from .utils import consolidate_schedule
 
 log = logging.getLogger("comfyui-prompt-control")
 
@@ -79,10 +83,51 @@ def lora_hooks_from_schedule(schedules, non_scheduled):
         return hooks
 
 
+class PCAttentionCoupleBatchNegative(ComfyNodeABC):
+    @classmethod
+    def INPUT_TYPES(cls) -> InputTypeDict:
+        return {
+            "required": {
+                "positive": (IO.CONDITIONING, {}),
+                "negative": (IO.CONDITIONING, {}),
+            },
+        }
+
+    RETURN_TYPES = (IO.CONDITIONING, IO.CONDITIONING)
+    RETURN_NAMES = ("positive", "negative")
+    CATEGORY = "promptcontrol/v2"
+    FUNCTION = "batch"
+    EXPERIMENTAL = True
+
+    # May cause side-effects?
+    # TODO: Support scheduling in negative prompt
+    def batch(self, positive, negative):
+        if len(negative) != 1:
+            log.warning("Batching scheduled negatives is not supported yet")
+            return (positive, negative)
+
+        negative_batch = []
+        for p in positive:
+            n = [negative[0][0], negative[0][1].copy()]
+            n_hook_group: comfy.hooks.HookGroup = n[1].get("hooks", comfy.hooks.HookGroup()).clone()
+            p_hook_group: comfy.hooks.HookGroup = p[1].get("hooks", comfy.hooks.HookGroup())
+            attn_couple = [hook for hook in p_hook_group.hooks if isinstance(hook, AttentionCoupleHook)]
+            for hook in attn_couple:
+                n_hook_group.add(hook)
+            n[1]["hooks"] = p_hook_group if n_hook_group.hooks == p_hook_group.hooks else n_hook_group
+            n[1]["start_percent"] = p[1].get("start_percent", 0.0)
+            n[1]["end_percent"] = p[1].get("end_percent", 1.0)
+            negative_batch.append(n)
+
+        return (positive, negative_batch)
+
+
 NODE_CLASS_MAPPINGS = {
     "PCLoraHooksFromText": PCLoraHooksFromText,
+    "PCAttentionCoupleBatchNegative": PCAttentionCoupleBatchNegative,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "PCLoraHooksFromText": "PC LoRA Hooks From Text (non-lazy)",
+    "PCLoraHooksFromText": "PC: LoRA Hooks From Text (non-lazy)",
+    "PCAttentionCoupleBatchNegative": "PC: Attention Couple (batch negative)",
 }

prompt_control/nodes_lazy.py

@@ -1,12 +1,31 @@
 import logging
 from .parser import parse_prompt_schedules
-from comfy_execution.graph_utils import GraphBuilder
+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")
 
 from .utils import consolidate_schedule, find_nonscheduled_loras
+import json
+
+
+def _cache_key(cachekey, inputs):
+    out = inputs.copy()
+    text = inputs.get("text")
+    if text is not None and not is_link(text):
+        out["text"] = cache_key_from_inputs(cachekey, **inputs)
+    return out
+
+
+def cache_key_prompt(inputs):
+    return _cache_key("prompt", inputs)
+
+
+def cache_key_lora(inputs):
+    return _cache_key("loras", inputs)
 
 
 def create_lora_loader_nodes(graph, model, clip, loras):
@@ -24,9 +43,10 @@ def create_lora_loader_nodes(graph, model, clip, loras):
 
 
 def create_hook_nodes_for_lora(graph, path, info, existing_node, start_pct, end_pct):
+    prev_keyframe = None
     next_keyframe = None
     if not existing_node:
-        log.debug("Creating hook for %s", path)
+        log.debug("Creating hook for %s, weight=%s, weight_clip=%s", path, info["weight"], info["weight_clip"])
         hook_node = graph.node("CreateHookLora")
         hook_node.set_input("lora_name", path)
         hook_node.set_input("strength_model", info["weight"])
@@ -59,7 +79,7 @@ def create_hook_nodes_for_lora(graph, path, info, existing_node, start_pct, end_
     next_keyframe.set_input("strength_mult", 1.0)
     prev_hook_kf = next_keyframe.out(0)
     if end_pct < 1.0:
-        log.debug("Creating end keyframe for %s, start=%s", path, start_pct)
+        log.debug("Creating end keyframe for %s, start=%s", path, end_pct)
         next_keyframe = graph.node("CreateHookKeyframe")
         next_keyframe.set_input("strength_mult", 0.0)
         next_keyframe.set_input("start_percent", end_pct)
@@ -67,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):
+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
@@ -100,39 +124,45 @@ def build_lora_schedule(graph, schedule, model, clip, apply_hooks):
     # Finally, combine all hooks and optionally apply
     if len(hooks) > 0:
         res = hooks[0]
-        for h in hooks[:1]:
+        for h in hooks[1:]:
             n = graph.node("CombineHooks2")
             n.set_input("hooks_A", res.out(0))
             n.set_input("hooks_B", h.out(0))
             res = n
         res = res.out(0)
-    if apply_hooks:
-        n = graph.node("SetClipHooks")
-        n.set_input("clip", clip)
-        n.set_input("hooks", res)
-        n.set_input("apply_to_conds", True)
-        n.set_input("schedule_clip", True)
-        clip = n.out(0)
+        if clip is not None and apply_hooks:
+            n = graph.node("SetClipHooks")
+            n.set_input("clip", clip)
+            n.set_input("hooks", res)
+            n.set_input("apply_to_conds", True)
+            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))
 
-    return {"result": (model, clip, res), "expand": r}
+    ret = (model, clip, res)
+
+    return {"result": ret, "expand": r}
 
 
 class PCLazyLoraLoaderAdvanced:
+    CACHE_KEY = cache_key_lora
+
     @classmethod
     def INPUT_TYPES(s):
         return {
-            "required": {
-                "text": ("STRING", {"multiline": True}),
+            "optional": {
                 "model": ("MODEL", {"rawLink": True}),
                 "clip": ("CLIP", {"rawLink": True}),
+                "text": ("STRING", {"multiline": True, "default": ""}),
                 "apply_hooks": ("BOOLEAN", {"default": True}),
-            },
-            "optional": {
                 "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"},
         }
@@ -142,40 +172,42 @@ class PCLazyLoraLoaderAdvanced:
     CATEGORY = "promptcontrol"
     FUNCTION = "apply"
 
-    def apply(self, model, clip, text, apply_hooks, unique_id, tags="", start=0.0, end=1.0):
-        schedule = parse_prompt_schedules(text).with_filters(filters=tags, start=start, end=end)
-        graph = GraphBuilder(f"PCLazyLoraLoaderAdvanced-{unique_id}")
-        return build_lora_schedule(graph, schedule, model, clip, apply_hooks)
+    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()
+        r = build_lora_schedule(graph, schedule, model, clip, apply_hooks=apply_hooks)
+        return r
 
 
-class PCLazyLoraLoader:
+class PCLazyLoraLoader(PCLazyLoraLoaderAdvanced):
     @classmethod
     def INPUT_TYPES(s):
         return {
-            "required": {
-                "text": ("STRING", {"multiline": True}),
+            "optional": {
                 "model": ("MODEL", {"rawLink": True}),
                 "clip": ("CLIP", {"rawLink": True}),
-                "apply_hooks": ("BOOLEAN", {"default": True}),
+                "text": ("STRING", {"multiline": True, "default": ""}),
             },
             "hidden": {"unique_id": "UNIQUE_ID"},
         }
 
-    RETURN_TYPES = ("MODEL", "CLIP", "HOOKS")
-    OUTPUT_TOOLTIPS = ("Returns a model and clip with LoRAs scheduled",)
+    RETURN_TYPES = (
+        "MODEL",
+        "CLIP",
+    )
     CATEGORY = "promptcontrol"
-    FUNCTION = "apply"
 
-    def apply(self, model, clip, text, apply_hooks, unique_id):
-        graph = GraphBuilder(f"PCLazyLoraLoader-{unique_id}")
-        schedule = parse_prompt_schedules(text)
-        return build_lora_schedule(graph, schedule, model, clip, apply_hooks)
+    def apply(self, *args, **kwargs):
+        r = super().apply(*args, **kwargs)
+        r["result"] = r["result"][:2]
+        return r
 
 
 def build_scheduled_prompts(graph, schedules, clip):
     nodes = []
     start_pct = 0.0
-    prompt_cache = {}
     for end_pct, c in schedules:
         p = c["prompt"]
         p, classnames = get_function(p, "NODE", ["PCTextEncode", "text"])
@@ -184,12 +216,9 @@ def build_scheduled_prompts(graph, schedules, clip):
         if classnames:
             classname = classnames[0][0]
             paramname = classnames[0][1]
-        node = prompt_cache.get((p, classname, paramname))
-        if not node:
-            node = graph.node(classname)
-            node.set_input("clip", clip)
-            node.set_input(paramname, p)
-            prompt_cache[(p, classname, paramname)] = node
+        node = graph.node(classname)
+        node.set_input("clip", clip)
+        node.set_input(paramname, p)
         timestep = graph.node("ConditioningSetTimestepRange")
         timestep.set_input("conditioning", node.out(0))
         timestep.set_input("start", start_pct)
@@ -203,30 +232,20 @@ def build_scheduled_prompts(graph, schedules, clip):
         combiner.set_input("conditioning_2", othernode.out(0))
         node = combiner
 
-    return {"result": (node.out(0),), "expand": graph.finalize()}
+    g = graph.finalize()
+    log.debug("Built graph: %s", json.dumps(g))
+
+    return {"result": (node.out(0),), "expand": g}
 
 
-class PCLazyTextEncode:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {"clip": ("CLIP", {"rawLink": True}), "text": ("STRING", {"multiline": True})},
-            # "optional": {"defaults": ("SCHEDULE_DEFAULTS",)},
-            "hidden": {"unique_id": "UNIQUE_ID"},
-        }
-
-    RETURN_TYPES = ("CONDITIONING",)
-    OUTPUT_TOOLTIPS = ("A fully encoded and scheduled conditioning",)
-    CATEGORY = "promptcontrol"
-    FUNCTION = "apply"
-
-    def apply(self, clip, text, unique_id):
-        schedules = parse_prompt_schedules(text)
-        graph = GraphBuilder(f"PCEncodeLazy-{unique_id}")
-        return build_scheduled_prompts(graph, schedules, clip)
+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 PCLazyTextEncodeAdvanced:
+    CACHE_KEY = cache_key_prompt
+
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -235,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"},
         }
@@ -243,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).with_filters(start=start, end=end, filters=tags)
-        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()
         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,8 +1,80 @@
 import logging
+from .parser import parse_prompt_schedules, expand_macros
+from .nodes_lazy import NODE_CLASS_MAPPINGS as LAZY_NODES
+from .utils import expand_graph
+import json
+import folder_paths
+from pathlib import Path
 
 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": "PROMPT",
+            },
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(self, input_types):
+        return True
+
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+    CATEGORY = "promptcontrol/tools"
+    DESCRIPTION = "Expands lazy prompt control nodes in the prompt and saves the expanded 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
+        )
+        expanded = expand_graph(LAZY_NODES, prompt)
+        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(expanded, f)
+
+        return ()
+
+
+class PCSetLogLevel:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clip": ("CLIP",),
+            },
+            "optional": {
+                "level": (["INFO", "DEBUG", "WARNING", "ERROR"], {"default": "INFO"}),
+            },
+        }
+
+    def apply(self, clip, level="INFO"):
+        log.setLevel(getattr(logging, level))
+        log.info("Set logging level to %s", level)
+        return (clip,)
+
+    RETURN_TYPES = ("CLIP",)
+    CATEGORY = "promptcontrol/tools"
+    DESCRIPTION = (
+        "A debug node to configure Prompt Control logging level. Pass a CLIP through it before you run any PC nodes"
+    )
+
+    FUNCTION = "apply"
+
+
 class PCAddMaskToCLIP:
     @classmethod
     def INPUT_TYPES(s):
@@ -14,8 +86,9 @@ class PCAddMaskToCLIP:
         }
 
     RETURN_TYPES = ("CLIP",)
-    CATEGORY = "promptcontrol/v2"
+    CATEGORY = "promptcontrol/tools"
     FUNCTION = "apply"
+    DESCRIPTION = "Attaches a mask to a CLIP object so that they can be referred to in a prompt using IMASK(). Using this node multiple times adds more masks rather than replacing existing ones."
 
     def apply(self, clip, mask=None):
         return PCAddMaskToCLIPMany().apply(clip, mask1=mask)
@@ -35,8 +108,9 @@ class PCAddMaskToCLIPMany:
         }
 
     RETURN_TYPES = ("CLIP",)
-    CATEGORY = "promptcontrol/v2"
+    CATEGORY = "promptcontrol/tools"
     FUNCTION = "apply"
+    DESCRIPTION = "Multi-input version of PCAddMaskToCLIP, for convenience"
 
     def apply(self, clip, mask1=None, mask2=None, mask3=None, mask4=None):
         clip = clip.clone()
@@ -52,7 +126,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}),
@@ -65,13 +138,13 @@ class PCSetPCTextEncodeSettings:
         }
 
     RETURN_TYPES = ("CLIP",)
-    CATEGORY = "promptcontrol/v2"
+    CATEGORY = "promptcontrol/tools"
     FUNCTION = "apply"
+    DESCRIPTION = "Configures default values for PCTextEncode"
 
     def apply(
         self,
         clip,
-        steps=0,
         mask_width=512,
         mask_height=512,
         sdxl_width=1024,
@@ -82,7 +155,6 @@ class PCSetPCTextEncodeSettings:
         sdxl_crop_h=0,
     ):
         settings = {
-            "steps": steps,
             "mask_width": mask_width,
             "mask_height": mask_height,
             "sdxl_width": sdxl_width,
@@ -97,14 +169,63 @@ class PCSetPCTextEncodeSettings:
         return (clip,)
 
 
+class PCExtractScheduledPrompt:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "text": ("STRING", {"multiline": True}),
+                "at": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 1.0, "step": 0.01}),
+            },
+            "optional": {"tags": ("STRING", {"default": ""})},
+        }
+
+    RETURN_TYPES = ("STRING",)
+    CATEGORY = "promptcontrol/tools"
+    FUNCTION = "apply"
+    DESCRIPTION = "Parses the input prompt and returns the prompt scheduled at the specified point"
+
+    def apply(self, text, at, tags=""):
+        schedule = parse_prompt_schedules(text, filters=tags)
+        _, entry = schedule.at_step(at, total_steps=1)
+        prompt_text = entry.get("prompt", "")
+        return (prompt_text,)
+
+
+class PCMacroExpand:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "text": ("STRING", {"multiline": True}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    CATEGORY = "promptcontrol/tools"
+    FUNCTION = "apply"
+    DESCRIPTION = "Expands DEF macros in a string and returns the result"
+
+    def apply(self, text):
+        return (expand_macros(text),)
+
+
 NODE_CLASS_MAPPINGS = {
     "PCSetPCTextEncodeSettings": PCSetPCTextEncodeSettings,
     "PCAddMaskToCLIP": PCAddMaskToCLIP,
     "PCAddMaskToCLIPMany": PCAddMaskToCLIPMany,
+    "PCSetLogLevel": PCSetLogLevel,
+    "PCExtractScheduledPrompt": PCExtractScheduledPrompt,
+    "PCSaveExpandedWorkflow": PCSaveExpandedWorkflow,
+    "PCMacroExpand": PCMacroExpand,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "PCSetTextEncodeSettings": "PC: Configure PCTextEncode",
+    "PCSetPCTextEncodeSettings": "PC: Configure PCTextEncode",
     "PCAddMaskToCLIP": "PC: Attach Mask",
     "PCAddMaskToCLIPMany": "PC: Attach Mask (multi)",
+    "PCSetLogLevel": "PC: Configure Logging (for debug)",
+    "PCExtractScheduledPrompt": "PC: Extract Scheduled Prompt",
+    "PCSaveExpandedWorkflow": "PC: Save Expanded Workflow (for debug)",
+    "PCMacroExpand": "PC: Expand Macros",
 }

prompt_control/parser.py

@@ -1,29 +1,70 @@
+# vim: sw=4 ts=4
 import lark
 import logging
 from math import ceil
 
 logging.basicConfig()
 log = logging.getLogger("comfyui-prompt-control")
+import re
+
+from functools import lru_cache
+from .utils import get_function, find_closing_paren
 
 if lark.__version__ == "0.12.0":
-    x = "Your lark package reports an ancient version (0.12.0) and will not work. If you have the 'lark-parser' package in your Python environment, remove that and *reinstall* lark!"
+    from sys import executable
+
+    x = "\n".join(
+        [
+            "Your lark package reports an ancient version (0.12.0) and will not work. If you have the 'lark-parser' package in your Python environment, remove that and *reinstall* lark!",
+            f"{executable} -m pip uninstall lark-parser lark",
+            f"{executable} -m pip install lark",
+        ]
+    )
     log.error(x)
     raise ImportError(x)
 
 
+ESCAPES = [
+    ("XxPCBackslashESCAPExX", "\\"),
+    ("XxPCColonESCAPExX", ":"),
+    ("XxPCCommentESCAPExX", "#"),
+]
+
+
+def escape_specials(string):
+    for ph, c in ESCAPES:
+        string = string.replace(rf"\{c}", ph)
+    return string
+
+
+def restore_escaped(string):
+    for ph, c in ESCAPES:
+        string = string.replace(ph, c)
+    return string
+
+
+def remove_comments(string):
+    r = []
+    for line in string.split("\n"):
+        comment = line.find("#")
+        if comment >= 0:
+            r.append(line[:comment])
+        else:
+            r.append(line)
+    return "\n".join(r)
+
+
 prompt_parser = lark.Lark(
     r"""
 !start: (prompt | /[][():|]/+)*
-prompt: (emphasized | embedding | scheduled | alternate | sequence | interpolate | loraspec | PLAIN | /</ | />/ | WHITESPACE)+
+prompt: (emphasized | embedding | scheduled | alternate | sequence | loraspec | PLAIN | | /\\:/ | /</ | />/ | WHITESPACE)+
 !emphasized: "(" prompt? ")"
         | "(" prompt ":" prompt ")"
         | "[" prompt "]"
-scheduled: "[" [prompt ":"] [prompt] ":" _WS? NUMBER ["," NUMBER] "]"
-        | "[" [prompt ":"] [prompt] ":" _WS? TAG "]"
-sequence:  "[SEQ" ":" [prompt] ":" NUMBER (":" [prompt] ":" NUMBER)+ "]"
-interpolate.100: "[INT" ":" interp_prompts ":" interp_steps "]"
-interp_prompts: prompt (":" [prompt])+
-interp_steps: NUMBER ("," NUMBER)+ [":" 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
@@ -39,6 +80,7 @@ TAG: /[A-Z_]+/
     lexer="dynamic",
 )
 
+
 cut_parser = lark.Lark(
     r"""
 !start: (prompt | /[][:()]/+)*
@@ -80,7 +122,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:
@@ -92,14 +134,25 @@ def clamp(a, b, c):
     return min(max(a, b), c)
 
 
-def get_steps(tree):
-    res = [100]
-    interpolation_steps = []
+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):
@@ -116,55 +169,60 @@ def get_steps(tree):
             for i, _ in enumerate(tree.children[:-1]):
                 tree.children[i] = tostep(tree.children[i])
 
-            interpolation_steps.append((tuple(tree.children[:-1]), tree.children[-1]))
             res.extend(tree.children[:-1])
 
         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)
 
-    return sorted(set(interpolation_steps)), sorted(set(res))
+    return sorted(set(res))
 
 
 def at_step(step, filters, tree):
     class AtStep(lark.Transformer):
         def scheduled(self, args):
+            before = None
+            during = None
+            after = None
             when_end = None
-            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 ""
 
@@ -180,24 +238,6 @@ def at_step(step, filters, tree):
                     previous_step = s
             return ""
 
-        def interpolate(self, args):
-            prompts, starts = args
-            starts = starts[:-1]
-            prev_prompt = None
-            if step < starts[0]:
-                return prompts[0]
-            for i, x in enumerate(starts):
-                prev_prompt = prompts[i]
-                if x >= step:
-                    break
-            return prev_prompt
-
-        def interp_steps(self, args):
-            return list(args)
-
-        def interp_prompts(self, args):
-            return ["".join(flatten(a or [])) for a in args]
-
         def alternate(self, args):
             step_size = args[-1]
             idx = ceil(step / step_size)
@@ -231,7 +271,7 @@ def at_step(step, filters, tree):
             return {"prompt": p, "loras": loraspecs}
 
         def PLAIN(self, args):
-            return args.replace("\\:", ":")
+            return restore_escaped(args)
 
         def FILENAME(self, value):
             return str(value)
@@ -268,68 +308,47 @@ def at_step(step, filters, tree):
 
 
 class PromptSchedule(object):
-    def __init__(self, prompt, filters="", start=0.0, end=1.0, defaults=None, masks=None):
+    # 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.prompt = prompt.strip()
+        self.num_steps = num_steps
+        # placeholder is restored on parse
+        self.prompt = remove_comments(escape_specials(prompt.strip()))
         self.defaults = {}
-        if defaults:
-            self.defaults = defaults
         self.loaded_loras = {}
 
-        self.interpolations = None
-        self.parsed_prompt = None
-        self.interpolations, self.parsed_prompt = self._parse()
-        self.masks = masks
-        if masks is None:
-            self.masks = []
+        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 = []
-            interpolations = set()
             tree = prompt_parser.parse(self.prompt)
-            interpolation_steps, steps = get_steps(tree)
-            log.debug("Interpolation steps: %s", interpolation_steps)
+            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)
-                for control_points, step in interpolation_steps:
-                    interp_start = None
-                    interp_end = None
-                    if t == control_points[-1]:
-                        interp_start = max(control_points[0], int(self.start * 100))
-                        interp_end = min(control_points[-1], int(self.end * 100))
-                        control_points = tuple(
-                            sorted(set(f(c) for c in control_points if c >= interp_start or c <= interp_end))
-                        )
-                    if interp_start is not None and interp_end is not None and interp_end > interp_start:
-                        interpolations.add((control_points, f(step)))
                 parsed.append([f(t), p])
 
         except lark.exceptions.LarkError as e:
             log.error("Prompt editing parse error: %s", e)
             parsed = [[1.0, {"prompt": self.prompt, "loras": {}}]]
+            raise
 
         # Tag filtering may return redundant prompts, so filter them out here
         res = []
-        prev_p = None
         prev_end = -1
 
         for end_at, p in parsed:
-            # Preserve prompt if it ends at the start of an interpolation, otherwise bump its end time
-            if p == prev_p and res[-1][0] not in [x[0][0] for x in interpolations]:
-                res[-1][0] = end_at
-                continue
             if end_at < self.start:
                 continue
             elif end_at <= self.end:
@@ -338,18 +357,20 @@ class PromptSchedule(object):
             elif end_at > self.end and prev_end < self.end:
                 res.append([end_at, p])
                 break
-            prev_p = p
 
         # Always use the last prompt if everything was filtered
         if len(res) == 0:
             res = [[1.0, parsed[-1][1]]]
 
-        return interpolations, res
+        final = [res[0]]
 
-    def add_masks(self, *masks):
-        for mask in masks:
-            if mask is not None:
-                self.masks.append(mask)
+        # 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()
@@ -363,8 +384,7 @@ class PromptSchedule(object):
             filters=ifspecified(filters, self.filters),
             start=ifspecified(start, self.start),
             end=ifspecified(end, self.end),
-            defaults=ifspecified(defaults, self.defaults),
-            masks=self.masks[:],
+            num_steps=self.num_steps,
         )
         return p
 
@@ -378,23 +398,76 @@ class PromptSchedule(object):
                 return i, x
         return len(self.parsed_prompt) - 1, self.parsed_prompt[-1]
 
-    def interpolation_at(self, step, total_steps=1):
-        i, x = self.at_step_idx(step, total_steps)
-        for y in self.parsed_prompt[i:]:
-            step = min(y[0], 1.0)
-            if x[1]["prompt"] != y[1]["prompt"]:
-                return step, y
-        return 1.0, self.parsed_prompt[-1]
 
-    def load_loras(self, lora_cache=None):
-        from .utils import Timer, load_loras_from_schedule
+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 lora_cache is not None:
-            self.loaded_loras = lora_cache
-        with Timer("PromptSchedule.load_loras()"):
-            self.loaded_loras = load_loras_from_schedule(self.parsed_prompt, self.loaded_loras)
-        return self.loaded_loras
+    if not name:
+        return None
+    args = args.strip()
+    # If using the form DEF(F()=$1) then the default value of $1 is the empty string
+    if arg_start > 0:
+        args = [a.strip() for a in args.split(";")]
+    else:
+        args = []
+    return name, args
 
 
-def parse_prompt_schedules(prompt):
-    return PromptSchedule(prompt)
+def expand_macros(text):
+    text, defs = get_function(text, "DEF", defaults=None)
+    res = text
+    prevres = text
+    replacements = []
+    for d in defs:
+        r = d.split("=", 1)
+        search = parse_search(r[0].strip())
+        if not search or len(r) != 2:
+            log.warning("Ignoring invalid DEF(%s)", d)
+            continue
+        replacements.append((search, r[1].strip()))
+    iterations = 0
+    while True:
+        iterations += 1
+        if iterations > 10:
+            raise ValueError("Unable to resolve DEFs, make sure there are no cycles!")
+            return text
+        for search, replace in replacements:
+            res = substitute_defcall(res, search, replace)
+        if res == prevres:
+            break
+        prevres = res
+    if res.strip() != text.strip():
+        res = res.strip()
+        log.info("DEFs expanded to: %s", res)
+    return res
+
+
+def substitute_defcall(text, search, replace):
+    name, default_args = search
+    text, defns = get_function(text, name, defaults=None, placeholder=f"DEFNCALL{name}", require_args=False)
+    for i, parameters in enumerate(defns):
+        ph = f"\0DEFNCALL{name}{i}\0"
+        paramvals = []
+        if parameters is not None:
+            paramvals = [x.strip() for x in parameters.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 = expand_macros(prompt)
+    return PromptSchedule(prompt, **kwargs)

prompt_control/prompts.py

@@ -1,14 +1,18 @@
 import logging
 import re
 import torch
+import math
 from functools import partial
 from comfy_extras.nodes_mask import FeatherMask, MaskComposite
+from nodes import ConditioningAverage
 
-from .utils import safe_float, get_function, parse_floats
+from .utils import safe_float, get_function, split_by_function, parse_floats, smarter_split, call_node
 from .adv_encode import advanced_encode_from_tokens
 from .cutoff import process_cuts
 from .parser import parse_cuts
 
+from .attention_couple_ppm import set_cond_attnmask
+
 log = logging.getLogger("comfyui-prompt-control")
 
 AVAILABLE_STYLES = ["comfy", "perp", "A1111", "compel", "comfy++", "down_weight"]
@@ -62,13 +66,14 @@ def get_style(text, default_style="comfy", default_normalization="none"):
     style, normalization = styles[0]
     style = style.strip()
     normalization = normalization.strip()
-    if style not in AVAILABLE_STYLES:
+    if style.replace("old+", "") not in AVAILABLE_STYLES:
         log.warning("Unrecognized prompt style: %s. Using %s", style, default_style)
-        style = default_style
 
-    if normalization not in AVAILABLE_NORMALIZATIONS:
-        log.warning("Unrecognized prompt normalization: %s. Using %s", normalization, default_normalization)
-        normalization = default_normalization
+    for part in normalization.split("+"):
+        if part not in AVAILABLE_NORMALIZATIONS:
+            log.warning("Unrecognized prompt normalization: %s. Using %s", normalization, default_normalization)
+            normalization = default_normalization
+            break
 
     return style, normalization, text
 
@@ -88,8 +93,9 @@ def shuffle_chunk(shuffle, c):
         "separator": separator,
     }.get(joiner, joiner)
 
-    log.info("%s arg=%s sep=%s join=%s", func, shuffle_count, separator, joiner)
-    separated = c.split(separator)
+    log.debug("%s arg=%s sep=%s join=%s", func, shuffle_count, separator, joiner)
+    separated = smarter_split(separator, c)
+    log.debug("Prompt split into %s", separated)
     if func == "SHIFT":
         shuffle_count = shuffle_count % len(separated)
         permutation = separated[shuffle_count:] + separated[:shuffle_count]
@@ -122,6 +128,83 @@ def fix_word_ids(tokens):
     return tokens
 
 
+def tokenize_chunks(clip, text, need_word_ids, can_break):
+    chunks = re.split(r"\bBREAK\b", text)
+    token_chunks = []
+    shuffled_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)
+        shuffled_chunks.append(r)
+        t = clip.tokenize(c, return_word_ids=need_word_ids)
+        token_chunks.append(t)
+
+    tokens = token_chunks[0]
+    full_prompt = "".join(shuffled_chunks)
+    full_tokenized = tokens
+    if len(chunks) > 1:
+        full_tokenized = clip.tokenize(full_prompt, return_word_ids=need_word_ids)
+        for key in tokens:
+            if not can_break.get(key):
+                log.warning("BREAK does not make sense for %s, tokenizing as one chunk. Use CAT instead.", key)
+                tokens[key] = full_tokenized[key]
+                continue
+            for c in token_chunks[1:]:
+                tokens[key].extend(c[key])
+
+    return tokens
+
+
+def tokenize(clip, text, can_break, empty_tokens):
+    # defaults=None means there is no argument parsing at all
+    text, l_prompts = get_function(text, "CLIP_L", defaults=None)
+    text, te_prompts = get_function(text, "TE", defaults=None)
+    need_word_ids = True
+    tokens = tokenize_chunks(clip, text, need_word_ids, can_break)
+
+    per_te_prompts = {}
+    if l_prompts:
+        log.warning("Note: CLIP_L is deprecated. Use TE(l=prompt) instead")
+        per_te_prompts["l"] = l_prompts
+
+    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 per_te_prompts:
+        for key in per_te_prompts:
+            prompt = " ".join(per_te_prompts[key])
+            tokens[key] = tokenize_chunks(clip, prompt, need_word_ids, can_break)[key]
+            log.info("Encoded prompt with TE '%s': %s", key, prompt)
+
+    maxlen = max([0] + [len(tokens[k]) for k in tokens if can_break[k]])
+    for k in tokens:
+        if not can_break[k]:
+            continue
+        while len(tokens[k]) < maxlen:
+            tokens[k] += empty_tokens[k]
+
+    return fix_word_ids(tokens)
+
+
 def encode_prompt_segment(
     clip,
     text,
@@ -139,52 +222,63 @@ def encode_prompt_segment(
     if cuts:
         extra["cuts"] = cuts
 
-    # 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 = []
-    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]
+    empty = clip.tokenize("", return_word_ids=True)
+    can_break = {}
+    for k in empty:
+        tokenizer = getattr(clip.tokenizer, f"clip_{k}", getattr(clip.tokenizer, k, None))
+        can_break[k] = tokenizer and tokenizer.pad_to_max_length
 
-    for key in tokens:
-        for c in token_chunks[1:]:
-            tokens[key].extend(c[key])
+    clip = hook_te(clip, empty.keys(), style, normalization, extra)
 
-    # 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"]
+    # Chunks to ConditioningAverage:
 
-    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"]
+    text, averages = split_by_function(text, "AVG", ["0.5"], require_args=False)
+    prompts_to_avg = []
+    for avg in averages:
+        w = safe_float(avg["args"][0], 0.5)
+        prompts_to_avg.append((text, w))
+        text = avg["text"]
+    prompts_to_avg.append((text, 1.0))
 
-    tokens = fix_word_ids(tokens)
+    conds_to_avg = []
+    for prompt, weight in prompts_to_avg:
+        conds_to_cat = []
+        chunks = re.split(r"\bCAT\b", prompt)
+        for c in chunks:
+            tokens = tokenize(clip, c, can_break, empty)
+            conds_to_cat.append(clip.encode_from_tokens_scheduled(tokens, add_dict=settings))
 
-    tes = []
-    for k in tokens:
-        if k in ["g", "l"]:
-            tes.append(f"clip_{k}")
-        else:
-            tes.append(k)
+        base = conds_to_cat[0]
+        for cond in conds_to_cat[1:]:
+            assert len(cond) == len(base), "Conditioning length mismatch"
+            # Pooled gets ignored
+            for i in range(len(base)):
+                c1 = base[i][0]
+                c2 = cond[i][0]
+                base[i][0] = torch.cat((c1, c2), 1)
+        conds_to_avg.append((base, weight))
 
-    clip = hook_te(clip, tes, style, normalization, extra)
+    base, w = conds_to_avg[0]
+    for cond, next_w in conds_to_avg[1:]:
+        assert len(base) == len(cond), "Conditioning length mismatch"
+        if w == 1.0:
+            w = next_w
+            continue
+        for i in range(len(base)):
+            (cond,) = call_node(ConditioningAverage, [base[i]], [cond[i]], w)
+            base[i] = cond[0]
+        w = next_w
 
-    return clip.encode_from_tokens_scheduled(tokens, add_dict=settings)
+    return base
+
+
+def calc_w(tensor, w):
+    if math.isclose(w, 0):
+        return torch.zeros_like(tensor)
+    elif math.isclose(w, 1.0):
+        return tensor
+    else:
+        return tensor * w
 
 
 def apply_weights(output, te_name, spec):
@@ -195,21 +289,29 @@ 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 = calc_w(out, w)
+            if pooled is not None:
+                pooled = calc_w(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 = calc_w(output, w)
         return output
 
 
@@ -230,15 +332,24 @@ def hook_te(clip, te_names, style, normalization, extra):
         return clip
     newclip = clip.clone()
     for te_name in te_names:
-        if hasattr(clip.patcher.model, te_name):
+        tokenizer = getattr(clip.tokenizer, f"clip_{te_name}", getattr(clip.tokenizer, te_name, None))
+        if tokenizer:
             x = extra.copy()
-            x["tokenizer"] = getattr(clip.tokenizer, te_name)
-            log.debug("Hooked into %s with style=%s, normalization=%s", te_name, style, normalization)
+            x["tokenizer"] = tokenizer
+            if not hasattr(clip.patcher.model, te_name):
+                te_name = "clip_" + te_name
+            if not hasattr(clip.patcher.model, te_name):
+                log.warning("TE model %s not found on model patcher. Skipping...", te_name)
+                continue
+
+            log.debug("Hooked into te=%s with style=%s, normalization=%s", te_name, style, normalization)
+            encode = clip.patcher.get_model_object(f"{te_name}.encode_token_weights")
+            x["has_negpip"] = clip.patcher.model_options.get("ppm_negpip", False)
             newclip.patcher.add_object_patch(
                 f"{te_name}.encode_token_weights",
                 make_patch(
                     te_name,
-                    clip.patcher.get_model_object(f"{te_name}.encode_token_weights"),
+                    encode,
                     normalization,
                     style,
                     x,
@@ -246,7 +357,7 @@ def hook_te(clip, te_names, style, normalization, extra):
             )
         # 'g' and 'l' exist in these are clip_g and clip_l
         else:
-            log.debug("Tokens contain items with key %s but no TE found on object with that name.", te_name)
+            log.warning("Tokens contain items with key %s but no tokenizer found on object with that name.", te_name)
     return newclip
 
 
@@ -312,7 +423,7 @@ def make_mask(args, size, weight):
     mask = torch.full((h, w), 0, dtype=torch.float32, device="cpu")
     mask[ys[0] : ys[1], xs[0] : xs[1]] = weight
     mask = mask.unsqueeze(0)
-    log.info("Mask xs=%s, ys=%s, shape=%s, weight=%s", xs, ys, mask.shape, weight)
+    log.debug("Mask xs=%s, ys=%s, shape=%s, weight=%s", xs, ys, mask.shape, weight)
     return mask
 
 
@@ -328,7 +439,7 @@ def get_mask(text, size, input_masks):
 
     def feather(f, mask):
         l, t, r, b, *_ = [int(x) for x in parse_floats(f[0], [0, 0, 0, 0], split_re="\\s+")]
-        mask = FeatherMask().feather(mask, l, t, r, b)[0]
+        mask = call_node(FeatherMask, mask, l, t, r, b)[0]
         log.info("FeatherMask l=%s, t=%s, r=%s, b=%s", l, t, r, b)
         return mask
 
@@ -346,7 +457,7 @@ def get_mask(text, size, input_masks):
         i += 1
         if mask is not None:
             log.info("MaskComposite op=%s", op)
-            mask = MaskComposite().combine(mask, nextmask, 0, 0, op)[0]
+            mask = call_node(MaskComposite, mask, nextmask, 0, 0, op)[0]
         else:
             mask = nextmask
 
@@ -361,7 +472,7 @@ def get_mask(text, size, input_masks):
             nextmask = feather(feathers[i], nextmask)
         i += 1
         if mask is not None:
-            mask = MaskComposite().combine(mask, nextmask, 0, 0, op)[0]
+            mask = call_node(MaskComposite, mask, nextmask, 0, 0, op)[0]
         else:
             mask = nextmask
 
@@ -402,6 +513,49 @@ def apply_noise(cond, weight, gen):
     return cond * (1 - weight) + n * weight
 
 
+def process_settings(prompt, defaults, masks, mask_size, sdxl_opts):
+    if "ATTN()" in prompt:
+        raise ValueError("ATTN() no longer works and has been replaced by COUPLE()")
+
+    def weight(t):
+        opts = {}
+        m = re.search(r":(-?\d\.?\d*)(![A-Za-z]+)?$", t.strip())
+        if not m:
+            return (None, opts, t)
+        w = float(m[1])
+        tag = m[2]
+        t = t[: m.span()[0]]
+        if tag == "!noscale":
+            opts["scale"] = 1
+
+        return w, opts, t
+
+    settings = {"prompt": prompt}
+
+    if "FILL()" in prompt:
+        prompt = prompt.replace("FILL()", "")
+        settings["x-promptcontrol.fill"] = True
+    prompt, mask, mask_weight = get_mask(prompt, mask_size, masks)
+    prompt, noise_w, generator = get_noise(prompt)
+    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 w is not None:
+        settings["strength"] = w
+    settings.update(sdxl_opts)
+    settings.update(local_sdxl_opts)
+    if area:
+        settings["area"] = area[0]
+        settings["strength"] = area[1]
+        settings["set_area_to_bounds"] = False
+    if mask is not None:
+        settings["mask"] = mask
+        settings["mask_strength"] = mask_weight
+
+    return prompt, settings
+
+
 def encode_prompt(clip, text, start_pct, end_pct, defaults, masks):
     # First style modifier applies to ANDed prompts too unless overridden
     style, normalization, text = get_style(text)
@@ -412,45 +566,61 @@ def encode_prompt(clip, text, start_pct, end_pct, defaults, masks):
     p, sdxl_opts = get_sdxl(prompts[0], defaults)
     prompts[0] = p
 
-    def weight(t):
-        opts = {}
-        m = re.search(r":(-?\d\.?\d*)(![A-Za-z]+)?$", t)
-        if not m:
-            return (1.0, opts, t)
-        w = float(m[1])
-        tag = m[2]
-        t = t[: m.span()[0]]
-        if tag == "!noscale":
-            opts["scale"] = 1
-
-        return w, opts, t
-
     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))
-    for prompt in prompts:
-        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)
-        settings = {"prompt": prompt}
-        settings["strength"] = w
-        settings.update(sdxl_opts)
-        settings.update(local_sdxl_opts)
-        if area:
-            settings["area"] = area[0]
-            settings["strength"] = area[1]
-            settings["set_area_to_bounds"] = False
-        if mask is not None:
-            settings["mask"] = mask
-            settings["mask_strength"] = mask_weight
 
-        settings["start_percent"] = start_pct
-        settings["end_percent"] = end_pct
-        x = encode_prompt_segment(clip, prompt, settings, style, normalization)
-        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
+
+    def couple_mask(args):
+        if args is None:
+            return ""
+        return f"MASK({args})"
+
+    for prompt in prompts:
+        base_prompt, attn_couple_prompts = split_by_function(prompt, "COUPLE", defaults=None, require_args=False)
+
+        prompts = [base_prompt] + [couple_mask(p["args"]) + p["text"] for p in attn_couple_prompts]
+        encoded = []
+        for p in prompts:
+            p, settings = process_settings(p, defaults, masks, mask_size, sdxl_opts)
+            if settings.get("strength") == 0:  # weight is explicitly set to 0, skip
+                continue
+            settings["start_percent"] = start_pct
+            settings["end_percent"] = end_pct
+            x = encode_prompt_segment(clip, p, settings, style, normalization)
+            encoded.append(x)
+
+        assert all(
+            len(c) == len(encoded[0]) for c in encoded
+        ), "All encoded prompts didn't produce the same number of conds, I don't know what to do in this situation."
+
+        # each call to encode_prompt_segment can produce a number of conds based on any
+        # scheduled LoRA hooks on the clip model. Zip them together with coupled prompts
+        base_cond = []
+        for base_cond, *attention_couple in zip(*encoded):
+            s = base_cond[1]
+            # If there are LoRAs on the CLIP, we need to fix start_percent and end_percent on the new conds for things to work properly.
+            s["start_percent"] = s.get("clip_start_percent", s["start_percent"])
+            s["end_percent"] = s.get("clip_end_percent", s["end_percent"])
+            s.pop("clip_start_percent", None)
+            s.pop("clip_end_percent", None)
+            base_cond = [base_cond]
+            if attention_couple:
+                fill = base_cond[0][1].get("x-promptcontrol.fill")
+                if not fill:
+                    ensure_mask(base_cond[0])
+                # else, set_cond_attnmask will have the base mask fill any unspecified areas
+                base_cond = set_cond_attnmask(
+                    base_cond,
+                    [ensure_mask(c) for c in attention_couple],
+                    fill=fill,
+                )
+        conds.extend(base_cond)
 
     return conds

prompt_control/test_encode.py

@@ -0,0 +1,192 @@
+import unittest
+import unittest.mock as mock
+import numpy.testing as npt
+from os import environ
+import nodes
+import comfy_extras.nodes_mask
+from .nodes_base import PCTextEncode
+
+clips = []
+
+import logging
+
+logging.basicConfig()
+
+
+def run(f, *args):
+    if hasattr(f, "execute"):
+        return f.execute(*args)
+    else:
+        return getattr(f, f.FUNCTION)(*args)
+
+
+@mock.patch("torch.cuda.current_device", lambda: "cpu")
+class TestEncode(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        global clips
+        print("Loading ComfyUI")
+        from comfy.sd import load_clip
+        from pathlib import Path
+
+        to_test = environ.get("TEST_TE", "clip_l").split()
+        model_dir = environ.get("COMFYUI_TE_DIR", ".")
+
+        te_root = Path(model_dir).resolve()
+
+        if "clip_l" in to_test:
+            clip_l = load_clip(
+                ckpt_paths=[str(te_root / "clip_l.safetensors")], clip_type="stable_diffusion", model_options={}
+            )
+            clips.append(("clip_l", clip_l))
+
+        if "t5" in to_test:
+            dual = load_clip(
+                [str(te_root / "clip_l.safetensors"), str(te_root / "t5xxl_fp16.safetensors")],
+                clip_type="flux",
+                model_options={},
+            )
+            clips.append(("clip_l+t5", dual))
+
+        print("Starting tests")
+
+    def tensorsEqual(self, t1, t2):
+        npt.assert_equal(t1.detach().numpy(), t2.detach().numpy())
+
+    def condEqual(self, c1, c2, key=None, key_assert=None):
+        self.assertEqual(len(c1), len(c2))
+        for i in range(len(c1)):
+            a, b = c1[i], c2[i]
+            if key:
+                (key_assert or self.assertEqual)(a[1].get(key), b[1].get(key))
+            else:
+                self.tensorsEqual(a[0], b[0])
+
+    def test_basic_encode(self):
+        pc = PCTextEncode()
+        comfy = nodes.CLIPTextEncode()
+        combine = nodes.ConditioningCombine()
+        average = nodes.ConditioningAverage()
+        concat = nodes.ConditioningConcat()
+        zeroout = nodes.ConditioningZeroOut()
+        for k, clip in clips:
+            with self.subTest(k):
+                with self.subTest("No exceptions"):
+                    run(
+                        pc,
+                        clip,
+                        "test AND test (test:1.2) BREAK test AND TE_WEIGHT(all=0) SDXL() AND AREA(,,) test CAT test",
+                    )
+                with self.subTest("Basic"):
+                    (c1,) = run(pc, clip, "test")
+                    (c2,) = run(comfy, clip, "test")
+                    c = c2  # Used in later tests
+                    self.condEqual(c1, c2)
+
+                with self.subTest("Function cornercase"):
+                    (c1,) = run(pc, clip, "test SDXL function")
+                    (c2,) = run(comfy, clip, "test SDXL function")
+                    (c3,) = run(pc, clip, "test SDXL() function")
+                    self.condEqual(c1, c2)
+
+                with self.subTest("Weights"):
+                    (c1,) = run(pc, clip, "(test:1.2) (test:0.6)")
+                    (c2,) = run(comfy, clip, "(test:1.2) (test:0.6)")
+                    self.condEqual(c1, c2)
+
+                with self.subTest("Concat"):
+                    (c1,) = run(pc, clip, "test CAT test")
+                    (c2,) = run(concat, c, c)
+                    self.condEqual(c1, c2)
+
+                with self.subTest("Combine"):
+                    (c1,) = run(pc, clip, "test AND test")
+                    (c2,) = run(combine, c, c)
+                    self.condEqual(c1, c2)
+
+                with self.subTest("Zero out"):
+                    (c1,) = run(pc, clip, "test TE_WEIGHT(all=0)")
+                    (c2,) = run(zeroout, c)
+                    self.condEqual(c1, c2)
+
+                with self.subTest("Average"):
+                    (c1,) = run(comfy, clip, "test1")
+                    (c2,) = run(comfy, clip, "test2")
+                    (c3,) = run(pc, clip, "test1 AVG() test2")
+                    (c4,) = run(pc, clip, "test1 AVG test2")
+                    (avg,) = run(average, c1, c2, 0.5)
+                    self.condEqual(avg, c3)
+                    self.condEqual(avg, c4)
+
+    @unittest.expectedFailure
+    def test_failure(self):
+        pc = PCTextEncode()
+        comfy = nodes.CLIPTextEncode()
+        for k, clip in clips:
+            with self.subTest(k):
+                (c1,) = run(comfy, clip, "test SDXL function")
+                (c2,) = run(pc, clip, "test SDXL() function")
+                self.condEqual(c1, c2)
+
+    def test_weight(self):
+        pc = PCTextEncode()
+        comfy = nodes.CLIPTextEncode()
+        combine = nodes.ConditioningCombine()
+        strength = nodes.ConditioningSetAreaStrength()
+        for k, clip in clips:
+            (c,) = run(comfy, clip, "test")
+            (c2,) = run(strength, c, 0.5)
+            with self.subTest(f"Testing {k}"):
+                with self.subTest("Conditioning weights"):
+                    (a,) = run(pc, clip, "test :0.5 AND test :0.5")
+                    (b,) = run(combine, c2, c2)
+                    self.condEqual(a, b)
+                    self.condEqual(a, b, "strength")
+                with self.subTest("Weight == 0"):
+                    (a,) = run(pc, clip, "test :0.5 AND test :0 AND test")
+                    (b,) = run(combine, c2, c)
+                    self.condEqual(a, b)
+                    self.condEqual(a, b, "strength")
+
+    def test_attn_couple(self):
+        pc = PCTextEncode()
+        for k, clip in clips:
+            with self.subTest(f"Testing {k}"):
+                (c,) = run(pc, clip, "test COUPLE prompt1 AND test2 COUPLE prompt2")
+                (c2,) = run(pc, clip, "test COUPLE prompt1 COUPLE test2 COUPLE prompt2")
+                self.assertTrue(len(c) == 2)
+                self.assertTrue(len(c2) == 1)
+
+    def test_styles(self):
+        pc = PCTextEncode()
+        comfy = nodes.CLIPTextEncode()
+        for k, clip in clips:
+            (no_weights,) = run(comfy, clip, "this prompt has no weights")
+            for style in ["comfy", "A1111", "comfy++", "compel", "down_weight", "perp"]:
+                with self.subTest(f"TE {k} style {style} no weights equal comfy"):
+                    (c,) = run(pc, clip, "this prompt has no weights")
+                    self.condEqual(no_weights, c)
+                with self.subTest(f"TE {k} style {style} does not fail when encoding weights"):
+                    for normalization in ["none", "mean", "length", "mean+length", "length+mean"]:
+                        with self.subTest(f"TE {k} style {style} normalization {normalization}"):
+                            (c,) = run(
+                                pc,
+                                clip,
+                                f"STYLE({style}, {normalization}) (this prompt) (has weights:0.9), (a:1.2) (b:1.2)",
+                            )
+
+    def test_masks(self):
+        pc = PCTextEncode()
+        comfy = nodes.CLIPTextEncode()
+        solidmask = comfy_extras.nodes_mask.SolidMask()
+        setMask = nodes.ConditioningSetMask()
+        for k, clip in clips:
+            (c1,) = run(pc, clip, "test MASK()")
+            (c2,) = run(comfy, clip, "test")
+            (c2,) = run(setMask, c2, run(solidmask, 1.0, 512, 512)[0], "default", 1.0)
+            self.condEqual(c1, c2)
+            self.condEqual(c1, c2, "mask", self.tensorsEqual)
+
+
+if __name__ == "__main__":
+    unittest.main()

prompt_control/test_graph.py

@@ -0,0 +1,244 @@
+import unittest
+import unittest.mock as mock
+import logging
+
+log = logging.getLogger("comfyui-prompt-control")
+
+
+def reset_graphbuilder_state():
+    from comfy_execution.graph_utils import GraphBuilder
+
+    GraphBuilder.set_default_prefix("UID", 0, 0)
+
+
+def find_file(name):
+    names = {"test": "test.safetensors", "other": "some/other.safetensors"}
+    return names.get(name)
+
+
+def loraloader(text, adv=False, **kwargs):
+    from .nodes_lazy import PCLazyLoraLoader, PCLazyLoraLoaderAdvanced
+
+    reset_graphbuilder_state()
+    if adv:
+        cls = PCLazyLoraLoader
+    else:
+        cls = PCLazyLoraLoaderAdvanced
+    model = [0, 1]
+    clip = [0, 0]
+    return cls().apply(unique_id="UID", model=model, clip=clip, text=text, **kwargs)
+
+
+def te(text, adv=False, **kwargs):
+    from .nodes_lazy import PCLazyTextEncode, PCLazyTextEncodeAdvanced
+
+    if adv:
+        cls = PCLazyTextEncode
+    else:
+        cls = PCLazyTextEncodeAdvanced
+    reset_graphbuilder_state()
+    clip = [0, 0]
+    return cls().apply(clip=clip, text=text, unique_id="UID", **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):
+        for p in ["test", "[test:0.2] test", "[test[test::0.5]]<lora:test:1>"]:
+            r1 = te(p)
+            r2 = te(p, adv=True)
+            with self.subTest(f"Expansion: {p}"):
+                self.assertEqual(r1, r2)
+
+        reset_graphbuilder_state()
+        with self.subTest("Expansion: LoRA"):
+            r = te("test<lora:test:1>")
+            self.assertEqual(
+                r,
+                {
+                    "result": (["UID.0.0.2", 0],),
+                    "expand": {
+                        "UID.0.0.1": {"class_type": "PCTextEncode", "inputs": {"clip": [0, 0], "text": "test"}},
+                        "UID.0.0.2": {
+                            "class_type": "ConditioningSetTimestepRange",
+                            "inputs": {"conditioning": ["UID.0.0.1", 0], "start": 0.0, "end": 1.0},
+                        },
+                    },
+                },
+            )
+        with self.subTest("Expansion: LoRA with schedule"):
+            r = te("simple [test:0.1,0.5] prompt<lora:test:1>")
+            self.assertEqual(
+                r,
+                {
+                    "result": (["UID.0.0.8", 0],),
+                    "expand": {
+                        "UID.0.0.1": {
+                            "class_type": "PCTextEncode",
+                            "inputs": {"clip": [0, 0], "text": "simple  prompt"},
+                        },
+                        "UID.0.0.2": {
+                            "class_type": "ConditioningSetTimestepRange",
+                            "inputs": {"conditioning": ["UID.0.0.1", 0], "start": 0.0, "end": 0.1},
+                        },
+                        "UID.0.0.3": {
+                            "class_type": "PCTextEncode",
+                            "inputs": {"clip": [0, 0], "text": "simple test prompt"},
+                        },
+                        "UID.0.0.4": {
+                            "class_type": "ConditioningSetTimestepRange",
+                            "inputs": {"conditioning": ["UID.0.0.3", 0], "start": 0.1, "end": 0.5},
+                        },
+                        "UID.0.0.5": {
+                            "class_type": "PCTextEncode",
+                            "inputs": {"clip": [0, 0], "text": "simple  prompt"},
+                        },
+                        "UID.0.0.6": {
+                            "class_type": "ConditioningSetTimestepRange",
+                            "inputs": {"conditioning": ["UID.0.0.5", 0], "start": 0.5, "end": 1.0},
+                        },
+                        "UID.0.0.7": {
+                            "class_type": "ConditioningCombine",
+                            "inputs": {"conditioning_1": ["UID.0.0.2", 0], "conditioning_2": ["UID.0.0.4", 0]},
+                        },
+                        "UID.0.0.8": {
+                            "class_type": "ConditioningCombine",
+                            "inputs": {"conditioning_1": ["UID.0.0.7", 0], "conditioning_2": ["UID.0.0.6", 0]},
+                        },
+                    },
+                },
+            )
+
+    @mock.patch("prompt_control.utils.lora_name_to_file", find_file)
+    def test_loraloader(self):
+        with self.assertLogs(log, level="WARNING") as cm:
+            result = loraloader("prompt here <lora:nonexistent:1.0:0.5>")["expand"]
+            result_adv = loraloader("prompt here <lora:nonexistent:1.0:0.5>", adv=True)["expand"]
+        self.assertIn("LoRA 'nonexistent' not found", cm.output[0])
+        self.assertEqual(result, {})
+        self.assertEqual(result_adv, {})
+
+        result = loraloader("<lora:test:1>")["expand"]
+        result2 = loraloader("prompt here <lora:test:1.0:0.5><lora:test:0:0.5>")["expand"]
+        result3 = loraloader("prompt here <lora:test:1.0:0.5><lora:test:0:0.5>", adv=True)["expand"]
+        self.assertEqual(result, result2)
+        self.assertEqual(result2, result3)
+        self.assertEqual(
+            result,
+            {
+                "UID.0.0.1": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": [0, 1],
+                        "clip": [0, 0],
+                        "strength_model": 1.0,
+                        "strength_clip": 1.0,
+                        "lora_name": "test.safetensors",
+                    },
+                }
+            },
+        )
+        result = loraloader("<lora:test:1><lora:other:0.5>")["expand"]
+        self.assertEqual(
+            result,
+            {
+                "UID.0.0.1": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": [0, 1],
+                        "clip": [0, 0],
+                        "strength_model": 1.0,
+                        "strength_clip": 1.0,
+                        "lora_name": "test.safetensors",
+                    },
+                },
+                "UID.0.0.2": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": ["UID.0.0.1", 0],
+                        "clip": ["UID.0.0.1", 1],
+                        "strength_model": 0.5,
+                        "strength_clip": 0.5,
+                        "lora_name": "some/other.safetensors",
+                    },
+                },
+            },
+        )
+
+        result = loraloader("prompt here <lora:test:1.0:0.5>")["expand"]
+        self.assertEqual(
+            result,
+            {
+                "UID.0.0.1": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": [0, 1],
+                        "clip": [0, 0],
+                        "strength_model": 1.0,
+                        "strength_clip": 0.5,
+                        "lora_name": "test.safetensors",
+                    },
+                }
+            },
+        )
+
+        result = loraloader("prompt [<lora:test:0.5>:0.5]")["expand"]
+        result2 = loraloader("prompt [<lora:test:0.5>:0.5]", adv=True)["expand"]
+        self.assertEqual(result, result2)
+        expected = {
+            "UID.0.0.1": {
+                "class_type": "CreateHookLora",
+                "inputs": {"lora_name": "test.safetensors", "strength_model": 0.5, "strength_clip": 0.5},
+            },
+            "UID.0.0.2": {
+                "class_type": "CreateHookKeyframe",
+                "inputs": {"strength_mult": 0.0, "start_percent": 0.0},
+            },
+            "UID.0.0.3": {
+                "class_type": "CreateHookKeyframe",
+                "inputs": {
+                    "start_percent": 0.5,
+                    "prev_hook_kf": ["UID.0.0.2", 0],
+                    "strength_mult": 1.0,
+                },
+            },
+            "UID.0.0.4": {
+                "class_type": "SetHookKeyframes",
+                "inputs": {"hooks": ["UID.0.0.1", 0], "hook_kf": ["UID.0.0.3", 0]},
+            },
+            "UID.0.0.5": {
+                "class_type": "SetClipHooks",
+                "inputs": {
+                    "clip": [0, 0],
+                    "hooks": ["UID.0.0.4", 0],
+                    "apply_to_conds": True,
+                    "schedule_clip": True,
+                },
+            },
+        }
+        self.assertEqual(result, expected)
+        result2 = loraloader("prompt [<lora:test:0.5>:0.5]", adv=True, start=0.6)["expand"]
+        self.assertEqual(
+            result2,
+            {
+                "UID.0.0.1": {
+                    "class_type": "LoraLoader",
+                    "inputs": {
+                        "model": [0, 1],
+                        "clip": [0, 0],
+                        "strength_model": 0.5,
+                        "strength_clip": 0.5,
+                        "lora_name": "test.safetensors",
+                    },
+                }
+            },
+        )
+        result2 = loraloader("prompt [<lora:test:0.5>:0.5]", end=0.5)["expand"]
+        self.assertEqual(result2, {})
+
+
+if __name__ == "__main__":
+    unittest.main()

prompt_control/test_parser.py

@@ -0,0 +1,240 @@
+import unittest
+from .parser import parse_prompt_schedules as parse, expand_macros
+
+
+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]"]],
+            [parse(p) for p in ["[before:during:after:0.1]", "[before:during:after:0.1,1.0]", "[before:during:0.1]"]],
+            [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]"]],
+            [parse(p) for p in ["[a:b:0.5]", "[a::b:0.5,0.5]"]],
+            [parse(p) for p in ["[a::0.5]", "[a:::0.5,0.5]"]],
+        ]
+        for group in eqs:
+            for p in group[1:]:
+                with self.subTest(p):
+                    self.assertEqual(group[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]")
+        with self.subTest("parameters"):
+            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]")
+        with self.subTest("two functions"):
+            self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+
+        p = expand_macros("DEF(X(a;b)=$1 $2 $3 d)X(A) X(A;B;C)")
+        with self.subTest("defaults"):
+            self.assertEqual(p, "A b $3 d A B C d")
+
+        p = expand_macros("DEF(MACRO()=[empty:$1:$2])MACRO MACRO(;) MACRO(;0.5) MACRO(a;0.5)")
+        with self.subTest("Empty default for $1"):
+            self.assertEqual(p, "[empty::$2] [empty::] [empty::0.5] [empty:a:0.5]")
+
+        p = expand_macros("DEF(X=$1)DEF(Y()=$1)[X Y][X() Y()][X(1) Y(1)]")
+        with self.subTest("defaults, DEF=X vs DEF=X()"):
+            self.assertEqual(p, "[$1 ][ ][1 1]")
+
+        p = parse("DEF(test(1)=prompt $1)DEF(test2((a); (test))=[$1:$2:0.5])test test2")
+        p2 = parse("prompt 1 [(a):(prompt 1):0.5]")
+        with self.subTest("defaults, nested parens"):
+            self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+
+        with self.assertRaises(ValueError) as c:
+            expand_macros("DEF(X=recurse Y) DEF(Y=recurse X) X")
+        self.assertTrue("Unable to resolve DEFs" in str(c.exception))
+
+    def test_escapes(self):
+        p = parse(r"[a:\:a:0.5] :\[a:b:0.5]")
+        self.assertPrompt(p, 0, 0.5, r"a :\[a:b:0.5]")
+        self.assertPrompt(p, 0.55, 1, r":a :\[a:b:0.5]")
+
+        p = parse(r"[embedding\:a:embedding\:b:0.1,0.5]")
+        self.assertPrompt(p, 0.15, 0.5, r"embedding:a")
+        self.assertPrompt(p, 0.55, 1, r"embedding:b")
+
+        p = parse(r"[embedding\:a:embedding\:b:embedding\:c:0.1,0.5]")
+        self.assertPrompt(p, 0.0, 0.1, r"embedding:a")
+        self.assertPrompt(p, 0.15, 0.5, r"embedding:b")
+        self.assertPrompt(p, 0.55, 1, r"embedding:c")
+
+        p = parse(r"[a\:b\\:c:0.5]")
+        self.assertPrompt(p, 0.0, 0.5, "a:b\\")
+        self.assertPrompt(p, 0.55, 1, r"c")
+
+        p = parse(r"[a:\#b:0.5]")
+        self.assertPrompt(p, 0.0, 0.5, "a")
+        self.assertPrompt(p, 0.55, 1, "#b")
+
+    def test_comments(self):
+        p = parse("this is a # comment")
+        self.assertPrompt(p, 0, 1.0, "this is a ")
+        p = parse("this is a [comment#:scheduled:0.6]")
+        self.assertPrompt(p, 0, 1.0, "this is a [comment")
+        p = parse(r"this is a [comment\#:scheduled:0.6]")
+        self.assertPrompt(p, 0, 0.6, "this is a comment#")
+        self.assertPrompt(p, 0.65, 1.0, "this is a scheduled")
+        p = parse("#this is a comment\nthis is a prompt")
+        self.assertPrompt(p, 0, 1.0, "\nthis is a prompt")
+
+    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)
+            with self.subTest(step):
+                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)
+            with self.subTest(step):
+                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

@@ -1,12 +1,29 @@
 from pathlib import Path
 import re
 import logging
+import copy
+
+# Allow testing
+try:
+    from folder_paths import get_filename_list
+except ImportError:
+
+    def get_filename_list(x):
+        raise NotImplementedError("How did you get here?")
 
-import folder_paths
 
 log = logging.getLogger("comfyui-prompt-control")
 
 
+def call_node(cls, *args, **kwargs):
+    if hasattr(cls, "execute"):
+        # v3 node
+        return cls.execute(*args, **kwargs)
+    else:
+        func = getattr(cls(), cls.FUNCTION)
+        return func(*args, **kwargs)
+
+
 def consolidate_schedule(prompt_schedule):
     prev_loras = {}
     not_found = []
@@ -35,7 +52,6 @@ def find_nonscheduled_loras(consolidated_schedule):
     if not consolidated_schedule:
         return {}
     last_end, candidate_loras = consolidated_schedule[0]
-    print(candidate_loras)
     to_remove = set()
     for candidate, weights in candidate_loras.items():
         for end, loras in consolidated_schedule[1:]:
@@ -48,6 +64,26 @@ def find_nonscheduled_loras(consolidated_schedule):
     return {k: v for (k, v) in candidate_loras.items() if k not in to_remove}
 
 
+def smarter_split(separator, string):
+    """Does not break () when splitting"""
+    splits = []
+    prev = 0
+    stack = 0
+    escape = False
+    for idx, x in enumerate(string):
+        if x == "(" and not escape:
+            stack += 1
+        elif x == ")" and not escape:
+            stack = max(0, stack - 1)
+        elif x == separator and stack == 0:
+            splits.append(string[prev:idx])
+            prev = idx + 1
+        escape = x == "\\"
+
+    splits.append(string[prev : idx + 1])
+    return splits
+
+
 def find_closing_paren(text, start):
     stack = 1
     for i, char in enumerate(text[start:]):
@@ -61,26 +97,70 @@ def find_closing_paren(text, start):
     return len(text)
 
 
-def get_function(text, func, defaults, return_func_name=False):
-    rex = re.compile(rf"\b{func}\(", re.MULTILINE)
+def get_function(text, func, defaults, return_func_name=False, placeholder="", return_dict=False, require_args=True):
+    if require_args:
+        rex = re.compile(rf"\b{func}\(", re.MULTILINE)
+    else:
+        rex = re.compile(rf"\b{func}\b", re.MULTILINE)
     instances = []
     match = rex.search(text)
+    count = 0
     while match:
         # Match start, content start
-        start, after_first_paren = match.span()
-        funcname = text[start : after_first_paren - 1]
-        end = find_closing_paren(text, after_first_paren)
-        args = parse_strings(text[after_first_paren:end], defaults)
-        if return_func_name:
+        start, at_paren = match.span()
+        if require_args:
+            at_paren = at_paren - 1
+        funcname = text[start:at_paren]
+        after_first_paren = at_paren + 1
+        if text[at_paren:after_first_paren] == "(":
+            end = find_closing_paren(text, after_first_paren)
+            args = parse_strings(text[after_first_paren:end], defaults)
+            end += 1
+        else:
+            end = at_paren
+            args = defaults
+        ph = None
+        if placeholder:
+            ph = f"\0{placeholder}{count}\0"
+        if return_dict:
+            instances.append(
+                {
+                    "name": funcname,
+                    "args": args,
+                    "position": start,
+                    "placeholder": ph,
+                }
+            )
+        elif return_func_name:
             instances.append((funcname, args))
         else:
             instances.append(args)
 
-        text = text[:start] + text[end + 1 :]
+        if placeholder:
+            text = text[:start] + f"\0{placeholder}{count}\0" + text[end:]
+        else:
+            text = text[:start] + text[end:]
         match = rex.search(text)
+        count += 1
     return text, instances
 
 
+def split_by_function(text, func, defaults=None, require_args=True):
+    """
+    Splits a string by function calls, returning the text preceding the first call and a list of dictionaries with a "text" key with the prompt before the next split or until hthe end of the text.
+    """
+    text, functions = get_function(text, func, defaults, return_dict=True, require_args=require_args)
+    chunks = []
+    prev = 0
+    for f in functions:
+        chunks.append(text[prev : f["position"]])
+        prev = f["position"]
+    chunks.append(text[prev:])
+    for i, f in enumerate(functions):
+        f["text"] = chunks[i + 1]
+    return chunks[0], functions
+
+
 def parse_args(strings, arg_spec, strip=True):
     args = [s[1] for s in arg_spec]
     for i, spec in list(enumerate(arg_spec))[: len(strings)]:
@@ -119,7 +199,7 @@ def safe_float(f, default):
 
 
 def lora_name_to_file(name):
-    filenames = folder_paths.get_filename_list("loras")
+    filenames = get_filename_list("loras")
     # Return exact matches as is
     if name in filenames:
         return name
@@ -129,4 +209,45 @@ def lora_name_to_file(name):
             p = Path(f).with_suffix("")
             if p.name == n or str(p) == n:
                 return f
+    # Finally, try to find unique match from parts
+    parts = name.split()
+    search = [f for f in filenames if all(p in f for p in parts)]
+    if len(search) == 1:
+        return search[0]
+
     return None
+
+
+def map_inputs(input_map, inputs):
+    new_inputs = {}
+    for k in inputs:
+        key = inputs[k]
+        new_inputs[k] = key
+        if isinstance(key, list):
+            key = tuple(key)
+            x = input_map.get(key, inputs[k])
+            new_inputs[k] = x
+    return new_inputs
+
+
+def expand_graph(node_mappings, graph):
+    input_map = {}
+    new_graph = copy.deepcopy(graph)
+    for k in graph:
+        data = graph[k]
+        if not isinstance(data, dict) or "class_type" not in data or data["class_type"] not in node_mappings:
+            continue
+        node = node_mappings[data["class_type"]]()
+        inputs = map_inputs(input_map, data["inputs"].copy())
+        inputs["unique_id"] = k
+        fn = getattr(node, getattr(node, "FUNCTION"))
+        expansion = fn(**inputs)
+        for i, v in enumerate(expansion["result"]):
+            input_map[(k, i)] = v
+        del new_graph[k]
+        new_graph.update(expansion["expand"])
+
+    for k in new_graph:
+        data = new_graph[k]
+        data["inputs"] = map_inputs(input_map, data["inputs"])
+    return new_graph

pyproject.toml

@@ -1,14 +1,13 @@
 [project]
 name = "comfyui-prompt-control"
-description = "Nodes for convenient prompt editing, making many common operations prompt-controllable"
-version = "1.2.1"
+description = "Provides nodes for prompt editing and LoRA scheduling, advanced regional prompting (including attention masking) and more, all controlled through your text prompt"
+version = "2.1.0"
 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"

tools/expand_graph.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+from prompt_control.utils import expand_graph
+from prompt_control.nodes_lazy import NODE_CLASS_MAPPINGS as LN
+import json
+import sys
+
+
+# Needs ComfyUI in Python path
+# Usage: PYTHONPATH=../..:. python tools/expand_graph < graph_in_api_format.json > out.json
+if __name__ == "__main__":
+    graph = json.load(sys.stdin)
+    new = expand_graph(LN, graph)
+    print(json.dumps(new))

web/docs/PCAddMaskToCLIP.md

@@ -0,0 +1,3 @@
+# PC: Attach Mask
+
+Attaches custom masks to a CLIP object so that they can be referred to in prompts using `PCTextEncode` or `PC: Schedule prompt`.

web/docs/PCAddMaskToCLIPMany.md

@@ -0,0 +1 @@
+PCAddMaskToCLIP.md

web/docs/PCAttentionCoupleBatchNegative.md

@@ -0,0 +1,7 @@
+# PC: Attention Couple (batch negative)
+
+This node applies an optimization that re-enables negative cond batching when Attention Couple is in use.
+
+It improves performance when negative prompts are not scheduled, but slightly affects outputs and is not required for Attention Couple to work.
+
+Simply add it to your workflow and pass in your positive and negative prompts. It is always safe to use, as it will not do anything when it detects that the optimization can't be applied (eg. when negative prompts contain schedules)

web/docs/PCLazyLoraLoader.md

@@ -0,0 +1,7 @@
+# PC: Schedule LoRAs
+
+This node is the core of Prompt Control. It evaluates a prompt schedule and dynamically expands into a scheduled workflow consisting of necessary calls to `LoRALoader` and `Create Hook LoRA` (for scheduled LoRAs).
+
+You can use it in place or in addition to your usual `LoRA Loader` nodes; just pass in a text prompt containing your LoRA schedule (it can be shared with `PC: Schedule Prompt`). Then connect your MODEL output as usual and the CLIP output to your `PC: Schedule Prompt` nodes.
+
+For documentation on syntax, for now see the [documentation on GitHub](https://github.com/asagi4/comfyui-prompt-control/blob/master/doc/schedules.md)

web/docs/PCLazyLoraLoaderAdvanced.md

@@ -0,0 +1 @@
+PCLazyLoraLoader.md

web/docs/PCLazyTextEncode.md

@@ -0,0 +1,7 @@
+# PC: Schedule Prompt
+
+This node is the core of Prompt Control. It evaluates a prompt schedule and dynamically expands into a scheduled workflow consisting of calls to `PCTextEncode`, `SetConditioningTimesteps` and other necessary nodes.
+
+To use it, simply replace your usual `CLIP Text Encode` nodes with `PC: Schedule Prompt` nodes. For LoRA Loading, you should use `PC: Schedule LoRAs` in place (or in addition to) of your usual LoRA Loader node.
+
+For documentation on syntax, for now see the [documentation on GitHub](https://github.com/asagi4/comfyui-prompt-control/blob/master/doc/schedules.md)

web/docs/PCLazyTextEncodeAdvanced.md

@@ -0,0 +1 @@
+PCLazyTextEncode.md

web/docs/PCLoraHooksFromText.md

@@ -0,0 +1,5 @@
+# PC: LoRA Hooks from Text (non-lazy)
+
+Creates cond hooks from a LoRA schedule, if you want to apply them manually for some reason.
+
+You should not need to use this. Use `PC: Schedule LoRAs`.

web/docs/PCMacroExpand.md

@@ -0,0 +1,5 @@
+# PC: Expand Macros
+
+Expands [prompt macros](https://github.com/asagi4/comfyui-prompt-control/blob/master/doc/macros.md)
+
+You should not need to use this directly. Use `PC: Schedule Prompt` instead.

web/docs/PCSetPCTextEncodeSettings.md

@@ -0,0 +1,7 @@
+# PC: Configure PCTextEncode
+
+Configures a CLIP object with new default values used by `PCTextEncode`. Apply it before everything else.
+
+This is needed if you want to do scheduling with steps instead of denoising percentages, but otherwise it's completely optional.
+
+Note that steps are simply syntactic sugar for percentages and may not correspond to actual steps depending on the scheduler used.

web/docs/PCTextEncode.md

@@ -0,0 +1,5 @@
+# PC: Text Encode (no scheduling)
+
+This node encodes text using some special syntax for advanced features. You should rarely need to use this node directly, and instead use `PC: Schedule Prompt` which uses this node under the hood.
+
+For documentation on syntax, see the [documentation on GitHub](https://github.com/asagi4/comfyui-prompt-control/blob/master/doc/basic.md)