Seunghyun JiandClaude Fable 5
2b978ba03e
The ~0.9GB per-apply cond KV cache had two OOM paths on tight cards: - swapping the reference image left the previous MODEL output (and its cache) alive in ComfyUI's execution cache while the new one built — _build_cond_kv now evicts every other state's cache first (they rebuild lazily if that model is sampled again); - the cache is allocated lazily inside the first DiT forward, where comfy's memory manager can no longer make room — a PREPARE_SAMPLING wrapper now enforces a free-VRAM floor (cache*1.25 + 768MB) before sampling, partially unloading model weights instead of OOMing. Verified: headless A->B->A swap run logs 'evicted 1 stale cond KV cache(s)' per swap, no OOM, outputs clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
795 lines
35 KiB
Python
795 lines
35 KiB
Python
"""EasyControl image conditioning for ComfyUI — KSampler-compatible.
|
||
|
||
This is a from-scratch reimplementation of anima_lora's EasyControl *inference*
|
||
path, rewritten against ComfyUI's native Anima/Cosmos DiT backbone
|
||
(``comfy/ldm/cosmos/predict2.py``) so it runs under the stock KSampler with no
|
||
dedicated sampler node.
|
||
|
||
Why a reimplementation rather than vendoring anima_lora's network module:
|
||
EasyControl reaches *inside* each DiT block's self-attention (it concatenates a
|
||
reference-image key/value stream into the target's attention). anima_lora's
|
||
training DiT uses a **fused** ``self_attn.qkv_proj`` (D->3D) and an
|
||
``adaln_fused_down`` / ``adaln_up_*`` AdaLN-LoRA topology; ComfyUI's backbone
|
||
uses **split** ``q_proj`` / ``k_proj`` / ``v_proj`` and ``adaln_modulation_*``
|
||
Sequentials. So the in-block splice cannot be reused as-is.
|
||
|
||
Correctness hinge: the trained cond-LoRA delta is a fused ``D->3D`` tensor whose
|
||
three chunks are (q, k, v) in that order — exactly the split layout ComfyUI
|
||
uses. The *base* projection weights are numerically identical between the two
|
||
DiTs (same pretrained Anima model, just fused-vs-split layout — see anima_lora
|
||
``networks/attn_fuse.py``), so adding the sliced fused delta onto ComfyUI's
|
||
split q/k/v projections reproduces what training saw.
|
||
|
||
Mechanism, per denoising step:
|
||
- The reference image is VAE-encoded once (node side) and the cond stream is
|
||
walked once through all blocks (lazily, on the first DiT forward) to produce
|
||
a per-block (cond_k, cond_v) cache. The cond stream is deterministic across
|
||
steps (cond t-embedding at t=0, frozen weights), so this is computed once
|
||
and reused for every step and every CFG branch.
|
||
- Each block's self-attention is replaced by an *extended* attention over
|
||
``[target_k ; cond_k]`` / ``[target_v ; cond_v]`` with a per-block additive
|
||
logit bias ``b_cond`` on the cond columns (the trained gate). Cross-attn and
|
||
MLP run baseline. This extended attention runs on flash-attn (via an exact
|
||
LSE decomposition of the joint softmax) when ComfyUI is launched with
|
||
``--use-flash-attention`` and ``flash_attn`` is installed; otherwise it uses
|
||
a portable float-mask SDPA path. See ``_extended_self_attn``.
|
||
|
||
Robustness: the per-block forward closures resolve their block from
|
||
``state.live_blocks`` (refreshed by the pre-hook from the live diffusion_model
|
||
each forward) rather than capturing a block reference, so a downstream node that
|
||
rebuilds the DiT (e.g. block-compile) doesn't strand the patch on a dead
|
||
instance — same hazard the in-tree Anima nodes document.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import math
|
||
import weakref
|
||
from collections import OrderedDict
|
||
from typing import Optional
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from einops import rearrange
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Every EasyControlState ever applied (weak — dead model outputs drop out).
|
||
# Used to evict stale cond KV caches before building a new one: when the user
|
||
# swaps the reference image, ComfyUI's execution cache can keep the previous
|
||
# MODEL output (and its ~0.9 GB cond cache) alive next to the new one — the
|
||
# classic image-swap OOM. See ``_build_cond_kv``.
|
||
_LIVE_STATES: "weakref.WeakSet[EasyControlState]" = weakref.WeakSet()
|
||
|
||
# ComfyUI backbone helpers. Imported at module load — this node only makes sense
|
||
# inside a ComfyUI install, where these are always importable.
|
||
#
|
||
# Cosmos/Anima rope application. ComfyUI shipped the standalone
|
||
# ``apply_rotary_pos_emb(t, freqs)`` helper from June 2025 until #14181
|
||
# ("Speed up anima a bit on nvidia", in 0.24.0), which deleted the function and
|
||
# inlined its two call sites as ``comfy.quant_ops.ck.apply_rope_split_half(q, k,
|
||
# rope_emb)``. The ``rope_emb`` *format* is unchanged across that refactor — only
|
||
# the helper moved — so we keep a vendored copy of the original (numerically
|
||
# identical, split-half convention) and prefer the in-tree symbol when present.
|
||
try: # ComfyUI < 0.24
|
||
from comfy.ldm.cosmos.predict2 import apply_rotary_pos_emb as _comfy_apply_rope
|
||
except ImportError: # ComfyUI >= 0.24 — standalone helper removed
|
||
_comfy_apply_rope = None
|
||
|
||
# Optional flash-attention fast path for the extended self-attention. Mirrors
|
||
# ComfyUI's own gating (comfy/ldm/modules/attention.py): the kernel is used only
|
||
# when (a) ``flash_attn`` imports and (b) the user launched with
|
||
# ``--use-flash-attention`` (surfaced as ``model_management.flash_attention_enabled()``).
|
||
# Otherwise we fall back to the float-mask SDPA path, which works everywhere.
|
||
try:
|
||
from flash_attn import flash_attn_func as _flash_attn_func
|
||
except ImportError:
|
||
_flash_attn_func = None
|
||
|
||
_FLASH_ENABLED_CACHE: Optional[bool] = None
|
||
|
||
|
||
def _flash_ext_enabled() -> bool:
|
||
"""True iff flash-attn is importable AND ComfyUI was launched with
|
||
``--use-flash-attention``. Result cached — the launch flag can't change
|
||
mid-process."""
|
||
global _FLASH_ENABLED_CACHE
|
||
if _FLASH_ENABLED_CACHE is None:
|
||
if _flash_attn_func is None:
|
||
_FLASH_ENABLED_CACHE = False
|
||
else:
|
||
try:
|
||
from comfy import model_management
|
||
|
||
_FLASH_ENABLED_CACHE = bool(model_management.flash_attention_enabled())
|
||
except Exception:
|
||
_FLASH_ENABLED_CACHE = False
|
||
if _FLASH_ENABLED_CACHE:
|
||
logger.info("EasyControl: flash-attention extended-attn path enabled.")
|
||
return _FLASH_ENABLED_CACHE
|
||
|
||
|
||
def apply_rotary_pos_emb(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
|
||
"""Apply cosmos split-half RoPE to a single [B, S, H, D] tensor.
|
||
|
||
Vendored from ComfyUI's pre-0.24 ``comfy.ldm.cosmos.predict2`` so the node
|
||
survives the helper's removal. The newer in-tree path applies rope to q and
|
||
k jointly (``apply_rope_split_half``); this per-tensor form is what the
|
||
extended-attention splice needs (cond q/k are produced separately from the
|
||
target stream), and it consumes the same ``freqs`` layout.
|
||
"""
|
||
if _comfy_apply_rope is not None:
|
||
return _comfy_apply_rope(t, freqs)
|
||
t_ = t.reshape(*t.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2).float()
|
||
t_out = freqs[..., 0] * t_[..., 0] + freqs[..., 1] * t_[..., 1]
|
||
t_out = t_out.movedim(-1, -2).reshape(*t.shape).type_as(t)
|
||
return t_out
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Weight loading
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _load_weights_and_meta(path: str):
|
||
from safetensors import safe_open
|
||
from safetensors.torch import load_file
|
||
|
||
sd = load_file(path)
|
||
meta = {}
|
||
if path.endswith(".safetensors"):
|
||
with safe_open(path, framework="pt") as f:
|
||
meta = f.metadata() or {}
|
||
return sd, meta
|
||
|
||
|
||
class _LoRA:
|
||
"""A single ``D->r->out`` cond-LoRA delta with an fp32 bottleneck.
|
||
|
||
Mirrors anima_lora ``networks/methods/easycontrol.py::_LoRAProj``: the
|
||
bottleneck runs in fp32 for bf16 stability, the ``alpha/r`` scale is folded
|
||
in here, and the caller-side ``cond_scale * strength`` is applied via
|
||
``eff_scale`` at the call site.
|
||
|
||
``inv_scale`` is the per-input-channel factor of a channel-scaled checkpoint
|
||
(anima_lora SmoothQuant-style rebalance: the calibration scale ``s`` is
|
||
absorbed into ``lora_down`` at init and ``inv_scale = 1/s`` is stored so the
|
||
forward stays ``x @ down`` in effect). The trained function is
|
||
``up(down((x * inv_scale)))`` — skipping it mis-scales every input channel
|
||
by ``s``.
|
||
"""
|
||
|
||
__slots__ = ("down", "up", "scale", "inv_scale")
|
||
|
||
def __init__(
|
||
self,
|
||
down: torch.Tensor,
|
||
up: torch.Tensor,
|
||
alpha_over_r: float,
|
||
inv_scale: Optional[torch.Tensor] = None,
|
||
):
|
||
self.down = down
|
||
self.up = up
|
||
self.scale = alpha_over_r
|
||
self.inv_scale = inv_scale
|
||
|
||
def to(self, device, dtype):
|
||
# Keep master copies in fp32 (the forward casts to fp32 anyway).
|
||
self.down = self.down.to(device=device, dtype=torch.float32)
|
||
self.up = self.up.to(device=device, dtype=torch.float32)
|
||
if self.inv_scale is not None:
|
||
self.inv_scale = self.inv_scale.to(device=device, dtype=torch.float32)
|
||
return self
|
||
|
||
def __call__(self, x: torch.Tensor) -> torch.Tensor:
|
||
h = x.float()
|
||
if self.inv_scale is not None:
|
||
h = h * self.inv_scale
|
||
h = F.linear(h, self.down)
|
||
h = F.linear(h, self.up)
|
||
return (h * self.scale).to(x.dtype)
|
||
|
||
|
||
class EasyControlState:
|
||
"""Holds the trained cond-LoRA tensors, the reference latent, and the lazily
|
||
built per-block KV cache. One per applied node instance; carried on the
|
||
ModelPatcher via a strong ref so it isn't GC'd."""
|
||
|
||
def __init__(self, sd: dict, meta: dict, strength: float,
|
||
cond_scale_override: Optional[float] = None):
|
||
self.num_blocks = int(meta.get("ss_num_blocks", 28))
|
||
self.r = int(meta.get("ss_cond_lora_dim", 16))
|
||
self.alpha = float(meta.get("ss_cond_lora_alpha", float(self.r)))
|
||
self.apply_ffn = bool(int(meta.get("ss_apply_ffn_lora", 1)))
|
||
cond_scale = (
|
||
float(cond_scale_override)
|
||
if cond_scale_override is not None
|
||
else float(meta.get("ss_cond_scale", 1.0))
|
||
)
|
||
# alpha/r folds into each _LoRA; cond_scale * strength is the call-site
|
||
# effective scale (matches anima_lora get_effective_scale + multiplier).
|
||
a_over_r = self.alpha / self.r if self.r > 0 else 1.0
|
||
self.eff_scale = cond_scale * float(strength)
|
||
self.strength = float(strength)
|
||
|
||
consumed = set()
|
||
|
||
def grab(prefix, idx, alpha_over_r=a_over_r):
|
||
dk = f"{prefix}.{idx}.lora_down.weight"
|
||
uk = f"{prefix}.{idx}.lora_up.weight"
|
||
sk = f"{prefix}.{idx}.inv_scale"
|
||
d, u, inv = sd.get(dk), sd.get(uk), sd.get(sk)
|
||
if d is None or u is None:
|
||
return None
|
||
consumed.update((dk, uk))
|
||
if inv is not None:
|
||
consumed.add(sk)
|
||
return _LoRA(d, u, alpha_over_r, inv_scale=inv)
|
||
|
||
self.lora_qkv = [grab("cond_lora_qkv", i) for i in range(self.num_blocks)]
|
||
self.lora_o = [grab("cond_lora_o", i) for i in range(self.num_blocks)]
|
||
if self.apply_ffn:
|
||
self.lora_ffn1 = [grab("cond_lora_ffn1", i) for i in range(self.num_blocks)]
|
||
self.lora_ffn2 = [grab("cond_lora_ffn2", i) for i in range(self.num_blocks)]
|
||
else:
|
||
self.lora_ffn1 = [None] * self.num_blocks
|
||
self.lora_ffn2 = [None] * self.num_blocks
|
||
|
||
# b_cond stored as a ParameterList → keys "b_cond.{i}" (0-d tensors).
|
||
self.b_cond = []
|
||
for i in range(self.num_blocks):
|
||
b = sd.get(f"b_cond.{i}")
|
||
if b is not None:
|
||
consumed.add(f"b_cond.{i}")
|
||
self.b_cond.append(
|
||
(b if b is not None else torch.tensor(-10.0)).float().reshape(())
|
||
)
|
||
|
||
if any(q is None for q in self.lora_qkv):
|
||
missing = [i for i, q in enumerate(self.lora_qkv) if q is None]
|
||
raise ValueError(
|
||
f"EasyControl checkpoint missing cond_lora_qkv for blocks {missing}. "
|
||
"Is this an EasyControl (networks.methods.easycontrol) checkpoint?"
|
||
)
|
||
|
||
# ---- target-stream adaln LoRA (train_adaln checkpoints) ----
|
||
# Per-block, per-branch delta on the AdaLN-LoRA up-projections
|
||
# (anima_lora ``adaln_up_{self_attn,cross_attn,mlp}``, 256 -> 3D). Comfy's
|
||
# split layout is ``adaln_modulation_{branch} = Sequential(SiLU,
|
||
# Linear(D->256), Linear(256->3D))`` with identical numerics, so each
|
||
# delta merges exactly onto ``adaln_modulation_{branch}.2.weight`` as an
|
||
# ordinary LoRA weight patch (see ``install_easycontrol``). The scale is
|
||
# ``strength * adaln_alpha / adaln_rank`` — training applies the delta at
|
||
# ``delta_scale = multiplier``, NOT ``cond_scale * multiplier``.
|
||
adaln_probe = sd.get("adaln_lora_self_attn.0.lora_down.weight")
|
||
self.train_adaln = adaln_probe is not None
|
||
self.adaln_rank = int(adaln_probe.shape[0]) if self.train_adaln else 0
|
||
adaln_alpha = float(meta.get("ss_adaln_alpha", 0.0))
|
||
if self.train_adaln and adaln_alpha <= 0.0:
|
||
# √r law fallback, mirroring anima_lora (alpha<=0 → derive from the
|
||
# cond LoRA's alpha/dim so alpha/√r stays consistent).
|
||
adaln_alpha = self.alpha * (self.adaln_rank / max(self.r, 1)) ** 0.5
|
||
self.adaln_alpha = adaln_alpha
|
||
self.adaln = {} # branch -> list[_LoRA | None], for the cond-prefill subtract
|
||
if self.train_adaln:
|
||
aa_over_r = adaln_alpha / self.adaln_rank
|
||
for branch in ("self_attn", "cross_attn", "mlp"):
|
||
mods = [
|
||
grab(f"adaln_lora_{branch}", i, alpha_over_r=aa_over_r)
|
||
for i in range(self.num_blocks)
|
||
]
|
||
if any(m is None for m in mods):
|
||
missing = [i for i, m in enumerate(mods) if m is None]
|
||
raise ValueError(
|
||
f"EasyControl checkpoint has adaln_lora keys but is missing "
|
||
f"adaln_lora_{branch} for blocks {missing} — truncated checkpoint?"
|
||
)
|
||
self.adaln[branch] = mods
|
||
elif bool(int(meta.get("ss_train_adaln", 0) or 0)):
|
||
raise ValueError(
|
||
"EasyControl checkpoint metadata says ss_train_adaln=1 but no "
|
||
"adaln_lora_* tensors are present — truncated checkpoint?"
|
||
)
|
||
|
||
# ---- strict key accounting ----
|
||
# A trained feature this node silently drops degrades output with no
|
||
# error (this is exactly how un-applied adaln checkpoints produced
|
||
# garbage) — so refuse checkpoints carrying tensors we don't map.
|
||
leftover = sorted(set(sd) - consumed)
|
||
if leftover:
|
||
shown = ", ".join(leftover[:8]) + (", …" if len(leftover) > 8 else "")
|
||
raise ValueError(
|
||
f"EasyControl checkpoint has {len(leftover)} tensor(s) this node "
|
||
f"does not implement ({shown}). The checkpoint was trained with a "
|
||
"feature newer than this node — update the node."
|
||
)
|
||
|
||
# Reference latent in *model space* (process_latent_in already applied by
|
||
# the node). 4D [1, C, H, W]; given a device/dtype at prefill time.
|
||
self.cond_latent: Optional[torch.Tensor] = None
|
||
# DiT hidden dim, filled in by install_easycontrol (for cache sizing).
|
||
self.hidden_dim: int = 0
|
||
|
||
# Built lazily on first DiT forward.
|
||
self.cond_kv_cache: Optional[list] = None
|
||
self._cache_device = None
|
||
# Refreshed each forward by the pre-hook from the live diffusion_model.
|
||
self.live_blocks = None
|
||
_LIVE_STATES.add(self)
|
||
|
||
def _to(self, device, dtype):
|
||
for lst in (self.lora_qkv, self.lora_o, self.lora_ffn1, self.lora_ffn2,
|
||
*self.adaln.values()):
|
||
for m in lst:
|
||
if m is not None:
|
||
m.to(device, dtype)
|
||
self.b_cond = [b.to(device=device) for b in self.b_cond]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Cond stream prefill (build per-block KV cache)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _cond_branch_modulation(seq, emb, adaln_delta, strength):
|
||
"""Modulation for one adaln branch of the *cond* stream.
|
||
|
||
Training applies the target-stream adaln delta only on the target side —
|
||
the cond stream always runs the frozen ``adaln_up_*``. With the delta merged
|
||
into ``seq[2].weight`` as a weight patch (see ``install_easycontrol``), the
|
||
live module would leak it into the cond stream, so run the bottleneck
|
||
stepwise and subtract the merged delta from the up output.
|
||
"""
|
||
if adaln_delta is None:
|
||
return seq(emb)
|
||
h = seq[1](seq[0](emb)) # SiLU -> down: [B, adaln_dim]
|
||
return seq[2](h) - strength * adaln_delta(h)
|
||
|
||
|
||
def _cond_qkv(attn, cond_normed, lora_qkv, eff_scale, rope):
|
||
"""Split-projection cond Q/K/V with the fused LoRA delta sliced into thirds,
|
||
then per-head norm + RoPE. Returns q, k, v as [B, S, H, D]."""
|
||
q = attn.q_proj(cond_normed)
|
||
k = attn.k_proj(cond_normed)
|
||
v = attn.v_proj(cond_normed)
|
||
if lora_qkv is not None:
|
||
dq, dk, dv = (eff_scale * lora_qkv(cond_normed)).chunk(3, dim=-1)
|
||
q = q + dq
|
||
k = k + dk
|
||
v = v + dv
|
||
h, d = attn.n_heads, attn.head_dim
|
||
q = rearrange(q, "b s (h d) -> b s h d", h=h, d=d)
|
||
k = rearrange(k, "b s (h d) -> b s h d", h=h, d=d)
|
||
v = rearrange(v, "b s (h d) -> b s h d", h=h, d=d)
|
||
q = attn.q_norm(q)
|
||
k = attn.k_norm(k)
|
||
v = attn.v_norm(v)
|
||
if rope is not None:
|
||
q = apply_rotary_pos_emb(q, rope)
|
||
k = apply_rotary_pos_emb(k, rope)
|
||
return q, k, v
|
||
|
||
|
||
def _self_attn_bshd(q, k, v):
|
||
"""Plain self-attention on [B, S, H, D] tensors → [B, S, H, D]."""
|
||
qt, kt, vt = (t.transpose(1, 2) for t in (q, k, v)) # [B,H,S,D]
|
||
out = F.scaled_dot_product_attention(qt, kt, vt)
|
||
return out.transpose(1, 2)
|
||
|
||
|
||
@torch.no_grad()
|
||
def _build_cond_kv(dit, state: EasyControlState, device, dtype):
|
||
"""Walk the cond stream once; cache per-block (cond_k, cond_v).
|
||
|
||
The cond stream is self-attention only (cross-attn is dropped on the cond
|
||
side, matching anima_lora's two-stream variant), evolving cond_x block by
|
||
block exactly as anima_lora ``precompute_cond_kv`` does, but against
|
||
ComfyUI's split-projection block API.
|
||
"""
|
||
# Evict caches left by earlier applies (e.g. the previous reference image —
|
||
# its MODEL output can stay referenced by ComfyUI's execution cache, so its
|
||
# cond cache would sit next to the one we're about to build). Dropped caches
|
||
# rebuild lazily if that older model is ever sampled again.
|
||
dropped = 0
|
||
for other in list(_LIVE_STATES):
|
||
if other is not state and other.cond_kv_cache is not None:
|
||
other.cond_kv_cache = None
|
||
other._cache_device = None
|
||
dropped += 1
|
||
if dropped and device.type == "cuda":
|
||
torch.cuda.empty_cache()
|
||
logger.info(f"EasyControl: evicted {dropped} stale cond KV cache(s)")
|
||
|
||
state._to(device, dtype)
|
||
cond_latent = state.cond_latent.to(device=device, dtype=dtype)
|
||
if cond_latent.ndim == 4:
|
||
cond_latent = cond_latent.unsqueeze(2) # [B,C,1,H,W]
|
||
|
||
cond_x5d, cond_rope_raw, cond_extra = dit.prepare_embedded_sequence(
|
||
cond_latent, fps=None, padding_mask=None
|
||
)
|
||
cond_x = rearrange(cond_x5d, "b t h w d -> b (t h w) d")
|
||
if cond_extra is not None:
|
||
# Anima ships extra_per_block_abs_pos_emb=False, so this is None. If a
|
||
# future config enables it, the per-block add would need mirroring here.
|
||
logger.warning(
|
||
"EasyControl: extra_per_block_abs_pos_emb is set; cond stream does "
|
||
"not replicate it. Results may drift."
|
||
)
|
||
cond_rope = cond_rope_raw.unsqueeze(1).unsqueeze(0) if cond_rope_raw is not None else None
|
||
|
||
# cond t-embedding at t=0 — identical to MiniTrainDIT._forward's t path.
|
||
B = cond_x.shape[0]
|
||
zeros = torch.zeros(B, 1, device=device, dtype=dtype)
|
||
cemb, cadaln = dit.t_embedder[1](dit.t_embedder[0](zeros).to(dtype))
|
||
cemb = dit.t_embedding_norm(cemb)
|
||
|
||
eff = state.eff_scale
|
||
cache = []
|
||
for idx in range(state.num_blocks):
|
||
block = dit.blocks[idx]
|
||
attn = block.self_attn
|
||
|
||
d_sa = state.adaln["self_attn"][idx] if state.train_adaln else None
|
||
d_mlp = state.adaln["mlp"][idx] if state.train_adaln else None
|
||
if block.use_adaln_lora:
|
||
s_sa, sc_sa, g_sa = (
|
||
_cond_branch_modulation(
|
||
block.adaln_modulation_self_attn, cemb, d_sa, state.strength
|
||
)
|
||
+ cadaln
|
||
).chunk(3, dim=-1)
|
||
s_mlp, sc_mlp, g_mlp = (
|
||
_cond_branch_modulation(
|
||
block.adaln_modulation_mlp, cemb, d_mlp, state.strength
|
||
)
|
||
+ cadaln
|
||
).chunk(3, dim=-1)
|
||
else:
|
||
s_sa, sc_sa, g_sa = block.adaln_modulation_self_attn(cemb).chunk(3, dim=-1)
|
||
s_mlp, sc_mlp, g_mlp = block.adaln_modulation_mlp(cemb).chunk(3, dim=-1)
|
||
|
||
# ---- self-attn (this is what we cache) ----
|
||
cond_normed = block.layer_norm_self_attn(cond_x) * (1 + sc_sa) + s_sa
|
||
cq, ck, cv = _cond_qkv(attn, cond_normed, state.lora_qkv[idx], eff, cond_rope)
|
||
cache.append((ck.detach(), cv.detach()))
|
||
|
||
# ---- evolve cond_x to feed the next block ----
|
||
attn_out = _self_attn_bshd(cq, ck, cv)
|
||
attn_out = rearrange(attn_out, "b s h d -> b s (h d)")
|
||
proj = attn.output_proj(attn_out)
|
||
if state.lora_o[idx] is not None:
|
||
proj = proj + eff * state.lora_o[idx](attn_out)
|
||
cond_x = cond_x + g_sa * proj
|
||
|
||
# ---- MLP ----
|
||
mlp_normed = block.layer_norm_mlp(cond_x) * (1 + sc_mlp) + s_mlp
|
||
hmid = block.mlp.layer1(mlp_normed)
|
||
if state.lora_ffn1[idx] is not None:
|
||
hmid = hmid + eff * state.lora_ffn1[idx](mlp_normed)
|
||
hmid = block.mlp.activation(hmid)
|
||
mlp_out = block.mlp.layer2(hmid)
|
||
if state.lora_ffn2[idx] is not None:
|
||
mlp_out = mlp_out + eff * state.lora_ffn2[idx](hmid)
|
||
cond_x = cond_x + g_mlp * mlp_out
|
||
|
||
state.cond_kv_cache = cache
|
||
state._cache_device = device
|
||
nbytes = sum(k.numel() + v.numel() for k, v in cache) * cache[0][0].element_size()
|
||
logger.info(
|
||
f"EasyControl: built cond KV cache ({len(cache)} blocks, "
|
||
f"S_c={cache[0][0].shape[1]}, {nbytes / 1e6:.0f} MB)"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Extended self-attention (target keys + cached cond keys, b_cond gate)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _extended_self_attn(tq, tk, tv, ck, cv, b_cond):
|
||
"""Attention of target queries over [target_k; cond_k] with an additive
|
||
``b_cond`` logit bias on the cond columns. All inputs [B, S, H, D] (cond may
|
||
be batch-1 and is broadcast). Returns [B, S, H, D].
|
||
|
||
Dispatches to the flash-attn LSE-decomposed path when enabled (see
|
||
``_extended_self_attn_flash``); otherwise uses the portable float-mask SDPA
|
||
path. The float ``attn_mask`` forces SDPA onto its mem-efficient/math
|
||
backend (flash SDPA can't take arbitrary float masks), so the flash path is
|
||
a genuine speedup when available."""
|
||
if (
|
||
_flash_ext_enabled()
|
||
and tq.is_cuda
|
||
and tq.dtype in (torch.float16, torch.bfloat16)
|
||
):
|
||
return _extended_self_attn_flash(tq, tk, tv, ck, cv, b_cond)
|
||
|
||
B, S = tq.shape[0], tq.shape[1]
|
||
if ck.shape[0] != B:
|
||
ck = ck.expand(B, -1, -1, -1)
|
||
cv = cv.expand(B, -1, -1, -1)
|
||
Sc = ck.shape[1]
|
||
k = torch.cat([tk, ck], dim=1).transpose(1, 2) # [B,H,S+Sc,D]
|
||
v = torch.cat([tv, cv], dim=1).transpose(1, 2)
|
||
q = tq.transpose(1, 2)
|
||
mask = torch.zeros(S + Sc, device=q.device, dtype=q.dtype)
|
||
mask[S:] = b_cond.to(q.dtype)
|
||
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask.view(1, 1, 1, S + Sc))
|
||
return out.transpose(1, 2) # [B,S,H,D]
|
||
|
||
|
||
def _extended_self_attn_flash(tq, tk, tv, ck, cv, b_cond):
|
||
"""Flash-attn equivalent of ``_extended_self_attn`` via LSE decomposition.
|
||
|
||
The joint softmax over ``[target_k ; cond_k + b_cond]`` is split into two
|
||
independent flash attentions — one over target keys, one over cond keys —
|
||
then recombined by log-sum-exp. The uniform scalar ``b_cond`` bias on the
|
||
cond group is exact under this split: it cancels inside the cond group's own
|
||
normalization (so ``out_c`` is computed bias-free) and re-enters only as a
|
||
constant shift ``+b_cond`` on the cond group's LSE during the combine. This
|
||
is the forward half of anima_lora ``_ExtendedSelfAttnLSEFunc``.
|
||
|
||
flash_attn_func consumes ``[B, S, H, D]`` directly (no transpose) and
|
||
returns ``softmax_lse`` as ``[B, H, S]`` in the natural-log domain with the
|
||
softmax_scale already applied — matching SDPA's default ``1/sqrt(D)`` scale,
|
||
so the two paths are numerically equivalent."""
|
||
B = tq.shape[0]
|
||
if ck.shape[0] != B:
|
||
# flash requires matching batch dims; expand+contiguous (B is tiny — the
|
||
# CFG fan, typically 1–2 — so the cond-KV duplication is cheap).
|
||
ck = ck.expand(B, -1, -1, -1).contiguous()
|
||
cv = cv.expand(B, -1, -1, -1).contiguous()
|
||
|
||
out_t, lse_t, _ = _flash_attn_func(tq, tk, tv, return_attn_probs=True)
|
||
out_c, lse_c, _ = _flash_attn_func(tq, ck, cv, return_attn_probs=True)
|
||
|
||
# lse_*: [B, H, S], fp32. Shift the cond group by the trained scalar gate.
|
||
lse_c_b = lse_c + b_cond.to(lse_c.dtype)
|
||
log_z = torch.logaddexp(lse_t, lse_c_b)
|
||
w_t = torch.exp(lse_t - log_z) # [B, H, S]
|
||
w_c = torch.exp(lse_c_b - log_z)
|
||
# weights → [B, S, H, 1] to scale the [B, S, H, D] flash outputs.
|
||
w_t = w_t.transpose(1, 2).unsqueeze(-1).to(out_t.dtype)
|
||
w_c = w_c.transpose(1, 2).unsqueeze(-1).to(out_c.dtype)
|
||
return out_t * w_t + out_c * w_c # [B, S, H, D]
|
||
|
||
|
||
def _make_self_attn_forward(state: EasyControlState, idx: int):
|
||
"""Replacement for one block's ``self_attn.forward`` (an ``Attention``).
|
||
|
||
Overriding only the self-attention — not the whole Block — keeps every other
|
||
part of ComfyUI's block native (AdaLN, cross-attn, MLP, residuals, dtype
|
||
handling), so there is nothing to drift and no block-level weights to strand
|
||
(``Attention`` holds no direct parameters; its q/k/v/o projections and norms
|
||
are submodules that keep their own cast-weights forward).
|
||
|
||
Resolves the live ``Attention`` from ``state.live_blocks[idx].self_attn``
|
||
each call, so a downstream DiT rebuild can't strand the patch on a dead
|
||
instance.
|
||
|
||
``torch._dynamo.disable`` keeps this an opaque eager region so the closure
|
||
composes with per-block ``torch.compile`` (e.g. the Anima Block Compile
|
||
node). Without it, dynamo tries to trace the ``state.live_blocks[idx]``
|
||
deref while compiling that very block and re-enters the module tree from a
|
||
non-local source, tripping ``"UnspecializedNNModuleVariable(Block) is
|
||
already tracked for mutation"``. The same live-block indirection that makes
|
||
this patch survive a DiT rebuild is what dynamo can't reconcile, so we just
|
||
graph-break around the extended attention (the rest of the block still
|
||
compiles).
|
||
"""
|
||
|
||
@torch._dynamo.disable
|
||
def forward(x, context=None, rope_emb=None, transformer_options={}):
|
||
attn = state.live_blocks[idx].self_attn
|
||
cache = state.cond_kv_cache
|
||
# Cond not ready, or this somehow ran as cross-attn → native attention.
|
||
if cache is None or context is not None:
|
||
return type(attn).forward(attn, x, context, rope_emb, transformer_options)
|
||
|
||
tq, tk, tv = attn.compute_qkv(x, None, rope_emb=rope_emb) # [B,S,H,D]
|
||
ck, cv = cache[idx]
|
||
out = _extended_self_attn(
|
||
tq, tk, tv, ck.to(tk.dtype), cv.to(tv.dtype), state.b_cond[idx]
|
||
)
|
||
out = rearrange(out, "b s h d -> b s (h d)")
|
||
return attn.output_dropout(attn.output_proj(out))
|
||
|
||
return forward
|
||
|
||
|
||
def _cond_cache_bytes(state: EasyControlState) -> int:
|
||
"""Estimated VRAM of the per-block cond KV cache (2-byte sampling dtype)."""
|
||
if state.cond_latent is None or not state.hidden_dim:
|
||
return 0
|
||
h, w = state.cond_latent.shape[-2:]
|
||
tokens = math.ceil(h / 2) * math.ceil(w / 2) # patch_spatial = 2
|
||
return state.num_blocks * 2 * tokens * state.hidden_dim * 2
|
||
|
||
|
||
def _make_prepare_sampling_wrapper(state: EasyControlState):
|
||
"""PREPARE_SAMPLING wrapper: reserve VRAM for the cond KV cache.
|
||
|
||
The cache (+ the one-off prefill walk) is allocated lazily inside the first
|
||
DiT forward, where comfy's memory manager can no longer make room — on
|
||
tight cards an image swap then OOMs mid-sampling. After comfy has sized and
|
||
loaded models for the run (the wrapped executor), ask ``free_memory`` to
|
||
hold the cache's worth of extra headroom free; it partially unloads loaded
|
||
models if needed, which is safe here (before any forward) and merely slower,
|
||
where the alternative was an OOM.
|
||
"""
|
||
|
||
def prepare_sampling_wrapper(executor, model, noise_shape, conds, **kwargs):
|
||
out = executor(model, noise_shape, conds, **kwargs)
|
||
if state.cond_kv_cache is None: # only when a prefill is pending
|
||
try:
|
||
from comfy import model_management
|
||
|
||
dev = model_management.get_torch_device()
|
||
extra = _cond_cache_bytes(state)
|
||
if extra and dev.type == "cuda":
|
||
# Absolute free-VRAM floor: cache (+25% for the prefill
|
||
# walk) + headroom for sampling activations. A no-op when
|
||
# that much is already free; under real pressure it
|
||
# partially unloads model weights instead of OOMing.
|
||
floor = int(extra * 1.25) + (768 << 20)
|
||
if model_management.get_free_memory(dev) < floor:
|
||
model_management.free_memory(floor, dev)
|
||
except Exception:
|
||
logger.warning(
|
||
"EasyControl: cond-cache VRAM reservation failed", exc_info=True
|
||
)
|
||
return out
|
||
|
||
return prepare_sampling_wrapper
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Pre-hook: lazy cond-KV prefill + live-block refresh
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_prefill_pre_hook(state: EasyControlState):
|
||
def pre_hook(module, args):
|
||
# module = live diffusion_model; args[0] = x (model-space latent).
|
||
state.live_blocks = module.blocks
|
||
x = args[0]
|
||
device, dtype = x.device, x.dtype
|
||
if state.cond_kv_cache is None or state._cache_device != device:
|
||
_build_cond_kv(module, state, device, dtype)
|
||
return None
|
||
|
||
return pre_hook
|
||
|
||
|
||
def _install_adaln_patches(model, diffusion_model, state: EasyControlState, strength):
|
||
"""Merge the target-stream adaln LoRA as ordinary ModelPatcher weight patches.
|
||
|
||
Each per-branch delta lands on ``adaln_modulation_{branch}.2.weight`` (the
|
||
256->3D up-projection — comfy's split-Sequential twin of anima_lora's
|
||
``adaln_up_{branch}``; both are bias-free, and the 256-dim bottleneck value
|
||
is identical between the fused and split layouts, so the weight merge is
|
||
exact). Going through ``add_patches`` keeps this comfy-native: it composes
|
||
with lowvram partial loading, block compile, and other LoRAs, and reverts on
|
||
unpatch. Patch strength = node strength, matching training's
|
||
``delta_scale = multiplier``. The cond stream must NOT see these patches —
|
||
``_cond_branch_modulation`` subtracts them during prefill.
|
||
"""
|
||
b0 = diffusion_model.blocks[0]
|
||
if not getattr(b0, "use_adaln_lora", False):
|
||
raise ValueError(
|
||
"EasyControl checkpoint carries target-stream adaln LoRA "
|
||
"(train_adaln) but the loaded DiT uses vanilla adaln MLPs "
|
||
"(use_adaln_lora=False). Wrong base model?"
|
||
)
|
||
probe = state.adaln["self_attn"][0]
|
||
seq = b0.adaln_modulation_self_attn
|
||
delta_shape = (probe.up.shape[0], probe.down.shape[1])
|
||
if len(seq) != 3 or tuple(seq[2].weight.shape) != delta_shape:
|
||
got = tuple(seq[2].weight.shape) if len(seq) == 3 else "not a bottleneck"
|
||
raise ValueError(
|
||
f"adaln shape mismatch: checkpoint delta is {delta_shape} but the "
|
||
f"DiT's adaln_modulation up-projection is {got}."
|
||
)
|
||
|
||
try: # current comfy: loras are WeightAdapterBase instances
|
||
from comfy.weight_adapter import LoRAAdapter
|
||
|
||
def make_patch(m):
|
||
return LoRAAdapter(
|
||
set(), (m.up, m.down, state.adaln_alpha, None, None, None)
|
||
)
|
||
except ImportError: # older comfy: ("lora", 6-tuple) patch format
|
||
|
||
def make_patch(m):
|
||
return ("lora", (m.up, m.down, state.adaln_alpha, None, None, None))
|
||
|
||
patches = {}
|
||
for branch in ("self_attn", "cross_attn", "mlp"):
|
||
for i, m in enumerate(state.adaln[branch]):
|
||
key = f"diffusion_model.blocks.{i}.adaln_modulation_{branch}.2.weight"
|
||
patches[key] = make_patch(m)
|
||
applied = model.add_patches(patches, strength_patch=float(strength))
|
||
if len(applied) != len(patches):
|
||
missing = sorted(set(patches) - set(applied))[:4]
|
||
raise ValueError(
|
||
f"adaln weight patches failed to attach ({len(applied)}/{len(patches)} "
|
||
f"applied; e.g. {missing}). ComfyUI key layout changed?"
|
||
)
|
||
logger.info(
|
||
f"EasyControl: merged target-stream adaln LoRA "
|
||
f"(r={state.adaln_rank}, alpha={state.adaln_alpha:g}, "
|
||
f"{len(patches)} patches, strength={float(strength):.3f})"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Public entry point — called by the node
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def install_easycontrol(model, weight_path, cond_latent, strength, cond_scale_override=None):
|
||
"""Install EasyControl onto a cloned ModelPatcher.
|
||
|
||
``cond_latent`` is the VAE-encoded reference in **model space** (the node
|
||
applies ``process_latent_in``), 4D ``[1, C, H, W]``.
|
||
"""
|
||
sd, meta = _load_weights_and_meta(weight_path)
|
||
state = EasyControlState(sd, meta, strength, cond_scale_override)
|
||
state.cond_latent = cond_latent
|
||
|
||
diffusion_model = model.get_model_object("diffusion_model")
|
||
nblocks = len(diffusion_model.blocks)
|
||
if nblocks != state.num_blocks:
|
||
raise ValueError(
|
||
f"EasyControl checkpoint is for {state.num_blocks} blocks but the "
|
||
f"loaded DiT has {nblocks}. Wrong base model?"
|
||
)
|
||
state.hidden_dim = int(diffusion_model.blocks[0].x_dim)
|
||
|
||
if state.train_adaln:
|
||
_install_adaln_patches(model, diffusion_model, state, strength)
|
||
|
||
# Reserve VRAM for the lazily-built cond KV cache before sampling starts
|
||
# (comfy's memory manager can't see it otherwise — image-swap OOM source).
|
||
try:
|
||
from comfy.patcher_extension import WrappersMP
|
||
|
||
model.add_wrapper_with_key(
|
||
WrappersMP.PREPARE_SAMPLING,
|
||
"easycontrol_cond_cache",
|
||
_make_prepare_sampling_wrapper(state),
|
||
)
|
||
except Exception:
|
||
logger.warning(
|
||
"EasyControl: could not register PREPARE_SAMPLING wrapper; "
|
||
"cond-cache VRAM will not be pre-reserved", exc_info=True
|
||
)
|
||
|
||
# Lazy-prefill pre-hook on diffusion_model (object-patched OrderedDict so it
|
||
# reverts on unpatch and composes with other pre-hooks).
|
||
pre_hook = _make_prefill_pre_hook(state)
|
||
new_pre = OrderedDict(diffusion_model._forward_pre_hooks)
|
||
new_pre[id(pre_hook)] = pre_hook
|
||
model.add_object_patch("diffusion_model._forward_pre_hooks", new_pre)
|
||
|
||
# Per-block self-attention replacement (object-patched → path-resolved to
|
||
# the live blocks; reverted on unpatch). Only self_attn is swapped; the rest
|
||
# of each block stays native.
|
||
for idx in range(state.num_blocks):
|
||
model.add_object_patch(
|
||
f"diffusion_model.blocks.{idx}.self_attn.forward",
|
||
_make_self_attn_forward(state, idx),
|
||
)
|
||
|
||
# Keep the state alive for the patcher's lifetime.
|
||
if not hasattr(model, "_easycontrol_states"):
|
||
model._easycontrol_states = []
|
||
model._easycontrol_states.append(state)
|
||
logger.info(
|
||
f"EasyControl: installed ({state.num_blocks} blocks, r={state.r}, "
|
||
f"ffn_lora={state.apply_ffn}, adaln={state.train_adaln}, "
|
||
f"eff_scale={state.eff_scale:.3f})"
|
||
)
|
||
return state
|