v3.0.0-beta.9

See: #150

README.md

@@ -11,6 +11,9 @@ A `Basic Text to Image` template is included with the extension, and can be load
 > The parser was rewritten using parsy. It is intended to have the same behaviour as the old parser, but is **significantly** faster.
 > Please report any bugs or incompatibilities you find.
 
+## Notable changes
+
+- `PC: Schedule Prompt` now strips surrounding whitespace by default, which may change some prompts. Add `NOSTRIP()` to your prompt to restore previous behaviour.
 
 ## What can it do?
 
@@ -23,6 +26,7 @@ A `Basic Text to Image` template is included with the extension, and can be load
   - Different weight interpretation types (ComfyUI, A1111, compel, etc.)
   - Prompt masking with an implementation of [cutoff](https://github.com/BlenderNeko/ComfyUI_Cutoff).
 - Organize complicated prompts with [segments and prompt macros](/doc/macros.md).
+- [Schedule your own encoder nodes](/doc/node_function.md), allowing prompt control of eg. video or audio models with non-text inputs.
 
 All features are fully schedulable unless otherwise stated. See the [scheduling syntax documentation](doc/schedules.md) to get started.
 
@@ -52,10 +56,6 @@ If you run into problems, update ComfyUI first.
 
 for example, if you first encode `[cat:dog:0.1]` and later change that to `[cat:dog:0.5]`, no re-encoding takes place.
 
-for added fun, put `NODE(NodeClassName, textinputname)` in a prompt to generate a graph using **any other node** that's compatible. The node can't have required parameters besides a single CLIP parameter (which must be named `clip`) and the text prompt, and it must return a `CONDITIONING` as its first return value. The "default" values are `PCTextEncode` and `text`.
-
-For example, if you for some reason do not want the advanced features of `PCTextEncode`, use `NODE(CLIPTextEncode)` in the prompt and you'll still get scheduling with ComfyUI's regular TE node.
-
 The advanced node enables filtering the prompt for multi-pass workflows.
 
 ## PCLazyLoraLoader and PCLazyLoraLoaderAdvanced

doc/attention_couple.md

@@ -24,6 +24,13 @@ See also the [regional prompting documentation](/doc/regional_prompts.md) for in
 
 You can use `COUPLE` to attach attention-coupled prompts to a base prompt:
 
+For example:
+```
+dog FILL() COUPLE(0.5 1) cat
+```
+
+The full syntax looks as follows (to use `IMASK` you need to attach a custom mask)
+
 `base_prompt COUPLE MASK(0 0.5) coupled prompt 1 with mask COUPLE IMASK(0) coupled prompt 2 with custom mask`
 
 as a shortcut, `COUPLE(maskparams)` is expanded to `COUPLE MASK(maskparams)`, so the above prompt can also be written as:
@@ -31,15 +38,17 @@ as a shortcut, `COUPLE(maskparams)` is expanded to `COUPLE MASK(maskparams)`, so
 `base_prompt COUPLE(0 0.5) coupled prompt 1 with mask COUPLE IMASK(0) coupled prompt 2 with custom mask`
 
 Behaviour:
-- If no mask is specified, an implicit `MASK()` is assumed.
+- If no mask is specified, an implicit `MASK()` is assumed, meaning that the prompt affects the entire image.
 
-- For the base prompt, you can also use `FILL()` to automatically mask all parts not masked by coupled prompts
+- For the base prompt, you can use `FILL()` to automatically mask all parts not masked by other coupled prompts
 
-- If the base prompt has weight set to zero (ie. ´:0` at the end), then the first coupled prompt with non-zero weight becomes the base prompt.
+- If the base prompt has weight set to zero (ie. ´:0` at the end), then the first coupled prompt with non-zero weight becomes the base prompt:
 
-For example:
 ```
-dog FILL() COUPLE(0.5 1) cat
+disabled prompt :0 COUPLE new base prompt COUPLE coupled prompt
 ```
 
-Note that because the generation still sees and diffuses the full latent, attention coupling is not guaranteed to perfectly limit the effect of your prompt to the masked area.
+You can also schedule the weight normally: `prompt :[1:0:0.35]`
+
+> ![NOTE]
+> Note that because the generation still sees and diffuses the full latent, attention coupling is not guaranteed to perfectly limit the effect of your prompt to the masked area.

doc/macros.md

@@ -86,7 +86,9 @@ Top panel: a sleeping cat. The cat has orange fur with white stripes
 Bottom panel: A dog chasing its
   tail in a living room.
 ```
-Unlike macros, SEG is processed *after* scheduling syntax.
+
+> [!NOTE]
+> Unlike macros, SEGs are processed *after* scheduling syntax has been expanded, except in the LoRA loader (this may change later, but requires a bit of refactoring)
 
 In this case, the first section before any `SEG` becomes the *template* and any text after a `SEG` call becomes part of that segment. Whitespace is stripped from the start and end of segments and the template.
 

doc/node_function.md

@@ -0,0 +1,41 @@
+# The NODE function
+
+The `NODE` function allows you to use any other text encoding node within `PC: Schedule Prompt`, replacing the default `PCTextEncode` and allowing for example video model scheduling.
+
+> [!NOTE]
+> When using NODE, you lose access to *all* special syntax provided by `PCTextEncode`. Only SEGs, macros and scheduling will continue to work since those are processed at graph expansion time before the text prompt is passed into the node.
+
+## Basic usage
+
+Use `NODE(NodeClassName, textinputname)` in a prompt to generate a graph using any node that's compatible. The requirements are as follows:
+- The node must have a CLIP parameter (which must be named `clip`)
+- It must have a text field
+- It must return a `CONDITIONING` as its first return value.
+
+For example, if you for some reason do not want the advanced features of `PCTextEncode`, use `NODE(CLIPTextEncode)` in the prompt and you'll still get scheduling with ComfyUI's regular TE node.
+
+The default parameters are `PCTextEncode` and `text`.
+
+## Advanced Usage with arbitrary parameters
+
+Advanced usage of `NODE` can be complicated. For an example, see [The H3 workflow](/example_workflows/Prompt%20Control%20with%20MiniMax%20H3.json?raw=1). You can also find it in the template library.
+
+The full synopsis of the function is `NODE(NodeClassName, textinputname, arg_spec)` where `arg_spec` is a semicolon-separated list of `parameter_name json_value` pairs. In raw form, it looks like this:
+```
+NODE(MiniMaxH3ImageToVideo, prompt, vae ["1", 0]; width 1024; height 1024; first_frame ["2", 0])
+```
+
+The names and inputs must match the ComfyUI API format which **may differ from frontend names**. You can export your workflow in API format and inspect it to see how inputs are passed in to nodes.
+
+The arrays are literal ComfyUI node links, meaning the `vae` parameter is taken from node ID "1" first output and `first_frame` from node ID "2" first output.
+
+The values are arbitrary JSON literals, meaning that you can also pass in constant values. To pass in literal strings for example, you need to use quotes `"like this"`.
+
+This is intended to be used with the helper node `PC: NODE Input Helper`, which can be used to pass arbitrary parameters (named `$a` to `$n`) to the encoder. The recommended pattern is to put something like the following:
+```
+SEG(node)
+NODE(MiniMaxH3ImageToVideo, prompt, vae $a; width $b; height $c; length $d; first_frame $e)
+```
+to the helper and then concatenate it at the end of your prompt (use whitespace as a separator). You can then trigger the node with `$node` in your prompt. (see [documentation](/doc/macros.md) for `SEG`)
+
+The helper will replace the parameters with the correct ComfyUI link values.

doc/schedules.md

@@ -14,6 +14,9 @@ Besides the syntax documented below, the [basic syntax](/doc/basic.md) and [prom
 a [large::0.1] [cat|dog:0.05] [<lora:somelora:0.5:0.6>::0.5]
 [in a park:in space:0.4]
 ```
+## Note on whitespace
+`PC: Schedule Prompt` will strip leading and following whitespace from the prompt automatically. If you really want whitespace in your prompt, include `NOSTRIP()` in your prompt.
+
 ## Comments and escaping
 
 In schedules, any text on a line following a `#` is considered a comment and removed, including the `#` character.

prompt_control/anima_couple.py

@@ -1,11 +1,13 @@
 # Adapted from https://github.com/pamparamm/ComfyUI-ppm
 import itertools
 from collections.abc import Callable
+from functools import partial
 from math import lcm
 
 import torch
 import torch.nn.functional as F
 from comfy.ldm.anima.model import Anima as AnimaDIT
+from comfy.ldm.cosmos.predict2 import Attention as CosmosAttention
 from comfy.patcher_extension import WrapperExecutor
 from comfy.sampler_helpers import convert_cond
 from comfy.samplers import process_conds
@@ -23,6 +25,23 @@ def reshape_mask(mask: torch.Tensor, size: tuple[int, int], bs: int, num_tokens:
     return mask_downsample_reshaped
 
 
+def wrap_forwards(anima_model):
+    backups = {}
+    for block_name, b in (
+        (n, b) for n, b in anima_model.named_modules() if "cross_attn" in n and isinstance(b, CosmosAttention)
+    ):
+        backups[block_name] = b.forward
+        b.forward = partial(cosmos_attention_forward_couple, b.forward)
+    return backups
+
+
+def unwrap_forwards(anima_model, backups):
+    for block_name, b in (
+        (n, b) for n, b in anima_model.named_modules() if "cross_attn" in n and isinstance(b, CosmosAttention)
+    ):
+        b.forward = backups[block_name]
+
+
 def anima_sample_wrapper(executor, *args, **kwargs):
     guider, _, extra_options, _, noise, latent_image, denoise_mask, *_ = args
     seed = extra_options["seed"]
@@ -67,7 +86,13 @@ def anima_forward_wrapper(executor: WrapperExecutor, *args, **kwargs):
     transformer_options["activations_shape"] = activations_shape
     kwargs["transformer_options"] = transformer_options
 
-    return executor(*args, **kwargs)
+    b = {}
+    if pc:
+        b = wrap_forwards(anima_model)
+    r = executor(*args, **kwargs)
+    if pc:
+        unwrap_forwards(anima_model, b)
+    return r
 
 
 def cosmos_attention_forward_couple(_forward: Callable, x, context, rope_emb, transformer_options):

prompt_control/macros.py

@@ -88,7 +88,6 @@ def expand_macros(text):
         iterations += 1
         if iterations > 10:
             raise ValueError("Unable to resolve DEFs, make sure there are no cycles!")
-            return text
         for search, replace in replacements:
             res = substitute_defcall(res, search, replace)
         if res == prevres:
@@ -96,7 +95,7 @@ def expand_macros(text):
         prevres = res
     if res.strip() != text.strip():
         res = res.strip()
-        log.info("DEFs expanded to: %s", res)
+        log.debug("DEFs expanded to: %s", res)
     return res
 
 

prompt_control/nodes_anima.py

@@ -3,7 +3,6 @@
 
 import comfy.model_management
 import comfy.patcher_extension
-from comfy.ldm.cosmos.predict2 import Attention as CosmosAttention
 from comfy.model_base import Anima
 from comfy.model_patcher import ModelPatcher
 from comfy_api.latest import io
@@ -11,20 +10,9 @@ from comfy_api.latest import io
 from .anima_couple import (
     anima_forward_wrapper,
     anima_sample_wrapper,
-    cosmos_attention_forward_couple,
 )
 
 
-class CoupleForward:
-    def __init__(self, fn, block):
-        self.fn = fn
-        self.block = block
-
-    def __call__(self, *args, **kwargs):
-        self.block.to("cuda")
-        return cosmos_attention_forward_couple(self.fn, *args, **kwargs)
-
-
 class PCAnimaAttnCouplePatch(io.ComfyNode):
     @classmethod
     def define_schema(cls) -> io.Schema:
@@ -47,7 +35,6 @@ class PCAnimaAttnCouplePatch(io.ComfyNode):
 
         if issubclass(model_type, Anima):
             m = model.clone()
-            anima_model = model.get_model_object("diffusion_model")
             m.add_wrapper_with_key(
                 comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL,
                 cls.__name__,
@@ -59,12 +46,6 @@ class PCAnimaAttnCouplePatch(io.ComfyNode):
                 anima_sample_wrapper,
             )
 
-            for block_name, b in (
-                (n, b) for n, b in anima_model.named_modules() if "cross_attn" in n and isinstance(b, CosmosAttention)
-            ):
-                attn_forward_prev = m.get_model_object(f"diffusion_model.{block_name}.forward")
-                m.add_object_patch(f"diffusion_model.{block_name}.forward", CoupleForward(attn_forward_prev, b))
-
         return io.NodeOutput(m)
 
 

prompt_control/nodes_lazy.py

@@ -8,6 +8,7 @@ from comfy_api.latest import io
 from comfy_execution.graph import ExecutionBlocker
 from comfy_execution.graph_utils import GraphBuilder
 
+from .macros import expand_macros, expand_segs
 from .parser import parse_prompt_schedules
 from .utils import consolidate_schedule, find_nonscheduled_loras, get_function
 
@@ -158,6 +159,7 @@ class PCLazyLoraLoaderAdvanced(io.ComfyNode):
 
     @classmethod
     def execute(cls, model=None, clip=None, text="", apply_hooks=True, tags="", start=0.0, end=1.0, num_steps=0):
+        text = expand_segs(expand_macros(text))
         schedule = parse_prompt_schedules(text, filters=tags, start=start, end=end, num_steps=num_steps)
         graph = GraphBuilder()
         r = build_lora_schedule(graph, schedule, model, clip, apply_hooks=apply_hooks)
@@ -195,14 +197,42 @@ def build_scheduled_prompts(graph, schedules, clip):
     start_pct = 0.0
     for end_pct, c in schedules:
         p = c["prompt"]
-        p, classnames = get_function(p, "NODE", ["PCTextEncode", "text"])
-        classname = "PCTextEncode"
-        paramname = "text"
+        # Need to explicitly expand SEGs here *before* NODE is processed
+        p = expand_segs(p)
+        p, classnames = get_function(p, "NODE", defaults=None)
+        realargs = ["PCTextEncode", "text", ""]
+        if len(classnames) > 1:
+            log.warning("You have more than one NODE call in your prompt. Only the first one will be used")
         if classnames:
-            classname, paramname = classnames[0].args
-        node = graph.node(classname)
+            args = classnames[0].args[0]
+            if not args.strip():
+                raise ValueError("NODE can't be empty!")
+            for i, v in enumerate(args.split(",", maxsplit=2)):
+                realargs[i] = v
+        classname, paramname, magic_spec = realargs
+        node = graph.node(classname.strip())
         node.set_input("clip", clip)
-        node.set_input(paramname, p)
+        # We should strip extra whitespace so that people don't have to worry about functions.
+        p = p.replace("NOSTRIP()", "") if "NOSTRIP()" in p else p.strip()
+        node.set_input(paramname.strip(), p)
+        magic_spec.replace(r"\;", "__ESCAPED_SEMICOLON__")
+        extra_inputs = magic_spec.split(";") if magic_spec.strip() else []
+        for e in extra_inputs:
+            e = e.strip()
+            if not e:
+                continue
+            e = e.replace("__ESCAPED_SEMICOLON__", ";")
+            name, jsondata = e.split(maxsplit=1)
+            jsondata = jsondata.strip()
+            if not jsondata.strip():
+                continue
+            # From helper node:
+            if jsondata == "__EMPTY__":
+                continue
+            try:
+                node.set_input(name.strip(), json.loads(jsondata.strip()))
+            except ValueError as e:
+                raise ValueError(f"Invalid JSON input: '{jsondata}'") from e
         timestep = graph.node("ConditioningSetTimestepRange")
         timestep.set_input("conditioning", node.out(0))
         timestep.set_input("start", start_pct)

prompt_control/nodes_tools.py

@@ -1,3 +1,4 @@
+import json
 import logging
 
 from comfy_api.latest import io
@@ -5,6 +6,7 @@ from comfy_api.latest import io
 from .macros import expand_macros as macroexpand
 from .macros import expand_segs as segexpand
 from .macros import expand_subs as subexpand
+from .macros import substitute_var
 from .parser import parse_prompt_schedules
 
 log = logging.getLogger("comfyui-prompt-control")
@@ -182,6 +184,68 @@ class PCMacroExpand(io.ComfyNode):
         return io.NodeOutput(macroexpand(text))
 
 
+class PCLinkHelper(io.ComfyNode):
+    # a-z
+    NAMES = [chr(97 + i) for i in range(26)]
+
+    @classmethod
+    def define_schema(cls):
+        t1 = io.Autogrow.TemplateNames(io.AnyType.Input("link", raw_link=True), min=0, names=cls.NAMES)
+        t2 = io.Autogrow.TemplateNames(
+            io.AnyType.Input("value", lazy=True), min=0, names=[f"var{i + 1}" for i in range(50)]
+        )
+        return io.Schema(
+            node_id="PCNODELinkHelper",
+            display_name="PC: Extra argument helper for NODE",
+            category="promptcontrol/tools",
+            description="Takes in arbitrary inputs and renders them as NODE-compatible values, replacing $a -> $z with JSON link values.",
+            is_experimental=True,
+            inputs=[
+                io.Autogrow.Input("links", template=t1),
+                io.Autogrow.Input(
+                    "vars",
+                    template=t2,
+                ),
+                io.String.Input(
+                    "template",
+                    tooltip="The variables $a to $z will be replaced in this text with their corresponding input's JSON link value",
+                    placeholder="In this text you can refer to the input links as $a, $b etc. and the var inputs as either $var1 or $json1 etc. (the latter will be rendered through Python's json.dumps function which will cause strings to be quoted)",
+                    multiline=True,
+                ),
+            ],
+            outputs=[io.String.Output()],
+        )
+
+    # This requires https://github.com/Comfy-Org/ComfyUI/pull/15103 to work properly
+    @classmethod
+    def check_lazy_status(cls, template, links, vars):
+        r = []
+        for name, (v, input_name) in vars.items():
+            if v is None and f"${name}" in template or v is None and f"$json{name[3:]}" in template:
+                r.append(input_name)
+        return r
+
+    @classmethod
+    def execute(cls, template, links, vars) -> io.NodeOutput:
+        text = template
+        for k in cls.NAMES:
+            v = "__EMPTY__"
+            if k in links:
+                v = json.dumps(links[k])
+            text = substitute_var(text, k, v)
+        for i in range(50):
+            v = "__EMPTY__"
+            k = f"var{i + 1}"
+            if k in vars:
+                v = vars[k]
+            text = substitute_var(text, k, str(v))
+            if f"$json{i + 1}" in text:
+                v = v if v == "__EMPTY__" else json.dumps(v)
+                text = substitute_var(text, f"json{i + 1}", v)
+
+        return io.NodeOutput(text)
+
+
 NODES = [
     PCSetPCTextEncodeSettings,
     PCAddMaskToCLIP,
@@ -189,4 +253,5 @@ NODES = [
     PCSetLogLevel,
     PCExtractScheduledPrompt,
     PCMacroExpand,
+    PCLinkHelper,
 ]

prompt_control/parser_parsy.py

@@ -370,7 +370,17 @@ ctlweight = seq(number, (string("@") >> number).optional(0)).sep_by(comma, min=1
 loractl = (string("<loractl:") >> filename * 1 + lora_weights(ctlweight) << string(">")).combine(LoRACTL)
 emb = (string("<emb:") >> filename << string(">")).map(lambda f: Text(f"embedding:{f}"))
 
-expr = escape | comment | non_special | bracketed | emphasis.combine(combine_prompt) | lora | loractl | emb
+expr = (
+    escape
+    | comment
+    | non_special
+    | bracketed
+    | emphasis.combine(combine_prompt)
+    | lora
+    | loractl
+    | emb
+    | char_from("<>").map(Text)
+)
 prompt_ = expr.at_least(1).combine(combine_prompt)
 prompt.become(prompt_)
 # Treat any character that isn't valid prompt syntax as just text

pyproject.toml

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

tests/test_graph.py

@@ -460,6 +460,102 @@ def test_textencode_lora_with_schedule():
     }
 
 
+def test_textencode_custom():
+    r = te("NODE(CLIPTextEncode)simple [test:0.1,0.5] $p SEG(p) prompt")
+    assert r == {
+        "result": (["UID.0.0.8", 0],),
+        "expand": {
+            "UID.0.0.1": {
+                "class_type": "CLIPTextEncode",
+                "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": "CLIPTextEncode",
+                "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": "CLIPTextEncode",
+                "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]},
+            },
+        },
+    }
+
+
+def test_textencode_custom_extra():
+    r = te(
+        'NODE(CustomTextEncode, prompt, image ["1", 0]; option "test"; float [10.0:__EMPTY__:0.5])simple [test:0.1,0.5] prompt'
+    )
+    assert r == {
+        "result": (["UID.0.0.8", 0],),
+        "expand": {
+            "UID.0.0.1": {
+                "class_type": "CustomTextEncode",
+                "inputs": {
+                    "clip": [0, 0],
+                    "prompt": "simple  prompt",
+                    "image": ["1", 0],
+                    "option": "test",
+                    "float": 10.0,
+                },
+            },
+            "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": "CustomTextEncode",
+                "inputs": {
+                    "clip": [0, 0],
+                    "prompt": "simple test prompt",
+                    "image": ["1", 0],
+                    "option": "test",
+                    "float": 10.0,
+                },
+            },
+            "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": "CustomTextEncode",
+                "inputs": {"clip": [0, 0], "prompt": "simple  prompt", "image": ["1", 0], "option": "test"},
+            },
+            "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]},
+            },
+        },
+    }
+
+
 def test_loraloader_empty(monkeypatch, caplog):
     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"]
@@ -582,3 +678,9 @@ def test_loraloader_adv_start():
 def test_loraloader_end_zero():
     result2 = loraloader("prompt [<lora:test:0.5>:0.5]", adv=True, end=0.5)["expand"]
     assert result2 == {}
+
+
+def test_loraloader_segs():
+    result = loraloader("prompt [<lora:test:0.5>:0.5]")["expand"]
+    result2 = loraloader("prompt [$lora:0.5]\nSEG(lora)<lora:test:0.5>\nSEG(lora2)<lora:ignored:1>")["expand"]
+    assert result == result2

tests/test_parser.py

@@ -342,6 +342,12 @@ def test_cornercase_corrected(parse):
     assert p.parsed_prompt[1:] == p2.parsed_prompt
 
 
+def test_ltgt_in_schedule(parse):
+    p = parse("This should [<parse> correctly:be <Picture 1>:0.1]<lora:test:1>")
+    assert_prompt(p, 0.1, 0.1, "This should <parse> correctly", ("test", 1.0, 1.0))
+    assert_prompt(p, 0.15, 1.0, "This should be <Picture 1>", ("test", 1.0, 1.0))
+
+
 def test_floats(parse):
     p = parse("[a:b:0.5] [c:d:e:0.2,0.7] <lora:test:-0.3>")
     p2 = parse("[a:b:.5] [c:d:e:.2,.7] <lora:test:-.3>")