v2.0.0-beta.5

Experimental attention masking

__init__.py

@@ -30,14 +30,21 @@ 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

@@ -195,3 +195,15 @@ 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 break without warning.
+
+## 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.

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_lazy.py

@@ -12,7 +12,8 @@ import json
 
 def _cache_key(cachekey, inputs):
     out = inputs.copy()
-    if not is_link(inputs["text"]):
+    text = inputs.get("text")
+    if text is not None and not is_link(text):
         out["text"] = cache_key_from_inputs(cachekey, **inputs)
     return out
 

prompt_control/nodes_tools.py

@@ -1,4 +1,5 @@
 import logging
+from .parser import parse_prompt_schedules
 
 log = logging.getLogger("comfyui-prompt-control")
 
@@ -42,6 +43,7 @@ class PCAddMaskToCLIP:
     RETURN_TYPES = ("CLIP",)
     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)
@@ -63,6 +65,7 @@ class PCAddMaskToCLIPMany:
     RETURN_TYPES = ("CLIP",)
     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()
@@ -93,6 +96,7 @@ class PCSetPCTextEncodeSettings:
     RETURN_TYPES = ("CLIP",)
     CATEGORY = "promptcontrol/tools"
     FUNCTION = "apply"
+    DESCRIPTION = "Configures default values for PCTextEncode"
 
     def apply(
         self,
@@ -123,11 +127,35 @@ 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,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -135,4 +163,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "PCAddMaskToCLIP": "PC: Attach Mask",
     "PCAddMaskToCLIPMany": "PC: Attach Mask (multi)",
     "PCSetLogLevel": "PC: Configure Logging (for debug)",
+    "PCExtractScheduledPrompt": "PC: Extract Scheduled Prompt",
 }

prompt_control/parser.py

@@ -4,8 +4,10 @@ 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":
     from sys import executable
@@ -46,6 +48,7 @@ TAG: /[A-Z_]+/
     lexer="dynamic",
 )
 
+
 cut_parser = lark.Lark(
     r"""
 !start: (prompt | /[][:()]/+)*
@@ -333,6 +336,34 @@ class PromptSchedule(object):
         return len(self.parsed_prompt) - 1, self.parsed_prompt[-1]
 
 
+def replace_defs(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:
+            log.error("Unable to resolve DEFs, make sure there are no cycles!")
+            return text
+        for search, replace in replacements:
+            res = re.sub(rf"\b{re.escape(search)}\b", replace, res)
+        if res == prevres:
+            break
+        prevres = res
+    if res != text:
+        log.info("DEFs expanded to: %s", res)
+    return res
+
+
 @lru_cache
 def parse_prompt_schedules(prompt, **kwargs):
+    prompt = replace_defs(prompt)
     return PromptSchedule(prompt, **kwargs)

prompt_control/prompts.py

@@ -9,6 +9,21 @@ 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"]
@@ -430,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)
@@ -452,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/utils.py

@@ -3,8 +3,13 @@ import re
 import logging
 
 # Allow testing
-if __name__ != "__main__":
-    import folder_paths
+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")
 
@@ -140,7 +145,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.4"
+version = "2.0.0-beta.5"
 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"]