628 lines
20 KiB
Python
628 lines
20 KiB
Python
"""
|
|
CaptionForge Global Model Cache Manager
|
|
|
|
- CaptionForge
|
|
- This module is part of **CaptionForge**, a model-agnostic captioning
|
|
framework for ComfyUI developed by **J. L. Córdova**.
|
|
|
|
- Repository:
|
|
https://github.com/Damkohler/CaptionForge
|
|
|
|
- CaptionForge 1.0 uses independent Pass-A witnesses, text-LLM synthesis,
|
|
image-aware validation, SHORT/TAGGY formatting, and JSONL audit trails to
|
|
produce grounded LoRA dataset captions.
|
|
|
|
- Engine Purpose
|
|
- The **CaptionForge Global Model Cache Manager** provides shared,
|
|
process-local residency management for heavyweight CaptionForge model
|
|
bundles.
|
|
|
|
- It supports local ComfyUI workflows where Joy, Qwen, and other
|
|
captioning engines may be available, but limited VRAM makes simultaneous
|
|
residency impractical.
|
|
|
|
- The cache manager is responsible for:
|
|
• stable cache-key construction
|
|
• reuse of already-loaded model bundles
|
|
• registration of newly loaded model handles
|
|
• eviction of incompatible or excess resident models
|
|
• optional engine-provided unload callbacks
|
|
• Python reference cleanup
|
|
• best-effort CUDA allocator cleanup
|
|
• lightweight cache diagnostics
|
|
|
|
- Cached objects may include model instances, processor/tokenizer bundles,
|
|
engine-specific dataclass containers, or dictionaries/tuples wrapping
|
|
those objects.
|
|
|
|
- CaptionForge Pipeline Role
|
|
- This module is shared infrastructure, not a captioning pass.
|
|
|
|
- It supports CaptionForge Pass A caption witness engines by reducing
|
|
avoidable model reloads while also protecting constrained-VRAM systems
|
|
from accidental multi-model co-residency.
|
|
|
|
- It does not own model loading logic. Engines remain responsible for
|
|
constructing their own models and may register unload callbacks when they
|
|
need custom teardown behavior.
|
|
|
|
- Execution Model
|
|
- The cache is global, process-local, and protected by a re-entrant lock.
|
|
|
|
- The default policy keeps only one heavyweight CaptionForge model resident
|
|
at a time.
|
|
|
|
- Cache keys include fields that materially affect model residency:
|
|
• role
|
|
• family
|
|
• model path
|
|
• device
|
|
• quantization mode
|
|
• dtype
|
|
• revision
|
|
|
|
- CUDA cleanup is attempted only when PyTorch and CUDA are available.
|
|
|
|
- Design Philosophy
|
|
- CaptionForge treats model loading as an engine concern and model residency
|
|
as shared infrastructure.
|
|
|
|
- This keeps the cache small, explicit, and dependency-light while allowing
|
|
multiple CaptionForge engines to share one conservative VRAM policy.
|
|
|
|
- The module favors predictable behavior over aggressive automatic memory
|
|
management.
|
|
|
|
- Production Status
|
|
- Active CaptionForge 1.0 runtime support. Its conservative single-model
|
|
residency policy coordinates heavyweight local Hugging Face witnesses;
|
|
it does not own caption semantics.
|
|
|
|
- Attribution & License
|
|
- Concept and implementation by **J. L. Córdova**
|
|
with development assistance from **ChatGPT (OpenAI)**.
|
|
|
|
- Designed for use with:
|
|
https://github.com/comfyanonymous/ComfyUI
|
|
|
|
- Copyright (c) 2026 J. L. Córdova
|
|
|
|
- Released under the **MIT License**.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from ..captionforge_version import CAPTIONFORGE_VERSION
|
|
|
|
MANIFEST = {
|
|
"name": "CaptionForge Global Model Cache Manager",
|
|
"version": CAPTIONFORGE_VERSION,
|
|
"author": "J. L. Córdova",
|
|
"description": (
|
|
"Shared process-local cache manager for heavyweight CaptionForge model "
|
|
"bundles. Provides stable cache keys, model registration, reuse, "
|
|
"eviction, optional unload callbacks, CUDA cleanup, and lightweight "
|
|
"diagnostics while keeping engine-specific loading logic separate from "
|
|
"shared residency policy."
|
|
),
|
|
}
|
|
|
|
import gc
|
|
import os
|
|
import time
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Callable, Dict, Optional
|
|
|
|
|
|
try:
|
|
import torch
|
|
except Exception: # Keeps import safe if torch is unavailable during metadata inspection.
|
|
torch = None
|
|
|
|
|
|
UnloadFn = Callable[[Any], None]
|
|
|
|
|
|
# Default: only one heavyweight CaptionForge model resident.
|
|
_MAX_LOADED_MODELS = 1
|
|
|
|
# Safe eviction/cuda-cleanup defaults.
|
|
# Existing fast policies use synchronization + cleanup but no cooldown.
|
|
# Explicit *_safe policies add a short cooldown to reduce driver/watchdog stress
|
|
# during rapid Qwen/Joy/bitsandbytes model swaps.
|
|
def _env_float(name: str, default: float) -> float:
|
|
try:
|
|
return float(os.environ.get(name, default))
|
|
except Exception:
|
|
return float(default)
|
|
|
|
|
|
_SAFE_EVICTION_COOLDOWN_SEC = max(0.0, _env_float("CAPTIONFORGE_SAFE_EVICTION_COOLDOWN_SEC", 1.25))
|
|
_FAST_EVICTION_COOLDOWN_SEC = max(0.0, _env_float("CAPTIONFORGE_FAST_EVICTION_COOLDOWN_SEC", 0.0))
|
|
_VERBOSE_SAFE_EVICTION = os.environ.get("CAPTIONFORGE_CACHE_VERBOSE", "0").strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
# Global process-local cache.
|
|
_CACHE_LOCK = threading.RLock()
|
|
_MODEL_CACHE: Dict[str, "CacheEntry"] = {}
|
|
|
|
|
|
@dataclass
|
|
class CacheEntry:
|
|
key: str
|
|
obj: Any
|
|
family: str
|
|
model_path: str
|
|
device: str
|
|
quantization: str = "none"
|
|
role: str = "caption"
|
|
unload_fn: Optional[UnloadFn] = None
|
|
created_at: float = field(default_factory=time.time)
|
|
last_used_at: float = field(default_factory=time.time)
|
|
keep: bool = False
|
|
|
|
def touch(self) -> None:
|
|
self.last_used_at = time.time()
|
|
|
|
|
|
def normalize_cache_part(value: Any) -> str:
|
|
if value is None:
|
|
return "none"
|
|
text = str(value).strip()
|
|
return text if text else "none"
|
|
|
|
|
|
def make_cache_key(
|
|
*,
|
|
family: str,
|
|
model_path: str,
|
|
device: str = "cuda",
|
|
quantization: str = "none",
|
|
dtype: str = "auto",
|
|
revision: str = "default",
|
|
role: str = "caption",
|
|
) -> str:
|
|
"""
|
|
Build a stable cache key.
|
|
|
|
Include fields that change the actual resident model object:
|
|
- family: qwen, joy, cleanup_llm, validator, etc.
|
|
- model_path: local path or HF repo id
|
|
- device: cuda/cpu
|
|
- quantization: fp16, bf16, fp8, 8bit, none, etc.
|
|
- dtype: bf16/fp16/fp32/auto
|
|
- revision: optional model revision
|
|
- role: caption/cleanup/validator
|
|
"""
|
|
parts = [
|
|
normalize_cache_part(role).lower(),
|
|
normalize_cache_part(family).lower(),
|
|
normalize_cache_part(model_path),
|
|
normalize_cache_part(device).lower(),
|
|
normalize_cache_part(quantization).lower(),
|
|
normalize_cache_part(dtype).lower(),
|
|
normalize_cache_part(revision),
|
|
]
|
|
return "::".join(parts)
|
|
|
|
|
|
def set_max_loaded_models(max_loaded_models: int) -> None:
|
|
"""
|
|
Adjust global cache capacity.
|
|
|
|
Normally leave this at 1 for 16 GB VRAM systems.
|
|
"""
|
|
global _MAX_LOADED_MODELS
|
|
|
|
value = int(max_loaded_models)
|
|
if value < 0:
|
|
raise ValueError("max_loaded_models must be >= 0")
|
|
|
|
with _CACHE_LOCK:
|
|
_MAX_LOADED_MODELS = value
|
|
_enforce_capacity_locked()
|
|
|
|
|
|
def get_max_loaded_models() -> int:
|
|
with _CACHE_LOCK:
|
|
return _MAX_LOADED_MODELS
|
|
|
|
|
|
def cache_size(*, include_keep: bool = True, role: str | None = None) -> int:
|
|
"""Return the number of resident CaptionForge Python model cache entries.
|
|
|
|
This is intentionally lightweight and side-effect free. Ollama-backed nodes
|
|
use it to decide whether a Python/HF/torch model is still resident before
|
|
handing the workflow over to the local Ollama server.
|
|
|
|
Ollama models are not registered in this process-local cache, so this check
|
|
will not cause eviction between Ollama models.
|
|
"""
|
|
role_norm = normalize_cache_part(role).lower() if role is not None else None
|
|
|
|
with _CACHE_LOCK:
|
|
count = 0
|
|
for entry in _MODEL_CACHE.values():
|
|
if entry.keep and not include_keep:
|
|
continue
|
|
if role_norm is not None and normalize_cache_part(entry.role).lower() != role_norm:
|
|
continue
|
|
count += 1
|
|
return count
|
|
|
|
|
|
def get_cached_model(key: str) -> Optional[Any]:
|
|
"""
|
|
Return cached model bundle/object if present.
|
|
"""
|
|
with _CACHE_LOCK:
|
|
entry = _MODEL_CACHE.get(key)
|
|
if entry is None:
|
|
return None
|
|
|
|
entry.touch()
|
|
print(f"[CaptionForge Cache] Reusing cached model: {key}")
|
|
return entry.obj
|
|
|
|
|
|
def register_model(
|
|
key: str,
|
|
obj: Any,
|
|
*,
|
|
family: str,
|
|
model_path: str,
|
|
device: str = "cuda",
|
|
quantization: str = "none",
|
|
role: str = "caption",
|
|
unload_fn: Optional[UnloadFn] = None,
|
|
keep: bool = False,
|
|
) -> Any:
|
|
"""
|
|
Register a loaded model bundle/object.
|
|
|
|
The object can be:
|
|
- a model
|
|
- a tuple/dict containing model + processor/tokenizer
|
|
- an engine-specific bundle dataclass
|
|
"""
|
|
with _CACHE_LOCK:
|
|
existing = _MODEL_CACHE.get(key)
|
|
if existing is not None:
|
|
existing.obj = obj
|
|
existing.family = family
|
|
existing.model_path = model_path
|
|
existing.device = device
|
|
existing.quantization = quantization
|
|
existing.role = role
|
|
existing.unload_fn = unload_fn
|
|
existing.keep = keep
|
|
existing.touch()
|
|
else:
|
|
_MODEL_CACHE[key] = CacheEntry(
|
|
key=key,
|
|
obj=obj,
|
|
family=family,
|
|
model_path=model_path,
|
|
device=device,
|
|
quantization=quantization,
|
|
role=role,
|
|
unload_fn=unload_fn,
|
|
keep=keep,
|
|
)
|
|
|
|
print(f"[CaptionForge Cache] Registered model: {key}")
|
|
_enforce_capacity_locked(protected_key=key)
|
|
return obj
|
|
|
|
|
|
def prepare_for_model_load(
|
|
key: str,
|
|
*,
|
|
policy: str = "evict_other_caption_models",
|
|
family: Optional[str] = None,
|
|
role: str = "caption",
|
|
max_loaded_models: Optional[int] = None,
|
|
) -> None:
|
|
"""
|
|
Call this immediately before loading a heavyweight model.
|
|
|
|
Policies:
|
|
- keep_this_model:
|
|
Do not evict anything immediately. Capacity is still enforced after register_model.
|
|
- evict_other_caption_models:
|
|
Evict other caption-role models before loading this one, then run standard CUDA cleanup.
|
|
- evict_other_caption_models_safe:
|
|
Same as evict_other_caption_models, but with extra synchronization and a short cooldown.
|
|
- unload_after_run:
|
|
Same pre-load behavior as evict_other_caption_models; caller should evict this model after run.
|
|
- unload_after_run_safe:
|
|
Same pre-load behavior as evict_other_caption_models_safe; caller should safely evict after run.
|
|
- none:
|
|
Do nothing before load.
|
|
"""
|
|
policy = (policy or "").strip() or "evict_other_caption_models"
|
|
safe_mode = policy.endswith("_safe")
|
|
|
|
with _CACHE_LOCK:
|
|
if max_loaded_models is not None:
|
|
set_max_loaded_models(max_loaded_models)
|
|
|
|
if policy == "none":
|
|
return
|
|
|
|
if policy == "keep_this_model":
|
|
return
|
|
|
|
if policy in {
|
|
"evict_other_caption_models",
|
|
"evict_other_caption_models_safe",
|
|
"unload_after_run",
|
|
"unload_after_run_safe",
|
|
}:
|
|
did_evict = False
|
|
for existing_key, entry in list(_MODEL_CACHE.items()):
|
|
if existing_key == key:
|
|
continue
|
|
|
|
# Future cleanup/validator models can use role="cleanup" or "validator".
|
|
# For now, the default policy primarily protects against Qwen/Joy co-residency.
|
|
if entry.role == role or entry.role == "caption":
|
|
did_evict = _evict_locked(existing_key, reason=f"policy={policy}", safe=safe_mode) or did_evict
|
|
|
|
_cuda_cleanup(
|
|
reason=f"post_prepare_for_model_load:{policy}",
|
|
synchronize=True,
|
|
cooldown_sec=_SAFE_EVICTION_COOLDOWN_SEC if safe_mode and did_evict else _FAST_EVICTION_COOLDOWN_SEC,
|
|
)
|
|
return
|
|
|
|
raise ValueError(
|
|
f"Unknown CaptionForge cache policy: {policy!r}. "
|
|
"Expected one of: keep_this_model, evict_other_caption_models, "
|
|
"evict_other_caption_models_safe, unload_after_run, unload_after_run_safe, none."
|
|
)
|
|
|
|
|
|
def evict_model(key: str, *, reason: str = "manual", safe: bool = False) -> bool:
|
|
"""
|
|
Evict one cached model by key.
|
|
"""
|
|
with _CACHE_LOCK:
|
|
did_evict = _evict_locked(key, reason=reason, safe=safe)
|
|
_cuda_cleanup(
|
|
reason=f"post_evict_model:{reason}",
|
|
synchronize=True,
|
|
cooldown_sec=_SAFE_EVICTION_COOLDOWN_SEC if safe and did_evict else _FAST_EVICTION_COOLDOWN_SEC,
|
|
)
|
|
return did_evict
|
|
|
|
|
|
def evict_family(family: str, *, reason: str = "manual_family_evict", safe: bool = False) -> int:
|
|
"""
|
|
Evict all models for a family, e.g. 'qwen' or 'joy'.
|
|
"""
|
|
family_norm = normalize_cache_part(family).lower()
|
|
count = 0
|
|
|
|
with _CACHE_LOCK:
|
|
for key, entry in list(_MODEL_CACHE.items()):
|
|
if entry.family.lower() == family_norm:
|
|
if _evict_locked(key, reason=reason, safe=safe):
|
|
count += 1
|
|
|
|
_cuda_cleanup(
|
|
reason=f"post_evict_family:{family_norm}",
|
|
synchronize=True,
|
|
cooldown_sec=_SAFE_EVICTION_COOLDOWN_SEC if safe and count else _FAST_EVICTION_COOLDOWN_SEC,
|
|
)
|
|
return count
|
|
|
|
|
|
def unload_all(*, include_keep: bool = True, reason: str = "manual_unload_all", safe: bool = False) -> int:
|
|
"""
|
|
Clear the entire CaptionForge cache.
|
|
"""
|
|
count = 0
|
|
|
|
with _CACHE_LOCK:
|
|
for key in list(_MODEL_CACHE.keys()):
|
|
entry = _MODEL_CACHE.get(key)
|
|
if entry is None:
|
|
continue
|
|
if entry.keep and not include_keep:
|
|
continue
|
|
|
|
if _evict_locked(key, reason=reason, safe=safe):
|
|
count += 1
|
|
|
|
_cuda_cleanup(
|
|
reason=f"post_unload_all:{reason}",
|
|
synchronize=True,
|
|
cooldown_sec=_SAFE_EVICTION_COOLDOWN_SEC if safe and count else _FAST_EVICTION_COOLDOWN_SEC,
|
|
)
|
|
return count
|
|
|
|
|
|
def unload_after_run(key: str, *, enabled: bool, safe: bool = False) -> None:
|
|
"""
|
|
Convenience helper for node/engine code after a captioning run.
|
|
"""
|
|
if enabled:
|
|
evict_model(key, reason="unload_after_run_safe" if safe else "unload_after_run", safe=safe)
|
|
|
|
|
|
def set_safe_eviction_cooldown(seconds: float) -> None:
|
|
"""
|
|
Adjust the cooldown used by explicit *_safe cache policies.
|
|
"""
|
|
global _SAFE_EVICTION_COOLDOWN_SEC
|
|
_SAFE_EVICTION_COOLDOWN_SEC = max(0.0, float(seconds))
|
|
|
|
|
|
def get_safe_eviction_cooldown() -> float:
|
|
return _SAFE_EVICTION_COOLDOWN_SEC
|
|
|
|
|
|
def set_fast_eviction_cooldown(seconds: float) -> None:
|
|
"""
|
|
Adjust the cooldown used by normal/fast eviction policies. Defaults to 0.0.
|
|
"""
|
|
global _FAST_EVICTION_COOLDOWN_SEC
|
|
_FAST_EVICTION_COOLDOWN_SEC = max(0.0, float(seconds))
|
|
|
|
|
|
def get_fast_eviction_cooldown() -> float:
|
|
return _FAST_EVICTION_COOLDOWN_SEC
|
|
|
|
|
|
def cache_info() -> Dict[str, Any]:
|
|
"""
|
|
Return lightweight cache diagnostics for console logging or future UI/debug nodes.
|
|
"""
|
|
with _CACHE_LOCK:
|
|
entries = []
|
|
for key, entry in _MODEL_CACHE.items():
|
|
entries.append(
|
|
{
|
|
"key": key,
|
|
"family": entry.family,
|
|
"model_path": entry.model_path,
|
|
"device": entry.device,
|
|
"quantization": entry.quantization,
|
|
"role": entry.role,
|
|
"keep": entry.keep,
|
|
"age_sec": round(time.time() - entry.created_at, 3),
|
|
"idle_sec": round(time.time() - entry.last_used_at, 3),
|
|
}
|
|
)
|
|
|
|
return {
|
|
"max_loaded_models": _MAX_LOADED_MODELS,
|
|
"loaded_count": len(_MODEL_CACHE),
|
|
"safe_eviction_cooldown_sec": _SAFE_EVICTION_COOLDOWN_SEC,
|
|
"fast_eviction_cooldown_sec": _FAST_EVICTION_COOLDOWN_SEC,
|
|
"entries": entries,
|
|
}
|
|
|
|
|
|
def _enforce_capacity_locked(protected_key: Optional[str] = None) -> None:
|
|
"""
|
|
Enforce max loaded model count.
|
|
|
|
Eviction order:
|
|
1. Non-keep entries first
|
|
2. Least recently used first
|
|
3. Avoid protected_key if possible
|
|
"""
|
|
if _MAX_LOADED_MODELS < 0:
|
|
return
|
|
|
|
while len(_MODEL_CACHE) > _MAX_LOADED_MODELS:
|
|
candidates = []
|
|
|
|
for key, entry in _MODEL_CACHE.items():
|
|
if key == protected_key:
|
|
continue
|
|
if entry.keep:
|
|
continue
|
|
candidates.append(entry)
|
|
|
|
# If everything is protected/keep, do not force-evict.
|
|
if not candidates:
|
|
break
|
|
|
|
victim = min(candidates, key=lambda e: e.last_used_at)
|
|
_evict_locked(victim.key, reason="capacity", safe=False)
|
|
_cuda_cleanup(reason="post_capacity_evict", synchronize=True, cooldown_sec=_FAST_EVICTION_COOLDOWN_SEC)
|
|
|
|
|
|
def _evict_locked(key: str, *, reason: str = "unspecified", safe: bool = False) -> bool:
|
|
entry = _MODEL_CACHE.pop(key, None)
|
|
if entry is None:
|
|
return False
|
|
|
|
print(f"[CaptionForge Cache] Evicting model: {key} | reason={reason}")
|
|
|
|
# Synchronize before tearing down CUDA-backed model objects. This adds little
|
|
# overhead when kernels are already idle, but gives bitsandbytes/accelerate
|
|
# paths a cleaner boundary before references are cleared.
|
|
_cuda_synchronize(reason=f"pre_unload:{reason}", verbose=safe or _VERBOSE_SAFE_EVICTION)
|
|
|
|
try:
|
|
if entry.unload_fn is not None:
|
|
entry.unload_fn(entry.obj)
|
|
except Exception as exc:
|
|
print(f"[CaptionForge Cache] Warning: unload_fn failed for {key}: {exc}")
|
|
|
|
try:
|
|
entry.obj = None
|
|
except Exception:
|
|
pass
|
|
|
|
del entry
|
|
gc.collect()
|
|
|
|
# A second synchronize catches any work triggered by engine unload hooks.
|
|
_cuda_synchronize(reason=f"post_unload:{reason}", verbose=safe or _VERBOSE_SAFE_EVICTION)
|
|
return True
|
|
|
|
|
|
def _cuda_synchronize(*, reason: str = "", verbose: bool = False) -> None:
|
|
if torch is None:
|
|
return
|
|
|
|
try:
|
|
if torch.cuda.is_available():
|
|
if verbose:
|
|
print(f"[CaptionForge Cache] CUDA synchronize start ({reason})")
|
|
torch.cuda.synchronize()
|
|
if verbose:
|
|
print(f"[CaptionForge Cache] CUDA synchronize done ({reason})")
|
|
except Exception as exc:
|
|
print(f"[CaptionForge Cache] Warning: CUDA synchronize failed ({reason}): {exc}")
|
|
|
|
|
|
def _cuda_cleanup(
|
|
*,
|
|
reason: str = "cuda_cleanup",
|
|
synchronize: bool = True,
|
|
cooldown_sec: float = 0.0,
|
|
) -> None:
|
|
"""
|
|
Clear Python and CUDA allocator leftovers after model eviction.
|
|
|
|
synchronize=True gives native CUDA/bitsandbytes/accelerate teardown a cleaner
|
|
boundary. cooldown_sec is intentionally opt-in/nonzero only for safe policies
|
|
unless the user sets environment overrides.
|
|
"""
|
|
gc.collect()
|
|
|
|
if torch is None:
|
|
if cooldown_sec > 0:
|
|
time.sleep(cooldown_sec)
|
|
return
|
|
|
|
if synchronize:
|
|
_cuda_synchronize(reason=f"pre_cleanup:{reason}", verbose=_VERBOSE_SAFE_EVICTION)
|
|
|
|
try:
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
try:
|
|
torch.cuda.ipc_collect()
|
|
except Exception:
|
|
# ipc_collect can fail harmlessly on some Windows/CUDA setups.
|
|
pass
|
|
except Exception as exc:
|
|
print(f"[CaptionForge Cache] Warning: CUDA cleanup failed ({reason}): {exc}")
|
|
|
|
gc.collect()
|
|
|
|
if synchronize:
|
|
_cuda_synchronize(reason=f"post_cleanup:{reason}", verbose=_VERBOSE_SAFE_EVICTION)
|
|
|
|
if cooldown_sec > 0:
|
|
if _VERBOSE_SAFE_EVICTION:
|
|
print(f"[CaptionForge Cache] Cooldown {cooldown_sec:.2f}s ({reason})")
|
|
time.sleep(float(cooldown_sec))
|