439 lines
18 KiB
Python
439 lines
18 KiB
Python
"""
|
||
Self-managed model lifecycle for Bernini-R.
|
||
|
||
Provides ``BerniniRModelHandle``: a lightweight object that knows how to load
|
||
and unload a Bernini-R / Wan model from disk on demand. LoRA weights are
|
||
merged inline when the model is loaded, so the handle never depends on
|
||
ComfyUI's ``load_lora_for_models`` / patcher mechanism for model management.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
from collections import OrderedDict
|
||
from typing import Any
|
||
|
||
import torch
|
||
import comfy.model_management as mm
|
||
|
||
from .vram import collect_garbage
|
||
|
||
from .log import get_logger as _get_logger
|
||
|
||
logger = _get_logger("ModelMgr")
|
||
# ---------------------------------------------------------------------------
|
||
# Lightweight proxy for latent_preview.prepare_callback compatibility
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class _LatentFormatProxy:
|
||
"""Minimal model-like object that provides ``latent_format`` without
|
||
loading the full model. ``prepare_callback`` only accesses:
|
||
|
||
model.load_device – set on the handle itself
|
||
model.model.latent_format – provided here
|
||
|
||
So callers can create a preview callback before the handle is loaded.
|
||
"""
|
||
|
||
def __init__(self):
|
||
import comfy.latent_formats
|
||
self.latent_format = comfy.latent_formats.Wan21()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Module-level model cache — avoids a full checkpoint disk read + build on re-runs
|
||
# ---------------------------------------------------------------------------
|
||
_MAX_MODEL_CACHE: int = 1
|
||
_model_cache: OrderedDict[str, Any] = OrderedDict()
|
||
|
||
|
||
def _make_cache_key(
|
||
model_path: str,
|
||
lora_specs: list | None,
|
||
compile_cfg: dict | None,
|
||
attn_backend_args: dict | None,
|
||
block_swap: bool = False,
|
||
) -> str:
|
||
"""Stable hash of the full model config for cache lookup.
|
||
|
||
``block_swap`` MUST be part of the key: it changes the patcher's
|
||
``load_device`` and attaches a ``BlockSwapManager`` to
|
||
``diffusion_model`` (see ``wan_model.load_bernini_model``). Reusing a
|
||
non-block-swap patcher for a block-swap request would leave the model on
|
||
the GPU with no manager, breaking block swapping at runtime.
|
||
"""
|
||
_loras = sorted([(str(p), float(s)) for p, s in (lora_specs or [])])
|
||
_compile = dict(sorted(compile_cfg.items())) if compile_cfg else {}
|
||
# attn_backend_args may be a BerniniAttention dataclass — not a dict.
|
||
_attn = {}
|
||
if attn_backend_args is not None:
|
||
import dataclasses as _dc
|
||
if _dc.is_dataclass(attn_backend_args) and not isinstance(attn_backend_args, dict):
|
||
_attn = dict(sorted(_dc.asdict(attn_backend_args).items()))
|
||
elif isinstance(attn_backend_args, dict):
|
||
_attn = dict(sorted(attn_backend_args.items()))
|
||
raw = json.dumps(
|
||
dict(path=model_path, loras=_loras, compile=_compile, attn=_attn,
|
||
block_swap=block_swap),
|
||
sort_keys=True,
|
||
)
|
||
return hashlib.sha256(raw.encode()).hexdigest()[:16]
|
||
|
||
|
||
def _cache_evict_patcher(patcher: Any) -> bool:
|
||
"""Remove a specific patcher from the LRU cache and force cleanup."""
|
||
removed = False
|
||
for key, value in list(_model_cache.items()):
|
||
if value is patcher:
|
||
_model_cache.pop(key, None)
|
||
logger.info("Evicted model from cache: %s", key[:8])
|
||
removed = True
|
||
if removed:
|
||
collect_garbage()
|
||
return removed
|
||
|
||
|
||
def _cache_put(key: str, patcher: Any) -> None:
|
||
"""Store a built patcher in the LRU cache, evicting oldest if full."""
|
||
_model_cache[key] = patcher
|
||
_model_cache.move_to_end(key)
|
||
while len(_model_cache) > _MAX_MODEL_CACHE:
|
||
oldest_key, oldest_patcher = _model_cache.popitem(last=False)
|
||
logger.info("Model cache evict: %s", oldest_key[:8])
|
||
try:
|
||
# Cancel any async block-swap transfers before unloading.
|
||
_dm = getattr(getattr(oldest_patcher, "model", None),
|
||
"diffusion_model", None)
|
||
_mgr = getattr(_dm, "_block_swap_mgr", None)
|
||
if _mgr is not None:
|
||
_mgr.evict_all()
|
||
try:
|
||
delattr(_dm, "_block_swap_mgr")
|
||
except Exception:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
try:
|
||
# Fully release from ComfyUI's model manager so VRAM is freed.
|
||
mm.unload_model_and_clones(oldest_patcher)
|
||
except Exception:
|
||
pass
|
||
del oldest_patcher
|
||
collect_garbage()
|
||
|
||
|
||
class BerniniRModelHandle:
|
||
"""Lazy-loading handle for a Bernini-R diffusion model.
|
||
|
||
The handle itself is tiny (only paths and configs). The actual model is
|
||
loaded only when ``load()`` is called, and fully released with
|
||
``unload()``. A module-level LRU cache (``_MAX_MODEL_CACHE = 1``) keeps
|
||
the most recently used patcher object alive so re-runs skip the build
|
||
step; the weights are offloaded to CPU by ComfyUI when not in use.
|
||
|
||
Attributes
|
||
----------
|
||
model_path:
|
||
Path to the base model checkpoint on disk.
|
||
attn_backend_args:
|
||
Optional attention backend config from ``BerniniR_AttentionConfig``.
|
||
lora_specs:
|
||
List of ``(lora_path, strength)`` to merge inline at load time.
|
||
compile_cfg:
|
||
Optional ``{"mode": ..., "fullgraph": ..., "dynamic_shapes": ...}``.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model_path: str,
|
||
attn_backend_args: dict | None = None,
|
||
lora_specs: list[tuple[str, float]] | None = None,
|
||
compile_cfg: dict | None = None,
|
||
block_swap: bool = False,
|
||
):
|
||
self.model_path = model_path
|
||
self.attn_backend_args = attn_backend_args
|
||
self.lora_specs = list(lora_specs) if lora_specs else []
|
||
self.compile_cfg = compile_cfg
|
||
self.block_swap = block_swap
|
||
self._model_patcher: Any | None = None
|
||
|
||
# Preview-compatible interface (no model load required):
|
||
# prepare_callback(model, steps) needs:
|
||
# model.load_device -> mm.get_torch_device()
|
||
# model.model.latent_format -> Wan21()
|
||
self.load_device = mm.get_torch_device()
|
||
self.model = _LatentFormatProxy()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Public API
|
||
# ------------------------------------------------------------------
|
||
|
||
def load(self, block_swap_config: Any = None) -> Any:
|
||
"""Load the model from disk, apply LoRAs/compile, and return patcher.
|
||
|
||
*block_swap_config* is a ``BerniniBlockSwap`` dataclass (or ``None``
|
||
to disable block swap). When present, the model always goes through
|
||
``BlockSwapManager`` (ring buffer + ``SlotEntry``) for GPU placement,
|
||
regardless of ``block_to_swap``. When ``block_to_swap > 0``, extra
|
||
blocks are offloaded to CPU/disk.
|
||
"""
|
||
has_config = block_swap_config is not None
|
||
need_swap = has_config and block_swap_config.block_to_swap > 0
|
||
self.block_swap = need_swap
|
||
block_swap = need_swap
|
||
lazy = block_swap_config.lazy if need_swap else False
|
||
|
||
# 1. Module-level cache — cross-execution reuse
|
||
cache_key = _make_cache_key(
|
||
self.model_path, self.lora_specs,
|
||
self.compile_cfg, self.attn_backend_args,
|
||
block_swap=self.block_swap,
|
||
)
|
||
cached: Any | None = _model_cache.get(cache_key)
|
||
if cached is not None:
|
||
_model_cache.move_to_end(cache_key) # LRU bump
|
||
self._model_patcher = cached
|
||
logger.debug("Cache hit: %s", cache_key[:8])
|
||
if need_swap:
|
||
# BlockSwapManager owns GPU placement — keep on offload device.
|
||
cached.load_device = cached.offload_device
|
||
# Ensure _block_swap_config is set on the model (may be missing
|
||
# if the cached patcher predates the unified config path).
|
||
dm = getattr(cached.model, "diffusion_model", None)
|
||
if dm is not None and not hasattr(dm, '_block_swap_config'):
|
||
dm._block_swap_config = block_swap_config
|
||
else:
|
||
mm.load_models_gpu([cached])
|
||
return cached
|
||
|
||
# 2. Guard: handle still holds a loaded patcher (not yet unloaded)
|
||
if self._model_patcher is not None:
|
||
if need_swap:
|
||
self._model_patcher.load_device = self._model_patcher.offload_device
|
||
dm = getattr(self._model_patcher.model, "diffusion_model", None)
|
||
if dm is not None and not hasattr(dm, '_block_swap_config'):
|
||
dm._block_swap_config = block_swap_config
|
||
else:
|
||
mm.load_models_gpu([self._model_patcher])
|
||
return self._model_patcher
|
||
|
||
logger.info(
|
||
"Loading model from disk: %s (loras=%d)",
|
||
self.model_path,
|
||
len(self.lora_specs),
|
||
)
|
||
|
||
# 1. Build model options (attention backend override).
|
||
model_opts = self._build_model_options()
|
||
|
||
# 2. In block-swap mode, create the RandomAccessBlockReader so the
|
||
# streaming loader skips block weights and defers to
|
||
# BlockSwapManager's _DiskPrefetcher for on-demand disk reads.
|
||
# Only needed when blocks are actually offloaded (block_to_swap > 0).
|
||
_reader = None
|
||
if has_config and block_swap_config.block_to_swap > 0:
|
||
from ..utils.block_reader import RandomAccessBlockReader
|
||
_reader = RandomAccessBlockReader(
|
||
self.model_path,
|
||
lora_specs=self.lora_specs or None,
|
||
)
|
||
|
||
# 3. Build the model. For .safetensors we use a streaming loader that
|
||
# never holds the full state dict in RAM, so the load peak drops
|
||
# from ~2x model size to ~1x model size + one block group.
|
||
from ..models.loader import load_bernini_model
|
||
patcher = load_bernini_model(
|
||
self.model_path,
|
||
model_options=model_opts,
|
||
state_dict=None,
|
||
lora_specs=self.lora_specs or None,
|
||
block_swap=need_swap,
|
||
lazy=lazy,
|
||
block_reader=_reader,
|
||
)
|
||
|
||
# 3. Set device targets.
|
||
# - Block swap config present: the model goes through
|
||
# BlockSwapManager (ring buffer + SlotEntry). We set
|
||
# load_device == offload_device so ComfyUI's executor never
|
||
# hoists the whole model onto the GPU first (that would create
|
||
# a second full copy and race BlockSwapManager's transfers).
|
||
# - No config: standard ComfyUI behaviour (model on GPU).
|
||
patcher.offload_device = mm.unet_offload_device()
|
||
if need_swap:
|
||
patcher.load_device = patcher.offload_device
|
||
else:
|
||
patcher.load_device = mm.get_torch_device()
|
||
|
||
# Store the full block swap config on the diffusion model so the
|
||
# forward pass can read it directly (only when block swap is active).
|
||
dm = getattr(patcher.model, "diffusion_model", None)
|
||
if dm is not None and need_swap:
|
||
dm._block_swap_config = block_swap_config
|
||
|
||
# 4. torch.compile if requested.
|
||
if self.compile_cfg and self.compile_cfg.get("mode", "none") != "none":
|
||
from ..models.wan_compile import compile_wan_model
|
||
compile_wan_model(
|
||
patcher,
|
||
compile_mode=self.compile_cfg["mode"],
|
||
fullgraph=self.compile_cfg.get("fullgraph", False),
|
||
dynamic=self.compile_cfg.get("dynamic_shapes", True),
|
||
)
|
||
|
||
# 5. Store in module-level cache so re-runs skip disk + build
|
||
_cache_put(cache_key, patcher)
|
||
self._model_patcher = patcher
|
||
return patcher
|
||
|
||
def unload(self) -> None:
|
||
"""Fully unload the model from VRAM / RAM and drop the patcher.
|
||
|
||
We always do a full release instead of caching the model on CPU:
|
||
moving a large model to CPU before deletion leaves PyTorch's CPU
|
||
allocator holding the RAM on Windows, which causes the next large
|
||
model load to run out of memory. This also ensures ComfyUI's own
|
||
model manager removes the patcher from its loaded-model list.
|
||
"""
|
||
if self._model_patcher is None:
|
||
return
|
||
|
||
logger.info("Unloading model from memory: %s", self.model_path)
|
||
|
||
patcher = self._model_patcher
|
||
|
||
# 0. Close block reader (disk fd) if in lazy mode.
|
||
base_model = getattr(patcher, "model", None)
|
||
dm = getattr(base_model, "diffusion_model", None)
|
||
if dm is not None:
|
||
_reader = getattr(dm, '_block_reader', None)
|
||
if _reader is not None:
|
||
try:
|
||
_reader.close()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
delattr(dm, '_block_reader')
|
||
except Exception:
|
||
pass
|
||
|
||
# 1. Shrink all parameter tensors to zero-size before ComfyUI's unload.
|
||
# mm.unload_model_and_clones only offloads to CPU — the PyTorch CPU
|
||
# allocator never returns pages to the OS on Windows.
|
||
from .vram import release_model_ram as _release_model_ram
|
||
_release_model_ram(patcher)
|
||
|
||
# 1b. Evict this patcher from the module-level cache so the next
|
||
# load() rebuilds a *fresh* model instead of reusing a patcher
|
||
# whose parameters were just shrunk to zero-size (1-D) tensors.
|
||
# Without this, load()'s cache hit returns the corpse and the
|
||
# forward pass crashes with "mat2 must be a matrix, got 1-D
|
||
# tensor" (e.g. text_embedding's F.linear). The dual-expert path
|
||
# already evicts its high patcher the same way right after freeing.
|
||
_cache_evict_patcher(patcher)
|
||
|
||
# 1. Evict block-swap blocks before dropping the model so in-flight
|
||
# async transfers don't touch freed memory. shutdown() cancels any
|
||
# prefetch on the transfer stream, synchronises *all* CUDA streams,
|
||
# and moves every block/peripheral back to CPU deterministically.
|
||
base_model = getattr(patcher, "model", None)
|
||
dm = getattr(base_model, "diffusion_model", None)
|
||
mgr = getattr(dm, "_block_swap_mgr", None)
|
||
if mgr is not None:
|
||
try:
|
||
mgr.shutdown()
|
||
except Exception as e:
|
||
logger.warning("BlockSwap shutdown failed: %s", e)
|
||
try:
|
||
delattr(dm, "_block_swap_mgr")
|
||
except Exception:
|
||
pass
|
||
|
||
# 2. Tell ComfyUI's model manager to free this model and remove it
|
||
# from the loaded-model list. Without this, ComfyUI keeps counting
|
||
# its VRAM as in-use and won't free memory for the next model.
|
||
try:
|
||
mm.unload_model_and_clones(patcher)
|
||
except Exception as e:
|
||
logger.warning("unload_model_and_clones failed: %s", e)
|
||
|
||
# 3. Drop our reference.
|
||
self._model_patcher = None
|
||
del patcher
|
||
|
||
# 4. Blocking CUDA cleanup.
|
||
if torch.cuda.is_available():
|
||
try:
|
||
torch.cuda.synchronize()
|
||
except Exception:
|
||
pass
|
||
collect_garbage()
|
||
|
||
def clone_with_lora(self, lora_path: str, strength: float) -> "BerniniRModelHandle":
|
||
"""Return a new handle with an additional LoRA spec.
|
||
|
||
The actual LoRA weights are not touched until ``load()`` is called.
|
||
"""
|
||
if strength == 0.0:
|
||
return self
|
||
new_specs = list(self.lora_specs)
|
||
new_specs.append((lora_path, strength))
|
||
return BerniniRModelHandle(
|
||
model_path=self.model_path,
|
||
attn_backend_args=self.attn_backend_args,
|
||
lora_specs=new_specs,
|
||
compile_cfg=self.compile_cfg,
|
||
block_swap=self.block_swap,
|
||
)
|
||
|
||
def clone_with_compile(
|
||
self,
|
||
mode: str,
|
||
fullgraph: bool,
|
||
dynamic_shapes: bool,
|
||
) -> "BerniniRModelHandle":
|
||
"""Return a new handle with compile config stored (lazy, like LoRA)."""
|
||
cfg = {"mode": mode, "fullgraph": fullgraph, "dynamic_shapes": dynamic_shapes}
|
||
return BerniniRModelHandle(
|
||
model_path=self.model_path,
|
||
attn_backend_args=self.attn_backend_args,
|
||
lora_specs=list(self.lora_specs),
|
||
compile_cfg=cfg,
|
||
block_swap=self.block_swap,
|
||
)
|
||
|
||
def is_loaded(self) -> bool:
|
||
"""Return True if the model is currently resident in memory."""
|
||
return self._model_patcher is not None
|
||
|
||
# ------------------------------------------------------------------
|
||
# Internal helpers
|
||
# ------------------------------------------------------------------
|
||
|
||
def _build_model_options(self) -> dict:
|
||
"""Build the model_options dict passed to ``load_bernini_model``."""
|
||
opts = {}
|
||
if self.attn_backend_args is not None:
|
||
from ..attention.backends import create_attention_override
|
||
ov = create_attention_override(
|
||
backend=self.attn_backend_args.backend,
|
||
force_backend=self.attn_backend_args.force_backend,
|
||
)
|
||
opts.setdefault("transformer_options", {})[
|
||
"optimized_attention_override"
|
||
] = ov
|
||
return opts
|
||
|
||
|
||
def load_model_handle(
|
||
model_path: str,
|
||
attn_backend_args: dict | None = None,
|
||
block_swap: bool = False,
|
||
) -> BerniniRModelHandle:
|
||
"""Factory for creating a lazy-loading Bernini-R model handle."""
|
||
return BerniniRModelHandle(
|
||
model_path=model_path,
|
||
attn_backend_args=attn_backend_args,
|
||
block_swap=block_swap,
|
||
) |