Gah, tests broke.

I hate mocking

.github/ISSUE_TEMPLATE/bug_report.md

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

.github/workflows/publish.yml

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

.github/workflows/tests.yml

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

.github/workflows/tests_with_comfy.yml

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

Makefile

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

README.md

@@ -2,14 +2,13 @@
 
 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.
 
+A `Basic Text to Image` template is included with the extension, and can be loaded from ComfyUI's template library.
+
 ## 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. This has some advantages:
+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
 
-- ComfyUI will not re-run unchanged parts of generated graphs. This is especially useful for two-pass workflows where previously you'd be forced to re-run the first sampling pass even with filtering. That is no longer the case and it does the right thing.
-- The generated graph is often exactly equivalent to a manually built workflow using native ComfyUI nodes. There are no more weird sampling hooks that could cause problems with other nodes
-
-Prompt Control also comes with `PCTextEncode`, which provides advanced text encoding with many additional features compared to ComfyUI's base `CLIPTextEncode`.
+Prompt Control comes with `PCTextEncode`, which provides advanced text encoding with many additional features compared to ComfyUI's base `CLIPTextEncode`.
 
 ### Removed features
 
@@ -17,10 +16,6 @@ Prompt Control also comes with `PCTextEncode`, which provides advanced text enco
 - LoRA block weight integration; ditto, for now.
 
 
-### Is it stable now?
-
-Unless I run into bugs or significant annoyances that require changing the interface, it probably won't change too much, but until I tag 2.0, everything can change.
-
 ### Everything broke, where are the old nodes?
 
 If you really need them, you can install the [legacy nodes](https://github.com/asagi4/comfyui-prompt-control-legacy). However, I will not fix bugs in those nodes, and I strongly recommend just migrating your workflows to the new nodes.
@@ -42,7 +37,6 @@ 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.
@@ -65,6 +59,8 @@ Then restart ComfyUI afterwards.
 
 # Core nodes
 
+**Note**: The documentation refers to the nodes with their internal names for consistency. The display name may change, but ComfyUI's search will always find the nodes with the internal name. `PCLazyTextEncode` and `PCLazyLoraLoader` are the main ones you'll want to use, also known as `PC: Schedule Prompt` and `PC: Schedule LoRas`.
+
 ## PCLazyTextEncode and PCLazyTextEncodeAdvanced
 
 `PCLazyTextEncode` uses ComfyUI's lazy graph execution mechanism to generate a graph of `PCTextEncode` and `SetConditioningTimestepRange` nodes from a prompt with schedules. This has the advantage that if a part of the schedule doesn't change, ComfyUI's caching mechanism allows you to avoid re-encoding the non-changed part.
@@ -168,4 +164,8 @@ The parameters affect how the masked and unmasked prompts are combined to produc
 
 # Known issues
 
-- None at the moment
+- ComfyUI's caching mechanism has an issue that makes it unnecessarily invalidate caches for certain inputs; you'll still get some benefit from the lazy nodes, but changing inputs that shouldn't affect downstream nodes (especially if using filtering) will still cause them to be recomputed because ComfyUI doesn't realize the inputs haven't changed.
+
+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.

__init__.py

@@ -15,7 +15,7 @@ log = logging.getLogger("comfyui-prompt-control")
 log.propagate = False
 if not log.handlers:
     h = logging.StreamHandler(sys.stdout)
-    h.setFormatter(logging.Formatter("[%(levelname)s] PromptControl: %(message)s"))
+    h.setFormatter(logging.Formatter("[PromptControl] %(levelname)s: %(message)s"))
     log.addHandler(h)
 
 if os.environ.get("PROMPTCONTROL_DEBUG"):
@@ -23,19 +23,28 @@ if os.environ.get("PROMPTCONTROL_DEBUG"):
 else:
     log.setLevel(logging.INFO)
 
+cache_hack = importlib.import_module(".prompt_control.cache_hack", package=__name__)
+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.append("hooks")
+    nodes.extend(["hooks"])
 else:
-    log.warning(
-        "Your ComfyUI version is too old, can't import comfy.hooks for PCEncodeSchedule and PCLoraHooksFromSchedule. Update your installation."
-    )
+    log.error("Your ComfyUI version is too old, can't import comfy.hooks. Update your installation.")
 
 for node in nodes:
     mod = importlib.import_module(f".prompt_control.nodes_{node}", package=__name__)
     NODE_CLASS_MAPPINGS.update(mod.NODE_CLASS_MAPPINGS)
     NODE_DISPLAY_NAME_MAPPINGS.update(mod.NODE_DISPLAY_NAME_MAPPINGS)
+
+for node in optional_nodes:
+    try:
+        mod = importlib.import_module(f".prompt_control.nodes_{node}", package=__name__)
+        NODE_CLASS_MAPPINGS.update(mod.NODE_CLASS_MAPPINGS)
+        NODE_DISPLAY_NAME_MAPPINGS.update(mod.NODE_DISPLAY_NAME_MAPPINGS)
+    except ImportError:
+        log.info(f"Could not import optional nodes: {node}; continuing anyway")

doc/syntax.md

@@ -24,12 +24,12 @@ a [red:[blue::0.7]:0.5] cat
 
 switches from `a red cat` to `a blue cat` at 0.5 and to `a cat` at 0.7
 
-
-**Note:** As a special case, `[cat:0.5]` is like `[:cat:0.5]` meaning it switches from empty to `cat` at 0.5. Currently, `[:cat:0.5]` doesn't actually parse correctly, so you **must** use the shortcut form
+**Note:** As a special case, `[cat:0.5]` is equivalent to `[:cat:0.5]` meaning it switches from empty to `cat` at 0.5.
 
 ### Range expressions
 
 You can also use `a [during:after:0.3,0.7]` as a shortcut. The prompt be `a` until 0.3, `a during` until 0.7, and then `a after`. This form is equivalent to `[[during:after:0.7]:0.3]`
+
 For convenience, `[during:0.1,0.4]` is equivalent to `[during::0.1,0.4]`
 
 ## Tag selection
@@ -79,7 +79,6 @@ Might be useful with Jinja templating (see https://github.com/asagi4/comfyui-uti
 ```
 generates a LoRA schedule based on a sinewave
 
-
 # Basic prompt syntax
 
 This syntax is also available in outside scheduled prompts, where applicable.
@@ -98,7 +97,6 @@ cat :1 AND dog :2
 ```
 The weight defaults to 1 and are normalized so that `a:2 AND b:2` is equal to `a AND b`. `AND` is processed after schedule parsing, so you can change the weight mid-prompt: `cat:[1:2:0.5] AND dog`
 
-
 ## Functions
 
 There are some "functions" that can be included in a prompt to do various things. 
@@ -135,7 +133,7 @@ These functions are applied to each prompt chunk **after** `BREAK`, `AND` etc. h
 
 Multiple instances of these functions are applied in the order they appear in the prompt.
 
-**NOTE:** These functions are *not* smart about syntax and will break emphasis if the separator occurs inside parentheses. I might fix this at some point, but for now, keep this in mind.
+**NOTE** To avoid breaking emphasis syntax, the functions ignore any separators inside parentheses
 
 For example:
 - `SHIFT(1) cat, dog, tiger, mouse` does a shift and results in `dog, tiger, mouse, cat`. (whitespace may vary)
@@ -195,3 +193,52 @@ The order of the `FEATHER` and `MASK` calls doesn't matter; you can have `FEATHE
 
 ## Miscellaneous
 - `<emb:xyz>` is alternative syntax for `embedding:xyz` to work around a syntax conflict with `[embedding:xyz:0.5]` which is parsed as a schedule that switches from `embedding` to `xyz`.
+
+# Experimental features
+
+Experimental features are unstable and may disappear or change without warning.
+
+## DEF
+
+You can define "prompt macros" by using `DEF`:
+```
+DEF(MYMACRO=this is a prompt)
+[(MYMACRO:0.6):(MYMACRO:1.1):0.5]
+```
+is equivalent to
+```
+[(this is a prompt:0.5):(this is a prompt:1.1):0.5]
+```
+
+It's also possible to give parameters to a macro:
+```
+DEF(MYMACRO=[(prompt $1:$2):(prompt $1:$3):$4])
+MYMACRO(test; 1.1; 0.7; 0.2)
+```
+gives
+```
+[(prompt test:1.1):(prompt test:0.7):0.2]
+```
+in this form, the variables $N (where N is any number corresponding to a positional parameter) will be replaced with the given parameter. The parameters must be separated with a semicolon.
+
+Note that unspecified parameters will not be substituted:
+```
+DEF(mything=a $1 b $2)
+mything
+mything(A)
+```
+gives
+```
+a $1 b $2
+a A b $2
+```
+
+Macros are expanded before any other parsing takes place. The expansion continues until no further changes occur. Recursion will raise an error.
+
+## Attention masking
+
+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.
+
+## 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.

example_workflows/Basic Text to Image.json

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

prompt_control/cache_hack.py

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

prompt_control/nodes_attnmask.py

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

prompt_control/nodes_base.py

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

prompt_control/nodes_hooks.py

@@ -84,5 +84,5 @@ NODE_CLASS_MAPPINGS = {
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "PCLoraHooksFromText": "PC LoRA Hooks From Text (non-lazy)",
+    "PCLoraHooksFromText": "PC: LoRA Hooks From Text (non-lazy)",
 }

prompt_control/nodes_lazy.py

@@ -1,12 +1,31 @@
 import logging
 from .parser import parse_prompt_schedules
-from comfy_execution.graph_utils import GraphBuilder
+from comfy_execution.graph_utils import GraphBuilder, is_link
 
-from .prompts import get_function
+from comfy_execution.graph import ExecutionBlocker
+
+from .utils import get_function
 
 log = logging.getLogger("comfyui-prompt-control")
 
 from .utils import consolidate_schedule, find_nonscheduled_loras
+import json
+
+
+def _cache_key(cachekey, inputs):
+    out = inputs.copy()
+    text = inputs.get("text")
+    if text is not None and not is_link(text):
+        out["text"] = cache_key_from_inputs(cachekey, **inputs)
+    return out
+
+
+def cache_key_prompt(inputs):
+    return _cache_key("prompt", inputs)
+
+
+def cache_key_lora(inputs):
+    return _cache_key("loras", inputs)
 
 
 def create_lora_loader_nodes(graph, model, clip, loras):
@@ -24,9 +43,10 @@ def create_lora_loader_nodes(graph, model, clip, loras):
 
 
 def create_hook_nodes_for_lora(graph, path, info, existing_node, start_pct, end_pct):
+    prev_keyframe = None
     next_keyframe = None
     if not existing_node:
-        log.debug("Creating hook for %s", path)
+        log.debug("Creating hook for %s, weight=%s, weight_clip=%s", path, info["weight"], info["weight_clip"])
         hook_node = graph.node("CreateHookLora")
         hook_node.set_input("lora_name", path)
         hook_node.set_input("strength_model", info["weight"])
@@ -59,7 +79,7 @@ def create_hook_nodes_for_lora(graph, path, info, existing_node, start_pct, end_
     next_keyframe.set_input("strength_mult", 1.0)
     prev_hook_kf = next_keyframe.out(0)
     if end_pct < 1.0:
-        log.debug("Creating end keyframe for %s, start=%s", path, start_pct)
+        log.debug("Creating end keyframe for %s, start=%s", path, end_pct)
         next_keyframe = graph.node("CreateHookKeyframe")
         next_keyframe.set_input("strength_mult", 0.0)
         next_keyframe.set_input("start_percent", end_pct)
@@ -70,8 +90,12 @@ def create_hook_nodes_for_lora(graph, path, info, existing_node, start_pct, end_
 def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_hooks=True):
     # This gets rid of non-existent LoRAs
     consolidated = consolidate_schedule(schedule)
-    non_scheduled = find_nonscheduled_loras(consolidated)
-    model, clip = create_lora_loader_nodes(graph, model, clip, non_scheduled)
+    if model is not None:
+        non_scheduled = find_nonscheduled_loras(consolidated)
+        model, clip = create_lora_loader_nodes(graph, model, clip, non_scheduled)
+    else:
+        non_scheduled = {}
+        model = ExecutionBlocker("No model provided to PCLazyLoRALoader or PCLazyLoRALoaderAdvanced")
 
     hook_nodes = {}
     start_pct = 0.0
@@ -100,21 +124,22 @@ def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_h
     # Finally, combine all hooks and optionally apply
     if len(hooks) > 0:
         res = hooks[0]
-        for h in hooks[:1]:
+        for h in hooks[1:]:
             n = graph.node("CombineHooks2")
             n.set_input("hooks_A", res.out(0))
             n.set_input("hooks_B", h.out(0))
             res = n
         res = res.out(0)
-    if apply_hooks:
-        n = graph.node("SetClipHooks")
-        n.set_input("clip", clip)
-        n.set_input("hooks", res)
-        n.set_input("apply_to_conds", True)
-        n.set_input("schedule_clip", True)
-        clip = n.out(0)
+        if apply_hooks:
+            n = graph.node("SetClipHooks")
+            n.set_input("clip", clip)
+            n.set_input("hooks", res)
+            n.set_input("apply_to_conds", True)
+            n.set_input("schedule_clip", True)
+            clip = n.out(0)
 
     r = graph.finalize()
+    log.debug("LazyLoraLoader built graph: %s", json.dumps(r))
 
     if return_hooks:
         ret = (model, clip, res)
@@ -125,16 +150,18 @@ def build_lora_schedule(graph, schedule, model, clip, apply_hooks=True, return_h
 
 
 class PCLazyLoraLoaderAdvanced:
+    CACHE_KEY = cache_key_lora
+
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
-                "text": ("STRING", {"multiline": True}),
                 "model": ("MODEL", {"rawLink": True}),
                 "clip": ("CLIP", {"rawLink": True}),
-                "apply_hooks": ("BOOLEAN", {"default": True}),
+                "text": ("STRING", {"multiline": True}),
             },
             "optional": {
+                "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}),
@@ -147,20 +174,22 @@ class PCLazyLoraLoaderAdvanced:
     CATEGORY = "promptcontrol"
     FUNCTION = "apply"
 
-    def apply(self, model, clip, text, apply_hooks, unique_id, tags="", start=0.0, end=1.0):
-        schedule = parse_prompt_schedules(text).with_filters(filters=tags, start=start, end=end)
-        graph = GraphBuilder(f"PCLazyLoraLoaderAdvanced-{unique_id}")
+    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)
 
 
 class PCLazyLoraLoader:
+    CACHE_KEY = cache_key_lora
+
     @classmethod
     def INPUT_TYPES(s):
         return {
-            "required": {
+            "optional": {
                 "model": ("MODEL", {"rawLink": True}),
                 "clip": ("CLIP", {"rawLink": True}),
-                "text": ("STRING", {"multiline": True}),
+                "text": ("STRING", {"multiline": True, "default": ""}),
             },
             "hidden": {"unique_id": "UNIQUE_ID"},
         }
@@ -173,16 +202,20 @@ class PCLazyLoraLoader:
     CATEGORY = "promptcontrol"
     FUNCTION = "apply"
 
-    def apply(self, model, clip, text, unique_id):
-        graph = GraphBuilder(f"PCLazyLoraLoader-{unique_id}")
+    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 build_scheduled_prompts(graph, schedules, clip):
     nodes = []
     start_pct = 0.0
-    prompt_cache = {}
     for end_pct, c in schedules:
         p = c["prompt"]
         p, classnames = get_function(p, "NODE", ["PCTextEncode", "text"])
@@ -191,12 +224,9 @@ def build_scheduled_prompts(graph, schedules, clip):
         if classnames:
             classname = classnames[0][0]
             paramname = classnames[0][1]
-        node = prompt_cache.get((p, classname, paramname))
-        if not node:
-            node = graph.node(classname)
-            node.set_input("clip", clip)
-            node.set_input(paramname, p)
-            prompt_cache[(p, classname, paramname)] = node
+        node = graph.node(classname)
+        node.set_input("clip", clip)
+        node.set_input(paramname, p)
         timestep = graph.node("ConditioningSetTimestepRange")
         timestep.set_input("conditioning", node.out(0))
         timestep.set_input("start", start_pct)
@@ -210,15 +240,24 @@ def build_scheduled_prompts(graph, schedules, clip):
         combiner.set_input("conditioning_2", othernode.out(0))
         node = combiner
 
-    return {"result": (node.out(0),), "expand": graph.finalize()}
+    g = graph.finalize()
+    log.debug("Built graph: %s", json.dumps(g))
+
+    return {"result": (node.out(0),), "expand": g}
+
+
+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)
+    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})},
-            # "optional": {"defaults": ("SCHEDULE_DEFAULTS",)},
             "hidden": {"unique_id": "UNIQUE_ID"},
         }
 
@@ -229,11 +268,13 @@ class PCLazyTextEncode:
 
     def apply(self, clip, text, unique_id):
         schedules = parse_prompt_schedules(text)
-        graph = GraphBuilder(f"PCEncodeLazy-{unique_id}")
+        graph = GraphBuilder(f"{unique_id}-")
         return build_scheduled_prompts(graph, schedules, clip)
 
 
 class PCLazyTextEncodeAdvanced:
+    CACHE_KEY = cache_key_prompt
+
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -251,8 +292,8 @@ class PCLazyTextEncodeAdvanced:
     FUNCTION = "apply"
 
     def apply(self, clip, text, unique_id, tags="", start=0.1, end=1.0):
-        schedules = parse_prompt_schedules(text).with_filters(start=start, end=end, filters=tags)
-        graph = GraphBuilder(f"PCLazyTextEncodeAdvanced-{unique_id}")
+        schedules = parse_prompt_schedules(text, filters=tags, start=start, end=end)
+        graph = GraphBuilder(f"{unique_id}-")
         return build_scheduled_prompts(graph, schedules, clip)
 
 

prompt_control/nodes_tools.py

@@ -1,8 +1,97 @@
 import logging
+from .parser import parse_prompt_schedules
+from .nodes_lazy import NODE_CLASS_MAPPINGS as LAZY_NODES
+import json
+import folder_paths
+from pathlib import Path
+from comfy_execution.graph_utils import is_link
 
 log = logging.getLogger("comfyui-prompt-control")
 
 
+class PCSaveExpandedWorkflow:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "any": ("*", {}),
+            },
+            "hidden": {
+                "prompt": "DYNPROMPT",
+            },
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(self, input_types):
+        return True
+
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+    CATEGORY = "promptcontrol/tools"
+    DESCRIPTION = "Saves the current expanded dynamic prompt into a JSON file"
+
+    FUNCTION = "apply"
+
+    def apply(self, any, prompt):
+        full_output_folder, filename, counter, subfolder, prefix = folder_paths.get_save_image_path(
+            "pc_workflow_debug", self.output_dir
+        )
+        p = {}
+        input_replace_map = {}
+        for node in prompt.all_node_ids():
+            n = prompt.get_node(node)
+            t = n["class_type"]
+            if t in LAZY_NODES:
+                expanded_prompt = LAZY_NODES[t]().apply(**n["inputs"], unique_id=node)
+                for k in expanded_prompt["expand"]:
+                    p[k] = expanded_prompt["expand"][k]
+                for i, _ in enumerate(expanded_prompt["result"]):
+                    input_replace_map[(node, i)] = [k, i]
+            else:
+                p[node] = n
+        for k in p:
+            for ik in p[k]["inputs"]:
+                x = p[k]["inputs"][ik]
+                if is_link(x) and tuple(x) in input_replace_map:
+                    p[k]["inputs"][ik] = input_replace_map[tuple(x)]
+        file = f"{filename}_{counter:05}_.json"
+        full_path = Path(full_output_folder) / file
+        with open(full_path, "w") as f:
+            log.info(f"Saving workflow to {full_path}")
+            json.dump(p, f)
+
+        return ()
+
+
+class PCSetLogLevel:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clip": ("CLIP",),
+            },
+            "optional": {
+                "level": (["INFO", "DEBUG", "WARNING", "ERROR"], {"default": "INFO"}),
+            },
+        }
+
+    def apply(self, clip, level="INFO"):
+        log.setLevel(getattr(logging, level))
+        log.info("Set logging level to %s", level)
+        return (clip,)
+
+    RETURN_TYPES = ("CLIP",)
+    CATEGORY = "promptcontrol/tools"
+    DESCRIPTION = (
+        "A debug node to configure Prompt Control logging level. Pass a CLIP through it before you run any PC nodes"
+    )
+
+    FUNCTION = "apply"
+
+
 class PCAddMaskToCLIP:
     @classmethod
     def INPUT_TYPES(s):
@@ -14,8 +103,9 @@ class PCAddMaskToCLIP:
         }
 
     RETURN_TYPES = ("CLIP",)
-    CATEGORY = "promptcontrol/v2"
+    CATEGORY = "promptcontrol/tools"
     FUNCTION = "apply"
+    DESCRIPTION = "Attaches a mask to a CLIP object so that they can be referred to in a prompt using IMASK(). Using this node multiple times adds more masks rather than replacing existing ones."
 
     def apply(self, clip, mask=None):
         return PCAddMaskToCLIPMany().apply(clip, mask1=mask)
@@ -35,8 +125,9 @@ class PCAddMaskToCLIPMany:
         }
 
     RETURN_TYPES = ("CLIP",)
-    CATEGORY = "promptcontrol/v2"
+    CATEGORY = "promptcontrol/tools"
     FUNCTION = "apply"
+    DESCRIPTION = "Multi-input version of PCAddMaskToCLIP, for convenience"
 
     def apply(self, clip, mask1=None, mask2=None, mask3=None, mask4=None):
         clip = clip.clone()
@@ -65,8 +156,9 @@ class PCSetPCTextEncodeSettings:
         }
 
     RETURN_TYPES = ("CLIP",)
-    CATEGORY = "promptcontrol/v2"
+    CATEGORY = "promptcontrol/tools"
     FUNCTION = "apply"
+    DESCRIPTION = "Configures default values for PCTextEncode"
 
     def apply(
         self,
@@ -97,14 +189,43 @@ class PCSetPCTextEncodeSettings:
         return (clip,)
 
 
+class PCExtractScheduledPrompt:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "text": ("STRING", {"multiline": True}),
+                "at": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 1.0, "step": 0.01}),
+            },
+            "optional": {"tags": ("STRING", {"default": ""})},
+        }
+
+    RETURN_TYPES = ("STRING",)
+    CATEGORY = "promptcontrol/tools"
+    FUNCTION = "apply"
+    DESCRIPTION = "Parses the input prompt and returns the prompt scheduled at the specified point"
+
+    def apply(self, text, at, tags=""):
+        schedule = parse_prompt_schedules(text, filters=tags)
+        _, entry = schedule.at_step(at, total_steps=1)
+        prompt_text = entry.get("prompt", "")
+        return (prompt_text,)
+
+
 NODE_CLASS_MAPPINGS = {
     "PCSetPCTextEncodeSettings": PCSetPCTextEncodeSettings,
     "PCAddMaskToCLIP": PCAddMaskToCLIP,
     "PCAddMaskToCLIPMany": PCAddMaskToCLIPMany,
+    "PCSetLogLevel": PCSetLogLevel,
+    "PCExtractScheduledPrompt": PCExtractScheduledPrompt,
+    "PCSaveExpandedWorkflow": PCSaveExpandedWorkflow,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "PCSetTextEncodeSettings": "PC: Configure PCTextEncode",
+    "PCSetPCTextEncodeSettings": "PC: Configure PCTextEncode",
     "PCAddMaskToCLIP": "PC: Attach Mask",
     "PCAddMaskToCLIPMany": "PC: Attach Mask (multi)",
+    "PCSetLogLevel": "PC: Configure Logging (for debug)",
+    "PCExtractScheduledPrompt": "PC: Extract Scheduled Prompt",
+    "PCSaveExpandedWorkflow": "PC: Save Expanded Workflow (for debug)",
 }

prompt_control/parser.py

@@ -1,12 +1,25 @@
+# vim: sw=4 ts=4
 import lark
 import logging
 from math import ceil
 
 logging.basicConfig()
 log = logging.getLogger("comfyui-prompt-control")
+import re
+
+from functools import lru_cache
+from .utils import get_function
 
 if lark.__version__ == "0.12.0":
-    x = "Your lark package reports an ancient version (0.12.0) and will not work. If you have the 'lark-parser' package in your Python environment, remove that and *reinstall* lark!"
+    from sys import executable
+
+    x = "\n".join(
+        [
+            "Your lark package reports an ancient version (0.12.0) and will not work. If you have the 'lark-parser' package in your Python environment, remove that and *reinstall* lark!",
+            f"{executable} -m pip uninstall lark-parser lark",
+            f"{executable} -m pip install lark",
+        ]
+    )
     log.error(x)
     raise ImportError(x)
 
@@ -18,9 +31,9 @@ prompt: (emphasized | embedding | scheduled | alternate | sequence | loraspec |
 !emphasized: "(" prompt? ")"
         | "(" prompt ":" prompt ")"
         | "[" prompt "]"
-scheduled: "[" [prompt ":"] [prompt] ":" _WS? NUMBER ["," NUMBER] "]"
-        | "[" [prompt ":"] [prompt] ":" _WS? TAG "]"
-sequence:  "[SEQ" ":" [prompt] ":" NUMBER (":" [prompt] ":" NUMBER)+ "]"
+scheduled: "[" [[prompt] ":"] [prompt] ":" _WS? NUMBER ["," NUMBER] "]"
+        | "[" [[prompt] ":"] [prompt] ":" _WS? TAG "]"
+sequence.5:  "[SEQ" ":" [prompt] ":" NUMBER (":" [prompt] ":" NUMBER)* "]"
 alternate: "[" [prompt] ("|" [prompt])+ [":" NUMBER] "]"
 loraspec.99: "<lora:" FILENAME lora_weights [lora_block_weights] ">"
 lora_weights.1: (":" _WS? NUMBER)~1..2
@@ -36,6 +49,7 @@ TAG: /[A-Z_]+/
     lexer="dynamic",
 )
 
+
 cut_parser = lark.Lark(
     r"""
 !start: (prompt | /[][:()]/+)*
@@ -295,7 +309,15 @@ class PromptSchedule(object):
         if len(res) == 0:
             res = [[1.0, parsed[-1][1]]]
 
-        return res
+        final = [res[0]]
+
+        # Clean up duplicates
+        for p in res[1:]:
+            if p[1] != final[-1][1]:
+                final.append(p)
+            else:
+                final[-1][0] = p[0]
+        return final
 
     def clone(self):
         return self.with_filters()
@@ -323,5 +345,52 @@ class PromptSchedule(object):
         return len(self.parsed_prompt) - 1, self.parsed_prompt[-1]
 
 
-def parse_prompt_schedules(prompt):
-    return PromptSchedule(prompt)
+def replace_def(text):
+    text, defs = get_function(text, "DEF", defaults=None)
+    res = text
+    prevres = text
+    replacements = []
+    for d in defs:
+        r = d.split("=", 1)
+        if len(r) != 2 or not r[0].strip():
+            log.warning("Ignoring invalid DEF(%s)", d)
+            continue
+        replacements.append((r[0].strip(), r[1].strip()))
+    iterations = 0
+    while True:
+        iterations += 1
+        if iterations > 10:
+            raise ValueError("Unable to resolve DEFs, make sure there are no cycles!")
+            return text
+        for search, replace in replacements:
+            res = substitute_defcall(res, search, replace)
+            res = substitute_def(res, search, replace)
+        if res == prevres:
+            break
+        prevres = res
+    if res.strip() != text.strip():
+        res = res.strip()
+        log.info("DEFs expanded to: %s", res)
+    return res
+
+
+def substitute_def(text, search, 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(";")]
+        r = replace
+        for i, v in enumerate(paramvals):
+            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)
+    return PromptSchedule(prompt, **kwargs)

prompt_control/prompts.py

@@ -4,11 +4,26 @@ import torch
 from functools import partial
 from comfy_extras.nodes_mask import FeatherMask, MaskComposite
 
-from .utils import safe_float, get_function, parse_floats
+from .utils import safe_float, get_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
+
+
 log = logging.getLogger("comfyui-prompt-control")
 
 AVAILABLE_STYLES = ["comfy", "perp", "A1111", "compel", "comfy++", "down_weight"]
@@ -88,8 +103,9 @@ def shuffle_chunk(shuffle, c):
         "separator": separator,
     }.get(joiner, joiner)
 
-    log.info("%s arg=%s sep=%s join=%s", func, shuffle_count, separator, joiner)
-    separated = c.split(separator)
+    log.debug("%s arg=%s sep=%s join=%s", func, shuffle_count, separator, joiner)
+    separated = smarter_split(separator, c)
+    log.debug("Prompt split into %s", separated)
     if func == "SHIFT":
         shuffle_count = shuffle_count % len(separated)
         permutation = separated[shuffle_count:] + separated[:shuffle_count]
@@ -429,6 +445,11 @@ def encode_prompt(clip, text, start_pct, end_pct, defaults, masks):
     # TODO: is this still needed?
     # scale = sum(abs(weight(p)[0]) for p in prompts if not ("AREA(" in p or "MASK(" in p))
     for prompt in prompts:
+        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)
@@ -451,6 +472,11 @@ def encode_prompt(clip, text, start_pct, end_pct, defaults, masks):
         settings["start_percent"] = start_pct
         settings["end_percent"] = end_pct
         x = encode_prompt_segment(clip, prompt, settings, style, normalization)
+        if attn and mask is not None:
+            mask = settings.pop("mask")
+            strength = settings.pop("mask_strength")
+            x = set_cond_attnmask(x, mask * strength)
+
         conds.extend(x)
 
     return conds

prompt_control/test_graph.py

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

prompt_control/test_parser.py

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

prompt_control/utils.py

@@ -2,7 +2,14 @@ from pathlib import Path
 import re
 import logging
 
-import folder_paths
+# Allow testing
+try:
+    from folder_paths import get_filename_list
+except ImportError:
+
+    def get_filename_list(x):
+        raise NotImplementedError("How did you get here?")
+
 
 log = logging.getLogger("comfyui-prompt-control")
 
@@ -35,7 +42,6 @@ def find_nonscheduled_loras(consolidated_schedule):
     if not consolidated_schedule:
         return {}
     last_end, candidate_loras = consolidated_schedule[0]
-    print(candidate_loras)
     to_remove = set()
     for candidate, weights in candidate_loras.items():
         for end, loras in consolidated_schedule[1:]:
@@ -48,6 +54,26 @@ def find_nonscheduled_loras(consolidated_schedule):
     return {k: v for (k, v) in candidate_loras.items() if k not in to_remove}
 
 
+def smarter_split(separator, string):
+    """Does not break () when splitting"""
+    splits = []
+    prev = 0
+    stack = 0
+    escape = False
+    for idx, x in enumerate(string):
+        if x == "(" and not escape:
+            stack += 1
+        elif x == ")" and not escape:
+            stack = max(0, stack - 1)
+        elif x == separator and stack == 0:
+            splits.append(string[prev:idx])
+            prev = idx + 1
+        escape = x == "\\"
+
+    splits.append(string[prev : idx + 1])
+    return splits
+
+
 def find_closing_paren(text, start):
     stack = 1
     for i, char in enumerate(text[start:]):
@@ -61,10 +87,11 @@ def find_closing_paren(text, start):
     return len(text)
 
 
-def get_function(text, func, defaults, return_func_name=False):
+def get_function(text, func, defaults, return_func_name=False, placeholder=""):
     rex = re.compile(rf"\b{func}\(", re.MULTILINE)
     instances = []
     match = rex.search(text)
+    count = 0
     while match:
         # Match start, content start
         start, after_first_paren = match.span()
@@ -76,8 +103,12 @@ def get_function(text, func, defaults, return_func_name=False):
         else:
             instances.append(args)
 
-        text = text[:start] + text[end + 1 :]
+        if placeholder:
+            text = text[:start] + f"\0{placeholder}{count}\0" + text[end + 1 :]
+        else:
+            text = text[:start] + text[end + 1 :]
         match = rex.search(text)
+        count += 1
     return text, instances
 
 
@@ -119,7 +150,7 @@ def safe_float(f, default):
 
 
 def lora_name_to_file(name):
-    filenames = folder_paths.get_filename_list("loras")
+    filenames = get_filename_list("loras")
     # Return exact matches as is
     if name in filenames:
         return name

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-prompt-control"
 description = "Nodes for convenient prompt editing, making many common operations prompt-controllable"
-version = "2.0.0-beta.2"
+version = "2.0.0-beta.9"
 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"]