Wen Li
91ce7daa7b
Seven-mode guidance (CFG/APG/RAAG/S2/Z2/STG_A/STG_R), context-window sampling, segment prompt-travel, lazy streaming load, layered VRAM management.
117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
"""
|
|
Shared text-embedding disk cache and CLIP encoding helpers.
|
|
|
|
Used by both ``PromptEmbedding`` and ``PromptSchedule`` nodes to avoid
|
|
duplicated encoding of identical prompts across workflow runs.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import comfy.sd
|
|
import folder_paths
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Disk cache
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_CACHE_DIR = Path(__file__).parent.parent / "text_embed_cache"
|
|
_MAX_CACHE_FILES = 200 # evict oldest files when exceeded
|
|
|
|
|
|
def _ensure_cache_dir() -> Path:
|
|
_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
|
return _CACHE_DIR
|
|
|
|
|
|
def _evict_oldest_if_needed() -> None:
|
|
"""Remove oldest cache files if count exceeds _MAX_CACHE_FILES."""
|
|
try:
|
|
files = sorted(_CACHE_DIR.glob("*.pt"), key=lambda p: p.stat().st_mtime)
|
|
while len(files) > _MAX_CACHE_FILES:
|
|
oldest = files.pop(0)
|
|
oldest.unlink(missing_ok=True)
|
|
logger.debug(f"[BerniniR] Evicted old cache file: {oldest.name}")
|
|
except Exception:
|
|
pass # never fail on cache maintenance
|
|
|
|
|
|
def _cache_key(prompt: str, tag: str = "") -> str:
|
|
payload = prompt if not tag else f"{tag}::{prompt}"
|
|
return hashlib.md5(payload.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _cache_path(prompt: str, tag: str = "") -> Path:
|
|
return _ensure_cache_dir() / f"{_cache_key(prompt, tag)}.pt"
|
|
|
|
|
|
def _load_cached_conditioning(prompt: str, tag: str = "") -> Optional[list]:
|
|
path = _cache_path(prompt, tag)
|
|
if not path.exists():
|
|
return None
|
|
try:
|
|
data = torch.load(path, map_location="cpu", weights_only=False)
|
|
cond = data.get("cond")
|
|
pooled = data.get("pooled", {})
|
|
if cond is None:
|
|
return None
|
|
return [[cond, pooled]]
|
|
except Exception as e:
|
|
logger.warning(f"[BerniniR] Corrupt cache {path.name}: {e}")
|
|
try:
|
|
path.unlink(missing_ok=True)
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def _save_cached_conditioning(prompt: str, conditioning: list, tag: str = "") -> None:
|
|
if not conditioning or len(conditioning[0]) < 2:
|
|
return
|
|
cond, pooled = conditioning[0]
|
|
safe_pooled = {}
|
|
if isinstance(pooled, dict):
|
|
for k, v in pooled.items():
|
|
if isinstance(v, (torch.Tensor, str, int, float, bool)) or v is None:
|
|
safe_pooled[k] = v
|
|
try:
|
|
path = _cache_path(prompt, tag)
|
|
torch.save({"cond": cond.cpu(), "pooled": safe_pooled}, path)
|
|
logger.info(f"[BerniniR] Saved text embed cache: {path.name}")
|
|
_evict_oldest_if_needed()
|
|
except Exception as e:
|
|
logger.warning(f"[BerniniR] Failed to save cache: {e}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Text encoding
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _encode_text(clip, text: str):
|
|
"""Encode text using ComfyUI's CLIPTextEncode logic."""
|
|
tokens = clip.tokenize(text)
|
|
return clip.encode_from_tokens_scheduled(tokens)
|
|
|
|
|
|
def _load_clip_internal(clip_name: str, clip_type: str = "wan", device: str = "cpu"):
|
|
"""Load a Wan-series text encoder on demand."""
|
|
clip_path = folder_paths.get_full_path_or_raise("text_encoders", clip_name)
|
|
clip_type_enum = getattr(comfy.sd.CLIPType, clip_type.upper(), comfy.sd.CLIPType.WAN)
|
|
model_options = {}
|
|
if device == "cpu":
|
|
model_options["load_device"] = torch.device("cpu")
|
|
model_options["offload_device"] = torch.device("cpu")
|
|
return comfy.sd.load_clip(
|
|
ckpt_paths=[clip_path],
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
clip_type=clip_type_enum,
|
|
model_options=model_options,
|
|
)
|