379 lines
15 KiB
Python
379 lines
15 KiB
Python
"""
|
|
Batch CLIP Encoding Module
|
|
Handles efficient batch encoding of prompts with caching.
|
|
Supports prompt combinators: AND, CAT, AVG(), BREAK (native).
|
|
"""
|
|
|
|
import re
|
|
import torch
|
|
import gc
|
|
import comfy.model_management
|
|
|
|
|
|
# ============================================================
|
|
# === PROMPT COMBINATOR ENCODING (AND / CAT / AVG / BREAK) ===
|
|
# ============================================================
|
|
|
|
def _encode_single(clip_model, text, clip_skip=0):
|
|
"""
|
|
Encode a single text segment. BREAK is handled natively by clip.tokenize().
|
|
|
|
Args:
|
|
clip_model: CLIP model instance
|
|
text: Prompt text (may contain BREAK keyword)
|
|
clip_skip: CLIP layer skip value
|
|
|
|
Returns:
|
|
Conditioning list: [[cond_tensor, pooled_dict]]
|
|
"""
|
|
original_layer = None
|
|
if clip_skip != 0 and hasattr(clip_model.cond_stage_model, 'clip_layer'):
|
|
original_layer = clip_model.cond_stage_model.clip_layer
|
|
clip_model.cond_stage_model.set_clip_options({"layer": clip_skip})
|
|
|
|
tokens = clip_model.tokenize(text) # BREAK handled natively here
|
|
pooled_dict = clip_model.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
|
|
cond = pooled_dict.pop("cond")
|
|
result = [[cond, pooled_dict]]
|
|
|
|
if original_layer is not None:
|
|
clip_model.cond_stage_model.set_clip_options({"layer": original_layer})
|
|
|
|
return result
|
|
|
|
|
|
def _parse_and_segments(prompt):
|
|
"""
|
|
Split prompt on AND keyword. Returns list of (text, weight) tuples.
|
|
|
|
Examples:
|
|
"a cat AND a dog" -> [("a cat", 1.0), ("a dog", 1.0)]
|
|
"a cat :1.5 AND a dog :0.5" -> [("a cat", 1.5), ("a dog", 0.5)]
|
|
"a cat AND a dog :0" -> [("a cat", 1.0), ("a dog", 0.0)]
|
|
"""
|
|
parts = re.split(r'\s+AND\s+', prompt)
|
|
segments = []
|
|
for part in parts:
|
|
part = part.strip()
|
|
# Check for trailing weight like "cat :1.5" (space before colon distinguishes from (word:weight) syntax)
|
|
weight_match = re.match(r'^(.*?)\s+:(\d+\.?\d*)\s*$', part)
|
|
if weight_match:
|
|
segments.append((weight_match.group(1).strip(), float(weight_match.group(2))))
|
|
else:
|
|
segments.append((part, 1.0))
|
|
return segments
|
|
|
|
|
|
def _parse_avg_segments(text):
|
|
"""
|
|
Split text on AVG(weight) markers. Returns list of (text, avg_weight) tuples.
|
|
The first segment has weight=None (it's the base). Subsequent segments have their AVG weight.
|
|
|
|
Examples:
|
|
"realistic AVG(0.3) anime" -> [("realistic", None), ("anime", 0.3)]
|
|
"a AVG() b" -> [("a", None), ("b", 0.5)] (default weight 0.5)
|
|
"no avg here" -> [("no avg here", None)]
|
|
"""
|
|
parts = re.split(r'\s*AVG\((\d*\.?\d*)\)\s*', text)
|
|
if len(parts) == 1:
|
|
return [(text, None)] # No AVG found
|
|
|
|
segments = []
|
|
for i in range(0, len(parts), 2):
|
|
segment_text = parts[i].strip() if i < len(parts) else ""
|
|
if not segment_text:
|
|
continue
|
|
if i == 0:
|
|
# First segment is the base — no weight
|
|
segments.append((segment_text, None))
|
|
else:
|
|
# Subsequent segments get the weight from the preceding capture group
|
|
weight_str = parts[i - 1] if i - 1 < len(parts) else ""
|
|
weight = float(weight_str) if weight_str else 0.5
|
|
segments.append((segment_text, weight))
|
|
|
|
return segments if segments else [(text, None)]
|
|
|
|
|
|
def _parse_cat_segments(text):
|
|
"""
|
|
Split text on CAT keyword. Returns list of text segments.
|
|
|
|
Example:
|
|
"a cat CAT in a garden" -> ["a cat", "in a garden"]
|
|
"""
|
|
return [s.strip() for s in re.split(r'\s+CAT\s+', text) if s.strip()]
|
|
|
|
|
|
def _conditioning_combine(cond1, cond2):
|
|
"""
|
|
ConditioningCombine: list concatenation.
|
|
Each conditioning entry becomes a separate guidance signal during sampling.
|
|
Same as ComfyUI's ConditioningCombine node (nodes.py).
|
|
"""
|
|
return cond1 + cond2
|
|
|
|
|
|
def _conditioning_concat(cond_to, cond_from):
|
|
"""
|
|
ConditioningConcat: tensor concatenation along dim 1.
|
|
Extends the token context window by appending from-conditioning tokens.
|
|
Same as ComfyUI's ConditioningConcat node (nodes.py).
|
|
"""
|
|
out = []
|
|
cond_from_tensor = cond_from[0][0]
|
|
for i in range(len(cond_to)):
|
|
t1 = cond_to[i][0]
|
|
tw = torch.cat((t1, cond_from_tensor), 1)
|
|
n = [tw, cond_to[i][1].copy()]
|
|
out.append(n)
|
|
return out
|
|
|
|
|
|
def _conditioning_average(cond_to, cond_from, strength):
|
|
"""
|
|
ConditioningAverage: weighted blend of conditioning tensors.
|
|
strength=1.0 means 100% cond_to, strength=0.0 means 100% cond_from.
|
|
Same as ComfyUI's ConditioningAverage node (nodes.py).
|
|
"""
|
|
out = []
|
|
cond_from_tensor = cond_from[0][0]
|
|
pooled_from = cond_from[0][1].get("pooled_output", None)
|
|
|
|
for i in range(len(cond_to)):
|
|
t1 = cond_to[i][0]
|
|
pooled_to = cond_to[i][1].get("pooled_output", pooled_from)
|
|
|
|
t0 = cond_from_tensor[:, :t1.shape[1]]
|
|
if t0.shape[1] < t1.shape[1]:
|
|
t0 = torch.cat([t0, torch.zeros((1, t1.shape[1] - t0.shape[1], t1.shape[2]))], dim=1)
|
|
|
|
tw = torch.mul(t1, strength) + torch.mul(t0, 1.0 - strength)
|
|
t_to = cond_to[i][1].copy()
|
|
|
|
if pooled_from is not None and pooled_to is not None:
|
|
t_to["pooled_output"] = torch.mul(pooled_to, strength) + torch.mul(pooled_from, 1.0 - strength)
|
|
elif pooled_from is not None:
|
|
t_to["pooled_output"] = pooled_from
|
|
|
|
out.append([tw, t_to])
|
|
return out
|
|
|
|
|
|
def encode_prompt_with_combinators(clip_model, prompt, clip_skip=0):
|
|
"""
|
|
Encode a prompt that may contain AND, CAT, or AVG() combinators.
|
|
|
|
Processing order (matching comfyui-prompt-control):
|
|
1. Split on AND (creates separate conditioning entries via ConditioningCombine)
|
|
2. Within each AND-segment, process AVG(weight) (weighted blend via ConditioningAverage)
|
|
3. Within each segment, process CAT (tensor concat via ConditioningConcat)
|
|
4. BREAK is handled natively by clip.tokenize() within each segment
|
|
|
|
Syntax:
|
|
- "prompt1 AND prompt2" → ConditioningCombine
|
|
- "prompt1 AND prompt2 :0.5" → AND with weight scaling
|
|
- "prompt1 CAT prompt2" → ConditioningConcat
|
|
- "prompt1 AVG(0.3) prompt2" → ConditioningAverage (0.3 = 30% prompt2)
|
|
- "prompt1 BREAK prompt2" → Separate 77-token chunks (native)
|
|
|
|
Args:
|
|
clip_model: CLIP model instance
|
|
prompt: Prompt text possibly containing combinators
|
|
clip_skip: CLIP layer skip value
|
|
|
|
Returns:
|
|
Conditioning list: [[cond_tensor, pooled_dict], ...]
|
|
"""
|
|
# Fast path: if no combinators present, use direct encoding
|
|
if ' AND ' not in prompt and ' CAT ' not in prompt and 'AVG(' not in prompt:
|
|
return _encode_single(clip_model, prompt, clip_skip)
|
|
|
|
# 1. Split on AND
|
|
and_segments = _parse_and_segments(prompt)
|
|
|
|
all_conditionings = []
|
|
for segment_text, and_weight in and_segments:
|
|
if and_weight == 0:
|
|
continue # Skip zero-weighted AND segments
|
|
|
|
# 2. Process AVG() within this AND-segment
|
|
avg_segments = _parse_avg_segments(segment_text)
|
|
|
|
if len(avg_segments) > 1:
|
|
# Has AVG — encode first segment as base, then blend with subsequent segments
|
|
result = _encode_single(clip_model, avg_segments[0][0], clip_skip)
|
|
for avg_text, avg_weight in avg_segments[1:]:
|
|
next_cond = _encode_single(clip_model, avg_text, clip_skip)
|
|
w = avg_weight if avg_weight is not None else 0.5
|
|
result = _conditioning_average(result, next_cond, w)
|
|
else:
|
|
# 3. Process CAT within this segment
|
|
cat_parts = _parse_cat_segments(segment_text)
|
|
if len(cat_parts) > 1:
|
|
result = _encode_single(clip_model, cat_parts[0], clip_skip)
|
|
for cat_text in cat_parts[1:]:
|
|
next_cond = _encode_single(clip_model, cat_text, clip_skip)
|
|
result = _conditioning_concat(result, next_cond)
|
|
else:
|
|
result = _encode_single(clip_model, segment_text, clip_skip)
|
|
|
|
# Apply AND weight by scaling the conditioning tensor
|
|
if and_weight != 1.0:
|
|
result = [[torch.mul(r[0], and_weight), r[1]] for r in result]
|
|
|
|
all_conditionings.append(result)
|
|
|
|
# 4. Combine all AND segments via ConditioningCombine
|
|
if not all_conditionings:
|
|
return _encode_single(clip_model, "", clip_skip)
|
|
|
|
final = all_conditionings[0]
|
|
for additional in all_conditionings[1:]:
|
|
final = _conditioning_combine(final, additional)
|
|
|
|
return final
|
|
|
|
|
|
def batch_encode_with_cache(clip_model, prompts, cond_cache, prompt_type="positive", batch_size=64, clip_skip=0):
|
|
"""
|
|
Batch encode prompts while checking persistent cache first.
|
|
Only encodes prompts that aren't already cached.
|
|
|
|
Args:
|
|
clip_model: CLIP model to use for encoding
|
|
prompts: Set or list of prompt strings to encode
|
|
cond_cache: ConditioningCache instance
|
|
prompt_type: "positive" or "negative"
|
|
batch_size: Number of prompts to encode at once
|
|
clip_skip: Number of CLIP layers to skip from the end (0 = use last layer, -1 = skip 1 layer, -2 = skip 2 layers)
|
|
|
|
Returns:
|
|
dict: Mapping of prompt text to conditioning tensors
|
|
"""
|
|
results = {}
|
|
prompts_to_encode = []
|
|
|
|
|
|
# Check cache first
|
|
print(f"[GridTester] 🔍 Checking cache for {len(prompts)} {prompt_type} prompts...")
|
|
for prompt in prompts:
|
|
cached = cond_cache.get(prompt, prompt_type)
|
|
if cached is not None:
|
|
results[prompt] = cached
|
|
else:
|
|
prompts_to_encode.append(prompt)
|
|
|
|
cache_hits = len(results)
|
|
cache_misses = len(prompts_to_encode)
|
|
|
|
print(f"[GridTester] 📊 Cache: {cache_hits} hits, {cache_misses} misses")
|
|
|
|
# Batch encode uncached prompts
|
|
if prompts_to_encode:
|
|
clip_skip_msg = f" (clip_skip={clip_skip})" if clip_skip != 0 else ""
|
|
print(f"[GridTester] 🚀 Batch encoding {len(prompts_to_encode)} {prompt_type} prompts{clip_skip_msg} (batch_size={batch_size})")
|
|
|
|
# Force model to stay in VRAM
|
|
comfy.model_management.load_models_gpu([clip_model.patcher])
|
|
|
|
prompts_list = list(prompts_to_encode)
|
|
total_batches = (len(prompts_list) + batch_size - 1) // batch_size
|
|
|
|
with torch.no_grad():
|
|
for batch_idx in range(0, len(prompts_list), batch_size):
|
|
batch_prompts = prompts_list[batch_idx:batch_idx + batch_size]
|
|
current_batch = (batch_idx // batch_size) + 1
|
|
|
|
# Encode batch
|
|
for prompt in batch_prompts:
|
|
# Check for interrupt before each prompt encoding
|
|
if comfy.model_management.processing_interrupted():
|
|
print(f"\n[GridTester] 🛑 INTERRUPTED during {prompt_type} encoding - Stopping all encoding")
|
|
raise comfy.model_management.InterruptProcessingException()
|
|
|
|
try:
|
|
# Encode with combinator support (AND/CAT/AVG/BREAK)
|
|
conditioning = encode_prompt_with_combinators(clip_model, prompt, clip_skip)
|
|
results[prompt] = conditioning
|
|
cond_cache.set(prompt, conditioning, prompt_type)
|
|
|
|
except comfy.model_management.InterruptProcessingException:
|
|
print(f"\n[GridTester] 🛑 INTERRUPTED during {prompt_type} encoding - Stopping all encoding")
|
|
raise # Re-raise to stop all encoding
|
|
|
|
except Exception as e:
|
|
print(f"[GridTester] ⚠️ Failed to encode {prompt_type} prompt: {e}")
|
|
print(f"[GridTester] ⚠️ Prompt was: {prompt[:80]}...")
|
|
import traceback
|
|
traceback.print_exc()
|
|
# Store None so we can detect this downstream with a clear error
|
|
results[prompt] = None
|
|
|
|
# Progress
|
|
encoded_count = min(batch_idx + batch_size, len(prompts_list))
|
|
if current_batch % 5 == 0 or current_batch == total_batches:
|
|
print(f"[GridTester] Batch {current_batch}/{total_batches} | Encoded {encoded_count}/{len(prompts_list)}")
|
|
|
|
print(f"[GridTester] ✅ Batch encoding complete!")
|
|
|
|
return results
|
|
|
|
|
|
def batch_encode_prompts(patched_clip, unique_positives, unique_negatives, cond_cache, clip_skip=0, enable_disk_cache=True):
|
|
"""
|
|
Batch encode both positive and negative prompts.
|
|
|
|
Args:
|
|
patched_clip: Patched CLIP model
|
|
unique_positives: Set of unique positive prompts
|
|
unique_negatives: Set of unique negative prompts
|
|
cond_cache: ConditioningCache instance
|
|
clip_skip: Number of CLIP layers to skip from the end
|
|
enable_disk_cache: Whether to save cache to disk (passed through to cache)
|
|
|
|
Returns:
|
|
dict: conditioning_cache with "positive" and "negative" keys
|
|
"""
|
|
conditioning_cache = {"positive": {}, "negative": {}}
|
|
|
|
print(f"[GridTester] 🧠 Found {len(unique_positives)} unique positive prompts")
|
|
print(f"[GridTester] 🧠 Found {len(unique_negatives)} unique negative prompts")
|
|
|
|
# Use ComfyUI's model management to prevent memory leak warnings during batch encoding
|
|
# The model needs to stay loaded for the entire encoding loop
|
|
with torch.no_grad():
|
|
# Encode all positive prompts
|
|
print(f"[GridTester] 🧠 Encoding {len(unique_positives)} unique positive prompts...")
|
|
conditioning_cache["positive"] = batch_encode_with_cache(
|
|
patched_clip,
|
|
unique_positives,
|
|
cond_cache,
|
|
prompt_type="positive",
|
|
batch_size=16,
|
|
clip_skip=clip_skip
|
|
)
|
|
|
|
conditioning_cache["negative"] = batch_encode_with_cache(
|
|
patched_clip,
|
|
unique_negatives,
|
|
cond_cache,
|
|
prompt_type="negative",
|
|
batch_size=16,
|
|
clip_skip=clip_skip
|
|
)
|
|
|
|
# Save and print cache stats (only saves if disk cache is enabled)
|
|
if cond_cache is not None:
|
|
cond_cache.save()
|
|
# cond_cache.print_stats() # for debugging
|
|
|
|
# Final cleanup after all encoding is complete
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
print(f"[GridTester] ✅ All prompts pre-encoded!")
|
|
print(f"[GridTester] 💾 Cache size: {len(conditioning_cache['positive'])} positive, {len(conditioning_cache['negative'])} negative")
|
|
|
|
return conditioning_cache |