Add single-pass Anima regional LoRAs

This commit is contained in:
BlackVortexAI
2026-08-20 08:52:00 +02:00
parent a5fc50677e
commit cfcb49209a
9 changed files with 603 additions and 14 deletions
+44
View File
@@ -124,6 +124,35 @@ single regional stack therefore normally doubles sampling work, while two differ
regional stacks normally triple it. Actual runtime and peak VRAM depend on the model,
resolution, sampler and ComfyUI memory management.
`BV Regional Anima Conditioning` also exposes the experimental
`regional_lora_mode` selector. Its default remains `multipass_legacy`, preserving
published Anima workflows and their image results. `token_gated_singlepass` applies
compatible model-side LoRA deltas directly to the assigned regional text and
spatiotemporal image tokens while regional attention routing is active. After
`end_percent`, regional text routing stops as configured, while the spatially masked
image-side LoRA deltas decay with the remaining sampling sigma from full strength
towards 35% at the end. This preserves late character identity while leaving the base
model more room for fine geometry and without extending the more aggressive regional
text context through the entire denoising schedule. CLIP-side LoRA strengths continue to be
applied while each regional prompt is encoded. Unsupported or non-maskable model
layers fail explicitly; select `multipass_legacy` for those LoRAs. When `base_ratio`
is greater than zero, Anima still evaluates its intentional unmodified baseline pass,
so "single-pass" refers specifically to eliminating the additional per-stack
regional LoRA passes.
| Anima regional LoRA mode | Recommended use | Trade-off |
| --- | --- | --- |
| `multipass_legacy` (default) | Published workflows, maximum compatibility and quality-sensitive interaction scenes | One additional model pass per distinct effective regional model stack |
| `token_gated_singlepass` (experimental) | Faster iteration and compositions whose subjects remain mostly spatially separated | Different image result; overlapping subjects and anatomy can receive competing regional LoRA influence |
Single-pass routing controls where compatible LoRA deltas are applied; it cannot
guarantee anatomy or resolve contradictory prompts beyond the underlying model's own
capabilities. Direct subject interaction, shared limbs, hands, embraces and broad mask
overlaps remain difficult cases even without regional modification. Use
`multipass_legacy` when those details matter more than sampling speed. Existing Anima
workflows that do not store `regional_lora_mode` continue to resolve to
`multipass_legacy`.
### Named LoRA stack interoperability
The incoming community `LORA_STACK` value itself is an unlabelled runtime list:
@@ -912,6 +941,21 @@ Technical references:
## Changelog
### 2026-08-20 — v0.13.0
- Add experimental `token_gated_singlepass` regional model-LoRA execution to
`BV Regional Anima Conditioning` while retaining `multipass_legacy` as the default
for backward compatibility and published-result reproduction.
- Apply compatible regional LoRA deltas to Anima text and spatiotemporal image tokens
inside the existing attention pass, eliminating additional per-stack model passes.
- Preserve global LoRAs, scoped CLIP encoding and the intentional unmodified
`base_ratio` pass; reject unsupported or non-maskable model layers explicitly.
- After the configured attention window, stop regional text routing and retain a
spatial image-LoRA tail that decays towards 35 percent to reduce late identity drift.
- Document that single-pass output intentionally differs from multipass and that
overlapping subjects, shared anatomy and contradictory prompts remain limited by
both competing regional influence and the underlying model's generative capability.
### 2026-08-20 — v0.12.0
- Add token-gated single-pass regional LoRA and LoKr execution to
+20 -4
View File
@@ -44,6 +44,12 @@ DEFAULT_JSON = json.dumps(default_document(), ensure_ascii=False, separators=(",
DEFAULT_LORA_BINDINGS_JSON = json.dumps(default_bindings(), ensure_ascii=False, separators=(",", ":"))
def apply_anima_token_lora_patch(*args, **kwargs):
from ..util.regional.anima_token_lora import apply_anima_token_lora_patch as apply_patch
return apply_patch(*args, **kwargs)
class BVRegionalPromptNode:
@classmethod
def INPUT_TYPES(cls):
@@ -516,6 +522,10 @@ 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}),
"regional_lora_mode": (["multipass_legacy", "token_gated_singlepass"], {
"default": "multipass_legacy",
"tooltip": "Legacy preserves published Anima results. Single-pass is an experimental token-gated model-LoRA path.",
}),
},
"optional": {
"lora_registry": (REGISTRY, {}),
@@ -529,7 +539,7 @@ class BVRegionalAnimaConditioningNode:
CATEGORY = CATEGORY
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."
"Legacy regional LoRA hook passes preserve published results; token-gated single-pass is experimental."
)
def apply(
@@ -546,9 +556,12 @@ class BVRegionalAnimaConditioningNode:
base_ratio,
cross_inject_every_n_blocks,
self_inject_every_n_blocks,
regional_lora_mode="multipass_legacy",
lora_registry=None,
lora_bindings=None,
):
if regional_lora_mode not in {"multipass_legacy", "token_gated_singlepass"}:
raise ValueError("regional_lora_mode must be multipass_legacy or token_gated_singlepass")
try:
from ..util.regional.anima_patcher import ApplyAnimaRegionalConditioningPatch
except ImportError as error:
@@ -563,9 +576,10 @@ class BVRegionalAnimaConditioningNode:
)
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
)
if regional_lora_mode == "multipass_legacy":
positive, negative = apply_attention_hook_passes(
positive, negative, document, scope_stacks, hook_groups
)
patched_model = ApplyAnimaRegionalConditioningPatch().apply(
model=model,
regions=regions,
@@ -580,6 +594,8 @@ class BVRegionalAnimaConditioningNode:
self_inject_every_n_blocks=self_inject_every_n_blocks,
background_conditioning=background,
)[0]
if regional_lora_mode == "token_gated_singlepass":
patched_model = apply_anima_token_lora_patch(patched_model, scope_stacks)
return patched_model, positive, negative
+2
View File
@@ -21,6 +21,7 @@ class AnimaRegionChain:
mask: torch.Tensor
conditioning: list
weight: float
scope: str | None = None
def flatten(self) -> list["AnimaRegionChain"]:
regions = []
@@ -103,6 +104,7 @@ def compile_anima_adapter(
mask=mask.detach().float().clamp(0.0, 1.0).cpu().contiguous(),
conditioning=_encode(clip, prompt["text"], scoped_hooks.get(region["id"])),
weight=float(region["strength"]),
scope=region["id"],
)
if chain is None:
+169 -4
View File
@@ -28,6 +28,7 @@ class AnimaConditioningRegionChain:
mask: torch.Tensor
conditioning: list
weight: float
scope: str | None = None
def flatten(self) -> list["AnimaConditioningRegionChain"]:
regions = []
@@ -383,8 +384,10 @@ class AnimaRegionalConditioningPatch:
# after token-grid resize, which matches Flux-style boolean routing.
self.region_masks: list[torch.Tensor] = []
self.region_weights: list[float] = []
self.region_scopes: list[str] = []
self.region_conditionings: list[tuple[torch.Tensor, dict]] = []
self.background_conditioning: Optional[tuple[torch.Tensor, dict]] = None
self.token_lora = None
if background_conditioning is not None:
cond, metadata = _extract_conditioning_parts(
@@ -400,6 +403,7 @@ class AnimaRegionalConditioningPatch:
)
self.region_masks.append(mask)
self.region_weights.append(weight)
self.region_scopes.append(region.scope or f"region_{idx}")
self.region_conditionings.append(
(cond.detach().float().cpu().contiguous(), metadata.copy())
)
@@ -488,6 +492,103 @@ def _validate_anima_model(model) -> torch.nn.Module:
# Main wrapper
# ---------------------------------------------------------------------------
_LATE_IMAGE_LORA_FINAL_SCALE = 0.35
def _sampling_sigma(transformer_options: dict) -> float | None:
sigmas = transformer_options.get("sigmas", None)
if sigmas is None or not torch.is_tensor(sigmas) or sigmas.numel() == 0:
return None
return float(sigmas.max().detach().cpu().item())
def _late_image_lora_scale(
patch: AnimaRegionalConditioningPatch,
transformer_options: dict,
) -> float | None:
"""Return regional image-LoRA scale after attention routing, or None before it."""
sigma = _sampling_sigma(transformer_options)
if sigma is None:
return None
low = min(patch.start_sigma, patch.end_sigma)
high = max(patch.start_sigma, patch.end_sigma)
if sigma > high:
return None
if sigma >= low:
return 1.0
if low <= 0.0:
return _LATE_IMAGE_LORA_FINAL_SCALE
remaining = max(0.0, min(sigma / low, 1.0))
return _LATE_IMAGE_LORA_FINAL_SCALE + (
1.0 - _LATE_IMAGE_LORA_FINAL_SCALE
) * remaining
def _execute_image_only_token_lora(
executor,
patch: AnimaRegionalConditioningPatch,
args: tuple,
kwargs: dict,
input_x: torch.Tensor,
context: torch.Tensor,
num_chunks: int,
batch_size: int,
regional_scale: float,
):
"""Keep regional model LoRAs spatially active after attention routing ends."""
from .krea2_token_lora import _RUNTIME, _RuntimeContext
diffusion_model = executor.class_obj
patch_spatial = int(getattr(diffusion_model, "patch_spatial", 2))
patch_temporal = int(getattr(diffusion_model, "patch_temporal", 1))
latent_t = int(input_x.shape[2])
latent_h = int(input_x.shape[-2])
latent_w = int(input_x.shape[-1])
padded_t = math.ceil(latent_t / patch_temporal) * patch_temporal
padded_h = math.ceil(latent_h / patch_spatial) * patch_spatial
padded_w = math.ceil(latent_w / patch_spatial) * patch_spatial
temporal_tokens = padded_t // patch_temporal
h_tokens = padded_h // patch_spatial
w_tokens = padded_w // patch_spatial
image_tokens = temporal_tokens * h_tokens * w_tokens
base_mask = torch.zeros((1, padded_h, padded_w), dtype=torch.float32)
spatial_masks = _masks_to_token_masks(
[base_mask] + patch.region_masks,
latent_h,
latent_w,
patch_spatial,
temporal_tokens,
)
text_tokens = int(context.shape[1])
token_masks = patch.token_lora.masks(
[text_tokens] + [0] * len(patch.region_scopes),
patch.region_scopes,
spatial_masks,
num_chunks,
batch_size,
context.device,
)
if regional_scale < 1.0:
global_uids = {
spec.uid for spec in patch.token_lora.specs if "global" in spec.scopes
}
token_masks = {
uid: mask if uid in global_uids else mask * regional_scale
for uid, mask in token_masks.items()
}
runtime = _RuntimeContext(
token_masks=token_masks,
text_tokens=text_tokens,
image_tokens=image_tokens,
text_layers=-1,
image_shape=(temporal_tokens, h_tokens, w_tokens),
)
runtime_token = _RUNTIME.set(runtime)
try:
return executor(*args, **kwargs)
finally:
_RUNTIME.reset(runtime_token)
def _diffusion_model_wrapper(executor, *args, **kwargs):
transformer_options = kwargs.get("transformer_options", None)
if not isinstance(transformer_options, dict):
@@ -496,9 +597,18 @@ def _diffusion_model_wrapper(executor, *args, **kwargs):
patch: Optional[AnimaRegionalConditioningPatch] = transformer_options.get(WRAPPER_KEY, None)
if patch is None:
return executor(*args, **kwargs)
if not patch.is_active(transformer_options):
attention_active = patch.is_active(transformer_options)
if not attention_active and patch.token_lora is None:
return executor(*args, **kwargs)
if patch.base_ratio >= 1.0 or (patch.cross_mask_strength <= 0.0 and patch.self_mask_strength <= 0.0):
late_image_lora_scale = None
if not attention_active:
late_image_lora_scale = _late_image_lora_scale(patch, transformer_options)
if late_image_lora_scale is None:
return executor(*args, **kwargs)
if attention_active and (
patch.base_ratio >= 1.0
or (patch.cross_mask_strength <= 0.0 and patch.self_mask_strength <= 0.0)
):
return executor(*args, **kwargs)
diffusion_model = executor.class_obj
@@ -539,6 +649,12 @@ def _diffusion_model_wrapper(executor, *args, **kwargs):
return executor(*args, **kwargs)
batch_size = B_total // num_chunks
if not attention_active:
return _execute_image_only_token_lora(
executor, patch, args, kwargs, input_x, context, num_chunks, batch_size,
late_image_lora_scale,
)
background_cond = patch.prepare_background_cond(diffusion_model, device, dtype)
background_cond_batched: Optional[torch.Tensor] = None
if background_cond is not None:
@@ -681,7 +797,47 @@ def _diffusion_model_wrapper(executor, *args, **kwargs):
self_parts.append(b.expand(batch_size, -1, -1, -1))
full_self_bias = torch.cat(self_parts, dim=0)
base_output = executor(*args, **kwargs) if patch.base_ratio > 0.0 else None
token_runtime = None
base_runtime = None
if patch.token_lora is not None:
from .krea2_token_lora import _RuntimeContext
h_tokens = padded_h // patch_spatial
w_tokens = padded_w // patch_spatial
image_tokens = temporal_tokens * h_tokens * w_tokens
token_runtime = _RuntimeContext(
token_masks=patch.token_lora.masks(
text_lengths, patch.region_scopes, token_masks,
num_chunks, batch_size, device,
),
text_tokens=S_total,
image_tokens=image_tokens,
text_layers=-1,
image_shape=(temporal_tokens, h_tokens, w_tokens),
)
base_runtime = _RuntimeContext(
token_masks=patch.token_lora.masks(
[S_base], [], torch.zeros((1, image_tokens), device=device),
num_chunks, batch_size, device, global_only=True,
),
text_tokens=S_base,
image_tokens=image_tokens,
text_layers=-1,
image_shape=(temporal_tokens, h_tokens, w_tokens),
)
base_output = None
if patch.base_ratio > 0.0:
if base_runtime is None:
base_output = executor(*args, **kwargs)
else:
from .krea2_token_lora import _RUNTIME
runtime_token = _RUNTIME.set(base_runtime)
try:
base_output = executor(*args, **kwargs)
finally:
_RUNTIME.reset(runtime_token)
# ---- Patch cross_attn.attn_op on all blocks ------------------------
# We replace the attn_op callable (normally torch_attention_op) with our
@@ -720,7 +876,16 @@ def _diffusion_model_wrapper(executor, *args, **kwargs):
kwargs["context"] = unified_context
args = tuple(args)
regional_output = executor(*args, **kwargs)
if token_runtime is None:
regional_output = executor(*args, **kwargs)
else:
from .krea2_token_lora import _RUNTIME
runtime_token = _RUNTIME.set(token_runtime)
try:
regional_output = executor(*args, **kwargs)
finally:
_RUNTIME.reset(runtime_token)
if base_output is not None and torch.is_tensor(regional_output) and torch.is_tensor(base_output):
return regional_output * (1.0 - patch.base_ratio) + base_output * patch.base_ratio
return regional_output
+132
View File
@@ -0,0 +1,132 @@
from __future__ import annotations
import copy
import logging
from dataclasses import dataclass, field
from typing import Any
import torch
from .anima_patcher import WRAPPER_KEY as ANIMA_WRAPPER_KEY
from .krea2_token_lora import (
TokenLoRAReport,
TokenLoRASpec,
_inject_adapter,
build_token_lora_specs,
)
BACKEND_ID = "anima_token_gated_singlepass"
ATTACHMENT_KEY = "bv_anima_token_lora"
@dataclass
class AnimaTokenLoRAPatch:
specs: list[TokenLoRASpec]
report: TokenLoRAReport
_cache: dict[tuple[Any, ...], dict[str, torch.Tensor]] = field(default_factory=dict, init=False)
def masks(
self,
text_lengths: list[int],
region_scopes: list[str],
spatial_masks: torch.Tensor,
batch_chunks: int,
batch_size: int,
device: torch.device,
*,
global_only: bool = False,
) -> dict[str, torch.Tensor]:
key = (
tuple(text_lengths), tuple(region_scopes), tuple(spatial_masks.shape),
batch_chunks, batch_size, device, global_only,
)
if key in self._cache:
return self._cache[key]
text_tokens = sum(text_lengths)
image_tokens = int(spatial_masks.shape[1])
masks: dict[str, torch.Tensor] = {}
for spec in self.specs:
per_chunk = torch.zeros(
batch_chunks, text_tokens + image_tokens, device=device, dtype=torch.float32
)
if "global" in spec.scopes:
per_chunk.fill_(1.0)
elif not global_only:
offset = text_lengths[0]
for scope, length in zip(region_scopes, text_lengths[1:]):
if scope in spec.scopes:
per_chunk[:, offset:offset + length] = 1.0
offset += length
image_mask = torch.zeros(image_tokens, device=device, dtype=torch.float32)
for index, scope in enumerate(region_scopes, start=1):
if scope in spec.scopes:
spatial_mask = spatial_masks[index].to(
device=device, dtype=image_mask.dtype
)
image_mask = torch.maximum(image_mask, spatial_mask)
per_chunk[:, text_tokens:] = image_mask
masks[spec.uid] = per_chunk.repeat_interleave(batch_size, dim=0).unsqueeze(-1)
self._cache[key] = masks
return masks
def apply_anima_token_lora_patch(
model: Any,
scope_stacks: dict[str, Any],
) -> Any:
specs = build_token_lora_specs(scope_stacks)
if not specs:
return model
try:
import comfy.utils
except ImportError as error:
raise RuntimeError("Anima token-gated LoRA requires a current ComfyUI installation") from error
result = model.clone()
report = TokenLoRAReport(mode="token_gated_singlepass", adapters=len(specs))
loaded: dict[str, dict[str, Any]] = {}
for spec in specs:
state = loaded.get(spec.path)
if state is None:
state = comfy.utils.load_torch_file(spec.path, safe_load=True)
loaded[spec.path] = state
patched, unmatched, unmaskable = _inject_adapter(
result, state, spec, excluded_fragments=("adaln_modulation",)
)
report.patched_layers += patched
if unmatched:
report.unmatched_keys[spec.path] = unmatched
if unmaskable:
report.unmaskable_keys[spec.path] = unmaskable
if patched == 0:
raise RuntimeError(
f"Anima token-gated LoRA matched no maskable model layers: {spec.path}. "
"Use multipass_legacy for this LoRA."
)
if report.unmaskable_keys:
examples = [f"{path}: {keys[:3]}" for path, keys in report.unmaskable_keys.items()]
raise RuntimeError(
"Anima token-gated LoRA found model layers without a spatial or regional text axis; "
"use multipass_legacy. " + "; ".join(examples)
)
if report.unmatched_keys:
logging.warning(
"BV Anima token LoRA ignored unmatched non-model keys: %s",
{path: keys[:3] for path, keys in report.unmatched_keys.items()},
)
attention_patch = result.model_options.get("transformer_options", {}).get(ANIMA_WRAPPER_KEY)
if attention_patch is None:
raise RuntimeError("Apply the Anima regional attention patch before token-gated LoRAs")
attention_patch = copy.copy(attention_patch)
attention_patch.token_lora = AnimaTokenLoRAPatch(specs, report)
result.model_options.setdefault("transformer_options", {})[ANIMA_WRAPPER_KEY] = attention_patch
result.set_attachments(ANIMA_WRAPPER_KEY, attention_patch)
result.set_attachments(ATTACHMENT_KEY, attention_patch.token_lora)
logging.info(
"BV Anima regional LoRA mode=%s adapters=%d patched_layers=%d",
report.mode, report.adapters, report.patched_layers,
)
return result
+24 -5
View File
@@ -43,6 +43,7 @@ class _RuntimeContext:
text_tokens: int
image_tokens: int
text_layers: int
image_shape: tuple[int, int, int] | None = None
_RUNTIME: contextvars.ContextVar[_RuntimeContext | None] = contextvars.ContextVar(
@@ -173,6 +174,13 @@ def _adapt_mask(mask: torch.Tensor, value: torch.Tensor) -> torch.Tensor | None:
if value.ndim == 4 and int(value.shape[1]) == runtime.text_tokens:
result = mask[:, :runtime.text_tokens]
return result[..., None] if int(result.shape[0]) in (1, int(value.shape[0])) else None
if value.ndim == 5 and runtime.image_shape is not None:
temporal, height, width = runtime.image_shape
if tuple(value.shape[1:4]) != (temporal, height, width):
return None
result = mask[:, runtime.text_tokens:runtime.text_tokens + runtime.image_tokens]
result = result.reshape(result.shape[0], temporal, height, width, 1)
return result if int(result.shape[0]) in (1, int(value.shape[0])) else None
return None
@@ -190,7 +198,7 @@ class _TokenLoRAMixin:
result = super().forward(value, *args, **kwargs)
adapters = getattr(self, "bv_krea2_token_adapters", None)
runtime = _RUNTIME.get()
if not adapters or runtime is None or not torch.is_tensor(value) or value.ndim not in (3, 4):
if not adapters or runtime is None or not torch.is_tensor(value) or value.ndim not in (3, 4, 5):
return result
cache = self.__dict__.setdefault("_bv_krea2_token_cast", {})
for uid, adapter in adapters.items():
@@ -241,7 +249,10 @@ def _patched_module(module: torch.nn.Module) -> torch.nn.Module:
return patched
def _module_lookup(diffusion_model: Any) -> tuple[dict[str, str], dict[str, str]]:
def _module_lookup(
diffusion_model: Any,
excluded_fragments: tuple[str, ...] = (),
) -> tuple[dict[str, str], dict[str, str]]:
maskable: dict[str, str] = {}
all_linear: dict[str, str] = {}
for name, module in diffusion_model.named_modules():
@@ -250,15 +261,23 @@ def _module_lookup(diffusion_model: Any) -> tuple[dict[str, str], dict[str, str]
continue
for key in (name, name.replace(".", "_")):
all_linear[key] = name
if name.startswith(("blocks.", "txtfusion.", "txtmlp", "first", "last")):
if (
name.startswith(("blocks.", "txtfusion.", "txtmlp", "first", "last"))
and not any(fragment in name for fragment in excluded_fragments)
):
maskable[key] = name
return maskable, all_linear
def _inject_adapter(model: Any, state: dict[str, Any], spec: TokenLoRASpec) -> tuple[int, list[str], list[str]]:
def _inject_adapter(
model: Any,
state: dict[str, Any],
spec: TokenLoRASpec,
excluded_fragments: tuple[str, ...] = (),
) -> tuple[int, list[str], list[str]]:
normalized = _normalize_lora_state(state)
diffusion_model = model.get_model_object("diffusion_model")
maskable, all_linear = _module_lookup(diffusion_model)
maskable, all_linear = _module_lookup(diffusion_model, excluded_fragments)
patched_count = 0
unmatched: list[str] = []
unmaskable: list[str] = []
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "bv_nodepack"
version = "0.12.0"
version = "0.13.0"
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"}
+142
View File
@@ -23,6 +23,16 @@ from util.regional.anima_patcher import ( # noqa: E402
WRAPPER_KEY,
_build_flux_cross_attention_bias,
_build_flux_self_attention_bias,
_diffusion_model_wrapper,
_late_image_lora_scale,
)
from util.regional.anima_token_lora import AnimaTokenLoRAPatch # noqa: E402
from util.regional.krea2_token_lora import ( # noqa: E402
TokenLoRAReport,
TokenLoRASpec,
_RUNTIME,
_RuntimeContext,
_patched_module,
)
@@ -30,6 +40,138 @@ class RegionalAnimaPatcherTests(unittest.TestCase):
def test_uses_bv_specific_wrapper_key(self):
self.assertEqual(WRAPPER_KEY, "bv_anima_regional_conditioning")
def test_token_lora_masks_region_text_and_spatial_tokens(self):
local = TokenLoRASpec("local", "local.safetensors", 1.0, frozenset({"left"}))
global_spec = TokenLoRASpec("global", "global.safetensors", 1.0, frozenset({"global"}))
patch = AnimaTokenLoRAPatch([local, global_spec], TokenLoRAReport())
spatial = torch.tensor([
[0.0, 0.0, 0.0, 0.0],
[1.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 1.0],
])
masks = patch.masks(
[2, 2, 1], ["left", "right"], spatial,
batch_chunks=2, batch_size=1, device=torch.device("cpu"),
)
self.assertEqual(
masks["local"][0, :, 0].tolist(),
[0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0],
)
self.assertTrue(torch.all(masks["global"] == 1.0))
def test_token_lora_moves_cpu_spatial_masks_to_the_runtime_device(self):
local = TokenLoRASpec("local", "local.safetensors", 1.0, frozenset({"left"}))
patch = AnimaTokenLoRAPatch([local], TokenLoRAReport())
spatial = torch.tensor([
[0.0, 0.0, 0.0, 0.0],
[1.0, 1.0, 0.0, 0.0],
], device="cpu")
masks = patch.masks(
[2, 2], ["left"], spatial,
batch_chunks=1, batch_size=1, device=torch.device("meta"),
)
self.assertEqual(masks["local"].device.type, "meta")
def test_token_lora_gates_five_dimensional_anima_activations(self):
linear = torch.nn.Linear(2, 2, bias=False)
linear.weight.data.zero_()
patched = _patched_module(linear)
patched.bv_krea2_token_adapters["adapter"] = (
"lora", torch.eye(2), torch.eye(2), 1.0,
)
runtime = _RuntimeContext(
token_masks={
"adapter": torch.tensor(
[[[0.0], [0.0], [1.0], [0.0], [0.5], [0.0]]]
)
},
text_tokens=2,
image_tokens=4,
text_layers=-1,
image_shape=(1, 2, 2),
)
token = _RUNTIME.set(runtime)
try:
output = patched(torch.ones((1, 1, 2, 2, 2)))
finally:
_RUNTIME.reset(token)
self.assertTrue(torch.equal(output[0, 0, 0, 0], torch.ones(2)))
self.assertTrue(torch.equal(output[0, 0, 0, 1], torch.zeros(2)))
self.assertTrue(torch.equal(output[0, 0, 1, 0], torch.full((2,), 0.5)))
def test_late_sampling_keeps_image_lora_without_regional_text_context(self):
local = TokenLoRASpec("local", "local.safetensors", 1.0, frozenset({"left"}))
global_spec = TokenLoRASpec("global", "global.safetensors", 1.0, frozenset({"global"}))
token_patch = AnimaTokenLoRAPatch([local, global_spec], TokenLoRAReport())
class LatePatch:
token_lora = token_patch
region_masks = [torch.ones((1, 4, 4))]
region_scopes = ["left"]
start_sigma = 10.0
end_sigma = 1.0
@staticmethod
def is_active(_options):
return False
class FakeModel:
patch_spatial = 2
patch_temporal = 1
class RecordingExecutor:
class_obj = FakeModel()
def __init__(self):
self.context_tokens = None
self.runtime = None
def __call__(self, *args, **_kwargs):
self.context_tokens = int(args[2].shape[1])
self.runtime = _RUNTIME.get()
return args[0]
executor = RecordingExecutor()
latent = torch.zeros((2, 4, 1, 4, 4))
context = torch.zeros((2, 3, 8))
options = {
WRAPPER_KEY: LatePatch(),
"cond_or_uncond": [0, 1],
"sigmas": torch.tensor([0.0]),
}
_diffusion_model_wrapper(
executor, latent, torch.zeros(2), context,
transformer_options=options,
)
self.assertEqual(executor.context_tokens, 3)
self.assertIsNotNone(executor.runtime)
self.assertEqual(executor.runtime.text_tokens, 3)
self.assertEqual(executor.runtime.image_shape, (1, 2, 2))
local_mask = executor.runtime.token_masks["local"]
self.assertTrue(torch.all(local_mask[:, :3] == 0))
self.assertTrue(torch.allclose(local_mask[:, 3:], torch.full_like(local_mask[:, 3:], 0.35)))
self.assertTrue(torch.all(executor.runtime.token_masks["global"] == 1))
def test_late_image_lora_scale_decays_linearly_after_attention_window(self):
class Patch:
start_sigma = 10.0
end_sigma = 1.0
self.assertIsNone(_late_image_lora_scale(Patch(), {"sigmas": torch.tensor([11.0])}))
self.assertEqual(_late_image_lora_scale(Patch(), {"sigmas": torch.tensor([1.0])}), 1.0)
self.assertAlmostEqual(
_late_image_lora_scale(Patch(), {"sigmas": torch.tensor([0.5])}),
0.675,
)
self.assertEqual(_late_image_lora_scale(Patch(), {"sigmas": torch.tensor([0.0])}), 0.35)
def test_cross_attention_bias_preserves_overlap_membership(self):
masks = torch.tensor(
[
+69
View File
@@ -237,6 +237,75 @@ class RegionalNodeTests(unittest.TestCase):
self.module.BVRegionalAnimaConditioningNode.RETURN_TYPES,
("MODEL", "CONDITIONING", "CONDITIONING"),
)
mode = self.module.BVRegionalAnimaConditioningNode.INPUT_TYPES()["required"]["regional_lora_mode"]
self.assertEqual(mode[0], ["multipass_legacy", "token_gated_singlepass"])
self.assertEqual(mode[1]["default"], "multipass_legacy")
def test_anima_singlepass_skips_legacy_model_hook_passes(self):
node = self.module.BVRegionalAnimaConditioningNode()
class FakeApply:
def apply(self, **_kwargs):
return ("attention-model",)
fake_patcher = types.ModuleType(f"{PACKAGE}.py.util.regional.anima_patcher")
fake_patcher.ApplyAnimaRegionalConditioningPatch = FakeApply
with (
unittest.mock.patch.dict(sys.modules, {fake_patcher.__name__: fake_patcher}),
unittest.mock.patch.object(self.module, "resolve_scope_stacks", return_value={}),
unittest.mock.patch.object(self.module, "resolve_stack_paths", return_value={"region-a": []}),
unittest.mock.patch.object(self.module, "create_hook_groups", return_value={}),
unittest.mock.patch.object(
self.module, "compile_anima_adapter",
return_value=(["positive"], ["negative"], "regions", ["background"]),
),
unittest.mock.patch.object(self.module, "apply_attention_hook_passes") as legacy,
unittest.mock.patch.object(
self.module, "apply_anima_token_lora_patch", return_value="singlepass-model"
) as singlepass,
):
result = node.apply(
"model", "clip", fixture(), "disabled", 0.2, 0.0, 0.35,
1.0, 0.2, 0.1, 1, 1,
regional_lora_mode="token_gated_singlepass",
)
legacy.assert_not_called()
singlepass.assert_called_once_with("attention-model", {"region-a": []})
self.assertEqual(result, ("singlepass-model", ["positive"], ["negative"]))
def test_anima_call_without_mode_preserves_legacy_hook_passes(self):
node = self.module.BVRegionalAnimaConditioningNode()
class FakeApply:
def apply(self, **_kwargs):
return ("attention-model",)
fake_patcher = types.ModuleType(f"{PACKAGE}.py.util.regional.anima_patcher")
fake_patcher.ApplyAnimaRegionalConditioningPatch = FakeApply
with (
unittest.mock.patch.dict(sys.modules, {fake_patcher.__name__: fake_patcher}),
unittest.mock.patch.object(self.module, "resolve_scope_stacks", return_value={}),
unittest.mock.patch.object(self.module, "resolve_stack_paths", return_value={"region-a": []}),
unittest.mock.patch.object(self.module, "create_hook_groups", return_value={}),
unittest.mock.patch.object(
self.module, "compile_anima_adapter",
return_value=(["positive"], ["negative"], "regions", ["background"]),
),
unittest.mock.patch.object(
self.module, "apply_attention_hook_passes",
return_value=(["legacy-positive"], ["legacy-negative"]),
) as legacy,
unittest.mock.patch.object(self.module, "apply_anima_token_lora_patch") as singlepass,
):
result = node.apply(
"model", "clip", fixture(), "disabled", 0.2, 0.0, 0.35,
1.0, 0.2, 0.1, 1, 1,
)
legacy.assert_called_once()
singlepass.assert_not_called()
self.assertEqual(result, ("attention-model", ["legacy-positive"], ["legacy-negative"]))
def test_color_control_node_is_registered_and_returns_image_with_legend(self):
node = self.module.BVRegionalColorControlImageNode()