From 6d7df7d5cecf7113e5020e9caebbba06c70ac8e4 Mon Sep 17 00:00:00 2001 From: BlackVortexAI Date: Wed, 19 Aug 2026 15:47:23 +0200 Subject: [PATCH] Add regional LoRA hooks to Anima attention --- README.md | 21 +++-- py/nodes/bv_regional.py | 22 ++++- py/util/regional/anima_adapter.py | 20 +++-- py/util/regional/clip_hooks.py | 19 +++++ py/util/regional/lora_hooks.py | 90 ++++++++++++++++++++ py/util/regional/native_conditioning.py | 19 +---- pyproject.toml | 2 +- tests/test_regional_anima_adapter.py | 14 ++++ tests/test_regional_lora_hooks.py | 104 +++++++++++++++++++++++- 9 files changed, 277 insertions(+), 34 deletions(-) create mode 100644 py/util/regional/clip_hooks.py diff --git a/README.md b/README.md index 296e615..0a262ee 100644 --- a/README.md +++ b/README.md @@ -107,12 +107,21 @@ the spatial guidance. 5. Connect `patched_model`/`positive`/`negative` to a standard KSampler. 6. Optionally return the latest image through **Preview Send** or **Save Send**. -Optional regional LoRA hooks are available on the native-conditioning path. Connect -an external `LORA_STACK` producer to **BV Named LoRA Stack**, chain its registry into -**BV Regional Native Conditioning**, and connect the prompt node's `lora_bindings` -sidecar output. The editor can then assign one unchanged live stack globally and one -additional stack per region. The registry and bindings inputs are optional; workflows -that do not connect them retain the previous `BV_REGIONAL` v1 behavior. +Optional regional LoRA hooks are available on the native-conditioning and Anima +attention paths. Connect an external `LORA_STACK` producer to **BV Named LoRA Stack**, +chain its registry into the conditioning node, and connect the prompt node's +`lora_bindings` sidecar output. The editor can then assign one unchanged live stack +globally and one additional stack per region. The registry and bindings inputs are +optional; workflows that do not connect them retain the previous `BV_REGIONAL` v1 +behavior. Empty stacks and entries whose model and CLIP strengths are both zero are +valid no-ops. + +Anima attention evaluates one full model pass for the global baseline plus one masked +pass for every distinct effective regional model stack. Regions with identical model +stacks share a pass, and CLIP-only strength differences do not add model passes. A +single regional stack therefore normally doubles sampling work, while two different +regional stacks normally triple it. Actual runtime and peak VRAM depend on the model, +resolution, sampler and ComfyUI memory management. ### Named LoRA stack interoperability diff --git a/py/nodes/bv_regional.py b/py/nodes/bv_regional.py index 85ed4ac..197e68c 100644 --- a/py/nodes/bv_regional.py +++ b/py/nodes/bv_regional.py @@ -21,6 +21,7 @@ from ..util.regional.lora_hooks import ( BINDINGS, REGISTRY, add_named_stack, + apply_attention_hook_passes, create_hook_groups, default_bindings, reconcile_bindings, @@ -472,14 +473,21 @@ class BVRegionalAnimaConditioningNode: "base_ratio": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}), "cross_inject_every_n_blocks": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}), "self_inject_every_n_blocks": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}), - } + }, + "optional": { + "lora_registry": (REGISTRY, {}), + "lora_bindings": (BINDINGS, {}), + }, } RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING") RETURN_NAMES = ("patched_model", "positive", "negative") FUNCTION = "apply" CATEGORY = CATEGORY - DESCRIPTION = "Compiles BV Regional and applies the built-in Anima attention patch for a standard KSampler." + DESCRIPTION = ( + "Compiles BV Regional and applies the built-in Anima attention patch for a standard KSampler. " + "Regional LoRA hook passes are experimental and add one model pass per distinct effective stack." + ) def apply( self, @@ -495,6 +503,8 @@ class BVRegionalAnimaConditioningNode: base_ratio, cross_inject_every_n_blocks, self_inject_every_n_blocks, + lora_registry=None, + lora_bindings=None, ): try: from ..util.regional.anima_patcher import ApplyAnimaRegionalConditioningPatch @@ -504,7 +514,13 @@ class BVRegionalAnimaConditioningNode: "Update ComfyUI and verify the dependencies reported by the original import error." ) from error - positive, negative, regions, background = compile_anima_adapter(regional, clip) + document = parse_document(regional) + scope_stacks = resolve_scope_stacks(lora_registry, lora_bindings, document) + hook_groups = create_hook_groups(scope_stacks) + positive, negative, regions, background = compile_anima_adapter(document, clip, hook_groups) + positive, negative = apply_attention_hook_passes( + positive, negative, document, scope_stacks, hook_groups + ) patched_model = ApplyAnimaRegionalConditioningPatch().apply( model=model, regions=regions, diff --git a/py/util/regional/anima_adapter.py b/py/util/regional/anima_adapter.py index 58ed2e3..cd24c96 100644 --- a/py/util/regional/anima_adapter.py +++ b/py/util/regional/anima_adapter.py @@ -5,6 +5,7 @@ from typing import Any, Optional import torch +from .clip_hooks import clip_with_hooks from .document import parse_document, selection_prompts from .mask_renderer import render_selection @@ -31,10 +32,11 @@ class AnimaRegionChain: return regions -def _encode(clip: Any, text: str) -> list: +def _encode(clip: Any, text: str, hooks: Any = None) -> list: if clip is None: raise ValueError("clip is required") - return clip.encode_from_tokens_scheduled(clip.tokenize(text)) + encoder = clip_with_hooks(clip, hooks) + return encoder.encode_from_tokens_scheduled(encoder.tokenize(text)) def _zero_out(conditioning: list) -> list: @@ -57,20 +59,26 @@ def _prompt_text(document: dict[str, Any], scope: str, polarity: str) -> tuple[s return prompt["source"], prompt["text"] -def compile_anima_adapter(document: Any, clip: Any) -> tuple[list, list, AnimaRegionChain, list]: +def compile_anima_adapter( + document: Any, + clip: Any, + hooks_by_scope: dict[str, Any] | None = None, +) -> tuple[list, list, AnimaRegionChain, list]: """Compile BV_REGIONAL into the Sen-sou Anima patcher's public node contract.""" clean = parse_document(document) + scoped_hooks = hooks_by_scope or {} width, height = clean["canvas"]["width"], clean["canvas"]["height"] global_positive_source, global_positive_text = _prompt_text(clean, "global", "positive") background_positive_source, background_positive_text = _prompt_text(clean, "background", "positive") base_text = global_positive_text if global_positive_source.strip() else background_positive_text - positive = _encode(clip, base_text) + positive = _encode(clip, base_text, scoped_hooks.get("global")) background = _encode( clip, background_positive_text if background_positive_source.strip() else base_text, + scoped_hooks.get("background"), ) chain: AnimaRegionChain | None = None @@ -87,7 +95,7 @@ def compile_anima_adapter(document: Any, clip: Any) -> tuple[list, list, AnimaRe chain = AnimaRegionChain( previous=chain, mask=mask.detach().float().clamp(0.0, 1.0).cpu().contiguous(), - conditioning=_encode(clip, prompt["text"]), + conditioning=_encode(clip, prompt["text"], scoped_hooks.get(region["id"])), weight=float(region["strength"]), ) @@ -99,6 +107,6 @@ def compile_anima_adapter(document: Any, clip: Any) -> tuple[list, list, AnimaRe if mode == "zero_out" or (mode == "auto" and not global_negative_source.strip()): negative = _zero_out(positive) else: - negative = _encode(clip, global_negative_text) + negative = _encode(clip, global_negative_text, scoped_hooks.get("global")) return positive, negative, chain, background diff --git a/py/util/regional/clip_hooks.py b/py/util/regional/clip_hooks.py new file mode 100644 index 0000000..45c163f --- /dev/null +++ b/py/util/regional/clip_hooks.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from typing import Any + + +def clip_with_hooks(clip: Any, hooks: Any) -> Any: + if hooks is None: + return clip + import comfy.hooks + + hooked = clip.clone(disable_dynamic=True) + hooked.apply_hooks_to_conds = hooks + hooked.patcher.forced_hooks = hooks.clone() + hooked.use_clip_schedule = False + hooked.patcher.forced_hooks.set_keyframes_on_hooks(None) + hooked.patcher.register_all_hook_patches( + hooks, comfy.hooks.create_target_dict(comfy.hooks.EnumWeightTarget.Clip) + ) + return hooked diff --git a/py/util/regional/lora_hooks.py b/py/util/regional/lora_hooks.py index 56f6005..8f05f94 100644 --- a/py/util/regional/lora_hooks.py +++ b/py/util/regional/lora_hooks.py @@ -141,6 +141,7 @@ def resolve_scope_stacks(registry: Any, bindings: Any, document: dict[str, Any]) result: dict[str, Any] = {} global_id = parsed["global_stack_id"] global_entries = stacks[global_id]["stack"] if global_id else [] + global_entries = [entry for entry in global_entries if entry[1] != 0.0 or entry[2] != 0.0] if global_entries: result["global"] = global_entries result["background"] = global_entries @@ -151,6 +152,7 @@ def resolve_scope_stacks(registry: Any, bindings: Any, document: dict[str, Any]) combined.extend(global_entries) if local_id: combined.extend(stacks[local_id]["stack"]) + combined = [entry for entry in combined if entry[1] != 0.0 or entry[2] != 0.0] if combined: result[region_id] = combined return result @@ -188,3 +190,91 @@ def create_hook_groups(scope_stacks: dict[str, Any]) -> dict[str, Any]: hooks.append(hook_group) groups[scope] = comfy.hooks.HookGroup.combine_all_hooks(hooks) return groups + + +def _model_stack_fingerprint(entries: Any) -> tuple[tuple[str, float], ...]: + return tuple( + (str(path), float(model_strength)) + for path, model_strength, _clip_strength in (entries or []) + if float(model_strength) != 0.0 + ) + + +def _clip_stack_fingerprint(entries: Any) -> tuple[tuple[str, float], ...]: + return tuple( + (str(path), float(clip_strength)) + for path, _model_strength, clip_strength in (entries or []) + if float(clip_strength) != 0.0 + ) + + +def _conditionings_with_pass(conditioning: list, hooks: Any, mask: Any = None) -> list: + result = [] + for embedding, metadata in conditioning: + values = metadata.copy() + if hooks is None: + values.pop("hooks", None) + else: + values["hooks"] = hooks + if mask is None: + values.pop("mask", None) + values.pop("mask_strength", None) + values.pop("set_area_to_bounds", None) + else: + values.update({"mask": mask, "mask_strength": 1.0, "set_area_to_bounds": False}) + result.append([embedding, values]) + return result + + +def apply_attention_hook_passes( + positive: list, + negative: list, + document: dict[str, Any], + scope_stacks: dict[str, Any], + hook_groups: dict[str, Any], +) -> tuple[list, list]: + """Split full attention conditioning into masked LoRA model passes.""" + if not scope_stacks: + return positive, negative + + import torch + + from .mask_renderer import render_selection + + width, height = document["canvas"]["width"], document["canvas"]["height"] + baseline_key = _model_stack_fingerprint(scope_stacks.get("global")) + grouped_masks: dict[tuple[tuple[str, float], ...], Any] = {} + representative_scope: dict[tuple[tuple[str, float], ...], str] = {} + + for region in document["regions"]: + if not region["enabled"]: + continue + scope = region["id"] + key = _model_stack_fingerprint(scope_stacks.get(scope)) + if key == baseline_key: + continue + selection = {"document": document, "scope": "region", "region_id": scope} + mask = render_selection(selection, width, height).detach().float().clamp(0.0, 1.0) + grouped_masks[key] = torch.maximum(grouped_masks[key], mask) if key in grouped_masks else mask + representative_scope.setdefault(key, scope) + + if not grouped_masks: + hooks = hook_groups.get("global") + return _conditionings_with_pass(positive, hooks), _conditionings_with_pass(negative, hooks) + + override_union = torch.zeros((1, height, width), dtype=torch.float32) + for mask in grouped_masks.values(): + override_union = torch.maximum(override_union, mask) + baseline_mask = 1.0 - override_union + + positive_passes: list = [] + negative_passes: list = [] + if bool(torch.any(baseline_mask > 0)): + baseline_hooks = hook_groups.get("global") + positive_passes.extend(_conditionings_with_pass(positive, baseline_hooks, baseline_mask)) + negative_passes.extend(_conditionings_with_pass(negative, baseline_hooks, baseline_mask)) + for key, mask in grouped_masks.items(): + hooks = hook_groups.get(representative_scope[key]) + positive_passes.extend(_conditionings_with_pass(positive, hooks, mask)) + negative_passes.extend(_conditionings_with_pass(negative, hooks, mask)) + return positive_passes, negative_passes diff --git a/py/util/regional/native_conditioning.py b/py/util/regional/native_conditioning.py index 6a6d578..eb63fd2 100644 --- a/py/util/regional/native_conditioning.py +++ b/py/util/regional/native_conditioning.py @@ -5,26 +5,11 @@ from typing import Any import torch +from .clip_hooks import clip_with_hooks from .document import parse_document, selection_prompts from .mask_renderer import render_selection -def _clip_with_hooks(clip: Any, hooks: Any) -> Any: - if hooks is None: - return clip - import comfy.hooks - - hooked = clip.clone(disable_dynamic=True) - hooked.apply_hooks_to_conds = hooks - hooked.patcher.forced_hooks = hooks.clone() - hooked.use_clip_schedule = False - hooked.patcher.forced_hooks.set_keyframes_on_hooks(None) - hooked.patcher.register_all_hook_patches( - hooks, comfy.hooks.create_target_dict(comfy.hooks.EnumWeightTarget.Clip) - ) - return hooked - - def _is_anima_clip(clip: Any) -> bool: cond_stage_model = getattr(clip, "cond_stage_model", None) model_type = type(cond_stage_model) @@ -37,7 +22,7 @@ def _is_anima_clip(clip: Any) -> bool: def _encode(clip: Any, text: str, hooks: Any = None) -> list: if clip is None: raise ValueError("clip is required") - encoder = _clip_with_hooks(clip, hooks) + encoder = clip_with_hooks(clip, hooks) conditioning = encoder.encode_from_tokens_scheduled(encoder.tokenize(text)) if hooks is None or getattr(encoder, "apply_hooks_to_conds", None) is hooks: return conditioning diff --git a/pyproject.toml b/pyproject.toml index ad8dacb..cfa381f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "bv_nodepack" -version = "0.10.0" +version = "0.10.1" description = "Visual regional prompting with regional LoRA hooks, native composition modes, shape authoring, mask previews and model-specific attention backends, plus Smart Pipes and ComfyUI workflow tools." authors = [ {name = "BlackVortex", email = "blackvortexx@proton.me"} diff --git a/tests/test_regional_anima_adapter.py b/tests/test_regional_anima_adapter.py index 0374f04..01eaa55 100644 --- a/tests/test_regional_anima_adapter.py +++ b/tests/test_regional_anima_adapter.py @@ -3,6 +3,7 @@ import json from pathlib import Path import sys import unittest +from unittest.mock import patch import torch @@ -65,6 +66,19 @@ class RegionalAnimaAdapterTests(unittest.TestCase): with self.assertRaisesRegex(ValueError, "at least one enabled region"): compile_anima_adapter(document, FakeClip()) + @patch("util.regional.anima_adapter.clip_with_hooks", side_effect=lambda clip, hooks: clip) + def test_encodes_each_scope_with_its_hook_group(self, clip_with_hooks): + document = fixture() + hooks = {"global": "global", "background": "background"} + hooks.update({region["id"]: region["id"] for region in document["regions"]}) + + compile_anima_adapter(document, FakeClip(), hooks) + + used_hooks = [call.args[1] for call in clip_with_hooks.call_args_list] + self.assertEqual(used_hooks[:2], ["global", "background"]) + self.assertEqual(used_hooks[2:5], [region["id"] for region in document["regions"]]) + self.assertEqual(used_hooks[5], "global") + if __name__ == "__main__": unittest.main() diff --git a/tests/test_regional_lora_hooks.py b/tests/test_regional_lora_hooks.py index 9b46304..9f2a4ca 100644 --- a/tests/test_regional_lora_hooks.py +++ b/tests/test_regional_lora_hooks.py @@ -1,7 +1,30 @@ import json import unittest +from unittest.mock import patch -from py.util.regional.lora_hooks import add_named_stack, default_bindings, parse_bindings, parse_registry, reconcile_bindings, resolve_scope_stacks +import torch + +from py.util.regional.lora_hooks import add_named_stack, apply_attention_hook_passes, default_bindings, parse_bindings, parse_registry, reconcile_bindings, resolve_scope_stacks + + +def pass_document(*regions): + return { + "canvas": {"width": 4, "height": 2}, + "regions": [{"id": region_id, "enabled": enabled} for region_id, enabled in regions], + } + + +def conditioning(): + return [[torch.ones((1, 1, 1)), {"source": "test"}]] + + +def selection_mask(selection, _width, _height): + masks = { + "left": torch.tensor([[[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]]), + "right": torch.tensor([[[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 1.0, 1.0]]]), + "overlap": torch.tensor([[[0.0, 0.5, 1.0, 0.0], [0.0, 0.5, 1.0, 0.0]]]), + } + return masks[selection["region_id"]] class RegionalLoraHookContractTests(unittest.TestCase): @@ -73,6 +96,15 @@ class RegionalLoraHookContractTests(unittest.TestCase): self.assertEqual(resolve_scope_stacks(registry, bindings, document), {}) + def test_fully_disabled_stack_is_a_valid_noop(self): + document = {"document_id": "doc-a", "regions": [{"id": "region-a"}]} + registry = add_named_stack(None, "disabled", "Disabled", [("missing.safetensors", 0, 0)]) + bindings = default_bindings("doc-a") + bindings["global_stack_id"] = "disabled" + bindings["regions"] = {"region-a": "disabled"} + + self.assertEqual(resolve_scope_stacks(registry, bindings, document), {}) + def test_orphaned_region_binding_does_not_require_its_old_stack_at_runtime(self): document = {"document_id": "doc-a", "regions": [{"id": "current-region"}]} registry = add_named_stack(None, "current", "Current", [("current.safetensors", 1, 0)]) @@ -97,6 +129,76 @@ class RegionalLoraHookContractTests(unittest.TestCase): self.assertEqual(reconciled["global_stack_id"], "global") self.assertEqual(reconciled["regions"], {"current-region": "current"}) + def test_no_lora_preserves_existing_conditioning(self): + positive, negative = conditioning(), conditioning() + result = apply_attention_hook_passes(positive, negative, pass_document(("left", True)), {}, {}) + self.assertIs(result[0], positive) + self.assertIs(result[1], negative) + + def test_global_lora_uses_one_unmasked_pass(self): + positive, negative = apply_attention_hook_passes( + conditioning(), conditioning(), pass_document(("left", True)), + {"global": [("base.safetensors", 1.0, 0.5)], "left": [("base.safetensors", 1.0, 0.5)]}, + {"global": "global-hooks"}, + ) + self.assertEqual((len(positive), len(negative)), (1, 1)) + self.assertEqual(positive[0][1]["hooks"], "global-hooks") + self.assertNotIn("mask", positive[0][1]) + + @patch("py.util.regional.mask_renderer.render_selection", side_effect=selection_mask) + def test_one_regional_lora_adds_baseline_and_override_passes(self, _render): + positive, negative = apply_attention_hook_passes( + conditioning(), conditioning(), pass_document(("left", True)), + {"left": [("local.safetensors", 0.8, 0.6)]}, {"left": "local-hooks"}, + ) + self.assertEqual((len(positive), len(negative)), (2, 2)) + self.assertNotIn("hooks", positive[0][1]) + self.assertEqual(positive[1][1]["hooks"], "local-hooks") + + @patch("py.util.regional.mask_renderer.render_selection", side_effect=selection_mask) + def test_identical_regional_model_stacks_share_one_override_pass(self, _render): + stack = [("local.safetensors", 0.8, 0.6)] + positive, _ = apply_attention_hook_passes( + conditioning(), conditioning(), pass_document(("left", True), ("overlap", True)), + {"left": stack, "overlap": list(stack)}, {"left": "left-hooks", "overlap": "overlap-hooks"}, + ) + expected = torch.maximum(selection_mask({"region_id": "left"}, 4, 2), selection_mask({"region_id": "overlap"}, 4, 2)) + self.assertEqual(len(positive), 2) + self.assertTrue(torch.equal(positive[1][1]["mask"], expected)) + + @patch("py.util.regional.mask_renderer.render_selection", side_effect=selection_mask) + def test_different_regional_model_stacks_use_separate_passes(self, _render): + positive, _ = apply_attention_hook_passes( + conditioning(), conditioning(), pass_document(("left", True), ("right", True)), + {"left": [("left.safetensors", 0.8, 0.6)], "right": [("right.safetensors", 0.8, 0.6)]}, + {"left": "left-hooks", "right": "right-hooks"}, + ) + self.assertEqual(len(positive), 2) + self.assertEqual({item[1]["hooks"] for item in positive}, {"left-hooks", "right-hooks"}) + + @patch("py.util.regional.mask_renderer.render_selection", side_effect=selection_mask) + def test_clip_only_differences_do_not_add_model_passes(self, render): + positive, _ = apply_attention_hook_passes( + conditioning(), conditioning(), pass_document(("left", True)), + {"global": [("style.safetensors", 1.0, 0.4)], "left": [("style.safetensors", 1.0, 0.9)]}, + {"global": "global-hooks", "left": "left-hooks"}, + ) + self.assertEqual(len(positive), 1) + self.assertEqual(positive[0][1]["hooks"], "global-hooks") + render.assert_not_called() + + @patch("py.util.regional.mask_renderer.render_selection", side_effect=selection_mask) + def test_overlapping_different_stacks_keep_both_override_masks(self, _render): + positive, _ = apply_attention_hook_passes( + conditioning(), conditioning(), pass_document(("left", True), ("overlap", True)), + {"left": [("left.safetensors", 0.8, 0.6)], "overlap": [("overlap.safetensors", 0.8, 0.6)]}, + {"left": "left-hooks", "overlap": "overlap-hooks"}, + ) + by_hook = {item[1].get("hooks"): item[1]["mask"] for item in positive} + self.assertEqual(by_hook["left-hooks"][0, 0, 1].item(), 1.0) + self.assertEqual(by_hook["overlap-hooks"][0, 0, 1].item(), 0.5) + self.assertEqual(by_hook[None][0, 0, 3].item(), 1.0) + if __name__ == "__main__": unittest.main()