nodes/Enferlain-ComfyUI-A1111-cond
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add model-aware encoders and selectable Anima emphasis modes

Browse Source
This commit is contained in:
imi
2026-11-09 01:04:41 +01:00
parent 28e6616028
commit af5b13addc
18 changed files with 1288 additions and 424 deletions
Download Patch File Download Diff File Expand all files Collapse all files
encoding/README.md
+47
View File
@@ -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.
encoding/__init__.py
+30
View File
@@ -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"]
encoding/anima.py
+46
View File
@@ -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
encoding/anima_attention.py
+104
View File
@@ -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
encoding/anima_tokenization.py
+81
View File
@@ -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,
}
encoding/base.py
+52
View File
@@ -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."""
encoding/clip.py
+213
View File
@@ -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)
nodes/negative_node.py
+10 -18
View File
@@ -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]],),
}
nodes/prompt_node.py
+32 -302
View File
@@ -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)
pyproject.toml
+2
View File
@@ -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"
research/anima_attention_bias.py
+53
View File
@@ -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
tests/comfy_stubs.py
+47
View File
@@ -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
tests/test_anima_attention_bias.py
+56
View File
@@ -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()
tests/test_anima_emphasis.py
+48
View File
@@ -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()
tests/test_anima_modes.py
+187
View File
@@ -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()
tests/test_encoder_metadata.py
+254
View File
@@ -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()
tests/test_nodes_logic.py
+5 -45
View File
@@ -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):
tests/test_prompt_node_cache.py
+21 -59
View File
@@ -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)