Add single-pass Anima regional LoRAs
This commit is contained in:
@@ -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
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
@@ -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"}
|
||||
|
||||
@@ -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(
|
||||
[
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user