v2.0.0-rc.8

Very basic test for COUPLE
Fix skipping prompt segments by setting weight to 0
2025-06-09 20:07:54 +03:00 · 2025-06-09 20:06:25 +03:00 · 2025-06-09 20:05:50 +03:00 · 2025-06-09 19:53:05 +03:00 · 2025-06-09 19:15:55 +03:00 · 2025-06-09 18:37:53 +03:00
27 changed files with 2557 additions and 1227 deletions
@@ -7,6 +7,9 @@ on:
    paths:
      - "pyproject.toml"

+permissions:
+  issues: write
+
 jobs:
  tests:
    uses: ./.github/workflows/tests.yml
@@ -15,6 +18,7 @@ jobs:
  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
@@ -22,5 +26,4 @@ jobs:
      - name: Publish Custom Node
        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 }}
@@ -1,11 +1,11 @@
 name: Run tests requiring ComfyUI
 on:
-  - workflow_call
-  - workflow_dispatch
-  - push:
-      paths:
-        - prompt_control/nodes_lazy.py
-        - prompt_control/utils.py
+  workflow_call:
+  workflow_dispatch:
+  push:
+    paths:
+      - prompt_control/nodes_lazy.py
+      - prompt_control/utils.py


 jobs:
@@ -23,5 +23,6 @@ jobs:
      - uses: actions/setup-python@v5
        with:
          python-version: '3.11'
-      - run: pip install -r requirements.txt
+      - run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+      - run: pip install -r requirements.txt -r ComfyUI/requirements.txt
      - run: PYTHONPATH=ComfyUI python -m prompt_control.test_graph
@@ -11,4 +11,15 @@ test:
 test_graph:
 	PYTHONPATH=../../ python -m prompt_control.test_graph

+test_encode:
+	PYTHONPATH=../../ python -m prompt_control.test_encode
+
+test_encode_both:
+	TEST_TE="clip_l t5" PYTHONPATH=../../ python -m prompt_control.test_encode
+
+test_heavy: test_graph test_encode_both
+
+manual_test:
+	PYTHONPATH=../../ python -im prompt_control.manual_test
+
 .PHONY: check format all
@@ -2,48 +2,51 @@

 Control LoRA and prompt scheduling, advanced text encoding, regional prompting, and much more, through your text prompt. Generates dynamic graphs that are literally identical to handcrafted noodle soup.

+Prompt Control comes with `PCTextEncode`, which provides advanced text encoding with many additional features compared to ComfyUI's base `CLIPTextEncode`.
+
 A `Basic Text to Image` template is included with the extension, and can be loaded from ComfyUI's template library.

+## What can it do?
+
+You can use text prompts to control the following:
+
+- A1111-style prompt scheduling and filtering without noodle soup.
+- LoRA loading and scheduling via ComfyUI's hook system
+- Masking, composition and area control (regional prompting) with an implementation of [Attention Couple](/doc/attention_couple.md), also fully schedulable.
+- 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 with `DEF`
+- And a bunch more
+
+All features are fully schedulable unless otherwise stated. See the [syntax documentation](doc/syntax.md) for details on how to use each feature.
+
+If you find prompt scheduling inconvenient for some reason, `PCTextEncode` can be used as a drop-in replacement for `CLIPTextEncode` to get everything else.
+
+[This workflow](example_workflows/Workflow%20Comparison.json?raw=1) shows LoRA scheduling and prompt editing and compares it with the same prompt implemented with built-in ComfyUI nodes. You can also find it in the template library.
+
+## Compatibility
+
+Prompt Control uses graph generation,  and tries to delegate functionality to core ComfyUI wherever possible, implementing any hooks and patches in a way that is maximally compatible. This means that it should just work in most cases, even with models and nodes not explicitly supported.
+
+If you encounter issues as a user or if you're a node developer and Prompt Control somehow breaks something, feel free to file a bug report.
+
 ## Prompt Control v2

 Prompt control has been almost completely rewritten. It now uses ComfyUI's lazy execution to build graphs from the text prompt at runtime. The generated graph is often exactly equivalent to a manually built workflow using native ComfyUI nodes. There are no more weird sampling hooks that could cause problems with other nodes

-Prompt Control 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.

-
 ### Everything broke, where are the old nodes?

 If you really need them, you can install the [legacy nodes](https://github.com/asagi4/comfyui-prompt-control-legacy). However, I will not fix bugs in those nodes, and I strongly recommend just migrating your workflows to the new nodes.

 You can have both installed at the same time; none of the nodes conflict.

-## What can it do?
-
-See [features](#features) below. Things you can control via the prompt:
- Prompt editing and filtering without noodle soup
- LoRA loading and scheduling via ComfyUI's hook system
- Masking, composition and area control (regional prompting)
- Prompt operations like `BREAK` and `AND`
- Weight interpretation types (comfy, A1111, etc.)
- Prompt masking with [cutoff](#cutoff)
- And a bunch more
-
-See the [syntax documentation](doc/syntax.md)
-
-If you find prompt scheduling inconvenient for some reason, `PCTextEncode` can be used as a drop-in replacement for `CLIPTextEncode` to get everything else.
-
-[This workflow](workflows/example-lazy.json?raw=1) shows LoRA scheduling and prompt editing and compares it with the same prompt implemented with built-in ComfyUI nodes.
-
-[Here](workflows/example-2pass.json?raw=1) is a two-pass workflow illustrating more features, including custom masks and filtering.
-
-The tools in this repository combine well with the macro and wildcard functionality in [comfyui-utility-nodes](https://github.com/asagi4/comfyui-utility-nodes)
-
-
 ## Requirements

 For LoRA scheduling to work, you'll need at least version 0.3.7 of ComfyUI (0.3.36 of ComfyUI desktop).
@@ -93,75 +96,6 @@ This node attaches masks to a `CLIP` model so that they can be referred to when

 This node configures `PCTextEncode` default values for some functions by attaching the information to a `CLIP` model.

-# Features
-## Scheduling and LoRA loading
-
-Prompt control provides a way to easily schedule different prompts and control LoRA loading.
-
-See the [syntax documentation](doc/syntax.md)
-
-### Note on how schedules work
-
-ComfyUI does not use the step number to determine whether to apply conds; instead, it uses the sampler's timestep value which is affected by the scheduler you're using. This means that when the sampler scheduler isn't linear, the schedules generated by prompt control will not be either.
-
-## Advanced CLIP encoding
-
-If you use `PCTextEncode`, advanced encodings are available automatically. Thanks to BlenderNeko for the original code.
-
-Use the syntax `STYLE(weight_interpretation, normalization)` in a prompt to affect how prompts are interpreted.
-
-The weight interpretations available are:
-  - comfy (default)
-  - comfy++
-  - compel
-  - down_weight
-  - A1111
-  - perp
-
-Normalizations are:
-  - none (default)
-  - length
-  - mean
-
-The normalization calculations are independent operations and you can combine them with `+`, eg `STYLE(A1111, length+mean)` or `STYLE(comfy, mean+length)`, or even something silly like `STYLE(perp, mean+length+mean+length)`
-
-The style can be specified separately for each AND:ed prompt, but the first prompt is special; later prompts will "inherit" it as default. For example:
-
-```
-STYLE(A1111) a (red:1.1) cat with (brown:0.9) spots and a long tail AND an (old:0.5) dog AND a (green:1.4) (balloon:1.1)
-```
-will interpret everything as A1111, but
-```
-a (red:1.1) cat with (brown:0.9) spots and a long tail AND STYLE(A1111) an (old:0.5) dog AND a (green:1.4) (balloon:1.1)
-```
-Will interpret the first one using the default ComfyUI behaviour, the second prompt with A1111 and the last prompt with the default again
-
-For things (ie. the code imports) to work, the nodes must be cloned in a directory named exactly `ComfyUI_ADV_CLIP_emb`.
-
-## Cutoff
-
-NOTE: Cutoff syntax might change at some point; it's pretty clunky.
-
-`PCTextEncode` reimplements cutoff from [ComfyUI Cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff).
-
-The syntax is
-```
-a group of animals, [CUT:white cat:white], [CUT:brown dog:brown:0.5:1.0:1.0:_]
-```
-You should read the prompt as `a group of animals, white cat, brown dog`, but CUT causes the tokens in `target_tokens` to be masked off from the base prompt in `region_text`, so that their effect can be isolated, and you're less likely to get brown cats or white dogs.
-
-Target tokens are treated individually, separated by space, for example, `[CUT:green apple, red apple, green leaf:green apple]` will mask *both* greens and the apple, giving you `+ +, red +, + leaf`. To mask out just `green apple`, use `[CUT:green apple, red apple:green_apple]` which will result in a masked prompt of `+ +, red apple`. Escape `_` with a `\`.
-
-the parameters in the `CUT` section are `region_text:target_tokens:weight;strict_mask:start_from_masked:padding_token` of which only the first two are required. The default values are `weight=1.0`, `strict_mask=1.0` `start_from_masked=1.0`, `padding_token=+`
-
-If `strict_mask`, `start_from_masked` or `padding_token` are specified in more than one CUT, the *last* one becomes the default for any CUTs afterwards that do not explicitly set the parameters. For example, in:
-
-`[CUT:white cat:white:0.5] and [CUT:black parrot, flying:black:1.0:0.5] and [CUT:green apple:green]`
-
-`white cat` will a weight of 0.5, and 1.0 for all parameters, and `black parrot` and `green apple` will *both* have a `strict_mask` parameter of 0.5.
-
-The parameters affect how the masked and unmasked prompts are combined to produce the final embedding. Just play around with them.
-
 # Known issues

 - ComfyUI's caching mechanism has an issue that makes it unnecessarily invalidate caches for certain inputs; you'll still get some benefit from the lazy nodes, but changing inputs that shouldn't affect downstream nodes (especially if using filtering) will still cause them to be recomputed because ComfyUI doesn't realize the inputs haven't changed.
@@ -169,3 +103,5 @@ The parameters affect how the masked and unmasked prompts are combined to produc
 If you want to enable a hack to fix this, set `PROMPTCONTROL_ENABLE_CACHE_HACK=1` in your environment. Unset it to disable.

 It's a purely optional performance optimization that allows Prompt Control nodes to override their cache keys in a way that should not interfere with other nodes. Note that the optimization only works if the text input to the lazy nodes is a constant (so either directly on the node or from a primitive); outputs from other nodes can't be optimized.
+
+- Cutoff does not work with models that use non-CLIP text encoders, like Flux. This might be fixable, but it's uncertain if cutoff even makes sense for those models.
@@ -29,22 +29,9 @@ cache_hack.init()
 NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}

-nodes = ["base", "lazy", "tools"]
-optional_nodes = ["attnmask"]
-if importlib.util.find_spec("comfy.hooks"):
-    nodes.extend(["hooks"])
-else:
-    log.error("Your ComfyUI version is too old, can't import comfy.hooks. Update your installation.")
+nodes = ["base", "lazy", "tools", "hooks"]

 for node in nodes:
    mod = importlib.import_module(f".prompt_control.nodes_{node}", package=__name__)
    NODE_CLASS_MAPPINGS.update(mod.NODE_CLASS_MAPPINGS)
    NODE_DISPLAY_NAME_MAPPINGS.update(mod.NODE_DISPLAY_NAME_MAPPINGS)
-
-for node in optional_nodes:
-    try:
-        mod = importlib.import_module(f".prompt_control.nodes_{node}", package=__name__)
-        NODE_CLASS_MAPPINGS.update(mod.NODE_CLASS_MAPPINGS)
-        NODE_DISPLAY_NAME_MAPPINGS.update(mod.NODE_DISPLAY_NAME_MAPPINGS)
-    except ImportError:
-        log.info(f"Could not import optional nodes: {node}; continuing anyway")
@@ -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
+```
@@ -1,6 +1,8 @@
-# Scheduling syntax
+# Prompt Control Syntax

-Syntax is like A1111 for now, but only fractions are supported for steps. LoRAs are scheduled by including them in a scheduling expression.
+If you're viewing this on GitHub, I recommend opening the outline by clicking the button in the top right corner of the text view (it is annoyingly easy to miss).
+
+Scheduling syntax is similar to A1111, but only fractions are supported for steps. LoRAs are scheduled by including them in a scheduling expression.

 ```
 a [large::0.1] [cat|dog:0.05] [<lora:somelora:0.5:0.6>::0.5]
@@ -12,7 +14,7 @@ a [large::0.1] [cat|dog:0.05] [<lora:somelora:0.5:0.6>::0.5]
 There are two forms of scheduled prompts.

 ### Basic scheduling expressions
-Basic expressions take the form `[before:after:X]` where `X` is the switch point, a decimal number between 0.0 and 1.0 inclusive, representing 0 to 100% of timesteps.
+Basic expressions take the form `[before:after:X]` where `X` is the switch point, a decimal number between 0.0 and 1.0 inclusive, representing 0 to 100% of timesteps. Either prompt can also be empty.
 For example:
 ```
 a [red:blue:0.5] cat
@@ -24,13 +26,24 @@ a [red:[blue::0.7]:0.5] cat

 switches from `a red cat` to `a blue cat` at 0.5 and to `a cat` at 0.7

-**Note:** As a special case, `[cat:0.5]` is equivalent to `[:cat:0.5]` meaning it switches from empty to `cat` at 0.5.
+For convenience `[cat:0.5]` is equivalent to `[:cat:0.5]` meaning it switches from empty to `cat` at 0.5.

 ### Range expressions

-You can also use `a [during:after:0.3,0.7]` as a shortcut. The prompt be `a` until 0.3, `a during` until 0.7, and then `a after`. This form is equivalent to `[[during:after:0.7]:0.3]`
+The most general form of a schedule is a range expression: For example, in `[before:during:after:0.3,0.7]`, The prompt be `a before` until 0.3, `a during` until 0.7, and then `a after`. This form is equivalent to `[before:[during:after:0.7]:0.3]`

-For convenience, `[during:0.1,0.4]` is equivalent to `[during::0.1,0.4]`
+For convenience, `[during:0.1,0.4]` is equivalent to `[:during::0.1,0.4]` and `[during:after:0.1,0.4]` is equivalent to `[:during:after:0.1,0.4]`.
+
+`[before:during:after:0.1]` is the same as `[before:during:after:0.1,1.0]` which is same as `[before:during:0.1]`
+
+
+### Using step numbers with the Advanced nodes
+
+If you provide a non-zero value to `num_steps` to the `Advanced` versions of the scheduling nodes, you will be able to use step numbers in prompts.
+
+For now, a value between 0 and 1.0 will be interpreted as a percentage if it contains a ., and as an absolute step otherwise.
+
+This is just syntactic sugar. Behind the scenes, the values are converted to percentages and have normal ComfyUI scheduling behaviour.

 ## Tag selection
 Using the `FilterSchedule` node, in addition to step percentages, you can use a *tag* to select part of an input:
@@ -47,6 +60,7 @@ a [black:blue:X] [cat:dog:Y] [walking:running:Z] in space
 ```
 with `tags` `x,z` would result in the prompt `a blue cat running in space`

+The three prompt form `[a:b:c:TAG]` is parsed, but ignores `b` and is equivalent to `[a:c:TAG]`.

 ## LoRA Scheduling
 When using the lazy graph building nodes, LoRAs can be scheduled by referring to them in a scheduling expression, like so:
@@ -81,48 +95,133 @@ generates a LoRA schedule based on a sinewave

 # Basic prompt syntax

-This syntax is also available in outside scheduled prompts, where applicable.
+This syntax is also available in outside scheduled with the `PCTextEncode` node, where applicable.

-## LoRA loading
+## Combining prompts

-The A111-style syntax `<lora:loraname:weight>` can be used to load LoRAs via the prompt. See LoRA scheduling above.
+### AND

-## Combining prompts, A1111-style
+`AND` can be used to create "prompt segments". By default, it works as if you had combined the different prompts with `ConditioningCombine`. 

- 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.
+It is also used with regional prompting, see `MASK` and `COUPLE` below.

-`AND` can be used to combine prompts. You can also use a weight at the end. It does a weighted sum of each prompt,
+Prompts can have a weight at the end:
 ```
 cat :1 AND dog :2
 ```
-The weight defaults to 1 and are normalized so that `a:2 AND b:2` is equal to `a AND b`. `AND` is processed after schedule parsing, so you can change the weight mid-prompt: `cat:[1:2:0.5] AND dog`
+`AND` is processed after schedule parsing, so you can change the weight mid-prompt: `cat:[1:2:0.5] AND dog`
+
+The weight defaults to 1. If a prompt's weight is set to 0, it's **skipped entirely.** This can be useful when scheduling to completely disable a prompt:
+
+```
+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 do various things. 
+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.
-Note: Whitespace is *not* stripped from string parameters by default. Commas can be escaped with `\,`
+Functions have the form `FUNCNAME(param1, param2, ...)`. How parameters are interpreted is up to the function. 

-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.
+In general, function parameters will have default values that are used if the parameter is left empty.

-### SDXL
+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`.

-To set the `clip_l` prompt, as with `CLIPTextEncodeSDXL`, use the function `CLIP_L(prompt text goes here)`.
+### TE: Per-encoder prompts for multi-encoder models
+
+You can specify per-encoder prompts using the `TE` function. The syntax is as follows:
+`TE(encoder_name=prompt)`. Whitespace surrounding the prompt and encoder name are ignored.
+
+For example:
+```
+TE(l=cat) TE(g = (dog:1.1)) TE(t5xxl=tiger)
+```
+The keys to use depend on what key ComfyUI uses for the encoder; for example `l` for CLIP L, `g` for CLIP G, and `t5xxl` for T5 XXL (Flux text encoder).
+
+Use `TE(help)` to print a help text listing available keys.

 Things to note:
- Multiple instances of `CLIP_L` are joined with a space. That is, `CLIP_L(foo)CLIP_L(bar)` is the same as `CLIP_L(foo bar)`
- Using `BREAK` isn't supported in it; it'll just parse as the plain word BREAK.
- similarly, `AND` inside `CLIP_L` does not do anything sensible; `CLIP_L(foo AND bar)` will parse as two prompts `CLIP_L(foo` and `bar)`
- `CLIP_L` and `SDXL` have no effect on SD 1.5.
- The rest of the prompt becomes the `clip_g` prompt.
- If there is no `CLIP_L` or `SDXL`, the prompts will work as with `CLIPTextEncode`.
+- If you set a prompt with `TE`, it will override the prompt outside the function for the specified text encoder.
+- Multiple instances of `TE` are joined with a space. That is, `TE(l=foo)TE(l=bar)` is the same as `TE(l=foo bar)`
+- `AND` 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
+### SHUFFLE and SHIFT: Create prompt permutations

 Default parameters: `SHUFFLE(seed=0, separator=,, joiner=,)`, `SHIFT(steps=0, separator=,, joiner=,)`

@@ -147,23 +246,32 @@ For example:
 Whitespace is *not* stripped and may also be used as a joiner or separator
 - `SHIFT(1,, ) cat,dog` results in `dog cat`

-### NOISE
+### NOISE: Add noise to a prompt

-The function `NOISE(weight, seed)` adds some random noise into the prompt. The seed is optional, and if not specified, the global RNG is used. `weight` should be between 0 and 1.
+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 also [Attention Couple](#attention-couple) below

 ### MASK, IMASK and AREA

 You can use `MASK(x1 x2, y1 y2, weight, op)` to specify a region mask for a prompt. The values are specified as a percentage with a float between `0` and `1`, or as absolute pixel values (these can't be mixed). `1` will be interpreted as a percentage instead of a pixel value.

+Multiple `MASK` or `IMASK` calls will be composited together using ComfyUI's `MaskComposite` node, using `op` as the `operation` parameter (defaulting to `multiply`).
+
 Similarly, you can use `AREA(x1 x2, y1 y2, weight)` to specify an area for the prompt (see ComfyUI's area composition examples). The area is calculated by ComfyUI relative to your latent size.

-#### Custom masks: IMASK and `PCAddMaskToCLIP`
+### Custom masks: IMASK and `PCAddMaskToCLIP`

 You can attach custom masks to a `CLIP` with the `PC: Attach Mask` nodes and then refer to those masks in the prompt using `IMASK(index, weight, op)`. Indexing starts from zero, so 0 is the first attached mask etc. `PCSCheduleAddMasks` ignores empty inputs, so if you only add a mask to the `mask4` input, it will still have index 0.

 Applying the nodes multiple times *appends* masks rather than overriding existing ones, so if you need more than 4, you can just use it more than once.

-#### Behaviour of masks
+### Behaviour of masks
 If multiple `MASK`s are specified, they are combined together with ComfyUI's `MaskComposite` node, with `op` specifying the operation to use (default `multiply`). In this case, the combined mask weight can be set with `MASKW(weight)` (defaults to 1.0).

 Masks assume a size of `(512, 512)`, unless overridden with `PC: Configure PCTextEncode` and pixel values will be relative to that. ComfyUI will scale the mask to match the image resolution. You can change it manually by using `MASK_SIZE(width, height)` anywhere in the prompt,
@@ -176,7 +284,7 @@ Note that because the default values are percentages, `MASK(0 256, 64 512)` is v

 Masking does not affect LoRA scheduling unless you set unet weights to 0 for a LoRA.

-### FEATHER
+### FEATHER: Mask operations

 When you use `MASK` or `IMASK`, you can also call `FEATHER(left top right bottom)` to apply feathering using ComfyUI's `FeatherMask` node. The values are in pixels and default to `0`.

@@ -191,6 +299,30 @@ gives you a mask that is a combination of 1, 2 and 3, where 1 and 3 are feathere

 The order of the `FEATHER` and `MASK` calls doesn't matter; you can have `FEATHER` before `MASK` or even interleave them.

+## Cutoff
+
+NOTE: Cutoff syntax might change at some point; it's pretty clunky.
+
+`PCTextEncode` reimplements cutoff from [ComfyUI Cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff).
+
+The syntax is
+```
+a group of animals, [CUT:white cat:white], [CUT:brown dog:brown:0.5:1.0:1.0:_]
+```
+You should read the prompt as `a group of animals, white cat, brown dog`, but CUT causes the tokens in `target_tokens` to be masked off from the base prompt in `region_text`, so that their effect can be isolated, and you're less likely to get brown cats or white dogs.
+
+Target tokens are treated individually, separated by space, for example, `[CUT:green apple, red apple, green leaf:green apple]` will mask *both* greens and the apple, giving you `+ +, red +, + leaf`. To mask out just `green apple`, use `[CUT:green apple, red apple:green_apple]` which will result in a masked prompt of `+ +, red apple`. Escape `_` with a `\`.
+
+the parameters in the `CUT` section are `region_text:target_tokens:weight;strict_mask:start_from_masked:padding_token` of which only the first two are required. The default values are `weight=1.0`, `strict_mask=1.0` `start_from_masked=1.0`, `padding_token=+`
+
+If `strict_mask`, `start_from_masked` or `padding_token` are specified in more than one CUT, the *last* one becomes the default for any CUTs afterwards that do not explicitly set the parameters. For example, in:
+
+`[CUT:white cat:white:0.5] and [CUT:black parrot, flying:black:1.0:0.5] and [CUT:green apple:green]`
+
+`white cat` will a weight of 0.5, and 1.0 for all parameters, and `black parrot` and `green apple` will *both* have a `strict_mask` parameter of 0.5.
+
+The parameters affect how the masked and unmasked prompts are combined to produce the final embedding. Just play around with them.
+
 ## Miscellaneous
 - `<emb:xyz>` is alternative syntax for `embedding:xyz` to work around a syntax conflict with `[embedding:xyz:0.5]` which is parsed as a schedule that switches from `embedding` to `xyz`.

@@ -198,9 +330,12 @@ The order of the `FEATHER` and `MASK` calls doesn't matter; you can have `FEATHE

 Experimental features are unstable and may disappear or change without warning.

-## DEF
+## 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`

-You can define "prompt macros" by using `DEF`:
 ```
 DEF(MYMACRO=this is a prompt)
 [(MYMACRO:0.6):(MYMACRO:1.1):0.5]
@@ -209,7 +344,7 @@ 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])
@@ -219,26 +354,51 @@ 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.
+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:

-Note that unspecified parameters will not be substituted:
 ```
-DEF(mything=a $1 b $2)
-mything
-mything(A)
+DEF(MACRO(example; 0; 1)=[$1:$2,$3])
+MACRO MACRO(test; 0.2)
 ```
 gives
 ```
-a $1 b $2
-a A b $2
+[example:0,1] [test:0.2,1]
 ```

-Macros are expanded before any other parsing takes place. The expansion continues until no further changes occur. Recursion will raise an error.
+```
+DEF(MACRO() = [a:$1:0.5])
+```
+sets the default value of `$1` to an empty string.

-## Attention masking
+### Unspecified parameters in macros

-Use `ATTN()` in combination with `MASK()` or `IMASK()` to enable attention masking. Currently, it's pretty slow and only works with SDXL. You need to have a recent enough version of ComfyUI for this to work.
+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"
+```
+
+## 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 `g`, `l`, or `t5xxl`. For example with SDXL, try `TE_WEIGHT(g=0.25, l=0.75`). The weights are applied as a multiplier to the TE output.
+For models using multiple text encoders, you can set weights per TE using the syntax `TE_WEIGHT(clipname=weight, clipname2=weight2, ...)` where `clipname` is one of the encoder names printed by `TE(help)`. For example with SDXL, try `TE_WEIGHT(g=0.25, l=0.75)`.
+
+The weights are applied as a multiplier to the TE output. You can also override pooled output multipliers using eg. `l_pooled`.
+
+To set a default value for all encoders, use `TE_WEIGHT(all=weight)`
@@ -1,4 +1,6 @@
 {
+  "id": "e820c2fb-9502-45b7-a864-684757dddcdf",
+  "revision": 0,
  "last_node_id": 18,
  "last_link_id": 20,
  "nodes": [
@@ -21,26 +23,26 @@
        {
          "name": "MODEL",
          "type": "MODEL",
+          "slot_index": 0,
          "links": [
            2
-          ],
-          "slot_index": 0
+          ]
        },
        {
          "name": "CLIP",
          "type": "CLIP",
+          "slot_index": 1,
          "links": [
            3
-          ],
-          "slot_index": 1
+          ]
        },
        {
          "name": "VAE",
          "type": "VAE",
+          "slot_index": 2,
          "links": [
-            16
-          ],
-          "slot_index": 2
+            18
+          ]
        }
      ],
      "properties": {
@@ -66,7 +68,7 @@
      "flags": {
        "collapsed": true
      },
-      "order": 8,
+      "order": 7,
      "mode": 0,
      "inputs": [
        {
@@ -87,10 +89,10 @@
        {
          "name": "CONDITIONING",
          "type": "CONDITIONING",
+          "slot_index": 0,
          "links": [
            12
-          ],
-          "slot_index": 0
+          ]
        }
      ],
      "title": "PC: Schedule Prompt (positive)",
@@ -109,24 +111,26 @@
      "id": 3,
      "type": "PCLazyLoraLoader",
      "pos": [
-        255,
-        -765
+        257.5,
+        -745
      ],
      "size": [
        210,
-        98
+        78
      ],
      "flags": {},
-      "order": 7,
+      "order": 6,
      "mode": 0,
      "inputs": [
        {
          "name": "model",
+          "shape": 7,
          "type": "MODEL",
          "link": 2
        },
        {
          "name": "clip",
+          "shape": 7,
          "type": "CLIP",
          "link": 3
        },
@@ -143,19 +147,19 @@
        {
          "name": "MODEL",
          "type": "MODEL",
+          "slot_index": 0,
          "links": [
            17
-          ],
-          "slot_index": 0
+          ]
        },
        {
          "name": "CLIP",
          "type": "CLIP",
+          "slot_index": 1,
          "links": [
            5,
            9
-          ],
-          "slot_index": 1
+          ]
        }
      ],
      "properties": {
@@ -181,7 +185,7 @@
        474
      ],
      "flags": {},
-      "order": 10,
+      "order": 9,
      "mode": 0,
      "inputs": [
        {
@@ -209,10 +213,10 @@
        {
          "name": "LATENT",
          "type": "LATENT",
+          "slot_index": 0,
          "links": [
            15
-          ],
-          "slot_index": 0
+          ]
        }
      ],
      "properties": {
@@ -242,13 +246,16 @@
        225
      ],
      "flags": {},
-      "order": 5,
+      "order": 1,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "STRING",
          "type": "STRING",
+          "widget": {
+            "name": "text"
+          },
          "links": [
            6,
            7
@@ -277,17 +284,20 @@
        225
      ],
      "flags": {},
-      "order": 4,
+      "order": 2,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "STRING",
          "type": "STRING",
+          "widget": {
+            "name": "text"
+          },
+          "slot_index": 0,
          "links": [
            8
-          ],
-          "slot_index": 0
+          ]
        }
      ],
      "title": "Negative prompt",
@@ -314,7 +324,7 @@
      "flags": {
        "collapsed": true
      },
-      "order": 9,
+      "order": 8,
      "mode": 0,
      "inputs": [
        {
@@ -335,10 +345,10 @@
        {
          "name": "CONDITIONING",
          "type": "CONDITIONING",
+          "slot_index": 0,
          "links": [
            13
-          ],
-          "slot_index": 0
+          ]
        }
      ],
      "title": "PC: Schedule Prompt (negative)",
@@ -365,7 +375,7 @@
        106
      ],
      "flags": {},
-      "order": 1,
+      "order": 3,
      "mode": 0,
      "inputs": [],
      "outputs": [
@@ -402,7 +412,7 @@
        46
      ],
      "flags": {},
-      "order": 11,
+      "order": 10,
      "mode": 0,
      "inputs": [
        {
@@ -420,53 +430,18 @@
        {
          "name": "IMAGE",
          "type": "IMAGE",
+          "slot_index": 0,
          "links": [
            20
-          ],
-          "slot_index": 0
+          ]
        }
      ],
      "properties": {
        "cnr_id": "comfy-core",
        "ver": "0.3.18",
        "Node name for S&R": "VAEDecode"
-      }
-    },
-    {
-      "id": 11,
-      "type": "Reroute",
-      "pos": [
-        1155,
-        -900
-      ],
-      "size": [
-        75,
-        26
-      ],
-      "flags": {},
-      "order": 6,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "",
-          "type": "*",
-          "link": 16
-        }
-      ],
-      "outputs": [
-        {
-          "name": "",
-          "type": "VAE",
-          "links": [
-            18
-          ],
-          "slot_index": 0
-        }
-      ],
-      "properties": {
-        "showOutputText": false,
-        "horizontal": false
-      }
+      },
+      "widgets_values": []
    },
    {
      "id": 13,
@@ -480,7 +455,7 @@
        105
      ],
      "flags": {},
-      "order": 3,
+      "order": 4,
      "mode": 0,
      "inputs": [],
      "outputs": [],
@@ -503,7 +478,7 @@
        210
      ],
      "flags": {},
-      "order": 2,
+      "order": 5,
      "mode": 0,
      "inputs": [],
      "outputs": [],
@@ -526,7 +501,7 @@
        405
      ],
      "flags": {},
-      "order": 12,
+      "order": 11,
      "mode": 0,
      "inputs": [
        {
@@ -634,14 +609,6 @@
      0,
      "LATENT"
    ],
-    [
-      16,
-      1,
-      2,
-      11,
-      0,
-      "*"
-    ],
    [
      17,
      3,
@@ -650,14 +617,6 @@
      0,
      "MODEL"
    ],
-    [
-      18,
-      11,
-      0,
-      10,
-      1,
-      "VAE"
-    ],
    [
      20,
      10,
@@ -665,6 +624,14 @@
      18,
      0,
      "IMAGE"
+    ],
+    [
+      18,
+      1,
+      2,
+      10,
+      1,
+      "VAE"
    ]
  ],
  "groups": [],
@@ -673,10 +640,28 @@
    "ds": {
      "scale": 0.8,
      "offset": [
-        588,
-        1260
+        591.75,
+        1235
      ]
-    }
+    },
+    "linkExtensions": [
+      {
+        "id": 18,
+        "parentId": 1
+      }
+    ],
+    "reroutes": [
+      {
+        "id": 1,
+        "pos": [
+          1273.75,
+          -879.5
+        ],
+        "linkIds": [
+          18
+        ]
+      }
+    ]
  },
  "version": 0.4,
  "models": [{
@@ -1,6 +1,17 @@
 import torch
 import numpy as np
+from math import copysign
+import logging
 import itertools
+from .adv_encode_old import old_advanced_encode_from_tokens
+
+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):
@@ -12,29 +23,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 +52,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 +61,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 +68,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 +86,307 @@ 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
+    if "old+" not in weight_interpretation:
+        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)
    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(""))
+        weight_interpretation = weight_interpretation.replace("old+", "")
+        log.warning("Using old implementation of %s", weight_interpretation)
+        return old_advanced_encode_from_tokens(
+            tokenized,
+            token_normalization,
+            weight_interpretation,
+            encode_func,
+            266,
+            return_pooled=return_pooled,
+            apply_to_pooled=apply_to_pooled,
        )
-
-    if return_pooled:
-        if apply_to_pooled:
-            return weighted_emb, pooled
-        else:
-            return weighted_emb, pooled_base
-    return weighted_emb, None
@@ -0,0 +1,235 @@
+import torch
+import numpy as np
+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)
+    while True:
+        chunk = list(itertools.islice(it, n))
+        if not chunk:
+            return
+        yield chunk
+
+
+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)
+    weights = [[w_max if id == 0 else (w / top) * w_max for w, id in zip(x, y)] for x, y in zip(weights, word_ids)]
+    return weights
+
+
+def mask_word_id(tokens, word_ids, target_id, mask_token):
+    new_tokens = [[mask_token if wid == target_id else t for t, wid in zip(x, y)] for x, y in zip(tokens, word_ids)]
+    mask = np.array(word_ids) == target_id
+    return (new_tokens, mask)
+
+
+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 = torch.tensor(weights, dtype=base_emb.dtype, device=base_emb.device)
+    weight_tensor = weight_tensor.reshape(1, -1, 1).expand(base_emb.shape)
+
+    # 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)
+
+        m = torch.tensor(m, dtype=base_emb.dtype, device=base_emb.device)
+        m = m.reshape(1, -1, 1).expand(base_emb.shape)
+        masks.append(m)
+
+        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)
+    new_tokens = [
+        [mask_token if i * clip_len + j in inds_set else t for j, t in enumerate(x)] for i, x in enumerate(tokens)
+    ]
+    return new_tokens
+
+
+def down_weight(tokens, weights, word_ids, base_emb, length, encode_func):
+    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 = (266, 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)
+    embeddings_final = (norm_base / norm_weighted) * weighted_emb
+    return embeddings_final
+
+
+# For verification
+def A1111_renorm(base_emb, weighted_emb):
+    embeddings_final = (base_emb.mean() / weighted_emb.mean()) * weighted_emb
+    return embeddings_final
+
+
+def from_zero(weights, base_emb):
+    weight_tensor = torch.tensor(weights, dtype=base_emb.dtype, device=base_emb.device)
+    weight_tensor = weight_tensor.reshape(1, -1, 1).expand(base_emb.shape)
+    return base_emb * weight_tensor
+
+
+def old_advanced_encode_from_tokens(
+    tokenized,
+    token_normalization,
+    weight_interpretation,
+    encode_func,
+    m_token=266,
+    w_max=1.0,
+    return_pooled=False,
+    apply_to_pooled=False,
+    **extra_args,
+):
+    length = 77
+    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]
+
+    # weight normalization
+    # ====================
+
+    # distribute down/up weights over word lengths
+    if token_normalization.startswith("length"):
+        weights = divide_length(word_ids, weights)
+
+    # make mean of word tokens 1
+    if token_normalization.endswith("mean"):
+        weights = shift_mean_weight(word_ids, weights)
+
+    # weight interpretation
+    # =====================
+    pooled = None
+
+    if weight_interpretation in ["comfy", "perp"]:
+        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 = from_zero(weights, base_emb)
+        weighted_emb = A1111_renorm(base_emb, weighted_emb)
+        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 return_pooled:
+        if apply_to_pooled:
+            return weighted_emb, pooled
+        else:
+            return weighted_emb, pooled_base
+    return weighted_emb, None
@@ -0,0 +1,252 @@
+# 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._base_cond = base_cond
+        self._conds = conds
+        self._fill = 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)
@@ -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
@@ -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
@@ -1,79 +0,0 @@
-import logging
-
-log = logging.getLogger("comfyui-prompt-control")
-
-from comfy.hooks import TransformerOptionsHook, HookGroup, EnumHookScope
-from comfy.ldm.modules.attention import optimized_attention
-import torch.nn.functional as F
-import torch
-from math import sqrt
-
-
-class MaskedAttn2:
-    def __init__(self, mask):
-        self.mask = mask
-
-    def __call__(self, q, k, v, extra_options):
-        mask = self.mask
-        orig_shape = extra_options["original_shape"]
-        _, _, oh, ow = orig_shape
-        seq_len = q.shape[1]
-        mask_h = oh / sqrt(oh * ow / seq_len)
-        mask_h = int(mask_h) + int((seq_len % int(mask_h)) != 0)
-        mask_w = seq_len // mask_h
-        r = optimized_attention(q, k, v, extra_options["n_heads"])
-        mask = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="nearest").squeeze(1)
-        mask = mask.view(mask.shape[0], -1, 1).repeat(1, 1, r.shape[2])
-
-        return mask * r
-
-
-def create_attention_hook(mask):
-    attn_replacements = {}
-    mask = mask.detach().to(device="cuda", dtype=torch.float16)
-
-    masked_attention = MaskedAttn2(mask)
-
-    for id in [4, 5, 7, 8]:  # id of input_blocks that have cross attention
-        block_indices = range(2) if id in [4, 5] else range(10)  # transformer_depth
-        for index in block_indices:
-            k = ("input", id, index)
-            attn_replacements[k] = masked_attention
-    for id in range(6):  # id of output_blocks that have cross attention
-        block_indices = range(2) if id in [3, 4, 5] else range(10)  # transformer_depth
-        for index in block_indices:
-            k = ("output", id, index)
-            attn_replacements[k] = masked_attention
-    for index in range(10):
-        k = ("middle", 1, index)
-        attn_replacements[k] = masked_attention
-
-    hook = TransformerOptionsHook(
-        transformers_dict={"patches_replace": {"attn2": attn_replacements}}, hook_scope=EnumHookScope.HookedOnly
-    )
-    group = HookGroup()
-    group.add(hook)
-
-    return group
-
-
-class AttentionMaskHookExperimental:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {"mask": ("MASK",)},
-        }
-
-    RETURN_TYPES = ("HOOKS",)
-    CATEGORY = "promptcontrol/_testing"
-    FUNCTION = "apply"
-    EXPERIMENTAL = True
-    DESCRIPTION = "Experimental attention masking hook. For testing only"
-
-    def apply(self, mask):
-        return (create_attention_hook(mask),)
-
-
-NODE_CLASS_MAPPINGS = {"AttentionMaskHookExperimental": AttentionMaskHookExperimental}
-
-NODE_DISPLAY_NAME_MAPPINGS = {}
@@ -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)",
+    "PCAttentionCoupleBatchNegative": "PC: Attention Couple (batch negative)",
 }
@@ -1,6 +1,7 @@
 import logging
 from .parser import parse_prompt_schedules
 from comfy_execution.graph_utils import GraphBuilder, is_link
+
 from comfy_execution.graph import ExecutionBlocker

 from .utils import get_function
@@ -86,7 +87,7 @@ def create_hook_nodes_for_lora(graph, path, info, existing_node, start_pct, end_
    return hook_node, next_keyframe


-def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_hooks=True):
+def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True):
    # This gets rid of non-existent LoRAs
    consolidated = consolidate_schedule(schedule)
    if model is not None:
@@ -129,7 +130,7 @@ def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_h
            n.set_input("hooks_B", h.out(0))
            res = n
        res = res.out(0)
-        if apply_hooks:
+        if clip is not None and apply_hooks:
            n = graph.node("SetClipHooks")
            n.set_input("clip", clip)
            n.set_input("hooks", res)
@@ -137,13 +138,12 @@ def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_h
            n.set_input("schedule_clip", True)
            clip = n.out(0)

+    if clip is None:
+        clip = ExecutionBlocker("No clip model provided to PCLazyLoRALoader or PCLazyLoRALoaderAdvanced")
    r = graph.finalize()
    log.debug("LazyLoraLoader built graph: %s", json.dumps(r))

-    if return_hooks:
-        ret = (model, clip, res)
-    else:
-        ret = (model, clip)
+    ret = (model, clip, res)

    return {"result": ret, "expand": r}

@@ -154,16 +154,15 @@ class PCLazyLoraLoaderAdvanced:
    @classmethod
    def INPUT_TYPES(s):
        return {
-            "required": {
-                "text": ("STRING", {"multiline": True}),
+            "optional": {
                "model": ("MODEL", {"rawLink": True}),
                "clip": ("CLIP", {"rawLink": True}),
-            },
-            "optional": {
+                "text": ("STRING", {"multiline": True, "default": ""}),
                "apply_hooks": ("BOOLEAN", {"default": True}),
                "tags": ("STRING", {"default": ""}),
                "start": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 0.0, "step": 0.01}),
                "end": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 1.0, "step": 0.01}),
+                "num_steps": ("INT", {"min": 0, "max": 10000, "default": 0, "step": 1}),
            },
            "hidden": {"unique_id": "UNIQUE_ID"},
        }
@@ -173,15 +172,16 @@ class PCLazyLoraLoaderAdvanced:
    CATEGORY = "promptcontrol"
    FUNCTION = "apply"

-    def apply(self, model, clip, text, unique_id, apply_hooks=True, tags="", start=0.0, end=1.0):
-        schedule = parse_prompt_schedules(text, filters=tags, start=start, end=end)
-        graph = GraphBuilder(f"{unique_id}-")
-        return build_lora_schedule(graph, schedule, model, clip, apply_hooks=apply_hooks, return_hooks=True)
+    def apply(
+        self, unique_id, model=None, clip=None, text="", apply_hooks=True, tags="", start=0.0, end=1.0, num_steps=0
+    ):
+        schedule = parse_prompt_schedules(text, filters=tags, start=start, end=end, num_steps=num_steps)
+        graph = GraphBuilder()
+        r = build_lora_schedule(graph, schedule, model, clip, apply_hooks=apply_hooks)
+        return r


-class PCLazyLoraLoader:
-    CACHE_KEY = cache_key_lora
-
+class PCLazyLoraLoader(PCLazyLoraLoaderAdvanced):
    @classmethod
    def INPUT_TYPES(s):
        return {
@@ -197,19 +197,12 @@ class PCLazyLoraLoader:
        "MODEL",
        "CLIP",
    )
-    OUTPUT_TOOLTIPS = ("Returns a model and clip with LoRAs scheduled",)
    CATEGORY = "promptcontrol"
-    FUNCTION = "apply"

-    def apply(self, unique_id, model=None, clip=None, text=""):
-        graph = GraphBuilder(f"{unique_id}-")
-        schedule = parse_prompt_schedules(text)
-        if model is None and clip is None:
-            return (
-                ExecutionBlocker("No model input provided to PCLazyLoraLoader"),
-                ExecutionBlocker("No clip input provided to PCLazyLoraLoader"),
-            )
-        return build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_hooks=False)
+    def apply(self, *args, **kwargs):
+        r = super().apply(*args, **kwargs)
+        r["result"] = r["result"][:2]
+        return r


 def build_scheduled_prompts(graph, schedules, clip):
@@ -245,32 +238,11 @@ def build_scheduled_prompts(graph, schedules, clip):
    return {"result": (node.out(0),), "expand": g}


-def cache_key_from_inputs(cachekey, text, tags="", start=0.0, end=1.0, **kwargs):
-    schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end)
+def cache_key_from_inputs(cachekey, text, tags="", start=0.0, end=1.0, num_steps=0, **kwargs):
+    schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end, num_steps=num_steps)
    return [(pct, s[cachekey]) for pct, s in schedules]


-class PCLazyTextEncode:
-    CACHE_KEY = cache_key_prompt
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {"clip": ("CLIP", {"rawLink": True}), "text": ("STRING", {"multiline": True})},
-            "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"{unique_id}-")
-        return build_scheduled_prompts(graph, schedules, clip)
-
-
 class PCLazyTextEncodeAdvanced:
    CACHE_KEY = cache_key_prompt

@@ -282,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"},
        }
@@ -290,12 +263,23 @@ class PCLazyTextEncodeAdvanced:
    CATEGORY = "promptcontrol"
    FUNCTION = "apply"

-    def apply(self, clip, text, unique_id, tags="", start=0.1, end=1.0):
-        schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end)
-        graph = GraphBuilder(f"{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,
@@ -1,5 +1,5 @@
 import logging
-from .parser import parse_prompt_schedules
+from .parser import parse_prompt_schedules, expand_macros
 from .nodes_lazy import NODE_CLASS_MAPPINGS as LAZY_NODES
 import json
 import folder_paths
@@ -143,7 +143,6 @@ class PCSetPCTextEncodeSettings:
        return {
            "required": {"clip": ("CLIP",)},
            "optional": {
-                "steps": ("INT", {"default": 0, "min": 0, "max": 10000}),
                "mask_width": ("INT", {"default": 512, "min": 64, "max": 4096 * 4}),
                "mask_height": ("INT", {"default": 512, "min": 64, "max": 4096 * 4}),
                "sdxl_width": ("INT", {"default": 1024, "min": 0, "max": 4096 * 4}),
@@ -163,7 +162,6 @@ class PCSetPCTextEncodeSettings:
    def apply(
        self,
        clip,
-        steps=0,
        mask_width=512,
        mask_height=512,
        sdxl_width=1024,
@@ -174,7 +172,6 @@ class PCSetPCTextEncodeSettings:
        sdxl_crop_h=0,
    ):
        settings = {
-            "steps": steps,
            "mask_width": mask_width,
            "mask_height": mask_height,
            "sdxl_width": sdxl_width,
@@ -212,6 +209,24 @@ class PCExtractScheduledPrompt:
        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,
@@ -219,6 +234,7 @@ NODE_CLASS_MAPPINGS = {
    "PCSetLogLevel": PCSetLogLevel,
    "PCExtractScheduledPrompt": PCExtractScheduledPrompt,
    "PCSaveExpandedWorkflow": PCSaveExpandedWorkflow,
+    "PCMacroExpand": PCMacroExpand,
 }

 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -228,4 +244,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "PCSetLogLevel": "PC: Configure Logging (for debug)",
    "PCExtractScheduledPrompt": "PC: Extract Scheduled Prompt",
    "PCSaveExpandedWorkflow": "PC: Save Expanded Workflow (for debug)",
+    "PCMacroExpand": "PC: Expand Macros",
 }
@@ -8,7 +8,7 @@ log = logging.getLogger("comfyui-prompt-control")
 import re

 from functools import lru_cache
-from .utils import get_function
+from .utils import get_function, find_closing_paren

 if lark.__version__ == "0.12.0":
    from sys import executable
@@ -31,8 +31,9 @@ prompt: (emphasized | embedding | scheduled | alternate | sequence | loraspec |
 !emphasized: "(" prompt? ")"
        | "(" prompt ":" prompt ")"
        | "[" prompt "]"
-scheduled: "[" [[prompt] ":"] [prompt] ":" _WS? NUMBER ["," NUMBER] "]"
-        | "[" [[prompt] ":"] [prompt] ":" _WS? TAG "]"
+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] ">"
@@ -91,7 +92,7 @@ def parse_cuts(text):


 def flatten(x):
-    if type(x) in [str, tuple] or isinstance(x, dict) and "type" in x:
+    if type(x) in [str, tuple, int, type(None)] or isinstance(x, dict) and "type" in x:
        yield x
    else:
        for g in x:
@@ -103,13 +104,25 @@ def clamp(a, b, c):
    return min(max(a, b), c)


-def get_steps(tree):
-    res = [100]
+def get_steps(tree, num_steps):
+    res = [num_steps or 100]

    def tostep(s):
-        w = float(s) * 100
-        w = int(clamp(0, w, 100))
-        return w
+        steps = num_steps or 100
+        if "." in str(s) or not num_steps:
+            w = float(s)
+            value = w * steps
+        else:
+            w = int(s)
+            value = w
+
+        if w > 1 and not num_steps:
+            log.warning(
+                "You haven't configured the number of steps for Prompt Control to use, %s will be clipped to 1.0", w
+            )
+            value = steps
+
+        return int(clamp(0, value, steps))

    class CollectSteps(lark.Visitor):
        def scheduled(self, tree):
@@ -131,15 +144,14 @@ def get_steps(tree):
        def sequence(self, tree):
            steps = tree.children[1::2]
            for i, steps in enumerate(steps):
-                w = float(tree.children[i * 2 + 1]) * 100
-                tree.children[i * 2 + 1] = clamp(0, w, 100)
+                w = tostep(tree.children[i * 2 + 1])
+                tree.children[i * 2 + 1] = w
                res.append(w)

        def alternate(self, tree):
-            step_size = int(round(float(tree.children[-1] or 0.1), 2) * 100)
-            step_size = clamp(1, step_size, 100)
+            step_size = tostep(round(float(tree.children[-1] or 0.1), 2))
            tree.children[-1] = step_size
-            res.extend([x for x in range(step_size, 100, step_size)])
+            res.extend([x for x in range(step_size, num_steps or 100, step_size)])

    CollectSteps().visit(tree)

@@ -149,31 +161,38 @@ def get_steps(tree):
 def at_step(step, filters, tree):
    class AtStep(lark.Transformer):
        def scheduled(self, args):
+            before = None
+            during = None
+            after = None
            when_end = None
-            before, after, when, *rest = args
-            if isinstance(when, str):
-                return before or "" if when not in filters else after or ""
-
+            pl, when, *rest = args
            if rest:
                when_end = rest[0]

-            if when_end is not None and step <= when and before is not None:
-                return ""
+            pl = list(pl)
+            if len(pl) == 1:
+                (during,) = pl  # [after:0.5] == [::after:0.5,0.5]
+                if when_end is None:
+                    when_end = when
+                    after = during
+            elif len(pl) == 2:
+                during, after = pl  # [during:after:0.5] = [before::after:0.5,0.5]
+                if when_end is None:
+                    when_end = when
+                    before = during
+            else:
+                before, during, after = pl  # [before:during:after:0.5,0.8]

-            if when_end is not None and (step > when and step <= when_end):
-                # handle [a:0,1]
-                if before is None:
-                    return after or ""
-                return before or ""
+            if isinstance(when, str):
+                return before or "" if when not in filters else after or ""

-            if when_end is not None and step >= when_end:
-                # handle [a:0,1]
-                if before is None:
-                    return ""
-                return after or ""
+            if when_end is None:
+                when_end = 1000_000

            if step <= when:
                return before or ""
+            if when < step <= when_end:
+                return during or ""
            else:
                return after or ""

@@ -259,29 +278,31 @@ def at_step(step, filters, tree):


 class PromptSchedule(object):
-    def __init__(self, prompt, filters="", start=0.0, end=1.0):
+    # 0 num_steps means unconfigured
+    def __init__(self, prompt, filters="", start=0.0, end=1.0, num_steps=0):
        self.filters = filters
        self.start = start
        self.end = end
+        self.num_steps = num_steps
        self.prompt = prompt.strip()
        self.defaults = {}
        self.loaded_loras = {}

-        self.parsed_prompt = self._parse()
+        self.parsed_prompt = self._parse(num_steps)

    def __iter__(self):
        # Filter out zero, it's only useful for interpolation
        return (x for x in self.parsed_prompt if x[0] != 0)

-    def _parse(self):
+    def _parse(self, num_steps):
        filters = [x.strip() for x in self.filters.upper().split(",")]
        try:
            parsed = []
            tree = prompt_parser.parse(self.prompt)
-            steps = get_steps(tree)
+            steps = get_steps(tree, num_steps=num_steps)

            def f(x):
-                return round(x / 100, 2)
+                return round(x / (num_steps or 100), 2)

            for t in steps:
                p = at_step(t, filters, tree)
@@ -290,6 +311,7 @@ class PromptSchedule(object):
        except lark.exceptions.LarkError as e:
            log.error("Prompt editing parse error: %s", e)
            parsed = [[1.0, {"prompt": self.prompt, "loras": {}}]]
+            raise

        # Tag filtering may return redundant prompts, so filter them out here
        res = []
@@ -331,6 +353,7 @@ class PromptSchedule(object):
            filters=ifspecified(filters, self.filters),
            start=ifspecified(start, self.start),
            end=ifspecified(end, self.end),
+            num_steps=self.num_steps,
        )
        return p

@@ -345,17 +368,38 @@ class PromptSchedule(object):
        return len(self.parsed_prompt) - 1, self.parsed_prompt[-1]


-def replace_def(text):
+def parse_search(search):
+    arg_start = search.find("(")
+    args = ""
+    name = search.strip()
+    if arg_start > 0:
+        arg_end = find_closing_paren(search, arg_start)
+        name = search[:arg_start].strip()
+        args = search[arg_start + 1 : arg_end - 1]
+
+    if not name:
+        return None
+    args = args.strip()
+    # If 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 expand_macros(text):
    text, defs = get_function(text, "DEF", defaults=None)
    res = text
    prevres = text
    replacements = []
    for d in defs:
        r = d.split("=", 1)
-        if len(r) != 2 or not r[0].strip():
+        search = parse_search(r[0].strip())
+        if not search or len(r) != 2:
            log.warning("Ignoring invalid DEF(%s)", d)
            continue
-        replacements.append((r[0].strip(), r[1].strip()))
+        replacements.append((search, r[1].strip()))
    iterations = 0
    while True:
        iterations += 1
@@ -364,7 +408,6 @@ def replace_def(text):
            return text
        for search, replace in replacements:
            res = substitute_defcall(res, search, replace)
-            res = substitute_def(res, search, replace)
        if res == prevres:
            break
        prevres = res
@@ -375,22 +418,32 @@ def replace_def(text):


 def substitute_def(text, search, replace):
+    search, default_args = search
+    for i, v in enumerate(default_args):
+        replace = re.sub(rf"\${i+1}\b", v, replace)
    return re.sub(rf"\b{re.escape(search)}\b", replace, text)


 def substitute_defcall(text, search, replace):
-    text, defns = get_function(text, search, defaults=None, placeholder=f"DEFNCALL{search}")
-    for i, defn in enumerate(defns):
-        ph = f"\0DEFNCALL{search}{i}\0"
-        paramvals = [x.strip() for x in defn.split(";")]
+    name, default_args = search
+    text, defns = get_function(text, name, defaults=None, placeholder=f"DEFNCALL{name}")
+    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 = replace_def(prompt)
+    prompt = expand_macros(prompt)
    return PromptSchedule(prompt, **kwargs)
@@ -3,26 +3,14 @@ import re
 import torch
 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, smarter_split
+from .utils import safe_float, get_function, split_by_function, parse_floats, smarter_split
 from .adv_encode import advanced_encode_from_tokens
 from .cutoff import process_cuts
 from .parser import parse_cuts

-try:
-    from .nodes_attnmask import create_attention_hook
-    from comfy.hooks import set_hooks_for_conditioning
-
-    def set_cond_attnmask(cond, mask):
-        hook = create_attention_hook(mask)
-        return set_hooks_for_conditioning(cond, hooks=hook)
-
-except ImportError:
-
-    def set_cond_attnmask(cond, mask):
-        log.info("Attention masking is not available")
-        return cond
-
+from .attention_couple_ppm import set_cond_attnmask

 log = logging.getLogger("comfyui-prompt-control")

@@ -77,13 +65,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

@@ -138,6 +127,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,
@@ -155,52 +221,54 @@ 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"])
+    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,) = ConditioningAverage.addWeighted(None, [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 apply_weights(output, te_name, spec):
@@ -211,21 +279,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 = out * w
+            if pooled is not None:
+                pooled = pooled * pooled_w

        return out, pooled
    else:
-        if te_name in spec:
-            log.info("Weighting %s output by %s", te_name, spec[te_name])
-            output = output * spec[te_name]
+        if te_name in spec or default is not None:
+            w = spec.get(te_name, default)
+            log.info("Weighting %s output by %s", te_name, w)
+            output = output * w
        return output


@@ -246,15 +322,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,
@@ -262,7 +347,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


@@ -328,7 +413,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


@@ -418,6 +503,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)
@@ -428,55 +556,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))
+
+    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:
-        attn = False
-        if "ATTN()" in prompt:
-            prompt = prompt.replace("ATTN()", "")
-            attn = True
-            log.info("Using attention masking for prompt segment")
-        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
+        base_prompt, attn_couple_prompts = split_by_function(prompt, "COUPLE", defaults=None)

-        settings["start_percent"] = start_pct
-        settings["end_percent"] = end_pct
-        x = encode_prompt_segment(clip, prompt, settings, style, normalization)
-        if attn and mask is not None:
-            mask = settings.pop("mask")
-            strength = settings.pop("mask_strength")
-            x = set_cond_attnmask(x, mask * strength)
+        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)

-        conds.extend(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
@@ -0,0 +1,157 @@
+import unittest
+import numpy.testing as npt
+from os import environ
+
+clips = []
+
+import logging
+
+logging.basicConfig()
+
+
+def run(f, *args):
+    return getattr(f, f.FUNCTION)(*args)
+
+
+class TestEncode(unittest.TestCase):
+    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()
+        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("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)
+
+    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__":
+    print("Loading ComfyUI")
+    from comfy.sd import load_clip
+    import nodes
+    import comfy_extras.nodes_mask
+    from .nodes_base import PCTextEncode
+    from pathlib import Path
+
+    to_test = environ.get("TEST_TE", "clip_l").split()
+    model_path = environ.get("COMFYUI_MODEL_ROOT", ".")
+
+    te_root = (Path(model_path) / "text_encoders").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")
+    unittest.main()
@@ -0,0 +1,56 @@
+import unittest
+import numpy.testing as npt
+
+clip_l = None
+dual = None
+
+
+def run(f, *args):
+    return getattr(f, f.FUNCTION)(*args)
+
+
+class TestEncode(unittest.TestCase):
+    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][key], b[1][key])
+            else:
+                self.tensorsEqual(a[0], b[0])
+
+    def test_styles(self):
+        pc = PCTextEncode()
+        for k, clip in [("l", clip_l), ("dual", dual)]:
+            for style in ["comfy++", "A1111", "comfy++", "compel", "down_weight"]:
+                with self.subTest(f"TE {k} style {style} does not fail when encoding weights"):
+                    for normalization in ["none", "mean", "length", "length+mean"]:
+                        with self.subTest(f"TE {k} style {style} normalization {normalization}"):
+                            (c,) = run(
+                                pc,
+                                clip,
+                                f"STYLE(old+{style}, {normalization}) this prompt has weights, (a:1.2) (b:1.2)",
+                            )
+                            (c2,) = run(
+                                pc,
+                                clip,
+                                f"STYLE({style}, {normalization}) this prompt has weights, (a:1.2) (b:1.2)",
+                            )
+                            self.condEqual(c, c2)
+
+
+if __name__ == "__main__":
+    print("Loading ComfyUI")
+    import main
+
+    id(main)  # get rid of flake warning
+    import nodes
+    from .nodes_base import PCTextEncode
+
+    (clip_l,) = nodes.CLIPLoader().load_clip("clip_l.safetensors")
+    (dual,) = nodes.DualCLIPLoader().load_clip("clip_l.safetensors", "clip_g.safetensors", "sdxl")
+    print("Starting tests")
+    unittest.main()
@@ -1,5 +1,14 @@
 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):
@@ -7,95 +16,120 @@ def find_file(name):
    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

-    @mock.patch("prompt_control.utils.lora_name_to_file", find_file)
    def test_textencode(self):
-        clip = [0, 0]
-        from .nodes_lazy import PCLazyTextEncode, PCLazyTextEncodeAdvanced
-
        for p in ["test", "[test:0.2] test", "[test[test::0.5]]<lora:test:1>"]:
-            r1 = PCLazyTextEncode().apply(clip, p, "UID")
-            r2 = PCLazyTextEncodeAdvanced().apply(clip, p, "UID")
-            self.assertEqual(r1, r2)
+            r1 = te(p)
+            r2 = te(p, adv=True)
+            with self.subTest(f"Expansion: {p}"):
+                self.assertEqual(r1, r2)

-        r = PCLazyTextEncode().apply(clip, "test<lora:test:1>", "UID")
-        self.assertEqual(
-            r,
-            {
-                "result": (["UID-2", 0],),
-                "expand": {
-                    "UID-1": {"class_type": "PCTextEncode", "inputs": {"clip": [0, 0], "text": "test"}},
-                    "UID-2": {
-                        "class_type": "ConditioningSetTimestepRange",
-                        "inputs": {"conditioning": ["UID-1", 0], "start": 0.0, "end": 1.0},
+        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},
+                        },
                    },
                },
-            },
-        )
-        r = PCLazyTextEncode().apply(clip, "simple [test:0.1,0.5] prompt<lora:test:1>", "UID")
-        self.assertEqual(
-            r,
-            {
-                "result": (["UID-8", 0],),
-                "expand": {
-                    "UID-1": {"class_type": "PCTextEncode", "inputs": {"clip": [0, 0], "text": "simple  prompt"}},
-                    "UID-2": {
-                        "class_type": "ConditioningSetTimestepRange",
-                        "inputs": {"conditioning": ["UID-1", 0], "start": 0.0, "end": 0.1},
-                    },
-                    "UID-3": {"class_type": "PCTextEncode", "inputs": {"clip": [0, 0], "text": "simple test prompt"}},
-                    "UID-4": {
-                        "class_type": "ConditioningSetTimestepRange",
-                        "inputs": {"conditioning": ["UID-3", 0], "start": 0.1, "end": 0.5},
-                    },
-                    "UID-5": {"class_type": "PCTextEncode", "inputs": {"clip": [0, 0], "text": "simple  prompt"}},
-                    "UID-6": {
-                        "class_type": "ConditioningSetTimestepRange",
-                        "inputs": {"conditioning": ["UID-5", 0], "start": 0.5, "end": 1.0},
-                    },
-                    "UID-7": {
-                        "class_type": "ConditioningCombine",
-                        "inputs": {"conditioning_1": ["UID-2", 0], "conditioning_2": ["UID-4", 0]},
-                    },
-                    "UID-8": {
-                        "class_type": "ConditioningCombine",
-                        "inputs": {"conditioning_1": ["UID-7", 0], "conditioning_2": ["UID-6", 0]},
+            )
+        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):
-        from .nodes_lazy import PCLazyLoraLoader, PCLazyLoraLoaderAdvanced
-
-        model = [0, 1]
-        clip = [0, 0]
-        with self.assertLogs("comfyui-prompt-control", level="WARNING") as cm:
-            result = PCLazyLoraLoader().apply(model, clip, "prompt here <lora:nonexistent:1.0:0.5>", "UID")["expand"]
-            result_adv = PCLazyLoraLoaderAdvanced().apply(model, clip, "prompt here <lora:nonexistent:1.0:0.5>", "UID")[
-                "expand"
-            ]
+        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.assertIn("LoRA 'nonexistent' not found", cm.output[1])
        self.assertEqual(result, {})
        self.assertEqual(result_adv, {})

-        result = PCLazyLoraLoader().apply(model, clip, "<lora:test:1>", "UID")["expand"]
-        result2 = PCLazyLoraLoader().apply(model, clip, "prompt here <lora:test:1.0:0.5><lora:test:0:0.5>", "UID")[
-            "expand"
-        ]
-        result3 = PCLazyLoraLoaderAdvanced().apply(
-            model, clip, "prompt here <lora:test:1.0:0.5><lora:test:0:0.5>", "UID"
-        )["expand"]
+        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-1": {
+                "UID.0.0.1": {
                    "class_type": "LoraLoader",
                    "inputs": {
                        "model": [0, 1],
@@ -107,11 +141,11 @@ class GraphTests(unittest.TestCase):
                }
            },
        )
-        result = PCLazyLoraLoader().apply(model, clip, "<lora:test:1><lora:other:0.5>", "UID")["expand"]
+        result = loraloader("<lora:test:1><lora:other:0.5>")["expand"]
        self.assertEqual(
            result,
            {
-                "UID-1": {
+                "UID.0.0.1": {
                    "class_type": "LoraLoader",
                    "inputs": {
                        "model": [0, 1],
@@ -121,11 +155,11 @@ class GraphTests(unittest.TestCase):
                        "lora_name": "test.safetensors",
                    },
                },
-                "UID-2": {
+                "UID.0.0.2": {
                    "class_type": "LoraLoader",
                    "inputs": {
-                        "model": ["UID-1", 0],
-                        "clip": ["UID-1", 1],
+                        "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",
@@ -134,11 +168,11 @@ class GraphTests(unittest.TestCase):
            },
        )

-        result = PCLazyLoraLoader().apply(model, clip, "prompt here <lora:test:1.0:0.5>", "UID")["expand"]
+        result = loraloader("prompt here <lora:test:1.0:0.5>")["expand"]
        self.assertEqual(
            result,
            {
-                "UID-1": {
+                "UID.0.0.1": {
                    "class_type": "LoraLoader",
                    "inputs": {
                        "model": [0, 1],
@@ -151,48 +185,46 @@ class GraphTests(unittest.TestCase):
            },
        )

-        result = PCLazyLoraLoader().apply(model, clip, "prompt [<lora:test:0.5>:0.5]", "UID")["expand"]
-        result2 = PCLazyLoraLoaderAdvanced().apply(model, clip, "prompt [<lora:test:0.5>:0.5]", "UID")["expand"]
+        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-1": {
+            "UID.0.0.1": {
                "class_type": "CreateHookLora",
                "inputs": {"lora_name": "test.safetensors", "strength_model": 0.5, "strength_clip": 0.5},
            },
-            "UID-2": {
+            "UID.0.0.2": {
                "class_type": "CreateHookKeyframe",
                "inputs": {"strength_mult": 0.0, "start_percent": 0.0},
            },
-            "UID-3": {
+            "UID.0.0.3": {
                "class_type": "CreateHookKeyframe",
                "inputs": {
                    "start_percent": 0.5,
-                    "prev_hook_kf": ["UID-2", 0],
+                    "prev_hook_kf": ["UID.0.0.2", 0],
                    "strength_mult": 1.0,
                },
            },
-            "UID-4": {
+            "UID.0.0.4": {
                "class_type": "SetHookKeyframes",
-                "inputs": {"hooks": ["UID-1", 0], "hook_kf": ["UID-3", 0]},
+                "inputs": {"hooks": ["UID.0.0.1", 0], "hook_kf": ["UID.0.0.3", 0]},
            },
-            "UID-5": {
+            "UID.0.0.5": {
                "class_type": "SetClipHooks",
                "inputs": {
                    "clip": [0, 0],
-                    "hooks": ["UID-4", 0],
+                    "hooks": ["UID.0.0.4", 0],
                    "apply_to_conds": True,
                    "schedule_clip": True,
                },
            },
        }
        self.assertEqual(result, expected)
-        result2 = PCLazyLoraLoaderAdvanced().apply(model, clip, "prompt [<lora:test:0.5>:0.5]", "UID", start=0.6)[
-            "expand"
-        ]
+        result2 = loraloader("prompt [<lora:test:0.5>:0.5]", adv=True, start=0.6)["expand"]
        self.assertEqual(
            result2,
            {
-                "UID-1": {
+                "UID.0.0.1": {
                    "class_type": "LoraLoader",
                    "inputs": {
                        "model": [0, 1],
@@ -204,9 +236,7 @@ class GraphTests(unittest.TestCase):
                }
            },
        )
-        result2 = PCLazyLoraLoaderAdvanced().apply(model, clip, "prompt [<lora:test:0.5>:0.5]", "UID", end=0.5)[
-            "expand"
-        ]
+        result2 = loraloader("prompt [<lora:test:0.5>:0.5]", end=0.5)["expand"]
        self.assertEqual(result2, {})


@@ -1,5 +1,5 @@
 import unittest
-from .parser import parse_prompt_schedules as parse
+from .parser import parse_prompt_schedules as parse, expand_macros


 def prompt(until, text, *loras):
@@ -18,6 +18,19 @@ class TestParser(unittest.TestCase):
        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]"
@@ -71,6 +84,19 @@ class TestParser(unittest.TestCase):
        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]"
@@ -101,14 +127,33 @@ class TestParser(unittest.TestCase):

        p = parse("DEF(X=[($1):($1:$2):$2])X(test;0.7)")
        p2 = parse("[(test):(test:0.7):0.7]")
-        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+        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]")
-        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
+        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:
-            parse("DEF(X=recurse Y) DEF(Y=recurse X) X")
+            expand_macros("DEF(X=recurse Y) DEF(Y=recurse X) X")
        self.assertTrue("Unable to resolve DEFs" in str(c.exception))

    def test_misc(self):
@@ -148,11 +193,13 @@ class TestParser(unittest.TestCase):
        self.assertEqual(p.parsed_prompt, p2.parsed_prompt)
        for i, x in enumerate(["cat", "wolf", "tiger", "cat", "dog", "tiger", "cat", "wolf", "tiger", "cat"]):
            step = round((i * 0.1) + 0.1, 2)
-            self.assertPrompt(p3, step, step, x)
+            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)
-            self.assertPrompt(p4, step, step, *x)
+            with self.subTest(step):
+                self.assertPrompt(p4, step, step, *x)
        self.assertPrompt(p4, 0.7, 0.8, "wolf", ("canine", 1.0, 1.0))


@@ -87,31 +87,65 @@ def find_closing_paren(text, start):
    return len(text)


-def get_function(text, func, defaults, return_func_name=False, placeholder=""):
-    rex = re.compile(rf"\b{func}\(", re.MULTILINE)
+def get_function(text, func, defaults, return_func_name=False, placeholder="", return_dict=False):
+    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()
+        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 = None
+        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)

        if placeholder:
-            text = text[:start] + f"\0{placeholder}{count}\0" + text[end + 1 :]
+            text = text[:start] + f"\0{placeholder}{count}\0" + text[end:]
        else:
-            text = text[:start] + text[end + 1 :]
+            text = text[:start] + text[end:]
        match = rex.search(text)
        count += 1
    return text, instances


+def split_by_function(text, func, defaults=None):
+    """
+    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)
+    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)]:
@@ -1,14 +1,13 @@
 [project]
 name = "comfyui-prompt-control"
 description = "Nodes for convenient prompt editing, making many common operations prompt-controllable"
-version = "2.0.0-beta.9"
+version = "2.0.0-rc.8"
 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"
Author	SHA1	Message	Date
asagi4	8e59f140ff	v2.0.0-rc.8	2025-06-09 20:07:54 +03:00
asagi4	44044e962c	Very basic test for COUPLE	2025-06-09 20:06:25 +03:00
asagi4	04f36687c5	Fix skipping prompt segments by setting weight to 0	2025-06-09 20:05:50 +03:00
asagi4	3a8a360d03	Make testing less stupid	2025-06-09 19:53:05 +03:00
asagi4	d76331315a	Deduplicate tests	2025-06-09 19:15:55 +03:00
asagi4	e3a6050536	Don't call to() on every iteration	2025-06-09 18:37:53 +03:00
asagi4	7b001ace7b	Fix Attention Couple when combined with hooks on the CLIP (eg. LoRAs) All clones of the AC hook must maintain the same state. This feels a bit hacky though; there should be a better way Fixes #119	2025-06-09 17:25:26 +03:00
asagi4	f1de65f257	Don't override existing hooks. Unfortunately, this doesn't make things quite work; hmm.	2025-06-09 16:22:25 +03:00
asagi4	11aaa0ac7b	Fix indexing error	2025-06-09 01:06:44 +03:00
asagi4	85de3ef0d3	Fix links	2025-06-08 23:09:16 +03:00
asagi4	a73260ff34	Split t5 tests	2025-06-08 23:05:11 +03:00
asagi4	50a2e0abbf	Change ATTN() to COUPLE() and remove need for AND	2025-06-08 23:01:59 +03:00
asagi4	5cf45ca264	Make functions generally callable without argument lists	2025-06-08 18:57:25 +03:00
asagi4	bd4a787400	Use a helper function to parse function splits	2025-06-08 18:56:30 +03:00
asagi4	c9e5bc25c3	Need to do imports after torch mock, otherwise running tests on CPU torch fails	2025-06-08 17:22:54 +03:00
asagi4	d11ffa6e25	Don't pass in a custom prefix to GraphBuilder It breaks when PCLazyTextEncode etc. are called with list inputs. Tests needed adjusting after the change. Fixes #117	2025-06-08 16:49:37 +03:00
asagi4	8cc73a2e49	v2.0.0-rc.7	2025-06-08 01:17:08 +03:00
asagi4	27ae5f683e	Of course I forgot to test NegPiP	2025-06-08 01:15:16 +03:00
asagi4	68cda3663e	v2.0.0-rc.6	2025-06-08 01:02:16 +03:00
asagi4	3d46f705b6	Test downweighting too, and normalizations	2025-06-08 01:01:57 +03:00
asagi4	eec4bc4da9	Split old code into its own file for easy removal later	2025-06-08 00:55:43 +03:00
asagi4	3d3218e831	Tests for verifying refactor	2025-06-08 00:55:43 +03:00
asagi4	f4e57ec514	Switch on new implementation by default	2025-06-08 00:55:40 +03:00
asagi4	74f65c1b31	Use old from_masked batching for now I can't figure out why from_masked works differently from down_weight which also used that batching function but could be replaced with a simple torch.cat	2025-06-08 00:50:26 +03:00
asagi4	4d94cca88f	Restore adv_encoding to original implementation to compare them	2025-06-08 00:44:56 +03:00
asagi4	4569ecccf9	Manual testing...	2025-06-08 00:44:31 +03:00
asagi4	192e6d30d4	refactor adv_encode	2025-06-07 21:41:13 +03:00
asagi4	67112f11e0	T5 makes STYLE(perp) return NaNs. Just replace them with 0	2025-06-07 21:41:13 +03:00
asagi4	f761dfac86	Fix comfy++ with more than one weight, see #115 I'm not sure if this is correct, but it at least doesn't fail.	2025-06-07 14:09:32 +03:00
asagi4	278e733835	Fix normalization validity check	2025-06-07 14:09:32 +03:00
asagi4	2bf65720eb	Fix STYLE(perp) exception, see #115	2025-06-07 14:09:32 +03:00
asagi4	99f3af92b7	Add tests for weighting and a way to run manual testing	2025-06-07 14:09:14 +03:00
asagi4	2437cd4daf	Merge pull request #116 from pamparamm/pooled_none_check Partially resolve #115	2025-06-07 13:30:01 +03:00
asagi4	d262a7dc7a	Remove an extra clone.	2025-06-07 13:23:06 +03:00
Pam	6c319ad5b4	Partially resolve #115	2025-06-07 09:38:13 +05:00
asagi4	34056cac19	v2.0.0-rc.5	2025-06-06 19:52:02 +03:00
asagi4	8ae436abf1	Merge pull request #114 from pamparamm/negpip_option Use ppm_negpip option to detect NegPiP	2025-06-06 16:49:20 +03:00
Pam	c2ce2ce023	Use ppm_negpip option to detect NegPiP	2025-06-06 16:24:15 +05:00
asagi4	7a9e69ec31	Fix errors found in testing Who knew tests could be useful, too?	2025-06-05 23:49:02 +03:00
asagi4	aaff8dc7da	Encoding tests	2025-06-05 23:49:02 +03:00
asagi4	3e4722278a	v2.0.0-rc.4	2025-06-05 21:39:47 +03:00
asagi4	11cb430396	Support NegPiP without requiring a monkeypatch.	2025-06-05 21:38:31 +03:00
asagi4	9f1cbfd11c	Add a test suite for text encoding Can only be run manually and imports ComfyUI main to configure search paths, but at least it works...	2025-06-05 21:05:39 +03:00
asagi4	5b3a914f1d	Try not to fail in cases where masks have irregular shapes	2025-06-04 22:34:19 +03:00
asagi4	a5ffa1acd7	Documentation	2025-06-04 21:23:52 +03:00
asagi4	7f8783147b	I keep forgetting expand only works on batch size 1, see #108	2025-06-04 21:12:52 +03:00
asagi4	3c9b806e5f	Remove useless import	2025-06-04 21:04:15 +03:00
asagi4	a485c2655a	Fix case where negative prompt size changes lcm of cond size Also don't mutate existing hooks on input negative prompts if they exist.	2025-06-04 21:02:48 +03:00
asagi4	f888e69b00	Merge pull request #112 from pamparamm/attn_couple_batch Add PCAttentionCoupleBatchNegative	2025-06-04 20:59:52 +03:00
asagi4	b4b0858214	Fix multi-TE failure See #113	2025-06-04 08:48:40 +03:00
Pam	7a1cb2cf51	Fix latent masking	2025-06-04 00:30:23 +05:00
Pam	fbb6b5c8fa	Fix some uncond edgecases	2025-06-03 07:50:45 +05:00
Pam	6115c095cb	Fix attn couple batching with multiple positive schedules	2025-06-03 07:42:43 +05:00
Pam	04d3d2e959	Missing space	2025-06-03 03:46:56 +05:00
Pam	e4f64837ef	Add PCAttentionCoupleBatchNegative; Revert some optimizations in AttentionCoupleHook	2025-06-03 03:17:47 +05:00
asagi4	4a785b294b	Avoid hardcoding length in adv encode Makes these not explode on T5 at least. They seems to produce the same results still.	2025-06-03 00:47:57 +03:00
asagi4	b3195a6297	Stop using batched_clip_encode, it doesn't do anything?	2025-06-02 23:32:19 +03:00
asagi4	5c52bffc9d	Docs	2025-06-02 22:16:42 +03:00
asagi4	ac2d275dfb	Clarify TE lookups	2025-06-02 21:37:27 +03:00
asagi4	399f992a26	Fix BREAK while maintaining old behaviour, add CAT instead	2025-06-02 20:36:27 +03:00
asagi4	672a2a09bb	BREAK is now ConditioningConcat	2025-06-02 19:36:29 +03:00
asagi4	e3629961ce	Add AVG(weight) to work as ConditioningAverage	2025-06-02 19:14:20 +03:00
asagi4	ddac624ad1	Add NBREAK (Warning: unstable. Name is likely to change) NBREAK should have the same behaviour as ComfyUI's ConditioningConcat See #111	2025-06-02 17:39:09 +03:00
asagi4	99ddfe357e	DEF docs	2025-06-01 22:44:06 +03:00
asagi4	110d5248a0	DEF tests	2025-06-01 22:34:39 +03:00
asagi4	61b1ecc88e	Clean up tests a bit	2025-06-01 22:34:39 +03:00
asagi4	05b0b2ad26	Set the default value of $1 to empty with DEF(MACRO()=)	2025-06-01 22:34:34 +03:00
asagi4	5831608c4e	Add a macro expansion node	2025-06-01 21:41:03 +03:00
asagi4	c815bb44f1	Reduce logging verbosity	2025-05-31 01:19:37 +03:00
asagi4	e55c50e9d7	Fix the case with more than one coupled conditioning	2025-05-31 01:09:29 +03:00
asagi4	1b0ff62d10	v2.0.0-rc.3	2025-05-31 00:25:05 +03:00
asagi4	cf93093d59	Fix long prompts with Attention Couple Broken by moving the LCM calculation outside the loop See #108	2025-05-31 00:22:42 +03:00
asagi4	fc15a89a2f	v2.0.0-rc.2	2025-05-30 22:41:47 +03:00
asagi4	57c092bccf	Doc reorganization, part 4	2025-05-30 22:41:27 +03:00
asagi4	88f77a8124	Doc reorganization, part 3	2025-05-30 22:22:54 +03:00
asagi4	a9c2487c0c	Doc reorganization, part 2	2025-05-30 22:17:56 +03:00
asagi4	75bced7d2b	Doc reorganization	2025-05-30 22:12:57 +03:00
asagi4	4b285be07e	Merge pull request #109 from asagi4/attn_couple_refactor Attention couple refactor	2025-05-30 21:23:01 +03:00
asagi4	200d9f9daf	Cleanup: remove debug function	2025-05-30 21:21:28 +03:00
asagi4	d33208b1c3	refactor: calculate conds_kv only once	2025-05-30 18:22:23 +03:00
asagi4	ffa64816c0	Refactor: Remove loop	2025-05-30 17:59:36 +03:00
asagi4	6ffbf05d7d	refactor debug: LCM debug prints	2025-05-30 17:43:46 +03:00
asagi4	892a70d53b	refactor: remove self.batch_size	2025-05-30 17:32:56 +03:00
asagi4	20711358a2	refactor: cond_kvs is never empty with the hook	2025-05-30 17:20:36 +03:00
asagi4	453580545c	pyflakes cleanup	2025-05-30 17:09:55 +03:00
asagi4	98d78df7ba	refactor 5: inline get_mask	2025-05-30 17:09:55 +03:00
asagi4	bdd56410dc	Refactor 4: This produces correct output	2025-05-30 17:09:51 +03:00
asagi4	0289564e55	refactor 3: cond_pos should not matter anymore	2025-05-30 17:09:51 +03:00
asagi4	1e05d1a8cc	Refactor 2: new cond amount can be calculated from num_conds	2025-05-30 17:09:47 +03:00
asagi4	dc6fd0fc63	Debug function	2025-05-30 15:24:15 +03:00
asagi4	a356bddcc7	Refactor 1: Remove UNCOND special casing	2025-05-30 15:14:22 +03:00
asagi4	b8081e5736	Revert for loop removals, they change batched outputs somehow and I can't figure out why. This reverts commit `5e3ab1f51a`. This reverts commit `b21de76cd5`.	2025-05-30 14:36:55 +03:00
asagi4	aa00c26365	Fix FILL()	2025-05-30 03:48:06 +03:00
asagi4	e913bad73c	Docs and some more tests	2025-05-30 03:37:35 +03:00
asagi4	5c1b739b82	Extend scheduling syntax with [before:during:after:0.5,0.7]	2025-05-30 02:51:51 +03:00
asagi4	5e3ab1f51a	Remove the other for loop too	2025-05-29 23:08:40 +03:00
asagi4	b21de76cd5	Simplify attention couple code because ComfyUI will handle unmixing cond/uncond for us	2025-05-29 22:12:47 +03:00
asagi4	e4a27d01ee	docs	2025-05-29 22:00:24 +03:00
asagi4	42cdfa0f5a	Properly supports prompt weights with attention masking.	2025-05-29 21:30:37 +03:00
asagi4	98292e2bc8	Reorder README a bit	2025-05-26 21:33:45 +03:00
asagi4	f0c8e2e873	Adjust syntax for attention couple to be a bit more convenient. See #108	2025-05-26 21:25:21 +03:00
asagi4	c4ac37333d	Switch Attention Couple implementation to one based on ppm Also removes compatibility code with older ComfyUI	2025-05-26 20:11:36 +03:00
asagi4	2534e002ad	Add default values to DEF	2025-05-26 00:09:22 +03:00
asagi4	fd4823fd75	Remove debug logging	2025-05-25 22:13:14 +03:00
asagi4	a6f230ff8b	Try to make ATTN() more like attention couple. See #107	2025-05-25 21:55:19 +03:00
asagi4	633b2f05e0	Release v2.0.0-rc.1 properly	2025-05-21 19:10:53 +03:00
asagi4	a15135ddc5	Remove misleading instruction that no longer applies	2025-05-21 19:09:08 +03:00
asagi4	0d7e2a4e60	No, I do not want CUDA	2025-05-20 21:46:25 +03:00
asagi4	c5495832c5	Use CPU torch	2025-05-20 21:35:32 +03:00
asagi4	63d2cb3e0c	Tests are broken again...	2025-05-20 21:28:11 +03:00
asagi4	95832e801b	Make TE_WEIGHT more convenient	2025-05-20 21:13:22 +03:00
asagi4	01bd5568d0	New function: TE Fixes #106	2025-05-20 20:47:53 +03:00
asagi4	36b3638f4e	refactor: tokenize_chunks	2025-05-20 19:45:16 +03:00
asagi4	d46000ef78	Allow TE_WEIGHT(all=1.1)	2025-05-20 19:21:01 +03:00
asagi4	aba246a33c	Add a note about compositing to clarify docs, fixes #101	2025-05-05 15:33:25 +03:00
asagi4	42ae22db83	Rename the 2pass workflow for now; it needs review	2025-04-30 16:39:48 +03:00
asagi4	cb6de285cb	Fix link in README	2025-04-03 20:51:59 +03:00
asagi4	49a073bb12	Update the comparison workflow	2025-04-03 20:45:15 +03:00
asagi4	e9afe779ae	Refresh template workflow	2025-04-03 20:09:21 +03:00
asagi4	fa3b4f7da3	Fix brain typo	2025-04-02 23:25:38 +03:00
asagi4	7a76cc8c72	v2.0.0-beta.11	2025-04-02 23:17:55 +03:00
asagi4	6b1e2a5a8a	Remove steps from PCSetPCTextEncodeSettings, it can't be used	2025-04-02 23:17:55 +03:00
asagi4	9aee531c09	Allow configuring a steps value via the Advanced nodes	2025-04-02 23:17:55 +03:00
asagi4	c08bf395a6	Test cleanup	2025-04-02 23:17:55 +03:00
asagi4	1964708997	Combine PCLazyTextEncode with PCLazyTextEncodeAdvanced	2025-04-02 23:17:55 +03:00
Robin Huang snomiao asagi4	1c4b5ce0c4	chore(publish): update GitHub Actions workflow for node publishing (#99 ) - Add permissions to allow issue writing - Update action version to v1 for publish-node-action - Add condition to run job only for 'asagi4' repository owner Co-authored-by: snomiao <snomiao+comfy-pr@gmail.com> Co-authored-by: asagi4 <130366179+asagi4@users.noreply.github.com>	2025-03-22 17:12:08 +02:00
asagi4	5eabbb419c	Re-enable tests	2025-03-10 21:40:55 +02:00
asagi4	53400a029b	v2.0.0-beta.10	2025-03-10 21:39:37 +02:00
asagi4	c39605eec4	Gah, tests broke. I hate mocking	2025-03-10 21:38:03 +02:00
asagi4	dc62e638ed	How on earth does this syntax work	2025-03-10 20:30:38 +02:00
asagi4	109cac16ef	...	2025-03-10 20:27:07 +02:00
asagi4	cf6c2b3e6a	Fix github actions	2025-03-10 20:22:36 +02:00