577 lines
20 KiB
Python
577 lines
20 KiB
Python
"""
|
|
Typed data-transfer payloads for Bernini-R nodes.
|
|
|
|
This module centralises the structured data that flows between custom nodes.
|
|
ComfyUI native containers (CONDITIONING, LATENT, CLIP, VAE, IMAGE) keep their
|
|
external shape so other ComfyUI nodes can still consume them; only the custom
|
|
``BERNINI_*`` sockets are represented as typed dataclasses.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any, Iterator
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
# Lazy import for comfy_kitchen (only resolved in ComfyUI environment).
|
|
try:
|
|
from comfy_kitchen.tensor.base import QuantizedTensor, get_layout_class as _get_layout_class
|
|
except ImportError:
|
|
QuantizedTensor = None # type: ignore[assignment]
|
|
_get_layout_class = None # type: ignore[assignment]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ComfySamplerShim — lightweight ModelPatcher proxy for k-diffusion SDE
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@dataclass
|
|
class ComfySamplerShim:
|
|
"""Lightweight proxy that satisfies k-diffusion's ``model_patcher`` contract.
|
|
|
|
SDE-family samplers (``dpmpp_2m_sde`` etc.) resolve ``model_sampling`` by
|
|
walking the object graph::
|
|
|
|
model.inner_model.model_patcher.get_model_object('model_sampling')
|
|
|
|
Our self-contained guider (``BerniniModelWrapper``) bypasses ComfyUI's
|
|
``CFGGuider`` which normally provides this reverse reference. This shim
|
|
re-establishes the chain so SDE samplers see the same ``model_sampling``
|
|
as the rest of the pipeline.
|
|
|
|
Attributes:
|
|
_patcher: The ``ModelPatcher`` whose ``.model`` holds ``model_sampling``.
|
|
"""
|
|
_patcher: Any
|
|
|
|
def get_model_object(self, name: str) -> Any:
|
|
"""Delegate to the underlying model (e.g. for ``model_sampling``)."""
|
|
inner = getattr(self._patcher, 'model', None)
|
|
if inner is not None:
|
|
return getattr(inner, name, None)
|
|
return None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Enums
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class ContextSchedule(str, Enum):
|
|
"""Context-window scheduling strategies."""
|
|
UNIFORM_STANDARD = "uniform_standard"
|
|
UNIFORM_LOOPED = "uniform_looped"
|
|
STATIC_STANDARD = "static_standard"
|
|
|
|
|
|
class FuseMethod(str, Enum):
|
|
"""Context-window fusion methods."""
|
|
LINEAR = "linear"
|
|
SMOOTH = "smooth"
|
|
PYRAMID = "pyramid"
|
|
NONE = "none"
|
|
|
|
|
|
class GuidanceMode(str, Enum):
|
|
"""Guidance strategies supported by BerniniModelWrapper."""
|
|
CFG = "CFG"
|
|
APG = "APG"
|
|
RAAG = "RAAG"
|
|
S2 = "S2"
|
|
Z2 = "Z2"
|
|
STG_A = "STG_A"
|
|
STG_R = "STG_R"
|
|
|
|
@property
|
|
def is_stg(self) -> bool:
|
|
return self.value.startswith("STG")
|
|
|
|
@property
|
|
def stg_variant(self) -> str:
|
|
"""Return 'A' or 'R' for STG modes, empty string otherwise."""
|
|
return self.value[4:] if self.is_stg else ""
|
|
|
|
|
|
class MaskMode(str, Enum):
|
|
"""Differential-diffusion mask behaviours."""
|
|
ANNEAL = "anneal"
|
|
FREEZE = "freeze"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Custom socket payloads
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@dataclass(frozen=True)
|
|
class BerniniContext:
|
|
"""Temporal context-window configuration consumed by the sampler.
|
|
|
|
Pixel-frame values are converted to latent-frame values once at construction
|
|
time so the sampler no longer repeats the ``// 4`` arithmetic on every run.
|
|
"""
|
|
schedule: ContextSchedule
|
|
context_frames: int
|
|
context_stride: int
|
|
context_overlap: int
|
|
freenoise: bool
|
|
fuse_method: FuseMethod
|
|
rope_ntk_scale: float
|
|
latent_frames: int = field(init=False)
|
|
latent_overlap: int = field(init=False)
|
|
latent_stride: int = field(init=False)
|
|
|
|
def __post_init__(self):
|
|
# Wan VAE uses a 4x temporal downscale. ``(F - 1) // 4 + 1`` matches
|
|
# ComfyUI's latent-frame calculation for the last partial group.
|
|
object.__setattr__(
|
|
self, "latent_frames",
|
|
max(1, (self.context_frames - 1) // 4 + 1)
|
|
)
|
|
object.__setattr__(
|
|
self, "latent_overlap",
|
|
max(0, self.context_overlap // 4)
|
|
)
|
|
# Stride is provided in log2 levels by the UI; keep at least 1.
|
|
object.__setattr__(
|
|
self, "latent_stride",
|
|
max(1, self.context_stride // 4)
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BerniniBlockSwap:
|
|
"""Block-swap / VRAM-offloading configuration."""
|
|
block_to_swap: int
|
|
prefetch: bool = True
|
|
prefetch_count: int = 1
|
|
pin_memory: bool = False
|
|
loading_mode: str = "Streaming" # "Full" or "Streaming"
|
|
disk_workers: int = 4 # _DiskPrefetcher threads (Streaming mode only)
|
|
|
|
@property
|
|
def lazy(self) -> bool:
|
|
"""Whether to load block weights on demand (streaming) or all at once.
|
|
|
|
- Streaming: blocks loaded from disk on demand by _DiskPrefetcher.
|
|
- Full: all block weights loaded into CPU RAM at startup, then
|
|
BlockSwapManager moves them to the GPU window as needed.
|
|
"""
|
|
return self.loading_mode == "Streaming"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DiskLoadRequest:
|
|
"""A single disk→RAM block-load job handed to the prefetch thread pool.
|
|
|
|
Carries everything ``_DiskPrefetcher`` needs to populate one transformer
|
|
block, so the forward thread dispatches work to the pool with the same
|
|
typed ``@dataclass(frozen=True)`` payload convention used everywhere else
|
|
in Bernini-R instead of passing a bare ``block_idx`` integer.
|
|
"""
|
|
|
|
block_idx: int
|
|
group_key: str
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Slot entry — unified type for ring-buffer slot pool
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _params_extra_fields(params) -> tuple[dict, dict]:
|
|
"""Split a layout ``Params`` dataclass into (tensor_extras, meta_extras).
|
|
|
|
``scale``/``orig_dtype``/``orig_shape`` are first-class SlotEntry fields;
|
|
anything else a layout carries (nvfp4 ``block_scale`` tensor, ``transposed``
|
|
flag, future layouts) is captured here so slot copies stay layout-agnostic.
|
|
"""
|
|
import dataclasses
|
|
tensor_extras: dict = {}
|
|
meta_extras: dict = {}
|
|
if not dataclasses.is_dataclass(params):
|
|
return tensor_extras, meta_extras
|
|
for f in dataclasses.fields(params):
|
|
if f.name in ("scale", "orig_dtype", "orig_shape"):
|
|
continue
|
|
v = getattr(params, f.name)
|
|
if isinstance(v, torch.Tensor):
|
|
tensor_extras[f.name] = v
|
|
else:
|
|
meta_extras[f.name] = v
|
|
return tensor_extras, meta_extras
|
|
|
|
|
|
class SlotEntry:
|
|
"""Unified slot entry for the engine's ring-buffer slot pool.
|
|
|
|
Wraps either a regular ``torch.Tensor`` or the internal components of a
|
|
comfy_kitchen ``QuantizedTensor`` (``_qdata`` + ``_params``), so
|
|
the swap engine never branches on parameter type. Extra layout params
|
|
(e.g. nvfp4's ``block_scale``) ride in ``extra`` / ``meta``.
|
|
"""
|
|
|
|
__slots__ = ('is_qt', 'data', 'scale', 'layout_cls', 'orig_dtype',
|
|
'orig_shape', 'lora', 'extra', 'meta')
|
|
|
|
def __init__(self, *, data: torch.Tensor,
|
|
scale: torch.Tensor | None = None,
|
|
layout_cls: str = '',
|
|
orig_dtype: torch.dtype | None = None,
|
|
orig_shape: tuple[int, ...] | None = None,
|
|
lora: list | None = None,
|
|
extra: dict | None = None,
|
|
meta: dict | None = None):
|
|
self.is_qt = scale is not None
|
|
self.data = data
|
|
self.scale = scale
|
|
self.layout_cls = layout_cls
|
|
self.orig_dtype = orig_dtype
|
|
self.orig_shape = orig_shape
|
|
# Co-located LoRA payload (list of entry dicts), or None when the
|
|
# block weights in this slot are already merged. The unified slot is
|
|
# ``(block, lora)`` before folding and ``(block)`` after — "folded"
|
|
# is inferred from whether ``lora`` is still present, so no external
|
|
# bookkeeping set is needed.
|
|
self.lora = lora
|
|
# Extra layout params: tensors (nvfp4 block_scale) and scalars/flags
|
|
# (nvfp4 transposed) needed to rebuild ``Params`` on assign_to.
|
|
self.extra = extra or {}
|
|
self.meta = meta or {}
|
|
|
|
@classmethod
|
|
def empty_like(cls, param: nn.Parameter, device, pin_memory: bool = False):
|
|
"""Pre-allocate a slot buffer matching *param* on *device*."""
|
|
d = param.data
|
|
if hasattr(d, '_qdata'):
|
|
pin = device == "cpu" and pin_memory
|
|
t_extras, m_extras = _params_extra_fields(d._params)
|
|
return cls(
|
|
data=torch.empty_like(d._qdata, device=device, pin_memory=pin),
|
|
scale=torch.empty_like(d._params.scale, device=device, pin_memory=pin),
|
|
layout_cls=d._layout_cls,
|
|
orig_dtype=d._params.orig_dtype,
|
|
orig_shape=tuple(d.shape),
|
|
extra={n: torch.empty_like(t, device=device, pin_memory=pin)
|
|
for n, t in t_extras.items()},
|
|
meta=m_extras,
|
|
)
|
|
return cls(
|
|
data=torch.empty_like(d, device=device,
|
|
pin_memory=(device == 'cpu' and pin_memory)),
|
|
)
|
|
|
|
@classmethod
|
|
def empty_like_entry(cls, entry: "SlotEntry", device,
|
|
pin_memory: bool = False):
|
|
"""Pre-allocate a slot buffer on *device* matching *entry*'s layout.
|
|
|
|
Unlike ``empty_like`` (which reads from an ``nn.Parameter``), this
|
|
reads from another ``SlotEntry`` — safe when the source parameter
|
|
may already be on a different device.
|
|
"""
|
|
if entry.is_qt:
|
|
pin = device == "cpu" and pin_memory
|
|
return cls(
|
|
data=torch.empty_like(entry.data, device=device, pin_memory=pin),
|
|
scale=torch.empty_like(entry.scale, device=device, pin_memory=pin),
|
|
layout_cls=entry.layout_cls,
|
|
orig_dtype=entry.orig_dtype,
|
|
orig_shape=entry.orig_shape,
|
|
extra={n: torch.empty_like(t, device=device, pin_memory=pin)
|
|
for n, t in entry.extra.items()},
|
|
meta=dict(entry.meta),
|
|
)
|
|
return cls(
|
|
data=torch.empty_like(
|
|
entry.data, device=device,
|
|
pin_memory=(device == 'cpu' and pin_memory)),
|
|
)
|
|
|
|
def copy_from(self, param: nn.Parameter, non_blocking: bool = False) -> None:
|
|
"""Copy *param*'s weight data into this pre-allocated buffer."""
|
|
d = param.data
|
|
if self.is_qt:
|
|
self.data.copy_(d._qdata, non_blocking=non_blocking)
|
|
self.scale.copy_(d._params.scale, non_blocking=non_blocking)
|
|
t_extras, _ = _params_extra_fields(d._params)
|
|
for n, t in t_extras.items():
|
|
if n in self.extra:
|
|
self.extra[n].copy_(t, non_blocking=non_blocking)
|
|
else:
|
|
self.data.copy_(d, non_blocking=non_blocking)
|
|
|
|
def assign_to(self, block: nn.Module, param_name: str) -> None:
|
|
"""Wire this buffer's data into *block*'s submodule ``_parameters``.
|
|
|
|
Quantized entries: reconstruct a fresh GPU ``QuantizedTensor`` and
|
|
replace the parameter via ``_parameters`` dict (bypasses the broken
|
|
``p.data = ...`` on wrapper subclasses).
|
|
Regular entries: ``_parameters[name].data = ...`` (works normally).
|
|
"""
|
|
*mod_path, pn = param_name.split('.')
|
|
sub = block
|
|
for part in mod_path:
|
|
sub = getattr(sub, part)
|
|
|
|
if self.is_qt:
|
|
layout_cls = _get_layout_class(self.layout_cls)
|
|
params = layout_cls.Params(
|
|
scale=self.scale,
|
|
orig_dtype=self.orig_dtype,
|
|
orig_shape=self.orig_shape,
|
|
**self.meta,
|
|
**self.extra,
|
|
)
|
|
qt = QuantizedTensor(self.data, self.layout_cls, params)
|
|
sub._parameters[pn] = nn.Parameter(qt, requires_grad=False)
|
|
else:
|
|
sub._parameters[pn].data = self.data
|
|
|
|
|
|
# ── LoraTensorEntry: typed metadata for a LoRA tensor in an offset plan ─────
|
|
|
|
|
|
@dataclass
|
|
class LoraTensorEntry:
|
|
"""Metadata for one LoRA tensor in an offset plan (no materialised tensor).
|
|
|
|
``kind`` is one of ``"A"``, ``"B"``, ``"alpha"``, ``"diff_b"``,
|
|
``"diff"``. ``spec_idx`` is the index into the owning
|
|
``LoraBlockReader._spec_strengths`` list.
|
|
"""
|
|
fd_idx: int
|
|
offset: int
|
|
length: int
|
|
dtype: torch.dtype
|
|
shape: tuple[int, ...]
|
|
kind: str
|
|
spec_idx: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BerniniTeaCache:
|
|
"""TeaCache block-caching configuration."""
|
|
max_skip_blocks: int
|
|
start_block: int
|
|
rel_l1_thresh: float
|
|
warmup_steps: int
|
|
cooldown_steps: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BerniniAttention:
|
|
"""Attention-backend selection returned by BerniniR_AttentionConfig."""
|
|
backend: str
|
|
force_backend: bool
|
|
available: tuple[str, ...]
|
|
best: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BerniniGuidance:
|
|
"""Per-step guidance scale schedule."""
|
|
values: list[float]
|
|
curve: str
|
|
steps: int
|
|
start: float
|
|
end: float
|
|
|
|
def __post_init__(self):
|
|
# Defensive: ensure the list length matches the declared step count.
|
|
if len(self.values) != self.steps:
|
|
object.__setattr__(
|
|
self, "values",
|
|
[self.start] * self.steps if self.steps > 0 else []
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BerniniGuidanceConfig:
|
|
"""Guidance strategy selection + per-mode hyper-parameters.
|
|
|
|
Replaces the 8 individual guidance widgets on the sampler node with
|
|
a single ``BERNINI_GUIDANCE_CONFIG`` socket, following the same pattern
|
|
as ``BERNINI_CTX`` / ``BERNINI_BLOCKSWAP`` / ``BERNINI_TEACACHE``.
|
|
"""
|
|
mode: GuidanceMode = GuidanceMode.CFG
|
|
|
|
# APG
|
|
apg_eta: float = 0.15
|
|
apg_rescale: bool = True
|
|
apg_momentum: float = 0.0
|
|
|
|
# RAAG
|
|
raag_alpha: float = 1.0
|
|
|
|
# S²
|
|
s2_omega: float = 1.0
|
|
|
|
# STG
|
|
stg_scale: float = 1.0
|
|
stg_block_idx: str = "10,20,27"
|
|
|
|
# Z²
|
|
z2_collapse: float = 0.3
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Segment schedule (hard-cut denoising)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@dataclass(frozen=True)
|
|
class SegmentSpec:
|
|
"""A single segment in a hard-cut prompt schedule.
|
|
|
|
Each segment has a single prompt and covers a contiguous frame range.
|
|
No interpolation — the prompt is constant across the segment.
|
|
"""
|
|
start_frame: int # 1-based pixel frame, inclusive
|
|
end_frame: int # 1-based pixel frame, inclusive
|
|
prompt: str # single prompt (no "/" syntax)
|
|
|
|
|
|
@dataclass
|
|
class EncodedSegment:
|
|
"""A segment with its CLIP-encoded prompt embedding.
|
|
|
|
Passed from ``BerniniR_SegmentSchedule`` to the sampler via the
|
|
``segment_specs`` key in the positive conditioning extra dict.
|
|
|
|
``start_latent`` / ``end_latent`` are the pre-computed latent-space frame
|
|
bounds (``end_latent`` exclusive) so the sampler's context-window framework
|
|
can build prompt-travel windows without repeating the pixel→latent math.
|
|
"""
|
|
start_frame: int
|
|
end_frame: int
|
|
embed: torch.Tensor
|
|
pooled: dict
|
|
start_latent: int = 0
|
|
end_latent: int = 0
|
|
|
|
|
|
@dataclass
|
|
class SegmentWindow:
|
|
"""A latent-space window paired with its prompt embedding.
|
|
|
|
Consumed by the context-window wrapper (``_build_context_window_wrapper``)
|
|
to run temporal prompt-travel: each window denoises with its own text
|
|
embedding, and adjacent windows crossfade in their overlap region. All Wan
|
|
text embeds share shape ``(1, 512, 4096)``, so swapping the embedding per
|
|
window is always shape-safe.
|
|
"""
|
|
start_latent: int
|
|
end_latent: int # exclusive
|
|
embed: torch.Tensor # (1, 512, 4096)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Conditioning helper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@dataclass
|
|
class Condition:
|
|
"""Single ComfyUI conditioning entry: a prompt embedding plus extra dict."""
|
|
embed: torch.Tensor
|
|
extra: dict[str, Any]
|
|
|
|
def to_comfy(self) -> list:
|
|
"""Return the native ComfyUI representation ``[tensor, dict]``."""
|
|
return [self.embed, self.extra]
|
|
|
|
|
|
@dataclass
|
|
class Conditioning:
|
|
"""List-compatible wrapper around ComfyUI's ``[[tensor, dict], ...]`` format.
|
|
|
|
Nodes still return the underlying list on the wire so ComfyUI and other
|
|
custom nodes can consume the output; internally we use this class to avoid
|
|
magic indices like ``positive[0][1]``.
|
|
"""
|
|
items: list[Condition]
|
|
|
|
@classmethod
|
|
def from_comfy(cls, comfy_cond: list | Conditioning | None) -> "Conditioning":
|
|
"""Build a ``Conditioning`` helper from a native ComfyUI list."""
|
|
if comfy_cond is None:
|
|
return cls([])
|
|
if isinstance(comfy_cond, Conditioning):
|
|
return cls(list(comfy_cond.items))
|
|
parsed = []
|
|
for entry in comfy_cond:
|
|
if isinstance(entry, Condition):
|
|
parsed.append(entry)
|
|
elif isinstance(entry, (list, tuple)) and len(entry) >= 2:
|
|
parsed.append(Condition(embed=entry[0], extra=dict(entry[1] or {})))
|
|
else:
|
|
raise ValueError(f"Invalid conditioning entry: {entry!r}")
|
|
return cls(parsed)
|
|
|
|
def to_comfy(self) -> list[list]:
|
|
"""Return the native ComfyUI representation."""
|
|
return [item.to_comfy() for item in self.items]
|
|
|
|
# -- list-like interface ------------------------------------------------
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.items)
|
|
|
|
def __iter__(self) -> Iterator[Condition]:
|
|
return iter(self.items)
|
|
|
|
def __getitem__(self, index: int | slice):
|
|
result = self.items[index]
|
|
if isinstance(result, list):
|
|
return Conditioning(result)
|
|
return result
|
|
|
|
def __bool__(self) -> bool:
|
|
return bool(self.items)
|
|
|
|
# -- convenience accessors ----------------------------------------------
|
|
|
|
@property
|
|
def first(self) -> Condition | None:
|
|
"""Return the first conditioning entry, or ``None`` if empty."""
|
|
return self.items[0] if self.items else None
|
|
|
|
@property
|
|
def embed(self) -> torch.Tensor | None:
|
|
"""Return the embedding of the first entry."""
|
|
first = self.first
|
|
return first.embed if first is not None else None
|
|
|
|
@property
|
|
def extra(self) -> dict[str, Any] | None:
|
|
"""Return the extra dict of the first entry."""
|
|
first = self.first
|
|
return first.extra if first is not None else None
|
|
|
|
def get_extra(self, key: str, default: Any = None) -> Any:
|
|
"""Read a value from the first entry's extra dict."""
|
|
extra = self.extra
|
|
if extra is None:
|
|
return default
|
|
return extra.get(key, default)
|
|
|
|
def set_extra(self, key: str, value: Any) -> "Conditioning":
|
|
"""Set a value on every entry's extra dict (in-place)."""
|
|
for item in self.items:
|
|
item.extra[key] = value
|
|
return self
|
|
|
|
def with_extra(self, **values) -> "Conditioning":
|
|
"""Return a new ``Conditioning`` with merged extra dicts."""
|
|
new_items = []
|
|
for item in self.items:
|
|
merged = dict(item.extra)
|
|
merged.update(values)
|
|
new_items.append(Condition(embed=item.embed, extra=merged))
|
|
return Conditioning(new_items)
|
|
|
|
def copy(self) -> "Conditioning":
|
|
"""Shallow copy: new ``Condition`` objects sharing the same tensors."""
|
|
return Conditioning([
|
|
Condition(embed=item.embed, extra=dict(item.extra))
|
|
for item in self.items
|
|
]) |