google-labs-jules[bot]
f482ab1fbd
1. Fix Normalization Bug: - The previous implementation normalized each token immediately after scaling, effectively cancelling the emphasis weight. - Updated logic to calculate global mean before scaling and restore it after scaling all tokens, correctly implementing A1111's 'EmphasisOriginal' behavior. 2. Improve Model Compatibility: - Modified token handling to preserve all keys in the `tokens` dictionary (e.g., `t5xxl` for Flux/SD3) instead of reconstructing a partial dictionary. This prevents silent failures for non-SDXL/SD1.5 models. - Added a safeguard in `_apply_direct_scaling_sdxl` to check if the embedding dimension matches the expected SDXL size (2048). If not (e.g., SD3), it falls back to generic scaling and logs a warning instead of producing incorrect output. 3. Note on hooks.py: - The `hooks.py` file remains unused as per current implementation status (placeholder for future sampler integration).
326 lines
12 KiB
Python
326 lines
12 KiB
Python
"""
|
|
A1111 Style Prompt Node (God Node)
|
|
|
|
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)
|
|
"""
|
|
|
|
import torch
|
|
import logging
|
|
import re
|
|
import comfy.model_management as model_management
|
|
from .parser import get_prompt_schedule, HAS_LARK
|
|
|
|
logger = logging.getLogger("A1111PromptNode")
|
|
|
|
|
|
class A1111PromptNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
},
|
|
"optional": {
|
|
"steps": (
|
|
"INT",
|
|
{
|
|
"default": 20,
|
|
"min": 1,
|
|
"max": 1000,
|
|
"tooltip": "Total sampling steps (for step-based scheduling like [thing:10])",
|
|
},
|
|
),
|
|
"normalization": (
|
|
"BOOLEAN",
|
|
{"default": False, "label_on": "Enable", "label_off": "Disable"},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
RETURN_NAMES = ("conditioning",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "conditioning/advanced"
|
|
|
|
def encode(self, clip, text, steps=20, normalization=False):
|
|
"""
|
|
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
|
|
- Scheduling/alternation evaluated per-step like A1111
|
|
"""
|
|
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"
|
|
)
|
|
|
|
logger.info(f"[God Node] ========== DEBUG START ==========")
|
|
logger.info(
|
|
f"[God Node] Input text: {text[:100]}..."
|
|
if len(text) > 100
|
|
else f"[God Node] Input text: {text}"
|
|
)
|
|
logger.info(f"[God Node] Model: {'SDXL' if is_sdxl else 'SD1.5'}")
|
|
logger.info(f"[God Node] Steps: {steps}")
|
|
logger.info(f"[God Node] Normalization: {'ON' if normalization else 'OFF'}")
|
|
logger.info(f"[God Node] Has Lark: {HAS_LARK}")
|
|
|
|
# Use A1111-style schedule generation
|
|
schedule = get_prompt_schedule(text, steps)
|
|
|
|
logger.info(f"[God Node] Raw schedule ({len(schedule)} entries):")
|
|
for end_step, prompt_text in schedule[:10]:
|
|
preview = prompt_text[:60] + "..." if len(prompt_text) > 60 else prompt_text
|
|
logger.info(f" Step {end_step}: '{preview}'")
|
|
|
|
# Group consecutive identical prompts
|
|
grouped_schedule = self._group_schedule(schedule, steps)
|
|
|
|
logger.info(f"[God Node] Grouped schedule ({len(grouped_schedule)} ranges):")
|
|
for start_pct, end_pct, prompt_text in grouped_schedule:
|
|
preview = prompt_text[:60] + "..." if len(prompt_text) > 60 else prompt_text
|
|
logger.info(f" {start_pct:.1%}-{end_pct:.1%}: '{preview}'")
|
|
|
|
# Collect unique prompts
|
|
unique_prompts = list(set(prompt for _, _, prompt in grouped_schedule))
|
|
logger.info(f"[God Node] Unique prompts to encode: {len(unique_prompts)}")
|
|
|
|
logger.info(f"[God Node] ========== DEBUG END ==========")
|
|
|
|
# Encode all unique prompts WITH BREAK ISOLATION
|
|
encoded_cache = {}
|
|
for prompt_text in unique_prompts:
|
|
if prompt_text in encoded_cache:
|
|
continue
|
|
|
|
# Handle BREAK segments with proper isolation
|
|
cond, pooled = self._encode_with_break_isolation(
|
|
clip, prompt_text, normalization, is_sdxl
|
|
)
|
|
encoded_cache[prompt_text] = (cond, pooled)
|
|
|
|
# Build final conditioning schedule
|
|
final_conditionings = []
|
|
for start_pct, end_pct, prompt_text in grouped_schedule:
|
|
if prompt_text not in encoded_cache:
|
|
continue
|
|
|
|
cond, pooled = encoded_cache[prompt_text]
|
|
|
|
cond_dict = {
|
|
"start_percent": start_pct,
|
|
"end_percent": end_pct,
|
|
}
|
|
if pooled is not None:
|
|
cond_dict["pooled_output"] = pooled
|
|
|
|
final_conditionings.append([cond, cond_dict])
|
|
|
|
if not final_conditionings:
|
|
tokens = clip.tokenize("")
|
|
result = clip.encode_from_tokens_scheduled(tokens)
|
|
return (result,)
|
|
|
|
return (final_conditionings,)
|
|
|
|
def _encode_with_break_isolation(self, clip, prompt_text, normalization, is_sdxl):
|
|
"""
|
|
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)
|
|
|
|
logger.info(f"[God Node] 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
|
|
)
|
|
|
|
# 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]
|
|
|
|
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
|
|
)
|
|
|
|
def _group_schedule(self, schedule, steps):
|
|
"""
|
|
Convert A1111-style schedule [(end_step, text), ...]
|
|
to ComfyUI-style [(start_pct, end_pct, text), ...].
|
|
"""
|
|
if not schedule:
|
|
return []
|
|
|
|
groups = []
|
|
prev_end_step = 0
|
|
|
|
for end_step, prompt_text in schedule:
|
|
start_pct = prev_end_step / steps
|
|
end_pct = end_step / steps
|
|
|
|
if groups and groups[-1][2] == prompt_text:
|
|
groups[-1] = (groups[-1][0], end_pct, prompt_text)
|
|
else:
|
|
groups.append((start_pct, end_pct, prompt_text))
|
|
|
|
prev_end_step = end_step
|
|
|
|
return groups
|
|
|
|
def _encode_with_direct_scaling(self, clip, tokens, normalization, is_sdxl):
|
|
"""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
|
|
)
|
|
else:
|
|
cond = self._apply_direct_scaling(cond, weights, normalization)
|
|
|
|
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):
|
|
"""Apply A1111-style direct scaling to embeddings."""
|
|
cond = cond.clone()
|
|
|
|
if normalization:
|
|
mean_before = cond.mean()
|
|
|
|
for batch_idx, batch_weights in enumerate(weights):
|
|
if batch_idx >= cond.shape[0]:
|
|
break
|
|
for token_idx, weight in enumerate(batch_weights):
|
|
if token_idx >= cond.shape[1]:
|
|
break
|
|
if weight != 1.0:
|
|
cond[batch_idx, token_idx] *= weight
|
|
|
|
if normalization:
|
|
mean_after = cond.mean()
|
|
if mean_after != 0:
|
|
cond *= mean_before / mean_after
|
|
|
|
return cond
|
|
|
|
def _apply_direct_scaling_sdxl(self, cond, weights_l, weights_g, normalization):
|
|
"""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"[God Node] Unexpected embedding dimension {cond.shape[-1]} (expected 2048 for SDXL). Applying symmetric scaling.")
|
|
# We don't know where the split is, so we just apply weights_g (usually the main one) or weights_l?
|
|
# Or effectively average them?
|
|
# Safest fallback: use 'g' weights if available, as that's usually the primary semantic encoder in ComfyUI flows
|
|
return self._apply_direct_scaling(cond, weights_g if weights_g else weights_l, normalization)
|
|
|
|
cond = cond.clone()
|
|
dim_l = 768
|
|
cond_l = cond[:, :, :dim_l]
|
|
cond_g = cond[:, :, dim_l:]
|
|
|
|
cond_l = self._apply_direct_scaling(cond_l, weights_l, normalization)
|
|
cond_g = self._apply_direct_scaling(cond_g, weights_g, normalization)
|
|
|
|
return torch.cat([cond_l, cond_g], dim=-1)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"A1111Prompt": A1111PromptNode}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"A1111Prompt": "A1111 Style Prompt (God Node)"}
|