diff --git a/encoding/README.md b/encoding/README.md new file mode 100644 index 0000000..dd769cd --- /dev/null +++ b/encoding/README.md @@ -0,0 +1,47 @@ +# Conditioning implementations + +Both prompt nodes use `get_encoder(clip)`. They do not branch on model names. + +| Module | Responsibility | +| --- | --- | +| `base.py` | Encoding contract and capabilities for BREAK, scheduling, and normalization | +| `clip.py` | CLIP/SDXL token batches, BREAK isolation, direct emphasis scaling and pooling | +| `anima.py` | Anima mode selection and preservation of the metadata consumed by its T5-token adapter | +| `anima_tokenization.py` | Single-pass Qwen/T5 tokenization and character-span weight mapping | +| `anima_attention.py` | Conditioning-scoped log-weight attention masks after key normalization | +| `__init__.py` | Explicit selection from supported token-stream contracts | + +An implementation returns `(conditioning_tensor, metadata_dict)`. Metadata must +include every field needed by the downstream model, including pooled output, +attention masks, IDs, and weights when present. Token streams can have independent +lengths. Implementations do not retain per-request state. + +The positive node owns wildcard expansion, schedule parsing, encoding caches and +per-step orchestration. The negative node selects the first parsed prompt and +uses the same encoding contract. The existing schedule hook switches embedding +tensors, while pooled output stays fixed at the first prompt's value. In +particular, SDXL pooled conditioning is not switched with later prompts. +Additional metadata is rejected when scheduling until it can be switched +correctly. The metadata guard runs before caching each new scheduled encoding. + +CLIP supports BREAK, direct weighting, normalization and the existing schedule +path. Anima supports static native, single-pass value, and attention weighting. Anima +BREAK is rejected because isolation through the downstream adapter is undefined; +encoding fragments separately before that adapter does not establish isolation. +Anima scheduling and normalization are also unsupported. + +Nodes expose `anima_emphasis`, defaulting to `native`, and include it in encoding +cache keys. Non-native choices on other model families fail explicitly. Attention +mode neutralizes adapter value weights and attaches a `HookedOnly` transformer +hook to that conditioning. Comfy batches different hook groups separately. A +diffusion wrapper creates local attn1/attn2 markers and an attention override; +the override adds `log(weight)` as a mask through Comfy's existing attention +backend, composing existing masks and overrides. Model modules and caller option +dictionaries are not mutated. Multiple Anima attention hooks merged into the +same conditioning are explicitly unsupported. Native mode imports neither the +single-pass tokenizer helper nor the attention hook module. + +To add another conditioning path, implement it in a separate module, declare its +capabilities, and add an explicit selector entry. Its tests should verify native +plain-text parity, weighting, metadata, and supported operations through the +consuming model. Unknown token streams must not fall back to CLIP. diff --git a/encoding/__init__.py b/encoding/__init__.py new file mode 100644 index 0000000..61724f8 --- /dev/null +++ b/encoding/__init__.py @@ -0,0 +1,30 @@ +"""Explicit selection of supported encoder/model conditioning contracts.""" + +from .anima import AnimaEncoder, ANIMA_EMPHASIS_MODES +from .clip import CLIPEncoder + +_ENCODERS = (AnimaEncoder(), CLIPEncoder(("l",)), CLIPEncoder(("l", "g"))) + + +def get_encoder(clip, anima_emphasis="native"): + """Inspect token streams without loading encoder weights; never guess CLIP.""" + if clip is None: + raise ValueError("clip input is None") + if anima_emphasis not in ANIMA_EMPHASIS_MODES: + raise ValueError(f"Unknown Anima emphasis mode: {anima_emphasis}") + tokens = clip.tokenize("") + streams = set(tokens) if isinstance(tokens, dict) else set() + for encoder in _ENCODERS: + if encoder.matches(clip, streams): + if isinstance(encoder, AnimaEncoder): + return encoder if anima_emphasis == "native" else AnimaEncoder(anima_emphasis) + if anima_emphasis != "native": + raise ValueError("Anima emphasis modes apply only to Anima. Select native for CLIP/SDXL.") + return encoder + raise ValueError( + f"Unsupported text encoder token streams: {sorted(streams)}. " + "Supported conditioning paths: SD1.5 CLIP, SDXL CLIP-L/G, and static Anima." + ) + + +__all__ = ["get_encoder", "ANIMA_EMPHASIS_MODES"] diff --git a/encoding/anima.py b/encoding/anima.py new file mode 100644 index 0000000..f4c4259 --- /dev/null +++ b/encoding/anima.py @@ -0,0 +1,46 @@ +"""Anima's native Qwen embeddings and T5-token conditioning adapter contract.""" + +from .base import EncoderCapabilities, PromptEncoder + +ANIMA_EMPHASIS_MODES = ("native", "single_pass", "attention") + + +class AnimaEncoder(PromptEncoder): + name = "Anima" + capabilities = EncoderCapabilities(breaks=False, scheduling=False, normalization=False) + + def __init__(self, emphasis="native"): + if emphasis not in ANIMA_EMPHASIS_MODES: + raise ValueError(f"Unknown Anima emphasis mode: {emphasis}") + self.emphasis = emphasis + + def matches(self, clip, streams): + return ( + streams == {"qwen3_06b", "t5xxl"} + and getattr(getattr(clip, "cond_stage_model", None), "clip_name", None) == "qwen3_06b" + ) + + def encode(self, clip, text, normalization=False, debug=False): + self.validate_request(text, normalization=normalization) + hooks = None + if self.emphasis == "native": + tokens = clip.tokenize(text) + else: + from .anima_tokenization import single_pass_tokens + + tokens, _ = single_pass_tokens(clip.tokenizer, text) + if self.emphasis == "attention": + from .anima_attention import attention_hooks + + hooks = attention_hooks([token[1] for token in tokens["t5xxl"][0]]) + tokens["t5xxl"] = [[(token[0], 1.0) for token in tokens["t5xxl"][0]]] + if not self.matches(clip, set(tokens)): + raise ValueError("Anima encoder returned unexpected token streams.") + # T5 IDs and weights are consumed by Anima's adapter, not by direct + # scaling of Qwen embeddings. BREAK isolation through that adapter and + # preprocessing-aware schedules have not been defined or qualified. + output = clip.encode_from_tokens(tokens, return_dict=True).copy() + if hooks is not None: + existing = output.get("hooks") + output["hooks"] = hooks if existing is None else existing.clone_and_combine(hooks) + return output.pop("cond"), output diff --git a/encoding/anima_attention.py b/encoding/anima_attention.py new file mode 100644 index 0000000..ae0783c --- /dev/null +++ b/encoding/anima_attention.py @@ -0,0 +1,104 @@ +"""Conditioning-scoped Anima attention bias, without changing model objects.""" + +import torch +from comfy.hooks import HookGroup, TransformerOptionsHook, EnumHookScope + + +class AnimaAttention: + """Install invocation-local cross-attention markers and a log-weight mask. + + Cosmos calls attn2 patches before projection and the attention override after + key normalization. Marking the call avoids guessing from sequence lengths. + HookedOnly makes Comfy batch this conditioning separately from other prompts. + """ + + def __init__(self, weights): + self.weights = torch.as_tensor(weights, dtype=torch.float32).clone() + if self.weights.ndim != 1 or not bool(torch.isfinite(self.weights).all()) or not bool((self.weights > 0).all()): + raise ValueError("Anima attention mode requires finite, strictly positive weights. Use single_pass or native for zero/negative weights.") + + def __call__(self, executor, *args, **kwargs): + from comfy.ldm.anima.model import Anima + + if not isinstance(executor.class_obj, Anima): + raise ValueError("Anima attention conditioning requires an Anima diffusion model.") + options = dict(kwargs.get("transformer_options", {})) + if options.get("a1111_anima_attention"): + raise ValueError("Combining multiple Anima attention hooks on one conditioning is unsupported.") + options["a1111_anima_attention"] = True + state = {"cross": False, "calls": 0, "patches": 0, "last_cross": False} + biases = {} + + def mark_self(*args, **kwargs): + state["cross"] = False + state["last_cross"] = False + return {} + + def mark_cross(*args, **kwargs): + state["cross"] = True + state["last_cross"] = True + state["patches"] += 1 + return {} + + patches = dict(options.get("patches", {})) + patches["attn1_patch"] = list(patches.get("attn1_patch", ())) + [mark_self] + patches["attn2_patch"] = list(patches.get("attn2_patch", ())) + [mark_cross] + options["patches"] = patches + previous = options.get("optimized_attention_override") + + def attention(original, q, k, v, heads, **attn_kwargs): + cross = state["cross"] + if state["last_cross"] and not cross: + raise ValueError("Unexpected repeated Anima cross-attention call. Update the node or select native.") + state["cross"] = False + if cross: + state["calls"] += 1 + if not attn_kwargs.get("skip_reshape") or k.ndim != 4: + raise ValueError("Unexpected Anima attention layout; update ComfyUI or use native mode.") + # Anima.preprocess_text_embeds pads T5 adapter output to 512. + expected = max(512, len(self.weights)) + if k.shape[-2] != expected: + raise ValueError("Anima attention weights do not match the prepared context length.") + key = (q.device, q.dtype) + if key not in biases: + weights = self.weights.to(device=q.device) + biases[key] = torch.nn.functional.pad(weights.log(), (0, expected - len(weights))).to(q.dtype)[None, None, None, :] + bias = biases[key] + mask = attn_kwargs.get("mask") + if mask is not None: + mask = mask.to(device=q.device) + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim == 3: + mask = mask.unsqueeze(1) + if mask.dtype == torch.bool: + mask = torch.zeros_like(mask, dtype=q.dtype).masked_fill(~mask, float("-inf")) + else: + mask = mask.to(dtype=q.dtype) + bias = bias + mask + attn_kwargs["mask"] = bias + # Keep the installed attention backend and any prior override. Mask + # support/fallback is handled by ComfyUI, including output layout. + target = previous if previous is not None else lambda f, *a, **kw: f(*a, **kw) + return target(original, q, k, v, heads, **attn_kwargs) + + options["optimized_attention_override"] = attention + kwargs["transformer_options"] = options + result = executor(*args, **kwargs) + if not state["calls"]: + raise ValueError("This ComfyUI attention implementation does not support Anima attention mode. Update ComfyUI or select native.") + if state["calls"] != state["patches"]: + raise ValueError("Anima attention patch/call mismatch. Update the node or select native.") + return result + + +def attention_hooks(weights): + attention = AnimaAttention(weights) + if bool((attention.weights == 1).all()): + return None + hooks = HookGroup() + hooks.add(TransformerOptionsHook( + transformers_dict={"wrappers": {"diffusion_model": {"a1111_anima_attention": [attention]}}}, + hook_scope=EnumHookScope.HookedOnly, + )) + return hooks diff --git a/encoding/anima_tokenization.py b/encoding/anima_tokenization.py new file mode 100644 index 0000000..c3cf823 --- /dev/null +++ b/encoding/anima_tokenization.py @@ -0,0 +1,81 @@ +"""Single-pass Anima tokenization with character-span emphasis mapping. + +Reuse ComfyUI's emphasis parser, reconstruct the complete text, then tokenize +each stream once. T5 fast-tokenizer character offsets attach weights without +fragmenting either token stream. Tokens crossing weight boundaries receive the +mean weight of their non-whitespace characters; special tokens stay neutral. +""" + +from dataclasses import dataclass +import math + + +@dataclass(frozen=True) +class WeightedText: + text: str + character_weights: tuple + + +def join_weighted_segments(segments): + text = [] + weights = [] + for segment, weight in segments: + weight = float(weight) + if not math.isfinite(weight): + raise ValueError("Emphasis weights must be finite.") + text.append(segment) + weights.extend([weight] * len(segment)) + return WeightedText("".join(text), tuple(weights)) + + +def parse_emphasis(text): + from comfy.sd1_clip import escape_important, token_weights, unescape_important + + if "embedding:" in text: + raise ValueError("Anima single-pass/attention modes do not support embedding expansion. Use native mode.") + segments = token_weights(escape_important(text), 1.0) + return join_weighted_segments( + (unescape_important(segment), weight) for segment, weight in segments + ) + + +def weights_for_offsets(weighted_text, offsets): + weights = [] + crossings = [] + for index, (start, end) in enumerate(offsets): + if not 0 <= start <= end <= len(weighted_text.text): + raise ValueError("Tokenizer offset is outside the reconstructed text.") + values = [ + weighted_text.character_weights[pos] + for pos in range(start, end) + if not weighted_text.text[pos].isspace() + ] + weights.append(sum(values) / len(values) if values else 1.0) + if len(set(values)) > 1: + crossings.append(index) + return weights, crossings + + +def single_pass_tokens(tokenizer, text): + weighted_text = parse_emphasis(text) + qwen = tokenizer.qwen3_06b.tokenize_with_weights( + weighted_text.text, disable_weights=True + ) + t5 = tokenizer.t5xxl.tokenize_with_weights( + weighted_text.text, disable_weights=True + ) + offsets = tokenizer.t5xxl.tokenizer( + weighted_text.text, return_offsets_mapping=True + ) + if len(t5) != 1 or [item[0] for item in t5[0]] != offsets["input_ids"]: + raise ValueError("T5 offset IDs do not match ComfyUI's full-text IDs.") + weights, crossings = weights_for_offsets(weighted_text, offsets["offset_mapping"]) + tokens = { + "qwen3_06b": qwen, + "t5xxl": [[(item[0], weight) for item, weight in zip(t5[0], weights)]], + } + return tokens, { + "clean_text": weighted_text.text, + "t5_offsets": list(offsets["offset_mapping"]), + "boundary_crossings": crossings, + } diff --git a/encoding/base.py b/encoding/base.py new file mode 100644 index 0000000..ec9641b --- /dev/null +++ b/encoding/base.py @@ -0,0 +1,52 @@ +"""Contract between prompt orchestration and a conditioning implementation.""" + +import re +from abc import ABC, abstractmethod +from dataclasses import dataclass + + +@dataclass(frozen=True) +class EncoderCapabilities: + breaks: bool + scheduling: bool + normalization: bool + + +class PromptEncoder(ABC): + """Return (conditioning tensor, complete metadata) without retaining state. + + Implementations own tokenization, weighting, segmentation and the metadata + required by the consuming model. Nodes own parsing, caching and schedules. + """ + + name: str + capabilities: EncoderCapabilities + + @abstractmethod + def matches(self, clip, streams): + """Whether the encoder/model contract supports these token streams.""" + + def validate_request(self, text="", *, normalization=False, scheduled=False): + if re.search(r"\bBREAK\b", text) and not self.capabilities.breaks: + raise ValueError(f"{self.name} BREAK is not supported. Use an unsplit prompt.") + if scheduled and not self.capabilities.scheduling: + raise ValueError( + f"{self.name} scheduling/alternation is not supported. Use a static prompt." + ) + if normalization and not self.capabilities.normalization: + raise ValueError( + f"{self.name} normalization is not supported. Disable normalization." + ) + + def validate_schedule_metadata(self, metadata): + # The shared hook swaps the embedding tensor only. pooled_output is + # permitted for compatibility but stays pinned to the first prompt. + if set(metadata) - {"pooled_output"}: + raise ValueError( + "Prompt scheduling cannot swap this encoder's auxiliary " + "conditioning fields yet. Use a static prompt." + ) + + @abstractmethod + def encode(self, clip, text, normalization=False, debug=False): + """Encode text after validating supported operations; preserve metadata.""" diff --git a/encoding/clip.py b/encoding/clip.py new file mode 100644 index 0000000..a54f473 --- /dev/null +++ b/encoding/clip.py @@ -0,0 +1,213 @@ +"""CLIP-family chunk isolation and A1111 direct embedding weights.""" + +import logging +import re + +import torch +import comfy.model_management as model_management + +from .base import EncoderCapabilities, PromptEncoder + +logger = logging.getLogger("A1111PromptNode") + + +class CLIPEncoder(PromptEncoder): + capabilities = EncoderCapabilities(breaks=True, scheduling=True, normalization=True) + + def __init__(self, channels): + self.channels = channels + self.name = "SDXL CLIP-L/G" if channels == ("l", "g") else "SD1.5 CLIP" + + def matches(self, clip, streams): + return streams == set(self.channels) + + def _tokenize(self, clip, text): + tokens = clip.tokenize(text) + if not isinstance(tokens, dict) or not self.matches(clip, set(tokens)): + raise ValueError(f"{self.name} returned unexpected token streams.") + return tokens + + def encode(self, clip, text, normalization=False, debug=False): + self.validate_request(text, normalization=normalization) + segments = re.split(r"\s*\bBREAK\b\s*", text) + if len(segments) == 1: + tokens = self._tokenize(clip, text) + else: + tokens = None + for segment in segments: + if not segment.strip(): + continue + segment_tokens = self._tokenize(clip, segment.strip()) + if tokens is None: + tokens = segment_tokens + else: + tokens = {key: tokens[key] + segment_tokens[key] for key in self.channels} + if tokens is None: + tokens = self._tokenize(clip, "") + return self._encode_with_direct_scaling(clip, tokens, normalization, debug) + + def _encode_with_direct_scaling( + self, clip, tokens, normalization, debug=False + ): + """Encode with direct scaling (anti-burn) like A1111's EmphasisOriginalNoNorm.""" + clip.load_model() + clip.cond_stage_model.reset_clip_options() + + if clip.layer_idx is not None: + clip.cond_stage_model.set_clip_options({"layer": clip.layer_idx}) + + device = clip.patcher.load_device + clip.cond_stage_model.set_clip_options({"execution_device": device}) + + if self.channels == ("l", "g"): + weights_l = self._extract_weights(tokens.get("l", [])) + weights_g = self._extract_weights(tokens.get("g", [])) + tokens_neutral = tokens.copy() + tokens_neutral["l"] = self._neutralize_weights(tokens.get("l", [])) + tokens_neutral["g"] = self._neutralize_weights(tokens.get("g", [])) + else: + weights = self._extract_weights(tokens["l"]) + tokens_neutral = {"l": self._neutralize_weights(tokens["l"])} + + o = clip.cond_stage_model.encode_token_weights(tokens_neutral) + cond, pooled = o[:2] + + if self.channels == ("l", "g"): + cond = self._apply_direct_scaling_sdxl( + cond, weights_l, weights_g, normalization, debug + ) + else: + cond = self._apply_direct_scaling(cond, weights, normalization, debug) + + if pooled is not None: + pooled = pooled.to(model_management.intermediate_device()) + cond = cond.to(model_management.intermediate_device()) + + metadata = dict(o[2]) if len(o) > 2 else {} + metadata["pooled_output"] = pooled + return cond, metadata + + def _extract_weights(self, token_batches): + all_weights = [] + for batch in token_batches: + if not batch: + all_weights.append([]) + continue + batch_weights = [ + t[1] if hasattr(t, "__getitem__") and len(t) > 1 else 1.0 for t in batch + ] + all_weights.append(batch_weights) + return all_weights + + def _neutralize_weights(self, token_batches): + neutralized = [] + for batch in token_batches: + if not batch: + neutralized.append([]) + continue + new_batch = [ + (t[0], 1.0) if hasattr(t, "__getitem__") else (t, 1.0) for t in batch + ] + neutralized.append(new_batch) + return neutralized + + def _apply_direct_scaling(self, cond, weights, normalization, debug=False): + """ + Apply A1111-style direct scaling to embeddings. + + Uses proper tensor broadcasting like A1111's emphasis.py: + self.z = self.z * self.multipliers.reshape(self.multipliers.shape + (1,)).expand(self.z.shape) + + Note: weights is List[List[float]] where each inner list is a BREAK chunk. + cond shape is [batch, seq, dim] where seq = sum of all chunk lengths. + We need to flatten weights to match the sequence dimension. + """ + cond = cond.clone() + + # Flatten weights from all chunks into a single sequence + # weights is List[List[float]] - one list per BREAK chunk (or token batch) + flat_weights = [] + for chunk_weights in weights: + flat_weights.extend(chunk_weights) + + # Match to cond's sequence length + seq_len = cond.shape[1] + if len(flat_weights) < seq_len: + flat_weights = flat_weights + [1.0] * (seq_len - len(flat_weights)) + elif len(flat_weights) > seq_len: + flat_weights = flat_weights[:seq_len] + + # Create tensor matching cond's batch dimension [batch, seq] + batch_size = cond.shape[0] + multipliers = torch.tensor( + [flat_weights] * batch_size, device=cond.device, dtype=cond.dtype + ) + + # Debug: show weight statistics + if debug: + non_unity = [w for w in flat_weights if abs(w - 1.0) > 0.001] + if non_unity: + logger.info( + f"[A1111 Prompt] Emphasis weights applied: {len(non_unity)} tokens modified" + ) + logger.info( + f"[A1111 Prompt] Weight range: [{min(non_unity):.3f}, {max(non_unity):.3f}]" + ) + # Show first few modified weights with positions + weight_samples = [ + (i, w) for i, w in enumerate(flat_weights) if abs(w - 1.0) > 0.001 + ][:5] + logger.info(f"[A1111 Prompt] Sample weights: {weight_samples}") + else: + logger.info("[A1111 Prompt] No emphasis weights (all tokens = 1.0)") + + if normalization: + # A1111 EmphasisOriginal: preserve original mean + original_mean = cond.mean() + + # A1111-style broadcasting: expand [batch, seq] -> [batch, seq, dim] + multipliers = multipliers.unsqueeze(-1).expand_as(cond) + cond = cond * multipliers + + if normalization: + new_mean = cond.mean() + if new_mean.abs() > 1e-8: + scale_factor = original_mean / new_mean + cond = cond * scale_factor + if debug: + logger.info( + f"[A1111 Prompt] Normalization applied: scale={scale_factor.item():.4f}" + ) + + return cond + + def _apply_direct_scaling_sdxl( + self, cond, weights_l, weights_g, normalization, debug=False + ): + """SDXL output is [batch, seq, 2048] where 2048 = 768 (L) + 1280 (G).""" + if cond.shape[-1] != 2048: + raise ValueError( + f"SDXL CLIP-L/G requires 2048 embedding dimensions; got {cond.shape[-1]}. " + "This conditioning layout needs a separate encoder implementation." + ) + + if debug: + logger.info(f"[A1111 Prompt] SDXL scaling: cond shape {cond.shape}") + logger.info( + "[A1111 Prompt] Applying weights to CLIP-L (768 dims) AND CLIP-G (1280 dims)" + ) + + cond = cond.clone() + dim_l = 768 + cond_l = cond[:, :, :dim_l] + cond_g = cond[:, :, dim_l:] + + # A1111 scales BOTH CLIP-L and CLIP-G embeddings + if debug: + logger.info("[A1111 Prompt] Scaling CLIP-L...") + cond_l = self._apply_direct_scaling(cond_l, weights_l, normalization, debug) + if debug: + logger.info("[A1111 Prompt] Scaling CLIP-G...") + cond_g = self._apply_direct_scaling(cond_g, weights_g, normalization, debug) + + return torch.cat([cond_l, cond_g], dim=-1) diff --git a/nodes/negative_node.py b/nodes/negative_node.py index 40dc16e..c604d2c 100644 --- a/nodes/negative_node.py +++ b/nodes/negative_node.py @@ -1,9 +1,7 @@ """ A1111 Style Negative Prompt Node -Same encoding logic as A1111PromptNode but WITHOUT model input. -This prevents alternation hooks from being registered on the model, -which would break generation if used for negative conditioning. +Uses the shared encoding policy without attaching a sampling schedule hook. If scheduling/alternation syntax is detected, it uses the first step's prompt and logs an info message. @@ -12,7 +10,7 @@ prompt and logs an info message. import logging from ..parser import expand_wildcards, get_prompt_schedule from ..metadata import record_prompt_metadata -from .prompt_node import A1111PromptNode +from ..encoding import get_encoder, ANIMA_EMPHASIS_MODES logger = logging.getLogger("A1111PromptNode") @@ -21,9 +19,7 @@ class A1111PromptNegative: """ A1111 Style Negative Prompt Node. - Same encoding logic as A1111PromptNode but WITHOUT model input. - This prevents alternation hooks from being registered on the model, - which would break generation if used for negative conditioning. + Uses the same encoder/model contract as the positive node. If scheduling/alternation syntax is detected, it uses the first step's prompt and logs an info message. @@ -45,6 +41,7 @@ class A1111PromptNegative: "BOOLEAN", {"default": False, "label_on": "Enable", "label_off": "Disable"}, ), + "anima_emphasis": (list(ANIMA_EMPHASIS_MODES), {"default": "native", "tooltip": "Anima only: native weighting, stable single-pass weighting, or attention bias. Applies to this negative prompt only."}), }, "hidden": { "unique_id": "UNIQUE_ID", @@ -66,6 +63,7 @@ class A1111PromptNegative: debug=False, unique_id=None, extra_pnginfo=None, + anima_emphasis="native", ): """ Encode negative prompt. Uses first step only if scheduling is detected. @@ -74,6 +72,8 @@ class A1111PromptNegative: raise RuntimeError("ERROR: clip input is None") expanded_text = expand_wildcards(text) + encoder = get_encoder(clip, anima_emphasis) + encoder.validate_request(expanded_text, normalization=normalization) record_prompt_metadata( extra_pnginfo, unique_id, @@ -82,10 +82,6 @@ class A1111PromptNegative: expanded_text, ) - is_sdxl = hasattr(clip.cond_stage_model, "clip_l") and hasattr( - clip.cond_stage_model, "clip_g" - ) - # Parse the prompt schedule (use dummy steps as it's not supported for negative) schedule = get_prompt_schedule(expanded_text, steps=20) @@ -113,7 +109,7 @@ class A1111PromptNegative: else f"[A1111 Prompt Negative] Effective text: {expanded_text}" ) logger.info( - f"[A1111 Prompt Negative] Model: {'SDXL' if is_sdxl else 'SD1.5'}" + f"[A1111 Prompt Negative] Encoder: {encoder.name}" ) if len(schedule) > 1: logger.info( @@ -123,13 +119,9 @@ class A1111PromptNegative: f"[A1111 Prompt Negative] Resolved prompt: {prompt_text[:100]}..." ) - # Reuse the encoding helper from A1111PromptNode - encoder = A1111PromptNode() - cond, pooled = encoder._encode_with_break_isolation( - clip, prompt_text, normalization, is_sdxl, debug - ) + cond, metadata = encoder.encode(clip, prompt_text, normalization, debug) return { "ui": {"text": [expanded_text]}, - "result": ([[cond, {"pooled_output": pooled}]],), + "result": ([[cond, metadata]],), } diff --git a/nodes/prompt_node.py b/nodes/prompt_node.py index 139e4d1..8aa0d49 100644 --- a/nodes/prompt_node.py +++ b/nodes/prompt_node.py @@ -1,12 +1,8 @@ """ A1111 Style Prompt Node (Positive) -Implements A1111-style prompt handling: -- Hard chunking with BREAK isolation (no bleed between segments) -- Direct scaling (no burn) instead of interpolation -- BREAK support with proper token isolation -- Scheduling [from:to:when] (Switch, Add, Remove) -- Alternation [A|B] (true per-step switching via schedule) +Orchestrates wildcard expansion, schedule parsing, encoding caches and per-step +conditioning. Encoder/model behavior and supported operations live in encoding/. """ import torch @@ -17,6 +13,7 @@ import comfy.model_management as model_management from ..parser import expand_wildcards, get_prompt_schedule from ..hooks import create_step_schedule_cond from ..metadata import record_prompt_metadata +from ..encoding import get_encoder, ANIMA_EMPHASIS_MODES logger = logging.getLogger("A1111PromptNode") @@ -36,10 +33,10 @@ class A1111PromptNode: self._encoded_cache = {} self._max_cache_size = 50 # Limit unique prompts cached across executions - def _get_cache_key(self, clip, prompt_text, normalization): + def _get_cache_key(self, clip, prompt_text, normalization, anima_emphasis="native"): """Create a robust cache key including CLIP model and settings.""" # Using id(clip) as a proxy for the specific model instance/state - return (id(clip), normalization, prompt_text) + return (id(clip), normalization, prompt_text, anima_emphasis) @classmethod def INPUT_TYPES(s): @@ -57,6 +54,7 @@ class A1111PromptNode: "BOOLEAN", {"default": False, "label_on": "Enable", "label_off": "Disable"}, ), + "anima_emphasis": (list(ANIMA_EMPHASIS_MODES), {"default": "native", "tooltip": "Anima only: native weighting, stable single-pass weighting, or attention bias. Keep native for CLIP/SDXL."}), }, "hidden": { "prompt": "PROMPT", @@ -80,13 +78,13 @@ class A1111PromptNode: prompt=None, unique_id=None, extra_pnginfo=None, + anima_emphasis="native", ): """ Main encode function - A1111 style with step-based scheduling. Key features: - - BREAK segments are tokenized SEPARATELY for isolation - - Each BREAK segment becomes its own 77-token batch + - Encoding policy defines BREAK support and token layout - Scheduling/alternation evaluated per-step like A1111 - Steps auto-detected from downstream sampler/scheduler """ @@ -95,6 +93,10 @@ class A1111PromptNode: logger.info(f"[A1111 Prompt] Node ID: {unique_id}") expanded_text = expand_wildcards(text) + if clip is None: + raise RuntimeError("ERROR: clip input is None") + encoder = get_encoder(clip, anima_emphasis) + encoder.validate_request(expanded_text, normalization=normalization) record_prompt_metadata( extra_pnginfo, unique_id, @@ -134,6 +136,7 @@ class A1111PromptNode: ) if uses_step_syntax and steps is None: + encoder.validate_request(scheduled=True) raise StepCountRequiredError() # If steps is None and no step-based syntax, use a reasonable default for parsing @@ -144,13 +147,6 @@ class A1111PromptNode: f"[A1111 Prompt] Parse steps: {parse_steps} ({'auto-detected' if steps is not None else 'default for percentage-only'})" ) - if clip is None: - raise RuntimeError("ERROR: clip input is None") - - is_sdxl = hasattr(clip.cond_stage_model, "clip_l") and hasattr( - clip.cond_stage_model, "clip_g" - ) - if debug: logger.info( f"[A1111 Prompt] Input text: {text[:100]}..." @@ -163,7 +159,7 @@ class A1111PromptNode: if len(expanded_text) > 100 else f"[A1111 Prompt] Effective text: {expanded_text}" ) - logger.info(f"[A1111 Prompt] Model: {'SDXL' if is_sdxl else 'SD1.5'}") + logger.info(f"[A1111 Prompt] Encoder: {encoder.name}") logger.info( f"[A1111 Prompt] Normalization: {'ON' if normalization else 'OFF'}" ) @@ -174,31 +170,7 @@ class A1111PromptNode: if debug: logger.info(f"[A1111 Prompt] Schedule generated: {len(schedule)} segments") - # Check if we need alternation hook (if any step has alternation) - # The parser returns a flat list of (end_step, prompt) - # If the prompts are DIFFERENT for almost every step, it's likely alternation or a gradient - # We can just check if we have many changes. - - # Actually, let's just checking if the schedule is complex. - # But wait, [A|B] generates a schedule that changes every step. - # So we can detect if len(schedule) is roughly equal to steps? - # Or simpler: always use the hook if there's more than 1 item in schedule? - # NO. Standard scheduling [from:to:when] creates few segments. - # Alternation [A|B] creates 'steps' segments. - - # If we just implemented the hook for EVERYTHING, would that be bad? - # It adds a python call per step. Might be slightly slower. - # But it guarantees correctness. - # However, standard ComfyUI scheduling (TIMESTEP range) is more efficient for simple [from:to]. - # Let's support both. - - # Strategy: - # 1. Flatten schedule to per-step list of prompts (size = parse_steps) - # 2. Encode all unique prompts - # 3. Create a list of (cond, pooled) for each step - # 4. Attach hook with this list - - # First, expand schedule to full per-step list + # Expand the parsed schedule to one prompt per step. full_step_prompts = [""] * parse_steps prev_end = 0 for end_step, prompt in schedule: @@ -216,18 +188,8 @@ class A1111PromptNode: if full_step_prompts[i] != full_step_prompts[i - 1]: transitions += 1 - # If few transitions (like < 5?), we can use standard Comfy scheduling ranges for efficiency. - # If many transitions (alternation [A|B] = steps/2 transitions), use Hook. - # Let's say if transitions > 4, or if user asks for it (we can add a force/debug toggle later). - # For now, let's be aggressive and use Hook for > 2 transitions to be safe for [A|B|C]. - - # TODO: A1111 "Scheduling" [from:to:when] is also technically 'step based'. - # ComfyUI Percentages are sigma-based approximations. - # To get TRUE A1111 parity for even simple scheduling, we should use the Hook! - # Because ComfyUI's percent->sigma mapping is non-linear and might drift from A1111's linear step count. - # So, ALWAYS USE HOOK if there is ANY schedule? - has_schedule = len(schedule) > 1 + encoder.validate_request(scheduled=has_schedule) if debug: unique_prompts = list(set(full_step_prompts)) @@ -269,9 +231,9 @@ class A1111PromptNode: prompt_text = schedule[0][1] if schedule else "" # Check cache for static prompt too - cache_key = self._get_cache_key(clip, prompt_text, normalization) + cache_key = self._get_cache_key(clip, prompt_text, normalization, anima_emphasis) if cache_key in self._encoded_cache: - cond, pooled = self._encoded_cache[cache_key] + cond, metadata = self._encoded_cache[cache_key] if debug: logger.info("[A1111 Prompt] Found static prompt in cache") else: @@ -279,10 +241,8 @@ class A1111PromptNode: if len(self._encoded_cache) >= self._max_cache_size: self._encoded_cache.clear() - cond, pooled = self._encode_with_break_isolation( - clip, prompt_text, normalization, is_sdxl, debug - ) - self._encoded_cache[cache_key] = (cond, pooled) + cond, metadata = encoder.encode(clip, prompt_text, normalization, debug) + self._encoded_cache[cache_key] = (cond, metadata) if debug: logger.info( @@ -290,7 +250,7 @@ class A1111PromptNode: ) return { "ui": {"text": [expanded_text]}, - "result": ([[cond, {"pooled_output": pooled}]],), + "result": ([[cond, metadata.copy()]],), } # Use instance-level cache with size limit @@ -304,10 +264,12 @@ class A1111PromptNode: missing_keys = [] unique_keys = [] for p in unique_prompts: - key = self._get_cache_key(clip, p, normalization) + key = self._get_cache_key(clip, p, normalization, anima_emphasis) unique_keys.append(key) if key not in self._encoded_cache: missing_keys.append((key, p)) + else: + encoder.validate_schedule_metadata(self._encoded_cache[key][1]) # Ensure capacity for all new prompts before we start encoding # This prevents mid-execution clearing that would lose prompts for the current run @@ -333,13 +295,12 @@ class A1111PromptNode: logger.info( f"[A1111 Prompt] Encoding missing prompt {i + 1}/{len(missing_keys)}: {prompt_text[:50]}..." ) - cond, pooled = self._encode_with_break_isolation( - clip, prompt_text, normalization, is_sdxl, debug - ) - self._encoded_cache[key] = (cond, pooled) + cond, metadata = encoder.encode(clip, prompt_text, normalization, debug) + encoder.validate_schedule_metadata(metadata) + self._encoded_cache[key] = (cond, metadata) # Temporary local mapping for the current run (to help with padding) - # keys here are original prompt_texts, values are (cond, pooled) + # Keys are prompt texts; values contain the full encoder metadata. run_embeddings = {} for p, key in zip(unique_prompts, unique_keys): run_embeddings[p] = self._encoded_cache[key] @@ -352,7 +313,7 @@ class A1111PromptNode: # Pad only the embeddings used in this run to the same length padded_count = 0 for prompt_text in unique_prompts: - cond, pooled = run_embeddings[prompt_text] + cond, metadata = run_embeddings[prompt_text] if cond.shape[1] < max_seq_len: pad_size = max_seq_len - cond.shape[1] padding = torch.zeros( @@ -364,7 +325,7 @@ class A1111PromptNode: ) cond = torch.cat([cond, padding], dim=1) # Update run mapping with padded version - run_embeddings[prompt_text] = (cond, pooled) + run_embeddings[prompt_text] = (cond, metadata) padded_count += 1 if debug: @@ -385,7 +346,8 @@ class A1111PromptNode: device = model_management.intermediate_device() step_embeddings = [] for prompt in full_step_prompts: - cond, pooled = run_embeddings[prompt] + cond, metadata = run_embeddings[prompt] + pooled = metadata.get("pooled_output") # Ensure they are on the right device before sampling starts cond = cond.to(device) if pooled is not None: @@ -599,235 +561,3 @@ class A1111PromptNode: ) logger.warning(f"[A1111 Prompt] Visited nodes: {sorted(visited)}") return None - - def _encode_with_break_isolation( - self, clip, prompt_text, normalization, is_sdxl, debug - ): - """ - Encode prompt with BREAK isolation. - - Each BREAK segment is tokenized separately, creating isolated - 77-token batches that prevent bleed-over between segments. - """ - # Split by BREAK - break_segments = re.split(r"\s*\bBREAK\b\s*", prompt_text) - - if debug: - logger.info( - f"[A1111 Prompt] Encoding {len(break_segments)} BREAK segment(s)" - ) - - if len(break_segments) == 1: - # No BREAK, encode normally - tokens = clip.tokenize(break_segments[0]) - return self._encode_with_direct_scaling( - clip, tokens, normalization, is_sdxl, debug - ) - - # Multiple segments - tokenize each separately for isolation - all_tokens = None - - for i, segment in enumerate(break_segments): - segment = segment.strip() - if not segment: - continue - - segment_tokens = clip.tokenize(segment) - - if all_tokens is None: - all_tokens = segment_tokens - else: - # Merge token batches - each segment becomes its own batch - for key in segment_tokens: - if key in all_tokens: - all_tokens[key] = all_tokens[key] + segment_tokens[key] - else: - all_tokens[key] = segment_tokens[key] - - if debug: - logger.info( - f" Segment {i}: '{segment[:40]}...' -> {len(segment_tokens.get('l', segment_tokens.get('g', [])))} batch(es)" - ) - - if all_tokens is None: - # Empty prompt - all_tokens = clip.tokenize("") - - return self._encode_with_direct_scaling( - clip, all_tokens, normalization, is_sdxl, debug - ) - - def _encode_with_direct_scaling( - self, clip, tokens, normalization, is_sdxl, debug=False - ): - """Encode with direct scaling (anti-burn) like A1111's EmphasisOriginalNoNorm.""" - clip.load_model() - clip.cond_stage_model.reset_clip_options() - - if clip.layer_idx is not None: - clip.cond_stage_model.set_clip_options({"layer": clip.layer_idx}) - - device = clip.patcher.load_device - clip.cond_stage_model.set_clip_options({"execution_device": device}) - - if is_sdxl: - weights_l = self._extract_weights(tokens.get("l", [])) - weights_g = self._extract_weights(tokens.get("g", [])) - tokens_neutral = tokens.copy() - tokens_neutral["l"] = self._neutralize_weights(tokens.get("l", [])) - tokens_neutral["g"] = self._neutralize_weights(tokens.get("g", [])) - else: - weights = self._extract_weights( - tokens.get("l", tokens) if isinstance(tokens, dict) else tokens - ) - tokens_neutral = self._neutralize_weights( - tokens.get("l", tokens) if isinstance(tokens, dict) else tokens - ) - if isinstance(tokens, dict): - tokens_neutral_dict = tokens.copy() - tokens_neutral_dict["l"] = tokens_neutral - tokens_neutral = tokens_neutral_dict - - o = clip.cond_stage_model.encode_token_weights(tokens_neutral) - cond, pooled = o[:2] - - if is_sdxl: - cond = self._apply_direct_scaling_sdxl( - cond, weights_l, weights_g, normalization, debug - ) - else: - cond = self._apply_direct_scaling(cond, weights, normalization, debug) - - if pooled is not None: - pooled = pooled.to(model_management.intermediate_device()) - cond = cond.to(model_management.intermediate_device()) - - return cond, pooled - - def _extract_weights(self, token_batches): - all_weights = [] - for batch in token_batches: - if not batch: - all_weights.append([]) - continue - batch_weights = [ - t[1] if hasattr(t, "__getitem__") and len(t) > 1 else 1.0 for t in batch - ] - all_weights.append(batch_weights) - return all_weights - - def _neutralize_weights(self, token_batches): - neutralized = [] - for batch in token_batches: - if not batch: - neutralized.append([]) - continue - new_batch = [ - (t[0], 1.0) if hasattr(t, "__getitem__") else (t, 1.0) for t in batch - ] - neutralized.append(new_batch) - return neutralized - - def _apply_direct_scaling(self, cond, weights, normalization, debug=False): - """ - Apply A1111-style direct scaling to embeddings. - - Uses proper tensor broadcasting like A1111's emphasis.py: - self.z = self.z * self.multipliers.reshape(self.multipliers.shape + (1,)).expand(self.z.shape) - - Note: weights is List[List[float]] where each inner list is a BREAK chunk. - cond shape is [batch, seq, dim] where seq = sum of all chunk lengths. - We need to flatten weights to match the sequence dimension. - """ - cond = cond.clone() - - # Flatten weights from all chunks into a single sequence - # weights is List[List[float]] - one list per BREAK chunk (or token batch) - flat_weights = [] - for chunk_weights in weights: - flat_weights.extend(chunk_weights) - - # Match to cond's sequence length - seq_len = cond.shape[1] - if len(flat_weights) < seq_len: - flat_weights = flat_weights + [1.0] * (seq_len - len(flat_weights)) - elif len(flat_weights) > seq_len: - flat_weights = flat_weights[:seq_len] - - # Create tensor matching cond's batch dimension [batch, seq] - batch_size = cond.shape[0] - multipliers = torch.tensor( - [flat_weights] * batch_size, device=cond.device, dtype=cond.dtype - ) - - # Debug: show weight statistics - if debug: - non_unity = [w for w in flat_weights if abs(w - 1.0) > 0.001] - if non_unity: - logger.info( - f"[A1111 Prompt] Emphasis weights applied: {len(non_unity)} tokens modified" - ) - logger.info( - f"[A1111 Prompt] Weight range: [{min(non_unity):.3f}, {max(non_unity):.3f}]" - ) - # Show first few modified weights with positions - weight_samples = [ - (i, w) for i, w in enumerate(flat_weights) if abs(w - 1.0) > 0.001 - ][:5] - logger.info(f"[A1111 Prompt] Sample weights: {weight_samples}") - else: - logger.info("[A1111 Prompt] No emphasis weights (all tokens = 1.0)") - - if normalization: - # A1111 EmphasisOriginal: preserve original mean - original_mean = cond.mean() - - # A1111-style broadcasting: expand [batch, seq] -> [batch, seq, dim] - multipliers = multipliers.unsqueeze(-1).expand_as(cond) - cond = cond * multipliers - - if normalization: - new_mean = cond.mean() - if new_mean.abs() > 1e-8: - scale_factor = original_mean / new_mean - cond = cond * scale_factor - if debug: - logger.info( - f"[A1111 Prompt] Normalization applied: scale={scale_factor.item():.4f}" - ) - - return cond - - def _apply_direct_scaling_sdxl( - self, cond, weights_l, weights_g, normalization, debug=False - ): - """SDXL output is [batch, seq, 2048] where 2048 = 768 (L) + 1280 (G).""" - if cond.shape[-1] != 2048: - # Fallback for non-SDXL models (e.g. SD3) that might trigger is_sdxl check - logger.warning( - f"[A1111 Prompt] Unexpected embedding dimension {cond.shape[-1]} (expected 2048 for SDXL). Applying symmetric scaling." - ) - return self._apply_direct_scaling( - cond, weights_g if weights_g else weights_l, normalization, debug - ) - - if debug: - logger.info(f"[A1111 Prompt] SDXL scaling: cond shape {cond.shape}") - logger.info( - "[A1111 Prompt] Applying weights to CLIP-L (768 dims) AND CLIP-G (1280 dims)" - ) - - cond = cond.clone() - dim_l = 768 - cond_l = cond[:, :, :dim_l] - cond_g = cond[:, :, dim_l:] - - # A1111 scales BOTH CLIP-L and CLIP-G embeddings - if debug: - logger.info("[A1111 Prompt] Scaling CLIP-L...") - cond_l = self._apply_direct_scaling(cond_l, weights_l, normalization, debug) - if debug: - logger.info("[A1111 Prompt] Scaling CLIP-G...") - cond_g = self._apply_direct_scaling(cond_g, weights_g, normalization, debug) - - return torch.cat([cond_l, cond_g], dim=-1) diff --git a/pyproject.toml b/pyproject.toml index 64d7cfd..ee093d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,12 +18,14 @@ DisplayName = "A1111 Prompt Node" Icon = "" includes = [ "nodes/", + "encoding/", "parser/", "api/", "js/", "data/", "__init__.py", "hooks.py", + "metadata.py", "LICENSE", "README.md", "requirements.txt" diff --git a/research/anima_attention_bias.py b/research/anima_attention_bias.py new file mode 100644 index 0000000..f9f6e1f --- /dev/null +++ b/research/anima_attention_bias.py @@ -0,0 +1,53 @@ +"""Experimental attention-odds weighting after Anima's key normalization.""" + +from contextlib import contextmanager + + +def biased_attention(original, weights): + import torch + import torch.nn.functional as F + + weights = torch.as_tensor(weights, dtype=torch.float32) + if weights.ndim != 1 or not bool(torch.isfinite(weights).all()) or not bool((weights > 0).all()): + raise ValueError("Attention-bias weights must be a finite positive vector.") + + def apply(q, k, v, transformer_options=None): + options = transformer_options or {} + groups = options.get("cond_or_uncond") + if bool((weights == 1).all()): + return original(q, k, v, transformer_options=options) + if not groups or any(group not in (0, 1) for group in groups) or q.shape[0] % len(groups): + raise ValueError("Cannot identify positive/negative attention batch groups.") + if 0 not in groups: + return original(q, k, v, transformer_options=options) + if k.shape[1] != len(weights): + raise ValueError("Attention weights do not match the prepared context length.") + batch_per_group = q.shape[0] // len(groups) + bias = torch.zeros((q.shape[0], 1, 1, k.shape[1]), device=q.device, dtype=q.dtype) + log_weights = weights.to(device=q.device, dtype=q.dtype).log() + for index, group in enumerate(groups): + if group == 0: + bias[index * batch_per_group:(index + 1) * batch_per_group] = log_weights + # Mixed batches use this research SDPA backend for both branches. + # Negative rows have zero bias, but bitwise parity with a different + # optimized original backend is not asserted. + output = F.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=bias + ) + return output.transpose(1, 2).flatten(2) + + return apply + + +@contextmanager +def temporary_attention_bias(model, weights): + originals = [] + try: + for block in model.model.diffusion_model.blocks: + attention = block.cross_attn + originals.append((attention, attention.attn_op)) + attention.attn_op = biased_attention(attention.attn_op, weights) + yield + finally: + for attention, original in originals: + attention.attn_op = original diff --git a/tests/comfy_stubs.py b/tests/comfy_stubs.py new file mode 100644 index 0000000..548c057 --- /dev/null +++ b/tests/comfy_stubs.py @@ -0,0 +1,47 @@ +"""Minimal ComfyUI dependency setup for isolated node/encoder unit tests.""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +mock_model_management = MagicMock() +mock_model_management.intermediate_device.return_value = "cpu" + +mock_comfy = MagicMock() +mock_comfy.model_management = mock_model_management + +mock_hooks = MagicMock() + + +class DummyHook: + def __init__(self, *args, **kwargs): + self.transformers_dict = kwargs.get("transformers_dict") + self.hook_scope = kwargs.get("hook_scope") + + def clone(self): + return self + + +class DummyHookGroup: + def add(self, hook): + self.hook = hook + + +mock_hooks.TransformerOptionsHook = DummyHook +mock_hooks.HookGroup = DummyHookGroup +mock_hooks.EnumHookScope = MagicMock() +mock_hooks.set_hooks_for_conditioning = ( + lambda conditioning, hooks=None, append_hooks=True: conditioning +) +mock_comfy.hooks = mock_hooks + +server = MagicMock() +server.PromptServer.instance.routes.post = lambda path: (lambda fn: fn) +server.PromptServer.instance.routes.get = lambda path: (lambda fn: fn) + +sys.modules["comfy"] = mock_comfy +sys.modules["comfy.model_management"] = mock_model_management +sys.modules["comfy.hooks"] = mock_hooks +sys.modules["server"] = server diff --git a/tests/test_anima_attention_bias.py b/tests/test_anima_attention_bias.py new file mode 100644 index 0000000..6cc6f82 --- /dev/null +++ b/tests/test_anima_attention_bias.py @@ -0,0 +1,56 @@ +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from research.anima_attention_bias import biased_attention, temporary_attention_bias + + +class TestExperimentalAttentionBias(unittest.TestCase): + def test_probability_weighting_is_positive_only_in_both_cfg_orders(self): + for groups in ([0, 1], [1, 0]): + for batch_per_group in (1, 2): + batch = len(groups) * batch_per_group + q = torch.zeros(batch, 1, 1, 1) + k = torch.zeros(batch, 2, 1, 1) + v = torch.tensor([1., 0.]).reshape(1, 2, 1, 1).repeat(batch, 1, 1, 1) + for weight in (2., .5): + output = biased_attention(Mock(), [weight, 1])(q, k, v, {"cond_or_uncond": groups}) + expected = torch.tensor([weight / (weight + 1) if group == 0 else .5 for group in groups]).repeat_interleave(batch_per_group) + self.assertTrue(torch.allclose(output.flatten(), expected)) + + def test_neutral_weights_and_negative_only_batches_preserve_backend(self): + original = Mock(return_value="unchanged") + q = torch.zeros(1, 1, 1, 1) + for weights, options in (([1], {}), ([2], {"cond_or_uncond": [1]})): + self.assertEqual(biased_attention(original, weights)(q, q, q, options), "unchanged") + self.assertEqual(original.call_count, 2) + + def test_missing_cfg_groups_and_mismatched_length_are_rejected(self): + q = torch.zeros(1, 1, 1, 1) + op = biased_attention(Mock(), [2, 1]) + for options in ({}, {"cond_or_uncond": [0]}, {"cond_or_uncond": [2]}): + with self.assertRaises(ValueError): + op(q, q, q, options) + for weights in ([0], [-1], [float("nan")]): + with self.assertRaises(ValueError): + biased_attention(Mock(), weights) + + def test_temporary_override_is_restored_after_error(self): + original = Mock() + attention = SimpleNamespace(attn_op=original) + model = SimpleNamespace(model=SimpleNamespace(diffusion_model=SimpleNamespace( + blocks=[SimpleNamespace(cross_attn=attention)]))) + with self.assertRaises(RuntimeError): + with temporary_attention_bias(model, [2]): + self.assertIsNot(attention.attn_op, original) + raise RuntimeError("probe") + self.assertIs(attention.attn_op, original) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_anima_emphasis.py b/tests/test_anima_emphasis.py new file mode 100644 index 0000000..8f4da28 --- /dev/null +++ b/tests/test_anima_emphasis.py @@ -0,0 +1,48 @@ +"""Pure span-mapping checks for the production Anima tokenizer helper.""" + +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "encoding")) +from anima_tokenization import join_weighted_segments, weights_for_offsets + + +class TestExperimentalEmphasis(unittest.TestCase): + def test_complete_artist_span_includes_at_prefix_and_multiple_words(self): + text = join_weighted_segments([ + ("portrait, ", 1), ("@artist name", 1.5), (", smile", 1) + ]) + weights, crossings = weights_for_offsets( + text, [(0, 8), (8, 9), (10, 11), (11, 17), (18, 22), (22, 23), (24, 29), (0, 0)] + ) + self.assertEqual(weights, [1, 1, 1.5, 1.5, 1.5, 1, 1, 1]) + self.assertEqual(crossings, []) + + def test_partial_token_weights_are_character_averaged_and_reported(self): + text = join_weighted_segments([("ar", 2), ("tist", 1)]) + weights, crossings = weights_for_offsets(text, [(0, 6)]) + self.assertAlmostEqual(weights[0], 4 / 3) + self.assertEqual(crossings, [0]) + + def test_whitespace_does_not_dilute_weight_of_content_token(self): + text = join_weighted_segments([(" ", 1), ("artist", 2)]) + self.assertEqual(weights_for_offsets(text, [(0, 7)])[0], [2]) + + def test_special_and_whitespace_only_tokens_are_neutral(self): + text = join_weighted_segments([(" ", 2)]) + self.assertEqual(weights_for_offsets(text, [(0, 0), (0, 1)])[0], [1, 1]) + + def test_unicode_offsets_use_characters(self): + text = join_weighted_segments([("@猫", 1.5), (",", 1)]) + self.assertEqual(weights_for_offsets(text, [(0, 1), (1, 2), (2, 3)])[0], [1.5, 1.5, 1]) + + def test_bad_offsets_and_nonfinite_weights_are_rejected(self): + with self.assertRaises(ValueError): + join_weighted_segments([("cat", float("nan"))]) + with self.assertRaises(ValueError): + weights_for_offsets(join_weighted_segments([("cat", 1)]), [(0, 4)]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_anima_modes.py b/tests/test_anima_modes.py new file mode 100644 index 0000000..aa3dc72 --- /dev/null +++ b/tests/test_anima_modes.py @@ -0,0 +1,187 @@ +"""Production mode selection, cache separation and scoped attention behavior.""" + +import sys +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch +import torch.nn.functional as F +import comfy_stubs +from A1111_Prompt_Node.encoding import get_encoder +from A1111_Prompt_Node.encoding.anima_attention import AnimaAttention, attention_hooks +from A1111_Prompt_Node.encoding.anima_tokenization import join_weighted_segments, weights_for_offsets +from A1111_Prompt_Node.nodes.prompt_node import A1111PromptNode +from A1111_Prompt_Node.nodes.negative_node import A1111PromptNegative +from test_encoder_metadata import AnimaClip, Clip + + +class Anima: + pass + + +class Executor: + class_obj = Anima() + + def __init__(self, forward): + self.forward = forward + + def __call__(self, *args, **kwargs): + return self.forward(kwargs["transformer_options"]) + + +def attention(q, k, v, heads, mask=None, skip_reshape=False, **kwargs): + return F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + + +class TestAnimaModes(unittest.TestCase): + def test_optional_choices_and_defaults_on_both_nodes(self): + for node in (A1111PromptNode, A1111PromptNegative): + choices, settings = node.INPUT_TYPES()["optional"]["anima_emphasis"] + self.assertIsInstance(choices, list) + self.assertEqual(choices, ["native", "single_pass", "attention"]) + self.assertEqual(settings["default"], "native") + + def test_other_encoders_reject_non_native_choices(self): + for choice in ("single_pass", "attention", "typo"): + with self.assertRaises(ValueError): + get_encoder(Clip(), choice) + self.assertEqual(get_encoder(AnimaClip()).emphasis, "native") + + def test_modes_do_not_share_cached_encodings(self): + clip, node = AnimaClip(), A1111PromptNode() + clip.tokenizer = object() + tokens = clip.tokenize("(cat:1.3)") + # Stub only the real-tokenizer boundary; encode and cache are production. + with patch("A1111_Prompt_Node.encoding.anima_tokenization.single_pass_tokens", return_value=(tokens, {})): + outputs = {} + for mode in ("native", "single_pass", "attention", "native", "attention"): + outputs[mode] = node.encode(clip, "(cat:1.3)", anima_emphasis=mode)["result"][0][0][1] + self.assertEqual(len(clip.calls), 3) + self.assertEqual(len(node._encoded_cache), 3) + self.assertNotIn("hooks", outputs["native"]) + self.assertNotIn("hooks", outputs["single_pass"]) + self.assertIn("hooks", outputs["attention"]) + self.assertTrue(torch.all(outputs["attention"]["t5xxl_weights"] == 1)) + self.assertTrue(torch.any(outputs["single_pass"]["t5xxl_weights"] != 1)) + + def test_negative_gets_its_own_scoped_hook(self): + clip = AnimaClip() + clip.tokenizer = object() + tokens = clip.tokenize("(cat:1.3)") + with patch("A1111_Prompt_Node.encoding.anima_tokenization.single_pass_tokens", return_value=(tokens, {})): + metadata = A1111PromptNegative().encode(clip, "(cat:1.3)", anima_emphasis="attention")["result"][0][0][1] + self.assertEqual(metadata["hooks"].hook.hook_scope, comfy_stubs.mock_hooks.EnumHookScope.HookedOnly) + self.assertTrue(torch.all(metadata["t5xxl_weights"] == 1)) + + def test_neutral_attention_needs_no_hook_and_bad_weights_fail(self): + self.assertIsNone(attention_hooks([1, 1])) + for weights in ([0], [-1], [float("nan")], [float("inf")], [[1]]): + with self.assertRaises(ValueError): + attention_hooks(weights) + + def test_negative_preserves_first_step_schedule_fallback(self): + for mode in ("single_pass", "attention"): + clip = AnimaClip() + clip.tokenizer = object() + with patch("A1111_Prompt_Node.encoding.anima_tokenization.single_pass_tokens", return_value=(clip.tokenize("cat"), {})) as tokenize: + A1111PromptNegative().encode(clip, "[cat:dog:0.5]", anima_emphasis=mode) + tokenize.assert_called_once_with(clip.tokenizer, "cat") + + def test_production_span_mapping_keeps_special_tokens_neutral(self): + text = join_weighted_segments([("@猫", 3), (" ", 1), ("art", .5), ("ist", 1)]) + weights, crossings = weights_for_offsets(text, [(0, 0), (0, 1), (1, 2), (2, 9)]) + self.assertEqual(weights, [1, 3, 3, .75]) + self.assertEqual(crossings, [3]) + with self.assertRaises(ValueError): + weights_for_offsets(text, [(0, 100)]) + + def test_all_modes_keep_anima_capability_limits(self): + for mode in ("native", "single_pass", "attention"): + for node in (A1111PromptNode(), A1111PromptNegative()): + for text, kwargs in (("cat BREAK dog", {}), ("cat", {"normalization": True})): + with self.subTest(mode=mode, node=type(node).__name__, text=text): + with self.assertRaises(ValueError): + node.encode(AnimaClip(), text, anima_emphasis=mode, **kwargs) + with self.assertRaisesRegex(ValueError, "scheduling/alternation"): + A1111PromptNode().encode(AnimaClip(), "[cat:dog:0.5]", anima_emphasis=mode) + + +class TestScopedAnimaAttention(unittest.TestCase): + def setUp(self): + self.enterContext(patch.dict(sys.modules, {"comfy.ldm.anima.model": SimpleNamespace(Anima=Anima)})) + self.q = torch.zeros(2, 1, 1, 1) + self.k = torch.zeros(2, 1, 512, 1) + self.v = self.k.clone() + self.v[:, :, 0] = 1 + + def call(self, options, kind="attn2_patch", **kwargs): + for marker in options["patches"][kind]: + marker(self.q, self.k, self.v, extra_options=options) + return options["optimized_attention_override"](attention, self.q, self.k, self.v, 1, skip_reshape=True, **kwargs) + + def test_self_attention_with_same_context_length_is_not_biased(self): + def forward(options): + self_out = self.call(options, "attn1_patch") + cross_out = self.call(options) + return self_out, cross_out + plain, weighted = AnimaAttention([2])(Executor(forward), transformer_options={}) + self.assertTrue(torch.allclose(plain, torch.full_like(plain, 1 / 512))) + self.assertTrue(torch.allclose(weighted, torch.full_like(weighted, 2 / 513))) + + def test_deemphasis_masks_and_previous_override_are_composed(self): + previous = Mock(side_effect=lambda original, *args, **kwargs: original(*args, **kwargs)) + for mask in (torch.tensor([[True] * 511 + [False]]), torch.tensor([[0.] * 511 + [float("-inf")]], dtype=torch.float64)): + output = AnimaAttention([.5])(Executor(lambda opts: self.call(opts, mask=mask)), + transformer_options={"optimized_attention_override": previous}) + self.assertTrue(torch.allclose(output, torch.full_like(output, .5 / 510.5))) + self.assertEqual(previous.call_count, 2) + + def test_options_not_mutated_and_repeated_calls_do_not_accumulate(self): + original = {"patches": {"unrelated": [object()]}} + wrapper = AnimaAttention([3]) + first = wrapper(Executor(self.call), transformer_options=original) + second = wrapper(Executor(self.call), transformer_options=original) + self.assertTrue(torch.equal(first, second)) + self.assertEqual(set(original), {"patches"}) + self.assertEqual(set(original["patches"]), {"unrelated"}) + + def test_missing_patch_support_and_non_anima_fail(self): + wrapper = AnimaAttention([2]) + with self.assertRaisesRegex(ValueError, "does not support"): + wrapper(Executor(lambda options: None), transformer_options={}) + executor = Executor(self.call) + executor.class_obj = object() + with self.assertRaisesRegex(ValueError, "Anima diffusion"): + wrapper(executor, transformer_options={}) + + def test_repeated_or_missing_cross_calls_fail_closed(self): + wrapper = AnimaAttention([2]) + + def repeated(options): + self.call(options) + return options["optimized_attention_override"](attention, self.q, self.k, self.v, 1, skip_reshape=True) + + with self.assertRaisesRegex(ValueError, "repeated"): + wrapper(Executor(repeated), transformer_options={}) + + def missing(options): + self.call(options) + for marker in options["patches"]["attn2_patch"]: + marker(self.q, self.k, self.v, extra_options=options) + + with self.assertRaisesRegex(ValueError, "patch/call mismatch"): + wrapper(Executor(missing), transformer_options={}) + + def test_context_mismatch_and_nested_hooks_fail(self): + wrapper = AnimaAttention([2]) + with self.assertRaisesRegex(ValueError, "multiple"): + wrapper(Executor(self.call), transformer_options={"a1111_anima_attention": True}) + self.k = self.k[:, :, :511] + self.v = self.v[:, :, :511] + with self.assertRaisesRegex(ValueError, "context length"): + wrapper(Executor(self.call), transformer_options={}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_encoder_metadata.py b/tests/test_encoder_metadata.py new file mode 100644 index 0000000..4df2605 --- /dev/null +++ b/tests/test_encoder_metadata.py @@ -0,0 +1,254 @@ +"""Verify real node outputs with independent Qwen/T5 streams and CLIP weights.""" + +import re +import unittest +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +import comfy_stubs # noqa: F401 - install isolated ComfyUI dependencies +from A1111_Prompt_Node.nodes.prompt_node import A1111PromptNode +from A1111_Prompt_Node.nodes.negative_node import A1111PromptNegative +from A1111_Prompt_Node.encoding import get_encoder + + +class AnimaClip: + def __init__(self): + self.cond_stage_model = SimpleNamespace(clip_name="qwen3_06b") + self.calls = [] + + def tokenize(self, text): + words = re.sub(r"[():\d.]", "", text).split() or [""] + qwen = [(i + 10, 1.0) for i in range(len(words))] + # The T5 stream intentionally has a different length and weighting. + t5 = [(i + 40, 1.3 if ":1.3" in text and i == 1 else 1.0) + for i in range(len(words) + 2)] + return {"qwen3_06b": [qwen], "t5xxl": [t5]} + + def encode_from_tokens(self, tokens, return_dict=False): + assert return_dict + assert set(tokens) == {"qwen3_06b", "t5xxl"} + assert len(tokens["qwen3_06b"]) == len(tokens["t5xxl"]) == 1 + self.calls.append(tokens) + qwen = tokens["qwen3_06b"][0] + t5 = tokens["t5xxl"][0] + return { + "cond": torch.tensor([[[float(token[0])] * 3 for token in qwen]]), + "pooled_output": None, + "attention_mask": torch.ones(1, len(qwen)), + "t5xxl_ids": torch.tensor([token[0] for token in t5]), + "t5xxl_weights": torch.tensor([token[1] for token in t5]), + } + + +class Clip: + def __init__(self, sdxl=False): + self.layer_idx = None + self.patcher = SimpleNamespace(load_device="cpu") + self.tokens_seen = None + self.cond_stage_model = SimpleNamespace( + reset_clip_options=lambda: None, + set_clip_options=lambda options: None, + encode_token_weights=self.encode_token_weights, + ) + if sdxl: + self.cond_stage_model.clip_l = object() + self.cond_stage_model.clip_g = object() + + def load_model(self): + pass + + def tokenize(self, text): + return {"l": [[(1, 1.3), (2, 1.0)]]} + + def encode_token_weights(self, tokens): + self.tokens_seen = tokens + return torch.ones(1, 2, 3), None, {"attention_mask": torch.ones(1, 2)} + + +class XLClip(Clip): + def __init__(self): + super().__init__(sdxl=True) + + def tokenize(self, text): + return {"l": [[(1, 1.3), (2, 1.0)]], "g": [[(3, 0.7), (4, 1.0)]]} + + def encode_token_weights(self, tokens): + self.tokens_seen = tokens + return torch.ones(1, 2, 2048), torch.full((1, 1280), 2.0) + + +class ChunkedClip(Clip): + def tokenize(self, text): + token_id = {"left": 11, "right": 22, "": 0}[text] + return {"l": [[(token_id, 1.0)] * 77]} + + def encode_token_weights(self, tokens): + self.tokens_seen = tokens + flat = [token_id for batch in tokens["l"] for token_id, _ in batch] + return torch.tensor(flat, dtype=torch.float32).reshape(1, -1, 1), None + + +class TestEncoderMetadata(unittest.TestCase): + def test_selection_declares_family_capabilities(self): + clip = get_encoder(Clip()) + xl = get_encoder(XLClip()) + anima = get_encoder(AnimaClip()) + self.assertEqual(clip.name, "SD1.5 CLIP") + self.assertEqual(xl.name, "SDXL CLIP-L/G") + self.assertTrue(clip.capabilities.breaks) + self.assertTrue(xl.capabilities.scheduling) + self.assertFalse(anima.capabilities.breaks) + self.assertFalse(anima.capabilities.scheduling) + self.assertFalse(anima.capabilities.normalization) + + def test_unknown_and_mixed_streams_do_not_fall_back_to_clip(self): + for streams in ({"t5xxl": []}, {"l": [], "g": [], "t5xxl": []}): + clip = Clip() + clip.tokenize = lambda text: streams + with self.subTest(streams=streams): + for node in (A1111PromptNode(), A1111PromptNegative()): + with self.assertRaisesRegex(ValueError, "Unsupported text encoder"): + node.encode(clip, "cat") + self.assertIsNone(clip.tokens_seen) + + def test_anima_selection_requires_its_consuming_model_contract(self): + clip = AnimaClip() + clip.cond_stage_model.clip_name = "another_qwen_model" + with self.assertRaisesRegex(ValueError, "Unsupported text encoder"): + get_encoder(clip) + + def test_missing_consuming_model_contract_reports_unsupported_encoder(self): + clip = AnimaClip() + del clip.cond_stage_model + with self.assertRaisesRegex(ValueError, "Unsupported text encoder"): + get_encoder(clip) + + def test_clip_break_preserves_separate_full_batches_for_both_nodes(self): + for node in (A1111PromptNode(), A1111PromptNegative()): + clip = ChunkedClip() + cond, _ = node.encode(clip, "left BREAK right")["result"][0][0] + self.assertEqual([len(batch) for batch in clip.tokens_seen["l"]], [77, 77]) + self.assertTrue(torch.equal(cond[0, :77, 0], torch.full((77,), 11.0))) + self.assertTrue(torch.equal(cond[0, 77:, 0], torch.full((77,), 22.0))) + + def test_clip_empty_break_preserves_empty_prompt_fallback(self): + clip = ChunkedClip() + cond, _ = A1111PromptNode().encode(clip, " BREAK ")["result"][0][0] + self.assertEqual(cond.shape[1], 77) + self.assertTrue(torch.equal(cond, torch.zeros_like(cond))) + + def test_encoder_rechecks_streams_for_actual_prompt(self): + clip = Clip() + encoder = get_encoder(clip) + clip.tokenize = lambda text: {"t5xxl": []} + with self.assertRaisesRegex(ValueError, "unexpected token streams"): + encoder.encode(clip, "cat") + + def test_anima_static_positive_matches_core_including_extras(self): + clip = AnimaClip() + expected = clip.encode_from_tokens(clip.tokenize("a cat sitting"), return_dict=True) + cond, metadata = A1111PromptNode().encode(clip, "a cat sitting")["result"][0][0] + self.assertTrue(torch.equal(cond, expected["cond"])) + self.assertEqual(set(metadata), set(expected) - {"cond"}) + for key in ("attention_mask", "t5xxl_ids", "t5xxl_weights"): + self.assertTrue(torch.equal(metadata[key], expected[key]), key) + + def test_anima_negative_preserves_native_emphasis_channel(self): + clip = AnimaClip() + text = "a (cat:1.3) sitting" + expected = clip.encode_from_tokens(clip.tokenize(text), return_dict=True) + cond, metadata = A1111PromptNegative().encode(clip, text)["result"][0][0] + self.assertTrue(torch.equal(cond, expected["cond"])) + self.assertTrue(torch.equal(metadata["t5xxl_weights"], expected["t5xxl_weights"])) + self.assertAlmostEqual(metadata["t5xxl_weights"][1].item(), 1.3) + + def test_anima_break_is_explicitly_unsupported_for_both_nodes(self): + for node in (A1111PromptNode(), A1111PromptNegative()): + for text in ("short BREAK much longer text", " BREAK "): + clip = AnimaClip() + with self.subTest(node=type(node).__name__, text=text): + with self.assertRaisesRegex(ValueError, "Anima BREAK is not supported"): + node.encode(clip, text) + self.assertEqual(clip.calls, []) + + def test_anima_schedule_and_normalization_are_explicitly_unsupported(self): + for text, normalization, message in ( + ("[cat:dog:0.5]", False, "scheduling/alternation"), + ("[cat|dog]", False, "scheduling/alternation"), + ("cat", True, "normalization"), + ): + clip = AnimaClip() + with self.subTest(text=text), self.assertRaisesRegex(ValueError, message): + A1111PromptNode().encode( + clip, text, normalization=normalization, + prompt={"2": {"class_type": "KSampler", "inputs": { + "positive": ["1", 0], "steps": 4, + }}}, unique_id="1", + ) + self.assertEqual(clip.calls, []) + + def test_cached_metadata_dictionary_is_not_exposed_to_consumers(self): + node = A1111PromptNode() + clip = AnimaClip() + first = node.encode(clip, "cat")["result"][0][0] + del first[1]["t5xxl_ids"] + second = node.encode(clip, "cat")["result"][0][0] + self.assertIn("t5xxl_ids", second[1]) + self.assertEqual(len(clip.calls), 1) + + def test_clip_direct_scaling_and_auxiliary_metadata_are_preserved(self): + clip = Clip() + cond, metadata = A1111PromptNode().encode(clip, "cat")["result"][0][0] + self.assertTrue(torch.allclose(cond[0, 0], torch.full((3,), 1.3))) + self.assertTrue(torch.equal(cond[0, 1], torch.ones(3))) + self.assertEqual(clip.tokens_seen["l"], [[(1, 1.0), (2, 1.0)]]) + self.assertTrue(torch.equal(metadata["attention_mask"], torch.ones(1, 2))) + + def test_auxiliary_metadata_cannot_silently_enter_tensor_only_schedule(self): + node = A1111PromptNode() + clip = Clip() + encode = Mock(wraps=clip.cond_stage_model.encode_token_weights) + clip.cond_stage_model.encode_token_weights = encode + with self.assertRaisesRegex(ValueError, "auxiliary conditioning fields"): + node.encode(clip, "[cat:dog:0.5]") + self.assertEqual(encode.call_count, 1) + self.assertEqual(node._encoded_cache, {}) + + def test_cached_auxiliary_metadata_is_rejected_before_new_encoding(self): + node = A1111PromptNode() + clip = Clip() + node.encode(clip, "cat") + encode = Mock(wraps=clip.cond_stage_model.encode_token_weights) + clip.cond_stage_model.encode_token_weights = encode + with self.assertRaisesRegex(ValueError, "auxiliary conditioning fields"): + node.encode(clip, "[cat:dog:0.5]") + encode.assert_not_called() + + def test_rejected_prompts_do_not_write_resolved_prompt_metadata(self): + for node in (A1111PromptNode(), A1111PromptNegative()): + info = {} + with self.assertRaisesRegex(ValueError, "Anima BREAK"): + node.encode(AnimaClip(), "cat BREAK dog", unique_id="1", extra_pnginfo=info) + self.assertEqual(info, {}) + + def test_sdxl_keeps_separate_l_and_g_scaling_and_pooled_output(self): + clip = XLClip() + cond, metadata = A1111PromptNode().encode(clip, "cat")["result"][0][0] + self.assertTrue(torch.allclose(cond[0, 0, :768], torch.full((768,), 1.3))) + self.assertTrue(torch.allclose(cond[0, 0, 768:], torch.full((1280,), 0.7))) + self.assertTrue(torch.equal(cond[0, 1], torch.ones(2048))) + self.assertTrue(torch.equal(metadata["pooled_output"], torch.full((1, 1280), 2.0))) + for channel in ("l", "g"): + self.assertTrue(all(weight == 1.0 for _, weight in clip.tokens_seen[channel][0])) + + def test_other_dual_stream_layouts_are_not_scaled_as_sdxl(self): + clip = XLClip() + clip.cond_stage_model.encode_token_weights = lambda tokens: (torch.ones(1, 2, 4096), None) + with self.assertRaisesRegex(ValueError, "separate encoder implementation"): + A1111PromptNode().encode(clip, "cat") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_nodes_logic.py b/tests/test_nodes_logic.py index 79ec080..a66da31 100644 --- a/tests/test_nodes_logic.py +++ b/tests/test_nodes_logic.py @@ -1,52 +1,12 @@ import unittest import torch -# We extract the pure logic from prompt_node.py to test it in isolation -# without dealing with ComfyUI's complex package/import system. +import comfy_stubs # noqa: F401 - install isolated ComfyUI dependencies +from A1111_Prompt_Node.encoding.clip import CLIPEncoder - -def extract_weights(token_batches): - all_weights = [] - for batch in token_batches: - if not batch: - all_weights.append([]) - continue - batch_weights = [ - t[1] if hasattr(t, "__getitem__") and len(t) > 1 else 1.0 for t in batch - ] - all_weights.append(batch_weights) - return all_weights - - -def apply_direct_scaling(cond, weights, normalization): - cond = cond.clone() - flat_weights = [] - for chunk_weights in weights: - flat_weights.extend(chunk_weights) - - seq_len = cond.shape[1] - if len(flat_weights) < seq_len: - flat_weights = flat_weights + [1.0] * (seq_len - len(flat_weights)) - elif len(flat_weights) > seq_len: - flat_weights = flat_weights[:seq_len] - - batch_size = cond.shape[0] - multipliers = torch.tensor( - [flat_weights] * batch_size, device=cond.device, dtype=cond.dtype - ) - - if normalization: - original_mean = cond.mean() - - multipliers = multipliers.unsqueeze(-1).expand_as(cond) - cond = cond * multipliers - - if normalization: - new_mean = cond.mean() - if new_mean.abs() > 1e-8: - scale_factor = original_mean / new_mean - cond = cond * scale_factor - return cond +encoder = CLIPEncoder(("l",)) +extract_weights = encoder._extract_weights +apply_direct_scaling = encoder._apply_direct_scaling class TestNodesLogic(unittest.TestCase): diff --git a/tests/test_prompt_node_cache.py b/tests/test_prompt_node_cache.py index 86e2792..a84f6aa 100644 --- a/tests/test_prompt_node_cache.py +++ b/tests/test_prompt_node_cache.py @@ -1,85 +1,47 @@ -import sys import unittest -from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import patch import torch - -sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - -mock_model_management = MagicMock() -mock_model_management.intermediate_device.return_value = "cpu" - -mock_comfy = MagicMock() -mock_comfy.model_management = mock_model_management - -mock_hooks = MagicMock() - - -class DummyHook: - def __init__(self, *args, **kwargs): - pass - - def clone(self): - return self - - -class DummyHookGroup: - def add(self, hook): - self.hook = hook - - -mock_hooks.TransformerOptionsHook = DummyHook -mock_hooks.HookGroup = DummyHookGroup -mock_hooks.EnumHookScope = MagicMock() -mock_hooks.set_hooks_for_conditioning = ( - lambda conditioning, hooks=None, append_hooks=True: conditioning -) -mock_comfy.hooks = mock_hooks - -server = MagicMock() -server.PromptServer.instance.routes.post = lambda path: (lambda fn: fn) -server.PromptServer.instance.routes.get = lambda path: (lambda fn: fn) - -sys.modules["comfy"] = mock_comfy -sys.modules["comfy.model_management"] = mock_model_management -sys.modules["comfy.hooks"] = mock_hooks -sys.modules["server"] = server - +import comfy_stubs # noqa: F401 - install isolated ComfyUI dependencies from A1111_Prompt_Node.nodes.prompt_node import A1111PromptNode - - -class FakeCondStageModel: - pass +from A1111_Prompt_Node.encoding.base import EncoderCapabilities, PromptEncoder class FakeClip: - def __init__(self): - self.cond_stage_model = FakeCondStageModel() + pass + + +class CacheProbeEncoder(PromptEncoder): + name = "Cache probe" + capabilities = EncoderCapabilities(breaks=True, scheduling=True, normalization=True) + + def matches(self, clip, streams): + return True + + def encode(self, clip, prompt_text, normalization=False, debug=False): + seq_len = 2 if prompt_text == "short" else 5 + cond = torch.ones((1, seq_len, 3)) + pooled = torch.ones((1, 3)) + return cond, {"pooled_output": pooled} class CacheProbePromptNode(A1111PromptNode): def _get_downstream_steps(self, prompt, start_node_id, debug=False): return 2 - def _encode_with_break_isolation(self, clip, prompt_text, normalization, is_sdxl, debug): - seq_len = 2 if prompt_text == "short" else 5 - cond = torch.ones((1, seq_len, 3)) - pooled = torch.ones((1, 3)) - return cond, pooled - class TestPromptNodeCache(unittest.TestCase): def test_scheduled_padding_does_not_mutate_encoded_cache(self): node = CacheProbePromptNode() clip = FakeClip() - node.encode(clip, "[short:long:0.5]", prompt={}, unique_id="1") + with patch("A1111_Prompt_Node.nodes.prompt_node.get_encoder", return_value=CacheProbeEncoder()): + node.encode(clip, "[short:long:0.5]", prompt={}, unique_id="1") cached_shapes = { key[2]: cond.shape[1] - for key, (cond, _pooled) in node._encoded_cache.items() + for key, (cond, _metadata) in node._encoded_cache.items() } self.assertEqual(cached_shapes["short"], 2) self.assertEqual(cached_shapes["long"], 5)