373 lines
14 KiB
Python
373 lines
14 KiB
Python
"""
|
||
Conditioning Cache System for SamplerGridTester
|
||
Saves and loads encoded CLIP conditioning to avoid re-encoding prompts
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import torch
|
||
import hashlib
|
||
import base64
|
||
import numpy as np
|
||
from typing import Dict, Tuple, Optional, Any
|
||
|
||
|
||
class ConditioningCache:
|
||
"""
|
||
Manages persistent caching of encoded conditioning tensors.
|
||
|
||
The cache stores conditioning tensors as base64-encoded numpy arrays in JSON,
|
||
keyed by a hash of the prompt text + CLIP model state.
|
||
"""
|
||
|
||
def __init__(self, cache_dir: str, clip_hash: str = "unknown", enable_disk_cache: bool = True):
|
||
"""
|
||
Initialize the conditioning cache.
|
||
|
||
Args:
|
||
cache_dir: Directory to store cache files (e.g., benchmarks/session_name/)
|
||
clip_hash: Hash of the CLIP model to ensure compatibility
|
||
enable_disk_cache: Whether to save/load cache to/from disk
|
||
"""
|
||
self.cache_dir = cache_dir
|
||
self.clip_hash = clip_hash
|
||
self.cache_file = os.path.join(cache_dir, "conditioning_cache.json")
|
||
self.enable_disk_cache = enable_disk_cache
|
||
|
||
# Track current LoRA configuration for cache key generation
|
||
self.current_lora_config = ""
|
||
|
||
# Statistics - MUST be initialized BEFORE _load_from_disk()
|
||
self.stats = {
|
||
"hits": 0,
|
||
"misses": 0,
|
||
"saves": 0,
|
||
"loads": 0
|
||
}
|
||
|
||
# In-memory cache for current session
|
||
self.memory_cache = {
|
||
"positive": {},
|
||
"negative": {}
|
||
}
|
||
|
||
# Load existing cache from disk (uses self.stats) - only if enabled
|
||
if self.enable_disk_cache:
|
||
self.disk_cache = self._load_from_disk()
|
||
else:
|
||
self.disk_cache = {
|
||
"version": "1.0",
|
||
"clip_hash": self.clip_hash,
|
||
"entries": {}
|
||
}
|
||
print(f"[CondCache] ℹ️ Disk cache disabled - using memory-only cache for this session")
|
||
|
||
def _get_clip_hash(self, clip_model) -> str:
|
||
"""
|
||
Generate a hash of the CLIP model state to ensure cache validity.
|
||
|
||
Args:
|
||
clip_model: The CLIP model object
|
||
|
||
Returns:
|
||
MD5 hash string of the model
|
||
"""
|
||
try:
|
||
# Try to hash the model's state dict
|
||
if hasattr(clip_model, 'state_dict'):
|
||
state_dict = clip_model.state_dict()
|
||
# Just hash the keys and shapes, not full tensors (too slow)
|
||
model_signature = str([(k, tuple(v.shape)) for k, v in list(state_dict.items())[:10]])
|
||
elif hasattr(clip_model, 'cond_stage_model'):
|
||
# Try alternative attribute
|
||
model_signature = str(type(clip_model.cond_stage_model))
|
||
else:
|
||
model_signature = str(type(clip_model))
|
||
|
||
return hashlib.md5(model_signature.encode()).hexdigest()[:16]
|
||
except Exception as e:
|
||
print(f"[CondCache] Warning: Could not hash CLIP model: {e}")
|
||
return "unknown"
|
||
|
||
def _prompt_key(self, prompt_text: str) -> str:
|
||
"""
|
||
Generate a cache key from prompt text + LoRA configuration.
|
||
|
||
Args:
|
||
prompt_text: The prompt string
|
||
|
||
Returns:
|
||
MD5 hash of the prompt + CLIP hash + LoRA config
|
||
"""
|
||
# Include LoRA configuration in the cache key to avoid incorrect cache hits
|
||
# when using the same prompt with different LoRAs or strengths
|
||
combined = f"{self.clip_hash}:{self.current_lora_config}:{prompt_text}"
|
||
return hashlib.md5(combined.encode()).hexdigest()
|
||
|
||
def set_lora_config(self, lora_config: str):
|
||
"""
|
||
Set the current LoRA configuration for cache key generation.
|
||
|
||
This must be called whenever the LoRA configuration changes to ensure
|
||
that cached conditioning is only reused when the CLIP model is in the
|
||
same state (same LoRAs with same strengths).
|
||
|
||
Args:
|
||
lora_config: String describing the LoRA configuration, e.g.:
|
||
"None" or "lora1.safetensors:0.8:0.6,lora2.safetensors:1.0:1.0"
|
||
"""
|
||
self.current_lora_config = lora_config
|
||
|
||
def _conditioning_to_dict(self, conditioning) -> Dict[str, Any]:
|
||
"""
|
||
Convert conditioning tensor to JSON-serializable format.
|
||
|
||
Args:
|
||
conditioning: ComfyUI conditioning format [[tensor, {pooled_output: tensor}]]
|
||
|
||
Returns:
|
||
Dictionary with base64-encoded tensors
|
||
"""
|
||
try:
|
||
cond_tensor = conditioning[0][0] # Main conditioning tensor
|
||
pooled_tensor = conditioning[0][1].get("pooled_output")
|
||
|
||
# Convert to numpy and then to base64
|
||
cond_np = cond_tensor.cpu().numpy()
|
||
cond_b64 = base64.b64encode(cond_np.tobytes()).decode('utf-8')
|
||
cond_shape = list(cond_np.shape)
|
||
cond_dtype = str(cond_np.dtype)
|
||
|
||
result = {
|
||
"cond": {
|
||
"data": cond_b64,
|
||
"shape": cond_shape,
|
||
"dtype": cond_dtype
|
||
}
|
||
}
|
||
|
||
if pooled_tensor is not None:
|
||
pooled_np = pooled_tensor.cpu().numpy()
|
||
pooled_b64 = base64.b64encode(pooled_np.tobytes()).decode('utf-8')
|
||
pooled_shape = list(pooled_np.shape)
|
||
pooled_dtype = str(pooled_np.dtype)
|
||
|
||
result["pooled"] = {
|
||
"data": pooled_b64,
|
||
"shape": pooled_shape,
|
||
"dtype": pooled_dtype
|
||
}
|
||
|
||
return result
|
||
except Exception as e:
|
||
print(f"[CondCache] Error serializing conditioning: {e}")
|
||
return None
|
||
|
||
def _dict_to_conditioning(self, data: Dict[str, Any]):
|
||
"""
|
||
Convert JSON dict back to conditioning tensor.
|
||
|
||
Args:
|
||
data: Dictionary with base64-encoded tensors
|
||
|
||
Returns:
|
||
ComfyUI conditioning format [[tensor, {pooled_output: tensor}]]
|
||
"""
|
||
try:
|
||
# Reconstruct main conditioning tensor
|
||
cond_b64 = data["cond"]["data"]
|
||
cond_shape = tuple(data["cond"]["shape"])
|
||
cond_dtype = np.dtype(data["cond"]["dtype"])
|
||
|
||
cond_bytes = base64.b64decode(cond_b64)
|
||
cond_np = np.frombuffer(cond_bytes, dtype=cond_dtype).reshape(cond_shape)
|
||
cond_tensor = torch.from_numpy(cond_np)
|
||
|
||
# Reconstruct pooled output if present
|
||
pooled_tensor = None
|
||
if "pooled" in data:
|
||
pooled_b64 = data["pooled"]["data"]
|
||
pooled_shape = tuple(data["pooled"]["shape"])
|
||
pooled_dtype = np.dtype(data["pooled"]["dtype"])
|
||
|
||
pooled_bytes = base64.b64decode(pooled_b64)
|
||
pooled_np = np.frombuffer(pooled_bytes, dtype=pooled_dtype).reshape(pooled_shape)
|
||
pooled_tensor = torch.from_numpy(pooled_np)
|
||
|
||
# Return in ComfyUI format
|
||
return [[cond_tensor, {"pooled_output": pooled_tensor}]]
|
||
except Exception as e:
|
||
print(f"[CondCache] Error deserializing conditioning: {e}")
|
||
return None
|
||
|
||
def _load_from_disk(self) -> Dict[str, Any]:
|
||
"""
|
||
Load cache from disk.
|
||
|
||
Returns:
|
||
Dictionary containing cached conditioning data
|
||
"""
|
||
if not os.path.exists(self.cache_file):
|
||
return {
|
||
"version": "1.0",
|
||
"clip_hash": self.clip_hash,
|
||
"entries": {}
|
||
}
|
||
|
||
try:
|
||
with open(self.cache_file, "r") as f:
|
||
data = json.load(f)
|
||
|
||
# Validate CLIP hash matches
|
||
if data.get("clip_hash") != self.clip_hash:
|
||
print(f"[CondCache] ⚠️ CLIP model changed (cache: {data.get('clip_hash')}, current: {self.clip_hash})")
|
||
print(f"[CondCache] 🔄 Invalidating old cache...")
|
||
return {
|
||
"version": "1.0",
|
||
"clip_hash": self.clip_hash,
|
||
"entries": {}
|
||
}
|
||
|
||
num_entries = len(data.get("entries", {}))
|
||
print(f"[CondCache] 📂 Loaded {num_entries} cached conditioning entries")
|
||
self.stats["loads"] = num_entries
|
||
|
||
# Debug: Show a sample of what's in the cache
|
||
if num_entries > 0:
|
||
sample_keys = list(data.get("entries", {}).keys())[:3]
|
||
print(f"[CondCache] 🔍 Sample cache keys: {sample_keys}")
|
||
|
||
return data
|
||
except Exception as e:
|
||
print(f"[CondCache] ⚠️ Error loading cache: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return {
|
||
"version": "1.0",
|
||
"clip_hash": self.clip_hash,
|
||
"entries": {}
|
||
}
|
||
|
||
def _save_to_disk(self):
|
||
"""Save cache to disk."""
|
||
try:
|
||
os.makedirs(self.cache_dir, exist_ok=True)
|
||
|
||
with open(self.cache_file, "w") as f:
|
||
json.dump(self.disk_cache, f, indent=2)
|
||
|
||
num_entries = len(self.disk_cache.get("entries", {}))
|
||
print(f"[CondCache] 💾 Saved {num_entries} conditioning entries to cache")
|
||
except Exception as e:
|
||
print(f"[CondCache] ⚠️ Error saving cache: {e}")
|
||
|
||
def get(self, prompt_text: str, prompt_type: str = "positive", debug: bool = False) -> Optional[Any]:
|
||
"""
|
||
Get cached conditioning for a prompt.
|
||
|
||
Args:
|
||
prompt_text: The prompt string
|
||
prompt_type: "positive" or "negative"
|
||
debug: Enable debug logging
|
||
|
||
Returns:
|
||
Conditioning tensor or None if not cached
|
||
"""
|
||
# Check memory cache first (uses full text as key)
|
||
if prompt_text in self.memory_cache[prompt_type]:
|
||
self.stats["hits"] += 1
|
||
if debug:
|
||
print(f"[CondCache] ✅ Memory cache hit for: {prompt_text[:50]}...")
|
||
return self.memory_cache[prompt_type][prompt_text]
|
||
|
||
# Check disk cache (uses hash as key) - only if enabled
|
||
if not self.enable_disk_cache:
|
||
self.stats["misses"] += 1
|
||
if debug:
|
||
print(f"[CondCache] ❌ Cache miss (disk cache disabled)")
|
||
return None
|
||
|
||
key = self._prompt_key(prompt_text)
|
||
if debug:
|
||
print(f"[CondCache] 🔍 Looking for key: {key}")
|
||
print(f"[CondCache] 🔍 Prompt: {prompt_text[:80]}...")
|
||
|
||
if key in self.disk_cache.get("entries", {}):
|
||
entry = self.disk_cache["entries"][key]
|
||
conditioning = self._dict_to_conditioning(entry["conditioning"])
|
||
|
||
if conditioning is not None:
|
||
# Store in memory cache for fast access (using full text)
|
||
self.memory_cache[prompt_type][prompt_text] = conditioning
|
||
self.stats["hits"] += 1
|
||
if debug:
|
||
print(f"[CondCache] ✅ Disk cache hit!")
|
||
return conditioning
|
||
|
||
self.stats["misses"] += 1
|
||
if debug:
|
||
print(f"[CondCache] ❌ Cache miss")
|
||
return None
|
||
|
||
def set(self, prompt_text: str, conditioning, prompt_type: str = "positive"):
|
||
"""
|
||
Cache conditioning for a prompt.
|
||
|
||
Args:
|
||
prompt_text: The prompt string
|
||
conditioning: The conditioning tensor
|
||
prompt_type: "positive" or "negative"
|
||
"""
|
||
# Store in memory cache (always enabled)
|
||
self.memory_cache[prompt_type][prompt_text] = conditioning
|
||
|
||
# Store in disk cache (only if enabled)
|
||
if not self.enable_disk_cache:
|
||
return
|
||
|
||
key = self._prompt_key(prompt_text)
|
||
serialized = self._conditioning_to_dict(conditioning)
|
||
|
||
if serialized is not None:
|
||
if "entries" not in self.disk_cache:
|
||
self.disk_cache["entries"] = {}
|
||
|
||
self.disk_cache["entries"][key] = {
|
||
"prompt": prompt_text[:200], # Truncate for readability
|
||
"type": prompt_type,
|
||
"conditioning": serialized
|
||
}
|
||
|
||
self.stats["saves"] += 1
|
||
|
||
def save(self):
|
||
"""Save the cache to disk (only if disk cache is enabled)."""
|
||
if self.enable_disk_cache:
|
||
self._save_to_disk()
|
||
else:
|
||
print(f"[CondCache] ℹ️ Skipping disk save (disk cache disabled)")
|
||
|
||
def get_stats(self) -> Dict[str, int]:
|
||
"""Get cache statistics."""
|
||
hit_rate = 0
|
||
if (self.stats["hits"] + self.stats["misses"]) > 0:
|
||
hit_rate = (self.stats["hits"] / (self.stats["hits"] + self.stats["misses"])) * 100
|
||
|
||
return {
|
||
**self.stats,
|
||
"hit_rate": round(hit_rate, 1),
|
||
"total_cached": len(self.disk_cache.get("entries", {}))
|
||
}
|
||
|
||
def print_stats(self):
|
||
"""Print cache statistics."""
|
||
stats = self.get_stats()
|
||
print(f"\n{'='*80}")
|
||
print(f"[CondCache] 📊 CONDITIONING CACHE STATISTICS")
|
||
print(f"{'='*80}")
|
||
print(f" Hits: {stats['hits']}")
|
||
print(f" Misses: {stats['misses']}")
|
||
print(f" Hit Rate: {stats['hit_rate']}%")
|
||
print(f" Saved: {stats['saves']}")
|
||
print(f" Total Cached: {stats['total_cached']}")
|
||
print(f"{'='*80}\n") |