diff --git a/README.md b/README.md index d6830f3..1fee09a 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/py/nodes/bv_regional.py b/py/nodes/bv_regional.py index fe42840..87cd52c 100644 --- a/py/nodes/bv_regional.py +++ b/py/nodes/bv_regional.py @@ -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 diff --git a/py/util/regional/anima_adapter.py b/py/util/regional/anima_adapter.py index 00168f9..94d9504 100644 --- a/py/util/regional/anima_adapter.py +++ b/py/util/regional/anima_adapter.py @@ -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: diff --git a/py/util/regional/anima_patcher.py b/py/util/regional/anima_patcher.py index 58d3bb3..11bc847 100644 --- a/py/util/regional/anima_patcher.py +++ b/py/util/regional/anima_patcher.py @@ -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 diff --git a/py/util/regional/anima_token_lora.py b/py/util/regional/anima_token_lora.py new file mode 100644 index 0000000..b8e20a2 --- /dev/null +++ b/py/util/regional/anima_token_lora.py @@ -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 diff --git a/py/util/regional/krea2_token_lora.py b/py/util/regional/krea2_token_lora.py index 4af375f..9a1a327 100644 --- a/py/util/regional/krea2_token_lora.py +++ b/py/util/regional/krea2_token_lora.py @@ -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] = [] diff --git a/pyproject.toml b/pyproject.toml index cd3aab4..e771ba4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"} diff --git a/tests/test_regional_anima_patcher.py b/tests/test_regional_anima_patcher.py index bc38d06..72921bb 100644 --- a/tests/test_regional_anima_patcher.py +++ b/tests/test_regional_anima_patcher.py @@ -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( [ diff --git a/tests/test_regional_nodes.py b/tests/test_regional_nodes.py index 259445a..9bd57a8 100644 --- a/tests/test_regional_nodes.py +++ b/tests/test_regional_nodes.py @@ -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()