v3.0.0-beta.6

This avoids using stale block references and seems to fix LoRAs

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 follow (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.
 

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/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)

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.6"
 license = { file = "LICENSE" }
 
 requires-python = ">= 3.10"